feat(batching): ProbeIndividuallyAfter(N) count-based batch member isolation (GH-3289 Phase 3c)

docs/guide/handlers/batching.md

@@ -370,10 +370,27 @@ the bad member. The isolation is bounded and one-time — a member that has alre
 batch is simply dead-lettered rather than probed again. On a message type that is **not** batched,
 `IsolateBatchMembers()` behaves like a plain move-to-error-queue (there is nothing to isolate).
 
-::: info
-A count-based variant — `ProbeIndividuallyAfter(attempts: N)`, which probes individually only after N
-whole-batch failures — is a planned follow-up (see [GH-3289](https://github.com/JasperFx/wolverine/issues/3289)).
-:::
+### Probing after N whole-batch failures
+
+When you don't have a specific exception type to key on — the batch just fails and you can't tell whether it's
+transient or a poison item — configure a **count-based** probe directly on the batch. `ProbeIndividuallyAfter(N)`
+retries the whole batch `N` times and *only then* falls back to isolating each member individually:
+
+```csharp
+opts.BatchMessagesOf<Order>(b =>
+{
+    b.BatchSize = 500;
+    b.TriggerTime = 10.Seconds();
+
+    // Retry the whole batch 3 times (in case the failure was transient); if it still fails, re-run each
+    // member on its own so only the genuinely bad one dead-letters.
+    b.ProbeIndividuallyAfter(attempts: 3);
+});
+```
+
+This gives transient failures a chance to clear on retry before paying for the per-member probe, and like
+`IsolateBatchMembers` it is bounded and one-time — once reduced to a size-1 batch, a failing member is
+dead-lettered rather than probed again.
 
 ## Custom Batching Strategies
 

src/Testing/CoreTests/Acceptance/batch_probe_individually_after.cs

@@ -0,0 +1,114 @@
+using JasperFx.Core;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Wolverine;
+using Wolverine.Persistence.Durability;
+using Xunit;
+
+namespace CoreTests.Acceptance;
+
+public class batch_probe_individually_after : IAsyncLifetime
+{
+    private IHost _host = null!;
+    private readonly CapturingDeadLetterInterceptor _deadLetters = new();
+
+    public async Task InitializeAsync()
+    {
+        ProbeAfterHandler.Reset();
+
+        _host = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Discovery.DisableConventionalDiscovery();
+                opts.Discovery.IncludeType<ProbeAfterHandler>();
+                opts.Services.AddSingleton<IDeadLetterInterceptor>(_deadLetters);
+
+                opts.BatchMessagesOf<ProbeAfterItem>(b =>
+                {
+                    b.BatchSize = 500;
+                    b.TriggerTime = 1.Seconds();
+
+                    // Retry the whole batch 3 times, then probe each member individually.
+                    b.ProbeIndividuallyAfter(3);
+                }).Sequential();
+            }).StartAsync();
+    }
+
+    public async Task DisposeAsync()
+    {
+        await _host.StopAsync();
+        _host.Dispose();
+    }
+
+    [Fact]
+    public async Task retries_the_whole_batch_then_probes_individually()
+    {
+        var bus = _host.MessageBus();
+        await bus.PublishAsync(new ProbeAfterItem("good1", false));
+        await bus.PublishAsync(new ProbeAfterItem("bad", true));
+        await bus.PublishAsync(new ProbeAfterItem("good2", false));
+
+        await ProbeAfterHandler.BothGoodsSucceeded.Task.WaitAsync(15.Seconds());
+        await _deadLetters.Signal.Task.WaitAsync(15.Seconds());
+
+        // The whole 3-item batch was retried exactly 3 times before the probe kicked in.
+        ProbeAfterHandler.WholeBatchAttempts.ShouldBe(3);
+
+        // Then the probe isolated the poison item: the healthy members succeeded as singletons...
+        ProbeAfterHandler.SucceededIds.OrderBy(x => x).ShouldBe(new[] { "good1", "good2" });
+
+        // ...and only the poison item was dead-lettered.
+        _deadLetters.DeadLettered.OfType<ProbeAfterItem>().Select(x => x.Id).ShouldBe(new[] { "bad" });
+    }
+}
+
+public record ProbeAfterItem(string Id, bool Poison);
+
+public class ProbeAfterFailure : Exception;
+
+public class ProbeAfterHandler
+{
+    private static readonly object _locker = new();
+
+    public static int WholeBatchAttempts;
+    public static List<string> SucceededIds { get; } = new();
+    public static TaskCompletionSource BothGoodsSucceeded = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+    public static void Reset()
+    {
+        WholeBatchAttempts = 0;
+        lock (_locker)
+        {
+            SucceededIds.Clear();
+        }
+
+        BothGoodsSucceeded = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+    }
+
+    public void Handle(ProbeAfterItem[] items)
+    {
+        if (items.Any(x => x.Poison))
+        {
+            // Count only the whole-batch attempts (the poison singleton also lands here on the probe).
+            if (items.Length > 1)
+            {
+                Interlocked.Increment(ref WholeBatchAttempts);
+            }
+
+            throw new ProbeAfterFailure();
+        }
+
+        lock (_locker)
+        {
+            foreach (var item in items)
+            {
+                SucceededIds.Add(item.Id);
+            }
+
+            if (SucceededIds.Count >= 2)
+            {
+                BothGoodsSucceeded.TrySetResult();
+            }
+        }
+    }
+}

src/Wolverine/Runtime/Batching/BatchingOptions.cs

@@ -62,6 +62,27 @@ public void CoalesceBy<T, TKey>(Func<T, TKey> keySelector)
         Batcher = new CoalescingMessageBatcher<T, TKey>(keySelector);
     }
 
+    internal int? ProbeIndividuallyAfterAttempts { get; private set; }
+
+    /// <summary>
+    /// After the whole batch has failed <paramref name="attempts"/> times, re-run each member as its own
+    /// size-1 batch so that only the message actually causing the failure is dead-lettered while the
+    /// healthy members succeed. Use this for <em>opaque</em> failures where the handler cannot name the bad
+    /// item. Bounded and one-time: a member reduced to a size-1 batch is dead-lettered rather than probed
+    /// again. For failures the handler can attribute to a specific exception type, prefer the
+    /// <c>OnException&lt;T&gt;().IsolateBatchMembers()</c> policy instead. See GH-3289.
+    /// </summary>
+    /// <param name="attempts">The number of whole-batch failures to allow before probing. Must be >= 1.</param>
+    public void ProbeIndividuallyAfter(int attempts)
+    {
+        if (attempts < 1)
+        {
+            throw new ArgumentOutOfRangeException(nameof(attempts), "The probe threshold must be at least 1");
+        }
+
+        ProbeIndividuallyAfterAttempts = attempts;
+    }
+
     /// <summary>
     /// The maximum size of the message batch. Default is 100.
     /// </summary>

src/Wolverine/Runtime/Batching/ProbeIndividuallyContinuation.cs

@@ -4,28 +4,43 @@
 namespace Wolverine.Runtime.Batching;
 
 /// <summary>
-/// Built-in continuation source for the batch member-isolation error policy
-/// (<c>OnException&lt;T&gt;().IsolateBatchMembers()</c>). Used for <em>opaque</em> batch failures where the
-/// handler cannot name the bad item: re-run each member as its own size-1 batch so only the one carrying
-/// the bad element dead-letters and the healthy members succeed. See GH-3289.
+/// Built-in continuation source for the batch member-isolation error handling. Used for <em>opaque</em>
+/// batch failures where the handler cannot name the bad item: re-run each member as its own size-1 batch
+/// so only the one carrying the bad element dead-letters and the healthy members succeed. See GH-3289.
+///
+/// Two triggers install it: the exception-type policy
+/// <c>OnException&lt;T&gt;().IsolateBatchMembers()</c> (probe on the first failure of that type) and the
+/// count-based <c>BatchMessagesOf&lt;T&gt;(b =&gt; b.ProbeIndividuallyAfter(N))</c> (retry the whole batch,
+/// then probe after N failures).
 /// </summary>
 internal class ProbeIndividuallyContinuationSource : IContinuationSource
 {
+    private readonly int _probeAfterAttempts;
+
+    // probeAfterAttempts is the attempt on which to probe; earlier attempts retry the whole batch. The
+    // default of 1 probes on the first failure (the IsolateBatchMembers policy).
+    public ProbeIndividuallyContinuationSource(int probeAfterAttempts = 1)
+    {
+        _probeAfterAttempts = probeAfterAttempts;
+    }
+
     public string Description => "Isolate the failing batch member by re-running each member individually";
 
     public IContinuation Build(Exception ex, Envelope envelope)
     {
-        return new ProbeIndividuallyContinuation(ex);
+        return new ProbeIndividuallyContinuation(ex, _probeAfterAttempts);
     }
 }
 
 internal class ProbeIndividuallyContinuation : IContinuation
 {
     private readonly Exception _exception;
+    private readonly int _probeAfterAttempts;
 
-    public ProbeIndividuallyContinuation(Exception exception)
+    public ProbeIndividuallyContinuation(Exception exception, int probeAfterAttempts = 1)
     {
         _exception = exception ?? throw new ArgumentNullException(nameof(exception));
+        _probeAfterAttempts = probeAfterAttempts;
     }
 
     public async ValueTask ExecuteAsync(IEnvelopeLifecycle lifecycle, IWolverineRuntime runtime,
@@ -48,6 +63,16 @@ public async ValueTask ExecuteAsync(IEnvelopeLifecycle lifecycle, IWolverineRunt
             return;
         }
 
+        // Count-based trigger (ProbeIndividuallyAfter): retry the WHOLE batch until it has failed
+        // probeAfterAttempts times, then probe. The Executor increments Attempts before each handling, so
+        // Attempts is the number of failures so far. IsolateBatchMembers uses the default of 1 (probe on
+        // the first failure).
+        if (batch.Attempts < _probeAfterAttempts)
+        {
+            await lifecycle.DeferAsync().ConfigureAwait(false);
+            return;
+        }
+
         // Re-run each member as its own size-1 batch (via the shared re-execution primitive) so only the
         // member carrying the bad element fails and dead-letters; the healthy members succeed. Bounded
         // and one-time: the singletons cannot be split further (handled by the guard above).

src/Wolverine/Runtime/WolverineRuntime.HostService.cs

@@ -7,6 +7,7 @@
 using Microsoft.Extensions.Logging;
 using Wolverine.Attributes;
 using Wolverine.Configuration;
+using Wolverine.ErrorHandling;
 using Wolverine.Persistence.Durability;
 using Wolverine.Runtime.Agents;
 using Wolverine.Runtime.Scheduled;
@@ -112,6 +113,10 @@ public async Task StartAsync(CancellationToken cancellationToken)
             // Warn loudly (or throw, if opted in) so the shadowing is not a silent surprise. GH-3289.
             warnOrAssertBatchHandlerConflicts();
 
+            // Apply BatchMessagesOf<T>(b => b.ProbeIndividuallyAfter(N)) as a failure rule on the batch
+            // handler chain. GH-3289.
+            applyBatchProbePolicies();
+
             // Pre-populate the message-type-name cache so the per-message ToMessageTypeName()
             // hot path inside Envelope construction never pays the first-occurrence reflection
             // cost (attribute reads, interface walks, generic-type pretty-printing).
@@ -649,6 +654,34 @@ private void warnOrAssertBatchHandlerConflicts()
         }
     }
 
+    private void applyBatchProbePolicies()
+    {
+        if (Options.BatchDefinitions.Count == 0)
+        {
+            return;
+        }
+
+        foreach (var batch in Options.BatchDefinitions)
+        {
+            if (batch.ProbeIndividuallyAfterAttempts is not { } attempts)
+            {
+                continue;
+            }
+
+            // The failure rule lives on the batch handler chain (the T[] handler), matching any exception:
+            // retry the whole batch until it has failed `attempts` times, then re-run each member as its
+            // own size-1 batch so only the failing one dead-letters.
+            var batchChain = Handlers.ChainFor(batch.Batcher.BatchMessageType);
+            if (batchChain == null)
+            {
+                continue;
+            }
+
+            batchChain.OnException<Exception>()
+                .ContinueWith(new Batching.ProbeIndividuallyContinuationSource(attempts));
+        }
+    }
+
     private void discoverListenersFromConventions()
     {
         // Let any registered routing conventions discover listener endpoints