Add explicit subscription and consumer drain#1177
Conversation
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.
|
Claude finished @mtmk's task in 4m 25s —— View job PR Review: Add explicit subscription drain API
SummaryThe refactor is well-structured: Three issues to address before merging: 1. Cancellation contract (moderate) — When the caller-supplied 2.
3.
4. Unused test field (minor) —
The logic for the SubscriptionManager drain path (idempotent |
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.
|
@claude review recent jetstream changes |
|
Claude finished @mtmk's task in 3m 17s —— View job Review: JetStream Changes
The JetStream changes span two commits — the rename-to- 1.
|
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.
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.
|
@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. |
|
Claude finished @mtmk's task in 6m 36s —— View job Review: Latest drain-subscriptions commits
Six commits reviewed: 1.
|
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.
|
@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. |
|
Claude finished @mtmk's task in 2m 33s —— View job Follow-up review: commits 1b258d3 and 1525e70
Completion guarantee — confirmedThe restructuring in 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
}
}
The commit message's point about the gate as the real protection is correct: once One point of correctness I verified: Fence-test assert split — confirmedThe change in 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 One pre-existing note (not introduced by these commits)
_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 Both commits do what they say. The completion guarantee is solid and no new issues are introduced. |
|
hey @colprog i'd appreciate if you can take a look and see if this works for you when you get a chance. thanks. |
|
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? |
Adds explicit, dispose-independent subscription draining.
INatsSub<T>.DrainAsyncdrains 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-inDrainOnCancelmakes 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.