From 118a7131dfdaf23441a7cb5454aa27aee33d5e08 Mon Sep 17 00:00:00 2001 From: Alexander Bohle Date: Mon, 31 Aug 2026 13:39:18 +0100 Subject: [PATCH 1/7] Serve saga loads from the database instead of the change tracker When a saga write conflicts on the version column, the consume-tier retry re-invokes the pipeline inside the same DI scope, so every attempt shares one DbContext. DbContextSagaStore loads with a tracked query, and EF identity resolution then serves the instance the failed attempt left behind rather than the committed row: the original version never advances, and every retry re-issues the same doomed update. Raising the attempt count only multiplies the conflicts, so under contention a saga settles on the order of its own timeout and loses. The tracked instance is also mutated by the failed save, so a retry reads back the state the failed attempt tried to write rather than what the load observed. And because LoadAsync disposes the JsonDocument that is the tracked entity's State property, a second load in the same scope, as happens when a retry follows a failure before any save, throws ObjectDisposedException. Detach any tracked instance for the saga key before the load, so each attempt materializes the committed row fresh. The save path is deliberately untouched: its query returning the load-tracked instance is what makes the update check the version observed by this attempt's load, and re-reading at save time would let a transition computed from older state overwrite a concurrent write. A guard test pins that a write landing between load and save still conflicts. A delete retried after a conflict is covered by the same change, since the saga reloads before it deletes. These are the first tests this store has; the internals-visible-to entry for the Postgres test project mirrors how the sibling Postgres store is tested. Co-Authored-By: Claude Fable 5 (1M context) --- .../src/Mocha.EntityFrameworkCore/Assembly.cs | 1 + .../Sagas/DbContextSagaStore.cs | 17 ++ .../DbContextSagaStoreConcurrencyTests.cs | 186 ++++++++++++++++++ 3 files changed, 204 insertions(+) create mode 100644 src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs diff --git a/src/Mocha/src/Mocha.EntityFrameworkCore/Assembly.cs b/src/Mocha/src/Mocha.EntityFrameworkCore/Assembly.cs index e830a08b3a0..f4e32ba40b4 100644 --- a/src/Mocha/src/Mocha.EntityFrameworkCore/Assembly.cs +++ b/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")] diff --git a/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs b/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs index 3755cc1f694..a2fd7c3ed2c 100644 --- a/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs +++ b/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs @@ -126,6 +126,11 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation /// The deserialized saga state, or default if no state is found for the given identifier. public async Task LoadAsync(Saga saga, Guid id, CancellationToken cancellationToken) { + // An in-scope retry re-reads through the same store after a conflicted save. The instance + // the failed attempt left in the change tracker would otherwise be served back by identity + // resolution, so it is detached first and the committed row is materialized fresh. + DetachTrackedState(saga.Name, id); + // as the state is scoped we load the whole saga state into memory for the concurrency check var sageState = await context .Set() @@ -148,6 +153,18 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation } } + private void DetachTrackedState(string sagaName, Guid id) + { + foreach (var entry in context.ChangeTracker.Entries()) + { + if (entry.Entity.Id == id && entry.Entity.SagaName == sagaName) + { + entry.State = EntityState.Detached; + break; + } + } + } + private JsonDocument ToJsonDocument(Saga saga, SagaStateBase state) { lock (_lock) diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs new file mode 100644 index 00000000000..e6b9f4b721b --- /dev/null +++ b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs @@ -0,0 +1,186 @@ +using Microsoft.EntityFrameworkCore; +using Mocha.EntityFrameworkCore.Postgres.Tests.Helpers; +using Mocha.Sagas.EfCore; + +namespace Mocha.EntityFrameworkCore.Postgres.Tests; + +/// +/// Tests that resolves optimistic concurrency conflicts across +/// in-scope retries. A retry re-invokes load and save on the same scoped store, so a load must +/// observe the committed row rather than the instance a failed attempt left in the change tracker. +/// +public sealed class DbContextSagaStoreConcurrencyTests(PostgresFixture fixture) + : IClassFixture, IAsyncLifetime +{ + private readonly List _disposables = []; + + public ValueTask InitializeAsync() => ValueTask.CompletedTask; + + [Fact] + public async Task SaveAsync_Should_Succeed_When_RetriedAfterConcurrencyConflict() + { + // arrange + // saga consumer A loads its state, then a concurrent consumer commits first + var (saga, id, connectionString) = await SeedAsync("Initial"); + var (contextA, storeA) = await CreateStoreAsync(connectionString); + var (contextB, storeB) = await CreateStoreAsync(connectionString); + + var stateA = await storeA.LoadAsync(saga, id, CancellationToken.None); + var stateB = await storeB.LoadAsync(saga, id, CancellationToken.None); + stateB!.Data = "winner"; + await storeB.SaveAsync(saga, stateB, CancellationToken.None); + + // act + // the first save conflicts, then the retry re-runs load and save on the same store + stateA!.Data = "loser"; + await Assert.ThrowsAsync( + () => storeA.SaveAsync(saga, stateA, CancellationToken.None)); + + var retried = await storeA.LoadAsync(saga, id, CancellationToken.None); + retried!.Data = "retried"; + await storeA.SaveAsync(saga, retried, CancellationToken.None); + + // assert + var (_, verify) = await CreateStoreAsync(connectionString); + var committed = await verify.LoadAsync(saga, id, CancellationToken.None); + Assert.Equal("retried", committed!.Data); + _ = contextA; + _ = contextB; + } + + [Fact] + public async Task LoadAsync_Should_ReturnCommittedState_When_CalledAgainAfterConflictedSave() + { + // arrange + var (saga, id, connectionString) = await SeedAsync("Initial"); + var (_, storeA) = await CreateStoreAsync(connectionString); + var (_, storeB) = await CreateStoreAsync(connectionString); + + var stateA = await storeA.LoadAsync(saga, id, CancellationToken.None); + var stateB = await storeB.LoadAsync(saga, id, CancellationToken.None); + stateB!.Data = "winner"; + await storeB.SaveAsync(saga, stateB, CancellationToken.None); + + stateA!.Data = "loser"; + await Assert.ThrowsAsync( + () => storeA.SaveAsync(saga, stateA, CancellationToken.None)); + + // act + // the retry's load must observe the committed row, not the failed attempt's write + var retried = await storeA.LoadAsync(saga, id, CancellationToken.None); + + // assert + Assert.Equal("winner", retried!.Data); + } + + [Fact] + public async Task LoadAsync_Should_ReturnState_When_CalledTwiceOnTheSameStore() + { + // arrange + // a retry after a failure before any save re-runs the load on the same store + var (saga, id, connectionString) = await SeedAsync("Initial", data: "twice"); + var (_, store) = await CreateStoreAsync(connectionString); + + // act + var first = await store.LoadAsync(saga, id, CancellationToken.None); + var second = await store.LoadAsync(saga, id, CancellationToken.None); + + // assert + Assert.Equal("twice", first!.Data); + Assert.Equal("twice", second!.Data); + } + + [Fact] + public async Task SaveAsync_Should_Conflict_When_StateChangedAfterLoad() + { + // arrange + // the guard for the fix: a save must still check the version observed by its own load, + // so a write that lands between load and save is detected instead of overwritten + var (saga, id, connectionString) = await SeedAsync("Initial"); + var (_, storeA) = await CreateStoreAsync(connectionString); + var (_, storeB) = await CreateStoreAsync(connectionString); + + var stateA = await storeA.LoadAsync(saga, id, CancellationToken.None); + var stateB = await storeB.LoadAsync(saga, id, CancellationToken.None); + stateB!.Data = "winner"; + await storeB.SaveAsync(saga, stateB, CancellationToken.None); + + // act & assert + stateA!.Data = "loser"; + await Assert.ThrowsAsync( + () => storeA.SaveAsync(saga, stateA, CancellationToken.None)); + + // assert - the concurrent write survived + var (_, verify) = await CreateStoreAsync(connectionString); + var committed = await verify.LoadAsync(saga, id, CancellationToken.None); + Assert.Equal("winner", committed!.Data); + } + + [Fact] + public async Task DeleteAsync_Should_Succeed_When_RetriedAfterConcurrencyConflict() + { + // arrange + // a saga that reaches its final state on a retry deletes through the same store + var (saga, id, connectionString) = await SeedAsync("Initial"); + var (_, storeA) = await CreateStoreAsync(connectionString); + var (_, storeB) = await CreateStoreAsync(connectionString); + + var stateA = await storeA.LoadAsync(saga, id, CancellationToken.None); + var stateB = await storeB.LoadAsync(saga, id, CancellationToken.None); + stateB!.Data = "winner"; + await storeB.SaveAsync(saga, stateB, CancellationToken.None); + + stateA!.Data = "loser"; + await Assert.ThrowsAsync( + () => storeA.SaveAsync(saga, stateA, CancellationToken.None)); + + // act - the retry loads fresh state, reaches a final state, and deletes + _ = await storeA.LoadAsync(saga, id, CancellationToken.None); + await storeA.DeleteAsync(saga, id, CancellationToken.None); + + // assert + var (_, verify) = await CreateStoreAsync(connectionString); + Assert.Null(await verify.LoadAsync(saga, id, CancellationToken.None)); + } + + private async Task<(TestSaga Saga, Guid Id, string ConnectionString)> SeedAsync( + string state, + string? data = null) + { + var connectionString = await fixture.CreateDatabaseAsync(); + var saga = new TestSaga(); + var id = Guid.NewGuid(); + + var (context, store) = await CreateStoreAsync(connectionString, ensureCreated: true); + await store.SaveAsync(saga, new TestSagaState(id, state) { Data = data }, CancellationToken.None); + _ = context; + + return (saga, id, connectionString); + } + + private async Task<(TestDbContext Context, DbContextSagaStore Store)> CreateStoreAsync( + string connectionString, + bool ensureCreated = false) + { + var options = new DbContextOptionsBuilder().UseTestNpgsql(connectionString).Options; + var context = new TestDbContext(options); + + if (ensureCreated) + { + await context.Database.EnsureCreatedAsync(TestContext.Current.CancellationToken); + } + + var store = new DbContextSagaStore(context); + _disposables.Add(context); + + return (context, store); + } + + public async ValueTask DisposeAsync() + { + foreach (var disposable in _disposables) + { + await disposable.DisposeAsync(); + } + } +} From ec132e13223ee445486a16c0dbfb8a73d8bbe317 Mon Sep 17 00:00:00 2001 From: Alexander Bohle Date: Thu, 3 Sep 2026 18:23:29 +0200 Subject: [PATCH 2/7] Run each in-process retry in a fresh service scope The consume-tier retry re-invoked the handler with the same context, inside the single scope the receive endpoint creates per delivery. Every attempt therefore shared the failed attempt's scoped services, and re-ran against whatever it had left behind: a change tracker still holding the entities it failed to save, a unit of work half applied. A plain handler that adds a row, fails once, and saves on the retry inserted the row twice. A saga whose write conflicted on its version column re-issued the same stale update on every attempt, so the retries could not succeed at all. A retry now runs in its own scope, with the ambient consume context set there so correlation still propagates, and the delivery scope restored afterwards so a later consumer on the same delivery is unaffected. The first attempt keeps the delivery scope, so the common path allocates nothing extra. A retry thus starts from the same point a redelivery would, one tier down. Features that cache a scoped service must drop it at the scope boundary, or a retry would reuse an instance from the disposed scope. IScopeBoundFeature marks them; SagaFeature, which caches the saga store, is the one such feature. It is reset before the attempt and again after it. Co-Authored-By: Claude Fable 5.1 (1M context) --- .../Features/IScopeBoundFeature.cs | 13 ++ .../Consume/Retry/ConsumerRetryMiddleware.cs | 52 ++++++- .../Sagas/Execution/ISagaExecutionContext.cs | 8 +- .../ConsumerRetryScopeTests.cs | 99 +++++++++++++ .../Mocha.EntityFrameworkCore.Tests.csproj | 1 + .../Mocha.Sagas.Tests/SagaRetryScopeTests.cs | 117 +++++++++++++++ .../Behaviors/RetryTests.cs | 140 ++++++++++++++++++ 7 files changed, 427 insertions(+), 3 deletions(-) create mode 100644 src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs create mode 100644 src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs create mode 100644 src/Mocha/test/Mocha.Sagas.Tests/SagaRetryScopeTests.cs diff --git a/src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs b/src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs new file mode 100644 index 00000000000..0dcc5579a7f --- /dev/null +++ b/src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs @@ -0,0 +1,13 @@ +namespace Mocha.Features; + +/// +/// A feature that caches state resolved from the current service scope and must drop it when the +/// scope it was resolved from is replaced. +/// +public interface IScopeBoundFeature +{ + /// + /// Drops state resolved from the previous service scope so it is resolved again from the current one. + /// + void ResetScope(); +} diff --git a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs index 5a55bf4ac41..39639c30215 100644 --- a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs @@ -6,7 +6,8 @@ namespace Mocha; /// /// A consumer middleware that implements in-process retry with configurable backoff strategies -/// when transient failures occur. +/// when transient failures occur. Each retry runs in its own service scope, so scoped state the +/// failed attempt left behind does not carry into the next one. /// internal sealed class ConsumerRetryMiddleware(ImmutableArray exceptionPolicyRules) { @@ -28,11 +29,58 @@ public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate nex await RetryExecutor.ExecuteAsync( exceptionPolicyRules, (next, context, retryState), - static (s) => s.next(s.context), + static (s) => s.retryState.ImmediateRetryCount == 0 + ? s.next(s.context) + : RetryInFreshScopeAsync(s.next, s.context), static (s, attempts) => s.retryState.ImmediateRetryCount = attempts, context.CancellationToken); } + /// + /// Runs one retry attempt in a new service scope, restoring the delivery scope afterwards. + /// + /// + /// The first attempt runs in the scope of the delivery. A retry re-invokes the same handler, + /// which would otherwise observe whatever the failed attempt left in scoped services: a + /// change tracker still holding the entities it failed to save, a unit of work half applied. + /// A fresh scope gives the retry the same starting point a redelivery would have. + /// + private static async ValueTask RetryInFreshScopeAsync(ConsumerDelegate next, IConsumeContext context) + { + var deliveryServices = context.Services; + + await using var scope = deliveryServices.CreateAsyncScope(); + var accessor = scope.ServiceProvider.GetRequiredService(); + + context.Services = scope.ServiceProvider; + accessor.Context = context; + ResetScopeBoundFeatures(context); + + try + { + await next(context); + } + finally + { + // Features may now hold services of the scope being disposed; drop them so a later + // consumer on the same delivery resolves from the restored delivery scope. + ResetScopeBoundFeatures(context); + accessor.Context = null; + context.Services = deliveryServices; + } + } + + private static void ResetScopeBoundFeatures(IConsumeContext context) + { + foreach (var (_, feature) in context.Features) + { + if (feature is IScopeBoundFeature scopeBound) + { + scopeBound.ResetScope(); + } + } + } + public static ConsumerMiddlewareConfiguration Create() => new( static (context, next) => diff --git a/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs b/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs index 0640df5bdda..9dd82bdaf5d 100644 --- a/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs +++ b/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs @@ -5,7 +5,7 @@ namespace Mocha.Sagas; /// /// A pooled feature that provides access to the saga store during saga event processing. /// -public class SagaFeature : IPooledFeature +public class SagaFeature : IPooledFeature, IScopeBoundFeature { /// /// Gets or sets the saga store used for persisting saga state. @@ -23,6 +23,12 @@ public void Reset() { Store = null!; } + + /// + public void ResetScope() + { + Store = null!; + } } internal static class ConsumeContextSagaExtensions diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs new file mode 100644 index 00000000000..773f42c7514 --- /dev/null +++ b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs @@ -0,0 +1,99 @@ +using Microsoft.Data.Sqlite; +using Microsoft.EntityFrameworkCore; +using Microsoft.Extensions.DependencyInjection; +using Mocha.Transport.InMemory; + +namespace Mocha.EntityFrameworkCore.Tests; + +/// +/// Tests that an in-process retry does not re-run a handler against the +/// the failed attempt left behind, which would replay its pending changes on top of the retry's. +/// +public sealed class ConsumerRetryScopeTests +{ + private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); + + [Fact] + public async Task Retry_Should_InsertOnce_When_HandlerFailsAfterAddAndSucceedsOnRetry() + { + // arrange + await using var connection = new SqliteConnection("DataSource=:memory:"); + await connection.OpenAsync(TestContext.Current.CancellationToken); + var saved = new TaskCompletionSource(); + var attempts = new AttemptCounter(); + + var services = new ServiceCollection(); + services.AddDbContext(o => o.UseSqlite(connection)); + services.AddSingleton(saved); + services.AddSingleton(attempts); + var builder = services.AddMessageBus() + .AddResilience(p => p.On().Retry(3, TimeSpan.FromMilliseconds(1), RetryBackoffType.Constant)) + .AddEventHandler(); + builder.AddInMemory(); + + await using var provider = services.BuildServiceProvider(); + using (var setup = provider.CreateScope()) + { + await setup.ServiceProvider.GetRequiredService().Database + .EnsureCreatedAsync(TestContext.Current.CancellationToken); + } + + var runtime = (MessagingRuntime)provider.GetRequiredService(); + await runtime.StartAsync(CancellationToken.None); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new CreateRow("only-once"), CancellationToken.None); + await saved.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + + // assert - the retry saved exactly one row, not the failed attempt's row as well + using var verify = provider.CreateScope(); + var rows = await verify.ServiceProvider.GetRequiredService().Rows + .Select(r => r.Name) + .ToListAsync(TestContext.Current.CancellationToken); + + Assert.Equal(2, attempts.Value); + Assert.Equal(["only-once"], rows); + } + + public sealed class AttemptCounter + { + public int Value; + } + + public sealed record CreateRow(string Name); + + public sealed class Row + { + public int Id { get; set; } + + public string Name { get; set; } = ""; + } + + public sealed class RetryDbContext(DbContextOptions options) : DbContext(options) + { + public DbSet Rows => Set(); + } + + /// + /// Adds its row, throws on the first attempt before saving, and saves on the retry. + /// + public sealed class AddThenThrowOnceHandler(RetryDbContext db, AttemptCounter attempts, TaskCompletionSource saved) + : IEventHandler + { + public async ValueTask HandleAsync(CreateRow message, CancellationToken cancellationToken) + { + db.Rows.Add(new Row { Name = message.Name }); + + if (Interlocked.Increment(ref attempts.Value) == 1) + { + throw new InvalidOperationException("transient"); + } + + await db.SaveChangesAsync(cancellationToken); + saved.TrySetResult(); + } + } +} diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/Mocha.EntityFrameworkCore.Tests.csproj b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/Mocha.EntityFrameworkCore.Tests.csproj index 30f63514077..4f2b9f98071 100644 --- a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/Mocha.EntityFrameworkCore.Tests.csproj +++ b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/Mocha.EntityFrameworkCore.Tests.csproj @@ -12,6 +12,7 @@ + diff --git a/src/Mocha/test/Mocha.Sagas.Tests/SagaRetryScopeTests.cs b/src/Mocha/test/Mocha.Sagas.Tests/SagaRetryScopeTests.cs new file mode 100644 index 00000000000..47f58c958cd --- /dev/null +++ b/src/Mocha/test/Mocha.Sagas.Tests/SagaRetryScopeTests.cs @@ -0,0 +1,117 @@ +using System.Collections.Concurrent; +using Microsoft.Extensions.DependencyInjection; +using Mocha.Transport.InMemory; + +namespace Mocha.Sagas.Tests; + +/// +/// Tests that a saga retried in-process resolves its store from the retry's own scope rather than +/// reusing the store instance the failed attempt cached on the saga feature. +/// +public class SagaRetryScopeTests +{ + private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); + + [Fact] + public async Task Saga_Should_UseFreshStore_When_SaveIsRetried() + { + // arrange + var capture = new StoreCapture(); + var services = new ServiceCollection(); + services.AddSingleton(capture); + // registered first so the in-memory registration's TryAdd keeps this store + services.AddScoped(); + services.AddInMemorySagas(); + var builder = services.AddMessageBus() + .AddResilience(p => p.On().Retry(3, TimeSpan.FromMilliseconds(1), RetryBackoffType.Constant)) + .AddSaga(); + builder.AddInMemory(); + + await using var provider = services.BuildServiceProvider(); + var runtime = (MessagingRuntime)provider.GetRequiredService(); + await runtime.StartAsync(CancellationToken.None); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + var storage = provider.GetRequiredService(); + + // act + await bus.PublishAsync(new StartRetrySaga(), CancellationToken.None); + await capture.Saved.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + + // the save is staged in the saga transaction and reaches storage when the attempt commits + var deadline = DateTime.UtcNow + s_timeout; + while (storage.Count == 0 && DateTime.UtcNow < deadline) + { + await Task.Delay(25, TestContext.Current.CancellationToken); + } + + // assert - the failed save and the retried save came from two different scoped stores + Assert.Equal(2, capture.StoreIds.Count); + Assert.Equal(2, capture.StoreIds.Distinct().Count()); + Assert.Equal(1, storage.Count); + } + + public sealed class StoreCapture + { + public ConcurrentQueue StoreIds { get; } = new(); + + public TaskCompletionSource Saved { get; } = new(); + + public int Failures; + } + + /// + /// Wraps the in-memory store, records which instance each save came from, and fails the first. + /// + public sealed class ThrowOnceSagaStore(InMemorySagaStateStorage storage, StoreCapture capture) : ISagaStore + { + private readonly InMemorySagaStore _inner = new(storage); + private readonly Guid _id = Guid.NewGuid(); + + public Task StartTransactionAsync(CancellationToken cancellationToken) + => _inner.StartTransactionAsync(cancellationToken); + + public async Task SaveAsync(Saga saga, T state, CancellationToken cancellationToken) + where T : SagaStateBase + { + capture.StoreIds.Enqueue(_id); + + if (Interlocked.Increment(ref capture.Failures) == 1) + { + throw new InvalidOperationException("transient"); + } + + await _inner.SaveAsync(saga, state, cancellationToken); + capture.Saved.TrySetResult(); + } + + public Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellationToken) + => _inner.DeleteAsync(saga, id, cancellationToken); + + public Task LoadAsync(Saga saga, Guid id, CancellationToken cancellationToken) + => _inner.LoadAsync(saga, id, cancellationToken); + } + + public sealed class RetrySagaState : SagaStateBase; + + public sealed class StartRetrySaga; + + public sealed class EndRetrySaga; + + public sealed class RetrySaga : Saga + { + protected override void Configure(ISagaDescriptor descriptor) + { + descriptor + .Initially() + .OnEvent() + .StateFactory(_ => new RetrySagaState()) + .TransitionTo("Started"); + + descriptor.During("Started").OnEvent().TransitionTo("Ended"); + + descriptor.Finally("Ended"); + } + } +} diff --git a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs index 5cacef0b215..2a32d5e7bca 100644 --- a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs +++ b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs @@ -277,6 +277,73 @@ await counter.WaitForCountAsync(4, s_timeout), $"Expected 4 invocations (1 original + 3 default retries), but got {counter.Count}"); } + [Fact] + public async Task Retry_Should_UseFreshScope_When_HandlerIsRetried() + { + // arrange + // A retry re-runs the handler; it must not observe scoped state the failed attempt left + // behind, so each attempt has to resolve its own scoped service instance. + var counter = new RetryInvocationCounter(); + var capture = new ScopeCapture(); + await using var provider = await new ServiceCollection() + .AddSingleton(counter) + .AddSingleton(capture) + .AddScoped() + .AddMessageBus() + .AddResilience(p => + { + p.On() + .Retry(3, TimeSpan.FromMilliseconds(1), RetryBackoffType.Constant); + }) + .AddEventHandler() + .AddInMemory() + .BuildServiceProvider(); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new OrderCreated { OrderId = "ORD-SCOPE" }, CancellationToken.None); + + // assert - two attempts, each with its own scoped instance + Assert.True(await counter.WaitForCountAsync(2, s_timeout), "Handler was not retried"); + Assert.Equal(2, capture.ProbeIds.Distinct().Count()); + } + + [Fact] + public async Task Retry_Should_PropagateConversationId_When_HandlerPublishesOnRetry() + { + // arrange + // Correlation propagation reads the ambient consume context from the scope, so the fresh + // scope of a retry has to carry it too, or a message published from a retry loses its lineage. + var counter = new RetryInvocationCounter(); + var capture = new ScopeCapture(); + await using var provider = await new ServiceCollection() + .AddSingleton(counter) + .AddSingleton(capture) + .AddMessageBus() + .AddResilience(p => + { + p.On() + .Retry(3, TimeSpan.FromMilliseconds(1), RetryBackoffType.Constant); + }) + .AddConsumer() + .AddConsumer() + .AddInMemory() + .BuildServiceProvider(); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new OrderCreated { OrderId = "ORD-CONV" }, CancellationToken.None); + + // assert - the follow-up published from the retry belongs to the original conversation + await capture.FollowUpReceived.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.NotNull(capture.OriginalConversationId); + Assert.Equal(capture.OriginalConversationId, capture.FollowUpConversationId); + } + // ============================================================ // Test Helpers // ============================================================ @@ -334,6 +401,24 @@ public async Task WaitForCountAsync(int targetCount, TimeSpan timeout) } } + private sealed class ScopeProbe + { + public Guid Id { get; } = Guid.NewGuid(); + } + + private sealed class ScopeCapture + { + public ConcurrentBag ProbeIds { get; } = []; + + public string? OriginalConversationId { get; set; } + + public string? FollowUpConversationId { get; set; } + + public TaskCompletionSource FollowUpReceived { get; } = new(); + } + + private sealed record RetryFollowUp; + // ============================================================ // Test Handlers // ============================================================ @@ -420,4 +505,59 @@ public ValueTask ConsumeAsync(IConsumeContext context) throw new InvalidOperationException("Fail to trigger retry"); } } + /// + /// Records the scoped probe it was given on every attempt and throws on the first one. + /// + private sealed class ScopedThrowOnceHandler( + ScopeProbe probe, + ScopeCapture capture, + RetryInvocationCounter counter) + : IEventHandler + { + public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellationToken) + { + capture.ProbeIds.Add(probe.Id); + var invocation = counter.Count; + counter.Increment(); + + if (invocation == 0) + { + throw new InvalidOperationException("Transient failure"); + } + + return default; + } + } + + /// + /// Throws on the first attempt and publishes a follow-up from the retry. + /// + private sealed class PublishOnRetryConsumer(ScopeCapture capture, RetryInvocationCounter counter) + : IConsumer + { + public async ValueTask ConsumeAsync(IConsumeContext context) + { + var invocation = counter.Count; + counter.Increment(); + + if (invocation == 0) + { + throw new InvalidOperationException("Transient failure"); + } + + capture.OriginalConversationId = context.ConversationId; + var bus = context.Services.GetRequiredService(); + await bus.PublishAsync(new RetryFollowUp(), context.CancellationToken); + } + } + + private sealed class FollowUpSpy(ScopeCapture capture) : IConsumer + { + public ValueTask ConsumeAsync(IConsumeContext context) + { + capture.FollowUpConversationId = context.ConversationId; + capture.FollowUpReceived.TrySetResult(); + return default; + } + } } From 8a76a6acab5661d4152c766cadd23bc9903bedc5 Mon Sep 17 00:00:00 2001 From: Alexander Bohle Date: Thu, 3 Sep 2026 18:23:29 +0200 Subject: [PATCH 3/7] Keep the saga document with its tracked entity With retries running in their own scope, the store no longer needs to detach the instance a failed attempt left in the change tracker before loading; no framework path reads the same saga key twice through one DbContext anymore. The detach also cost a DetectChanges scan over the application's context on every load. What remained worth fixing is that LoadAsync disposed the JsonDocument owned by the entity it had just tracked, so any later read of that entity in the same scope failed on a disposed document. The document now stays with the tracked entity, and SaveAsync disposes the one it replaces, at the point the store owns the swap. A load that is never followed by a save keeps its document until the scope ends, a pool miss on the failure path rather than a leak. The tests that exercised a retry through the same store are removed with the behavior; that guarantee now lives in the pipeline and is tested there. The lost-update guard stays, since the save must keep checking the version its own load observed, as does the double-load test, which was the disposed-document failure on main. Co-Authored-By: Claude Fable 5.1 (1M context) --- .../Sagas/DbContextSagaStore.cs | 33 ++---- .../DbContextSagaStoreConcurrencyTests.cs | 106 ++++-------------- 2 files changed, 28 insertions(+), 111 deletions(-) diff --git a/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs b/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs index a2fd7c3ed2c..f19e8cc5818 100644 --- a/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs +++ b/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs @@ -86,10 +86,14 @@ public async Task SaveAsync(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,12 +130,8 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation /// The deserialized saga state, or default if no state is found for the given identifier. public async Task LoadAsync(Saga saga, Guid id, CancellationToken cancellationToken) { - // An in-scope retry re-reads through the same store after a conflicted save. The instance - // the failed attempt left in the change tracker would otherwise be served back by identity - // resolution, so it is detached first and the committed row is materialized fresh. - DetachTrackedState(saga.Name, id); - - // 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. var sageState = await context .Set() .AsTracking() @@ -143,26 +143,7 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation return default; } - try - { - return FromJsonDocument(saga, document); - } - finally - { - document.Dispose(); - } - } - - private void DetachTrackedState(string sagaName, Guid id) - { - foreach (var entry in context.ChangeTracker.Entries()) - { - if (entry.Entity.Id == id && entry.Entity.SagaName == sagaName) - { - entry.State = EntityState.Detached; - break; - } - } + return FromJsonDocument(saga, document); } private JsonDocument ToJsonDocument(Saga saga, SagaStateBase state) diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs index e6b9f4b721b..f1526ebd34c 100644 --- a/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs +++ b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs @@ -5,9 +5,9 @@ namespace Mocha.EntityFrameworkCore.Postgres.Tests; /// -/// Tests that resolves optimistic concurrency conflicts across -/// in-scope retries. A retry re-invokes load and save on the same scoped store, so a load must -/// observe the committed row rather than the instance a failed attempt left in the change tracker. +/// Tests the concurrency contract of : a save checks the version its +/// own load observed, and loading twice on one store stays safe because the tracked entity keeps +/// ownership of its document. A retry runs in its own scope and never re-reads through a used store. /// public sealed class DbContextSagaStoreConcurrencyTests(PostgresFixture fixture) : IClassFixture, IAsyncLifetime @@ -17,77 +17,40 @@ public sealed class DbContextSagaStoreConcurrencyTests(PostgresFixture fixture) public ValueTask InitializeAsync() => ValueTask.CompletedTask; [Fact] - public async Task SaveAsync_Should_Succeed_When_RetriedAfterConcurrencyConflict() - { - // arrange - // saga consumer A loads its state, then a concurrent consumer commits first - var (saga, id, connectionString) = await SeedAsync("Initial"); - var (contextA, storeA) = await CreateStoreAsync(connectionString); - var (contextB, storeB) = await CreateStoreAsync(connectionString); - - var stateA = await storeA.LoadAsync(saga, id, CancellationToken.None); - var stateB = await storeB.LoadAsync(saga, id, CancellationToken.None); - stateB!.Data = "winner"; - await storeB.SaveAsync(saga, stateB, CancellationToken.None); - - // act - // the first save conflicts, then the retry re-runs load and save on the same store - stateA!.Data = "loser"; - await Assert.ThrowsAsync( - () => storeA.SaveAsync(saga, stateA, CancellationToken.None)); - - var retried = await storeA.LoadAsync(saga, id, CancellationToken.None); - retried!.Data = "retried"; - await storeA.SaveAsync(saga, retried, CancellationToken.None); - - // assert - var (_, verify) = await CreateStoreAsync(connectionString); - var committed = await verify.LoadAsync(saga, id, CancellationToken.None); - Assert.Equal("retried", committed!.Data); - _ = contextA; - _ = contextB; - } - - [Fact] - public async Task LoadAsync_Should_ReturnCommittedState_When_CalledAgainAfterConflictedSave() + public async Task LoadAsync_Should_ReturnState_When_CalledTwiceOnTheSameStore() { // arrange - var (saga, id, connectionString) = await SeedAsync("Initial"); - var (_, storeA) = await CreateStoreAsync(connectionString); - var (_, storeB) = await CreateStoreAsync(connectionString); - - var stateA = await storeA.LoadAsync(saga, id, CancellationToken.None); - var stateB = await storeB.LoadAsync(saga, id, CancellationToken.None); - stateB!.Data = "winner"; - await storeB.SaveAsync(saga, stateB, CancellationToken.None); - - stateA!.Data = "loser"; - await Assert.ThrowsAsync( - () => storeA.SaveAsync(saga, stateA, CancellationToken.None)); + // the tracked entity keeps its document after a load, so a second load must not touch a + // disposed document + var (saga, id, connectionString) = await SeedAsync("Initial", data: "twice"); + var (_, store) = await CreateStoreAsync(connectionString); // act - // the retry's load must observe the committed row, not the failed attempt's write - var retried = await storeA.LoadAsync(saga, id, CancellationToken.None); + var first = await store.LoadAsync(saga, id, CancellationToken.None); + var second = await store.LoadAsync(saga, id, CancellationToken.None); // assert - Assert.Equal("winner", retried!.Data); + Assert.Equal("twice", first!.Data); + Assert.Equal("twice", second!.Data); } [Fact] - public async Task LoadAsync_Should_ReturnState_When_CalledTwiceOnTheSameStore() + public async Task LoadAsync_Should_ReturnSavedState_When_CalledAfterSaveOnTheSameStore() { // arrange - // a retry after a failure before any save re-runs the load on the same store - var (saga, id, connectionString) = await SeedAsync("Initial", data: "twice"); + // a save replaces and releases the loaded document; the entity must remain readable after + var (saga, id, connectionString) = await SeedAsync("Initial", data: "before"); var (_, store) = await CreateStoreAsync(connectionString); + var loaded = await store.LoadAsync(saga, id, CancellationToken.None); + loaded!.Data = "after"; + await store.SaveAsync(saga, loaded, CancellationToken.None); + // act - var first = await store.LoadAsync(saga, id, CancellationToken.None); - var second = await store.LoadAsync(saga, id, CancellationToken.None); + var reloaded = await store.LoadAsync(saga, id, CancellationToken.None); // assert - Assert.Equal("twice", first!.Data); - Assert.Equal("twice", second!.Data); + Assert.Equal("after", reloaded!.Data); } [Fact] @@ -116,33 +79,6 @@ await Assert.ThrowsAsync( Assert.Equal("winner", committed!.Data); } - [Fact] - public async Task DeleteAsync_Should_Succeed_When_RetriedAfterConcurrencyConflict() - { - // arrange - // a saga that reaches its final state on a retry deletes through the same store - var (saga, id, connectionString) = await SeedAsync("Initial"); - var (_, storeA) = await CreateStoreAsync(connectionString); - var (_, storeB) = await CreateStoreAsync(connectionString); - - var stateA = await storeA.LoadAsync(saga, id, CancellationToken.None); - var stateB = await storeB.LoadAsync(saga, id, CancellationToken.None); - stateB!.Data = "winner"; - await storeB.SaveAsync(saga, stateB, CancellationToken.None); - - stateA!.Data = "loser"; - await Assert.ThrowsAsync( - () => storeA.SaveAsync(saga, stateA, CancellationToken.None)); - - // act - the retry loads fresh state, reaches a final state, and deletes - _ = await storeA.LoadAsync(saga, id, CancellationToken.None); - await storeA.DeleteAsync(saga, id, CancellationToken.None); - - // assert - var (_, verify) = await CreateStoreAsync(connectionString); - Assert.Null(await verify.LoadAsync(saga, id, CancellationToken.None)); - } - private async Task<(TestSaga Saga, Guid Id, string ConnectionString)> SeedAsync( string state, string? data = null) From a37145c036a0b03e9f08be1bdbeb8b709f0f8a37 Mon Sep 17 00:00:00 2001 From: Pascal Senn Date: Sat, 5 Sep 2026 20:48:58 +0200 Subject: [PATCH 4/7] Isolate consumer retry attempts --- .../MessageBusHostBuilderExtensions.cs | 62 ++-- .../Features/IScopeBoundFeature.cs | 13 - .../Features/PooledFeatureCollection.cs | 1 + .../Abstractions/Context/IConsumeContext.cs | 13 +- .../Consumers/Batching/BatchConsumeContext.cs | 38 +- .../ConsumerExecutionStrategyFeature.cs | 12 + .../DirectConsumerExecutionStrategy.cs | 9 + .../Consumers/IConsumerExecutionStrategy.cs | 15 + .../Implementations/BatchConsumer.cs | 21 +- .../src/Mocha/Context/ConsumeContext~1.cs | 3 + .../Consume/Retry/ConsumerRetryMiddleware.cs | 103 +++--- .../src/Mocha/Middlewares/ReceiveContext.cs | 38 ++ .../Sagas/Execution/ISagaExecutionContext.cs | 8 +- .../ConsumerResilienceMiddlewareTests.cs | 342 ++++++++++++++++++ .../TestConsumeContext.cs | 29 ++ .../Consumers/Batching/BatchCollectorTests.cs | 3 + .../Consumers/Batching/ConsumeContextTests.cs | 12 + .../Context/ConsumeContextCloneTests.cs | 165 +++++++++ .../Inbox/ConsumeInboxMiddlewareTests.cs | 2 - .../ConsumerInstrumentationMiddlewareTests.cs | 3 + .../Behaviors/BatchingTests.cs | 119 ++++++ .../Behaviors/RetryTests.cs | 252 +++++++++++++ 22 files changed, 1136 insertions(+), 127 deletions(-) delete mode 100644 src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs create mode 100644 src/Mocha/src/Mocha/Consumers/ConsumerExecutionStrategyFeature.cs create mode 100644 src/Mocha/src/Mocha/Consumers/DirectConsumerExecutionStrategy.cs create mode 100644 src/Mocha/src/Mocha/Consumers/IConsumerExecutionStrategy.cs create mode 100644 src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs create mode 100644 src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs diff --git a/src/Mocha/src/Mocha.EntityFrameworkCore/MessageBusHostBuilderExtensions.cs b/src/Mocha/src/Mocha.EntityFrameworkCore/MessageBusHostBuilderExtensions.cs index a382a4773ad..a70dacc3836 100644 --- a/src/Mocha/src/Mocha.EntityFrameworkCore/MessageBusHostBuilderExtensions.cs +++ b/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 /// The same instance for chaining. 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 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() - .Get() - ?.ContextType - ?? throw new InvalidOperationException( - "No Entity Framework Core DbContext type has been configured. " - + "Call AddEntityFramework() 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) diff --git a/src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs b/src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs deleted file mode 100644 index 0dcc5579a7f..00000000000 --- a/src/Mocha/src/Mocha.Utilities/Features/IScopeBoundFeature.cs +++ /dev/null @@ -1,13 +0,0 @@ -namespace Mocha.Features; - -/// -/// A feature that caches state resolved from the current service scope and must drop it when the -/// scope it was resolved from is replaced. -/// -public interface IScopeBoundFeature -{ - /// - /// Drops state resolved from the previous service scope so it is resolved again from the current one. - /// - void ResetScope(); -} diff --git a/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs b/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs index 807b87a280a..e35da3d7649 100644 --- a/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs +++ b/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs @@ -151,6 +151,7 @@ public void Initialize(IFeatureCollection? defaults = null) /// public void Reset() { + _defaults = null; foreach (var item in _features) { if (item.Value is IPooledFeature pooledFeature) diff --git a/src/Mocha/src/Mocha/Abstractions/Context/IConsumeContext.cs b/src/Mocha/src/Mocha/Abstractions/Context/IConsumeContext.cs index e5bb1049356..e6b13fa5f40 100644 --- a/src/Mocha/src/Mocha/Abstractions/Context/IConsumeContext.cs +++ b/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. /// -public interface IConsumeContext : IMessageContext, IExecutionContext; +public interface IConsumeContext : IMessageContext, IExecutionContext +{ + /// + /// Creates an isolated copy of this context using the specified services. + /// + /// The service provider for the cloned execution. + /// + /// 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. + /// + IConsumeContext Clone(IServiceProvider services); +} diff --git a/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs b/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs index f9c987abdd9..a766a88977e 100644 --- a/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs +++ b/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs @@ -5,8 +5,8 @@ namespace Mocha; internal sealed class BatchConsumeContext : IBatchConsumeContext { - private readonly Headers _headers = new(0); - private readonly FeatureCollection _features = new(); + private readonly Headers _headers; + private readonly FeatureCollection _features; public BatchConsumeContext( MessageBatch 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 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 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(this, services); } diff --git a/src/Mocha/src/Mocha/Consumers/ConsumerExecutionStrategyFeature.cs b/src/Mocha/src/Mocha/Consumers/ConsumerExecutionStrategyFeature.cs new file mode 100644 index 00000000000..9be985b12fc --- /dev/null +++ b/src/Mocha/src/Mocha/Consumers/ConsumerExecutionStrategyFeature.cs @@ -0,0 +1,12 @@ +namespace Mocha; + +/// +/// A message bus feature that selects the used to run consumer attempts. +/// +public sealed class ConsumerExecutionStrategyFeature +{ + /// + /// Gets or sets the strategy that executes consumer attempts. + /// + public required IConsumerExecutionStrategy Strategy { get; set; } +} diff --git a/src/Mocha/src/Mocha/Consumers/DirectConsumerExecutionStrategy.cs b/src/Mocha/src/Mocha/Consumers/DirectConsumerExecutionStrategy.cs new file mode 100644 index 00000000000..ec3437b4a7f --- /dev/null +++ b/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 executeAttempt) + => executeAttempt(context.CancellationToken); +} diff --git a/src/Mocha/src/Mocha/Consumers/IConsumerExecutionStrategy.cs b/src/Mocha/src/Mocha/Consumers/IConsumerExecutionStrategy.cs new file mode 100644 index 00000000000..9b20dd65924 --- /dev/null +++ b/src/Mocha/src/Mocha/Consumers/IConsumerExecutionStrategy.cs @@ -0,0 +1,15 @@ +namespace Mocha; + +/// +/// Executes awaited, sequential consumer attempts, optionally repeating them according to an infrastructure policy. +/// Implementations must not retain the callback after execution completes. +/// +public interface IConsumerExecutionStrategy +{ + /// + /// Executes the supplied consumer attempt. + /// + /// The consume context the attempts are created from. + /// The callback that creates and executes one consumer attempt. + ValueTask ExecuteAsync(IConsumeContext context, Func executeAttempt); +} diff --git a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs index cb7bfaec94b..8a254596bf4 100644 --- a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs +++ b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs @@ -93,32 +93,15 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok try { _logger.DispatchingBatch(batch.Count, batch.CompletionMode); - - await using var scope = _applicationServices.CreateAsyncScope(); - var batchContext = new BatchConsumeContext( batch, - scope.ServiceProvider, + _applicationServices, batch.GetContext(0), Guid.NewGuid().ToString(), _itemMessageType, cancellationToken); - var consumerFeature = batchContext.Features.GetOrSet(); - consumerFeature.CurrentConsumer = this; - - var accessor = scope.ServiceProvider.GetRequiredService(); - var previousContext = accessor.Context; - accessor.Context = batchContext; - - try - { - await Pipeline(batchContext); - } - finally - { - accessor.Context = previousContext; - } + await Pipeline(batchContext); foreach (var entry in batch.Entries) { diff --git a/src/Mocha/src/Mocha/Context/ConsumeContext~1.cs b/src/Mocha/src/Mocha/Context/ConsumeContext~1.cs index 8549247f7b8..0255a931405 100644 --- a/src/Mocha/src/Mocha/Context/ConsumeContext~1.cs +++ b/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(Inner.Clone(services)); + public void Dispose() { _inner = null; diff --git a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs index 39639c30215..3db1afd391d 100644 --- a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs @@ -1,15 +1,19 @@ using System.Collections.Immutable; using Microsoft.Extensions.DependencyInjection; using Mocha.Features; +using Mocha.Middlewares; namespace Mocha; /// /// A consumer middleware that implements in-process retry with configurable backoff strategies -/// when transient failures occur. Each retry runs in its own service scope, so scoped state the -/// failed attempt left behind does not carry into the next one. +/// when transient failures occur. Each execution uses a fresh service scope and context, +/// while retry state is shared across the operation. /// -internal sealed class ConsumerRetryMiddleware(ImmutableArray exceptionPolicyRules) +internal sealed class ConsumerRetryMiddleware( + ImmutableArray exceptionPolicyRules, + Consumer consumer, + IConsumerExecutionStrategy executionStrategy) { public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate next) { @@ -28,55 +32,59 @@ public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate nex await RetryExecutor.ExecuteAsync( exceptionPolicyRules, - (next, context, retryState), - static (s) => s.retryState.ImmediateRetryCount == 0 - ? s.next(s.context) - : RetryInFreshScopeAsync(s.next, s.context), - static (s, attempts) => s.retryState.ImmediateRetryCount = attempts, + (next, context, retryState, consumer, executionStrategy), + static (s) => s.executionStrategy.ExecuteAsync( + s.context, + ct => ExecuteAttemptAsync(s.next, s.context, s.retryState, s.consumer, ct)), + onRetry: null, context.CancellationToken); } /// - /// Runs one retry attempt in a new service scope, restoring the delivery scope afterwards. + /// Executes one consumer attempt on a clone of the context in a new service scope. /// - /// - /// The first attempt runs in the scope of the delivery. A retry re-invokes the same handler, - /// which would otherwise observe whatever the failed attempt left in scoped services: a - /// change tracker still holding the entities it failed to save, a unit of work half applied. - /// A fresh scope gives the retry the same starting point a redelivery would have. - /// - private static async ValueTask RetryInFreshScopeAsync(ConsumerDelegate next, IConsumeContext context) + private static async ValueTask ExecuteAttemptAsync( + ConsumerDelegate next, + IConsumeContext context, + RetryFeature retryState, + Consumer consumer, + CancellationToken cancellationToken) { - var deliveryServices = context.Services; - - await using var scope = deliveryServices.CreateAsyncScope(); + await using var scope = context.Services.CreateAsyncScope(); var accessor = scope.ServiceProvider.GetRequiredService(); - - context.Services = scope.ServiceProvider; - accessor.Context = context; - ResetScopeBoundFeatures(context); + var pool = scope.ServiceProvider.GetRequiredService().ReceiveContext; + ReceiveContext? pooledContext = null; + var attempt = context is ReceiveContext receiveContext + ? pooledContext = receiveContext.CopyTo(pool.Get(), scope.ServiceProvider) + : context.Clone(scope.ServiceProvider); try { - await next(context); + // The attempt reads the delivery's features through its clone, but needs its own + // consumer feature: a pooled clone may still carry a reset one from its previous use + // that would shadow the delivery's, and CurrentConsumer must be set for this attempt + // without touching the delivery's consumer set. Set() initializes pooled features, + // which clears CurrentConsumer, so it is assigned after the feature is added. + var consumerFeature = new ReceiveConsumerFeature(); + attempt.Features.Set(consumerFeature); + consumerFeature.CurrentConsumer = consumer; + attempt.CancellationToken = cancellationToken; + accessor.Context = attempt; + await next(attempt); + } + catch + { + // Counts every failed attempt, including the ones the execution strategy repeats. + retryState.ImmediateRetryCount++; + throw; } finally { - // Features may now hold services of the scope being disposed; drop them so a later - // consumer on the same delivery resolves from the restored delivery scope. - ResetScopeBoundFeatures(context); accessor.Context = null; - context.Services = deliveryServices; - } - } - private static void ResetScopeBoundFeatures(IConsumeContext context) - { - foreach (var (_, feature) in context.Features) - { - if (feature is IScopeBoundFeature scopeBound) + if (pooledContext is not null) { - scopeBound.ResetScope(); + pool.Return(pooledContext); } } } @@ -86,14 +94,10 @@ public static ConsumerMiddlewareConfiguration Create() static (context, next) => { var feature = context.GetExceptionPolicyFeature(); - - if (feature is null) - { - // No exception policy configured - skip retry middleware entirely. - return next; - } - - var middleware = new ConsumerRetryMiddleware(feature.Rules.ToImmutableArray()); + var middleware = new ConsumerRetryMiddleware( + feature?.Rules.ToImmutableArray() ?? ImmutableArray.Empty, + context.Consumer, + context.GetConsumerExecutionStrategy()); return ctx => middleware.InvokeAsync(ctx, next); }, @@ -116,4 +120,15 @@ file static class Extensions return null; } + + /// + /// Resolves the bus-level consumer execution strategy, falling back to direct execution. + /// + public static IConsumerExecutionStrategy GetConsumerExecutionStrategy(this ConsumerMiddlewareFactoryContext context) + { + var busFeatures = context.Services.GetRequiredService(); + + return busFeatures.Get()?.Strategy + ?? DirectConsumerExecutionStrategy.Instance; + } } diff --git a/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs b/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs index 7f9a6e1921f..ce4b2f69036 100644 --- a/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs +++ b/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs @@ -148,6 +148,44 @@ public ReceiveContext() /// public CancellationToken CancellationToken { get; set; } + /// + public IConsumeContext Clone(IServiceProvider services) + => CopyTo(new ReceiveContext(), services); + + /// + /// Copies this context into using the specified services. + /// + /// The target context, typically taken from a pool. + /// The service provider for the cloned execution. + /// The instance. + public ReceiveContext CopyTo(ReceiveContext clone, IServiceProvider services) + { + clone.Services = services; + clone.Runtime = Runtime; + clone.Transport = Transport; + clone.Endpoint = Endpoint; + clone.Envelope = Envelope is { } envelope ? new MessageEnvelope(envelope) : null; + clone.ContentType = ContentType; + clone.MessageType = MessageType; + clone.MessageId = MessageId; + clone.CorrelationId = CorrelationId; + clone.ConversationId = ConversationId; + clone.CausationId = CausationId; + clone.SourceAddress = SourceAddress; + clone.DestinationAddress = DestinationAddress; + clone.ResponseAddress = ResponseAddress; + clone.FaultAddress = FaultAddress; + clone.SentAt = SentAt; + clone.DeliverBy = DeliverBy; + clone.DeliveryCount = DeliveryCount; + clone.Body = Body; + clone.Host = Host; + clone.CancellationToken = CancellationToken; + clone._features.Initialize(_features); + clone.Headers.AddRange(Headers); + return clone; + } + /// /// Resets all properties, headers, features, and the body to their default state /// so the instance can be returned to the object pool. diff --git a/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs b/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs index 9dd82bdaf5d..0640df5bdda 100644 --- a/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs +++ b/src/Mocha/src/Mocha/Sagas/Execution/ISagaExecutionContext.cs @@ -5,7 +5,7 @@ namespace Mocha.Sagas; /// /// A pooled feature that provides access to the saga store during saga event processing. /// -public class SagaFeature : IPooledFeature, IScopeBoundFeature +public class SagaFeature : IPooledFeature { /// /// Gets or sets the saga store used for persisting saga state. @@ -23,12 +23,6 @@ public void Reset() { Store = null!; } - - /// - public void ResetScope() - { - Store = null!; - } } internal static class ConsumeContextSagaExtensions diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs new file mode 100644 index 00000000000..5a035f8ce89 --- /dev/null +++ b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs @@ -0,0 +1,342 @@ +using Microsoft.EntityFrameworkCore; +using Microsoft.EntityFrameworkCore.Infrastructure; +using Microsoft.EntityFrameworkCore.Storage; +using Microsoft.Extensions.DependencyInjection; +using Mocha.Features; +using Mocha.Transport.InMemory; + +namespace Mocha.EntityFrameworkCore.Tests; + +public sealed class ConsumerResilienceMiddlewareTests +{ + private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); + + [Fact] + public async Task Resilience_Should_UseFreshCompleteAttemptEnvironment_When_StrategyRetries() + { + // arrange + var capture = new ResilienceCapture(); + capture.FailuresBeforeSuccess = 3; + var services = new ServiceCollection(); + services.AddSingleton(capture); + services.AddScoped(); + services.AddDbContext(options => + options + .UseSqlite("Data Source=:memory:") + .ReplaceService()); + var builder = services.AddMessageBus() + .AddEntityFramework(p => p.UseResilience()) + .AddResilience(p => p.On().Retry(1)) + .AddConsumer(); + builder.ConfigureMessageBus(b => b.UseReceive(new ReceiveMiddlewareConfiguration( + (_, next) => async context => + { + context.Features.Set(capture.ReceiveFeature); + await next(context); + }, + "ReceiveFeature"))); + builder.AddInMemory(); + + await using var provider = services.BuildServiceProvider(); + var runtime = (MessagingRuntime)provider.GetRequiredService(); + await runtime.StartAsync(CancellationToken.None); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync( + new ResilienceMessage { Value = "original" }, + CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + await capture.Disposed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.Equal( + (4, 4, 4), + ( + capture.DbContextIds.Distinct().Count(), + capture.ProbeIds.Distinct().Count(), + capture.DisposedProbeIds.Distinct().Count())); + Assert.Equal( + [(true, false, true), (true, false, true), (true, false, true), (true, false, true)], + capture.AccessorWasValid + .Zip(capture.FeatureWasPresent, capture.ReceiveFeatureWasShared) + .Select(static values => (values.First, values.Second, values.Third))); + Assert.Equal([0, 1, 2, 3], capture.ImmediateRetryCounts); + Assert.Equal([true, true, true, true], capture.ConsumerWasPresent); + } + + [Fact] + public async Task Resilience_Should_PreserveTypedBatchContextAndFreshScope_When_StrategyRetries() + { + // arrange + var capture = new ResilienceCapture(); + var services = new ServiceCollection(); + services.AddSingleton(capture); + services.AddScoped(); + services.AddDbContext(options => + options + .UseSqlite("Data Source=:memory:") + .ReplaceService()); + var builder = services.AddMessageBus() + .AddEntityFramework(p => p.UseResilience()) + .AddBatchHandler(options => options.MaxBatchSize = 1); + builder.AddInMemory(); + + await using var provider = services.BuildServiceProvider(); + var runtime = (MessagingRuntime)provider.GetRequiredService(); + await runtime.StartAsync(CancellationToken.None); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync( + new ResilienceMessage { Value = "original" }, + CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + await capture.Disposed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.Equal( + (2, 2), + (capture.DbContextIds.Distinct().Count(), capture.ProbeIds.Distinct().Count())); + Assert.Equal( + [(true, false, 0, true), (true, false, 1, true)], + capture.AccessorWasValid + .Select((value, index) => + ( + value, + capture.FeatureWasPresent[index], + capture.ImmediateRetryCounts[index], + capture.ConsumerWasPresent[index]))); + } + + [Theory] + [InlineData("EntityFrameworkTransaction", "Inbox", "Custom")] + [InlineData("Inbox", "Custom", "EntityFrameworkTransaction")] + [InlineData("EntityFrameworkTransaction", "Custom", "Inbox")] + [InlineData("Custom", "Inbox", "EntityFrameworkTransaction")] + [InlineData("Inbox", "EntityFrameworkTransaction", "Custom")] + [InlineData("Custom", "EntityFrameworkTransaction", "Inbox")] + public async Task Pipeline_Should_PreserveRegistrationOrder_When_UsingPersistenceMiddlewareKeys( + string first, + string second, + string third) + { + // arrange + var capture = new MiddlewareOrderCapture(); + var services = new ServiceCollection(); + services.AddSingleton(capture); + var builder = services.AddMessageBus() + .AddEventHandler(); + builder.ConfigureMessageBus(b => + { + b.UseConsume(CreateOrderMiddleware(first)); + b.UseConsume(CreateOrderMiddleware(second)); + b.UseConsume(CreateOrderMiddleware(third)); + }); + builder.AddInMemory(); + + await using var provider = services.BuildServiceProvider(); + var runtime = (MessagingRuntime)provider.GetRequiredService(); + await runtime.StartAsync(CancellationToken.None); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new MiddlewareOrderMessage(), CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.Equal( + [first, second, third, "Handler"], + capture.Entries); + } + + private static ConsumerMiddlewareConfiguration CreateOrderMiddleware(string key) + => new( + (_, next) => context => InvokeOrderMiddlewareAsync(key, context, next), + key); + + private static async ValueTask InvokeOrderMiddlewareAsync( + string key, + IConsumeContext context, + ConsumerDelegate next) + { + context.Services.GetRequiredService().Entries.Add(key); + await next(context); + } + + public sealed class ResilienceDbContext(DbContextOptions options) + : DbContext(options); + + public sealed class RetryExecutionStrategyFactory(ExecutionStrategyDependencies dependencies) + : IExecutionStrategyFactory + { + public IExecutionStrategy Create() => new RetryExecutionStrategy(dependencies); + } + + private sealed class RetryExecutionStrategy(ExecutionStrategyDependencies dependencies) + : ExecutionStrategy(dependencies, 1, TimeSpan.Zero) + { + protected override bool ShouldRetryOn(Exception exception) + => exception is RetryExecutionException; + } + + public sealed class ResilienceConsumer( + ResilienceDbContext dbContext, + ResilienceProbe probe, + ConsumeContextAccessor accessor, + ResilienceCapture capture) + : IConsumer + { + public ValueTask ConsumeAsync(IConsumeContext context) + { + var receiveFeature = context.Features.Get(); + capture.ReceiveFeatureWasShared.Add( + ReferenceEquals(receiveFeature, capture.ReceiveFeature) + && receiveFeature.Value == capture.Attempts); + receiveFeature!.Value++; + capture.DbContextIds.Add(dbContext.ContextId.InstanceId); + capture.ProbeIds.Add(probe.Id); + capture.AccessorWasValid.Add( + accessor.Context?.Services == context.Services + && ReferenceEquals( + context.Services.GetRequiredService(), + dbContext)); + capture.ImmediateRetryCounts.Add(context.Features.Get()!.ImmediateRetryCount); + capture.ConsumerWasPresent.Add( + context.Features.Get()?.CurrentConsumer is not null); + capture.FeatureWasPresent.Add( + context.Features.Get() is not null); + context.Features.Set(new ResilienceAttemptFeature()); + + if (Interlocked.Increment(ref capture.Attempts) <= capture.FailuresBeforeSuccess) + { + throw new RetryExecutionException(); + } + + capture.Completed.TrySetResult(); + return default; + } + } + + public sealed class ResilienceProbe(ResilienceCapture capture) : IDisposable + { + public Guid Id { get; } = Guid.NewGuid(); + + public void Dispose() + { + capture.DisposedProbeIds.Add(Id); + + if (capture.DisposedProbeIds.Count == capture.FailuresBeforeSuccess + 1) + { + capture.Disposed.TrySetResult(); + } + } + } + + public sealed class BatchResilienceHandler( + ResilienceDbContext dbContext, + ResilienceProbe probe, + ConsumeContextAccessor accessor, + ResilienceCapture capture) + : IBatchEventHandler + { + public ValueTask HandleAsync( + IMessageBatch batch, + CancellationToken cancellationToken) + { + var context = accessor.Context!; + capture.DbContextIds.Add(dbContext.ContextId.InstanceId); + capture.ProbeIds.Add(probe.Id); + capture.AccessorWasValid.Add( + context is IBatchConsumeContext + && ReferenceEquals( + context.Services.GetRequiredService(), + dbContext)); + capture.ImmediateRetryCounts.Add(context.Features.Get()!.ImmediateRetryCount); + capture.ConsumerWasPresent.Add( + context.Features.Get()?.CurrentConsumer is not null); + capture.FeatureWasPresent.Add( + context.Features.Get() is not null); + context.Features.Set(new ResilienceAttemptFeature()); + + if (Interlocked.Increment(ref capture.Attempts) <= capture.FailuresBeforeSuccess) + { + throw new RetryExecutionException(); + } + + capture.Completed.TrySetResult(); + return default; + } + } + + public sealed class ResilienceCapture + { + public SharedReceiveFeature ReceiveFeature { get; } = new(); + + public List ReceiveFeatureWasShared { get; } = []; + + public List DbContextIds { get; } = []; + + public List ProbeIds { get; } = []; + + public List DisposedProbeIds { get; } = []; + + public List AccessorWasValid { get; } = []; + + public List FeatureWasPresent { get; } = []; + + public List ImmediateRetryCounts { get; } = []; + + public List ConsumerWasPresent { get; } = []; + + public TaskCompletionSource Completed { get; } = new(); + + public TaskCompletionSource Disposed { get; } = new(); + + public int Attempts; + + public int FailuresBeforeSuccess { get; set; } = 1; + } + + public sealed class MiddlewareOrderHandler(MiddlewareOrderCapture capture) + : IEventHandler + { + public ValueTask HandleAsync( + MiddlewareOrderMessage message, + CancellationToken cancellationToken) + { + capture.Entries.Add("Handler"); + capture.Completed.TrySetResult(); + return default; + } + } + + public sealed class MiddlewareOrderCapture + { + public List Entries { get; } = []; + + public TaskCompletionSource Completed { get; } = new(); + } + + public sealed class ResilienceMessage + { + public required string Value { get; set; } + } + + public sealed class MiddlewareOrderMessage; + + private sealed class ResilienceAttemptFeature; + + public sealed class SharedReceiveFeature + { + public int Value { get; set; } + } + + private sealed class RetryExecutionException : Exception; +} diff --git a/src/Mocha/test/Mocha.Sagas.TestHelpers/TestConsumeContext.cs b/src/Mocha/test/Mocha.Sagas.TestHelpers/TestConsumeContext.cs index 6f156d1bc60..93784ac3d2c 100644 --- a/src/Mocha/test/Mocha.Sagas.TestHelpers/TestConsumeContext.cs +++ b/src/Mocha/test/Mocha.Sagas.TestHelpers/TestConsumeContext.cs @@ -43,4 +43,33 @@ public sealed class TestConsumeContext : IConsumeContext /// Provides mutable access to headers for test setup. /// public Headers MutableHeaders => _headers; + + public IConsumeContext Clone(IServiceProvider services) + { + var clone = new TestConsumeContext + { + Runtime = Runtime, + CancellationToken = CancellationToken, + Services = services, + Transport = Transport, + Endpoint = Endpoint, + MessageId = MessageId, + CorrelationId = CorrelationId, + ConversationId = ConversationId, + CausationId = CausationId, + SourceAddress = SourceAddress, + DestinationAddress = DestinationAddress, + ResponseAddress = ResponseAddress, + FaultAddress = FaultAddress, + ContentType = ContentType, + MessageType = MessageType, + SentAt = SentAt, + DeliverBy = DeliverBy, + DeliveryCount = DeliveryCount, + Envelope = Envelope is { } envelope ? new MessageEnvelope(envelope) : null, + Host = Host + }; + clone.MutableHeaders.AddRange(Headers); + return clone; + } } diff --git a/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs b/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs index f99ba324f5a..6af189dc9b4 100644 --- a/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs +++ b/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs @@ -262,6 +262,9 @@ private sealed class StubConsumeContext(TestEvent message, string? messageId = n public IMessagingRuntime Runtime { get; set; } = null!; public CancellationToken CancellationToken { get; set; } public IServiceProvider Services { get; set; } = null!; + + public IConsumeContext Clone(IServiceProvider services) + => new StubConsumeContext(message, MessageId) { Services = services }; } /// diff --git a/src/Mocha/test/Mocha.Tests/Consumers/Batching/ConsumeContextTests.cs b/src/Mocha/test/Mocha.Tests/Consumers/Batching/ConsumeContextTests.cs index 6675dd0c38d..0275d64ba8a 100644 --- a/src/Mocha/test/Mocha.Tests/Consumers/Batching/ConsumeContextTests.cs +++ b/src/Mocha/test/Mocha.Tests/Consumers/Batching/ConsumeContextTests.cs @@ -172,5 +172,17 @@ public void SetMessage(object message) public IMessagingRuntime Runtime { get; set; } = null!; public CancellationToken CancellationToken { get; set; } public IServiceProvider Services { get; set; } = null!; + + public IConsumeContext Clone(IServiceProvider services) + { + var clone = new StubConsumeContext { Services = services }; + + if (_features.Get() is { Message: { } message }) + { + clone.SetMessage(message); + } + + return clone; + } } } diff --git a/src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs b/src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs new file mode 100644 index 00000000000..f433dac262e --- /dev/null +++ b/src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs @@ -0,0 +1,165 @@ +using Microsoft.Extensions.DependencyInjection; +using Mocha.Features; +using Mocha.Middlewares; + +namespace Mocha.Tests.Context; + +public sealed class ConsumeContextCloneTests +{ + [Fact] + public void Clone_Should_PreserveCurrentMetadataAndShareReceiveFeatures_When_ReceiveContextIsCloned() + { + // arrange + using var services = new ServiceCollection().BuildServiceProvider(); + var envelope = new MessageEnvelope { MessageId = "envelope-message" }; + var context = new ReceiveContext + { + Services = null!, + MessageId = "current-message", + Envelope = envelope + }; + context.Headers.Set("x-test", "original"); + context.Features.Set(new CloneFeature()); + + // act + var clone = (ReceiveContext)context.Clone(services); + + // assert + Assert.Equal( + ("current-message", "envelope-message", true), + (clone.MessageId, clone.Envelope?.MessageId, ReferenceEquals(services, clone.Services))); + Assert.True(clone.Headers.TryGetValue("x-test", out var header) && Equals(header, "original")); + Assert.NotSame(envelope, clone.Envelope); + Assert.NotSame(context.Headers, clone.Headers); + Assert.Same(context.Features.Get(), clone.Features.Get()); + } + + [Fact] + public void Clone_Should_PreserveTypedContext_When_TypedContextIsCloned() + { + // arrange + using var services = new ServiceCollection().BuildServiceProvider(); + var context = new ConsumeContext(new ReceiveContext()); + + // act + var clone = (ConsumeContext)context.Clone(services); + + // assert + Assert.Same(services, clone.Services); + } + + [Fact] + public void Clone_Should_ReuseBatchAndShareInputFeatures_When_BatchContextIsCloned() + { + // arrange + using var services = new ServiceCollection().BuildServiceProvider(); + var firstContext = new ReceiveContext(); + var batch = new MessageBatch( + [new BufferedEntry(new ConsumeContext(firstContext))], + BatchCompletionMode.Size); + var context = new BatchConsumeContext( + batch, + null!, + firstContext, + "batch", + null, + CancellationToken.None); + context.Features.Set(new CloneFeature()); + + // act + var clone = (BatchConsumeContext)context.Clone(services); + + // assert + Assert.Same(batch, clone.Message); + Assert.Same(services, clone.Services); + Assert.Same(context.Features.Get(), clone.Features.Get()); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public void Clone_Should_KeepAttemptFeaturesOffTheInput_When_CloningAnExecution(bool batch) + { + // arrange + using var services = new ServiceCollection().BuildServiceProvider(); + var delivery = new ReceiveContext(); + IConsumeContext input = batch + ? new BatchConsumeContext( + new MessageBatch( + [new BufferedEntry(new ConsumeContext(delivery))], + BatchCompletionMode.Size), + services, + delivery, + "batch", + null, + CancellationToken.None) + : delivery; + var shared = new CloneFeature(); + input.Features.Set(shared); + + // act + var first = input.Clone(services); + first.Features.Get()!.Value = 42; + first.Features.Set(new ExecutionFeature()); + var sibling = input.Clone(services); + first.Features.Set(new CloneFeature { Value = 99 }); + + // assert + Assert.NotSame(input.Features, first.Features); + Assert.Equal(42, input.Features.Get()!.Value); + Assert.Null(input.Features.Get()); + Assert.Same(shared, sibling.Features.Get()); + Assert.Null(sibling.Features.Get()); + } + + [Fact] + public void Reset_Should_LeaveReceiveFeatureUntouched_When_CloneIsReset() + { + // arrange + using var services = new ServiceCollection().BuildServiceProvider(); + var delivery = new ReceiveContext(); + var shared = delivery.Features.GetOrSet(); + shared.MessageConsumed = true; + var clone = (ReceiveContext)delivery.Clone(services); + + // act + clone.Reset(); + + // assert + Assert.True(shared.MessageConsumed); + Assert.True(clone.Features.IsEmpty); + } + + [Fact] + public void Clone_Should_ReuseResetContext_When_ReturnedToPool() + { + // arrange + using var services = new ServiceCollection().BuildServiceProvider(); + var pool = new ReceiveContextPool(1); + var firstDelivery = new ReceiveContext(); + var shared = firstDelivery.Features.GetOrSet(); + shared.MessageConsumed = true; + var first = firstDelivery.CopyTo(pool.Get(), services); + first.Headers.Set("old", "value"); + + // act + pool.Return(first); + var secondDelivery = new ReceiveContext { MessageId = "second" }; + var second = secondDelivery.CopyTo(pool.Get(), services); + + // assert + Assert.Same(first, second); + Assert.Equal("second", second.MessageId); + Assert.True(second.Features.IsEmpty); + Assert.Empty(second.Headers); + Assert.True(shared.MessageConsumed); + pool.Return(second); + } + + private sealed class CloneFeature + { + public int Value { get; set; } + } + + private sealed class ExecutionFeature; +} diff --git a/src/Mocha/test/Mocha.Tests/Inbox/ConsumeInboxMiddlewareTests.cs b/src/Mocha/test/Mocha.Tests/Inbox/ConsumeInboxMiddlewareTests.cs index 3b20a4fe7c7..bbf3de25203 100644 --- a/src/Mocha/test/Mocha.Tests/Inbox/ConsumeInboxMiddlewareTests.cs +++ b/src/Mocha/test/Mocha.Tests/Inbox/ConsumeInboxMiddlewareTests.cs @@ -203,8 +203,6 @@ private static ReceiveContext CreateConsumeContext( context.MessageId = messageId; context.Services = provider; - // Set up the ReceiveConsumerFeature with a mock consumer identity - // so the inbox middleware can determine the consumer type. var consumerFeature = context.Features.GetOrSet(); consumerFeature.CurrentConsumer = new TestConsumer(); diff --git a/src/Mocha/test/Mocha.Tests/Middlewares/Consume/ConsumerInstrumentationMiddlewareTests.cs b/src/Mocha/test/Mocha.Tests/Middlewares/Consume/ConsumerInstrumentationMiddlewareTests.cs index 97b2cb6c008..566189e7c55 100644 --- a/src/Mocha/test/Mocha.Tests/Middlewares/Consume/ConsumerInstrumentationMiddlewareTests.cs +++ b/src/Mocha/test/Mocha.Tests/Middlewares/Consume/ConsumerInstrumentationMiddlewareTests.cs @@ -200,6 +200,9 @@ private sealed class StubConsumeContext : IConsumeContext public ReadOnlyMemory Body => ReadOnlyMemory.Empty; public MessageEnvelope? Envelope { get; set; } public IRemoteHostInfo Host { get; set; } = null!; + + public IConsumeContext Clone(IServiceProvider services) + => new StubConsumeContext { Services = services }; } private sealed class MockMessagingDiagnosticEvents : IMessagingDiagnosticEvents diff --git a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/BatchingTests.cs b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/BatchingTests.cs index d2b1dbabb8d..93b14deab47 100644 --- a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/BatchingTests.cs +++ b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/BatchingTests.cs @@ -1,3 +1,4 @@ +using CookieCrumble; using Microsoft.Extensions.DependencyInjection; using Mocha.Transport.InMemory.Tests.Helpers; @@ -94,6 +95,56 @@ public async Task Handler_Should_ReceiveMultiMessageBatch_When_ConcurrentDeliver Assert.Equal(BatchCompletionMode.Size, batch.CompletionMode); } + [Fact] + public async Task BatchHandler_Should_UseFreshScopeAndFeatures_When_Retried() + { + // arrange + var capture = new BatchRetryCapture(); + await using var provider = await InMemoryBusFixture.CreateBusAsync(b => + { + b.Services.AddSingleton(capture); + b.Services.AddScoped(); + b.AddResilience(p => p.On() + .Retry(1, TimeSpan.Zero, RetryBackoffType.Constant)); + b.AddBatchHandler(options => options.MaxBatchSize = 1); + }); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync( + new MutableBatchMessage { Value = "original" }, + CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + await capture.Disposed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + new + { + ProbeScopeCount = capture.ProbeIds.Distinct().Count(), + capture.RetryCounts, + RetryFeatureCount = capture.RetryFeatures.Distinct().Count(), + capture.RootFeatureWasPresent, + DisposedProbeCount = capture.DisposedProbeIds.Count + }.MatchInlineSnapshot( + """ + { + "ProbeScopeCount": 2, + "RetryCounts": [ + 0, + 1 + ], + "RetryFeatureCount": 1, + "RootFeatureWasPresent": [ + false, + false + ], + "DisposedProbeCount": 2 + } + """); + } + public sealed class TestBatchHandler(BatchMessageRecorder recorder) : IBatchEventHandler { public ValueTask HandleAsync(IMessageBatch batch, CancellationToken cancellationToken) @@ -102,4 +153,72 @@ public ValueTask HandleAsync(IMessageBatch batch, CancellationToke return default; } } + + public sealed class RetryOnceBatchHandler( + BatchScopeProbe probe, + ConsumeContextAccessor accessor, + BatchRetryCapture capture) + : IBatchEventHandler + { + public ValueTask HandleAsync( + IMessageBatch batch, + CancellationToken cancellationToken) + { + var context = accessor.Context!; + capture.ProbeIds.Add(probe.Id); + capture.RetryCounts.Add(context.Features.Get()!.ImmediateRetryCount); + capture.RetryFeatures.Add(context.Features.Get()!); + capture.RootFeatureWasPresent.Add(context.Features.Get() is not null); + context.Features.Set(new BatchAttemptFeature()); + + if (Interlocked.Increment(ref capture.Attempts) == 1) + { + throw new InvalidOperationException("Retry batch once."); + } + + capture.Completed.TrySetResult(); + return default; + } + } + + public sealed class BatchScopeProbe(BatchRetryCapture capture) : IDisposable + { + public Guid Id { get; } = Guid.NewGuid(); + + public void Dispose() + { + capture.DisposedProbeIds.Add(Id); + + if (capture.DisposedProbeIds.Count == 2) + { + capture.Disposed.TrySetResult(); + } + } + } + + public sealed class BatchRetryCapture + { + public List ProbeIds { get; } = []; + + public List DisposedProbeIds { get; } = []; + + public List RetryCounts { get; } = []; + + public List RetryFeatures { get; } = []; + + public List RootFeatureWasPresent { get; } = []; + + public TaskCompletionSource Completed { get; } = new(); + + public TaskCompletionSource Disposed { get; } = new(); + + public int Attempts; + } + + public sealed class MutableBatchMessage + { + public required string Value { get; set; } + } + + public sealed class BatchAttemptFeature; } diff --git a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs index 2a32d5e7bca..937650feca2 100644 --- a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs +++ b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs @@ -1,5 +1,6 @@ using System.Collections.Concurrent; using Microsoft.Extensions.DependencyInjection; +using Mocha.Middlewares; using Mocha.Transport.InMemory.Tests.Helpers; namespace Mocha.Transport.InMemory.Tests.Behaviors; @@ -310,6 +311,88 @@ public async Task Retry_Should_UseFreshScope_When_HandlerIsRetried() Assert.Equal(2, capture.ProbeIds.Distinct().Count()); } + [Fact] + public async Task Consumer_Should_UseDifferentScopes_When_NoPolicyAndMultipleConsumersHandleDelivery() + { + // arrange + var capture = new NoPolicyScopeCapture(); + await using var provider = await new ServiceCollection() + .AddSingleton(capture) + .AddScoped() + .AddMessageBus() + .AddEventHandler() + .AddEventHandler() + .AddInMemory() + .BuildServiceProvider(); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new OrderCreated { OrderId = "ORD-NO-POLICY-SCOPE" }, CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.Equal(2, capture.ProbeIds.Distinct().Count()); + Assert.Equal([true, true], capture.RetryFeatureWasPresent); + } + + [Fact] + public async Task ConsumerAttempt_Should_NotMutateDeliveryOrNextAttempt_When_FirstAttemptFails() + { + // arrange + var capture = new AttemptMutationCapture(); + var services = new ServiceCollection().AddSingleton(capture); + var builder = services.AddMessageBus() + .AddResilience(p => p.On() + .Retry(1, TimeSpan.Zero, RetryBackoffType.Constant)) + .AddConsumer(); + builder.ConfigureMessageBus(b => b.UseReceive(DeliveryStateCaptureMiddleware.Create(capture))); + await using var provider = await builder + .AddInMemory() + .BuildServiceProvider(); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new MutableRetryMessage { Value = "original" }, CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + await capture.DeliveryCaptured.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.Equal("original", capture.SecondAttemptValue); + Assert.False(capture.SecondAttemptFeatureWasPresent); + Assert.Equal(capture.DeliveryBefore, capture.DeliveryAfter); + } + + [Fact] + public async Task ConsumerAttempt_Should_UseAndDisposeScopeDistinctFromDelivery_When_FirstAttemptRuns() + { + // arrange + var capture = new ScopeLifecycleCapture(); + var services = new ServiceCollection() + .AddSingleton(capture) + .AddScoped(); + var builder = services.AddMessageBus() + .AddEventHandler(); + builder.ConfigureMessageBus(b => b.UseReceive(DeliveryScopeProbeMiddleware.Create())); + await using var provider = await builder + .AddInMemory() + .BuildServiceProvider(); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new OrderCreated { OrderId = "ORD-LIFECYCLE" }, CancellationToken.None); + + // assert + await capture.Disposed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.NotEqual(capture.DeliveryProbeId, capture.AttemptProbeId); + Assert.Equal(2, capture.DisposedProbeIds.Count); + } + [Fact] public async Task Retry_Should_PropagateConversationId_When_HandlerPublishesOnRetry() { @@ -417,6 +500,64 @@ private sealed class ScopeCapture public TaskCompletionSource FollowUpReceived { get; } = new(); } + private sealed class NoPolicyScopeCapture + { + public ConcurrentBag ProbeIds { get; } = []; + + public TaskCompletionSource Completed { get; } = new(); + + public ConcurrentBag RetryFeatureWasPresent { get; } = []; + + public void Record(Guid probeId, bool retryFeatureWasPresent) + { + RetryFeatureWasPresent.Add(retryFeatureWasPresent); + ProbeIds.Add(probeId); + + if (ProbeIds.Count == 2) + { + Completed.TrySetResult(); + } + } + } + + private sealed class AttemptMutationCapture + { + public (string? MessageId, string? CorrelationId, MessageEnvelope? Envelope) DeliveryBefore { get; set; } + + public (string? MessageId, string? CorrelationId, MessageEnvelope? Envelope) DeliveryAfter { get; set; } + + public string? SecondAttemptValue { get; set; } + + public bool SecondAttemptFeatureWasPresent { get; set; } + + public TaskCompletionSource Completed { get; } = new(); + + public TaskCompletionSource DeliveryCaptured { get; } = new(); + + public int Attempts; + } + + private sealed class ScopeLifecycleCapture + { + public Guid DeliveryProbeId { get; set; } + + public Guid AttemptProbeId { get; set; } + + public List DisposedProbeIds { get; } = []; + + public TaskCompletionSource Disposed { get; } = new(); + + public void RecordDisposed(Guid id) + { + DisposedProbeIds.Add(id); + + if (DisposedProbeIds.Count == 2) + { + Disposed.TrySetResult(); + } + } + } + private sealed record RetryFollowUp; // ============================================================ @@ -529,6 +670,117 @@ public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellatio } } + private sealed class FirstNoPolicyScopedHandler( + ScopeProbe probe, + ConsumeContextAccessor accessor, + NoPolicyScopeCapture capture) + : IEventHandler + { + public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellationToken) + { + capture.Record(probe.Id, accessor.Context!.Features.Get() is not null); + return default; + } + } + + private sealed class SecondNoPolicyScopedHandler( + ScopeProbe probe, + ConsumeContextAccessor accessor, + NoPolicyScopeCapture capture) + : IEventHandler + { + public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellationToken) + { + capture.Record(probe.Id, accessor.Context!.Features.Get() is not null); + return default; + } + } + + private sealed class MutatingRetryConsumer(AttemptMutationCapture capture) + : IConsumer + { + public ValueTask ConsumeAsync(IConsumeContext context) + { + var attempt = (IConsumeContext)context; + + if (Interlocked.Increment(ref capture.Attempts) == 1) + { + attempt.MessageId = "mutated-message"; + attempt.CorrelationId = "mutated-correlation"; + attempt.Envelope = null; + attempt.Features.Set(new AttemptMutationFeature()); + context.Message.Value = "mutated"; + throw new InvalidOperationException("Retry after mutation."); + } + + capture.SecondAttemptValue = context.Message.Value; + capture.SecondAttemptFeatureWasPresent = + attempt.Features.Get() is not null; + capture.Completed.TrySetResult(); + return default; + } + } + + private sealed class LifecycleProbe(ScopeLifecycleCapture capture) : IDisposable + { + public Guid Id { get; } = Guid.NewGuid(); + + public void Dispose() => capture.RecordDisposed(Id); + } + + private sealed class LifecycleProbeHandler( + LifecycleProbe probe, + ScopeLifecycleCapture capture) + : IEventHandler + { + public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellationToken) + { + capture.AttemptProbeId = probe.Id; + return default; + } + } + + private sealed class DeliveryStateCaptureMiddleware(AttemptMutationCapture capture) + { + public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next) + { + capture.DeliveryBefore = (context.MessageId, context.CorrelationId, context.Envelope); + await next(context); + capture.DeliveryAfter = (context.MessageId, context.CorrelationId, context.Envelope); + capture.DeliveryCaptured.TrySetResult(); + } + + public static ReceiveMiddlewareConfiguration Create(AttemptMutationCapture capture) + => new( + (_, next) => + { + var middleware = new DeliveryStateCaptureMiddleware(capture); + return context => middleware.InvokeAsync(context, next); + }, + "DeliveryStateCapture"); + } + + private sealed class DeliveryScopeProbeMiddleware + { + public static ReceiveMiddlewareConfiguration Create() + => new( + static (_, next) => async context => + { + var probe = context.Services.GetRequiredService(); + var capture = context.Services.GetRequiredService(); + capture.DeliveryProbeId = probe.Id; + await next(context); + }, + "DeliveryScopeProbe"); + } + + private sealed class MutableRetryMessage + { + public string Value { get; set; } = ""; + } + + private sealed class AttemptMutationFeature; + /// /// Throws on the first attempt and publishes a follow-up from the retry. /// From 8970e57a571c177063693396b69142ca299beb01 Mon Sep 17 00:00:00 2001 From: Pascal Senn Date: Sat, 5 Sep 2026 22:25:31 +0200 Subject: [PATCH 5/7] clenaup --- .../ConsumerResilienceMiddlewareTests.cs | 342 ------------------ .../ConsumerRetryScopeTests.cs | 99 ----- .../Context/ConsumeContextCloneTests.cs | 165 --------- 3 files changed, 606 deletions(-) delete mode 100644 src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs delete mode 100644 src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs delete mode 100644 src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs deleted file mode 100644 index 5a035f8ce89..00000000000 --- a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerResilienceMiddlewareTests.cs +++ /dev/null @@ -1,342 +0,0 @@ -using Microsoft.EntityFrameworkCore; -using Microsoft.EntityFrameworkCore.Infrastructure; -using Microsoft.EntityFrameworkCore.Storage; -using Microsoft.Extensions.DependencyInjection; -using Mocha.Features; -using Mocha.Transport.InMemory; - -namespace Mocha.EntityFrameworkCore.Tests; - -public sealed class ConsumerResilienceMiddlewareTests -{ - private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); - - [Fact] - public async Task Resilience_Should_UseFreshCompleteAttemptEnvironment_When_StrategyRetries() - { - // arrange - var capture = new ResilienceCapture(); - capture.FailuresBeforeSuccess = 3; - var services = new ServiceCollection(); - services.AddSingleton(capture); - services.AddScoped(); - services.AddDbContext(options => - options - .UseSqlite("Data Source=:memory:") - .ReplaceService()); - var builder = services.AddMessageBus() - .AddEntityFramework(p => p.UseResilience()) - .AddResilience(p => p.On().Retry(1)) - .AddConsumer(); - builder.ConfigureMessageBus(b => b.UseReceive(new ReceiveMiddlewareConfiguration( - (_, next) => async context => - { - context.Features.Set(capture.ReceiveFeature); - await next(context); - }, - "ReceiveFeature"))); - builder.AddInMemory(); - - await using var provider = services.BuildServiceProvider(); - var runtime = (MessagingRuntime)provider.GetRequiredService(); - await runtime.StartAsync(CancellationToken.None); - - using var scope = provider.CreateScope(); - var bus = scope.ServiceProvider.GetRequiredService(); - - // act - await bus.PublishAsync( - new ResilienceMessage { Value = "original" }, - CancellationToken.None); - - // assert - await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - await capture.Disposed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - Assert.Equal( - (4, 4, 4), - ( - capture.DbContextIds.Distinct().Count(), - capture.ProbeIds.Distinct().Count(), - capture.DisposedProbeIds.Distinct().Count())); - Assert.Equal( - [(true, false, true), (true, false, true), (true, false, true), (true, false, true)], - capture.AccessorWasValid - .Zip(capture.FeatureWasPresent, capture.ReceiveFeatureWasShared) - .Select(static values => (values.First, values.Second, values.Third))); - Assert.Equal([0, 1, 2, 3], capture.ImmediateRetryCounts); - Assert.Equal([true, true, true, true], capture.ConsumerWasPresent); - } - - [Fact] - public async Task Resilience_Should_PreserveTypedBatchContextAndFreshScope_When_StrategyRetries() - { - // arrange - var capture = new ResilienceCapture(); - var services = new ServiceCollection(); - services.AddSingleton(capture); - services.AddScoped(); - services.AddDbContext(options => - options - .UseSqlite("Data Source=:memory:") - .ReplaceService()); - var builder = services.AddMessageBus() - .AddEntityFramework(p => p.UseResilience()) - .AddBatchHandler(options => options.MaxBatchSize = 1); - builder.AddInMemory(); - - await using var provider = services.BuildServiceProvider(); - var runtime = (MessagingRuntime)provider.GetRequiredService(); - await runtime.StartAsync(CancellationToken.None); - - using var scope = provider.CreateScope(); - var bus = scope.ServiceProvider.GetRequiredService(); - - // act - await bus.PublishAsync( - new ResilienceMessage { Value = "original" }, - CancellationToken.None); - - // assert - await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - await capture.Disposed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - Assert.Equal( - (2, 2), - (capture.DbContextIds.Distinct().Count(), capture.ProbeIds.Distinct().Count())); - Assert.Equal( - [(true, false, 0, true), (true, false, 1, true)], - capture.AccessorWasValid - .Select((value, index) => - ( - value, - capture.FeatureWasPresent[index], - capture.ImmediateRetryCounts[index], - capture.ConsumerWasPresent[index]))); - } - - [Theory] - [InlineData("EntityFrameworkTransaction", "Inbox", "Custom")] - [InlineData("Inbox", "Custom", "EntityFrameworkTransaction")] - [InlineData("EntityFrameworkTransaction", "Custom", "Inbox")] - [InlineData("Custom", "Inbox", "EntityFrameworkTransaction")] - [InlineData("Inbox", "EntityFrameworkTransaction", "Custom")] - [InlineData("Custom", "EntityFrameworkTransaction", "Inbox")] - public async Task Pipeline_Should_PreserveRegistrationOrder_When_UsingPersistenceMiddlewareKeys( - string first, - string second, - string third) - { - // arrange - var capture = new MiddlewareOrderCapture(); - var services = new ServiceCollection(); - services.AddSingleton(capture); - var builder = services.AddMessageBus() - .AddEventHandler(); - builder.ConfigureMessageBus(b => - { - b.UseConsume(CreateOrderMiddleware(first)); - b.UseConsume(CreateOrderMiddleware(second)); - b.UseConsume(CreateOrderMiddleware(third)); - }); - builder.AddInMemory(); - - await using var provider = services.BuildServiceProvider(); - var runtime = (MessagingRuntime)provider.GetRequiredService(); - await runtime.StartAsync(CancellationToken.None); - - using var scope = provider.CreateScope(); - var bus = scope.ServiceProvider.GetRequiredService(); - - // act - await bus.PublishAsync(new MiddlewareOrderMessage(), CancellationToken.None); - - // assert - await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - Assert.Equal( - [first, second, third, "Handler"], - capture.Entries); - } - - private static ConsumerMiddlewareConfiguration CreateOrderMiddleware(string key) - => new( - (_, next) => context => InvokeOrderMiddlewareAsync(key, context, next), - key); - - private static async ValueTask InvokeOrderMiddlewareAsync( - string key, - IConsumeContext context, - ConsumerDelegate next) - { - context.Services.GetRequiredService().Entries.Add(key); - await next(context); - } - - public sealed class ResilienceDbContext(DbContextOptions options) - : DbContext(options); - - public sealed class RetryExecutionStrategyFactory(ExecutionStrategyDependencies dependencies) - : IExecutionStrategyFactory - { - public IExecutionStrategy Create() => new RetryExecutionStrategy(dependencies); - } - - private sealed class RetryExecutionStrategy(ExecutionStrategyDependencies dependencies) - : ExecutionStrategy(dependencies, 1, TimeSpan.Zero) - { - protected override bool ShouldRetryOn(Exception exception) - => exception is RetryExecutionException; - } - - public sealed class ResilienceConsumer( - ResilienceDbContext dbContext, - ResilienceProbe probe, - ConsumeContextAccessor accessor, - ResilienceCapture capture) - : IConsumer - { - public ValueTask ConsumeAsync(IConsumeContext context) - { - var receiveFeature = context.Features.Get(); - capture.ReceiveFeatureWasShared.Add( - ReferenceEquals(receiveFeature, capture.ReceiveFeature) - && receiveFeature.Value == capture.Attempts); - receiveFeature!.Value++; - capture.DbContextIds.Add(dbContext.ContextId.InstanceId); - capture.ProbeIds.Add(probe.Id); - capture.AccessorWasValid.Add( - accessor.Context?.Services == context.Services - && ReferenceEquals( - context.Services.GetRequiredService(), - dbContext)); - capture.ImmediateRetryCounts.Add(context.Features.Get()!.ImmediateRetryCount); - capture.ConsumerWasPresent.Add( - context.Features.Get()?.CurrentConsumer is not null); - capture.FeatureWasPresent.Add( - context.Features.Get() is not null); - context.Features.Set(new ResilienceAttemptFeature()); - - if (Interlocked.Increment(ref capture.Attempts) <= capture.FailuresBeforeSuccess) - { - throw new RetryExecutionException(); - } - - capture.Completed.TrySetResult(); - return default; - } - } - - public sealed class ResilienceProbe(ResilienceCapture capture) : IDisposable - { - public Guid Id { get; } = Guid.NewGuid(); - - public void Dispose() - { - capture.DisposedProbeIds.Add(Id); - - if (capture.DisposedProbeIds.Count == capture.FailuresBeforeSuccess + 1) - { - capture.Disposed.TrySetResult(); - } - } - } - - public sealed class BatchResilienceHandler( - ResilienceDbContext dbContext, - ResilienceProbe probe, - ConsumeContextAccessor accessor, - ResilienceCapture capture) - : IBatchEventHandler - { - public ValueTask HandleAsync( - IMessageBatch batch, - CancellationToken cancellationToken) - { - var context = accessor.Context!; - capture.DbContextIds.Add(dbContext.ContextId.InstanceId); - capture.ProbeIds.Add(probe.Id); - capture.AccessorWasValid.Add( - context is IBatchConsumeContext - && ReferenceEquals( - context.Services.GetRequiredService(), - dbContext)); - capture.ImmediateRetryCounts.Add(context.Features.Get()!.ImmediateRetryCount); - capture.ConsumerWasPresent.Add( - context.Features.Get()?.CurrentConsumer is not null); - capture.FeatureWasPresent.Add( - context.Features.Get() is not null); - context.Features.Set(new ResilienceAttemptFeature()); - - if (Interlocked.Increment(ref capture.Attempts) <= capture.FailuresBeforeSuccess) - { - throw new RetryExecutionException(); - } - - capture.Completed.TrySetResult(); - return default; - } - } - - public sealed class ResilienceCapture - { - public SharedReceiveFeature ReceiveFeature { get; } = new(); - - public List ReceiveFeatureWasShared { get; } = []; - - public List DbContextIds { get; } = []; - - public List ProbeIds { get; } = []; - - public List DisposedProbeIds { get; } = []; - - public List AccessorWasValid { get; } = []; - - public List FeatureWasPresent { get; } = []; - - public List ImmediateRetryCounts { get; } = []; - - public List ConsumerWasPresent { get; } = []; - - public TaskCompletionSource Completed { get; } = new(); - - public TaskCompletionSource Disposed { get; } = new(); - - public int Attempts; - - public int FailuresBeforeSuccess { get; set; } = 1; - } - - public sealed class MiddlewareOrderHandler(MiddlewareOrderCapture capture) - : IEventHandler - { - public ValueTask HandleAsync( - MiddlewareOrderMessage message, - CancellationToken cancellationToken) - { - capture.Entries.Add("Handler"); - capture.Completed.TrySetResult(); - return default; - } - } - - public sealed class MiddlewareOrderCapture - { - public List Entries { get; } = []; - - public TaskCompletionSource Completed { get; } = new(); - } - - public sealed class ResilienceMessage - { - public required string Value { get; set; } - } - - public sealed class MiddlewareOrderMessage; - - private sealed class ResilienceAttemptFeature; - - public sealed class SharedReceiveFeature - { - public int Value { get; set; } - } - - private sealed class RetryExecutionException : Exception; -} diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs deleted file mode 100644 index 773f42c7514..00000000000 --- a/src/Mocha/test/Mocha.EntityFrameworkCore.Tests/ConsumerRetryScopeTests.cs +++ /dev/null @@ -1,99 +0,0 @@ -using Microsoft.Data.Sqlite; -using Microsoft.EntityFrameworkCore; -using Microsoft.Extensions.DependencyInjection; -using Mocha.Transport.InMemory; - -namespace Mocha.EntityFrameworkCore.Tests; - -/// -/// Tests that an in-process retry does not re-run a handler against the -/// the failed attempt left behind, which would replay its pending changes on top of the retry's. -/// -public sealed class ConsumerRetryScopeTests -{ - private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); - - [Fact] - public async Task Retry_Should_InsertOnce_When_HandlerFailsAfterAddAndSucceedsOnRetry() - { - // arrange - await using var connection = new SqliteConnection("DataSource=:memory:"); - await connection.OpenAsync(TestContext.Current.CancellationToken); - var saved = new TaskCompletionSource(); - var attempts = new AttemptCounter(); - - var services = new ServiceCollection(); - services.AddDbContext(o => o.UseSqlite(connection)); - services.AddSingleton(saved); - services.AddSingleton(attempts); - var builder = services.AddMessageBus() - .AddResilience(p => p.On().Retry(3, TimeSpan.FromMilliseconds(1), RetryBackoffType.Constant)) - .AddEventHandler(); - builder.AddInMemory(); - - await using var provider = services.BuildServiceProvider(); - using (var setup = provider.CreateScope()) - { - await setup.ServiceProvider.GetRequiredService().Database - .EnsureCreatedAsync(TestContext.Current.CancellationToken); - } - - var runtime = (MessagingRuntime)provider.GetRequiredService(); - await runtime.StartAsync(CancellationToken.None); - - using var scope = provider.CreateScope(); - var bus = scope.ServiceProvider.GetRequiredService(); - - // act - await bus.PublishAsync(new CreateRow("only-once"), CancellationToken.None); - await saved.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - - // assert - the retry saved exactly one row, not the failed attempt's row as well - using var verify = provider.CreateScope(); - var rows = await verify.ServiceProvider.GetRequiredService().Rows - .Select(r => r.Name) - .ToListAsync(TestContext.Current.CancellationToken); - - Assert.Equal(2, attempts.Value); - Assert.Equal(["only-once"], rows); - } - - public sealed class AttemptCounter - { - public int Value; - } - - public sealed record CreateRow(string Name); - - public sealed class Row - { - public int Id { get; set; } - - public string Name { get; set; } = ""; - } - - public sealed class RetryDbContext(DbContextOptions options) : DbContext(options) - { - public DbSet Rows => Set(); - } - - /// - /// Adds its row, throws on the first attempt before saving, and saves on the retry. - /// - public sealed class AddThenThrowOnceHandler(RetryDbContext db, AttemptCounter attempts, TaskCompletionSource saved) - : IEventHandler - { - public async ValueTask HandleAsync(CreateRow message, CancellationToken cancellationToken) - { - db.Rows.Add(new Row { Name = message.Name }); - - if (Interlocked.Increment(ref attempts.Value) == 1) - { - throw new InvalidOperationException("transient"); - } - - await db.SaveChangesAsync(cancellationToken); - saved.TrySetResult(); - } - } -} diff --git a/src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs b/src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs deleted file mode 100644 index f433dac262e..00000000000 --- a/src/Mocha/test/Mocha.Tests/Context/ConsumeContextCloneTests.cs +++ /dev/null @@ -1,165 +0,0 @@ -using Microsoft.Extensions.DependencyInjection; -using Mocha.Features; -using Mocha.Middlewares; - -namespace Mocha.Tests.Context; - -public sealed class ConsumeContextCloneTests -{ - [Fact] - public void Clone_Should_PreserveCurrentMetadataAndShareReceiveFeatures_When_ReceiveContextIsCloned() - { - // arrange - using var services = new ServiceCollection().BuildServiceProvider(); - var envelope = new MessageEnvelope { MessageId = "envelope-message" }; - var context = new ReceiveContext - { - Services = null!, - MessageId = "current-message", - Envelope = envelope - }; - context.Headers.Set("x-test", "original"); - context.Features.Set(new CloneFeature()); - - // act - var clone = (ReceiveContext)context.Clone(services); - - // assert - Assert.Equal( - ("current-message", "envelope-message", true), - (clone.MessageId, clone.Envelope?.MessageId, ReferenceEquals(services, clone.Services))); - Assert.True(clone.Headers.TryGetValue("x-test", out var header) && Equals(header, "original")); - Assert.NotSame(envelope, clone.Envelope); - Assert.NotSame(context.Headers, clone.Headers); - Assert.Same(context.Features.Get(), clone.Features.Get()); - } - - [Fact] - public void Clone_Should_PreserveTypedContext_When_TypedContextIsCloned() - { - // arrange - using var services = new ServiceCollection().BuildServiceProvider(); - var context = new ConsumeContext(new ReceiveContext()); - - // act - var clone = (ConsumeContext)context.Clone(services); - - // assert - Assert.Same(services, clone.Services); - } - - [Fact] - public void Clone_Should_ReuseBatchAndShareInputFeatures_When_BatchContextIsCloned() - { - // arrange - using var services = new ServiceCollection().BuildServiceProvider(); - var firstContext = new ReceiveContext(); - var batch = new MessageBatch( - [new BufferedEntry(new ConsumeContext(firstContext))], - BatchCompletionMode.Size); - var context = new BatchConsumeContext( - batch, - null!, - firstContext, - "batch", - null, - CancellationToken.None); - context.Features.Set(new CloneFeature()); - - // act - var clone = (BatchConsumeContext)context.Clone(services); - - // assert - Assert.Same(batch, clone.Message); - Assert.Same(services, clone.Services); - Assert.Same(context.Features.Get(), clone.Features.Get()); - } - - [Theory] - [InlineData(false)] - [InlineData(true)] - public void Clone_Should_KeepAttemptFeaturesOffTheInput_When_CloningAnExecution(bool batch) - { - // arrange - using var services = new ServiceCollection().BuildServiceProvider(); - var delivery = new ReceiveContext(); - IConsumeContext input = batch - ? new BatchConsumeContext( - new MessageBatch( - [new BufferedEntry(new ConsumeContext(delivery))], - BatchCompletionMode.Size), - services, - delivery, - "batch", - null, - CancellationToken.None) - : delivery; - var shared = new CloneFeature(); - input.Features.Set(shared); - - // act - var first = input.Clone(services); - first.Features.Get()!.Value = 42; - first.Features.Set(new ExecutionFeature()); - var sibling = input.Clone(services); - first.Features.Set(new CloneFeature { Value = 99 }); - - // assert - Assert.NotSame(input.Features, first.Features); - Assert.Equal(42, input.Features.Get()!.Value); - Assert.Null(input.Features.Get()); - Assert.Same(shared, sibling.Features.Get()); - Assert.Null(sibling.Features.Get()); - } - - [Fact] - public void Reset_Should_LeaveReceiveFeatureUntouched_When_CloneIsReset() - { - // arrange - using var services = new ServiceCollection().BuildServiceProvider(); - var delivery = new ReceiveContext(); - var shared = delivery.Features.GetOrSet(); - shared.MessageConsumed = true; - var clone = (ReceiveContext)delivery.Clone(services); - - // act - clone.Reset(); - - // assert - Assert.True(shared.MessageConsumed); - Assert.True(clone.Features.IsEmpty); - } - - [Fact] - public void Clone_Should_ReuseResetContext_When_ReturnedToPool() - { - // arrange - using var services = new ServiceCollection().BuildServiceProvider(); - var pool = new ReceiveContextPool(1); - var firstDelivery = new ReceiveContext(); - var shared = firstDelivery.Features.GetOrSet(); - shared.MessageConsumed = true; - var first = firstDelivery.CopyTo(pool.Get(), services); - first.Headers.Set("old", "value"); - - // act - pool.Return(first); - var secondDelivery = new ReceiveContext { MessageId = "second" }; - var second = secondDelivery.CopyTo(pool.Get(), services); - - // assert - Assert.Same(first, second); - Assert.Equal("second", second.MessageId); - Assert.True(second.Features.IsEmpty); - Assert.Empty(second.Headers); - Assert.True(shared.MessageConsumed); - pool.Return(second); - } - - private sealed class CloneFeature - { - public int Value { get; set; } - } - - private sealed class ExecutionFeature; -} From c02acf9edc6ae17820b7e2e5e0f1c22567b7b3f8 Mon Sep 17 00:00:00 2001 From: Pascal Senn Date: Sat, 5 Sep 2026 23:42:46 +0200 Subject: [PATCH 6/7] Keep receive context features visible to pooled attempt clones A ReceiveContext rented as an attempt clone no longer restores its parked pooled features, which shadowed the receive context's instances through the fallback. The parsing feature is created on the receive context so the message is deserialized once per receive. Batch consumers get a receive scope again, the saga store disposes the documents it loads, and RetryFeature docs match the always-on retry middleware. Claude-Session: https://claude.ai/code/session_01SRDEKC2sASpSfqXSHrLJBV --- .../Sagas/DbContextSagaStore.cs | 14 +- .../Features/PooledFeatureCollection.cs | 27 +- .../Implementations/BatchConsumer.cs | 7 +- .../src/Mocha/Endpoints/ReceiveEndpoint.cs | 4 + .../Consume/Retry/ConsumerRetryMiddleware.cs | 9 +- .../Middlewares/Consume/Retry/RetryFeature.cs | 10 +- .../Behaviors/RetryTests.cs | 236 +++++++++++++++--- 7 files changed, 259 insertions(+), 48 deletions(-) diff --git a/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs b/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs index f19e8cc5818..78b161f9804 100644 --- a/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs +++ b/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 _documents = []; private PooledArrayWriter? _arrayWriter; private List? _buffers; @@ -71,6 +72,7 @@ public async Task SaveAsync(Saga saga, T state, CancellationToken cancellatio .FirstOrDefaultAsync(cancellationToken); var document = ToJsonDocument(saga, state); + _documents.Add(document); if (sagaState is null) { @@ -131,7 +133,8 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation public async Task LoadAsync(Saga saga, Guid id, CancellationToken cancellationToken) { // 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. + // 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() .AsTracking() @@ -143,6 +146,7 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation return default; } + _documents.Add(document); return FromJsonDocument(saga, document); } @@ -196,10 +200,16 @@ private static Guid NewVersion() } /// - /// 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. /// public void Dispose() { + foreach (var document in _documents) + { + document.Dispose(); + } + + _documents.Clear(); _arrayWriter?.Dispose(); if (_buffers is not null) { diff --git a/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs b/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs index e35da3d7649..dbe15741493 100644 --- a/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs +++ b/src/Mocha/src/Mocha.Utilities/Features/PooledFeatureCollection.cs @@ -131,12 +131,18 @@ public void Set(TFeature? instance) /// Initializes the feature collection with the specified 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. /// 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); @@ -156,14 +162,31 @@ public void Reset() { 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; + } + /// public IEnumerator> GetEnumerator() { diff --git a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs index 8a254596bf4..5306465823e 100644 --- a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs +++ b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs @@ -93,9 +93,14 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok try { _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( batch, - _applicationServices, + scope.ServiceProvider, batch.GetContext(0), Guid.NewGuid().ToString(), _itemMessageType, diff --git a/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs b/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs index ddcf70aea1d..027f496e4cc 100644 --- a/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs +++ b/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs @@ -355,6 +355,10 @@ public ReceiveEndpointDescription Describe() private static async ValueTask DefaultPipeline(IReceiveContext context) { + // The parsing feature has to live on the receive context so consumer attempts, which run + // on clones that fall back to the receive context's features, share one lazily deserialized message. + context.Features.GetOrSet(); + var feature = context.Features.GetOrSet(); var consumers = feature.Consumers; diff --git a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs index 3db1afd391d..911382e37b7 100644 --- a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs @@ -60,11 +60,10 @@ private static async ValueTask ExecuteAttemptAsync( try { - // The attempt reads the delivery's features through its clone, but needs its own - // consumer feature: a pooled clone may still carry a reset one from its previous use - // that would shadow the delivery's, and CurrentConsumer must be set for this attempt - // without touching the delivery's consumer set. Set() initializes pooled features, - // which clears CurrentConsumer, so it is assigned after the feature is added. + // The attempt reads the receive context's features through its clone, but needs its + // own consumer feature so CurrentConsumer can be set for this attempt without touching + // the receive context's consumer set. Set() initializes pooled features, which clears + // CurrentConsumer, so it is assigned after the feature is added. var consumerFeature = new ReceiveConsumerFeature(); attempt.Features.Set(consumerFeature); consumerFeature.CurrentConsumer = consumer; diff --git a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryFeature.cs b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryFeature.cs index 03af7aa0955..f8ff427afec 100644 --- a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryFeature.cs +++ b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryFeature.cs @@ -2,14 +2,16 @@ namespace Mocha; /// /// Provides retry state information to message handlers via -/// context.Features.Get<RetryState>(). -/// Null if retry is not configured via AddResilience. +/// context.Features.Get<RetryFeature>(). +/// Always present on a consumer attempt, even when no retry policy is configured via AddResilience. /// public sealed class RetryFeature { /// - /// Number of immediate retries already attempted for this delivery round. - /// 0 on the first (original) attempt. + /// Number of failed immediate attempts that preceded the current one in this delivery round, + /// including attempts repeated by an . + /// 0 on the first (original) attempt. Once the round has failed for good it equals the total + /// number of failed attempts. /// public int ImmediateRetryCount { get; internal set; } diff --git a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs index 937650feca2..c8219f70fe9 100644 --- a/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs +++ b/src/Mocha/test/Mocha.Transport.InMemory.Tests/Behaviors/RetryTests.cs @@ -1,5 +1,6 @@ using System.Collections.Concurrent; using Microsoft.Extensions.DependencyInjection; +using Mocha.Features; using Mocha.Middlewares; using Mocha.Transport.InMemory.Tests.Helpers; @@ -282,8 +283,8 @@ await counter.WaitForCountAsync(4, s_timeout), public async Task Retry_Should_UseFreshScope_When_HandlerIsRetried() { // arrange - // A retry re-runs the handler; it must not observe scoped state the failed attempt left - // behind, so each attempt has to resolve its own scoped service instance. + // Each attempt runs in its own scope, so a retry never sees what the failed attempt left + // in scoped services, like a DbContext with pending changes. var counter = new RetryInvocationCounter(); var capture = new ScopeCapture(); await using var provider = await new ServiceCollection() @@ -312,9 +313,11 @@ public async Task Retry_Should_UseFreshScope_When_HandlerIsRetried() } [Fact] - public async Task Consumer_Should_UseDifferentScopes_When_NoPolicyAndMultipleConsumersHandleDelivery() + public async Task Consumer_Should_UseDifferentScopes_When_NoPolicyAndMultipleConsumersHandleSameMessage() { // arrange + // The retry middleware creates the attempt scope and runs even without a policy, so two + // handlers on one receive context get separate scopes and RetryFeature is always present. var capture = new NoPolicyScopeCapture(); await using var provider = await new ServiceCollection() .AddSingleton(capture) @@ -338,16 +341,20 @@ public async Task Consumer_Should_UseDifferentScopes_When_NoPolicyAndMultipleCon } [Fact] - public async Task ConsumerAttempt_Should_NotMutateDeliveryOrNextAttempt_When_FirstAttemptFails() + public async Task ConsumerAttempt_Should_NotMutateReceiveContextOrNextAttempt_When_FirstAttemptFails() { // arrange + // Attempts run on a clone of the receive context: metadata is copied and features the + // attempt sets stay in the clone, so neither the receive context nor the next attempt sees + // them. Features read through the fallback and the message instance are the receive + // context's own and are shared. var capture = new AttemptMutationCapture(); var services = new ServiceCollection().AddSingleton(capture); var builder = services.AddMessageBus() .AddResilience(p => p.On() .Retry(1, TimeSpan.Zero, RetryBackoffType.Constant)) .AddConsumer(); - builder.ConfigureMessageBus(b => b.UseReceive(DeliveryStateCaptureMiddleware.Create(capture))); + builder.ConfigureMessageBus(b => b.UseReceive(ReceiveStateCaptureMiddleware.Create(capture))); await using var provider = await builder .AddInMemory() .BuildServiceProvider(); @@ -356,27 +363,29 @@ public async Task ConsumerAttempt_Should_NotMutateDeliveryOrNextAttempt_When_Fir var bus = scope.ServiceProvider.GetRequiredService(); // act - await bus.PublishAsync(new MutableRetryMessage { Value = "original" }, CancellationToken.None); + await bus.PublishAsync(new MutableRetryMessage("original"), CancellationToken.None); // assert await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - await capture.DeliveryCaptured.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - Assert.Equal("original", capture.SecondAttemptValue); + await capture.ReceiveCaptured.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); Assert.False(capture.SecondAttemptFeatureWasPresent); - Assert.Equal(capture.DeliveryBefore, capture.DeliveryAfter); + Assert.NotNull(capture.ReceiveBefore.Envelope); + Assert.Equal(capture.ReceiveBefore, capture.ReceiveAfter); } [Fact] - public async Task ConsumerAttempt_Should_UseAndDisposeScopeDistinctFromDelivery_When_FirstAttemptRuns() + public async Task ConsumerAttempt_Should_UseAndDisposeScopeDistinctFromReceiveScope_When_FirstAttemptRuns() { // arrange + // The attempt scope is separate from the receive scope and is disposed when the attempt + // ends, so scoped resources of a failed attempt do not outlive it into the retry. var capture = new ScopeLifecycleCapture(); var services = new ServiceCollection() .AddSingleton(capture) .AddScoped(); var builder = services.AddMessageBus() .AddEventHandler(); - builder.ConfigureMessageBus(b => b.UseReceive(DeliveryScopeProbeMiddleware.Create())); + builder.ConfigureMessageBus(b => b.UseReceive(ReceiveScopeProbeMiddleware.Create())); await using var provider = await builder .AddInMemory() .BuildServiceProvider(); @@ -389,16 +398,78 @@ public async Task ConsumerAttempt_Should_UseAndDisposeScopeDistinctFromDelivery_ // assert await capture.Disposed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); - Assert.NotEqual(capture.DeliveryProbeId, capture.AttemptProbeId); - Assert.Equal(2, capture.DisposedProbeIds.Count); + Assert.NotEqual(capture.ReceiveProbeId, capture.AttemptProbeId); + } + + [Fact] + public async Task ExecutionStrategy_Should_RepeatAttemptAndCountIt_When_AttemptFails() + { + // arrange + // An IConsumerExecutionStrategy (the EF Core resilience integration is one) repeats the + // attempt callback on its own. The retry middleware creates scope and clone per invocation + // and counts each failed one, so a repetition gives the handler a new scoped instance and + // an incremented ImmediateRetryCount like any other retry. + var capture = new StrategyCapture(); + var services = new ServiceCollection() + .AddSingleton(capture) + .AddScoped(); + var builder = services.AddMessageBus() + .AddEventHandler(); + builder.ConfigureMessageBus(b => b.ConfigureFeature(f => f.Set( + new ConsumerExecutionStrategyFeature { Strategy = new RepeatOnceStrategy() }))); + await using var provider = await builder + .AddInMemory() + .BuildServiceProvider(); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new OrderCreated { OrderId = "ORD-STRATEGY" }, CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.Equal([0, 1], capture.RetryCounts); + Assert.Equal(2, capture.ProbeIds.Distinct().Count()); + } + + [Fact] + public async Task Attempt_Should_SeeReceiveContextPooledFeature_When_CloneWasPooledBefore() + { + // arrange + // Attempt clones are pooled ReceiveContexts whose feature reads fall back to the receive + // context. Returning a context parks its reset IPooledFeature instances; renting it as a + // clone must leave them parked, otherwise an empty instance shadows the receive context's. + // That only shows on a reused instance, so the handler throws twice: the first attempt + // parks a pooled feature on the clone and the next two reuse the clone. + var capture = new PooledFeatureCapture(); + var services = new ServiceCollection().AddSingleton(capture); + var builder = services.AddMessageBus() + .AddResilience(p => p.On() + .Retry(2, TimeSpan.Zero, RetryBackoffType.Constant)) + .AddEventHandler(); + builder.ConfigureMessageBus(b => b.UseReceive(ReceivePooledFeatureMiddleware.Create())); + await using var provider = await builder + .AddInMemory() + .BuildServiceProvider(); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new OrderCreated { OrderId = "ORD-POOLED" }, CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.Equal(["ORD-POOLED", "ORD-POOLED", "ORD-POOLED"], capture.ObservedValues); } [Fact] public async Task Retry_Should_PropagateConversationId_When_HandlerPublishesOnRetry() { // arrange - // Correlation propagation reads the ambient consume context from the scope, so the fresh - // scope of a retry has to carry it too, or a message published from a retry loses its lineage. + // ConsumeContextAccessor is scoped, so the attempt scope has its own instance and the retry + // middleware has to set it, or a message published from a retry starts a new conversation. var counter = new RetryInvocationCounter(); var capture = new ScopeCapture(); await using var provider = await new ServiceCollection() @@ -522,24 +593,22 @@ public void Record(Guid probeId, bool retryFeatureWasPresent) private sealed class AttemptMutationCapture { - public (string? MessageId, string? CorrelationId, MessageEnvelope? Envelope) DeliveryBefore { get; set; } - - public (string? MessageId, string? CorrelationId, MessageEnvelope? Envelope) DeliveryAfter { get; set; } + public (string? MessageId, string? CorrelationId, MessageEnvelope? Envelope) ReceiveBefore { get; set; } - public string? SecondAttemptValue { get; set; } + public (string? MessageId, string? CorrelationId, MessageEnvelope? Envelope) ReceiveAfter { get; set; } public bool SecondAttemptFeatureWasPresent { get; set; } public TaskCompletionSource Completed { get; } = new(); - public TaskCompletionSource DeliveryCaptured { get; } = new(); + public TaskCompletionSource ReceiveCaptured { get; } = new(); public int Attempts; } private sealed class ScopeLifecycleCapture { - public Guid DeliveryProbeId { get; set; } + public Guid ReceiveProbeId { get; set; } public Guid AttemptProbeId { get; set; } @@ -560,6 +629,51 @@ public void RecordDisposed(Guid id) private sealed record RetryFollowUp; + private sealed class StrategyCapture + { + public List RetryCounts { get; } = []; + + public List ProbeIds { get; } = []; + + public TaskCompletionSource Completed { get; } = new(); + } + + private sealed class PooledFeatureCapture + { + public List ObservedValues { get; } = []; + + public TaskCompletionSource Completed { get; } = new(); + + public int Attempts; + } + + /// + /// Runs the attempt again once when it fails, the way a database execution strategy would. + /// + private sealed class RepeatOnceStrategy : IConsumerExecutionStrategy + { + public async ValueTask ExecuteAsync(IConsumeContext context, Func executeAttempt) + { + try + { + await executeAttempt(context.CancellationToken); + } + catch (InvalidOperationException) + { + await executeAttempt(context.CancellationToken); + } + } + } + + private sealed class ReceivePooledFeature : IPooledFeature + { + public string? Value { get; set; } + + public void Initialize(object state) => Value = null; + + public void Reset() => Value = null; + } + // ============================================================ // Test Handlers // ============================================================ @@ -709,11 +823,9 @@ public ValueTask ConsumeAsync(IConsumeContext context) attempt.CorrelationId = "mutated-correlation"; attempt.Envelope = null; attempt.Features.Set(new AttemptMutationFeature()); - context.Message.Value = "mutated"; throw new InvalidOperationException("Retry after mutation."); } - capture.SecondAttemptValue = context.Message.Value; capture.SecondAttemptFeatureWasPresent = attempt.Features.Get() is not null; capture.Completed.TrySetResult(); @@ -740,27 +852,27 @@ public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellatio } } - private sealed class DeliveryStateCaptureMiddleware(AttemptMutationCapture capture) + private sealed class ReceiveStateCaptureMiddleware(AttemptMutationCapture capture) { public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next) { - capture.DeliveryBefore = (context.MessageId, context.CorrelationId, context.Envelope); + capture.ReceiveBefore = (context.MessageId, context.CorrelationId, context.Envelope); await next(context); - capture.DeliveryAfter = (context.MessageId, context.CorrelationId, context.Envelope); - capture.DeliveryCaptured.TrySetResult(); + capture.ReceiveAfter = (context.MessageId, context.CorrelationId, context.Envelope); + capture.ReceiveCaptured.TrySetResult(); } public static ReceiveMiddlewareConfiguration Create(AttemptMutationCapture capture) => new( (_, next) => { - var middleware = new DeliveryStateCaptureMiddleware(capture); + var middleware = new ReceiveStateCaptureMiddleware(capture); return context => middleware.InvokeAsync(context, next); }, - "DeliveryStateCapture"); + "ReceiveStateCapture"); } - private sealed class DeliveryScopeProbeMiddleware + private sealed class ReceiveScopeProbeMiddleware { public static ReceiveMiddlewareConfiguration Create() => new( @@ -768,17 +880,73 @@ public static ReceiveMiddlewareConfiguration Create() { var probe = context.Services.GetRequiredService(); var capture = context.Services.GetRequiredService(); - capture.DeliveryProbeId = probe.Id; + capture.ReceiveProbeId = probe.Id; await next(context); }, - "DeliveryScopeProbe"); + "ReceiveScopeProbe"); + } + + private sealed class StrategyThrowOnceHandler( + ScopeProbe probe, + ConsumeContextAccessor accessor, + StrategyCapture capture) + : IEventHandler + { + public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellationToken) + { + capture.ProbeIds.Add(probe.Id); + capture.RetryCounts.Add(accessor.Context!.Features.Get()!.ImmediateRetryCount); + + if (capture.ProbeIds.Count == 1) + { + throw new InvalidOperationException("Repeat via strategy"); + } + + capture.Completed.TrySetResult(); + return default; + } + } + + /// + /// Records the pooled feature value it sees, sets its own so the clone parks one on return, + /// and throws twice so the clone is reused. + /// + private sealed class PooledFeatureHandler(ConsumeContextAccessor accessor, PooledFeatureCapture capture) + : IEventHandler + { + public ValueTask HandleAsync(OrderCreated message, CancellationToken cancellationToken) + { + var context = accessor.Context!; + capture.ObservedValues.Add(context.Features.Get()?.Value); + context.Features.Set(new ReceivePooledFeature()); + + if (Interlocked.Increment(ref capture.Attempts) < 3) + { + throw new InvalidOperationException("Reuse the pooled clone"); + } + + capture.Completed.TrySetResult(); + return default; + } } - private sealed class MutableRetryMessage + private sealed class ReceivePooledFeatureMiddleware { - public string Value { get; set; } = ""; + public static ReceiveMiddlewareConfiguration Create() + => new( + static (_, next) => context => + { + // Set() initializes pooled features, which clears Value, so it is assigned afterwards. + var feature = new ReceivePooledFeature(); + context.Features.Set(feature); + feature.Value = "ORD-POOLED"; + return next(context); + }, + "ReceivePooledFeature"); } + private sealed record MutableRetryMessage(string Value); + private sealed class AttemptMutationFeature; /// From d7f6c7654d0cab6f1318675e6482696882a5bfda Mon Sep 17 00:00:00 2001 From: Pascal Senn Date: Sun, 6 Sep 2026 00:10:05 +0200 Subject: [PATCH 7/7] Let retry attempts read the receive context's consumer feature The attempt clone only creates its own ReceiveConsumerFeature when the context has none, so consumers that inspect the consumer set, like the reply consumer, see the receive context's instance. Claude-Session: https://claude.ai/code/session_01SRDEKC2sASpSfqXSHrLJBV --- .../Consume/Retry/ConsumerRetryMiddleware.cs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs index 911382e37b7..97f5035e5a8 100644 --- a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs @@ -60,13 +60,17 @@ private static async ValueTask ExecuteAttemptAsync( try { - // The attempt reads the receive context's features through its clone, but needs its - // own consumer feature so CurrentConsumer can be set for this attempt without touching - // the receive context's consumer set. Set() initializes pooled features, which clears + // The receive pipeline sets CurrentConsumer on the receive context's consumer feature, + // which the attempt reads through its clone. A batch context has no consumer feature, + // so the attempt gets one of its own. Set() initializes pooled features, which clears // CurrentConsumer, so it is assigned after the feature is added. - var consumerFeature = new ReceiveConsumerFeature(); - attempt.Features.Set(consumerFeature); - consumerFeature.CurrentConsumer = consumer; + if (attempt.Features.Get() is null) + { + var consumerFeature = new ReceiveConsumerFeature(); + attempt.Features.Set(consumerFeature); + consumerFeature.CurrentConsumer = consumer; + } + attempt.CancellationToken = cancellationToken; accessor.Context = attempt; await next(attempt);