Stop Balanced-mode host shutdown paying a full agent reply window (GH-3781)

.github/workflows/slow-tests.yml

@@ -0,0 +1,52 @@
+name: slow tests
+
+# Manual only, on purpose.
+#
+# `SlowTests` holds the only real-host reproductions of the GH-3753 agent assignment chain --
+# three multi-node classes driving a 480-agent universe across real Postgres-backed hosts -- plus
+# the TCP and shared-memory transport compliance batteries. Those tests are wall-clock bound by
+# design: they wait out health checks, assignment evaluations and agent starts, so unlike every
+# other suite there is no amount of sharding or parallelism that makes them cheap. The first CI
+# run of the whole project was still going when the standard 20 minute cap killed it.
+#
+# So it does not belong in the PR path, where it would add ~20+ minutes to every push for a suite
+# whose subject barely changes. It belongs HERE, run deliberately: before a release, after anything
+# that touches agent assignment / node lifecycle / durability, and when triaging a report like
+# #3753 or #3781.
+#
+# Run it from the Actions tab, or: gh workflow run "slow tests" --ref <branch>
+on:
+  workflow_dispatch:
+
+env:
+  config: Release
+  disable_test_parallelization: true
+
+jobs:
+  slow-tests:
+    name: CISlowTests
+    runs-on: ubuntu-latest
+    # Deliberately past the 20 minute cap that tests.yml enforces. That cap is a real signal for the
+    # PR matrix -- a job that needs longer is a job that needs splitting -- but it is the wrong tool
+    # here: this suite is slow because the behaviour it reproduces is slow, and it is not gating a
+    # pull request. Still bounded, so a genuine wedge (see #3781) fails the job instead of burning a
+    # runner for six hours.
+    timeout-minutes: 60
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 1
+
+      - name: Setup .NET 10
+        uses: actions/setup-dotnet@v5
+        with:
+          dotnet-version: 10.0.x
+
+      - name: Run Tests
+        run: ./build.sh CISlowTests --framework net9.0
+
+      - name: Stop containers
+        if: always()
+        run: docker compose down

.github/workflows/tests.yml

@@ -61,6 +61,8 @@ jobs:
           - CISqlServer
           - CICircuitBreaking
           - CIAotSmoke
+          # NOTE: CISlowTests is deliberately NOT here. It is wall-clock bound by design and blew
+          # through the 20 minute cap on its first run; it lives in slow-tests.yml as a manual job.
         include:
           # CIMarten was twice killed outright by the runner ~15 minutes in ("The runner has received a
           # shutdown signal"), on unrelated diffs, at almost exactly the same elapsed time on both

build/CITargets.cs

@@ -750,6 +750,36 @@ void runAwsShard(AbsolutePath project, Func<WorkerTest, bool> testFilter = null)
             RunTestProject(tests);
         });
 
+    /// <summary>
+    /// GH-3779 / GH-3781. <c>SlowTests</c> has never run in <b>any</b> CI workflow — no job, no Nuke
+    /// target — even though it holds the only real-host reproductions of the GH-3753 agent
+    /// assignment chain (<c>SlowTests/Agents</c>: three multi-node classes over a 480-agent universe).
+    /// A reproduction nothing runs is a reproduction that rots, and GH-3781 is exactly what it caught
+    /// the moment anyone did run it: a Balanced-mode host that would not finish <c>StopAsync</c>.
+    ///
+    /// <para>Postgres is the only infrastructure the project needs — the SqlServer and Kafka project
+    /// references are transitive, and nothing here opens either.</para>
+    ///
+    /// <para>This is deliberately one unsharded job to begin with: the point is to <i>measure</i> what
+    /// the suite costs on a hosted runner, against the same 20 minute cap every other job answers to.
+    /// If it does not fit, the answer is to shard it the way CIMarten and CIPolecat were sharded (see
+    /// #3350), balanced on the measured per-class durations this job prints — not to raise the cap.</para>
+    /// </summary>
+    Target CISlowTests => _ => _
+        .ProceedAfterFailure()
+        .Executes(() =>
+        {
+            var slowTests = RootDirectory / "src" / "Testing" / "SlowTests" / "SlowTests.csproj";
+
+            // The agent scale classes are wall-clock bound, so overlap the container boot with the
+            // compile rather than paying them serially.
+            LaunchDockerServices("postgresql");
+            BuildTestProjects(slowTests);
+            AwaitDockerServices("postgresql");
+
+            RunTestProject(slowTests);
+        });
+
     // ─── AOT Smoke ──────────────────────────────────────────────────────
     //
     // Builds the Wolverine.AotSmoke project, which sets IsAotCompatible=true +

src/Testing/CoreTests/Runtime/Agents/per_destination_lane_dispatch.cs

@@ -1,4 +1,5 @@
 using System.Collections.Concurrent;
+using System.Diagnostics;
 using JasperFx.Core;
 using Microsoft.Extensions.Logging.Abstractions;
 using Shouldly;
@@ -278,6 +279,100 @@ public async Task a_failing_command_releases_its_claims_so_the_work_can_be_retri
         attempts.ShouldBe(2);
     }
 
+    /// <summary>
+    /// GH-3781. Completing a channel writer does not throw away what is already buffered, so the old
+    /// DisposeAsync executed every command still queued for a node the cluster was leaving -- each costing
+    /// its own AgentBatchTimeouts reply window (25.5 minutes at AgentStartBatchSize = 50) inside
+    /// IHost.StopAsync.
+    /// </summary>
+    [Fact]
+    public async Task disposal_abandons_the_commands_still_queued()
+    {
+        var log = new ConcurrentQueue<string>();
+        var running = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var dispatcher = dispatcherFor(async (command, token) =>
+        {
+            running.TrySetResult();
+            return await execute(command, token);
+        });
+
+        // Both aimed at the same node, so the second sits in the lane behind the first.
+        dispatcher.Enqueue(new GatedCommand("first", NodeA, gate.Task, log));
+        dispatcher.Enqueue(new GatedCommand("second", NodeA, Task.CompletedTask, log));
+
+        await running.Task.WaitAsync(10.Seconds(), TestContext.Current.CancellationToken);
+
+        // DisposeAsync latches and empties the lane queues synchronously, before its first await, so
+        // releasing the gate afterwards is deterministic rather than a race. Ordered this way the
+        // unfixed code fails this test on the assertion below instead of deadlocking the runner --
+        // which matters for a regression test whose subject is a shutdown that never returns.
+        var disposal = withLaneShutdownTimeout(500.Milliseconds(),
+            async () => await dispatcher.DisposeAsync());
+        gate.SetResult();
+        await disposal.WaitAsync(10.Seconds(), TestContext.Current.CancellationToken);
+
+        // Generous, deliberately: the point is that "second" never runs, not that it is merely late.
+        await Task.Delay(500, TestContext.Current.CancellationToken);
+        log.ShouldContain("exit:first");
+        log.ShouldNotContain("enter:second");
+
+        // And nothing is left claimed, or a later leader would suppress re-issuing this work.
+        dispatcher.InFlightAgents.ShouldBeEmpty();
+    }
+
+    /// <summary>
+    /// GH-3781, the backstop. The wedge was one lane parked on a reply from a node that had already gone,
+    /// with teardownAgentsAsync -- and so the node's own deregistration -- queued behind it. Disposal has to
+    /// give up on a lane rather than hold IHost.StopAsync() with it, however that lane came to be stuck.
+    /// </summary>
+    [Fact]
+    public async Task disposal_gives_up_on_a_lane_that_ignores_cancellation()
+    {
+        var log = new ConcurrentQueue<string>();
+        var running = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var never = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var dispatcher = dispatcherFor(async (command, token) =>
+        {
+            running.TrySetResult();
+            // Deliberately not token-aware -- this is the shape of an InvokeAsync sitting out a reply window.
+            await never.Task;
+            return AgentCommands.Empty;
+        });
+
+        dispatcher.Enqueue(new GatedCommand("wedged", NodeA, Task.CompletedTask, log));
+        await running.Task.WaitAsync(10.Seconds(), TestContext.Current.CancellationToken);
+
+        var stopwatch = Stopwatch.StartNew();
+
+        // The outer WaitAsync is what keeps the UNFIXED code from wedging this runner the way it wedges
+        // IHost.StopAsync -- it turns the defect into a failed assertion instead of a hung CI job.
+        await withLaneShutdownTimeout(500.Milliseconds(),
+            async () => await dispatcher.DisposeAsync().AsTask()
+                .WaitAsync(10.Seconds(), TestContext.Current.CancellationToken));
+        stopwatch.Stop();
+
+        stopwatch.Elapsed.ShouldBeLessThan(5.Seconds());
+
+        never.SetResult();
+    }
+
+    private static async Task withLaneShutdownTimeout(TimeSpan timeout, Func<Task> action)
+    {
+        var previous = AgentCommandDispatcher.LaneShutdownTimeout;
+        AgentCommandDispatcher.LaneShutdownTimeout = timeout;
+        try
+        {
+            await action();
+        }
+        finally
+        {
+            AgentCommandDispatcher.LaneShutdownTimeout = previous;
+        }
+    }
+
     [Fact]
     public async Task a_cascade_is_routed_to_the_lane_of_the_node_it_targets()
     {

src/Testing/CoreTests/Runtime/ResponseReply/response_handling.cs

@@ -89,6 +89,55 @@ public async Task timeout_failure()
 #pragma warning restore VSTHRD003 // Avoid awaiting foreign Tasks
 
 
+        _theListener.HasListener(envelope.Id).ShouldBeFalse();
+    }
+
+    /// <summary>
+    /// GH-3781. A caller handing in a token that is ALREADY cancelled has to fail immediately, not sit out
+    /// the reply window. CancellationTokenRegistration runs its callback synchronously for a cancelled
+    /// token, and the listener used to register the caller's token before assigning _completion -- so the
+    /// callback's `_completion?.TrySetException(...)` fired against null and did nothing at all. Every
+    /// shutdown path passes a cancelled token, and for a batched agent command the window it then waited
+    /// out is AgentBatchTimeouts.ReplyWindowFor(chunk): 25.5 minutes at the shipped AgentStartBatchSize of
+    /// 50, paid inside IHost.StopAsync.
+    /// </summary>
+    [Fact]
+    public async Task a_token_that_is_already_cancelled_fails_the_listener_immediately()
+    {
+        var envelope = ObjectMother.Envelope();
+
+        using var cancellation = new CancellationTokenSource();
+        await cancellation.CancelAsync();
+
+        // A window long enough that "waited it out" and "failed fast" cannot be confused.
+        var waiter = _theListener.RegisterListener<Message1>(envelope, cancellation.Token, 30.Seconds());
+
+        waiter.IsCompleted.ShouldBeTrue();
+#pragma warning disable VSTHRD003 // Avoid awaiting foreign Tasks
+        await Should.ThrowAsync<TimeoutException>(() => waiter);
+#pragma warning restore VSTHRD003 // Avoid awaiting foreign Tasks
+
+        // ...and it must not be left in the dictionary: it completed before it was ever registered, so
+        // nothing would take it out again.
+        _theListener.HasListener(envelope.Id).ShouldBeFalse();
+    }
+
+    [Fact]
+    public async Task a_token_cancelled_after_registration_fails_the_listener()
+    {
+        var envelope = ObjectMother.Envelope();
+
+        using var cancellation = new CancellationTokenSource();
+
+        var waiter = _theListener.RegisterListener<Message1>(envelope, cancellation.Token, 30.Seconds());
+        waiter.Status.ShouldBe(TaskStatus.WaitingForActivation);
+
+        await cancellation.CancelAsync();
+
+#pragma warning disable VSTHRD003 // Avoid awaiting foreign Tasks
+        await Should.ThrowAsync<TimeoutException>(() => waiter);
+#pragma warning restore VSTHRD003 // Avoid awaiting foreign Tasks
+
         _theListener.HasListener(envelope.Id).ShouldBeFalse();
     }
 }

src/Wolverine/Runtime/Agents/AgentCommandDispatcher.cs

@@ -1,5 +1,6 @@
 using System.Collections.Concurrent;
 using System.Threading.Channels;
+using JasperFx.Core;
 using Microsoft.Extensions.Logging;
 
 namespace Wolverine.Runtime.Agents;
@@ -50,6 +51,19 @@ internal class AgentCommandDispatcher : IAsyncDisposable
     // out. Same semantics as NodeAgentController's pending-assignment ledger.
     private readonly ConcurrentDictionary<Uri, Guid> _inFlight = new();
 
+    // GH-3781: latched by DisposeAsync so a lane stops picking work up. Completing a channel writer does
+    // NOT discard what is already buffered -- ReadAsync keeps handing it out -- so without this, shutdown
+    // executed every command still queued for a node that had already gone, one reply window at a time.
+    private volatile bool _disposing;
+
+    /// <summary>
+    ///     How long <see cref="DisposeAsync" /> will wait on a single lane before giving up on it. A lane
+    ///     that honours the cancellation token unwinds in microseconds; this only exists so that a future
+    ///     non-cancellable await inside a command can never again wedge <c>IHost.StopAsync()</c>. Settable
+    ///     for tests.
+    /// </summary>
+    internal static TimeSpan LaneShutdownTimeout { get; set; } = 5.Seconds();
+
     internal AgentCommandDispatcher(
         Func<IAgentCommand, CancellationToken, Task<AgentCommands?>> executor,
         ILogger logger,
@@ -86,7 +100,7 @@ public bool TryFindPendingDestination(Uri agentUri, out Guid nodeId)
     /// </summary>
     public void Enqueue(IAgentCommand command)
     {
-        if (_cancellation.IsCancellationRequested) return;
+        if (_cancellation.IsCancellationRequested || _disposing) return;
 
         var destination = command.DestinationNodeId ?? SharedLane;
 
@@ -169,7 +183,7 @@ private Lane laneFor(Guid destination)
 
     private async Task runLaneAsync(Lane lane, Guid destination)
     {
-        while (!_cancellation.IsCancellationRequested)
+        while (!_cancellation.IsCancellationRequested && !_disposing)
         {
             IAgentCommand command;
             try
@@ -185,6 +199,15 @@ private async Task runLaneAsync(Lane lane, Guid destination)
                 return;
             }
 
+            // GH-3781: completing the writer wakes this read with whatever is still buffered, so the
+            // shutdown latch has to be re-checked HERE and not only in the loop condition. Anything
+            // still queued when the node is going down is work for a cluster this node is leaving.
+            if (_disposing)
+            {
+                release(command, destination);
+                return;
+            }
+
             try
             {
                 var cascaded = await _executor(command, _cancellation);
@@ -217,17 +240,36 @@ private async Task runLaneAsync(Lane lane, Guid destination)
 
     public async ValueTask DisposeAsync()
     {
-        foreach (var lane in _lanes.Values)
+        _disposing = true;
+
+        foreach (var pair in _lanes)
         {
-            lane.Queue.Writer.TryComplete();
+            pair.Value.Queue.Writer.TryComplete();
+
+            // Abandon whatever is still buffered rather than executing it on the way out. Each of these
+            // would otherwise cost its own reply window -- AgentBatchTimeouts.ReplyWindowFor(50) is 25.5
+            // minutes -- and they are aimed at a cluster this node is in the middle of leaving.
+            while (pair.Value.Queue.Reader.TryRead(out var abandoned))
+            {
+                release(abandoned, pair.Key);
+            }
         }
 
-        foreach (var lane in _lanes.Values)
+        foreach (var pair in _lanes)
         {
             try
             {
-                var worker = lane.Worker;
-                if (worker != null) await worker;
+                var worker = pair.Value.Worker;
+                if (worker != null) await worker.WaitAsync(LaneShutdownTimeout);
+            }
+            catch (TimeoutException)
+            {
+                // GH-3781: a lane still holding on past the budget must not hold IHost.StopAsync() with it.
+                // The whole wedge was one lane parked on a reply from a node that had already gone, with
+                // teardownAgentsAsync -- and therefore the node's own deregistration -- queued behind it.
+                _logger.LogWarning(
+                    "Agent command lane for node {NodeId} did not finish within {Timeout} during shutdown; abandoning it",
+                    pair.Key == SharedLane ? null : pair.Key, LaneShutdownTimeout);
             }
             catch (Exception)
             {

src/Wolverine/Runtime/RemoteInvocation/ReplyListener.cs

@@ -16,15 +16,23 @@ public ReplyListener(Envelope envelope, ReplyTracker parent, TimeSpan timeout, C
         RequestId = envelope.Id;
         RequestType = envelope.MessageType;
         Parent = parent ?? throw new ArgumentNullException(nameof(parent));
-        cancellationToken.Register(onCancellation);
 
+        // GH-3781: every field onCancellation touches has to be set BEFORE either token is registered.
+        // CancellationTokenRegistration invokes the callback SYNCHRONOUSLY when the token is already
+        // cancelled, so registering the caller's token first meant onCancellation ran against a null
+        // _completion and its `_completion?.TrySetException(...)` silently did nothing -- the listener
+        // then sat out its whole reply window instead of failing fast. Callers pass an already-cancelled
+        // token on every shutdown path, and for a batched agent command that window is
+        // AgentBatchTimeouts.ReplyWindowFor(chunk) -- 25.5 minutes at the shipped AgentStartBatchSize of
+        // 50, paid inside IHost.StopAsync.
         _completion = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);
-        _cancellation = new CancellationTokenSource(timeout);
+        _timeout = timeout;
+        _resultTypes = resultTypes;
 
+        _cancellation = new CancellationTokenSource(timeout);
         _cancellation.Token.Register(onCancellation);
 
-        _timeout = timeout;
-        _resultTypes = resultTypes;
+        cancellationToken.Register(onCancellation);
     }
 
     public string? RequestType { get; set; }