GH-4397: count every batched member per pipeline, not just listener arrivals

docs/guide/handlers/batching.md

@@ -253,6 +253,43 @@ that refers to the original message is completely processed. See [Durability and
 settlement](#durability-and-message-settlement) above for the full settlement model and why a durable listener
 is required for guaranteed delivery.
 
+## Waiting for batches to finish <Badge type="tip" text="6.35" />
+
+Let's say you've got an integration test that causes some batched messages, and before the next test resets
+the database you want to be sure every one of those batches has actually run. Otherwise a batch sitting out its
+`TriggerTime` can fire *after* the reset and fail against data that's no longer there.
+
+Wolverine keeps a running count of every member message that's somewhere inside a batching pipeline -- waiting
+to be batched, sitting in an assembled batch, or executing in your batch handler -- and you can ask for it through
+`BatchingPendingCounts` on the Wolverine runtime:
+
+```csharp
+var counts = host.GetRuntime().BatchingPendingCounts;
+
+// Members of BatchMessagesOf<Item>() that haven't finished yet, however they arrived
+var pending = counts.PendingForBatchedMessage<Item>();
+
+// Or across every batching pipeline in the application
+var total = counts.TotalPendingBatchMembers;
+```
+
+A member is released when the batch it belongs to reaches its terminal, success or failure, so a count of
+zero means the batch handler has finished (including committing its transaction if you're using transactional
+middleware).
+
+::: warning
+A member is only counted once the batching processor picks it up from the element type's local queue. Right after
+you publish, the count can still read zero because the message hasn't gotten there yet, so wait for the count to
+go up before you wait for it to come back down -- or just use a [tracked session](/guide/testing), which handles
+all of this for you.
+:::
+
+::: info
+`BatchingPendingCounts.PendingFor(Uri)` is a different number. It only counts members that arrived through a
+specific transport listener, because it exists to feed that listener's back-pressure. Members you publish
+in-process, including cascading messages, never show up there.
+:::
+
 ## Message expiration (`DeliverBy` / `DeliverWithin`) <Badge type="tip" text="6.26" />
 
 Batched messages honor member-level [message expiration](/guide/messages) the same way unbatched messages do,

src/Persistence/MartenTests/batch_processing.cs

@@ -100,6 +100,82 @@ public async Task end_to_end_with_durable()
         }
     }
 
+    [Fact]
+    public async Task pending_count_covers_locally_published_members_on_durable_local_queues()
+    {
+        // GH-4397, the reporter's shape: in-process publishes onto durable local queues. Those members have
+        // no listener, so only the per-pipeline count sees them, and it must not reach zero until the batch
+        // handler's transaction has committed.
+        using var theHost = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Services.AddMarten(m =>
+                {
+                    m.DisableNpgsqlLogging = true;
+                    m.DatabaseSchemaName = "batch_pending";
+                    m.Connection(Servers.PostgresConnectionString);
+                }).IntegrateWithWolverine();
+
+                opts.Policies.AutoApplyTransactions();
+                opts.Policies.UseDurableLocalQueues();
+
+                opts.BatchMessagesOf<BatchItem>(batching =>
+                {
+                    batching.TriggerTime = 1.Seconds();
+                    batching.BatchSize = 8;
+                }).Sequential();
+
+                opts.Discovery.DisableConventionalDiscovery()
+                    .IncludeType(typeof(BatchItemHandler));
+                opts.Durability.Mode = DurabilityMode.Solo;
+            }).StartAsync(cancellationToken: TestContext.Current.CancellationToken);
+
+        await theHost.CleanAllMartenDataAsync();
+        await theHost.ResetResourceState(cancellation: TestContext.Current.CancellationToken);
+
+        var counts = theHost.GetRuntime().BatchingPendingCounts;
+        var items = new[]
+        {
+            new BatchItem("one", Guid.NewGuid()),
+            new BatchItem("two", Guid.NewGuid()),
+            new BatchItem("three", Guid.NewGuid())
+        };
+
+        var bus = theHost.MessageBus();
+        foreach (var item in items)
+        {
+            await bus.PublishAsync(item);
+        }
+
+        await waitUntil(() => counts.PendingForBatchedMessage<BatchItem>() == items.Length,
+            "every published member to be counted as pending");
+
+        await waitUntil(() => counts.PendingForBatchedMessage<BatchItem>() == 0,
+            "the batch to reach its terminal");
+
+        // Zero means the batch handler ran AND its transaction committed
+        using var session = theHost.DocumentStore().LightweightSession();
+        var ids = items.Select(x => x.Id).ToArray();
+        var stored = await session.Query<BatchItem>().Where(x => x.Id.IsOneOf(ids))
+            .CountAsync(token: TestContext.Current.CancellationToken);
+
+        stored.ShouldBe(items.Length);
+    }
+
+    private static async Task waitUntil(Func<bool> condition, string description)
+    {
+        var deadline = DateTime.UtcNow.AddSeconds(30);
+        while (!condition())
+        {
+            if (DateTime.UtcNow > deadline)
+            {
+                throw new TimeoutException($"Timed out waiting for {description}");
+            }
+
+            await Task.Delay(10);
+        }
+    }
+
     [Fact]
     public async Task end_to_end_with_tenancy()
     {
@@ -122,9 +198,9 @@ public async Task end_to_end_with_tenancy()
                     batching.TriggerTime = 1.Seconds();
                     batching.BatchSize = 8;
                     batching.LocalExecutionQueueName = "items";
-                }).UseDurableInbox();
-
-                opts.Discovery.DisableConventionalDiscovery()
+                }).UseDurableInbox();
+
+                opts.Discovery.DisableConventionalDiscovery()
                     .IncludeType(typeof(BatchItemHandler));
                 opts.Durability.Mode = DurabilityMode.Solo;
             }).StartAsync(cancellationToken: TestContext.Current.CancellationToken);

src/Testing/CoreTests/Acceptance/batching_pending_counts_for_local_members.cs

@@ -0,0 +1,128 @@
+using JasperFx.Core;
+using JasperFx.Core.Reflection;
+using Microsoft.Extensions.Hosting;
+using System.Collections.Concurrent;
+using Wolverine.Runtime;
+using Wolverine.Tracking;
+using Xunit;
+
+namespace CoreTests.Acceptance;
+
+/// <summary>
+/// GH-4397. Members published in-process -- here, cascaded from another handler -- have no listener, so the
+/// listener-keyed back-pressure count never sees them. The per-pipeline count has to, because "has every
+/// batch this test caused finished?" is exactly what an integration test harness needs to ask.
+/// </summary>
+public class batching_pending_counts_for_local_members : IAsyncLifetime
+{
+    private IHost _host = null!;
+
+    public async ValueTask InitializeAsync()
+    {
+        _host = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Discovery.DisableConventionalDiscovery()
+                    .IncludeType(typeof(StartEnrollmentsHandler))
+                    .IncludeType(typeof(PendingEnrollmentBatchHandler));
+
+                // Well past the time it takes to observe the members as pending, so the batch fires on its
+                // trigger time rather than racing the assertions
+                opts.BatchMessagesOf<PendingEnrollment>(batching =>
+                {
+                    batching.BatchSize = 100;
+                    batching.TriggerTime = 1.Seconds();
+                }).Sequential();
+            }).StartAsync();
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _host.StopAsync();
+        _host.Dispose();
+    }
+
+    [Fact]
+    public async Task a_cascaded_member_is_pending_until_its_batch_handler_has_finished()
+    {
+        var counts = _host.GetRuntime().BatchingPendingCounts;
+        var id = Guid.NewGuid();
+
+        counts.PendingForBatchedMessage<PendingEnrollment>().ShouldBe(0);
+
+        await _host.MessageBus().PublishAsync(new StartEnrollments(id, 3));
+
+        await waitUntil(() => counts.PendingForBatchedMessage<PendingEnrollment>() == 3,
+            "all three cascaded members to be counted as pending");
+
+        counts.TotalPendingBatchMembers.ShouldBe(3);
+
+        // The back-pressure view is unchanged: none of these members came through a listener
+        var elementQueue = _host.GetRuntime().RoutingFor(typeof(PendingEnrollment)).Routes.Single()
+            .As<Wolverine.Runtime.Routing.MessageRoute>().Sender.Destination;
+        counts.PendingFor(elementQueue).ShouldBe(0);
+
+        await waitUntil(() => counts.PendingForBatchedMessage<PendingEnrollment>() == 0,
+            "the batch to reach its terminal");
+
+        // Zero only once the batch handler has finished -- and while it ran, its members were still pending
+        PendingEnrollmentBatchHandler.Finished(id).ShouldBeTrue();
+        PendingEnrollmentBatchHandler.PendingSeenWhileExecuting(id).ShouldBe(3);
+        counts.TotalPendingBatchMembers.ShouldBe(0);
+    }
+
+    private static async Task waitUntil(Func<bool> condition, string description)
+    {
+        var deadline = DateTime.UtcNow.AddSeconds(15);
+        while (!condition())
+        {
+            if (DateTime.UtcNow > deadline)
+            {
+                throw new TimeoutException($"Timed out waiting for {description}");
+            }
+
+            await Task.Delay(10);
+        }
+    }
+}
+
+public record StartEnrollments(Guid Id, int Count);
+
+public record PendingEnrollment(Guid Id, int Number);
+
+public static class StartEnrollmentsHandler
+{
+    public static OutgoingMessages Handle(StartEnrollments command)
+    {
+        var messages = new OutgoingMessages();
+        for (var i = 0; i < command.Count; i++)
+        {
+            messages.Add(new PendingEnrollment(command.Id, i));
+        }
+
+        return messages;
+    }
+}
+
+public static class PendingEnrollmentBatchHandler
+{
+    private static readonly ConcurrentDictionary<Guid, int> _pendingWhileExecuting = new();
+    private static readonly ConcurrentDictionary<Guid, bool> _finished = new();
+
+    public static int PendingSeenWhileExecuting(Guid id) => _pendingWhileExecuting.GetValueOrDefault(id);
+
+    public static bool Finished(Guid id) => _finished.ContainsKey(id);
+
+    public static async Task Handle(PendingEnrollment[] enrollments, IWolverineRuntime runtime)
+    {
+        var id = enrollments[0].Id;
+        _pendingWhileExecuting[id] = runtime.As<WolverineRuntime>().BatchingPendingCounts
+            .PendingForBatchedMessage<PendingEnrollment>();
+
+        // Long enough that a count released at the start of execution, rather than at the terminal,
+        // would be caught by the test's wait for zero
+        await Task.Delay(250);
+
+        _finished[id] = true;
+    }
+}