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;
+ }
+ }
}