[Mocha] Fix the outbox worker missing messages written inside a TransactionScope

src/Mocha/src/Mocha.EntityFrameworkCore.Postgres/Outbox/PostgresMessageOutbox.cs

@@ -3,6 +3,7 @@
 using Microsoft.EntityFrameworkCore.Storage;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
+using Mocha.EntityFrameworkCore;
 using Mocha.Middlewares;
 using Mocha.Utils;
 using Npgsql;
@@ -16,8 +17,8 @@ namespace Mocha.Outbox;
 /// </summary>
 /// <remarks>
 /// When a database transaction is active on the DbContext, the insert participates in that transaction
-/// and the outbox signal is deferred until commit. Without an active transaction, the signal fires
-/// immediately after insert to wake the outbox processor.
+/// and the outbox signal is deferred until commit. Inside an ambient <see cref="System.Transactions.Transaction"/>,
+/// the signal fires when that transaction completes. Otherwise, the signal fires immediately after insert.
 /// </remarks>
 internal sealed class PostgresMessageOutbox : IMessageOutbox, IDisposable
 {
@@ -64,8 +65,9 @@ public static PostgresMessageOutbox Create(Type contextType, string optionsName,
     /// Serializes the message envelope and inserts it into the Postgres outbox table.
     /// </summary>
     /// <remarks>
-    /// If no database transaction is active, the outbox signal is raised immediately to wake the
-    /// processor. Otherwise, the signal is deferred to the transaction commit interceptor.
+    /// With an active database transaction, the signal is deferred to the transaction commit interceptor.
+    /// With an ambient <see cref="System.Transactions.Transaction"/>, it is raised when that transaction completes.
+    /// Otherwise, it is raised immediately.
     /// </remarks>
     /// <param name="envelope">The message envelope to persist in the outbox.</param>
     /// <param name="cancellationToken">A token to observe for cancellation.</param>
@@ -98,7 +100,7 @@ public async ValueTask PersistAsync(MessageEnvelope envelope, CancellationToken
 
                 if (_originalDbContext.Database.CurrentTransaction?.GetDbTransaction() is not NpgsqlTransaction)
                 {
-                    _signal.Set();
+                    _signal.SetAfterAmbientTransaction();
                 }
             }
             finally

src/Mocha/src/Mocha.EntityFrameworkCore/Assembly.cs

@@ -2,3 +2,4 @@
 
 [assembly: InternalsVisibleTo("Mocha.Tests")]
 [assembly: InternalsVisibleTo("Mocha.EntityFrameworkCore.Postgres.Tests")]
+[assembly: InternalsVisibleTo("Mocha.EntityFrameworkCore.Postgres")]

src/Mocha/src/Mocha.EntityFrameworkCore/Outbox/OutboxSaveChangesInterceptor.cs

@@ -28,7 +28,7 @@ public override int SavedChanges(SaveChangesCompletedEventData eventData, int re
 
         if (context.Database.CurrentTransaction is null)
         {
-            signal.Set();
+            signal.SetAfterAmbientTransaction();
         }
 
         return result;

src/Mocha/src/Mocha.EntityFrameworkCore/Outbox/OutboxSignalExtensions.cs

@@ -0,0 +1,26 @@
+using System.Transactions;
+using Mocha.Outbox;
+
+namespace Mocha.EntityFrameworkCore;
+
+internal static class OutboxSignalExtensions
+{
+    /// <summary>
+    /// Sets the signal when the ambient <see cref="Transaction"/> completes, or immediately when
+    /// there is no ambient transaction.
+    /// </summary>
+    /// <param name="signal">The outbox signal to set.</param>
+    public static void SetAfterAmbientTransaction(this IOutboxSignal signal)
+    {
+        if (Transaction.Current is { } transaction)
+        {
+            SetOnCompletion(signal, transaction);
+            return;
+        }
+
+        signal.Set();
+    }
+
+    private static void SetOnCompletion(IOutboxSignal signal, Transaction transaction)
+        => transaction.TransactionCompleted += (_, _) => signal.Set();
+}

src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/PostgresOutboxAmbientTransactionSignalTests.cs

@@ -0,0 +1,267 @@
+using System.Collections.Concurrent;
+using System.Text.Json;
+using System.Transactions;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Mocha.EntityFrameworkCore.Postgres.Tests.Helpers;
+using Mocha.Outbox;
+using Mocha.Transport.InMemory;
+using Npgsql;
+
+namespace Mocha.EntityFrameworkCore.Postgres.Tests;
+
+public sealed class PostgresOutboxAmbientTransactionSignalTests(PostgresFixture fixture)
+    : IClassFixture<PostgresFixture>
+{
+    private static CancellationToken TestToken => TestContext.Current.CancellationToken;
+
+    [Fact]
+    public async Task PersistAsync_Should_SignalOnlyAfterCommit_When_AmbientTransactionIsUsed()
+    {
+        // arrange
+        var connectionString = await fixture.CreateDatabaseAsync();
+        var signal = new VisibilityProbeSignal(connectionString);
+        await using var provider = await CreateProviderAsync(connectionString, signal, addHandler: false);
+        await EnsureCreatedAsync(provider);
+        signal.Events.Clear();
+
+        // act
+        using (var transaction = new TransactionScope(TransactionScopeAsyncFlowOption.Enabled))
+        {
+            await using var scope = provider.CreateAsyncScope();
+            var bus = scope.ServiceProvider.GetRequiredService<IMessageBus>();
+            await bus.PublishAsync(new TestEvent { Payload = "ambient" }, TestToken);
+            signal.Mark("before-complete");
+            transaction.Complete();
+        }
+
+        // assert
+        AssertSignaledOnlyAfterCommit(signal);
+    }
+
+    [Fact]
+    public async Task PersistAsync_Should_SignalOnlyAfterCommit_When_ManyMessagesShareAmbientTransaction()
+    {
+        // arrange
+        var connectionString = await fixture.CreateDatabaseAsync();
+        var signal = new VisibilityProbeSignal(connectionString);
+        await using var provider = await CreateProviderAsync(connectionString, signal, addHandler: false);
+        await EnsureCreatedAsync(provider);
+        signal.Events.Clear();
+
+        // act
+        using (var transaction = new TransactionScope(TransactionScopeAsyncFlowOption.Enabled))
+        {
+            await using var scope = provider.CreateAsyncScope();
+            var bus = scope.ServiceProvider.GetRequiredService<IMessageBus>();
+            for (var i = 0; i < 5; i++)
+            {
+                await bus.PublishAsync(new TestEvent { Payload = $"ambient-{i}" }, TestToken);
+            }
+
+            signal.Mark("before-complete");
+            transaction.Complete();
+        }
+
+        // assert
+        AssertSignaledOnlyAfterCommit(signal);
+    }
+
+    [Fact]
+    public async Task SaveChanges_Should_SignalOnlyAfterCommit_When_AmbientTransactionIsUsed()
+    {
+        // arrange
+        var connectionString = await fixture.CreateDatabaseAsync();
+        var signal = new VisibilityProbeSignal(connectionString);
+        await using var provider = await CreateProviderAsync(connectionString, signal, addHandler: false);
+        await EnsureCreatedAsync(provider);
+        signal.Events.Clear();
+
+        // act
+        using (var transaction = new TransactionScope(TransactionScopeAsyncFlowOption.Enabled))
+        {
+            await using var scope = provider.CreateAsyncScope();
+            var db = scope.ServiceProvider.GetRequiredService<TestDbContext>();
+            db.Add(new OutboxMessage(Guid.NewGuid(), JsonDocument.Parse("{}")));
+            await db.SaveChangesAsync(TestToken);
+            signal.Mark("before-complete");
+            transaction.Complete();
+        }
+
+        // assert
+        AssertSignaledOnlyAfterCommit(signal);
+    }
+
+    [Fact]
+    public async Task PersistAsync_Should_SignalOnCompletion_When_AmbientTransactionRollsBack()
+    {
+        // arrange
+        var connectionString = await fixture.CreateDatabaseAsync();
+        var signal = new VisibilityProbeSignal(connectionString);
+        await using var provider = await CreateProviderAsync(connectionString, signal, addHandler: false);
+        await EnsureCreatedAsync(provider);
+        signal.Events.Clear();
+
+        // act
+        using (new TransactionScope(TransactionScopeAsyncFlowOption.Enabled))
+        {
+            await using var scope = provider.CreateAsyncScope();
+            var bus = scope.ServiceProvider.GetRequiredService<IMessageBus>();
+            await bus.PublishAsync(new TestEvent { Payload = "ambient" }, TestToken);
+            signal.Mark("before-dispose");
+        }
+
+        signal.Mark("after-rollback");
+
+        // assert
+        Assert.Equal(
+            ["before-dispose", "set(visible=0)", "after-rollback"],
+            signal.Events.Distinct().ToArray());
+    }
+
+    [Fact]
+    public async Task Worker_Should_DeliverMessage_When_PublishedInsideAmbientTransaction()
+    {
+        // arrange
+        var connectionString = await fixture.CreateDatabaseAsync();
+        var recorder = new MessageRecorder();
+        await using var provider = await CreateProviderAsync(
+            connectionString,
+            new ResilientOutboxSignal(),
+            addHandler: true,
+            recorder);
+        await EnsureCreatedAsync(provider);
+        var hostedServices = provider.GetServices<IHostedService>().ToList();
+        foreach (var service in hostedServices)
+        {
+            await service.StartAsync(TestToken);
+        }
+
+        try
+        {
+            // worker drained the empty table and is parked on the signal
+            await PublishAsync(provider, "warm-up");
+            Assert.True(await recorder.WaitAsync(TimeSpan.FromSeconds(10)));
+
+            // act
+            using (var transaction = new TransactionScope(TransactionScopeAsyncFlowOption.Enabled))
+            {
+                await PublishAsync(provider, "ambient");
+
+                // give the prematurely woken worker time to find nothing and park again
+                await Task.Delay(TimeSpan.FromSeconds(1), TestToken);
+                transaction.Complete();
+            }
+
+            var delivered = await recorder.WaitAsync(TimeSpan.FromSeconds(10));
+
+            // assert
+            Assert.True(delivered, "message committed inside a TransactionScope was not delivered");
+        }
+        finally
+        {
+            foreach (var service in hostedServices)
+            {
+                await service.StopAsync(CancellationToken.None);
+            }
+
+            await Task.Delay(250, CancellationToken.None);
+        }
+    }
+
+    private static void AssertSignaledOnlyAfterCommit(VisibilityProbeSignal signal)
+    {
+        var events = signal.Events.ToArray();
+        var setsBeforeComplete = events.TakeWhile(e => e != "before-complete").Count(e => e.StartsWith("set"));
+        var setsSeeingNoRows = events.Count(e => e == "set(visible=0)");
+        var setsAfterComplete = events.SkipWhile(e => e != "before-complete").Count(e => e.StartsWith("set"));
+        Assert.Equal((0, 0), (setsBeforeComplete, setsSeeingNoRows));
+        Assert.NotEqual(0, setsAfterComplete);
+    }
+
+    private static async Task PublishAsync(IServiceProvider provider, string payload)
+    {
+        await using var scope = provider.CreateAsyncScope();
+        var bus = scope.ServiceProvider.GetRequiredService<IMessageBus>();
+        await bus.PublishAsync(new TestEvent { Payload = payload }, TestToken);
+    }
+
+    private static async Task EnsureCreatedAsync(IServiceProvider provider)
+    {
+        await using var scope = provider.CreateAsyncScope();
+        var db = scope.ServiceProvider.GetRequiredService<TestDbContext>();
+        await db.Database.EnsureCreatedAsync(TestToken);
+    }
+
+    private static async Task<ServiceProvider> CreateProviderAsync(
+        string connectionString,
+        IOutboxSignal signal,
+        bool addHandler,
+        MessageRecorder? recorder = null)
+    {
+        var services = new ServiceCollection();
+        services.AddSingleton(recorder ?? new MessageRecorder());
+        services.AddLogging();
+        services.AddDbContext<TestDbContext>(o => o.UseTestNpgsql(connectionString));
+        services.AddSingleton(signal);
+
+        var builder = services.AddMessageBus();
+        builder.AddEntityFramework<TestDbContext>(ef => ef.UsePostgresOutbox());
+        if (addHandler)
+        {
+            builder.AddEventHandler<TestEventHandler>();
+        }
+
+        builder.AddInMemory();
+
+        var provider = services.BuildServiceProvider();
+        var runtime = (MessagingRuntime)provider.GetRequiredService<IMessagingRuntime>();
+        await runtime.StartAsync(TestToken);
+        return provider;
+    }
+
+    public sealed class TestEvent
+    {
+        public required string Payload { get; init; }
+    }
+
+    public sealed class TestEventHandler(MessageRecorder recorder) : IEventHandler<TestEvent>
+    {
+        public ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken)
+        {
+            recorder.Record(message);
+            return default;
+        }
+    }
+
+    public sealed class MessageRecorder
+    {
+        private readonly SemaphoreSlim _semaphore = new(0);
+
+        public void Record(object message) => _semaphore.Release();
+
+        public Task<bool> WaitAsync(TimeSpan timeout) => _semaphore.WaitAsync(timeout);
+    }
+
+    private sealed class VisibilityProbeSignal(string connectionString) : IOutboxSignal
+    {
+        private readonly string _connectionString =
+            new NpgsqlConnectionStringBuilder(connectionString) { Enlist = false }.ConnectionString;
+
+        public ConcurrentQueue<string> Events { get; } = new();
+
+        public void Mark(string label) => Events.Enqueue(label);
+
+        public void Set()
+        {
+            using var connection = new NpgsqlConnection(_connectionString);
+            connection.Open();
+            using var command = connection.CreateCommand();
+            command.CommandText = "SELECT COUNT(*) FROM outbox_messages";
+            Events.Enqueue($"set(visible={command.ExecuteScalar()})");
+        }
+
+        public Task WaitAsync(CancellationToken cancellationToken)
+            => Task.Delay(Timeout.Infinite, cancellationToken);
+    }
+}