feat(batching): ApplyItemException poisons every member of a coalesced key (GH-3289 Phase 3d)

docs/guide/handlers/batching.md

@@ -346,6 +346,10 @@ Throwing the exception *is* the opt-in — there is no configuration to enable.
 their original messages by reference identity, so throw the factory with the actual object(s) handed to your
 handler.
 
+When combined with [`CoalesceBy`](#de-duplicating-a-batch-with-coalesceby), flagging a coalesced item
+poisons **every** original message that collapsed into that key — the handler only sees one message per key,
+but all the de-duplicated members for a poisoned key are dead-lettered together.
+
 `ApplyItemException` is for failures the handler can *name*. For **opaque** failures — where the handler
 throws but cannot tell which item was the culprit — use the `IsolateBatchMembers()` error policy below.
 

src/Testing/CoreTests/Acceptance/batch_coalesce_poison.cs

@@ -0,0 +1,102 @@
+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_coalesce_poison : IAsyncLifetime
+{
+    private IHost _host = null!;
+    private readonly CapturingDeadLetterInterceptor _deadLetters = new();
+
+    public async Task InitializeAsync()
+    {
+        CoalPoisonHandler.Reset();
+
+        _host = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Discovery.DisableConventionalDiscovery();
+                opts.Discovery.IncludeType<CoalPoisonHandler>();
+                opts.Services.AddSingleton<IDeadLetterInterceptor>(_deadLetters);
+
+                opts.BatchMessagesOf<CoalItem>(b =>
+                {
+                    b.BatchSize = 500;
+                    b.TriggerTime = 1.Seconds();
+                    b.CoalesceBy((CoalItem x) => x.Key);
+                }).Sequential();
+            }).StartAsync();
+    }
+
+    public async Task DisposeAsync()
+    {
+        await _host.StopAsync();
+        _host.Dispose();
+    }
+
+    [Fact]
+    public async Task poisoning_a_coalesced_item_dead_letters_every_member_of_that_key()
+    {
+        var bus = _host.MessageBus();
+        // Key "A" has three members; the last-wins (v3) is the poison the handler sees. Key "B" is healthy.
+        await bus.PublishAsync(new CoalItem("A", 1, false));
+        await bus.PublishAsync(new CoalItem("A", 2, false));
+        await bus.PublishAsync(new CoalItem("A", 3, true));
+        await bus.PublishAsync(new CoalItem("B", 1, false));
+
+        await CoalPoisonHandler.SurvivorSucceeded.Task.WaitAsync(10.Seconds());
+        await _deadLetters.Signal.Task.WaitAsync(10.Seconds());
+
+        // Every member that collapsed into the poisoned key "A" is dead-lettered - all three versions.
+        var deadLettered = _deadLetters.DeadLettered.OfType<CoalItem>().ToArray();
+        deadLettered.Where(x => x.Key == "A").Select(x => x.Version).OrderBy(x => x)
+            .ShouldBe(new[] { 1, 2, 3 });
+        deadLettered.ShouldNotContain(x => x.Key == "B");
+
+        // The healthy key survived and was replayed to success.
+        var survivor = CoalPoisonHandler.SuccessfulRuns.ShouldHaveSingleItem();
+        survivor.Select(x => x.Key).ShouldBe(new[] { "B" });
+    }
+}
+
+public record CoalItem(string Key, int Version, bool Poison);
+
+public class CoalPoisonHandler
+{
+    private static readonly object _locker = new();
+
+    public static List<CoalItem[]> SuccessfulRuns { get; } = new();
+    public static TaskCompletionSource SurvivorSucceeded = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+    public static void Reset()
+    {
+        lock (_locker)
+        {
+            SuccessfulRuns.Clear();
+        }
+
+        SurvivorSucceeded = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+    }
+
+    public void Handle(CoalItem[] items)
+    {
+        var poison = items.FirstOrDefault(x => x.Poison);
+        if (poison != null)
+        {
+            // Flag the single coalesced (last-wins) instance the handler sees; Wolverine expands it to
+            // every member that collapsed into that key.
+            throw ApplyItemException.DeadLetterAndReplayOthers(poison);
+        }
+
+        lock (_locker)
+        {
+            SuccessfulRuns.Add(items);
+        }
+
+        SurvivorSucceeded.TrySetResult();
+    }
+}

src/Wolverine/Envelope.Internals.cs

@@ -103,6 +103,12 @@ internal Envelope(object message, IMessageSerializer writer)
     [JsonIgnore]
     internal bool InBatch { get; set; }
 
+    // GH-3289: identifies which coalesced group a batch member belongs to when the batch was assembled by
+    // CoalescingMessageBatcher (CoalesceBy). Members sharing a key get the same id, so ApplyItemException
+    // can poison every member that collapsed into a flagged key. Null for non-coalesced batches.
+    [JsonIgnore]
+    internal int? BatchGroupId { get; set; }
+
     /// <summary>
     /// Set to <c>true</c> by <see cref="Runtime.WolverineRuntime.AcquireInternalEnvelope"/>
     /// when this envelope was pulled from <see cref="Runtime.WolverineRuntime.EnvelopePool"/>.
@@ -524,6 +530,7 @@ internal void Reset()
         Status = default;
         OwnerId = 0;
         InBatch = false;
+        BatchGroupId = null;
         FromPool = false;
         Sender = null;
         Listener = null;

src/Wolverine/Runtime/Batching/ApplyItemContinuation.cs

@@ -57,13 +57,28 @@ public async ValueTask ExecuteAsync(IEnvelopeLifecycle lifecycle, IWolverineRunt
             return;
         }
 
-        // Partition members by reference identity: a thrown poison item is matched to the member
-        // envelope whose Message IS that instance. This is exact for a normal (non-coalesced) batch.
-        // NOTE (GH-3289): under CoalesceBy the handler sees the last-wins instance for a key, which is
-        // the Message of only ONE member, so only that member is dead-lettered here; the earlier
-        // same-key members fall into the survivor set (acked/replayed). Poisoning every member that
-        // collapsed into a coalesced key is a deliberate follow-up (items 2/3).
-        var poison = members.Where(m => containsByReference(_exception.PoisonItems, m.Message)).ToArray();
+        // Partition members by reference identity: a thrown poison item is matched to the member envelope
+        // whose Message IS that instance. This is exact for a normal (non-coalesced) batch.
+        var poison = members.Where(m => containsByReference(_exception.PoisonItems, m.Message)).ToList();
+
+        // GH-3289: under CoalesceBy the handler only sees the last-wins instance for a key (the Message of
+        // ONE member), but every member that collapsed into that key is "about" the same flagged item, so
+        // poison them all. BatchGroupId is stamped per coalesced key by CoalescingMessageBatcher; it is
+        // null for a normal batch, in which case this expansion is a no-op.
+        var poisonGroups = poison.Where(m => m.BatchGroupId.HasValue).Select(m => m.BatchGroupId!.Value)
+            .ToHashSet();
+        if (poisonGroups.Count > 0)
+        {
+            foreach (var member in members)
+            {
+                if (member.BatchGroupId.HasValue && poisonGroups.Contains(member.BatchGroupId.Value)
+                                                 && !poison.Contains(member))
+                {
+                    poison.Add(member);
+                }
+            }
+        }
+
         var survivors = members.Where(m => !poison.Contains(m)).ToArray();
 
         var replay = _exception.Disposition switch
@@ -76,7 +91,7 @@ public async ValueTask ExecuteAsync(IEnvelopeLifecycle lifecycle, IWolverineRunt
         };
 
         // 1. Dead-letter ONLY the poison members (the final CompleteAsync then settles them too).
-        if (poison.Length > 0)
+        if (poison.Count > 0)
         {
             if (lifecycle is MessageContext context)
             {

src/Wolverine/Runtime/Batching/CoalescingMessageBatcher.cs

@@ -49,19 +49,26 @@ public IEnumerable<Envelope> Group(IReadOnlyList<Envelope> envelopes)
                 var key = _keySelector(typed);
 
                 // A null key can't index a Dictionary; never coalesce null-keyed items - keep each.
+                int groupIndex;
                 if (key is not null && indexByKey.TryGetValue(key, out var index))
                 {
                     coalesced[index] = typed;
+                    groupIndex = index;
                 }
                 else
                 {
+                    groupIndex = coalesced.Count;
                     if (key is not null)
                     {
-                        indexByKey[key] = coalesced.Count;
+                        indexByKey[key] = groupIndex;
                     }
 
                     coalesced.Add(typed);
                 }
+
+                // Stamp the member with its coalesced-group index so ApplyItemException can poison every
+                // member that collapsed into a flagged key (GH-3289). Unique per key within this batch.
+                envelope.BatchGroupId = groupIndex;
             }
 
             // ...but EVERY member envelope stays on the batch, so settlement (inbox/outbox tracking and