An initial design for an Audit.Polly extension for Audit.NET. Data Providers with resilience

Initial design Audit.Polly for Audit.NET

using Audit.Core;
using Polly;
using Polly.Fallback;

namespace test
{
    public static class FallbackActionArgumentsExtensions
    {
        public static async ValueTask<Outcome<object>> FallbackToDataProvider(this FallbackActionArguments<object> args, AuditDataProvider fallbackDataProvider)
        {
            var auditEvent = args.Context.Properties.GetValue<AuditEvent>(new ResiliencePropertyKey<AuditEvent>("AuditEvent"), null);

            if (auditEvent == null)
            {
                return await Outcome.FromResultAsValueTask(new object());
            }

            switch (args.Context.OperationKey)
            {
                case nameof(fallbackDataProvider.InsertEvent):
                {
                    var result = fallbackDataProvider.InsertEvent(auditEvent);

                    return await Outcome.FromResultAsValueTask(result);
                }
                case nameof(fallbackDataProvider.InsertEventAsync):
                {
                    var result = await fallbackDataProvider.InsertEventAsync(auditEvent, args.Context.CancellationToken);

                    return await Outcome.FromResultAsValueTask(result);
                }
                case nameof(fallbackDataProvider.ReplaceEvent):
                {
                    var eventId = args.Context.Properties.GetValue<object?>(new ResiliencePropertyKey<object?>("EventId"), null);

                    fallbackDataProvider.ReplaceEvent(eventId, auditEvent);

                    return await Outcome.FromResultAsValueTask(new object());
                }
                case nameof(fallbackDataProvider.ReplaceEventAsync):
                {
                    var eventId = args.Context.Properties.GetValue<object?>(new ResiliencePropertyKey<object?>("EventId"), null);
                
                    await fallbackDataProvider.ReplaceEventAsync(eventId, auditEvent, args.Context.CancellationToken);

                    return await Outcome.FromResultAsValueTask(new object());
                }
                default:
                    return await Outcome.FromResultAsValueTask(new object());
            }
        }
    }
    public class PollyDataProvider : AuditDataProvider
    {
        private AuditDataProvider _innerDataProvider;
        private ResiliencePipeline<object> _pipeline;

        public PollyDataProvider()
        {
        }

        public PollyDataProvider(ResiliencePipeline<object> resiliencePipeline, AuditDataProvider dataProvider)
        {
            _pipeline = resiliencePipeline;
            _innerDataProvider = dataProvider;
        }

        /// <summary>
        /// Creates a new resilience context for Insert or Replace operations
        /// </summary>
        /// <param name="operationKey">The operation key</param>
        /// <param name="auditEvent">The Audit Event</param>
        /// <param name="eventId">The Event ID in case of Replace</param>
        /// <param name="cancellationToken">The cancellation token</param>
        protected virtual ResilienceContext CreateResilienceContext(string operationKey, AuditEvent auditEvent, object? eventId, CancellationToken cancellationToken)
        {
            var context = ResilienceContextPool.Shared.Get(operationKey, cancellationToken);

            context.Properties.Set(new ResiliencePropertyKey<object?>("EventId"), eventId);
            context.Properties.Set(new ResiliencePropertyKey<AuditEvent>("AuditEvent"), auditEvent);

            return context;
        }

        /// <inheritdoc />
        public override object InsertEvent(AuditEvent auditEvent)
        {
            var context = CreateResilienceContext(nameof(InsertEvent), auditEvent, null, new CancellationToken());
            
            return _pipeline.Execute<object>(ctx => _innerDataProvider.InsertEvent(auditEvent), context);
        }

        /// <inheritdoc />
        public override async Task<object> InsertEventAsync(AuditEvent auditEvent, CancellationToken cancellationToken = default)
        {
            var context = CreateResilienceContext(nameof(InsertEventAsync), auditEvent, null, cancellationToken);
            
            return await _pipeline.ExecuteAsync<object>(async ctx => await _innerDataProvider.InsertEventAsync(auditEvent, ctx.CancellationToken), context);
        }

        /// <inheritdoc />
        public override void ReplaceEvent(object eventId, AuditEvent auditEvent)
        {
            var context = CreateResilienceContext(nameof(ReplaceEvent), auditEvent, eventId, new CancellationToken());

            _pipeline.Execute<object>(ctx =>
            {
                _innerDataProvider.ReplaceEvent(eventId, auditEvent);

                return new object();
            }, context);
        }

        /// <inheritdoc />
        public override async Task ReplaceEventAsync(object eventId, AuditEvent auditEvent, CancellationToken cancellationToken = default)
        {
            var context = CreateResilienceContext(nameof(ReplaceEventAsync), auditEvent, eventId, cancellationToken);
            
            await _pipeline.ExecuteAsync<object>(async ctx =>
            {
                await _innerDataProvider.ReplaceEventAsync(eventId, auditEvent, ctx.CancellationToken);

                return new object();
            }, context);
        }

        /// <inheritdoc />
        public override object CloneValue<T>(T value, AuditEvent auditEvent)
        {
            return _innerDataProvider.CloneValue(value, auditEvent);
        }

        /// <inheritdoc />
        public override T GetEvent<T>(object eventId)
        {
            return _innerDataProvider.GetEvent<T>(eventId);
        }

        /// <inheritdoc />
        public override Task<T> GetEventAsync<T>(object eventId, CancellationToken cancellationToken = default)
        {
            return _innerDataProvider.GetEventAsync<T>(eventId, cancellationToken);
        }
    }
}