Run in-process retries in a fresh scope so saga concurrency conflicts resolve

src/Mocha/src/Mocha.EntityFrameworkCore/Assembly.cs

@@ -1,3 +1,4 @@
 using System.Runtime.CompilerServices;
 
 [assembly: InternalsVisibleTo("Mocha.Tests")]
+[assembly: InternalsVisibleTo("Mocha.EntityFrameworkCore.Postgres.Tests")]

src/Mocha/src/Mocha.EntityFrameworkCore/MessageBusHostBuilderExtensions.cs

@@ -1,6 +1,8 @@
 using Microsoft.EntityFrameworkCore;
 using Microsoft.EntityFrameworkCore.Infrastructure;
+using Microsoft.EntityFrameworkCore.Storage;
 using Microsoft.Extensions.DependencyInjection;
+using System.Runtime.ExceptionServices;
 
 namespace Mocha.EntityFrameworkCore;
 
@@ -64,7 +66,12 @@ public static class EntityFrameworkCoreBuilderExtensions
     /// <returns>The same <paramref name="builder"/> instance for chaining.</returns>
     public static IEntityFrameworkCoreBuilder UseResilience(this IEntityFrameworkCoreBuilder builder)
     {
-        builder.HostBuilder.ConfigureMessageBus(x => x.UseConsume(EntityFrameworkResilienceConsumeMiddleware.Create()));
+        builder.HostBuilder.ConfigureMessageBus(x =>
+            x.ConfigureFeature(f =>
+                f.Set(new ConsumerExecutionStrategyFeature
+                {
+                    Strategy = new EntityFrameworkConsumerExecutionStrategy(builder.ContextType)
+                })));
 
         return builder;
     }
@@ -89,55 +96,32 @@ internal sealed class EntityFrameworkConfigurationFeature
     public Type? ContextType { get; set; }
 }
 
-internal sealed class EntityFrameworkResilienceConsumeMiddleware(Type contextType)
+internal sealed class EntityFrameworkConsumerExecutionStrategy(Type contextType) : IConsumerExecutionStrategy
 {
-    public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate next)
+    public async ValueTask ExecuteAsync(IConsumeContext context, Func<CancellationToken, ValueTask> executeAttempt)
     {
-        var dbContext = (DbContext)context.Services.GetRequiredService(contextType);
+        await using var scope = context.Services.CreateAsyncScope();
+        var dbContext = (DbContext)scope.ServiceProvider.GetRequiredService(contextType);
         var strategy = dbContext.Database.CreateExecutionStrategy();
 
         if (!strategy.RetriesOnFailure)
         {
-            await next(context);
+            await executeAttempt(context.CancellationToken);
             return;
         }
 
-        var originalServices = context.Services;
-
-        await strategy.ExecuteAsync(async () =>
+        try
         {
-            await using var scope = originalServices.CreateAsyncScope();
-            context.Services = scope.ServiceProvider;
-
-            try
-            {
-                await next(context);
-            }
-            finally
-            {
-                context.Services = originalServices;
-            }
-        });
+            await strategy.ExecuteAsync(
+                executeAttemptToken => executeAttempt(executeAttemptToken).AsTask(),
+                context.CancellationToken);
+        }
+        catch (RetryLimitExceededException exception) when (exception.InnerException is not null)
+        {
+            ExceptionDispatchInfo.Capture(exception.InnerException).Throw();
+            throw;
+        }
     }
-
-    public static ConsumerMiddlewareConfiguration Create()
-        => new(
-            static (context, next) =>
-            {
-                var contextType =
-                    context
-                        .Services.GetRequiredService<IFeatureCollection>()
-                        .Get<EntityFrameworkConfigurationFeature>()
-                        ?.ContextType
-                    ?? throw new InvalidOperationException(
-                        "No Entity Framework Core DbContext type has been configured. "
-                        + "Call AddEntityFramework<TContext>() on the message bus host builder "
-                        + "before using UseResilience().");
-
-                var middleware = new EntityFrameworkResilienceConsumeMiddleware(contextType);
-                return ctx => middleware.InvokeAsync(ctx, next);
-            },
-            "EntityFrameworkResilience");
 }
 
 internal sealed class EntityFrameworkTransactionConsumeMiddleware(Type contextType)

src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs

@@ -20,6 +20,7 @@ namespace Mocha.Sagas.EfCore;
 internal sealed class DbContextSagaStore(DbContext context) : ISagaStore, IDisposable
 {
     private readonly object _lock = new();
+    private readonly List<JsonDocument> _documents = [];
     private PooledArrayWriter? _arrayWriter;
     private List<byte[]>? _buffers;
 
@@ -71,6 +72,7 @@ public async Task SaveAsync<T>(Saga saga, T state, CancellationToken cancellatio
             .FirstOrDefaultAsync(cancellationToken);
 
         var document = ToJsonDocument(saga, state);
+        _documents.Add(document);
 
         if (sagaState is null)
         {
@@ -86,10 +88,14 @@ public async Task SaveAsync<T>(Saga saga, T state, CancellationToken cancellatio
         }
         else
         {
+            // The tracked entity owns its document; the one being replaced is released here rather
+            // than on load, so the entity never holds a disposed document while it is tracked.
+            var previous = sagaState.State;
             sagaState.State = document;
             sagaState.UpdatedAt = DateTimeOffset.UtcNow;
             sagaState.Version = NewVersion();
             set.Entry(sagaState).Property(x => x.State).IsModified = true;
+            previous.Dispose();
         }
 
         await context.SaveChangesAsync(cancellationToken);
@@ -126,7 +132,9 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation
     /// <returns>The deserialized saga state, or <c>default</c> if no state is found for the given identifier.</returns>
     public async Task<T?> LoadAsync<T>(Saga saga, Guid id, CancellationToken cancellationToken)
     {
-        // as the state is scoped we load the whole saga state into memory for the concurrency check
+        // The state is loaded tracked so the save that follows checks the version observed here.
+        // Its document stays with the tracked entity and is released when a save replaces it or
+        // when this store is disposed.
         var sageState = await context
             .Set<SagaState>()
             .AsTracking()
@@ -138,14 +146,8 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation
             return default;
         }
 
-        try
-        {
-            return FromJsonDocument<T>(saga, document);
-        }
-        finally
-        {
-            document.Dispose();
-        }
+        _documents.Add(document);
+        return FromJsonDocument<T>(saga, document);
     }
 
     private JsonDocument ToJsonDocument(Saga saga, SagaStateBase state)
@@ -198,10 +200,16 @@ private static Guid NewVersion()
     }
 
     /// <summary>
-    /// Releases the pooled array writer and returns all rented byte buffers to the pool.
+    /// Releases the documents held by tracked saga states, the pooled array writer, and all rented byte buffers.
     /// </summary>
     public void Dispose()
     {
+        foreach (var document in _documents)
+        {
+            document.Dispose();
+        }
+
+        _documents.Clear();
         _arrayWriter?.Dispose();
         if (_buffers is not null)
         {

src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs

@@ -131,12 +131,18 @@ public void Set<TFeature>(TFeature? instance)
     /// Initializes the feature collection with the specified defaults.
     /// </summary>
     /// <param name="defaults">
-    /// The defaults for the feature collection.
+    /// The defaults for the feature collection. When specified, pooled features retained from a
+    /// previous use stay parked so that they do not shadow the defaults.
     /// </param>
     public void Initialize(IFeatureCollection? defaults = null)
     {
         _defaults = defaults;
 
+        if (defaults is not null)
+        {
+            return;
+        }
+
         foreach (var pooledFeature in _pooledFeatures)
         {
             _features.Add(pooledFeature.Key, pooledFeature.Value);
@@ -151,18 +157,36 @@ public void Initialize(IFeatureCollection? defaults = null)
     /// </summary>
     public void Reset()
     {
+        _defaults = null;
         foreach (var item in _features)
         {
             if (item.Value is IPooledFeature pooledFeature)
             {
-                _pooledFeatures.Add(item);
                 pooledFeature.Reset();
+
+                if (!IsParked(item.Key))
+                {
+                    _pooledFeatures.Add(item);
+                }
             }
         }
 
         _features.Clear();
     }
 
+    private bool IsParked(Type key)
+    {
+        foreach (var parked in _pooledFeatures)
+        {
+            if (parked.Key == key)
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     /// <inheritdoc />
     public IEnumerator<KeyValuePair<Type, object>> GetEnumerator()
     {

src/Mocha/src/Mocha/Abstractions/Context/IConsumeContext.cs

@@ -4,4 +4,15 @@ namespace Mocha;
 /// Represents the context available during message consumption, combining message metadata with
 /// execution capabilities.
 /// </summary>
-public interface IConsumeContext : IMessageContext, IExecutionContext;
+public interface IConsumeContext : IMessageContext, IExecutionContext
+{
+    /// <summary>
+    /// Creates an isolated copy of this context using the specified services.
+    /// </summary>
+    /// <param name="services">The service provider for the cloned execution.</param>
+    /// <returns>
+    /// A context of the same concrete type with copied message metadata and headers, and its own feature
+    /// collection that falls back to the features of this context.
+    /// </returns>
+    IConsumeContext Clone(IServiceProvider services);
+}

src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs

@@ -5,8 +5,8 @@ namespace Mocha;
 
 internal sealed class BatchConsumeContext<TMessage> : IBatchConsumeContext<TMessage>
 {
-    private readonly Headers _headers = new(0);
-    private readonly FeatureCollection _features = new();
+    private readonly Headers _headers;
+    private readonly FeatureCollection _features;
 
     public BatchConsumeContext(
         MessageBatch<TMessage> message,
@@ -16,6 +16,8 @@ public BatchConsumeContext(
         MessageType? itemMessageType,
         CancellationToken cancellationToken)
     {
+        _headers = new Headers(0);
+        _features = new FeatureCollection();
         Message = message;
         Services = services;
         Runtime = firstContext.Runtime;
@@ -27,6 +29,35 @@ public BatchConsumeContext(
         CancellationToken = cancellationToken;
     }
 
+    private BatchConsumeContext(BatchConsumeContext<TMessage> context, IServiceProvider services)
+    {
+        _headers = new Headers(context._headers);
+        _features = new FeatureCollection(context._features);
+        Message = context.Message;
+        BatchId = context.BatchId;
+        ItemMessageType = context.ItemMessageType;
+        Services = services;
+        Runtime = context.Runtime;
+        Transport = context.Transport;
+        Endpoint = context.Endpoint;
+        Host = context.Host;
+        MessageId = context.MessageId;
+        CorrelationId = context.CorrelationId;
+        ConversationId = context.ConversationId;
+        CausationId = context.CausationId;
+        SourceAddress = context.SourceAddress;
+        DestinationAddress = context.DestinationAddress;
+        ResponseAddress = context.ResponseAddress;
+        FaultAddress = context.FaultAddress;
+        ContentType = context.ContentType;
+        MessageType = context.MessageType;
+        SentAt = context.SentAt;
+        DeliverBy = context.DeliverBy;
+        DeliveryCount = context.DeliveryCount;
+        Envelope = context.Envelope is { } envelope ? new MessageEnvelope(envelope) : null;
+        CancellationToken = context.CancellationToken;
+    }
+
     public IMessageBatch<TMessage> Message { get; }
 
     public string BatchId { get; }
@@ -78,4 +109,7 @@ public BatchConsumeContext(
     public CancellationToken CancellationToken { get; set; }
 
     public IServiceProvider Services { get; set; }
+
+    public IConsumeContext Clone(IServiceProvider services)
+        => new BatchConsumeContext<TMessage>(this, services);
 }

src/Mocha/src/Mocha/Consumers/ConsumerExecutionStrategyFeature.cs

@@ -0,0 +1,12 @@
+namespace Mocha;
+
+/// <summary>
+/// A message bus feature that selects the <see cref="IConsumerExecutionStrategy"/> used to run consumer attempts.
+/// </summary>
+public sealed class ConsumerExecutionStrategyFeature
+{
+    /// <summary>
+    /// Gets or sets the strategy that executes consumer attempts.
+    /// </summary>
+    public required IConsumerExecutionStrategy Strategy { get; set; }
+}

src/Mocha/src/Mocha/Consumers/DirectConsumerExecutionStrategy.cs

@@ -0,0 +1,9 @@
+namespace Mocha;
+
+internal sealed class DirectConsumerExecutionStrategy : IConsumerExecutionStrategy
+{
+    public static DirectConsumerExecutionStrategy Instance { get; } = new();
+
+    public ValueTask ExecuteAsync(IConsumeContext context, Func<CancellationToken, ValueTask> executeAttempt)
+        => executeAttempt(context.CancellationToken);
+}

src/Mocha/src/Mocha/Consumers/IConsumerExecutionStrategy.cs

@@ -0,0 +1,15 @@
+namespace Mocha;
+
+/// <summary>
+/// Executes awaited, sequential consumer attempts, optionally repeating them according to an infrastructure policy.
+/// Implementations must not retain the callback after execution completes.
+/// </summary>
+public interface IConsumerExecutionStrategy
+{
+    /// <summary>
+    /// Executes the supplied consumer attempt.
+    /// </summary>
+    /// <param name="context">The consume context the attempts are created from.</param>
+    /// <param name="executeAttempt">The callback that creates and executes one consumer attempt.</param>
+    ValueTask ExecuteAsync(IConsumeContext context, Func<CancellationToken, ValueTask> executeAttempt);
+}

src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs

@@ -94,6 +94,8 @@ private async Task ProcessBatchAsync(MessageBatch<TEvent> batch, CancellationTok
         {
             _logger.DispatchingBatch(batch.Count, batch.CompletionMode);
 
+            // The batch has no receive scope of its own, so it gets one here, mirroring the scope
+            // ReceiveEndpoint creates per receive context. Each consumer attempt then runs in a child scope.
             await using var scope = _applicationServices.CreateAsyncScope();
 
             var batchContext = new BatchConsumeContext<TEvent>(
@@ -104,21 +106,7 @@ private async Task ProcessBatchAsync(MessageBatch<TEvent> batch, CancellationTok
                 _itemMessageType,
                 cancellationToken);
 
-            var consumerFeature = batchContext.Features.GetOrSet<ReceiveConsumerFeature>();
-            consumerFeature.CurrentConsumer = this;
-
-            var accessor = scope.ServiceProvider.GetRequiredService<ConsumeContextAccessor>();
-            var previousContext = accessor.Context;
-            accessor.Context = batchContext;
-
-            try
-            {
-                await Pipeline(batchContext);
-            }
-            finally
-            {
-                accessor.Context = previousContext;
-            }
+            await Pipeline(batchContext);
 
             foreach (var entry in batch.Entries)
             {

src/Mocha/src/Mocha/Context/ConsumeContext~1.cs

@@ -143,6 +143,9 @@ public IServiceProvider Services
         set => Inner.Services = value;
     }
 
+    public IConsumeContext Clone(IServiceProvider services)
+        => new ConsumeContext<TMessage>(Inner.Clone(services));
+
     public void Dispose()
     {
         _inner = null;