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/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.EntityFrameworkCore/Sagas/DbContextSagaStore.cs b/src/Mocha/src/Mocha.EntityFrameworkCore/Sagas/DbContextSagaStore.cs index 3755cc1f694..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) { @@ -86,10 +88,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,7 +132,9 @@ 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) { - // as the state is scoped we load the whole saga state into memory for the concurrency check + // The state is loaded tracked so the save that follows checks the version observed here. + // Its document stays with the tracked entity and is released when a save replaces it or + // when this store is disposed. var sageState = await context .Set() .AsTracking() @@ -138,14 +146,8 @@ public async Task DeleteAsync(Saga saga, Guid id, CancellationToken cancellation return default; } - try - { - return FromJsonDocument(saga, document); - } - finally - { - document.Dispose(); - } + _documents.Add(document); + return FromJsonDocument(saga, document); } private JsonDocument ToJsonDocument(Saga saga, SagaStateBase state) @@ -198,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 807b87a280a..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); @@ -151,18 +157,36 @@ public void Initialize(IFeatureCollection? defaults = null) /// public void Reset() { + _defaults = null; foreach (var item in _features) { if (item.Value is IPooledFeature pooledFeature) { - _pooledFeatures.Add(item); pooledFeature.Reset(); + + if (!IsParked(item.Key)) + { + _pooledFeatures.Add(item); + } } } _features.Clear(); } + private bool IsParked(Type key) + { + foreach (var parked in _pooledFeatures) + { + if (parked.Key == key) + { + return true; + } + } + + return false; + } + /// public IEnumerator> GetEnumerator() { 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..5306465823e 100644 --- a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs +++ b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs @@ -94,6 +94,8 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok { _logger.DispatchingBatch(batch.Count, batch.CompletionMode); + // The batch has no receive scope of its own, so it gets one here, mirroring the scope + // ReceiveEndpoint creates per receive context. Each consumer attempt then runs in a child scope. await using var scope = _applicationServices.CreateAsyncScope(); var batchContext = new BatchConsumeContext( @@ -104,21 +106,7 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok _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/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 5a55bf4ac41..97f5035e5a8 100644 --- a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/ConsumerRetryMiddleware.cs @@ -1,14 +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. +/// 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) { @@ -27,25 +32,75 @@ public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate nex await RetryExecutor.ExecuteAsync( exceptionPolicyRules, - (next, context, retryState), - static (s) => 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); } + /// + /// Executes one consumer attempt on a clone of the context in a new service scope. + /// + private static async ValueTask ExecuteAttemptAsync( + ConsumerDelegate next, + IConsumeContext context, + RetryFeature retryState, + Consumer consumer, + CancellationToken cancellationToken) + { + await using var scope = context.Services.CreateAsyncScope(); + var accessor = scope.ServiceProvider.GetRequiredService(); + 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 + { + // 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. + 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); + } + catch + { + // Counts every failed attempt, including the ones the execution strategy repeats. + retryState.ImmediateRetryCount++; + throw; + } + finally + { + accessor.Context = null; + + if (pooledContext is not null) + { + pool.Return(pooledContext); + } + } + } + public static ConsumerMiddlewareConfiguration Create() => new( 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); }, @@ -68,4 +123,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/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/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/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs new file mode 100644 index 00000000000..f1526ebd34c --- /dev/null +++ b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/DbContextSagaStoreConcurrencyTests.cs @@ -0,0 +1,122 @@ +using Microsoft.EntityFrameworkCore; +using Mocha.EntityFrameworkCore.Postgres.Tests.Helpers; +using Mocha.Sagas.EfCore; + +namespace Mocha.EntityFrameworkCore.Postgres.Tests; + +/// +/// 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 +{ + private readonly List _disposables = []; + + public ValueTask InitializeAsync() => ValueTask.CompletedTask; + + [Fact] + public async Task LoadAsync_Should_ReturnState_When_CalledTwiceOnTheSameStore() + { + // arrange + // 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 + 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 LoadAsync_Should_ReturnSavedState_When_CalledAfterSaveOnTheSameStore() + { + // arrange + // 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 reloaded = await store.LoadAsync(saga, id, CancellationToken.None); + + // assert + Assert.Equal("after", reloaded!.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); + } + + 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(); + } + } +} 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.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.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.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/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 5cacef0b215..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,7 @@ using System.Collections.Concurrent; using Microsoft.Extensions.DependencyInjection; +using Mocha.Features; +using Mocha.Middlewares; using Mocha.Transport.InMemory.Tests.Helpers; namespace Mocha.Transport.InMemory.Tests.Behaviors; @@ -277,6 +279,225 @@ 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 + // 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() + .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 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) + .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_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(ReceiveStateCaptureMiddleware.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("original"), CancellationToken.None); + + // assert + await capture.Completed.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + await capture.ReceiveCaptured.Task.WaitAsync(s_timeout, TestContext.Current.CancellationToken); + Assert.False(capture.SecondAttemptFeatureWasPresent); + Assert.NotNull(capture.ReceiveBefore.Envelope); + Assert.Equal(capture.ReceiveBefore, capture.ReceiveAfter); + } + + [Fact] + 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(ReceiveScopeProbeMiddleware.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.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 + // 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() + .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 +555,125 @@ 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 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) ReceiveBefore { get; set; } + + public (string? MessageId, string? CorrelationId, MessageEnvelope? Envelope) ReceiveAfter { get; set; } + + public bool SecondAttemptFeatureWasPresent { get; set; } + + public TaskCompletionSource Completed { get; } = new(); + + public TaskCompletionSource ReceiveCaptured { get; } = new(); + + public int Attempts; + } + + private sealed class ScopeLifecycleCapture + { + public Guid ReceiveProbeId { 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; + + 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 // ============================================================ @@ -420,4 +760,224 @@ 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; + } + } + + 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()); + throw new InvalidOperationException("Retry after mutation."); + } + + 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 ReceiveStateCaptureMiddleware(AttemptMutationCapture capture) + { + public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next) + { + capture.ReceiveBefore = (context.MessageId, context.CorrelationId, context.Envelope); + await next(context); + capture.ReceiveAfter = (context.MessageId, context.CorrelationId, context.Envelope); + capture.ReceiveCaptured.TrySetResult(); + } + + public static ReceiveMiddlewareConfiguration Create(AttemptMutationCapture capture) + => new( + (_, next) => + { + var middleware = new ReceiveStateCaptureMiddleware(capture); + return context => middleware.InvokeAsync(context, next); + }, + "ReceiveStateCapture"); + } + + private sealed class ReceiveScopeProbeMiddleware + { + public static ReceiveMiddlewareConfiguration Create() + => new( + static (_, next) => async context => + { + var probe = context.Services.GetRequiredService(); + var capture = context.Services.GetRequiredService(); + capture.ReceiveProbeId = probe.Id; + await next(context); + }, + "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 ReceivePooledFeatureMiddleware + { + 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; + + /// + /// 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; + } + } }