GH-5171 Coalesce durable outbox/inbox writes per tenant database

src/Testing/CoreTests/Persistence/Durability/per_tenant_outbox_coalescing_5171.cs

@@ -0,0 +1,210 @@
+using CoreTests.Runtime;
+using JasperFx;
+using JasperFx.MultiTenancy;
+using Microsoft.Extensions.Logging.Abstractions;
+using NSubstitute;
+using Shouldly;
+using Wolverine;
+using Wolverine.ComplianceTests;
+using Wolverine.Configuration;
+using Wolverine.Logging;
+using Wolverine.Persistence.Durability;
+using Wolverine.Runtime;
+using Wolverine.Transports;
+using Wolverine.Transports.Sending;
+using Xunit;
+
+namespace CoreTests.Persistence.Durability;
+
+/// <summary>
+/// GH-5171. With a database per tenant, the durable sending agent's outbox coalescers used to be
+/// shared by every tenant sending to the same destination, one flush at a time. One tenant's slow
+/// database -- an exhausted connection pool, say -- then held up every other tenant's outbox writes
+/// and deletes, and so every other tenant's request. Each tenant database now gets its own lane.
+/// </summary>
+public class per_tenant_outbox_coalescing_5171 : IAsyncDisposable
+{
+    private readonly IMessageStore theMainStore = Substitute.For<IMessageStore>();
+    private readonly IMessageStore theSlowStore = Substitute.For<IMessageStore>();
+    private readonly IMessageStore theFastStore = Substitute.For<IMessageStore>();
+    private readonly ITenantedMessageSource theSource = Substitute.For<ITenantedMessageSource>();
+    private readonly MultiTenantedMessageStore theStore;
+
+    private readonly Uri theDestination = new("stub://one");
+    private readonly ISender theSender = Substitute.For<ISender>();
+    private DurableSendingAgent? _agent;
+
+    public per_tenant_outbox_coalescing_5171()
+    {
+        var runtime = new MockWolverineRuntime();
+
+        // Keep GetDatabaseAsync from trying to migrate a substitute on first resolution
+        runtime.Options.AutoBuildMessageStorageOnStartup = AutoCreate.None;
+
+        theSource.FindAsync("slow").Returns(theSlowStore);
+        theSource.FindAsync("fast").Returns(theFastStore);
+
+        theStore = new MultiTenantedMessageStore(theMainStore, runtime, theSource);
+
+        theSender.Destination.Returns(theDestination);
+        theSender.SupportsNativeScheduledSend.Returns(true);
+    }
+
+    private DurableSendingAgent agent()
+    {
+        // The default batch size, which is what turns the coalescers on
+        var settings = new DurabilitySettings { StoreOutgoingBatchSize = 100 };
+        _agent = new DurableSendingAgent(theSender, settings, NullLogger.Instance,
+            Substitute.For<IMessageTracker>(), theStore, new StubEndpoint(theDestination));
+        return _agent;
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (_agent != null) await _agent.DisposeAsync();
+    }
+
+    private Envelope envelopeFor(string tenantId)
+    {
+        var envelope = ObjectMother.Envelope();
+        envelope.TenantId = tenantId;
+        envelope.Destination = theDestination;
+        return envelope;
+    }
+
+    /// <summary>
+    /// Every write to the slow tenant's database waits until the test lets it go.
+    /// </summary>
+    private TaskCompletionSource slowDownTheSlowTenant()
+    {
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        theSlowStore.Outbox.DeleteOutgoingAsync(Arg.Any<Envelope>()).Returns(_ => gate.Task);
+        theSlowStore.Outbox.DeleteOutgoingAsync(Arg.Any<Envelope[]>()).Returns(_ => gate.Task);
+        theSlowStore.Outbox.StoreOutgoingAsync(Arg.Any<Envelope>(), Arg.Any<int>()).Returns(_ => gate.Task);
+        theSlowStore.Outbox.StoreOutgoingAsync(Arg.Any<IReadOnlyList<Envelope>>(), Arg.Any<int>())
+            .Returns(_ => gate.Task);
+
+        return gate;
+    }
+
+    [Fact]
+    public async Task a_slow_tenant_database_does_not_delay_another_tenants_successful_send()
+    {
+        var theAgent = agent();
+        var gate = slowDownTheSlowTenant();
+
+        // The slow tenant's delete is in flight, and more of its deletes pile up behind it
+        var slowDeletes = Enumerable.Range(0, 3)
+            .Select(_ => theAgent.MarkSuccessfulAsync(envelopeFor("slow")))
+            .ToArray();
+
+        // The other tenant's request is waiting on this to return
+        var fast = envelopeFor("fast");
+        await theAgent.MarkSuccessfulAsync(fast)
+            .WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+
+        await theFastStore.Outbox.Received(1).DeleteOutgoingAsync(fast);
+        slowDeletes.ShouldAllBe(x => !x.IsCompleted);
+
+        // And the slow tenant still gets all of its rows deleted once its database catches up
+        gate.SetResult();
+        await Task.WhenAll(slowDeletes).WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+    }
+
+    [Fact]
+    public async Task a_slow_tenant_database_does_not_delay_another_tenants_outbox_store()
+    {
+        var theAgent = agent();
+        var gate = slowDownTheSlowTenant();
+
+        var slowStores = Enumerable.Range(0, 3)
+            .Select(_ => theAgent.TryStoreOutgoingAsync(envelopeFor("slow")).AsTask())
+            .ToArray();
+
+        var fast = envelopeFor("fast");
+        (await theAgent.TryStoreOutgoingAsync(fast).AsTask()
+                .WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken))
+            .ShouldBeTrue();
+
+        await theFastStore.Outbox.Received(1).StoreOutgoingAsync(fast, Arg.Any<int>());
+        slowStores.ShouldAllBe(x => !x.IsCompleted);
+
+        gate.SetResult();
+        await Task.WhenAll(slowStores).WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+    }
+
+    [Fact]
+    public async Task tenants_that_share_a_database_still_share_a_batch()
+    {
+        // Two tenant ids, one database behind them
+        theSource.FindAsync("fast-too").Returns(theFastStore);
+        await theStore.GetDatabaseAsync("fast");
+        await theStore.GetDatabaseAsync("fast-too");
+
+        var theAgent = agent();
+
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var first = true;
+        theFastStore.Outbox.DeleteOutgoingAsync(Arg.Any<Envelope>()).Returns(_ =>
+        {
+            if (!first) return Task.CompletedTask;
+            first = false;
+            started.TrySetResult();
+            return gate.Task;
+        });
+
+        var lead = theAgent.MarkSuccessfulAsync(envelopeFor("fast"));
+        await started.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+
+        var others = new[] { envelopeFor("fast"), envelopeFor("fast-too"), envelopeFor("fast-too") };
+        var pending = others.Select(theAgent.MarkSuccessfulAsync).ToArray();
+
+        gate.SetResult();
+        await Task.WhenAll(pending.Append(lead))
+            .WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+
+        await theFastStore.Outbox.Received(1).DeleteOutgoingAsync(
+            Arg.Is<Envelope[]>(x => x.Length == 3 && others.All(x.Contains)));
+    }
+
+    [Fact]
+    public void a_single_database_outbox_keeps_the_single_lane_coalescer()
+    {
+        // No per-envelope partition lookup at all when there is only one database to write to
+        EnvelopeStoreCoalescing.PartitionerFor(Substitute.For<IMessageOutbox>()).ShouldBeNull();
+        EnvelopeStoreCoalescing.PartitionerFor(Substitute.For<IMessageInbox>()).ShouldBeNull();
+
+        EnvelopeStoreCoalescing.Build(null, _ => Task.CompletedTask, _ => Task.CompletedTask, 100, theDestination,
+            NullLogger.Instance).ShouldBeOfType<EnvelopeStoreCoalescer>();
+    }
+
+    [Fact]
+    public async Task the_default_tenant_and_the_main_store_share_one_lane()
+    {
+        await theStore.GetDatabaseAsync("fast");
+
+        theStore.CoalescingPartitionFor(null).ShouldBeSameAs(theMainStore);
+        theStore.CoalescingPartitionFor(StorageConstants.DefaultTenantId).ShouldBeSameAs(theMainStore);
+        theStore.CoalescingPartitionFor(TransportConstants.Default).ShouldBeSameAs(theMainStore);
+        theStore.CoalescingPartitionFor("fast").ShouldBeSameAs(theFastStore);
+
+        // A tenant nobody has resolved yet gets a lane by its id until its store is known
+        theStore.CoalescingPartitionFor("unseen").ShouldBe("unseen");
+    }
+
+    private class StubEndpoint : Endpoint
+    {
+        public StubEndpoint(Uri uri) : base(uri, EndpointRole.Application)
+        {
+        }
+
+        public override ValueTask<IListener> BuildListenerAsync(IWolverineRuntime runtime, IReceiver receiver)
+            => throw new NotSupportedException();
+
+        protected override ISender CreateSender(IWolverineRuntime runtime) => throw new NotSupportedException();
+
+        protected override bool supportsMode(EndpointMode mode) => true;
+    }
+}

src/Wolverine/Persistence/Durability/DelegatingMessageInbox.cs

@@ -15,6 +15,12 @@ public DelegatingMessageInbox(IMessageInbox inner, MessageStoreCollection stores
         _stores = stores;
     }
 
+    // GH-5171. Same routing as the writes below: Envelope.Store first, then the inner inbox
+    internal object CoalescingPartitionFor(Envelope envelope)
+    {
+        return EnvelopeStoreCoalescing.PartitionFor(envelope, _inner);
+    }
+
     /// <summary>
     /// Each envelope's row lives in the store it came from, so the heartbeat goes to each store with its own
     /// envelopes. A store without the capability is skipped.

src/Wolverine/Persistence/Durability/DelegatingMessageOutbox.cs

@@ -11,6 +11,12 @@ public DelegatingMessageOutbox(IMessageOutbox inner, MessageStoreCollection stor
         _stores = stores;
     }
 
+    // GH-5171. Same routing as the writes below: Envelope.Store first, then the inner outbox
+    internal object CoalescingPartitionFor(Envelope envelope)
+    {
+        return EnvelopeStoreCoalescing.PartitionFor(envelope, _inner);
+    }
+
     public Task<IReadOnlyList<Envelope>> LoadOutgoingAsync(Uri destination)
     {
         return _inner.LoadOutgoingAsync(destination);

src/Wolverine/Persistence/Durability/DurableSendingAgent.cs

@@ -21,8 +21,12 @@ internal class DurableSendingAgent : SendingAgent
     // GH-4319. Null when the batch sizes are 1. The store gates the send, so these coalescers are
     // strictly timer-free: a lone envelope is written immediately and batches only ever form behind a
     // flush that was already in flight.
-    private readonly EnvelopeStoreCoalescer? _storeCoalescer;
-    private readonly EnvelopeStoreCoalescer? _deleteCoalescer;
+    //
+    // GH-5171. When the outbox fans out to several databases (a database per tenant, or ancillary
+    // stores), each database gets its own lane, so one slow tenant database cannot hold back the sends
+    // of every other tenant on this endpoint.
+    private readonly IEnvelopeStoreCoalescer? _storeCoalescer;
+    private readonly IEnvelopeStoreCoalescer? _deleteCoalescer;
 
     private IList<Envelope> _queued = new List<Envelope>();
 
@@ -52,7 +56,9 @@ public DurableSendingAgent(ISender sender, DurabilitySettings settings, ILogger
 
         if (settings.StoreOutgoingBatchSize > 1)
         {
-            _storeCoalescer = new EnvelopeStoreCoalescer(
+            var partitionFor = EnvelopeStoreCoalescing.PartitionerFor(_outbox);
+
+            _storeCoalescer = EnvelopeStoreCoalescing.Build(partitionFor,
                 envelopes => _outbox.StoreOutgoingAsync(envelopes, settings.OutgoingOwnerId),
                 envelope => _outbox.StoreOutgoingAsync(envelope, settings.OutgoingOwnerId),
                 settings.StoreOutgoingBatchSize, endpoint.Uri, logger);
@@ -61,7 +67,7 @@ public DurableSendingAgent(ISender sender, DurabilitySettings settings, ILogger
             // used it since it existed -- but a single-envelope send paid its own round trip to remove
             // one row. Same coalescer, same guarantee: MarkSuccessfulAsync does not return until this
             // envelope's row is gone, which is what makes releasing a pooled envelope afterwards safe.
-            _deleteCoalescer = new EnvelopeStoreCoalescer(
+            _deleteCoalescer = EnvelopeStoreCoalescing.Build(partitionFor,
                 envelopes => _outbox.DeleteOutgoingAsync(envelopes as Envelope[] ?? envelopes.ToArray()),
                 envelope => _deleteOutgoingOne.PostAsync(envelope),
                 settings.StoreOutgoingBatchSize, endpoint.Uri, logger);

src/Wolverine/Persistence/Durability/EnvelopeStoreCoalescer.cs

@@ -2,6 +2,16 @@
 
 namespace Wolverine.Persistence.Durability;
 
+/// <summary>
+///     What the durable sending agent and the durable local queue need from a coalescer: write one
+///     envelope, and drain whatever is in flight at shutdown.
+/// </summary>
+internal interface IEnvelopeStoreCoalescer
+{
+    Task StoreAsync(Envelope envelope);
+    Task DrainAsync();
+}
+
 /// <summary>
 ///     GH-4319. Coalesces concurrent single-envelope durability writes into one batched round trip,
 ///     without ever making a caller wait for other work to show up. This is the insert-side twin of
@@ -29,7 +39,7 @@ namespace Wolverine.Persistence.Durability;
 ///     exception for exactly that envelope, and one poisoned envelope cannot fail its neighbours.
 ///     </para>
 /// </remarks>
-internal class EnvelopeStoreCoalescer
+internal class EnvelopeStoreCoalescer : IEnvelopeStoreCoalescer
 {
     private readonly Func<IReadOnlyList<Envelope>, Task> _storeBatch;
     private readonly Func<Envelope, Task> _storeOne;

src/Wolverine/Persistence/Durability/MultiTenantedMessageStore.cs

@@ -1126,6 +1126,23 @@ public IReadOnlyList<IMessageStore> ActiveDatabases()
         return databases().ToArray();
     }
 
+    /// <summary>
+    ///     GH-5171. The key of the database an envelope for this tenant is written to, without any I/O:
+    ///     the resolved store once the tenant has been seen, otherwise the tenant id itself. Used to give
+    ///     each tenant database its own durability write coalescer, so one slow database cannot hold back
+    ///     another tenant's writes.
+    /// </summary>
+    internal object CoalescingPartitionFor(string? tenantId)
+    {
+        if (tenantId.IsDefaultTenant() || tenantId!.EqualsIgnoreCase(TransportConstants.Default) ||
+            tenantId!.EqualsIgnoreCase(StorageConstants.Main))
+        {
+            return Main;
+        }
+
+        return _byTenant.TryFind(tenantId!, out var store) ? store : tenantId!;
+    }
+
     public async ValueTask<IMessageStore> GetDatabaseAsync(string? tenantId)
     {
         if (tenantId.IsDefaultTenant())