Skip to content

Add explicit subscription and consumer drain#1177

Merged
mtmk merged 17 commits into
release/3.0from
drain-subscriptions
Jun 12, 2026
Merged

Add explicit subscription and consumer drain#1177
mtmk merged 17 commits into
release/3.0from
drain-subscriptions

Conversation

@mtmk

@mtmk mtmk commented Jun 2, 2026

Copy link
Copy Markdown
Member

Adds explicit, dispose-independent subscription draining. INatsSub<T>.DrainAsync drains a single core subscription: stop new deliveries, fence in-flight messages with a PING/PONG, complete the channel, and keep the connection open. For JetStream pull consumers, a new opt-in DrainOnCancel makes cancelling a consume drain gracefully (stop pulling, deliver buffered messages so handlers can still ack on the open connection) instead of stopping abruptly; the default preserves the previous behavior.

Also runs the xUnit suite single-threaded (MaxParallelThreads 1) to remove the cross-collection parallelism behind several earlier integration-test flaps.

Fixes #1176.

Expose DrainAsync on INatsSub<T> so callers can stop receiving new
messages and finish reading already-buffered messages while keeping the
connection usable, instead of only getting drain behaviour on the
dispose path behind DrainSubscriptionsOnDispose.

The drain mechanics (UNSUB, PING/PONG fence, complete channel) move into
a shared core. The dispose-path entry point is renamed to
DrainOnDisposeAsync and keeps the opt-in flag gate and reader wait.
@synadia-claude-reviewer

synadia-claude-reviewer Bot commented Jun 2, 2026

Copy link
Copy Markdown
Contributor

Claude finished @mtmk's task in 4m 25s —— View job


PR Review: Add explicit subscription drain API

  • Gather context and understand the changes
  • Read changed files and analyze for correctness, security, performance
  • Post inline comments on specific issues
  • Write overall review summary

Summary

The refactor is well-structured: DrainCoreAsync cleanly centralises the UNSUB + PING/PONG fence + TryComplete() sequence, and the rename to DrainOnDisposeAsync clarifies the internal contract. The SubscriptionManager split (UnsubscribeForDrainAsync → PING/PONG → RemoveFromRouting) is correct — the SID stays live in the routing table during the fence so in-flight messages aren't dropped, and the cleanup is complete before TryComplete() fires.

Three issues to address before merging:

1. Cancellation contract (moderate) — NatsSubBase.cs:684

When the caller-supplied cancellationToken fires, OperationCanceledException is silently swallowed and TryComplete() still runs. The subscription is committed to the drained state regardless. This violates the standard CancellationToken contract: callers who cancel cannot abort the drain or recover the subscription. At minimum the DrainAsync doc-comment needs a clear warning that cancellation only shortens the PING/PONG fence — the drain is always committed once the method is entered.

2. virtual DrainAsync bypassed by DrainOnDisposeAsync (moderate) — NatsSubBase.cs:294

DrainOnDisposeAsync calls DrainCoreAsync directly, skipping the virtual DrainAsync override. Any JetStream subclass that overrides DrainAsync to add extra pre-drain logic (ack flushing, pull-request cancellation, etc.) will have that logic silently bypassed on the dispose path. This should be resolved before the follow-up JetStream consumer drain PR.

3. DrainPingTimeout scope undocumented (minor) — NatsSubBase.cs:667

NatsOpts.DrainPingTimeout is currently documented as only effective when DrainSubscriptionsOnDispose is true, but DrainCoreAsync now uses it unconditionally for the explicit DrainAsync path. NatsOpts.DrainPingTimeout's doc-comment should be updated accordingly.

4. Unused test field (minor) — SubscriptionDrainTest.cs:8

_output is injected but never used — will produce a compiler warning.


The logic for the SubscriptionManager drain path (idempotent _unsubscribed guard, correct DisposeAsync interaction after DrainAsync, external-manager fallback) is sound. The three test cases cover the key scenarios well.

@github-actions

github-actions Bot commented Jun 2, 2026

Copy link
Copy Markdown

Comment thread src/NATS.Client.Core/NatsSubBase.cs Outdated
Comment thread src/NATS.Client.Core/NatsSubBase.cs
Comment thread tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs Outdated
Comment thread src/NATS.Client.Core/NatsSubBase.cs
mtmk added 3 commits June 3, 2026 08:30
Cancelling ConsumeAsync was an abrupt stop that abandoned messages already buffered in the consumer, so handlers could not finish post-yield work like acking on a still-open connection. Add NatsJSConsumeOpts.DrainOnCancel: when set, cancellation drains the consumer (stop pulling, fence in-flight with PING/PONG, complete the channel) and keeps the connection usable. Default false preserves the previous behavior.
@mtmk

mtmk commented Jun 10, 2026

Copy link
Copy Markdown
Member Author

@claude review recent jetstream changes

@synadia-claude-reviewer

synadia-claude-reviewer Bot commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

Claude finished @mtmk's task in 3m 17s —— View job


Review: JetStream Changes

  • Gather context
  • Read JetStream changed files
  • Analyze git diff for JetStream changes
  • Post review findings

The JetStream changes span two commits — the rename-to-DrainOnDisposeAsync sweep (51b1b16) and the new DrainOnCancel feature (53eaacc). The ConsumeAsync drain loop logic is correct overall, but there are three things worth addressing.


1. _output still unused in ConsumeDrainTest.cs (minor — same issue as before)

ConsumeDrainTest.cs:9-13 injects and stores _output but never calls it:

private readonly ITestOutputHelper _output;   // line 9 — never read

public ConsumeDrainTest(ITestOutputHelper output, NatsServerFixture server)
{
    _output = output;   // line 13 — assigned but unused

The previous review caught this exact issue on SubscriptionDrainTest.cs and mtmk fixed it there — but the same mistake was reintroduced in this new file. Either wire it up (_output.WriteLine(...) in test loops, which would help debug flaps) or remove the field and parameter. Fix this →


2. Negative test assertion is too weak to catch regressions (minor)

ConsumeDrainTest.cs:113:

Assert.True(received < total, $"expected fewer than {total} messages, got {received}");

Cancellation fires after received == 1, so without drain the loop should stop at 1–2 messages. received < 100 would pass even if 98 messages were delivered — a near-complete failure of the default-stop behaviour would go undetected. Consider tightening to something like Assert.True(received <= 5, ...) to actually guard the invariant (prompt stop, not just "less than all of them"). The sibling SubscriptionDrainTest uses < total as well, so this is consistent, but in both tests the cancel-on-first-message scenario makes a tighter bound easily defensible.


3. DrainAsync initiated but WaitToReadAsync(CancellationToken.None) keeps reading on connection failure (design note)

In NatsJSConsumer.ConsumeAsync, once draining = true, the loop switches to WaitToReadAsync(CancellationToken.None). The NatsConnectionFailedException catch block re-throws unconditionally:

catch (NatsConnectionFailedException)
{
    // Connection has permanently failed, stop consuming and rethrow
    throw;
}

This is correct behaviour — if the connection dies during drain, we can't fence in-flight messages and the caller should see the failure. However, this path is not tested: ConsumeDrainTest only covers the happy path and the default-stop path. A test that kills the connection while DrainOnCancel is in progress would confirm that the connection failure propagates cleanly rather than causing the loop to hang on WaitToReadAsync(CancellationToken.None) indefinitely (the channel writer completing with the exception is what unblocks it). This is low priority but would give confidence in the finally-block / drain-task interplay under failure.


What's correct

  • NatsJSConsume.DrainAsync override (NatsJSConsume.cs:199-203) is correct: stopping the heartbeat timer before the base DrainAsync prevents the timer callback from issuing new PullRequests during the UNSUB/PING-PONG fence window. The TryComplete override already completes _pullRequests.Writer, which causes PullLoop to exit cleanly.
  • CallMsgNextAsync guard if (_cancellationToken.IsCancellationRequested) return default ensures no new pull requests are issued once the consumer token is cancelled, even before TryComplete fires. This is a clean early-out.
  • DrainOnDisposeAsync rename in NatsJSFetch / NatsJSOrderedConsume / NatsJSConsume is the right fix: the dispose path now calls the non-virtual DrainOnDisposeAsync directly, avoiding accidental dispatch through the new virtual DrainAsync override.
  • Double-drain safety: the _unsubscribed idempotency guard in SubscriptionManager means DisposeAsync (called when the await using cc scope exits after drain) is safe to call after DrainAsync already completed.
  • drainTask final-block await is correctly guarded by if (draining) — on the default path this is a no-op (default(ValueTask) is already complete), and the comment accurately describes that DrainAsync swallows errors internally so the await won't throw.

mtmk added 2 commits June 10, 2026 11:14
Drop the unused ITestOutputHelper field and tighten the no-drain assertion: cancellation fires after the first message, so assert a prompt stop rather than just fewer than the total.
Set MaxParallelThreads to 1 (was 1.0x, i.e. one thread per core). Integration tests share a server per collection and several historic flaps came from collections racing each other under parallelism; serialising removes that axis at the cost of wall-clock time.
@mtmk mtmk changed the title Add explicit subscription drain API Add explicit subscription and consumer drain Jun 10, 2026
mtmk added 8 commits June 10, 2026 12:34
The reconnects counter is incremented just after the ConnectionOpened event is queued, so the test, which waited on the event and read the metric once, could observe zero. Poll for the measurement after the reconnect, and make MeterTracker hand out locked snapshots since callbacks fire on background threads. Reproduced in isolation (failed 6/6, now passes).
DrainCoreAsync could skip TryComplete() when the UNSUB or PING threw an
unexpected exception (e.g. ObjectDisposedException if the connection is
disposed mid-drain), leaving the message channel forever incomplete; a
drain-on-cancel consumer reading with CancellationToken.None then hangs.
Wrap the unsubscribe and ping in one try, treat the dispose race as
best-effort, and complete the channel and clear routing in a finally so
a reader is never left waiting. Disarm the lifecycle timers before
taking the unsubscribe gate so a timeout callback is less likely to win
the gate and complete the channel without the drain fence.

For JetStream consume, the idle-heartbeat timer could re-arm during a
drain and complete the user channel ahead of the PING/PONG fence,
dropping in-flight messages that drain-on-cancel is meant to deliver.
Track a draining flag so the heartbeat callback, its re-arm, and
CompleteStop defer to the drain.
The base drain only fenced the UNSUB with a PING/PONG; stopping a
subclass delivery engine (the JetStream pull loop and idle-heartbeat
timer) was done ad hoc in NatsJSConsume's DrainAsync override, so only
the explicit cancel path stopped it. The dispose-drain path left the
heartbeat running, where it could complete the user channel ahead of
the fence and drop in-flight messages.

Add a StopDelivery hook on NatsSubBase, called at the start of the
shared drain core so both the explicit and dispose-drain paths stop the
engine before the fence. NatsJSConsume overrides the hook instead of the
whole DrainAsync.
The existing drain tests buffer every message before asserting (publish
then PING, or a count that fits a single pull), so they pass even if the
UNSUB+PING fence is removed. Add a test that publishes without a
pre-drain PING, so the tail is still in flight when the UNSUB is sent,
and asserts all of it is delivered.
The consume loop started the drain in two places: a top-of-loop check
for an already-cancelled token and the WaitToReadAsync cancellation
catch. WaitToReadAsync throws OperationCanceledException when the token
is already cancelled (even with buffered messages), so the catch already
covers the cancel-before-read case identically. Drop the top-of-loop
block and let the catch be the single place that reacts to cancellation.
@mtmk

mtmk commented Jun 11, 2026

Copy link
Copy Markdown
Member Author

@claude review the latest commits on this branch. They fix a drain hang and two message-loss races in subscription/consumer drain, route both the explicit and dispose drain paths through a StopDelivery hook so the consumer's pull/heartbeat engine is stopped before the fence, add an in-flight fence regression test, and drop a redundant drain-start branch in the consume loop. Please focus on the cancellation and teardown paths.

@synadia-claude-reviewer

synadia-claude-reviewer Bot commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Claude finished @mtmk's task in 6m 36s —— View job


Review: Latest drain-subscriptions commits

  • Gather context
  • Read recent commits and diff
  • Analyze cancellation paths
  • Analyze teardown paths
  • Analyze StopDelivery hook routing
  • Check in-flight fence regression test
  • Check redundant drain-start removal
  • Post review findings

Six commits reviewed: d0936a0 (hang fix + message-loss races), 0c99a60 (StopDelivery hook), ed2d86b (in-flight fence test), 01b7eda (redundant drain-start removal), f20b9be (doc scope), 533ced0 (doc clarity).


1. StopDelivery() is outside the exception-safe zone — potential re-hang (moderate)

DrainCoreAsync (NatsSubBase.cs:649-724) now looks like:

private async ValueTask DrainCoreAsync(CancellationToken cancellationToken)
{
    StopDelivery();                    // ← OUTSIDE try/finally
    _timeoutTimer?.Change(...);        // ← OUTSIDE try/finally
    ...
    if (!needsUnsub) return;           // ← early return before try/finally

    try { UNSUB + PING }
    catch (OperationCanceledException) { }
    catch (NatsException) { }
    catch (ObjectDisposedException) { }
    finally
    {
        mgr?.RemoveFromRouting(sid);
        TryComplete();                 // ← only called if we reach try block
    }
}

d0936a0 fixed the hang by wrapping UNSUB and PING inside a try/finally so TryComplete() always runs. But StopDelivery() was placed before that try block. In NatsJSConsume.StopDelivery(), StopHeartbeatTimer() calls _timer.Change(Timeout.Infinite, Timeout.Infinite). A Timer.Change() on an already-disposed timer throws ObjectDisposedException. If that happens, DrainCoreAsync exits before reaching the try block, so TryComplete() is never called, and a drain-on-cancel consumer blocked on WaitToReadAsync(CancellationToken.None) hangs forever — exactly the hang d0936a0 was supposed to prevent.

The timer race is narrow (requires _timer.Dispose() to run before StopDelivery(); in the normal lifecycle DisposeAsync disposes the timer after DrainOnDisposeAsync completes). But StopDelivery() is virtual — future overrides may do more fallible work. The conservative fix is to move StopDelivery() inside the try block (it's idempotent) or catch ObjectDisposedException in NatsJSConsume.StopDelivery() itself.

The same concern applies to _timeoutTimer?.Change(...) etc. immediately after, though those timers are owned by NatsSubBase and disposed in DisposeAsyncbase.DisposeAsync(), so the race window is similar.

Fix this →


2. Cancellation path — correct

ConsumeAsync cancel → drain-on-cancel flow:

  1. WaitToReadAsync(cancellationToken) throws OperationCanceledException
  2. Catch: draining = true; drainTask = cc.DrainAsync(); continue
  3. DrainAsyncDrainCoreAsyncStopDelivery() synchronously sets _draining = true before first await
  4. By the time continue loops back, _draining is visible to timer callbacks; cc.ResetHeartbeatTimer() (called after yield return) no-ops correctly
  5. WaitToReadAsync(CancellationToken.None) blocks until TryComplete() fires from the drain's finally
  6. Channel completes → yield breakfinallyawait drainTask (already complete)

The single-cancel-point simplification from 01b7eda is correct: WaitToReadAsync throws OperationCanceledException even when the token is already cancelled with buffered data, so the catch is authoritative. Removing the pre-check eliminates a duplicated code path that could have diverged.


3. Dispose-drain teardown path — correct

NatsJSConsume.DisposeAsync() ordering:

  1. DrainOnDisposeAsync()DrainCoreAsync(CancellationToken.None)StopDelivery() (timer alive) → UNSUB + PING → TryComplete() in finally
  2. base.DisposeAsync()UnsubscribeAsync() (gate already taken, no-op)
  3. finally: _timer.Dispose(), _pullRequests.Writer.TryComplete()

_timer is disposed only after StopDelivery() and TryComplete() have already run. The _draining flag blocks CompleteStop from racing TryComplete() on the timer thread. Double-drain (explicit drain then dispose) hits the !needsUnsub early return and is safe because the channel is already complete.


4. Heartbeat-racing-fence fix — correct

The three-site _draining guard in NatsJSConsume:

  • Timer callback (line 118): returns early if _draining, skipping CompleteStop notification
  • ResetHeartbeatTimer (line 210): returns early if _draining, so the timer can't re-arm mid-drain
  • CompleteStop callback (line 557): returns early if _draining, leaving completion to the drain's finally

The volatile qualifier on _draining ensures the write from StopDelivery() is visible to timer-thread readers without a memory-ordering window. All three guards are necessary because the timer callback, its re-arm, and the CompleteStop queued callback each represent independent paths that could complete the channel ahead of the fence.


5. In-flight fence regression test — good design, minor note

Drain_delivers_messages_still_in_flight_after_unsub correctly distinguishes itself from the existing test: no pre-drain PingAsync, so the server's socket buffer still contains messages when the UNSUB is sent. The 1000-message count stays well below SubPendingChannelCapacity (correct observation in the comment), exercising the PING/PONG fence rather than back-pressure. Assert.Equal(Enumerable.Range(0, count), received) verifies both count and order — order is guaranteed by NATS per-subscription, so this is tight and appropriate.

Minor: the test is order-sensitive. If NATS ever reorders (it shouldn't on a single subscription), the test would fail with a confusing equality message. Assert.Equal(count, received.Count); Assert.Equal(Enumerable.Range(0, count), received) makes the count failure surface separately. Low priority.


What's correct overall

  • TryComplete() in finally means a reader on the drain-on-cancel path can no longer hang due to UNSUB or PING throwing — this directly fixes the reported hang
  • StopDelivery() hook centralises engine teardown so both explicit and dispose-drain paths stop the heartbeat before the fence — this is the right abstraction and resolves the earlier DrainOnDisposeAsync bypass concern from the first review
  • The redundant top-of-loop cancel check removal is a clean simplification
  • _output unused field is gone from ConsumeDrainTest.cs — fixed
  • drainTask is observed in the finally after MarkReaderInactive(), so drain errors surface to the caller correctly on the non-exceptional path

The one actionable item is moving StopDelivery() (and the timer disarms) inside the try block to close the narrow exception-safety gap it introduces.

mtmk added 2 commits June 11, 2026 12:44
The hang fix put TryComplete() in a finally, but the later StopDelivery
hook and the lifecycle-timer disarm ran before the try, so if any of
them threw (StopDelivery is virtual and calls Timer.Change, which can
throw on a timer disposed by a racing DisposeAsync) the finally was
never reached and a drain-on-cancel reader blocked on the channel could
hang again. Take the unsubscribe gate first, then run the teardown
(decrement, StopDelivery, timer disarms, UNSUB, PING) inside the try so
TryComplete() always runs.

The timers are now disarmed after the gate rather than before it; the
gate was already the real protection (once it is taken a later timer
callback bails on the unsubscribe gate), so the prior before-gate disarm
only narrowed a microscopic window and is not worth skipping the
completion guard for.
@mtmk

mtmk commented Jun 11, 2026

Copy link
Copy Markdown
Member Author

@claude follow-up review of the last two commits (1b258d3, 1525e70), which address your previous review. The first moves StopDelivery and the lifecycle-timer disarms inside DrainCoreAsync's try block so the finally's TryComplete() always runs even if teardown throws (the exception-safety gap you flagged); timers now disarm after the unsubscribe gate, which the gate already protects. The second splits the fence-test assert so a count mismatch surfaces before an order mismatch. Please confirm the completion guarantee holds and flag anything new.

@synadia-claude-reviewer

synadia-claude-reviewer Bot commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Claude finished @mtmk's task in 2m 33s —— View job


Follow-up review: commits 1b258d3 and 1525e70

  • Gather context
  • Read DrainCoreAsync in NatsSubBase.cs — verify StopDelivery + timers inside try block
  • Read fence-test assert split in SubscriptionDrainTest.cs
  • Confirm completion guarantee and flag anything new
  • Post review

Completion guarantee — confirmed

The restructuring in 1b258d3 is correct. The DrainCoreAsync structure is now:

private async ValueTask DrainCoreAsync(CancellationToken cancellationToken)
{
    // Gate taken under lock — atomic, cannot throw
    lock (_gate) { ... needsUnsub ... }

    if (!needsUnsub)
        return;                      // safe: another caller owns TryComplete()

    var sid = -1;
    try
    {
        DecrementActiveSubscription(); // Interlocked.Exchange — effectively no-throw
        StopDelivery();                // virtual; ObjectDisposedException possible → caught below
        _timeoutTimer?.Change(...)×3   // same: caught below
        UNSUB + PING
    }
    catch (OperationCanceledException) { }
    catch (NatsException) { }
    catch (ObjectDisposedException) { }
    finally
    {
        if (sid != -1) mgr?.RemoveFromRouting(sid);
        TryComplete();               // always runs when needsUnsub was true
    }
}

TryComplete() now fires unconditionally whenever needsUnsub is true, both when the catches swallow a failure and when an uncaught exception propagates (C# finally semantics guarantee this). The three catch clauses cover the realistic failure modes from StopDelivery() and the timer Change() calls, so in those paths the method returns normally after TryComplete() — no exception surfaces to the caller.

The commit message's point about the gate as the real protection is correct: once _unsubscribed = true is set under the lock, any racing timer callback that tries to take the gate bails without calling TryComplete(), so moving the timer disarms to after the gate (but inside the try) doesn't open a correctness gap.

One point of correctness I verified: DecrementActiveSubscription uses Interlocked.Exchange(ref _telemetryActive, 0), which is idempotent — it will only fire the telemetry decrement once per subscription lifetime regardless of which path takes the gate. Moving it inside the try doesn't change that invariant.


Fence-test assert split — confirmed

The change in 1525e70:

Assert.Equal(count, received.Count);           // ← new: count surfaces cleanly
Assert.Equal(Enumerable.Range(0, count), received);

Correct. Without the first assertion, a dropped message failure manifests as a confusing sequence-equality diff. With it, the failure message is "expected 1000, actual 999" before the sequence comparison even runs. The same pattern should be applied in ConsumeDrainTest.cs if the fence test there has a similar trailing-only Assert.Equal(Enumerable.Range(...), received) shape.


One pre-existing note (not introduced by these commits)

UnsubscribeAsync (NatsSubBase.cs:249) still has its timer disarms outside the try { TryComplete() } block:

_timeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite);   // outside try
// ... 
try { TryComplete(); } catch { }

This was not changed by these commits and doesn't affect the drain-on-cancel hang path (the WaitToReadAsync(CancellationToken.None) pattern is only in DrainCoreAsync). Flagging for awareness only — it's pre-existing and out of scope here.


Both commits do what they say. The completion guarantee is solid and no new issues are introduced.

@mtmk
mtmk requested a review from scottf June 11, 2026 12:28
@mtmk

mtmk commented Jun 11, 2026

Copy link
Copy Markdown
Member Author

hey @colprog i'd appreciate if you can take a look and see if this works for you when you get a chance. thanks.

@colprog

colprog commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Hi monty, Thanks for your great work. The API looks very good and easy to work with. I assume the same option will be honored by PCGroupConsumer in the next release?

@scottf scottf left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add an explicit API to drain all active subscriptions without disposing the connection

3 participants