Fixed in-memory outbox telemetry to capture ExecutionContext, which should properly flow activity traces

Author: phatboyg

src/MassTransit/Middleware/InMemoryOutbox/InMemoryOutboxConsumeContext.cs

@@ -14,8 +14,8 @@ public class InMemoryOutboxConsumeContext :
         OutboxContext
     {
         readonly TaskCompletionSource<InMemoryOutboxConsumeContext> _clearToSend;
+        readonly InMemoryOutboxDeferredMethodCollection _deferredMethods;
         readonly InMemoryOutboxMessageSchedulerContext _outboxSchedulerContext;
-        readonly List<Func<Task>> _pendingActions;
 
         protected InMemoryOutboxConsumeContext(ConsumeContext context)
             : base(context)
@@ -27,9 +27,10 @@ protected InMemoryOutboxConsumeContext(ConsumeContext context)
             ReceiveContext = outboxReceiveContext;
             PublishEndpointProvider = outboxReceiveContext.PublishEndpointProvider;
 
-            _pendingActions = [];
             _clearToSend = TaskUtil.GetTask<InMemoryOutboxConsumeContext>();
 
+            _deferredMethods = new InMemoryOutboxDeferredMethodCollection(_clearToSend.Task);
+
             if (context.TryGetPayload(out MessageSchedulerContext schedulerContext))
             {
                 _outboxSchedulerContext = (InMemoryOutboxMessageSchedulerContext)context.AddOrUpdatePayload<MessageSchedulerContext>(
@@ -44,45 +45,14 @@ protected InMemoryOutboxConsumeContext(ConsumeContext context)
 
         public Task Add(Func<Task> method)
         {
-            if (_clearToSend.Task.IsCompleted)
-                return method();
-
-            lock (_pendingActions)
-            {
-                _pendingActions.Add(method);
-
-                return Task.CompletedTask;
-            }
+            return _deferredMethods.Add(method);
         }
 
         public virtual async Task ExecutePendingActions(bool concurrentMessageDelivery)
         {
             _clearToSend.TrySetResult(this);
 
-            Func<Task>[] pendingActions;
-            lock (_pendingActions)
-                pendingActions = _pendingActions.ToArray();
-
-            if (pendingActions.Length > 0)
-            {
-                if (concurrentMessageDelivery)
-                {
-                    var collection = new PendingTaskCollection(pendingActions.Length);
-
-                    collection.Add(pendingActions.Select(action => action()));
-
-                    await collection.Completed().ConfigureAwait(false);
-                }
-                else
-                {
-                    foreach (Func<Task> action in pendingActions)
-                    {
-                        var task = action();
-                        if (task != null)
-                            await task.ConfigureAwait(false);
-                    }
-                }
-            }
+            await _deferredMethods.Execute(concurrentMessageDelivery).ConfigureAwait(false);
 
             if (_outboxSchedulerContext != null)
             {
@@ -99,8 +69,7 @@ public virtual async Task ExecutePendingActions(bool concurrentMessageDelivery)
 
         public virtual async Task DiscardPendingActions()
         {
-            lock (_pendingActions)
-                _pendingActions.Clear();
+            await _deferredMethods.Discard().ConfigureAwait(false);
 
             if (_outboxSchedulerContext != null)
             {

src/MassTransit/Middleware/InMemoryOutbox/InMemoryOutboxDeferredMethod.cs

@@ -0,0 +1,43 @@
+namespace MassTransit.Middleware.InMemoryOutbox;
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+
+public class InMemoryOutboxDeferredMethod :
+    IDisposable
+{
+    readonly Func<Task> _method;
+    ExecutionContext _executionContext;
+
+    public InMemoryOutboxDeferredMethod(ExecutionContext executionContext, Func<Task> method)
+    {
+        _executionContext = executionContext;
+        _method = method;
+    }
+
+    public void Dispose()
+    {
+        _executionContext?.Dispose();
+    }
+
+    public async Task Run()
+    {
+        try
+        {
+            using var ec = _executionContext.CreateCopy();
+
+            Task task = null;
+
+            ExecutionContext.Run(ec, _ => task = _method(), null);
+
+            await task.ConfigureAwait(false);
+        }
+        finally
+        {
+            _executionContext?.Dispose();
+            _executionContext = null;
+        }
+    }
+}

src/MassTransit/Middleware/InMemoryOutbox/InMemoryOutboxDeferredMethodCollection.cs

@@ -0,0 +1,86 @@
+#nullable enable
+namespace MassTransit.Middleware.InMemoryOutbox;
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Util;
+
+
+public class InMemoryOutboxDeferredMethodCollection
+{
+    readonly Task? _clearToSend;
+    readonly List<InMemoryOutboxDeferredMethod> _pendingMethods;
+
+    public InMemoryOutboxDeferredMethodCollection(Task? clearToSend = null)
+    {
+        _clearToSend = clearToSend;
+        _pendingMethods = [];
+    }
+
+    public Task Add(Func<Task> method)
+    {
+        if (_clearToSend?.IsCompleted ?? false)
+            return method();
+
+        var executionContext = ExecutionContext.Capture();
+
+        lock (_pendingMethods)
+        {
+            _pendingMethods.Add(new InMemoryOutboxDeferredMethod(executionContext, method));
+
+            return Task.CompletedTask;
+        }
+    }
+
+    public async Task Execute(bool concurrent)
+    {
+        InMemoryOutboxDeferredMethod[] pendingActions;
+        lock (_pendingMethods)
+            pendingActions = _pendingMethods.ToArray();
+
+        try
+        {
+            if (pendingActions.Length > 0)
+            {
+                if (concurrent)
+                {
+                    var collection = new PendingTaskCollection(pendingActions.Length);
+
+                    collection.Add(pendingActions.Select(method => method.Run()));
+
+                    await collection.Completed().ConfigureAwait(false);
+                }
+                else
+                {
+                    foreach (var method in pendingActions)
+                    {
+                        var task = method.Run();
+                        if (task != null)
+                            await task.ConfigureAwait(false);
+                    }
+                }
+            }
+        }
+        finally
+        {
+            foreach (var deferredMethod in pendingActions)
+                deferredMethod.Dispose();
+        }
+    }
+
+    public async Task Discard()
+    {
+        InMemoryOutboxDeferredMethod[] pendingMethods;
+        lock (_pendingMethods)
+        {
+            pendingMethods = _pendingMethods.ToArray();
+            _pendingMethods.Clear();
+        }
+
+        foreach (var method in pendingMethods)
+            method.Dispose();
+    }
+}

src/MassTransit/Middleware/InMemoryOutbox/InMemoryOutboxMessageSchedulerContext.cs

@@ -10,7 +10,7 @@
     public class InMemoryOutboxMessageSchedulerContext :
         MessageSchedulerContext
     {
-        readonly List<Func<Task>> _cancelMessages;
+        readonly InMemoryOutboxDeferredMethodCollection _cancelMessages;
         readonly Task _clearToSend;
         readonly Uri _inputAddress;
         readonly object _listLock = new object();
@@ -26,8 +26,8 @@ public InMemoryOutboxMessageSchedulerContext(ConsumeContext consumeContext, Mess
 
             _scheduler = new Lazy<IMessageScheduler>(() => schedulerFactory(consumeContext));
 
-            _scheduledMessages = new List<ScheduledMessage>();
-            _cancelMessages = new List<Func<Task>>();
+            _scheduledMessages = [];
+            _cancelMessages = new InMemoryOutboxDeferredMethodCollection();
         }
 
         public MessageSchedulerFactory SchedulerFactory { get; }
@@ -422,20 +422,7 @@ public Task CancelAllScheduledMessages()
 
         public Task ExecutePendingActions()
         {
-            Func<Task>[] cancelMessages;
-            lock (_listLock)
-            {
-                if (_cancelMessages.Count == 0)
-                    return Task.CompletedTask;
-
-                cancelMessages = _cancelMessages.ToArray();
-            }
-
-            var tasks = new PendingTaskCollection(cancelMessages.Length);
-            foreach (Func<Task> cancel in cancelMessages)
-                tasks.Add(cancel());
-
-            return tasks.Completed();
+            return _cancelMessages.Execute(true);
         }
     }
 }

tests/MassTransit.Tests/ContainerTests/Batch_Specs.cs

@@ -1,6 +1,7 @@
 namespace MassTransit.Tests.ContainerTests
 {
     using System;
+    using System.Diagnostics;
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
@@ -49,6 +50,7 @@ public class When_retry_and_in_memory_outbox_are_used_with_batch_consumers
         public async Task Should_deliver_the_batch_to_the_consumer()
         {
             await using var provider = new ServiceCollection()
+                .AddTelemetryListener()
                 .AddMassTransitTestHarness(x =>
                 {
                     x.AddConsumer<TestOutboxBatchConsumer>();
@@ -141,6 +143,8 @@ class TestOutboxBatchConsumer :
         {
             public Task Consume(ConsumeContext<Batch<BatchItem>> context)
             {
+                using var activity = new Activity("TestOutboxBatchConsumer").Start();
+
                 if (context.TryGetPayload<InMemoryOutboxConsumeContext>(out _))
                 {
                     context.Respond(new BatchResult