From 51b1b1616456de8e33b6f9a1a3a3d776f150f459 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Tue, 2 Jun 2026 10:30:19 +0100 Subject: [PATCH 01/14] core: add explicit subscription drain API Expose DrainAsync on INatsSub 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. --- src/NATS.Client.Core/INatsSub.cs | 15 ++ src/NATS.Client.Core/NatsConnection.cs | 2 +- src/NATS.Client.Core/NatsSubBase.cs | 145 +++++++++++------- .../Internal/NatsJSConsume.cs | 2 +- .../Internal/NatsJSFetch.cs | 2 +- .../Internal/NatsJSOrderedConsume.cs | 2 +- .../SubscriptionDrainTest.cs | 92 +++++++++++ 7 files changed, 198 insertions(+), 62 deletions(-) create mode 100644 tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs diff --git a/src/NATS.Client.Core/INatsSub.cs b/src/NATS.Client.Core/INatsSub.cs index a9b7a2c36..9ce8940f2 100644 --- a/src/NATS.Client.Core/INatsSub.cs +++ b/src/NATS.Client.Core/INatsSub.cs @@ -27,4 +27,19 @@ public interface INatsSub : IAsyncDisposable /// /// A that represents the asynchronous server UNSUB operation. public ValueTask UnsubscribeAsync(); + + /// + /// Drain the subscription: stop receiving new messages while letting messages + /// already buffered be read, keeping the connection usable. + /// + /// + /// Sends an unsubscribe to the server, waits for a PING/PONG round-trip so messages + /// already in flight are delivered, then completes . Keep reading + /// until it completes to process the buffered messages. Unlike + /// , drain does not drop messages still in the socket + /// buffer. The connection stays open. + /// + /// Bounds the best-effort PING/PONG fence wait. + /// A that represents the asynchronous drain operation. + public ValueTask DrainAsync(CancellationToken cancellationToken = default); } diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index f80b556c9..fdcad7d45 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/src/NATS.Client.Core/NatsConnection.cs @@ -1235,7 +1235,7 @@ private ValueTask DrainRemainingParticipantsAsync() var tasks = new Task[remaining.Length]; for (var i = 0; i < remaining.Length; i++) - tasks[i] = remaining[i].DrainAsync().AsTask(); + tasks[i] = remaining[i].DrainOnDisposeAsync().AsTask(); return new ValueTask(Task.WhenAll(tasks)); } diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index 26893e799..87d6a54ce 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -276,6 +276,23 @@ public ValueTask UnsubscribeAsync() return _manager.RemoveAsync(this); } + /// + /// Drains the subscription: sends UNSUB so the server stops delivering new + /// messages, waits for a PING/PONG round-trip so messages already in flight + /// are received, then completes the message channel. Buffered messages stay + /// available to read and the connection remains usable. + /// + /// + /// Unlike , which completes the channel right + /// away and can drop messages still in the socket buffer, drain fences inbound + /// delivery with a PING/PONG so buffered messages are preserved. The fence is + /// best-effort and bounded by and + /// . + /// + /// Bounds the best-effort PING/PONG fence wait. + /// A that represents the asynchronous drain operation. + public virtual ValueTask DrainAsync(CancellationToken cancellationToken = default) => DrainCoreAsync(cancellationToken); + /// /// Disposes the instance asynchronously. /// @@ -403,10 +420,9 @@ internal static bool IsHeader503(ReadOnlySequence? headersBuffer) => internal void ClearException() => Interlocked.Exchange(ref _exception, null); /// - /// Sends UNSUB to the server, waits for PING/PONG round-trip to ensure - /// all in-flight messages are processed, then completes the channel. - /// This avoids the race where TryComplete() closes the channel while - /// the socket reader is still delivering messages. + /// Dispose-path drain. Runs the same UNSUB plus PING/PONG fence as + /// , then waits for an active user reader to drain + /// the channel before returning. /// /// /// No-op unless is @@ -415,64 +431,12 @@ internal static bool IsHeader503(ReadOnlySequence? headersBuffer) => /// ). /// /// A that represents the asynchronous drain operation. - internal async ValueTask DrainAsync() + internal async ValueTask DrainOnDisposeAsync() { if (!Connection.Opts.DrainSubscriptionsOnDispose) return; - var needsUnsub = false; - lock (_gate) - { - if (!_unsubscribed) - { - _unsubscribed = true; - needsUnsub = true; - } - } - - if (needsUnsub) - { - DecrementActiveSubscription(); - - _timeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - _idleTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - _startUpTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - - // For the built-in SubscriptionManager: send UNSUB but keep the SID in the routing - // table until after the PING/PONG round-trip, so messages already in the socket - // buffer are delivered before the channel is completed. - // For external managers: fall back to RemoveAsync (routing removed before PING). - var mgr = _manager as SubscriptionManager; - var sid = mgr != null - ? await mgr.UnsubscribeForDrainAsync(this).ConfigureAwait(false) - : -1; - if (mgr == null) - await _manager.RemoveAsync(this).ConfigureAwait(false); - - // PING/PONG round-trip: after we get the PONG back, we know the server - // has processed our UNSUB and the socket reader has processed all messages - // that were in-flight before the UNSUB was received by the server. - try - { - using var pingCts = new CancellationTokenSource(Connection.Opts.DrainPingTimeout); - await Connection.PingAsync(pingCts.Token).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - // Timeout via DrainPingTimeout. Expected on dispose path. - } - catch (NatsException) - { - // Connection failed or disconnected during drain. Drain is best-effort. - } - - // Remove from routing now that no more messages can arrive for this SID. - if (sid != -1) - mgr?.RemoveFromRouting(sid); - - // Now it's safe to complete the channel, no more messages will arrive. - TryComplete(); - } + await DrainCoreAsync(CancellationToken.None).ConfigureAwait(false); // Always wait for an active user reader to drain the channel before // returning, even if the subscription was already ended (e.g. via @@ -655,6 +619,71 @@ private void DecrementActiveSubscription() Telemetry.ActiveSubscriptions.Add(-1, Telemetry.BuildMetricTags(Connection, Telemetry.Constants.OpSub)); } + /// + /// Sends UNSUB to the server, waits for a PING/PONG round-trip to ensure all + /// in-flight messages are processed, then completes the channel. This avoids + /// the race where TryComplete() closes the channel while the socket reader is + /// still delivering messages. Idempotent: a no-op once the subscription has + /// already been unsubscribed. + /// + private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) + { + var needsUnsub = false; + lock (_gate) + { + if (!_unsubscribed) + { + _unsubscribed = true; + needsUnsub = true; + } + } + + if (!needsUnsub) + return; + + DecrementActiveSubscription(); + + _timeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + _idleTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + _startUpTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + + // For the built-in SubscriptionManager: send UNSUB but keep the SID in the routing + // table until after the PING/PONG round-trip, so messages already in the socket + // buffer are delivered before the channel is completed. + // For external managers: fall back to RemoveAsync (routing removed before PING). + var mgr = _manager as SubscriptionManager; + var sid = mgr != null + ? await mgr.UnsubscribeForDrainAsync(this).ConfigureAwait(false) + : -1; + if (mgr == null) + await _manager.RemoveAsync(this).ConfigureAwait(false); + + // PING/PONG round-trip: after we get the PONG back, we know the server + // has processed our UNSUB and the socket reader has processed all messages + // that were in-flight before the UNSUB was received by the server. + try + { + using var pingCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + pingCts.CancelAfter(Connection.Opts.DrainPingTimeout); + await Connection.PingAsync(pingCts.Token).ConfigureAwait(false); + } + catch (OperationCanceledException) + { + // Drain ping timed out (DrainPingTimeout) or was cancelled. Drain is best-effort. + } + catch (NatsException) + { + // Connection failed or disconnected during drain. Drain is best-effort. + } + + // Remove from routing now that no more messages can arrive for this SID. + if (sid != -1) + mgr?.RemoveFromRouting(sid); + + // Now it's safe to complete the channel, no more messages will arrive. + TryComplete(); + } + private async ValueTask WaitForReaderDrainCoreAsync(Task task, TimeSpan timeout) { try diff --git a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs index ca276d7c1..eb8055986 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs @@ -222,7 +222,7 @@ public override async ValueTask DisposeAsync() { try { - await DrainAsync().ConfigureAwait(false); + await DrainOnDisposeAsync().ConfigureAwait(false); await base.DisposeAsync().ConfigureAwait(false); } finally diff --git a/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs b/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs index 5dcf6d429..6abc50b9a 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs @@ -163,7 +163,7 @@ public override async ValueTask DisposeAsync() { try { - await DrainAsync().ConfigureAwait(false); + await DrainOnDisposeAsync().ConfigureAwait(false); await base.DisposeAsync().ConfigureAwait(false); } finally diff --git a/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs index cc2cacf4a..e595cc348 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs @@ -140,7 +140,7 @@ public override async ValueTask DisposeAsync() _context.Connection.ConnectionDisconnected -= ConnectionOnConnectionDisconnected; try { - await DrainAsync().ConfigureAwait(false); + await DrainOnDisposeAsync().ConfigureAwait(false); await base.DisposeAsync().ConfigureAwait(false); } finally diff --git a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs new file mode 100644 index 000000000..c162f2d6a --- /dev/null +++ b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs @@ -0,0 +1,92 @@ +using NATS.Client.Core2.Tests; + +namespace NATS.Client.Core.Tests; + +[Collection("nats-server")] +public class SubscriptionDrainTest +{ + private readonly ITestOutputHelper _output; + private readonly NatsServerFixture _server; + + public SubscriptionDrainTest(ITestOutputHelper output, NatsServerFixture server) + { + _output = output; + _server = server; + } + + [Fact] + public async Task Drain_preserves_buffered_messages_and_completes_channel() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + const int count = 10; + for (var i = 0; i < count; i++) + await nats.PublishAsync(subject, i, cancellationToken: cancellationToken); + + // PING/PONG round-trip guarantees the published messages have been + // delivered back to our subscription channel before we drain. + await nats.PingAsync(cancellationToken); + + await sub.DrainAsync(cancellationToken); + + // All buffered messages are still readable and the channel completes. + var received = new List(); + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken)) + received.Add(msg.Data); + + Assert.Equal(Enumerable.Range(0, count), received); + } + + [Fact] + public async Task Drain_stops_new_messages_but_connection_stays_usable() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + await nats.PublishAsync(subject, 1, cancellationToken: cancellationToken); + await nats.PingAsync(cancellationToken); + + await sub.DrainAsync(cancellationToken); + + // Published after the UNSUB; must not reach the drained subscription. + await nats.PublishAsync(subject, 2, cancellationToken: cancellationToken); + await nats.PingAsync(cancellationToken); + + var received = new List(); + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken)) + received.Add(msg.Data); + + Assert.Equal(new[] { 1 }, received); + + // The connection is still usable for a fresh subscription and publish. + var sub2 = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + await nats.PublishAsync(subject, 42, cancellationToken: cancellationToken); + var next = await sub2.Msgs.ReadAsync(cancellationToken); + Assert.Equal(42, next.Data); + await sub2.DisposeAsync(); + } + + [Fact] + public async Task Drain_is_idempotent_and_dispose_after_drain_is_safe() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + await sub.DrainAsync(cancellationToken); + await sub.DrainAsync(cancellationToken); + await sub.DisposeAsync(); + } +} From 130041bc19588b566119f93802909975182f5b80 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 3 Jun 2026 08:30:21 +0100 Subject: [PATCH 02/14] core: clarify drain cancellation and DrainPingTimeout docs --- src/NATS.Client.Core/NatsOpts.cs | 7 ++++--- src/NATS.Client.Core/NatsSubBase.cs | 8 +++++++- tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs | 4 +--- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src/NATS.Client.Core/NatsOpts.cs b/src/NATS.Client.Core/NatsOpts.cs index d5c59e9f7..49a482a2f 100644 --- a/src/NATS.Client.Core/NatsOpts.cs +++ b/src/NATS.Client.Core/NatsOpts.cs @@ -191,14 +191,15 @@ public sealed record NatsOpts public TimeSpan? ConsumerDrainOnDisposeTimeout { get; init; } = null; /// - /// Per-subscription PING/PONG timeout used during drain on dispose. + /// Per-subscription PING/PONG timeout used during drain. /// (default: 5 seconds) /// /// /// Bounds how long drain will wait for the server to ack the UNSUB before /// completing the user channel anyway. Increase for high-latency networks - /// where a single round-trip can exceed the default. Only effective when - /// is true. + /// where a single round-trip can exceed the default. Applies both to drain on + /// dispose (when is true) and to the + /// explicit per-subscription drain API. /// public TimeSpan DrainPingTimeout { get; init; } = TimeSpan.FromSeconds(5); diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index 87d6a54ce..e1a294cb5 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -288,8 +288,14 @@ public ValueTask UnsubscribeAsync() /// delivery with a PING/PONG so buffered messages are preserved. The fence is /// best-effort and bounded by and /// . + /// + /// Cancellation only shortens the PING/PONG fence wait; it cannot abort the + /// drain. Once is entered the subscription is always + /// committed to the drained state, and cancelling the token does not throw or + /// keep the subscription alive. + /// /// - /// Bounds the best-effort PING/PONG fence wait. + /// Bounds the best-effort PING/PONG fence wait. Cancellation does not abort the drain. /// A that represents the asynchronous drain operation. public virtual ValueTask DrainAsync(CancellationToken cancellationToken = default) => DrainCoreAsync(cancellationToken); diff --git a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs index c162f2d6a..73966b2d3 100644 --- a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs +++ b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs @@ -5,12 +5,10 @@ namespace NATS.Client.Core.Tests; [Collection("nats-server")] public class SubscriptionDrainTest { - private readonly ITestOutputHelper _output; private readonly NatsServerFixture _server; - public SubscriptionDrainTest(ITestOutputHelper output, NatsServerFixture server) + public SubscriptionDrainTest(NatsServerFixture server) { - _output = output; _server = server; } From 53eaacc67496b1f4ab62ac0a412fd07cadf280c3 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 10 Jun 2026 10:41:01 +0100 Subject: [PATCH 03/14] jetstream: add opt-in consume drain on cancel 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. --- .../Internal/NatsJSConsume.cs | 6 + src/NATS.Client.JetStream/NatsJSConsumer.cs | 40 +++++- src/NATS.Client.JetStream/NatsJSOpts.cs | 8 ++ .../ConsumeDrainTest.cs | 115 ++++++++++++++++++ 4 files changed, 166 insertions(+), 3 deletions(-) create mode 100644 tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs diff --git a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs index eb8055986..5faa169b2 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs @@ -196,6 +196,12 @@ public ValueTask CallMsgNextAsync(string origin, ConsumerGetnextRequest request, cancellationToken: cancellationToken); } + public override ValueTask DrainAsync(CancellationToken cancellationToken = default) + { + StopHeartbeatTimer(); + return base.DrainAsync(cancellationToken); + } + public void StopHeartbeatTimer() => _timer.Change(Timeout.Infinite, Timeout.Infinite); public void ResetHeartbeatTimer() => _timer.Change(_hbTimeout, _hbTimeout); diff --git a/src/NATS.Client.JetStream/NatsJSConsumer.cs b/src/NATS.Client.JetStream/NatsJSConsumer.cs index dd66e0b43..39a825758 100644 --- a/src/NATS.Client.JetStream/NatsJSConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSConsumer.cs @@ -66,20 +66,44 @@ public async IAsyncEnumerable> ConsumeAsync( [EnumeratorCancellation] CancellationToken cancellationToken = default) { opts ??= _context.Opts.DefaultConsumeOpts; + var drainOnCancel = opts.DrainOnCancel; await using var cc = await ConsumeInternalAsync(serializer, opts, cancellationToken).ConfigureAwait(false); cc.MarkReaderActive(); + + // When draining on cancel we keep reading buffered messages until the channel + // completes instead of bailing out immediately. The drain stops pulling, fences + // in-flight messages with a PING/PONG and completes the channel, leaving the + // connection usable so handlers can still ACK the messages they receive. + var draining = false; + var drainTask = default(ValueTask); try { - while (!cancellationToken.IsCancellationRequested) + while (true) { + if (cancellationToken.IsCancellationRequested && !draining) + { + if (!drainOnCancel) + yield break; + + draining = true; + drainTask = cc.DrainAsync(); + } + // We have to check calls individually since we can't use yield return in try-catch blocks. bool ready; try { - ready = await cc.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + ready = await cc.Msgs.WaitToReadAsync(draining ? CancellationToken.None : cancellationToken).ConfigureAwait(false); } catch (OperationCanceledException) { + if (drainOnCancel && !draining) + { + draining = true; + drainTask = cc.DrainAsync(); + continue; + } + ready = false; } catch (NatsConnectionFailedException) @@ -96,8 +120,13 @@ public async IAsyncEnumerable> ConsumeAsync( if (!ready) yield break; - while (!cancellationToken.IsCancellationRequested) + while (true) { + // Stop promptly on cancel unless we are draining, in which case we + // keep reading the buffered messages until the channel completes. + if (!draining && cancellationToken.IsCancellationRequested) + break; + bool read; NatsJSMsg jsMsg; try @@ -135,6 +164,11 @@ public async IAsyncEnumerable> ConsumeAsync( finally { cc.MarkReaderInactive(); + + // Observe the drain completion. DrainAsync swallows cancellation and + // connection errors internally, so this does not throw on the normal path. + if (draining) + await drainTask.ConfigureAwait(false); } } diff --git a/src/NATS.Client.JetStream/NatsJSOpts.cs b/src/NATS.Client.JetStream/NatsJSOpts.cs index e36868d90..6a78b2dbb 100644 --- a/src/NATS.Client.JetStream/NatsJSOpts.cs +++ b/src/NATS.Client.JetStream/NatsJSOpts.cs @@ -160,6 +160,14 @@ public record NatsJSConsumeOpts /// Set to -1 to disable this check. (default: 10) /// public int MaxConsecutive503Errors { get; init; } = 10; + + /// + /// When true, cancelling the token passed to ConsumeAsync drains the consumer + /// gracefully (stop pulling, deliver buffered messages so handlers can ACK, + /// then complete) and leaves the connection usable, instead of stopping + /// immediately. (default: false) + /// + public bool DrainOnCancel { get; init; } } /// diff --git a/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs b/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs new file mode 100644 index 000000000..ec3e92714 --- /dev/null +++ b/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs @@ -0,0 +1,115 @@ +using NATS.Client.Core2.Tests; +using NATS.Client.TestUtilities2; + +namespace NATS.Client.JetStream.Tests; + +[Collection("nats-server")] +public class ConsumeDrainTest +{ + private readonly ITestOutputHelper _output; + private readonly NatsServerFixture _server; + + public ConsumeDrainTest(ITestOutputHelper output, NatsServerFixture server) + { + _output = output; + _server = server; + } + + [Fact] + public async Task Consume_drain_on_cancel_delivers_buffered_and_keeps_connection() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + await nats.ConnectRetryAsync(); + var prefix = _server.GetNextId(); + + var js = new NatsJSContext(nats); + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + + await js.CreateStreamAsync($"{prefix}s1", new[] { $"{prefix}s1.*" }, cts.Token); + await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token); + + const int total = 100; + for (var i = 0; i < total; i++) + { + var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + } + + // MaxMsgs above the published count so a single pull buffers all messages. + var consumerOpts = new NatsJSConsumeOpts { MaxMsgs = 100, DrainOnCancel = true }; + var consumer = (NatsJSConsumer)await js.GetConsumerAsync($"{prefix}s1", $"{prefix}c1", cts.Token); + + var received = 0; + await foreach (var msg in consumer.ConsumeAsync(serializer: TestDataJsonSerializer.Default, consumerOpts, cancellationToken: cts.Token)) + { + // Simulate some processing time. + await Task.Delay(10, CancellationToken.None); + + // ACK after cancellation must still succeed since the connection stays open. + await msg.AckAsync(cancellationToken: CancellationToken.None); + received++; + + // Cancel once consumption is under way; drain must still deliver the + // remaining buffered/in-flight messages instead of dropping them. + if (received == 1) +#if NET8_0_OR_GREATER + await cts.CancelAsync(); +#else + cts.Cancel(); +#endif + } + + Assert.Equal(total, received); + + // Connection is left usable after the drain completes. + Assert.Equal(NatsConnectionState.Open, nats.ConnectionState); + await nats.PingAsync(CancellationToken.None); + var ack2 = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = total }, serializer: TestDataJsonSerializer.Default, cancellationToken: CancellationToken.None); + ack2.EnsureSuccess(); + } + + [Fact] + public async Task Consume_without_drain_on_cancel_drops_buffered_messages() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + await nats.ConnectRetryAsync(); + var prefix = _server.GetNextId(); + + var js = new NatsJSContext(nats); + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + + await js.CreateStreamAsync($"{prefix}s1", new[] { $"{prefix}s1.*" }, cts.Token); + await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token); + + const int total = 100; + for (var i = 0; i < total; i++) + { + var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + } + + // DrainOnCancel defaults to false: cancelling stops promptly and does not + // deliver the messages still buffered in the consumer. + var consumerOpts = new NatsJSConsumeOpts { MaxMsgs = 100 }; + var consumer = (NatsJSConsumer)await js.GetConsumerAsync($"{prefix}s1", $"{prefix}c1", cts.Token); + + var received = 0; + await foreach (var msg in consumer.ConsumeAsync(serializer: TestDataJsonSerializer.Default, consumerOpts, cancellationToken: cts.Token)) + { + await msg.AckAsync(cancellationToken: CancellationToken.None); + received++; + + if (received == 1) +#if NET8_0_OR_GREATER + await cts.CancelAsync(); +#else + cts.Cancel(); +#endif + } + + // Without drain the buffered messages are abandoned when the token is cancelled. + Assert.True(received < total, $"expected fewer than {total} messages, got {received}"); + } +} From b3f35e6dff566a2c0d99bee5665c4a82ed9456c5 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 10 Jun 2026 11:14:41 +0100 Subject: [PATCH 04/14] test: tidy consume drain tests 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. --- .../NATS.Client.JetStream.Tests/ConsumeDrainTest.cs | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs b/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs index ec3e92714..182fa5633 100644 --- a/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs +++ b/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs @@ -6,14 +6,9 @@ namespace NATS.Client.JetStream.Tests; [Collection("nats-server")] public class ConsumeDrainTest { - private readonly ITestOutputHelper _output; private readonly NatsServerFixture _server; - public ConsumeDrainTest(ITestOutputHelper output, NatsServerFixture server) - { - _output = output; - _server = server; - } + public ConsumeDrainTest(NatsServerFixture server) => _server = server; [Fact] public async Task Consume_drain_on_cancel_delivers_buffered_and_keeps_connection() @@ -109,7 +104,8 @@ public async Task Consume_without_drain_on_cancel_drops_buffered_messages() #endif } - // Without drain the buffered messages are abandoned when the token is cancelled. - Assert.True(received < total, $"expected fewer than {total} messages, got {received}"); + // Without drain the loop stops promptly on cancel (fired after the first + // message) and abandons the buffered messages rather than delivering them. + Assert.True(received <= 5, $"expected a prompt stop (<= 5 messages), got {received} of {total}"); } } From f1a4d15d65d12c48fd99ddf326834a05a0e80b0f Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 10 Jun 2026 11:14:48 +0100 Subject: [PATCH 05/14] test: run xunit tests single-threaded 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. --- tests/xunit.runsettings | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/xunit.runsettings b/tests/xunit.runsettings index 1bc835271..6c4ab22e3 100644 --- a/tests/xunit.runsettings +++ b/tests/xunit.runsettings @@ -4,6 +4,6 @@ 600000 - 1.0x + 1 From 6fa64f3ed6c23a30c090aa21af2503eb440ecd7b Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Wed, 10 Jun 2026 13:52:17 +0100 Subject: [PATCH 06/14] test: fix Reconnect_counter racing the metric increment 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). --- .../OpenTelemetryTest.cs | 52 ++++++++++++++++--- 1 file changed, 45 insertions(+), 7 deletions(-) diff --git a/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs b/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs index 1046fdc9c..d38c29145 100644 --- a/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs +++ b/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs @@ -295,9 +295,22 @@ public async Task Reconnect_counter() await nats.ReconnectAsync(); await openedTwice.Task.WaitAsync(TimeSpan.FromSeconds(30)); - var reconnects = meter.LongMeasurements - .Where(m => m.Name == "nats.client.reconnects") - .ToList(); + // The reconnects counter is incremented just after the ConnectionOpened event is + // queued, so the measurement can land slightly after the event we awaited. Poll + // for it rather than reading once and racing the increment. + List<(string Name, long Value, KeyValuePair[] Tags)> reconnects = new(); + var deadline = DateTime.UtcNow + TimeSpan.FromSeconds(10); + while (DateTime.UtcNow < deadline) + { + reconnects = meter.LongMeasurements + .Where(m => m.Name == "nats.client.reconnects") + .ToList(); + + if (reconnects.Count > 0) + break; + + await Task.Delay(25); + } reconnects.Sum(m => m.Value).Should().Be(1); @@ -536,6 +549,9 @@ public void AssertAllStopped() private sealed class MeterTracker : IDisposable { private readonly MeterListener _listener; + private readonly object _sync = new(); + private readonly List<(string Name, long Value, KeyValuePair[] Tags)> _longMeasurements = new(); + private readonly List<(string Name, double Value, KeyValuePair[] Tags)> _doubleMeasurements = new(); public MeterTracker() { @@ -548,17 +564,39 @@ public MeterTracker() }, }; + // Measurement callbacks fire on background threads (e.g. the reconnect loop), + // so guard the lists and hand out snapshots to readers. _listener.SetMeasurementEventCallback((inst, val, tags, _) => - LongMeasurements.Add((inst.Name, val, tags.ToArray()))); + { + lock (_sync) + _longMeasurements.Add((inst.Name, val, tags.ToArray())); + }); _listener.SetMeasurementEventCallback((inst, val, tags, _) => - DoubleMeasurements.Add((inst.Name, val, tags.ToArray()))); + { + lock (_sync) + _doubleMeasurements.Add((inst.Name, val, tags.ToArray())); + }); _listener.Start(); } - public List<(string Name, long Value, KeyValuePair[] Tags)> LongMeasurements { get; } = new(); + public IReadOnlyList<(string Name, long Value, KeyValuePair[] Tags)> LongMeasurements + { + get + { + lock (_sync) + return _longMeasurements.ToArray(); + } + } - public List<(string Name, double Value, KeyValuePair[] Tags)> DoubleMeasurements { get; } = new(); + public IReadOnlyList<(string Name, double Value, KeyValuePair[] Tags)> DoubleMeasurements + { + get + { + lock (_sync) + return _doubleMeasurements.ToArray(); + } + } public void Dispose() => _listener.Dispose(); } From d0936a0d582da6b78a8e66ddb0a8a147a3517894 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 10:35:14 +0100 Subject: [PATCH 07/14] core: fix subscription drain hang and message loss 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. --- src/NATS.Client.Core/NatsSubBase.cs | 61 ++++++++++++------- .../Internal/NatsJSConsume.cs | 25 +++++++- 2 files changed, 62 insertions(+), 24 deletions(-) diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index 9ed623609..c435f236a 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -636,6 +636,15 @@ private void DecrementActiveSubscription() /// private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) { + // Disarm the lifecycle timers before taking the unsubscribe gate so a + // Timeout/IdleTimeout/StartUp callback that is about to fire is less likely to win + // the gate and complete the channel without the drain fence. The gate below still + // guarantees correctness if a callback was already in flight; this just narrows the + // window where an explicit drain races those timers. + _timeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + _idleTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + _startUpTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + var needsUnsub = false; lock (_gate) { @@ -651,26 +660,22 @@ private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) DecrementActiveSubscription(); - _timeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - _idleTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - _startUpTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - - // For the built-in SubscriptionManager: send UNSUB but keep the SID in the routing - // table until after the PING/PONG round-trip, so messages already in the socket - // buffer are delivered before the channel is completed. - // For external managers: fall back to RemoveAsync (routing removed before PING). var mgr = _manager as SubscriptionManager; - var sid = mgr != null - ? await mgr.UnsubscribeForDrainAsync(this).ConfigureAwait(false) - : -1; - if (mgr == null) - await _manager.RemoveAsync(this).ConfigureAwait(false); - - // PING/PONG round-trip: after we get the PONG back, we know the server - // has processed our UNSUB and the socket reader has processed all messages - // that were in-flight before the UNSUB was received by the server. + var sid = -1; try { + // For the built-in SubscriptionManager: send UNSUB but keep the SID in the routing + // table until after the PING/PONG round-trip, so messages already in the socket + // buffer are delivered before the channel is completed. + // For external managers: fall back to RemoveAsync (routing removed before PING). + if (mgr != null) + sid = await mgr.UnsubscribeForDrainAsync(this).ConfigureAwait(false); + else + await _manager.RemoveAsync(this).ConfigureAwait(false); + + // PING/PONG round-trip: after we get the PONG back, we know the server + // has processed our UNSUB and the socket reader has processed all messages + // that were in-flight before the UNSUB was received by the server. using var pingCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); pingCts.CancelAfter(Connection.Opts.DrainPingTimeout); await Connection.PingAsync(pingCts.Token).ConfigureAwait(false); @@ -683,13 +688,23 @@ private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) { // Connection failed or disconnected during drain. Drain is best-effort. } + catch (ObjectDisposedException) + { + // The connection or its command writer was disposed mid-drain (e.g. a connection + // dispose racing an explicit drain). Drain is best-effort; fall through and still + // complete the channel below. + } + finally + { + // Remove from routing now that no more messages can arrive for this SID, then + // complete the channel. This runs even if UNSUB or PING threw, so a reader + // blocked on the channel (e.g. the drain-on-cancel consume loop reading with + // CancellationToken.None) is never left waiting forever. + if (sid != -1) + mgr?.RemoveFromRouting(sid); - // Remove from routing now that no more messages can arrive for this SID. - if (sid != -1) - mgr?.RemoveFromRouting(sid); - - // Now it's safe to complete the channel, no more messages will arrive. - TryComplete(); + TryComplete(); + } } private async ValueTask WaitForReaderDrainCoreAsync(Task task, TimeSpan timeout) diff --git a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs index 5faa169b2..ff846f24d 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs @@ -50,6 +50,7 @@ internal class NatsJSConsume : NatsSubBase private long _pendingMsgs; private long _pendingBytes; private int _consecutive503Errors; + private volatile bool _draining; public NatsJSConsume( long maxMsgs, @@ -114,6 +115,13 @@ public NatsJSConsume( static state => { var self = (NatsJSConsume)state!; + + // A drain is completing the subscription behind a PING/PONG fence. Don't let + // the heartbeat callback complete the channel (CompleteStop) and skip the + // fence, which would drop in-flight messages the drain is meant to preserve. + if (self._draining) + return; + self._notificationChannel?.Notify(NatsJSTimeoutNotification.Default); if (self._cancellationToken.IsCancellationRequested) @@ -198,13 +206,22 @@ public ValueTask CallMsgNextAsync(string origin, ConsumerGetnextRequest request, public override ValueTask DrainAsync(CancellationToken cancellationToken = default) { + _draining = true; StopHeartbeatTimer(); return base.DrainAsync(cancellationToken); } public void StopHeartbeatTimer() => _timer.Change(Timeout.Infinite, Timeout.Infinite); - public void ResetHeartbeatTimer() => _timer.Change(_hbTimeout, _hbTimeout); + public void ResetHeartbeatTimer() + { + // Once draining, the heartbeat timer stays stopped so it can't re-arm and complete + // the channel ahead of the drain fence. + if (_draining) + return; + + _timer.Change(_hbTimeout, _hbTimeout); + } public void Delivered(int msgSize) { @@ -535,6 +552,12 @@ private void CompleteStop() state => { var self = (NatsJSConsume)state!; + + // If a drain started after this stop was queued, leave completion to the + // drain fence rather than dropping in-flight messages here. + if (self._draining) + return; + self._userMsgs.Writer.TryComplete(); self.EndSubscription(NatsSubEndReason.None); }, From 0c99a606ed31a2274650f4d4c8176bc0c3da1682 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 11:09:04 +0100 Subject: [PATCH 08/14] jetstream: drain consumer delivery via overridable hook 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. --- src/NATS.Client.Core/NatsSubBase.cs | 17 +++++++++++++++++ .../Internal/NatsJSConsume.cs | 16 +++++++++------- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index c435f236a..e4e8cf36e 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -514,6 +514,18 @@ internal void MarkReaderInactive() /// protected abstract ValueTask ReceiveInternalAsync(string subject, string? replyTo, ReadOnlySequence? headersBuffer, ReadOnlySequence payloadBuffer); + /// + /// Stops any subclass-owned delivery machinery (e.g. a JetStream pull loop or + /// idle-heartbeat timer) as the first step of a drain, before the UNSUB and the + /// PING/PONG fence. The base implementation is a no-op; subclasses that keep + /// pulling or run a timer that can complete the message channel must override this + /// so drain does not leave that machinery running or let it skip the fence. + /// Called on both the explicit drain and the dispose-drain path; must be idempotent. + /// + protected virtual void StopDelivery() + { + } + /// /// Waits for the user-side reader to exit so buffered messages can be /// drained on dispose. No-op when the reader was never marked active or @@ -636,6 +648,11 @@ private void DecrementActiveSubscription() /// private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) { + // Stop any subclass-owned delivery engine (e.g. a JetStream pull/heartbeat loop) + // first so it can't keep pulling or complete the channel ahead of the fence below. + // Runs on both the explicit drain and the dispose-drain path. + StopDelivery(); + // Disarm the lifecycle timers before taking the unsubscribe gate so a // Timeout/IdleTimeout/StartUp callback that is about to fire is less likely to win // the gate and complete the channel without the drain fence. The gate below still diff --git a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs index ff846f24d..ed812c721 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs @@ -204,13 +204,6 @@ public ValueTask CallMsgNextAsync(string origin, ConsumerGetnextRequest request, cancellationToken: cancellationToken); } - public override ValueTask DrainAsync(CancellationToken cancellationToken = default) - { - _draining = true; - StopHeartbeatTimer(); - return base.DrainAsync(cancellationToken); - } - public void StopHeartbeatTimer() => _timer.Change(Timeout.Infinite, Timeout.Infinite); public void ResetHeartbeatTimer() @@ -322,6 +315,15 @@ await commandWriter.PublishAsync( } } + protected override void StopDelivery() + { + // Mark draining first so the heartbeat callback, its re-arm, and CompleteStop + // defer to the drain fence, then stop the timer so it can't pull or complete the + // channel during the drain. + _draining = true; + StopHeartbeatTimer(); + } + protected override async ValueTask ReceiveInternalAsync( string subject, string? replyTo, From ed2d86b856306b36e3697d03f4dd3a8caddc4925 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 11:10:59 +0100 Subject: [PATCH 09/14] test: cover drain in-flight message fence 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. --- .../SubscriptionDrainTest.cs | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs index 73966b2d3..0b6f1e2a6 100644 --- a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs +++ b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs @@ -40,6 +40,35 @@ public async Task Drain_preserves_buffered_messages_and_completes_channel() Assert.Equal(Enumerable.Range(0, count), received); } + [Fact] + public async Task Drain_delivers_messages_still_in_flight_after_unsub() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + // Kept below the default SubPendingChannelCapacity so nothing is dropped for being + // over capacity; what is under test is the drain fence, not back-pressure. + const int count = 1000; + for (var i = 0; i < count; i++) + await nats.PublishAsync(subject, i, cancellationToken: cancellationToken); + + // Unlike the test above there is no PING/PONG before draining, so messages the + // server sent before it processed our UNSUB are still in flight (not yet in the + // channel). The drain fences inbound delivery with its own PING/PONG so they all + // land before the channel completes; without that fence the tail would be dropped. + await sub.DrainAsync(cancellationToken); + + var received = new List(); + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken)) + received.Add(msg.Data); + + Assert.Equal(Enumerable.Range(0, count), received); + } + [Fact] public async Task Drain_stops_new_messages_but_connection_stays_usable() { From f20b9bea3bfc2e8544003925f305c73971533dd0 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 11:11:24 +0100 Subject: [PATCH 10/14] jetstream: scope DrainOnCancel docs to the pull consumer --- src/NATS.Client.JetStream/NatsJSOpts.cs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/NATS.Client.JetStream/NatsJSOpts.cs b/src/NATS.Client.JetStream/NatsJSOpts.cs index 6a78b2dbb..ef7cf614e 100644 --- a/src/NATS.Client.JetStream/NatsJSOpts.cs +++ b/src/NATS.Client.JetStream/NatsJSOpts.cs @@ -167,6 +167,11 @@ public record NatsJSConsumeOpts /// then complete) and leaves the connection usable, instead of stopping /// immediately. (default: false) /// + /// + /// Only the standard pull consumer's ConsumeAsync honors this. Ordered consumers + /// (see ) do not ACK and always stop promptly + /// on cancel, so the option has no effect there. + /// public bool DrainOnCancel { get; init; } } From 533ced0b45be4c9b034430a5a0e65a66f09f571c Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 11:25:00 +0100 Subject: [PATCH 11/14] core: clarify DrainAsync fence timeout docs --- src/NATS.Client.Core/INatsSub.cs | 2 +- src/NATS.Client.Core/NatsSubBase.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/NATS.Client.Core/INatsSub.cs b/src/NATS.Client.Core/INatsSub.cs index 9ce8940f2..b4e939573 100644 --- a/src/NATS.Client.Core/INatsSub.cs +++ b/src/NATS.Client.Core/INatsSub.cs @@ -39,7 +39,7 @@ public interface INatsSub : IAsyncDisposable /// , drain does not drop messages still in the socket /// buffer. The connection stays open. /// - /// Bounds the best-effort PING/PONG fence wait. + /// Bounds the best-effort PING/PONG fence wait, which lasts at most ; the token can only shorten it, not extend it past that timeout. /// A that represents the asynchronous drain operation. public ValueTask DrainAsync(CancellationToken cancellationToken = default); } diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index e4e8cf36e..b9d863fe1 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -297,7 +297,7 @@ public ValueTask UnsubscribeAsync() /// keep the subscription alive. /// /// - /// Bounds the best-effort PING/PONG fence wait. Cancellation does not abort the drain. + /// Bounds the best-effort PING/PONG fence wait, which lasts at most ; the token can only shorten it, not extend it past that timeout. Cancellation does not abort the drain. /// A that represents the asynchronous drain operation. public virtual ValueTask DrainAsync(CancellationToken cancellationToken = default) => DrainCoreAsync(cancellationToken); From 01b7eda99f8f4598d74c9e72666d958fae29ed03 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 11:42:01 +0100 Subject: [PATCH 12/14] jetstream: remove redundant drain-start in consume loop 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. --- src/NATS.Client.JetStream/NatsJSConsumer.cs | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/NATS.Client.JetStream/NatsJSConsumer.cs b/src/NATS.Client.JetStream/NatsJSConsumer.cs index 39a825758..bc08edf43 100644 --- a/src/NATS.Client.JetStream/NatsJSConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSConsumer.cs @@ -80,16 +80,10 @@ public async IAsyncEnumerable> ConsumeAsync( { while (true) { - if (cancellationToken.IsCancellationRequested && !draining) - { - if (!drainOnCancel) - yield break; - - draining = true; - drainTask = cc.DrainAsync(); - } - // We have to check calls individually since we can't use yield return in try-catch blocks. + // WaitToReadAsync throws if the token is already canceled (even with buffered + // messages), so the catch below is the single place that reacts to cancellation: + // start the drain when DrainOnCancel is set, otherwise stop. bool ready; try { From 1b258d3a463f9d234386e77b0fadfdf19b879644 Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 12:44:53 +0100 Subject: [PATCH 13/14] core: complete drain channel even if teardown throws 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. --- src/NATS.Client.Core/NatsSubBase.cs | 30 +++++++++++++---------------- 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index b9d863fe1..b8b655a13 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -648,20 +648,6 @@ private void DecrementActiveSubscription() /// private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) { - // Stop any subclass-owned delivery engine (e.g. a JetStream pull/heartbeat loop) - // first so it can't keep pulling or complete the channel ahead of the fence below. - // Runs on both the explicit drain and the dispose-drain path. - StopDelivery(); - - // Disarm the lifecycle timers before taking the unsubscribe gate so a - // Timeout/IdleTimeout/StartUp callback that is about to fire is less likely to win - // the gate and complete the channel without the drain fence. The gate below still - // guarantees correctness if a callback was already in flight; this just narrows the - // window where an explicit drain races those timers. - _timeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - _idleTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - _startUpTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); - var needsUnsub = false; lock (_gate) { @@ -675,12 +661,22 @@ private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) if (!needsUnsub) return; - DecrementActiveSubscription(); - var mgr = _manager as SubscriptionManager; var sid = -1; try { + DecrementActiveSubscription(); + + // Stop the subclass-owned delivery engine (e.g. a JetStream pull/heartbeat loop) + // and disarm the lifecycle timers before the fence so they can't keep pulling or + // complete the channel ahead of it. These run inside the try so that if any of + // them throw (e.g. a timer disposed by a racing DisposeAsync) the finally below + // still completes the channel and no reader is left waiting forever. + StopDelivery(); + _timeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + _idleTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + _startUpTimeoutTimer?.Change(Timeout.Infinite, Timeout.Infinite); + // For the built-in SubscriptionManager: send UNSUB but keep the SID in the routing // table until after the PING/PONG round-trip, so messages already in the socket // buffer are delivered before the channel is completed. @@ -714,7 +710,7 @@ private async ValueTask DrainCoreAsync(CancellationToken cancellationToken) finally { // Remove from routing now that no more messages can arrive for this SID, then - // complete the channel. This runs even if UNSUB or PING threw, so a reader + // complete the channel. This runs even if anything in the try threw, so a reader // blocked on the channel (e.g. the drain-on-cancel consume loop reading with // CancellationToken.None) is never left waiting forever. if (sid != -1) From 1525e70ea84ea3647b66a490e28d97d8b9a5bfdb Mon Sep 17 00:00:00 2001 From: Ziya Suzen Date: Thu, 11 Jun 2026 12:45:07 +0100 Subject: [PATCH 14/14] test: assert drain message count before order --- tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs index 0b6f1e2a6..56b8d63d0 100644 --- a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs +++ b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs @@ -66,6 +66,9 @@ public async Task Drain_delivers_messages_still_in_flight_after_unsub() await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken)) received.Add(msg.Data); + // Check the count first so a dropped message surfaces as a count mismatch rather + // than a confusing order-equality failure. + Assert.Equal(count, received.Count); Assert.Equal(Enumerable.Range(0, count), received); }