diff --git a/src/NATS.Client.Core/Internal/SubscriptionManager.cs b/src/NATS.Client.Core/Internal/SubscriptionManager.cs index ce0ee5bcb..928d80acc 100644 --- a/src/NATS.Client.Core/Internal/SubscriptionManager.cs +++ b/src/NATS.Client.Core/Internal/SubscriptionManager.cs @@ -149,18 +149,23 @@ public async ValueTask DisposeAsync() _cts.Cancel(); #endif - WeakReference[] subRefs; + var subs = new List(); lock (_gate) { - subRefs = _bySid.Values.Select(m => m.WeakReference).ToArray(); - _bySid.Clear(); + foreach (var sidMetadata in _bySid.Values) + { + if (sidMetadata.WeakReference.TryGetTarget(out var sub)) + subs.Add(sub); + } } - foreach (var subRef in subRefs) - { - if (subRef.TryGetTarget(out var sub)) - await sub.DisposeAsync().ConfigureAwait(false); - } + if (subs.Count == 0) + return; + + var tasks = new Task[subs.Count]; + for (var i = 0; i < subs.Count; i++) + tasks[i] = subs[i].DisposeAsync().AsTask(); + await Task.WhenAll(tasks).ConfigureAwait(false); } public ValueTask RemoveAsync(NatsSubBase sub) @@ -188,6 +193,34 @@ public ValueTask RemoveAsync(NatsSubBase sub) return _connection.UnsubscribeAsync(subMetadata.Sid); } + // Send UNSUB and remove from _bySub, but keep in _bySid until after PING/PONG. + // Returns the SID so the caller can call RemoveFromRouting after the round-trip. + internal async ValueTask UnsubscribeForDrainAsync(NatsSubBase sub) + { + SubscriptionMetadata? subMetadata; + lock (_gate) + { + if (!_bySub.TryGetValue(sub, out subMetadata)) + return -1; + _bySub.Remove(sub); + } + + if (_debug) + _logger.LogDebug(NatsLogEvents.Subscription, "Drain-unsubscribing {Subject}/{Sid}", sub.Subject, subMetadata.Sid); + + await _connection.UnsubscribeAsync(subMetadata.Sid).ConfigureAwait(false); + return subMetadata.Sid; + } + + // Remove from _bySid after the PING/PONG round-trip completes. + internal void RemoveFromRouting(int sid) + { + lock (_gate) + { + _bySid.TryRemove(sid, out _); + } + } + /// /// Returns commands for all the live subscriptions to be used on reconnect so that they can rebuild their connection state on the server. /// diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index b814a3676..ae2aaad2e 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/src/NATS.Client.Core/NatsConnection.cs @@ -48,6 +48,8 @@ public partial class NatsConnection : INatsConnection private readonly ClientOpts _clientOpts; private readonly SubscriptionManager _subscriptionManager; private readonly ReplyTaskFactory _replyTaskFactory; + private readonly object _drainParticipantsGate = new(); + private readonly HashSet _drainParticipants = new(); private ServerInfo? _writableServerInfo; private int _pongCount; @@ -280,7 +282,6 @@ public virtual async ValueTask DisposeAsync() IsDisposed = true; _logger.Log(LogLevel.Information, NatsLogEvents.Connection, "Disposing connection {Name}", _name); - await DisposeSocketAsync(false).ConfigureAwait(false); if (_pingTimerCancellationTokenSource != null) { #if NET8_0_OR_GREATER @@ -290,8 +291,28 @@ public virtual async ValueTask DisposeAsync() #endif } - await _subscriptionManager.DisposeAsync().ConfigureAwait(false); - await CommandWriter.DisposeAsync().ConfigureAwait(false); + if (Opts.DrainSubscriptionsOnDispose) + { + // Drain subs and flush the writer first so UNSUB/PING/PONG + // and any pending acks can land before the socket closes. + try + { + await _subscriptionManager.DisposeAsync().ConfigureAwait(false); + } + finally + { + await DrainRemainingParticipantsAsync().ConfigureAwait(false); + await CommandWriter.DisposeAsync().ConfigureAwait(false); + await DisposeSocketAsync(false).ConfigureAwait(false); + } + } + else + { + await DisposeSocketAsync(false).ConfigureAwait(false); + await _subscriptionManager.DisposeAsync().ConfigureAwait(false); + await CommandWriter.DisposeAsync().ConfigureAwait(false); + } + _waitForOpenConnection.TrySetCanceled(); #if NET8_0_OR_GREATER await _disposedCts.CancelAsync().ConfigureAwait(false); @@ -361,6 +382,18 @@ internal void ResetPongCount() // called only internally internal ValueTask SubscribeCoreAsync(int sid, string subject, string? queueGroup, int? maxMsgs, CancellationToken cancellationToken) => CommandWriter.SubscribeAsync(sid, subject, queueGroup, maxMsgs, cancellationToken); + internal void RegisterDrainParticipant(NatsSubBase sub) + { + lock (_drainParticipantsGate) + _drainParticipants.Add(sub); + } + + internal void UnregisterDrainParticipant(NatsSubBase sub) + { + lock (_drainParticipantsGate) + _drainParticipants.Remove(sub); + } + internal ValueTask UnsubscribeAsync(int sid) { try @@ -1140,4 +1173,26 @@ private void ThrowIfDisposed() if (IsDisposed) throw new ObjectDisposedException(null); } + + // Drain subs that already removed themselves from the SubscriptionManager + // (e.g. JetStream Fetch on natural completion) but still have a user + // iterator reading the channel. Their own DrainAsync waits on the + // reader-exited signal. Run in parallel so a slow/dead server doesn't + // multiply the dispose wait by the number of subs. + private ValueTask DrainRemainingParticipantsAsync() + { + NatsSubBase[] remaining; + lock (_drainParticipantsGate) + { + if (_drainParticipants.Count == 0) + return default; + remaining = _drainParticipants.ToArray(); + } + + var tasks = new Task[remaining.Length]; + for (var i = 0; i < remaining.Length; i++) + tasks[i] = remaining[i].DrainAsync().AsTask(); + + return new ValueTask(Task.WhenAll(tasks)); + } } diff --git a/src/NATS.Client.Core/NatsOpts.cs b/src/NATS.Client.Core/NatsOpts.cs index 9f26b5ec4..d5c59e9f7 100644 --- a/src/NATS.Client.Core/NatsOpts.cs +++ b/src/NATS.Client.Core/NatsOpts.cs @@ -153,6 +153,55 @@ public sealed record NatsOpts public bool WaitUntilSent { get; init; } = false; + /// + /// When true, drains subscriptions + /// and flushes the command writer before closing the socket. (default: false) + /// + /// + /// + /// With the default (false), the socket is closed first, which means UNSUB + /// and any pending writes (e.g. JetStream acks queued by a consume loop) + /// are dropped. Enable to let subscriptions drain cleanly on dispose. + /// Required for to be effective. + /// + /// + /// Drain mechanism: each subscription sends UNSUB, then a PING and waits + /// for the PONG (bounded by ). Once the + /// PONG arrives the server has processed UNSUB and the socket reader has + /// delivered any messages the server sent before that, so completing the + /// user channel afterwards won't drop in-flight messages. Subscriptions + /// drain in parallel on connection dispose, so total dispose time is + /// bounded by one round-trip plus the + /// wait if set, not multiplied by the number of subscriptions. + /// + /// + public bool DrainSubscriptionsOnDispose { get; init; } = false; + + /// + /// When set, JetStream consumer/fetch dispose waits up to this timeout for + /// the user's consume loop to finish acking buffered messages before + /// returning. (default: null, no wait) + /// + /// + /// Only effective when is also + /// true. Without this wait, messages already delivered into the consumer + /// channel but not yet yielded by the user loop are dropped on dispose + /// and remain NumAckPending until AckWait expires. + /// + public TimeSpan? ConsumerDrainOnDisposeTimeout { get; init; } = null; + + /// + /// Per-subscription PING/PONG timeout used during drain on dispose. + /// (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. + /// + public TimeSpan DrainPingTimeout { get; init; } = TimeSpan.FromSeconds(5); + /// /// Maximum number of reconnect attempts. (default: -1, unlimited) /// diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index 72f3ec8c4..afa3feb70 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -50,6 +50,7 @@ public abstract class NatsSubBase private int _pendingMsgs; private Exception? _exception; private int _isSlowConsumer; + private TaskCompletionSource? _readerExited; /// /// Creates a new instance of . @@ -358,6 +359,82 @@ 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. + /// + /// + /// No-op unless is + /// enabled. The connection-level dispose path also gates entry into + /// this method on the same flag (see ordering in + /// ). + /// + /// A that represents the asynchronous drain operation. + internal async ValueTask DrainAsync() + { + if (!Connection.Opts.DrainSubscriptionsOnDispose) + return; + + var needsUnsub = false; + lock (_gate) + { + if (!_unsubscribed) + { + _unsubscribed = true; + needsUnsub = true; + } + } + + if (needsUnsub) + { + _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(); + } + + // Always wait for an active user reader to drain the channel before + // returning, even if the subscription was already ended (e.g. via + // EndSubscription / UnsubscribeAsync) before drain was called. + await WaitForReaderDrainAsync().ConfigureAwait(false); + } + /// /// Marks this subscription as a slow consumer. Returns true if this was a state transition /// (i.e., the subscription was not previously marked as a slow consumer). @@ -377,6 +454,36 @@ internal static bool IsHeader503(ReadOnlySequence? headersBuffer) => /// ValueTask internal virtual ValueTask WriteReconnectCommandsAsync(CommandWriter commandWriter, int sid) => commandWriter.SubscribeAsync(sid, Subject, QueueGroup, PendingMsgs, CancellationToken.None); + /// + /// Marks the user-side reader as active so dispose can wait for it to + /// drain buffered messages before completing. Pairs with + /// . + /// + internal void MarkReaderActive() + { + if (!Connection.Opts.DrainSubscriptionsOnDispose) + return; + + if (_readerExited is null) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + Interlocked.CompareExchange(ref _readerExited, tcs, null); + } + + if (Connection is NatsConnection nc) + nc.RegisterDrainParticipant(this); + } + + /// + /// Signals the user-side reader has exited. Pairs with . + /// + internal void MarkReaderInactive() + { + if (Connection is NatsConnection nc) + nc.UnregisterDrainParticipant(this); + _readerExited?.TrySetResult(); + } + /// /// Invoked when a MSG or HMSG arrives for the subscription. /// @@ -390,6 +497,21 @@ internal static bool IsHeader503(ReadOnlySequence? headersBuffer) => /// protected abstract ValueTask ReceiveInternalAsync(string subject, string? replyTo, ReadOnlySequence? headersBuffer, ReadOnlySequence payloadBuffer); + /// + /// 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 + /// when is unset. + /// + protected virtual ValueTask WaitForReaderDrainAsync() + { + var tcs = _readerExited; + if (tcs is null || tcs.Task.IsCompleted) + return default; + if (Connection.Opts.ConsumerDrainOnDisposeTimeout is not { } timeout) + return default; + return WaitForReaderDrainCoreAsync(tcs.Task, timeout); + } + /// /// Resets the slow consumer state if the channel has drained, allowing another /// slow consumer event to be raised if the subscription becomes slow again. @@ -481,4 +603,17 @@ protected void EndSubscription(NatsSubEndReason reason) #pragma warning restore VSTHRD110 #pragma warning restore CA2012 } + + private async ValueTask WaitForReaderDrainCoreAsync(Task task, TimeSpan timeout) + { + try + { + using var cts = new CancellationTokenSource(timeout); + await task.WaitAsync(cts.Token).ConfigureAwait(false); + } + catch (OperationCanceledException) + { + _logger.LogWarning(NatsLogEvents.Subscription, "Timeout waiting for subscription reader to exit on dispose"); + } + } } diff --git a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs index 273290325..09e295c36 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs @@ -45,7 +45,6 @@ internal class NatsJSConsume : NatsSubBase private readonly object _pendingGate = new(); private long _pendingMsgs; private long _pendingBytes; - private int _disposed; private int _consecutive503Errors; public NatsJSConsume( @@ -217,9 +216,9 @@ public void Delivered(int msgSize) public override async ValueTask DisposeAsync() { - Interlocked.Exchange(ref _disposed, 1); try { + await DrainAsync().ConfigureAwait(false); await base.DisposeAsync().ConfigureAwait(false); } finally @@ -458,17 +457,12 @@ protected override async ValueTask ReceiveInternalAsync( NatsJSExtensionsInternal.TrySetPinIdFromHeaders(msg.Headers, _jsConsumer); - // Stop feeding the user if we are disposed. - // We need to exit as soon as possible. - if (Volatile.Read(ref _disposed) == 0) - { - // We can't pass cancellation token here because we need to hand - // the message to the user to be processed. Writer will be completed - // when the user calls Stop() or when the subscription is closed. - await _userMsgs.Writer.WriteAsync(msg).ConfigureAwait(false); + // We can't pass cancellation token here because we need to hand + // the message to the user to be processed. Writer will be completed + // when the user calls Stop() or when the subscription is closed. + await _userMsgs.Writer.WriteAsync(msg).ConfigureAwait(false); - ResetSlowConsumer(_userMsgs.Reader.Count); - } + ResetSlowConsumer(_userMsgs.Reader.Count); } } diff --git a/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs b/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs index f14dba6fc..5dcf6d429 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs @@ -32,7 +32,6 @@ internal class NatsJSFetch : NatsSubBase private long _pendingMsgs; private long _pendingBytes; - private int _disposed; public NatsJSFetch( long maxMsgs, @@ -162,9 +161,9 @@ public void ResetHeartbeatTimer() public override async ValueTask DisposeAsync() { - Interlocked.Exchange(ref _disposed, 1); try { + await DrainAsync().ConfigureAwait(false); await base.DisposeAsync().ConfigureAwait(false); } finally @@ -294,14 +293,9 @@ protected override async ValueTask ReceiveInternalAsync( _pendingMsgs--; _pendingBytes -= msg.Size; - // Stop feeding the user if we are disposed. - // We need to exit as soon as possible. - if (Volatile.Read(ref _disposed) == 0) - { - await _userMsgs.Writer.WriteAsync(msg).ConfigureAwait(false); + await _userMsgs.Writer.WriteAsync(msg).ConfigureAwait(false); - ResetSlowConsumer(_userMsgs.Reader.Count); - } + ResetSlowConsumer(_userMsgs.Reader.Count); } if (_maxBytes > 0 && _pendingBytes <= 0) diff --git a/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs index d327b9e3e..900275d85 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs @@ -34,7 +34,6 @@ internal class NatsJSOrderedConsume : NatsSubBase private readonly object _pendingGate = new(); private long _pendingMsgs; private long _pendingBytes; - private int _disposed; public NatsJSOrderedConsume( long maxMsgs, @@ -134,11 +133,10 @@ public ValueTask CallMsgNextAsync(string origin, ConsumerGetnextRequest request, public override async ValueTask DisposeAsync() { - Interlocked.Exchange(ref _disposed, 1); - _context.Connection.ConnectionDisconnected -= ConnectionOnConnectionDisconnected; try { + await DrainAsync().ConfigureAwait(false); await base.DisposeAsync().ConfigureAwait(false); } finally @@ -306,17 +304,12 @@ protected override async ValueTask ReceiveInternalAsync( } } - // Stop feeding the user if we are disposed. - // We need to exit as soon as possible. - if (Volatile.Read(ref _disposed) == 0) - { - // We can't pass cancellation token here because we need to hand - // the message to the user to be processed. Writer will be completed - // when the user calls Stop() or when the subscription is closed. - await _userMsgs.Writer.WriteAsync(msg).ConfigureAwait(false); + // We can't pass cancellation token here because we need to hand + // the message to the user to be processed. Writer will be completed + // when the user calls Stop() or when the subscription is closed. + await _userMsgs.Writer.WriteAsync(msg).ConfigureAwait(false); - ResetSlowConsumer(_userMsgs.Reader.Count); - } + ResetSlowConsumer(_userMsgs.Reader.Count); } CheckPending(); diff --git a/src/NATS.Client.JetStream/NatsJSConsumer.cs b/src/NATS.Client.JetStream/NatsJSConsumer.cs index 996bfacf2..e10958563 100644 --- a/src/NATS.Client.JetStream/NatsJSConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSConsumer.cs @@ -63,45 +63,20 @@ public async IAsyncEnumerable> ConsumeAsync( { opts ??= _context.Opts.DefaultConsumeOpts; await using var cc = await ConsumeInternalAsync(serializer, opts, cancellationToken).ConfigureAwait(false); - - while (!cancellationToken.IsCancellationRequested) + cc.MarkReaderActive(); + try { - // 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); - } - catch (OperationCanceledException) - { - ready = false; - } - catch (NatsConnectionFailedException) - { - // Connection has permanently failed, stop consuming and rethrow - throw; - } - catch (NatsJSException) - { - // Consumer-related errors (like 503 threshold exceeded), stop consuming and rethrow - throw; - } - - if (!ready) - yield break; - while (!cancellationToken.IsCancellationRequested) { - bool read; - NatsJSMsg jsMsg; + // We have to check calls individually since we can't use yield return in try-catch blocks. + bool ready; try { - read = cc.Msgs.TryRead(out jsMsg); + ready = await cc.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false); } catch (OperationCanceledException) { - read = false; - jsMsg = default; + ready = false; } catch (NatsConnectionFailedException) { @@ -114,17 +89,49 @@ public async IAsyncEnumerable> ConsumeAsync( throw; } - if (!read) - break; + if (!ready) + yield break; - // if yield is blocked by the application code, we don't want - // heartbeat timer kicking in and issuing unnecessary pull requests. - cc.StopHeartbeatTimer(); - yield return jsMsg; - cc.ResetHeartbeatTimer(); - cc.Delivered(jsMsg.Size); + while (!cancellationToken.IsCancellationRequested) + { + bool read; + NatsJSMsg jsMsg; + try + { + read = cc.Msgs.TryRead(out jsMsg); + } + catch (OperationCanceledException) + { + read = false; + jsMsg = default; + } + catch (NatsConnectionFailedException) + { + // Connection has permanently failed, stop consuming and rethrow + throw; + } + catch (NatsJSException) + { + // Consumer-related errors (like 503 threshold exceeded), stop consuming and rethrow + throw; + } + + if (!read) + break; + + // if yield is blocked by the application code, we don't want + // heartbeat timer kicking in and issuing unnecessary pull requests. + cc.StopHeartbeatTimer(); + yield return jsMsg; + cc.ResetHeartbeatTimer(); + cc.Delivered(jsMsg.Size); + } } } + finally + { + cc.MarkReaderInactive(); + } } /// @@ -177,9 +184,17 @@ public async IAsyncEnumerable> ConsumeAsync( serializer, cancellationToken: cancellationToken).ConfigureAwait(false); - await foreach (var natsJSMsg in f.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + f.MarkReaderActive(); + try { - return natsJSMsg; + await foreach (var natsJSMsg in f.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + { + return natsJSMsg; + } + } + finally + { + f.MarkReaderInactive(); } return null; @@ -196,47 +211,54 @@ public async IAsyncEnumerable> FetchAsync( serializer ??= _context.Connection.Opts.SerializerRegistry.GetDeserializer(); await using var fc = await FetchInternalAsync(opts, serializer, cancellationToken).ConfigureAwait(false); - - while (!cancellationToken.IsCancellationRequested) + fc.MarkReaderActive(); + try { - // We have to check calls individually since we can't use yield return in try-catch blocks. - bool ready; - try - { - ready = await fc.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - ready = false; - } - - if (!ready) - yield break; - while (!cancellationToken.IsCancellationRequested) { - bool read; - NatsJSMsg jsMsg; + // We have to check calls individually since we can't use yield return in try-catch blocks. + bool ready; try { - read = fc.Msgs.TryRead(out jsMsg); + ready = await fc.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false); } catch (OperationCanceledException) { - read = false; - jsMsg = default; + ready = false; } - if (!read) - break; + if (!ready) + yield break; - // if yield is blocked by the application code, we don't want - // heartbeat timer kicking in and issuing unnecessary pull requests. - fc.StopHeartbeatTimer(); - yield return jsMsg; - fc.ResetHeartbeatTimer(); + while (!cancellationToken.IsCancellationRequested) + { + bool read; + NatsJSMsg jsMsg; + try + { + read = fc.Msgs.TryRead(out jsMsg); + } + catch (OperationCanceledException) + { + read = false; + jsMsg = default; + } + + if (!read) + break; + + // if yield is blocked by the application code, we don't want + // heartbeat timer kicking in and issuing unnecessary pull requests. + fc.StopHeartbeatTimer(); + yield return jsMsg; + fc.ResetHeartbeatTimer(); + } } } + finally + { + fc.MarkReaderInactive(); + } } /// @@ -294,13 +316,21 @@ public async IAsyncEnumerable> FetchNoWaitAsync( serializer ??= _context.Connection.Opts.SerializerRegistry.GetDeserializer(); await using var fc = await FetchInternalAsync(opts with { NoWait = true }, serializer, cancellationToken).ConfigureAwait(false); - await foreach (var jsMsg in fc.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + fc.MarkReaderActive(); + try + { + await foreach (var jsMsg in fc.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + { + // if yield is blocked by the application code, we don't want + // heartbeat timer kicking in and issuing unnecessary pull requests. + fc.StopHeartbeatTimer(); + yield return jsMsg; + fc.ResetHeartbeatTimer(); + } + } + finally { - // if yield is blocked by the application code, we don't want - // heartbeat timer kicking in and issuing unnecessary pull requests. - fc.StopHeartbeatTimer(); - yield return jsMsg; - fc.ResetHeartbeatTimer(); + fc.MarkReaderInactive(); } } diff --git a/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs b/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs index e8d51e3d5..769853ca9 100644 --- a/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs @@ -76,6 +76,9 @@ public async IAsyncEnumerable> ConsumeAsync( ulong seq = 0; while (!cancellationToken.IsCancellationRequested) { + if (_context.Connection is NatsConnection { IsDisposed: true }) + yield break; + var consumer = await RecreateConsumer(consumerName, seq, cancellationToken); consumerName = consumer.Info.Name; _logger.LogInformation(NatsJSLogEvents.NewConsumer, "Created {ConsumerName} with sequence {Seq}", consumerName, seq); @@ -85,72 +88,80 @@ public async IAsyncEnumerable> ConsumeAsync( await using (var cc = await consumer.OrderedConsumeInternalAsync(serializer, opts, cancellationToken)) { - while (true) + cc.MarkReaderActive(); + try { - // We have to check every call to WaitToReadAsync and TryRead for - // protocol exceptions individually because we can't yield return - // within try-catch. - try - { - var read = await cc.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false); - if (!read) - break; - } - catch (OperationCanceledException) - { - break; - } - catch (NatsJSProtocolException pe) - { - protocolException = pe; - goto CONSUME_LOOP; - } - catch (NatsConnectionFailedException) - { - // Connection has permanently failed, stop consuming and rethrow - throw; - } - catch (NatsJSException e) when (e is not NatsJSProtocolException and not NatsJSConnectionException and not NatsJSTimeoutException) - { - // Consumer-related errors (like 503 threshold exceeded), stop consuming and rethrow - throw; - } - catch (NatsJSConnectionException e) - { - _logger.LogWarning(NatsJSLogEvents.Retry, "{Error}. Retrying...", e.Message); - goto CONSUME_LOOP; - } - catch (NatsJSTimeoutException e) - { - notificationHandler?.Invoke(NatsJSTimeoutNotification.Default, cancellationToken); - _logger.LogWarning(NatsJSLogEvents.Retry, "{Error}. Retrying...", e.Message); - goto CONSUME_LOOP; - } - while (true) { - NatsJSMsg msg; - - var canRead = cc.Msgs.TryRead(out msg); - if (!canRead) + // We have to check every call to WaitToReadAsync and TryRead for + // protocol exceptions individually because we can't yield return + // within try-catch. + try + { + var read = await cc.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + if (!read) + break; + } + catch (OperationCanceledException) + { break; - - if (msg.Metadata is not { } metadata) - continue; - - var expected = cseq + 1; - if (metadata.Sequence.Consumer != expected) + } + catch (NatsJSProtocolException pe) + { + protocolException = pe; + goto CONSUME_LOOP; + } + catch (NatsConnectionFailedException) + { + // Connection has permanently failed, stop consuming and rethrow + throw; + } + catch (NatsJSException e) when (e is not NatsJSProtocolException and not NatsJSConnectionException and not NatsJSTimeoutException) + { + // Consumer-related errors (like 503 threshold exceeded), stop consuming and rethrow + throw; + } + catch (NatsJSConnectionException e) { - _logger.LogWarning(NatsJSLogEvents.Retry, "Consumer sequence mismatch. Expected {Expected}, was {SequenceConsumer}. Retrying...", expected, metadata.Sequence.Consumer); + _logger.LogWarning(NatsJSLogEvents.Retry, "{Error}. Retrying...", e.Message); goto CONSUME_LOOP; } + catch (NatsJSTimeoutException e) + { + notificationHandler?.Invoke(NatsJSTimeoutNotification.Default, cancellationToken); + _logger.LogWarning(NatsJSLogEvents.Retry, "{Error}. Retrying...", e.Message); + goto CONSUME_LOOP; + } + + while (true) + { + NatsJSMsg msg; + + var canRead = cc.Msgs.TryRead(out msg); + if (!canRead) + break; - seq = metadata.Sequence.Stream; - cseq = metadata.Sequence.Consumer; + if (msg.Metadata is not { } metadata) + continue; - yield return msg; + var expected = cseq + 1; + if (metadata.Sequence.Consumer != expected) + { + _logger.LogWarning(NatsJSLogEvents.Retry, "Consumer sequence mismatch. Expected {Expected}, was {SequenceConsumer}. Retrying...", expected, metadata.Sequence.Consumer); + goto CONSUME_LOOP; + } + + seq = metadata.Sequence.Stream; + cseq = metadata.Sequence.Consumer; + + yield return msg; + } } } + finally + { + cc.MarkReaderInactive(); + } } CONSUME_LOOP: diff --git a/tests/NATS.Client.JetStream.Tests/ConsumerConsumeTest.cs b/tests/NATS.Client.JetStream.Tests/ConsumerConsumeTest.cs index 13ee99293..1b2fbc7a7 100644 --- a/tests/NATS.Client.JetStream.Tests/ConsumerConsumeTest.cs +++ b/tests/NATS.Client.JetStream.Tests/ConsumerConsumeTest.cs @@ -277,7 +277,11 @@ await Retry.Until( [Fact] public async Task Consume_dispose_test() { - await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + await using var nats = new NatsConnection(new NatsOpts + { + Url = _server.Url, + DrainSubscriptionsOnDispose = true, + }); var prefix = _server.GetNextId(); var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60)); @@ -357,6 +361,80 @@ await Retry.Until( Assert.Equal(0, consumer.Info.NumAckPending); } + [Fact] + public async Task Consume_connection_dispose_drains_buffered_messages() + { + // Disposing the connection while a consume loop is mid-flight must + // drain the buffered messages and let the loop ack them, not race + // CommandWriter teardown and throw ObjectDisposedException. + const int totalMsgs = 100; + const int pullBatch = 20; + const int bailAt = 10; + + var prefix = _server.GetNextId(); + var streamName = $"{prefix}s1"; + var subject = $"{prefix}s1.x"; + var consumerName = $"{prefix}c1"; + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60)); + + await using (var setup = new NatsConnection(new NatsOpts { Url = _server.Url })) + { + var setupJs = new NatsJSContext(setup); + var stream = await setupJs.CreateStreamAsync(new StreamConfig(streamName, [subject]), cts.Token); + + for (var i = 0; i < totalMsgs; i++) + (await setupJs.PublishAsync(subject, $"msg-{i}", cancellationToken: cts.Token)).EnsureSuccess(); + + await stream.CreateOrUpdateConsumerAsync(new ConsumerConfig(consumerName), cts.Token); + } + + var nats = new NatsConnection(new NatsOpts + { + Url = _server.Url, + DrainSubscriptionsOnDispose = true, + ConsumerDrainOnDisposeTimeout = TimeSpan.FromSeconds(30), + }); + var js = new NatsJSContext(nats); + var consumer = await js.GetConsumerAsync(streamName, consumerName, cts.Token); + + var reachedBail = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var consumeTask = Task.Run( + async () => + { + var count = 0; + await foreach (var msg in consumer.ConsumeAsync(opts: new NatsJSConsumeOpts { MaxMsgs = pullBatch }, cancellationToken: cts.Token)) + { + count++; + await msg.AckAsync(cancellationToken: cts.Token); + if (count >= bailAt) + reachedBail.TrySetResult(); + await Task.Delay(50, cts.Token); + } + + return count; + }, + cts.Token); + + await reachedBail.Task.WaitAsync(cts.Token); + + await nats.DisposeAsync(); + + var consumed = await consumeTask; + + await using var check = new NatsConnection(new NatsOpts { Url = _server.Url }); + var checkJs = new NatsJSContext(check); + var info = (await checkJs.GetConsumerAsync(streamName, consumerName, cts.Token)).Info; + + // Whatever the reader processed must all be acked: nothing left + // pending, AckFloor matches the client-side count, NumPending is + // the rest of the stream. + Assert.True(consumed >= bailAt, $"consumed {consumed} should be >= {bailAt}"); + Assert.Equal(0, info.NumAckPending); + Assert.Equal((ulong)consumed, (ulong)info.AckFloor.ConsumerSeq); + Assert.Equal((ulong)(totalMsgs - consumed), info.NumPending); + } + [Fact] public async Task Consume_stop_test() { diff --git a/tests/NATS.Client.JetStream.Tests/ConsumerFetchTest.cs b/tests/NATS.Client.JetStream.Tests/ConsumerFetchTest.cs index ed1ed755d..c67d7f124 100644 --- a/tests/NATS.Client.JetStream.Tests/ConsumerFetchTest.cs +++ b/tests/NATS.Client.JetStream.Tests/ConsumerFetchTest.cs @@ -1,5 +1,6 @@ using NATS.Client.Core.Tests; using NATS.Client.Core2.Tests; +using NATS.Client.JetStream.Models; namespace NATS.Client.JetStream.Tests; @@ -77,6 +78,8 @@ public async Task FetchNoWait_test(NatsRequestReplyMode mode) Assert.Equal(10, count); } + // Canary for the dispose race: messages in-flight when DisposeAsync is + // called must all reach the channel before it completes. [Theory] // TODO: Fix this test @@ -84,14 +87,25 @@ public async Task FetchNoWait_test(NatsRequestReplyMode mode) [InlineData(NatsRequestReplyMode.SharedInbox)] public async Task Fetch_dispose_test(NatsRequestReplyMode mode) { - await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url, RequestReplyMode = mode }); + await using var nats = new NatsConnection(new NatsOpts + { + Url = _server.Url, + RequestReplyMode = mode, + DrainSubscriptionsOnDispose = true, + }); var prefix = _server.GetNextId(); var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60)); var js = new NatsJSContext(nats); - var stream = await js.CreateStreamAsync($"{prefix}s1", new[] { $"{prefix}s1.*" }, cts.Token); - var consumer = (NatsJSConsumer)await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token); + + var streamName = $"{prefix}s1"; + var subjectSub = $"{prefix}s1.*"; + var consumerName = $"{prefix}c1"; + var subject = $"{prefix}s1.foo"; + + await js.CreateStreamAsync(streamName, [subjectSub], cts.Token); + var consumer = (NatsJSConsumer)await js.CreateOrUpdateConsumerAsync(streamName, consumerName, cancellationToken: cts.Token); var fetchOpts = new NatsJSFetchOpts { @@ -102,7 +116,7 @@ public async Task Fetch_dispose_test(NatsRequestReplyMode mode) for (var i = 0; i < 100; i++) { - var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer.Default, cancellationToken: cts.Token); + var ack = await js.PublishAsync(subject, new TestData { Test = i }, serializer: TestDataJsonSerializer.Default, cancellationToken: cts.Token); ack.EnsureSuccess(); } @@ -133,7 +147,7 @@ await Retry.Until( "ack pending 9", async () => { - var c = await js.GetConsumerAsync($"{prefix}s1", $"{prefix}c1", cts.Token); + var c = await js.GetConsumerAsync(streamName, consumerName, cts.Token); _output.WriteLine($"pend1:{c.Info.NumAckPending}"); return c.Info.NumAckPending == 9; }, @@ -150,7 +164,7 @@ await Retry.Until( "ack pending 0", async () => { - var c = await js.GetConsumerAsync($"{prefix}s1", $"{prefix}c1", cts.Token); + var c = await js.GetConsumerAsync(streamName, consumerName, cts.Token); _output.WriteLine($"pend:{c.Info.NumAckPending}"); return c.Info.NumAckPending == 0; }, @@ -159,4 +173,73 @@ await Retry.Until( await consumer.RefreshAsync(cts.Token); Assert.Equal(0, consumer.Info.NumAckPending); } + + [Fact] + public async Task Fetch_connection_dispose_drains_buffered_messages() + { + const int totalMsgs = 100; + const int pullBatch = 20; + const int bailAt = 10; + + var prefix = _server.GetNextId(); + var streamName = $"{prefix}s1"; + var subject = $"{prefix}s1.x"; + var consumerName = $"{prefix}c1"; + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60)); + + await using (var setup = new NatsConnection(new NatsOpts { Url = _server.Url })) + { + var setupJs = new NatsJSContext(setup); + var stream = await setupJs.CreateStreamAsync(new StreamConfig(streamName, [subject]), cts.Token); + + for (var i = 0; i < totalMsgs; i++) + (await setupJs.PublishAsync(subject, $"msg-{i}", cancellationToken: cts.Token)).EnsureSuccess(); + + await stream.CreateOrUpdateConsumerAsync(new ConsumerConfig(consumerName), cts.Token); + } + + var nats = new NatsConnection(new NatsOpts + { + Url = _server.Url, + DrainSubscriptionsOnDispose = true, + ConsumerDrainOnDisposeTimeout = TimeSpan.FromSeconds(30), + }); + var js = new NatsJSContext(nats); + var consumer = await js.GetConsumerAsync(streamName, consumerName, cts.Token); + + var reachedBail = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var fetchTask = Task.Run( + async () => + { + var count = 0; + var fetchOpts = new NatsJSFetchOpts { MaxMsgs = pullBatch, Expires = TimeSpan.FromSeconds(30) }; + await foreach (var msg in consumer.FetchAsync(fetchOpts, cancellationToken: cts.Token)) + { + count++; + await msg.AckAsync(cancellationToken: cts.Token); + if (count >= bailAt) + reachedBail.TrySetResult(); + await Task.Delay(50, cts.Token); + } + + return count; + }, + cts.Token); + + await reachedBail.Task.WaitAsync(cts.Token); + + await nats.DisposeAsync(); + + var consumed = await fetchTask; + + await using var check = new NatsConnection(new NatsOpts { Url = _server.Url }); + var checkJs = new NatsJSContext(check); + var info = (await checkJs.GetConsumerAsync(streamName, consumerName, cts.Token)).Info; + + Assert.True(consumed >= bailAt, $"consumed {consumed} should be >= {bailAt}"); + Assert.Equal(0, info.NumAckPending); + Assert.Equal((ulong)consumed, (ulong)info.AckFloor.ConsumerSeq); + Assert.Equal((ulong)(totalMsgs - consumed), info.NumPending); + } } diff --git a/tests/NATS.Client.JetStream.Tests/OrderedConsumerTest.cs b/tests/NATS.Client.JetStream.Tests/OrderedConsumerTest.cs index 85235fffb..dd240f95b 100644 --- a/tests/NATS.Client.JetStream.Tests/OrderedConsumerTest.cs +++ b/tests/NATS.Client.JetStream.Tests/OrderedConsumerTest.cs @@ -494,4 +494,75 @@ public async Task Fetch_maxbytes_with_consumer_deletion() _output.WriteLine($"Received {count} messages with MaxBytes mode"); Assert.True(count > 5, "Should have received messages after consumer deletion"); } + + [Fact] + public async Task OrderedConsume_connection_dispose_drains_buffered_messages() + { + const int totalMsgs = 100; + const int pullBatch = 20; + const int bailAt = 10; + + var prefix = _server.GetNextId(); + var streamName = $"{prefix}s1"; + var subject = $"{prefix}s1.x"; + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60)); + + await using (var setup = new NatsConnection(new NatsOpts { Url = _server.Url })) + { + var setupJs = new NatsJSContext(setup); + await setupJs.CreateStreamAsync(new StreamConfig(streamName, [subject]), cts.Token); + + for (var i = 0; i < totalMsgs; i++) + (await setupJs.PublishAsync(subject, i, cancellationToken: cts.Token)).EnsureSuccess(); + } + + var nats = new NatsConnection(new NatsOpts + { + Url = _server.Url, + DrainSubscriptionsOnDispose = true, + ConsumerDrainOnDisposeTimeout = TimeSpan.FromSeconds(30), + }); + var js = new NatsJSContext(nats); + var stream = await js.GetStreamAsync(streamName, cancellationToken: cts.Token); + var consumer = await stream.CreateOrderedConsumerAsync(cancellationToken: cts.Token); + + var reachedBail = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var lastSeq = -1; + Exception? readerError = null; + var consumeTask = Task.Run( + async () => + { + var count = 0; + try + { + var consumeOpts = new NatsJSConsumeOpts { MaxMsgs = pullBatch, Expires = TimeSpan.FromSeconds(30) }; + await foreach (var msg in consumer.ConsumeAsync(opts: consumeOpts, cancellationToken: cts.Token)) + { + Volatile.Write(ref lastSeq, msg.Data); + count++; + if (count >= bailAt) + reachedBail.TrySetResult(); + await Task.Delay(50, cts.Token); + } + } + catch (Exception e) when (e is not OperationCanceledException) + { + readerError = e; + } + + return count; + }, + cts.Token); + + await reachedBail.Task.WaitAsync(cts.Token); + + await nats.DisposeAsync(); + + var consumed = await consumeTask; + + Assert.Null(readerError); + Assert.True(consumed >= bailAt, $"consumed {consumed} should be >= {bailAt}"); + Assert.Equal(consumed - 1, Volatile.Read(ref lastSeq)); + } } diff --git a/tests/xunit.runsettings b/tests/xunit.runsettings index d9dcca60a..db71244a6 100644 --- a/tests/xunit.runsettings +++ b/tests/xunit.runsettings @@ -1,7 +1,7 @@  - 300000 + 600000 4.0x diff --git a/tools/site_src/blog/2026-04-27-drain-on-dispose.md b/tools/site_src/blog/2026-04-27-drain-on-dispose.md new file mode 100644 index 000000000..d45543368 --- /dev/null +++ b/tools/site_src/blog/2026-04-27-drain-on-dispose.md @@ -0,0 +1,52 @@ +--- +title: Opt-in drain on dispose for JetStream consumers +date: 2026-04-27 +--- + +# Opt-in drain on dispose for JetStream consumers + +_2026-04-27_ + +Disposing a `NatsConnection` while a JetStream `Consume`, `Fetch`, or +ordered-consumer loop was still buffering messages used to drop pending +acks. The socket got closed before the command writer flushed, so the +`UNSUB` and any acks the user loop had queued went nowhere. The messages +stayed `NumAckPending` on the server until `AckWait` expired and were +then redelivered. + +There are now two opt-in options on `NatsOpts` to control this: + +- `DrainSubscriptionsOnDispose` (default `false`). On dispose, drain + subscriptions and flush the command writer before closing the socket. +- `ConsumerDrainOnDisposeTimeout` (default `null`). When set, JetStream + consumer and fetch dispose waits up to this timeout for the user's + consume loop to finish acking the buffered messages before returning. + Only effective when `DrainSubscriptionsOnDispose` is also `true`. +- `DrainPingTimeout` (default `5s`). Per-subscription PING/PONG timeout + used during drain. Increase for high-latency networks where a single + round-trip can exceed the default. + +```csharp +var opts = NatsOpts.Default with +{ + DrainSubscriptionsOnDispose = true, + ConsumerDrainOnDisposeTimeout = TimeSpan.FromSeconds(5), +}; + +await using var nats = new NatsConnection(opts); +``` + +Both default to off, so existing apps see no behavior change. Turn them +on if you want a clean shutdown without redeliveries. + +## A note on resilience + +You do not lose messages either way. JetStream redelivers anything that +wasn't acked once `AckWait` expires, which is the whole point of +at-least-once delivery. The old default is still correct, just noisier: +a graceful shutdown can cause redeliveries you didn't actually need. +For most apps that's fine. These options are for when you want shutdown +to be precise rather than just safe. + +Applies to `Consume`, `Fetch`, and ordered consumer loops in +`NATS.Client.JetStream`. Lands in `2.8.0`. diff --git a/tools/site_src/blog/index.md b/tools/site_src/blog/index.md new file mode 100644 index 000000000..3be7f64be --- /dev/null +++ b/tools/site_src/blog/index.md @@ -0,0 +1,5 @@ +# Blog + +Announcements and short notes about NATS .NET. + +- [Opt-in drain on dispose for JetStream consumers](2026-04-27-drain-on-dispose.md) (2026-04-27) diff --git a/tools/site_src/blog/toc.yml b/tools/site_src/blog/toc.yml new file mode 100644 index 000000000..dfb8e6fa9 --- /dev/null +++ b/tools/site_src/blog/toc.yml @@ -0,0 +1,2 @@ +- name: Opt-in drain on dispose for JetStream consumers (2026-04-27) + href: 2026-04-27-drain-on-dispose.md diff --git a/tools/site_src/docfx.json b/tools/site_src/docfx.json index 79189d8cd..578493469 100644 --- a/tools/site_src/docfx.json +++ b/tools/site_src/docfx.json @@ -29,6 +29,8 @@ "files": [ "documentation/**.md", "documentation/**/toc.yml", + "blog/**.md", + "blog/**/toc.yml", "toc.yml", "*.md" ], diff --git a/tools/site_src/toc.yml b/tools/site_src/toc.yml index b49e5aebc..9b91ec115 100644 --- a/tools/site_src/toc.yml +++ b/tools/site_src/toc.yml @@ -5,6 +5,10 @@ href: api/ homepage: api/index.md +- name: Blog + href: blog/ + homepage: blog/index.md + - name: GitHub href: https://github.com/nats-io/nats.net