From dbf3971398022a17c826c5a8cc8fe48384959545 Mon Sep 17 00:00:00 2001 From: Pascal Senn Date: Mon, 22 Jun 2026 01:41:15 +0200 Subject: [PATCH 1/4] Fixed Batch Consumer Middleware Pipeline --- .../src/Mocha.Inbox/ConsumeInboxMiddleware.cs | 95 ++++++++ .../Context/IBatchConsumeContext.cs | 17 ++ .../Context/IBatchConsumeContext~1.cs | 18 ++ .../src/Mocha/Abstractions/IMessageBatch.cs | 36 ++- .../Consumers/Batching/BatchConsumeContext.cs | 81 +++++++ src/Mocha/src/Mocha/Consumers/Consumer.cs | 8 +- .../Implementations/BatchConsumer.cs | 47 +++- .../ActivityMessagingDiagnosticListener.cs | 59 ++++- .../Mocha/Middlewares/DefaultMessageBus.cs | 4 +- .../Batching/BatchConsumerIntegrationTests.cs | 220 ++++++++++++++++++ .../Inbox/InboxIntegrationTests.cs | 67 +++++- .../Outbox/OutboxIntegrationTests.cs | 99 +++++++- .../Telemetry/OpenTelemetryTests.cs | 69 +++++- 13 files changed, 801 insertions(+), 19 deletions(-) create mode 100644 src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext.cs create mode 100644 src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext~1.cs create mode 100644 src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs diff --git a/src/Mocha/src/Mocha.Inbox/ConsumeInboxMiddleware.cs b/src/Mocha/src/Mocha.Inbox/ConsumeInboxMiddleware.cs index de75fd029d1..fa1e81a5bf2 100644 --- a/src/Mocha/src/Mocha.Inbox/ConsumeInboxMiddleware.cs +++ b/src/Mocha/src/Mocha.Inbox/ConsumeInboxMiddleware.cs @@ -61,6 +61,12 @@ public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate nex return; } + if (context is IBatchConsumeContext batchContext) + { + await InvokeBatchAsync(batchContext, context, next, activity); + return; + } + var messageId = context.MessageId; if (messageId is null) @@ -133,6 +139,95 @@ public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate nex await next(context); } + private async ValueTask InvokeBatchAsync( + IBatchConsumeContext batchContext, + IConsumeContext context, + ConsumerDelegate next, + Activity? activity) + { + var consumerType = GetConsumerType(context); + var inbox = context.Services.GetRequiredService(); + var claimed = 0; + var skipped = 0; + var untracked = 0; + + var batch = batchContext.Message; + + for (var i = 0; i < batch.Count; i++) + { + var itemContext = batch.GetContext(i); + var messageId = itemContext.MessageId; + + if (messageId is null) + { + untracked++; + continue; + } + + if (itemContext.Envelope is not { } envelope) + { + try + { + if (await inbox.ExistsAsync(messageId, consumerType, context.CancellationToken)) + { + logger.MessageSkippedDueToInboxExists(messageId, consumerType); + skipped++; + } + else + { + untracked++; + } + } + catch (OperationCanceledException) + { + throw; + } + catch (Exception ex) + { + logger.InboxExistsCheckFailed(messageId, ex); + untracked++; + } + + continue; + } + + try + { + if (await inbox.TryClaimAsync(envelope, consumerType, context.CancellationToken)) + { + claimed++; + } + else + { + logger.MessageSkippedDueToInboxExists(messageId, consumerType); + skipped++; + } + } + catch (OperationCanceledException) + { + throw; + } + catch (Exception ex) + { + logger.InboxClaimFailed(messageId, ex); + untracked++; + } + } + + activity?.SetTag("inbox.batch.claimed", claimed); + activity?.SetTag("inbox.batch.skipped", skipped); + activity?.SetTag("inbox.batch.untracked", untracked); + + if (claimed == 0 && untracked == 0) + { + activity?.SetTag("inbox.claimed", false); + return; + } + + activity?.SetTag("inbox.claimed", true); + await next(context); + } + /// /// Gets the consumer type name from the current context by reading the /// identity. diff --git a/src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext.cs b/src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext.cs new file mode 100644 index 00000000000..3ddca396b56 --- /dev/null +++ b/src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext.cs @@ -0,0 +1,17 @@ +namespace Mocha; + +/// +/// Represents the consume context for a batch of messages. +/// +public interface IBatchConsumeContext : IConsumeContext +{ + /// + /// Gets the logical identifier of this batch operation. + /// + string BatchId { get; } + + /// + /// Gets the message type metadata for each item in the batch. + /// + MessageType? ItemMessageType { get; } +} diff --git a/src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext~1.cs b/src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext~1.cs new file mode 100644 index 00000000000..64e1de9be69 --- /dev/null +++ b/src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext~1.cs @@ -0,0 +1,18 @@ +namespace Mocha; + +/// +/// Represents the consume context for a typed batch of messages. +/// +/// The message type contained in the batch. +public interface IBatchConsumeContext : + IBatchConsumeContext, + IConsumeContext> +{ + /// + /// Gets the typed message batch. + /// + new IMessageBatch Message { get; } + + /// + IMessageBatch IConsumeContext.Message => Message; +} diff --git a/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs b/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs index 58cee63295e..1fb4bfdea27 100644 --- a/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs +++ b/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs @@ -1,21 +1,47 @@ namespace Mocha; /// -/// An immutable batch of messages delivered to an . -/// Provides indexed access, count, enumeration, and batch metadata. +/// An immutable batch of messages delivered to a batch handler. /// -/// The type of event in the batch. -public interface IMessageBatch : IReadOnlyList +public interface IMessageBatch { /// /// Gets the reason this batch was dispatched. /// BatchCompletionMode CompletionMode { get; } + /// + /// Gets the number of messages in the batch. + /// + int Count { get; } + /// /// Gets the consume context for a specific message in the batch. /// /// The zero-based index of the message. /// The consume context for the message. - IConsumeContext GetContext(int index); + IConsumeContext GetContext(int index); +} + +/// +/// An immutable batch of messages delivered to an . +/// Provides indexed access, count, enumeration, and batch metadata. +/// +/// The type of event in the batch. +public interface IMessageBatch : IMessageBatch, IReadOnlyList +{ + /// + /// Gets the number of messages in the batch. + /// + new int Count { get; } + + /// + /// Gets the typed consume context for a specific message in the batch. + /// + /// The zero-based index of the message. + /// The consume context for the message. + new IConsumeContext GetContext(int index); + + /// + IConsumeContext IMessageBatch.GetContext(int index) => GetContext(index); } diff --git a/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs b/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs new file mode 100644 index 00000000000..74b5b31e3db --- /dev/null +++ b/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs @@ -0,0 +1,81 @@ +using Mocha.Features; +using Mocha.Middlewares; + +namespace Mocha; + +internal sealed class BatchConsumeContext : IBatchConsumeContext +{ + private readonly Headers _headers = new(0); + private readonly FeatureCollection _features = new(); + + public BatchConsumeContext( + MessageBatch message, + IServiceProvider services, + IConsumeContext firstContext, + string batchId, + MessageType? itemMessageType, + CancellationToken cancellationToken) + { + Message = message; + Services = services; + Runtime = firstContext.Runtime; + Transport = firstContext.Transport; + Endpoint = firstContext.Endpoint; + Host = Runtime.Host; + BatchId = batchId; + ItemMessageType = itemMessageType; + CancellationToken = cancellationToken; + } + + public IMessageBatch Message { get; } + + public string BatchId { get; } + + public MessageType? ItemMessageType { get; } + + public IFeatureCollection Features => _features; + + public MessagingTransport Transport { get; set; } + + public ReceiveEndpoint Endpoint { get; set; } + + public string? MessageId { get; set; } + + public string? CorrelationId { get; set; } + + public string? ConversationId { get; set; } + + public string? CausationId { get; set; } + + public Uri? SourceAddress { get; set; } + + public Uri? DestinationAddress { get; set; } + + public Uri? ResponseAddress { get; set; } + + public Uri? FaultAddress { get; set; } + + public MessageContentType? ContentType { get; set; } + + public MessageType? MessageType { get; set; } + + public IReadOnlyHeaders Headers => _headers; + + public DateTimeOffset? SentAt { get; set; } + + public DateTimeOffset? DeliverBy { get; set; } + + public int? DeliveryCount { get; set; } + + public ReadOnlyMemory Body => ReadOnlyMemory.Empty; + + public MessageEnvelope? Envelope { get; set; } + + public IRemoteHostInfo Host { get; set; } + + public IMessagingRuntime Runtime { get; set; } + + public CancellationToken CancellationToken { get; set; } + + public IServiceProvider Services { get; set; } +} diff --git a/src/Mocha/src/Mocha/Consumers/Consumer.cs b/src/Mocha/src/Mocha/Consumers/Consumer.cs index 9afc99cec40..a52416c3902 100644 --- a/src/Mocha/src/Mocha/Consumers/Consumer.cs +++ b/src/Mocha/src/Mocha/Consumers/Consumer.cs @@ -56,7 +56,7 @@ protected virtual void Configure(IConsumerDescriptor descriptor) { } /// public Type Identity { get; private set; } = null!; - private ConsumerDelegate _pipeline = null!; + private protected ConsumerDelegate Pipeline { get; private set; } = null!; /// /// Handles an incoming message after the consume middleware pipeline has completed. @@ -84,14 +84,14 @@ protected virtual void Configure(IConsumerDescriptor descriptor) { } /// /// Thrown when the context does not implement . /// - public async ValueTask ProcessAsync(IReceiveContext context) + public virtual async ValueTask ProcessAsync(IReceiveContext context) { if (context is not IConsumeContext handlerContext) { throw ThrowHelper.InvalidHandlerContext(); } - await _pipeline(handlerContext); + await Pipeline(handlerContext); } /// @@ -176,7 +176,7 @@ internal void Complete(IMessagingSetupContext context) Consumer = this }; - _pipeline = MiddlewareCompiler.CompileHandler( + Pipeline = MiddlewareCompiler.CompileHandler( middlewareFactoryContext, ConsumeAsync, [context.GetConsumerMiddlewares(), Configuration!.ConsumerMiddlewares], diff --git a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs index 017e0505f0a..4c886cf93cc 100644 --- a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs +++ b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs @@ -1,6 +1,8 @@ using System.Threading.Channels; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Logging; +using Mocha.Features; +using Mocha.Middlewares; using Mocha.Threading; namespace Mocha; @@ -14,13 +16,15 @@ namespace Mocha; /// without any modifications to the middleware chain. /// internal sealed class BatchConsumer( - Action? configure = null) : Consumer where THandler : IBatchEventHandler + Action? configure = null) + : Consumer where THandler : IBatchEventHandler { private BatchCollector _collector = null!; private Channel> _channel = null!; private ChannelProcessor> _processor = null!; private IServiceProvider _applicationServices = null!; private ILogger _logger = null!; + private MessageType? _itemMessageType; protected override void Configure(IConsumerDescriptor descriptor) { @@ -57,13 +61,26 @@ protected override void OnAfterInitialize(IMessagingSetupContext context) options.MaxConcurrentBatches); _collector = new BatchCollector(options, batch => _channel.Writer.WriteAsync(batch), timeProvider); + _itemMessageType = context.Messages.GetMessageType(typeof(TEvent)); } protected override async ValueTask ConsumeAsync(IConsumeContext context) { + var batchContext = (IBatchConsumeContext)context; + var handler = context.Services.GetRequiredService(); + await handler.HandleAsync(batchContext.Message, context.CancellationToken); + } + + public override async ValueTask ProcessAsync(IReceiveContext context) + { + if (context is not IConsumeContext consumeContext) + { + throw ThrowHelper.InvalidHandlerContext(); + } + // we dispose the consume context so the reference is free as soon as we leave consume // as then the context will be returned to the pool - using var batchContext = new ConsumeContext(context); + using var batchContext = new ConsumeContext(consumeContext); // force deserialization to keep the work outside of the batch and also verify that // the message can be deserialized before adding to the batch @@ -80,8 +97,30 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok _logger.DispatchingBatch(batch.Count, batch.CompletionMode); await using var scope = _applicationServices.CreateAsyncScope(); - var handler = scope.ServiceProvider.GetRequiredService(); - await handler.HandleAsync(batch, cancellationToken); + + var batchContext = new BatchConsumeContext( + batch, + scope.ServiceProvider, + batch.GetContext(0), + Guid.NewGuid().ToString(), + _itemMessageType, + cancellationToken); + + var consumerFeature = batchContext.Features.GetOrSet(); + consumerFeature.CurrentConsumer = this; + + var accessor = scope.ServiceProvider.GetRequiredService(); + var previousContext = accessor.Context; + accessor.Context = batchContext; + + try + { + await Pipeline(batchContext); + } + finally + { + accessor.Context = previousContext; + } foreach (var entry in batch.Entries) { diff --git a/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs b/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs index 2179eb3b907..5cf03c653c7 100644 --- a/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs +++ b/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs @@ -1,4 +1,6 @@ +using System.Collections.Immutable; using System.Diagnostics; +using Mocha.Features; using Mocha.Middlewares; namespace Mocha; @@ -67,6 +69,11 @@ private ReceiveActivity(IReceiveContext context) activity ??= OpenTelemetry.Source.StartActivity($"{context.Endpoint.Address} receive", ActivityKind.Client); _activity = activity; + + if (activity is not null) + { + context.Features.GetOrSet().ActivityContext = activity.Context; + } } public void Dispose() @@ -95,7 +102,14 @@ private ConsumerActivity(IConsumeContext context) // TODO this can be done better! var currentConsumer = context.Features.Get()?.CurrentConsumer?.Name ?? "unknown"; - _activity = OpenTelemetry.Source.StartActivity($"{currentConsumer} consume", ActivityKind.Consumer); + _activity = context is IBatchConsumeContext batchContext + ? OpenTelemetry.Source.StartActivity( + $"{currentConsumer} consume", + ActivityKind.Consumer, + default(ActivityContext), + null, + CreateLinks(batchContext)) + : OpenTelemetry.Source.StartActivity($"{currentConsumer} consume", ActivityKind.Consumer); } public void Dispose() @@ -103,19 +117,60 @@ public void Dispose() // Enrich activity with context state after all middlewares have run if (_activity is not null) { - var consumerName = _context.MessageType is not null ? _context.MessageType.Identity : "unknown"; + var batchConsumeContext = _context as IBatchConsumeContext; + var consumerName = batchConsumeContext?.ItemMessageType?.Identity + ?? _context.MessageType?.Identity + ?? "unknown"; _activity .EnrichMessageDefault() .SetMessageId(_context.MessageId ?? string.Empty) .SetConversationId(_context.CorrelationId ?? string.Empty) .SetConsumerName(consumerName); + + if (batchConsumeContext is not null) + { + _activity.SetTag("messaging.batch.id", batchConsumeContext.BatchId); + _activity.SetTag("messaging.batch.message_count", batchConsumeContext.Message.Count); + } } _activity?.Dispose(); } public static ConsumerActivity Create(IConsumeContext context) => new(context); + + private static ImmutableArray CreateLinks(IBatchConsumeContext batchContext) + { + var links = ImmutableArray.CreateBuilder(batchContext.Message.Count); + var batch = batchContext.Message; + + for (var i = 0; i < batch.Count; i++) + { + if (batch.GetContext(i).Features.TryGet(out var linkContext) + && linkContext.ActivityContext is { } activityContext) + { + links.Add(new ActivityLink(activityContext)); + } + } + + return links.ToImmutable(); + } + } + + private sealed class ReceiveActivityFeature : IPooledFeature + { + public ActivityContext? ActivityContext { get; set; } + + public void Initialize(object state) + { + ActivityContext = null; + } + + public void Reset() + { + ActivityContext = null; + } } private sealed class DispatchActivity : IDisposable diff --git a/src/Mocha/src/Mocha/Middlewares/DefaultMessageBus.cs b/src/Mocha/src/Mocha/Middlewares/DefaultMessageBus.cs index 2bb6363b01b..b0a09fde230 100644 --- a/src/Mocha/src/Mocha/Middlewares/DefaultMessageBus.cs +++ b/src/Mocha/src/Mocha/Middlewares/DefaultMessageBus.cs @@ -436,7 +436,9 @@ private void PropagateCorrelationIds(DispatchContext context) if (consumeContextAccessor.Context is { } ambient) { context.ConversationId ??= ambient.ConversationId; - context.CausationId ??= ambient.MessageId; + context.CausationId ??= ambient is IBatchConsumeContext batchContext + ? batchContext.BatchId + : ambient.MessageId; } } } diff --git a/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchConsumerIntegrationTests.cs b/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchConsumerIntegrationTests.cs index 935e1ca2691..0121db5025b 100644 --- a/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchConsumerIntegrationTests.cs +++ b/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchConsumerIntegrationTests.cs @@ -116,6 +116,76 @@ public async Task Handler_Should_ReceiveMetadata_When_BatchDelivered() Assert.Equal("meta-1", batch[0].Id); } + [Fact] + public async Task ConsumePipeline_Should_UseBatchContext_When_BatchDelivered() + { + // arrange + var recorder = new BatchConsumeContextRecorder(); + await using var provider = await CreateBusAsync( + b => + { + b.Services.AddSingleton(recorder); + b.AddBatchHandler(opts => + { + opts.MaxBatchSize = 2; + opts.BatchTimeout = TimeSpan.FromSeconds(5); + }); + b.ConfigureMessageBus(h => + h.UseConsume( + new ConsumerMiddlewareConfiguration( + static (_, next) => + async ctx => + { + var recorder = ctx.Services.GetRequiredService(); + recorder.RecordMiddleware(ctx); + await next(ctx); + }, + "BatchContextCapture"))); + }, + t => t.Endpoint("batch-ep").Handler().MaxConcurrency(2)); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync( + new TestBatchEvent { Id = "ctx-1" }, + new PublishOptions { Headers = new Dictionary { ["x-entry"] = "first" } }, + CancellationToken.None); + await bus.PublishAsync( + new TestBatchEvent { Id = "ctx-2" }, + new PublishOptions { Headers = new Dictionary { ["x-entry"] = "second" } }, + CancellationToken.None); + + // assert + Assert.True(await recorder.WaitForHandlerAsync(Timeout), "Batch handler was not invoked within timeout"); + + var middleware = Assert.Single(recorder.MiddlewareSnapshots); + Assert.True(middleware.IsBatchContext); + Assert.True(middleware.IsTypedBatchContext); + Assert.True(middleware.NonGenericMessageIsContextMessage); + Assert.Equal(2, middleware.BatchCount); + Assert.NotNull(middleware.BatchId); + Assert.NotNull(middleware.ItemMessageTypeIdentity); + Assert.Null(middleware.MessageId); + Assert.Null(middleware.CorrelationId); + Assert.Null(middleware.ConversationId); + Assert.Null(middleware.CausationId); + Assert.Null(middleware.MessageTypeIdentity); + Assert.True(middleware.EnvelopeIsNull); + Assert.True(middleware.BodyIsEmpty); + Assert.Equal(0, middleware.HeaderCount); + Assert.Equal(2, middleware.EntryMessageIds.Count); + Assert.All(middleware.EntryMessageIds, Assert.NotNull); + Assert.Equal(["first", "second"], middleware.EntryHeaderValues.OrderBy(static x => x).ToArray()); + Assert.All(middleware.EntryMessageTypeIdentities, Assert.NotNull); + + var handler = Assert.Single(recorder.HandlerSnapshots); + Assert.True(handler.AccessorContextIsBatchContext); + Assert.True(handler.AccessorMessageIsHandlerBatch); + Assert.True(handler.NonGenericMessageIsHandlerBatch); + } + [Fact] public async Task Handler_Should_ResolveDependencies_When_InvokedThroughDI() { @@ -327,6 +397,18 @@ public ValueTask HandleAsync(IMessageBatch batch, CancellationTo } } + public sealed class ContextCapturingBatchHandler( + BatchConsumeContextRecorder recorder, + ConsumeContextAccessor accessor) + : IBatchEventHandler + { + public ValueTask HandleAsync(IMessageBatch batch, CancellationToken cancellationToken) + { + recorder.RecordHandler(batch, accessor.Context); + return default; + } + } + public sealed class ThrowingBatchHandler([FromKeyedServices("throwing")] BatchMessageRecorder recorder) : IBatchEventHandler { @@ -346,4 +428,142 @@ public ValueTask HandleAsync(IMessageBatch batch, CancellationT return default; } } + + public sealed class BatchConsumeContextRecorder + { + private readonly object _sync = new(); + private readonly SemaphoreSlim _handlerSemaphore = new(0); + private readonly List _middlewareSnapshots = []; + private readonly List _handlerSnapshots = []; + + public IReadOnlyList MiddlewareSnapshots + { + get + { + lock (_sync) + { + return _middlewareSnapshots.ToList(); + } + } + } + + public IReadOnlyList HandlerSnapshots + { + get + { + lock (_sync) + { + return _handlerSnapshots.ToList(); + } + } + } + + public void RecordMiddleware(IConsumeContext context) + { + lock (_sync) + { + _middlewareSnapshots.Add(BatchContextSnapshot.From(context)); + } + } + + public void RecordHandler(IMessageBatch batch, IConsumeContext? ambientContext) + { + lock (_sync) + { + _handlerSnapshots.Add(BatchHandlerContextSnapshot.From(batch, ambientContext)); + } + + _handlerSemaphore.Release(); + } + + public async Task WaitForHandlerAsync(TimeSpan timeout) + { + return await _handlerSemaphore.WaitAsync(timeout); + } + } + + public sealed record BatchContextSnapshot( + bool IsBatchContext, + bool IsTypedBatchContext, + bool NonGenericMessageIsContextMessage, + int BatchCount, + string? BatchId, + string? ItemMessageTypeIdentity, + string? MessageId, + string? CorrelationId, + string? ConversationId, + string? CausationId, + string? MessageTypeIdentity, + bool EnvelopeIsNull, + bool BodyIsEmpty, + int HeaderCount, + IReadOnlyList EntryMessageIds, + IReadOnlyList EntryHeaderValues, + IReadOnlyList EntryMessageTypeIdentities) + { + public static BatchContextSnapshot From(IConsumeContext context) + { + var batchContext = context as IBatchConsumeContext; + var typedBatchContext = context as IBatchConsumeContext; + var batch = batchContext?.Message; + var entryMessageIds = new List(); + var entryHeaderValues = new List(); + var entryMessageTypeIdentities = new List(); + + if (batch is not null) + { + for (var i = 0; i < batch.Count; i++) + { + var itemContext = batch.GetContext(i); + entryMessageIds.Add(itemContext.MessageId); + entryHeaderValues.Add(itemContext.Headers.GetValue("x-entry")); + entryMessageTypeIdentities.Add(itemContext.MessageType?.Identity); + } + } + + return new BatchContextSnapshot( + batchContext is not null, + typedBatchContext is not null, + batchContext is IConsumeContext nonGenericContext + && ReferenceEquals(nonGenericContext.Message, batch), + batch?.Count ?? 0, + batchContext?.BatchId, + batchContext?.ItemMessageType?.Identity, + context.MessageId, + context.CorrelationId, + context.ConversationId, + context.CausationId, + context.MessageType?.Identity, + context.Envelope is null, + context.Body.IsEmpty, + context.Headers.Count(), + entryMessageIds, + entryHeaderValues, + entryMessageTypeIdentities); + } + } + + public sealed record BatchHandlerContextSnapshot( + bool AccessorContextIsBatchContext, + bool AccessorMessageIsHandlerBatch, + bool NonGenericMessageIsHandlerBatch) + { + public static BatchHandlerContextSnapshot From( + IMessageBatch batch, + IConsumeContext? ambientContext) + { + var batchContext = ambientContext as IBatchConsumeContext; + IMessageBatch? nonGenericMessage = null; + + if (batchContext is not null) + { + nonGenericMessage = ((IConsumeContext)batchContext).Message; + } + + return new BatchHandlerContextSnapshot( + batchContext is not null, + batchContext is not null && ReferenceEquals(batchContext.Message, batch), + ReferenceEquals(nonGenericMessage, batch)); + } + } } diff --git a/src/Mocha/test/Mocha.Tests/Inbox/InboxIntegrationTests.cs b/src/Mocha/test/Mocha.Tests/Inbox/InboxIntegrationTests.cs index 8112c2f4a96..10e9c8fb949 100644 --- a/src/Mocha/test/Mocha.Tests/Inbox/InboxIntegrationTests.cs +++ b/src/Mocha/test/Mocha.Tests/Inbox/InboxIntegrationTests.cs @@ -97,6 +97,43 @@ await recorder.WaitAsync(s_timeout, expectedCount: 3), Assert.Equal(3, inbox.RecordedEnvelopes.Count); } + [Fact] + public async Task Inbox_Should_RecordEachBatchEntry_When_BatchReceived() + { + // arrange + var inbox = new InMemoryMessageInbox(); + var recorder = new BatchMessageRecorder(); + await using var provider = await CreateBusWithInboxAsync( + inbox, + b => + { + b.Services.AddSingleton(recorder); + b.AddBatchHandler(opts => + { + opts.MaxBatchSize = 2; + opts.BatchTimeout = TimeSpan.FromSeconds(5); + }); + }, + t => t.Endpoint("batch-ep").Handler().MaxConcurrency(2)); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new InboxBatchTestEvent { Payload = "batch-1" }, CancellationToken.None); + await bus.PublishAsync(new InboxBatchTestEvent { Payload = "batch-2" }, CancellationToken.None); + + // assert + Assert.True(await recorder.WaitAsync(s_timeout), "Batch handler did not receive the events within timeout"); + + var batch = Assert.IsAssignableFrom>(Assert.Single(recorder.Batches)); + Assert.Equal(2, batch.Count); + + await WaitUntilAsync(() => inbox.RecordedEnvelopes.Count >= 2, s_timeout); + Assert.Equal(2, inbox.RecordedEnvelopes.Count); + Assert.Equal(2, inbox.RecordedEnvelopes.Select(static e => e.MessageId).Distinct().Count()); + } + [Fact] public async Task Inbox_Should_SkipRecording_When_SkipInboxFeatureSet() { @@ -156,6 +193,12 @@ private static async Task WaitUntilAsync(Func condition, TimeSpan timeout) private static async Task CreateBusWithInboxAsync( InMemoryMessageInbox inbox, Action configure) + => await CreateBusWithInboxAsync(inbox, configure, configureTransport: null); + + private static async Task CreateBusWithInboxAsync( + InMemoryMessageInbox inbox, + Action configure, + Action? configureTransport) { var services = new ServiceCollection(); services.AddSingleton(inbox); @@ -164,7 +207,14 @@ private static async Task CreateBusWithInboxAsync( builder.UseInboxCore(); configure(builder); - builder.AddInMemory(); + if (configureTransport is null) + { + builder.AddInMemory(); + } + else + { + builder.AddInMemory(configureTransport); + } var provider = services.BuildServiceProvider(); var runtime = (MessagingRuntime)provider.GetRequiredService(); @@ -181,6 +231,11 @@ public sealed class InboxTestEvent public required string Payload { get; init; } } + public sealed class InboxBatchTestEvent + { + public required string Payload { get; init; } + } + public sealed class InboxTestEventHandler(MessageRecorder recorder) : IEventHandler { public ValueTask HandleAsync(InboxTestEvent message, CancellationToken cancellationToken) @@ -189,4 +244,14 @@ public ValueTask HandleAsync(InboxTestEvent message, CancellationToken cancellat return default; } } + + public sealed class InboxBatchTestEventHandler(BatchMessageRecorder recorder) + : IBatchEventHandler + { + public ValueTask HandleAsync(IMessageBatch batch, CancellationToken cancellationToken) + { + recorder.Record(batch); + return default; + } + } } diff --git a/src/Mocha/test/Mocha.Tests/Outbox/OutboxIntegrationTests.cs b/src/Mocha/test/Mocha.Tests/Outbox/OutboxIntegrationTests.cs index 2a430751870..6c505885bcc 100644 --- a/src/Mocha/test/Mocha.Tests/Outbox/OutboxIntegrationTests.cs +++ b/src/Mocha/test/Mocha.Tests/Outbox/OutboxIntegrationTests.cs @@ -126,6 +126,46 @@ public async Task Outbox_Should_SignalWorker_When_MessagePersisted() Assert.True(signal.SignalCount >= 1, "Signal should have been set at least once"); } + [Fact] + public async Task Outbox_Should_UseBatchIdAsCausation_When_PublishedFromBatchHandler() + { + // arrange + var outbox = new InMemoryMessageOutbox(); + var recorder = new BatchOutboxRecorder(); + await using var provider = await CreateBusWithOutboxAsync( + outbox, + b => + { + b.Services.AddSingleton(recorder); + b.AddBatchHandler(opts => + { + opts.MaxBatchSize = 2; + opts.BatchTimeout = TimeSpan.FromSeconds(5); + }); + }, + t => t.Endpoint("batch-ep").Handler().MaxConcurrency(2)); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + var inputOptions = new PublishOptions + { + Headers = new Dictionary { ["x-skip-outbox"] = "true" } + }; + + // act + await bus.PublishAsync(new OutboxBatchInputEvent { Payload = "batch-1" }, inputOptions, CancellationToken.None); + await bus.PublishAsync(new OutboxBatchInputEvent { Payload = "batch-2" }, inputOptions, CancellationToken.None); + + // assert + Assert.True(await recorder.WaitAsync(s_timeout), "Batch handler did not publish from the batch"); + await WaitUntilAsync(() => !outbox.Envelopes.IsEmpty, s_timeout); + + var envelope = Assert.Single(outbox.Envelopes); + Assert.NotNull(recorder.BatchId); + Assert.Equal(2, recorder.BatchCount); + Assert.Equal(recorder.BatchId, envelope.CausationId); + } + // ══════════════════════════════════════════════════════════════════════ // Helpers // ══════════════════════════════════════════════════════════════════════ @@ -142,6 +182,12 @@ private static async Task WaitUntilAsync(Func condition, TimeSpan timeout) private static async Task CreateBusWithOutboxAsync( InMemoryMessageOutbox outbox, Action configure) + => await CreateBusWithOutboxAsync(outbox, configure, configureTransport: null); + + private static async Task CreateBusWithOutboxAsync( + InMemoryMessageOutbox outbox, + Action configure, + Action? configureTransport) { var services = new ServiceCollection(); services.AddSingleton(outbox); @@ -167,7 +213,14 @@ private static async Task CreateBusWithOutboxAsync( ); configure(builder); - builder.AddInMemory(); + if (configureTransport is null) + { + builder.AddInMemory(); + } + else + { + builder.AddInMemory(configureTransport); + } var provider = services.BuildServiceProvider(); var runtime = (MessagingRuntime)provider.GetRequiredService(); @@ -184,6 +237,11 @@ public sealed class OutboxTestEvent public required string Payload { get; init; } } + public sealed class OutboxBatchInputEvent + { + public required string Payload { get; init; } + } + public sealed class OutboxTestEventHandler(MessageRecorder recorder) : IEventHandler { public ValueTask HandleAsync(OutboxTestEvent message, CancellationToken cancellationToken) @@ -193,6 +251,45 @@ public ValueTask HandleAsync(OutboxTestEvent message, CancellationToken cancella } } + public sealed class OutboxBatchPublishingHandler( + IMessageBus bus, + ConsumeContextAccessor accessor, + BatchOutboxRecorder recorder) + : IBatchEventHandler + { + public async ValueTask HandleAsync(IMessageBatch batch, CancellationToken cancellationToken) + { + if (accessor.Context is not IBatchConsumeContext batchContext) + { + throw new InvalidOperationException("Expected an ambient batch consume context."); + } + + await bus.PublishAsync(new OutboxTestEvent { Payload = "from-batch" }, cancellationToken); + recorder.Record(batchContext.BatchId, batch.Count); + } + } + + public sealed class BatchOutboxRecorder + { + private readonly SemaphoreSlim _semaphore = new(0); + + public string? BatchId { get; private set; } + + public int BatchCount { get; private set; } + + public void Record(string batchId, int batchCount) + { + BatchId = batchId; + BatchCount = batchCount; + _semaphore.Release(); + } + + public async Task WaitAsync(TimeSpan timeout) + { + return await _semaphore.WaitAsync(timeout); + } + } + /// /// In-memory outbox that stores envelopes for test assertions. /// Optionally signals an on persist, diff --git a/src/Mocha/test/Mocha.Tests/Telemetry/OpenTelemetryTests.cs b/src/Mocha/test/Mocha.Tests/Telemetry/OpenTelemetryTests.cs index bc8ccb6cdbf..98767445b0b 100644 --- a/src/Mocha/test/Mocha.Tests/Telemetry/OpenTelemetryTests.cs +++ b/src/Mocha/test/Mocha.Tests/Telemetry/OpenTelemetryTests.cs @@ -125,6 +125,47 @@ public async Task MultiplePublishes_Should_CreateMultipleActivities_When_Publish Assert.True(activities.Count >= 3, $"Expected at least 3 activities but got {activities.Count}"); } + [Fact] + public async Task BatchConsume_Should_LinkReceiveActivities_When_BatchDelivered() + { + // arrange + var activities = new ConcurrentBag(); + using var listener = CreateListener(activities); + + var recorder = new BatchMessageRecorder(); + await using var provider = await CreateBusAsync( + b => + { + b.Services.AddSingleton(recorder); + b.AddBatchHandler(opts => + { + opts.MaxBatchSize = 2; + opts.BatchTimeout = TimeSpan.FromSeconds(5); + }); + }, + t => t.Endpoint("batch-ep").Handler().MaxConcurrency(2)); + + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + + // act + await bus.PublishAsync(new TracedBatchEvent { Data = "batch-1" }, CancellationToken.None); + await bus.PublishAsync(new TracedBatchEvent { Data = "batch-2" }, CancellationToken.None); + + Assert.True(await recorder.WaitAsync(s_timeout), "Batch handler did not receive the events"); + + // Task.Delay: ActivityListener callbacks fire asynchronously; brief wait lets all callbacks complete + await Task.Delay(200, TestContext.Current.CancellationToken); + + // assert + var batchActivity = activities.SingleOrDefault( + static a => Equals(a.GetTagItem("messaging.batch.message_count"), 2)); + + Assert.NotNull(batchActivity); + Assert.NotNull(batchActivity.GetTagItem("messaging.batch.id")); + Assert.Equal(2, batchActivity.Links.Count()); + } + [Fact] public async Task MessageBus_Should_ProcessMessages_When_NoListenerRegistered() { @@ -842,12 +883,24 @@ public void RecordConsumeMessage_With_Null_Subscription_Does_Not_Throw() private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); private static async Task CreateBusAsync(Action configure) + => await CreateBusAsync(configure, configureTransport: null); + + private static async Task CreateBusAsync( + Action configure, + Action? configureTransport) { var services = new ServiceCollection(); var builder = services.AddMessageBus(); configure(builder); builder.AddInstrumentation(); - builder.AddInMemory(); + if (configureTransport is null) + { + builder.AddInMemory(); + } + else + { + builder.AddInMemory(configureTransport); + } var provider = services.BuildServiceProvider(); var runtime = (MessagingRuntime)provider.GetRequiredService(); @@ -860,6 +913,11 @@ public sealed class TracedEvent public required string Data { get; init; } } + public sealed class TracedBatchEvent + { + public required string Data { get; init; } + } + public sealed class TracedRequest : IEventRequest { public required string Query { get; init; } @@ -879,6 +937,15 @@ public ValueTask HandleAsync(TracedEvent message, CancellationToken ct) } } + public sealed class TracedBatchHandler(BatchMessageRecorder recorder) : IBatchEventHandler + { + public ValueTask HandleAsync(IMessageBatch batch, CancellationToken ct) + { + recorder.Record(batch); + return default; + } + } + public sealed class TracedRequestHandler(MessageRecorder recorder) : IEventRequestHandler { From 81fff5d51bb807bb27647d7d9568c7194d75b166 Mon Sep 17 00:00:00 2001 From: Pascal Senn Date: Mon, 22 Jun 2026 01:44:40 +0200 Subject: [PATCH 2/4] Add generic IMessageBatch interface for typed message handling --- .../src/Mocha/Abstractions/IMessageBatch.cs | 23 ------------------ .../src/Mocha/Abstractions/IMessageBatch~1.cs | 24 +++++++++++++++++++ 2 files changed, 24 insertions(+), 23 deletions(-) create mode 100644 src/Mocha/src/Mocha/Abstractions/IMessageBatch~1.cs diff --git a/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs b/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs index 1fb4bfdea27..9c35cca23a5 100644 --- a/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs +++ b/src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs @@ -22,26 +22,3 @@ public interface IMessageBatch /// The consume context for the message. IConsumeContext GetContext(int index); } - -/// -/// An immutable batch of messages delivered to an . -/// Provides indexed access, count, enumeration, and batch metadata. -/// -/// The type of event in the batch. -public interface IMessageBatch : IMessageBatch, IReadOnlyList -{ - /// - /// Gets the number of messages in the batch. - /// - new int Count { get; } - - /// - /// Gets the typed consume context for a specific message in the batch. - /// - /// The zero-based index of the message. - /// The consume context for the message. - new IConsumeContext GetContext(int index); - - /// - IConsumeContext IMessageBatch.GetContext(int index) => GetContext(index); -} diff --git a/src/Mocha/src/Mocha/Abstractions/IMessageBatch~1.cs b/src/Mocha/src/Mocha/Abstractions/IMessageBatch~1.cs new file mode 100644 index 00000000000..efd9b6a2b37 --- /dev/null +++ b/src/Mocha/src/Mocha/Abstractions/IMessageBatch~1.cs @@ -0,0 +1,24 @@ +namespace Mocha; + +/// +/// An immutable batch of messages delivered to an . +/// Provides indexed access, count, enumeration, and batch metadata. +/// +/// The type of event in the batch. +public interface IMessageBatch : IMessageBatch, IReadOnlyList +{ + /// + /// Gets the number of messages in the batch. + /// + new int Count { get; } + + /// + /// Gets the typed consume context for a specific message in the batch. + /// + /// The zero-based index of the message. + /// The consume context for the message. + new IConsumeContext GetContext(int index); + + /// + IConsumeContext IMessageBatch.GetContext(int index) => GetContext(index); +} From b86bda18bfdca63a19bcde0a88dec76790f7db95 Mon Sep 17 00:00:00 2001 From: PascalSenn Date: Mon, 22 Jun 2026 01:55:30 +0200 Subject: [PATCH 3/4] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../Instrumentation/ActivityMessagingDiagnosticListener.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs b/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs index 5cf03c653c7..2f5cb3bc78c 100644 --- a/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs +++ b/src/Mocha/src/Mocha/Instrumentation/ActivityMessagingDiagnosticListener.cs @@ -124,8 +124,8 @@ public void Dispose() _activity .EnrichMessageDefault() - .SetMessageId(_context.MessageId ?? string.Empty) - .SetConversationId(_context.CorrelationId ?? string.Empty) + .SetMessageId(_context.MessageId) + .SetConversationId(_context.ConversationId) .SetConsumerName(consumerName); if (batchConsumeContext is not null) From 11c088e1929e7b9dda4b91e5b83de2dabaa450b3 Mon Sep 17 00:00:00 2001 From: PascalSenn Date: Mon, 22 Jun 2026 01:55:36 +0200 Subject: [PATCH 4/4] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs b/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs index 74b5b31e3db..f9c987abdd9 100644 --- a/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs +++ b/src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs @@ -21,7 +21,7 @@ public BatchConsumeContext( Runtime = firstContext.Runtime; Transport = firstContext.Transport; Endpoint = firstContext.Endpoint; - Host = Runtime.Host; + Host = firstContext.Host; BatchId = batchId; ItemMessageType = itemMessageType; CancellationToken = cancellationToken;