Fix message loss on consumer dispose

[mtmk]

JetStream consumer dispose was dropping in-flight messages. The per-sub path completed the channel before all wire messages had landed, and connection dispose killed the socket before subs got a chance to drain, leaving NumAckPending pinned until AckWait. NatsSubBase now has a DrainAsync (UNSUB -> PING/PONG -> TryComplete) that per-sub dispose calls to flush the wire cleanly. On connection dispose, drain participants run in parallel so the wait is one round-trip total, not multiplied by the number of subs.

New NatsOpts:

• DrainSubscriptionsOnDispose (default false) - opt-in for the dispose reorder.
• ConsumerDrainOnDisposeTimeout (default null) - wait for active consume loops to finish acking.
• DrainPingTimeout (default 5s) - per-sub UNSUB ack budget.

Defaults preserve current behavior.

[github-actions]

🚀 Deployed on https://69f20098da466c3a8ffff1f1--teal-klepon-7e92d0.netlify.app

[synadia-claude-reviewer]

Claude finished @mtmk's task in 9m 23s —— View job

---

PR #1085 Review: Fix message loss on consumer dispose

• Read CLAUDE.md for repo-specific guidelines
• Examine PR diff
• Review NatsSubBase.cs - DrainAsync implementation
• Review NatsConnection.cs - drain participant changes
• Review NatsOpts.cs - new options
• Review JetStream consumer changes
• Review tests
• Post final review

---

Summary

The PR addresses a real and important bug (channel completed before wire messages land, acks dropped on dispose). The MarkReaderActive/MarkReaderInactive + WaitForReaderDrainAsync mechanism, the new DrainRemainingParticipantsAsync parallel path, and the PING/PONG-gated TryComplete are all sound in concept. A few correctness gaps and one documentation/design mismatch are worth addressing before merge.

---

Issues

1. Routing table cleared before PING/PONG – drain protocol is incomplete

SubscriptionManager.DisposeAsync() calls _bySid.Clear() inside the lock before calling sub.DisposeAsync() (which runs DrainAsync with its PING/PONG round-trip):

// SubscriptionManager.DisposeAsync()
lock (_gate) {
    subRefs = _bySid.Values.Select(m => m.WeakReference).ToArray();
    _bySid.Clear();   // ← routing gone NOW
}
foreach (var subRef in subRefs)
    await sub.DisposeAsync();  // DrainAsync/PING runs here, routing already gone

PublishToClientHandlersAsync also takes lock (_gate) for its _bySid.TryGetValue lookup. Any message the socket reader processes after the lock is released (routing cleared) is silently dropped with a debug log, even though a PING/PONG round-trip follows. The PONG confirms "all bytes have been read from the socket," but it can't deliver messages that were already dropped at the routing layer.

The drain works well for the test scenarios (all batch messages were already in _userMsgs before dispose was called). But for cases where messages are still in the TCP socket buffer at the time the routing table is cleared, they will be dropped — exactly the scenario the PR aims to fix.

A correct drain order would keep the sub in _bySid until the PONG arrives. This likely requires a new DrainAsync-aware path in SubscriptionManager that defers _bySid removal until after the round-trip, rather than using _manager.RemoveAsync(this) (which removes from routing and sends UNSUB atomically).

Fix this →

2. Sequential drain contradicts the "parallel" documentation

The PR description says:

drain participants run in parallel so the wait is one round-trip total, not multiplied by the number of subs

But SubscriptionManager.DisposeAsync() uses a plain foreach + await:

foreach (var subRef in subRefs)
    await sub.DisposeAsync().ConfigureAwait(false);  // sequential – N subs = N × RTT

Only DrainRemainingParticipantsAsync() (for subs already removed from the manager) is actually parallel. For active JetStream subscriptions at connection dispose time, total drain time is N × DrainPingTimeout in the worst case, not one round-trip.

Fix this →

3. PingAsync exception handling is too narrow

In DrainAsync(), only OperationCanceledException is caught:

catch (OperationCanceledException)
{
    // Timeout reached or connection cancelled mid-dispose. Expected on dispose path.
}

If the connection is in Failed or Connecting state, PingAsync first calls ConnectAsync().WaitAsync(cancellationToken). A NatsConnectionFailedException or NatsNoRespondersException from that path would propagate uncaught through sub.DisposeAsync() and crash SubscriptionManager.DisposeAsync(). Consider catching Exception broadly here (or at least NatsException) and logging, since drain is best-effort on dispose.

4. MarkReaderActive() allocates even when drain is disabled

MarkReaderActive() unconditionally creates a TaskCompletionSource and calls RegisterDrainParticipant, even when DrainSubscriptionsOnDispose = false. Since DrainAsync() exits early in that case and never calls WaitForReaderDrainAsync(), the allocation and registration are wasted on every ConsumeAsync/FetchAsync call for users who haven't opted in.

internal void MarkReaderActive()
{
    if (_readerExited is null)
    {
        var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        Interlocked.CompareExchange(ref _readerExited, tcs, null);
    }
    if (Connection is NatsConnection nc)
        nc.RegisterDrainParticipant(this);
}

A guard if (!Connection.Opts.DrainSubscriptionsOnDispose) return; would eliminate the overhead for the common case.

---

Smaller observations (no action required)

• Ordered consumer dispose guard (NatsJSOrderedConsumer.cs:79): The IsDisposed check to break the outer recreate loop is a nice touch; without it, dispose during a consumer-recreate cycle could loop indefinitely.
• ChannelClosedException swallowing (NatsSubBase.ReceiveAsync, line ~330) was already present and correctly handles the case where WriteAsync is called on a completed channel after TryComplete. The removal of the _disposed guard in ReceiveInternalAsync is safe.
• WaitAsync + OperationCanceledException in WaitForReaderDrainCoreAsync: correct pattern; the warning log is helpful.
• Test doubling of TestSessionTimeout (300 s → 600 s): reasonable given the new drain wait in tests.
• DrainRemainingParticipantsAsync parallel path: the Task.WhenAll approach is correct for subs already removed from the manager (naturally-ended fetch/consume loops still being iterated by the user).

[synadia-claude-reviewer]

Claude finished @mtmk's task in 8m 27s —— View job

---

PR #1085 Review: Fix message loss on consumer dispose

• Read CLAUDE.md for repo guidelines
• Examine current PR diff (19 commits, 767+/176-)
• Review NatsSubBase.cs - DrainAsync, MarkReaderActive/Inactive
• Review SubscriptionManager.cs - parallel dispose, UnsubscribeForDrainAsync
• Review NatsConnection.cs - drain path ordering, DrainRemainingParticipantsAsync
• Review JetStream consumer changes and tests
• Post final review

---

Summary

All four issues from the previous review are addressed. The routing-before-PING fix (UnsubscribeForDrainAsync/RemoveFromRouting), the parallel Task.WhenAll in SubscriptionManager.DisposeAsync, the broader exception catch in DrainAsync, and the MarkReaderActive allocation guard are all correctly implemented. The drain ordering (UNSUB → PING/PONG → remove-from-routing → TryComplete) and the overall connection dispose sequence (drain subs → drain remaining participants → flush writer → close socket) are sound. A couple of smaller points below.

---

Issues

1. RemoveFromRouting(-1) when sub was already removed from _bySub

UnsubscribeForDrainAsync returns -1 when the sub is not found in _bySub (already removed via RemoveAsync on natural completion). The SID is then passed to RemoveFromRouting(-1), which calls _bySid.TryRemove(-1, ...) — a harmless no-op, but it's an unexpected sentinel flowing into a routing-table operation. A minimal guard makes intent explicit:

// src/NATS.Client.Core/NatsSubBase.cs  – inside DrainAsync(), after the PING/PONG
mgr?.RemoveFromRouting(sid);   // today: called even when sid == -1

// suggested
if (sid != -1)
    mgr?.RemoveFromRouting(sid);

Fix this →

2. NatsConnection.DisposeAsync drain path skips socket/writer cleanup on sub-dispose exception

If any sub's stored Exception (set during a JetStream protocol error) causes NatsSubBase.DisposeAsync() to re-throw, that exception propagates through Task.WhenAll inside SubscriptionManager.DisposeAsync(), and the drain-path sequence in NatsConnection.DisposeAsync is abandoned:

// NatsConnection.cs — drain path (no try/finally)
await _subscriptionManager.DisposeAsync().ConfigureAwait(false);  // can throw
await DrainRemainingParticipantsAsync().ConfigureAwait(false);    // skipped on throw
await CommandWriter.DisposeAsync().ConfigureAwait(false);          // skipped on throw
await DisposeSocketAsync(false).ConfigureAwait(false);             // skipped — socket leaks

This is pre-existing (the sequential path had the same skip-on-throw behaviour), and Task.WhenAll actually improves things by still disposing every sub before propagating. But it's now more likely to surface since the parallel path is the drain-path default, so it's worth wrapping in a try/finally:

try { await _subscriptionManager.DisposeAsync().ConfigureAwait(false); }
finally
{
    await DrainRemainingParticipantsAsync().ConfigureAwait(false);
    await CommandWriter.DisposeAsync().ConfigureAwait(false);
    await DisposeSocketAsync(false).ConfigureAwait(false);
}

Fix this →

---

Smaller observations (no action required)

• RemoveFromRouting uses lock(_gate) around a ConcurrentDictionary.TryRemove: Technically redundant for thread-safety (the ConcurrentDictionary is already safe), but it coordinates correctly with PublishToClientHandlersAsync which also takes _gate before reading _bySid. Belt-and-suspenders is fine here.
• MarkReaderInactive() has no DrainSubscriptionsOnDispose guard: It calls UnregisterDrainParticipant and TrySetResult unconditionally, but both are no-ops when drain is disabled (_drainParticipants never contains the sub; _readerExited is null). Harmless.
• Double-drain via DrainRemainingParticipantsAsync when ConsumerDrainOnDisposeTimeout = null: If a sub's WaitForReaderDrainAsync() returns immediately (no timeout), the sub may still be in _drainParticipants when DrainRemainingParticipantsAsync runs. The second DrainAsync call hits the _unsubscribed guard immediately and falls straight through to WaitForReaderDrainAsync() again — a no-op. Acceptable.
• IsDisposed check before RecreateConsumer in NatsJSOrderedConsumer: Correctly prevents the outer recreate loop from spinning indefinitely when the connection is disposed mid-loop.
• Tests: The three new *_connection_dispose_drains_buffered_messages tests cover Consume, Fetch, and OrderedConsumer with proper bail-at / total-msgs structure and post-dispose ack-floor assertions. Good coverage of the actual bug.

[synadia-claude-reviewer]

Claude finished @mtmk's task in 7m 51s —— View job

---

PR #1085 Review: Fix message loss on consumer dispose

• Read CLAUDE.md for repo guidelines
• Examine PR diff
• Review NatsSubBase.cs - DrainAsync, MarkReaderActive/Inactive
• Review SubscriptionManager.cs - UnsubscribeForDrainAsync, parallel dispose
• Review NatsConnection.cs - drain path ordering, try/finally
• Review NatsOpts.cs - new options
• Review JetStream consumer changes
• Review tests
• Post final review

---

Summary

All six issues from the two previous reviews are correctly addressed. The drain ordering (UnsubscribeForDrainAsync → PING/PONG → RemoveFromRouting → TryComplete), the parallel Task.WhenAll in SubscriptionManager.DisposeAsync, the try/finally in NatsConnection.DisposeAsync, the sid != -1 guard, the exception catch broadening, and the MarkReaderActive allocation guard are all correctly implemented.

One new issue and one smaller observation below.

---

Issues

1. _disposed field is dead code in all three JetStream sub classes

NatsJSConsume, NatsJSFetch, and NatsJSOrderedConsume each declare private int _disposed and write it via Interlocked.Exchange(ref _disposed, 1) in DisposeAsync(), but the return value is discarded and the field is never read again. Its original purpose was to gate ReceiveInternalAsync message delivery (if (Volatile.Read(ref _disposed) == 0) ...), which was correctly removed by this PR to ensure all buffered messages are delivered. The field is now purely vestigial.

Because the return value is discarded, two concurrent DisposeAsync() calls both proceed through DrainAsync() (safe via the _unsubscribed lock guard) and both reach base.DisposeAsync() (safe via NatsSubBase._disposed bool guard). No bug results, but the exchange without a read is misleading — it implies idempotency protection that isn't actually provided by that field.

// NatsJSConsume.DisposeAsync (same pattern in Fetch and OrderedConsume)
Interlocked.Exchange(ref _disposed, 1);  // ← return value discarded, field never read

Either remove the field entirely (the base class _disposed bool already provides idempotency), or add the guard:

if (Interlocked.Exchange(ref _disposed, 1) != 0)
{
    await base.DisposeAsync().ConfigureAwait(false);
    return;
}

Fix this →

---

Smaller observations (no action required)

• Parallel Task.WhenAll now throws AggregateException: Previously, sequential disposal propagated the first sub exception directly. Now multiple sub exceptions are wrapped in AggregateException. Any callers catching specific exception types (e.g., NatsSubException) from nats.DisposeAsync() may need updating, though this affects a very narrow API surface.

• External ISubscriptionManager drain limitation acknowledged: The fallback to _manager.RemoveAsync(this) for non-SubscriptionManager managers (removes routing before PING/PONG) is correctly noted in comments. The full drain guarantee only holds for the built-in manager.

• _pullTask await ordering in NatsJSConsume.DisposeAsync finally: DrainAsync() → TryComplete() → _pullRequests.Writer.TryComplete() causes PullLoop to exit before await _pullTask in finally runs. Ordering is correct.

• Tests are well-structured: The three *_connection_dispose_drains_buffered_messages tests cover all three consumer types with proper bail-at/total-msgs structure and post-dispose NumAckPending == 0 assertions.

[synadia-claude-reviewer]

The int _disposed field is now dead code — Interlocked.Exchange is called but the return value is discarded and _disposed is never read. Its original use in ReceiveInternalAsync (if (Volatile.Read(ref _disposed) == 0)) was correctly removed by this PR. Idempotency is provided by NatsSubBase._disposed (in base.DisposeAsync()). Either drop the field or guard on the return value:

Suggested change

	Interlocked.Exchange(ref _disposed, 1);
	if (Interlocked.Exchange(ref _disposed, 1) != 0)
	{
	await base.DisposeAsync().ConfigureAwait(false);
	return;
	}

[synadia-claude-reviewer]

Same dead-code pattern as NatsJSConsume — _disposed is written but never read. Consider adding the early-exit guard or removing the field:

Suggested change

	Interlocked.Exchange(ref _disposed, 1);
	if (Interlocked.Exchange(ref _disposed, 1) != 0)
	{
	await base.DisposeAsync().ConfigureAwait(false);
	return;
	}

[mtmk]

@claude fixed your comments, review again

[synadia-claude-reviewer]

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

---

PR #1085 Review: Fix message loss on consumer dispose

• Read CLAUDE.md for repo guidelines
• Examine current PR diff
• Review NatsJSConsume.cs - _disposed dead code removed
• Review NatsJSFetch.cs - _disposed dead code removed
• Review NatsJSOrderedConsume.cs - _disposed dead code removed
• Review NatsSubBase.cs - DrainAsync, MarkReaderActive/Inactive
• Review SubscriptionManager.cs - UnsubscribeForDrainAsync, parallel dispose
• Review NatsConnection.cs - drain path ordering, try/finally
• Post final review

---

Summary

All issues from the three previous review rounds are correctly addressed. The latest commit (fd26de8) removes the dead _disposed field and its useless Interlocked.Exchange call from NatsJSConsume, NatsJSFetch, and NatsJSOrderedConsume — the cleanest possible fix (field removal rather than adding the guard, since NatsSubBase._disposed already provides idempotency).

The complete drain ordering is now correct end-to-end:

• UnsubscribeForDrainAsync: removes from _bySub, sends UNSUB, retains in _bySid
• PING/PONG: confirms server processed UNSUB and socket reader delivered all in-flight messages
• RemoveFromRouting(sid): gated on sid != -1, removes from _bySid
• TryComplete(): safe to close the channel now
• WaitForReaderDrainAsync: waits for the user iterator to drain buffered messages (bounded by ConsumerDrainOnDisposeTimeout)
• SubscriptionManager.DisposeAsync: Task.WhenAll makes all sub drains run in parallel
• NatsConnection.DisposeAsync: try/finally ensures socket and writer cleanup even if a sub throws

No new issues found. The PR is in good shape.

---

Smaller observations (no action required)

• Exception catch breadth in DrainAsync: The two-catch pattern (OperationCanceledException then NatsException) is correct. OperationCanceledException is not a NatsException so both are needed. Non-NatsException exceptions from CommandWriter.PingAsync (e.g. hypothetical ChannelClosedException) would propagate through Task.WhenAll and then be absorbed by the try/finally in NatsConnection.DisposeAsync, so cleanup always runs regardless.

• WaitForReaderDrainAsync called when needsUnsub = false: When a sub already completed naturally (_unsubscribed = true via EndSubscription/UnsubscribeAsync), DrainAsync skips the UNSUB block and falls through to WaitForReaderDrainAsync. This is correct: if the reader is still active, waiting up to ConsumerDrainOnDisposeTimeout is exactly the desired behavior.

• SubscriptionManager.DisposeAsync no longer pre-clears _bySid: The new code leaves _bySid entries intact until each sub's RemoveFromRouting call, so messages can be delivered during the drain window. Subs that already completed naturally are already gone from _bySid (via RemoveAsync), so they won't be double-disposed.

• AggregateException from parallel Task.WhenAll: Pre-existing callers catching specific exceptions (e.g. NatsSubException) from nats.DisposeAsync() may need updating, but this affects a very narrow API surface.

[scottf]

LGTM