Fix message loss on consumer dispose

src/NATS.Client.Core/Internal/SubscriptionManager.cs

@@ -149,18 +149,23 @@ public async ValueTask DisposeAsync()
         _cts.Cancel();
 #endif
 
-        WeakReference<NatsSubBase>[] subRefs;
+        var subs = new List<NatsSubBase>();
         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<int> 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 _);
+        }
+    }
+
     /// <summary>
     /// Returns commands for all the live subscriptions to be used on reconnect so that they can rebuild their connection state on the server.
     /// </summary>

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<NatsSubBase> _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));
+    }
 }

src/NATS.Client.Core/NatsOpts.cs

@@ -153,6 +153,55 @@ public sealed record NatsOpts
 
     public bool WaitUntilSent { get; init; } = false;
 
+    /// <summary>
+    /// When true, <see cref="NatsConnection.DisposeAsync"/> drains subscriptions
+    /// and flushes the command writer before closing the socket. (default: false)
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// 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 <see cref="ConsumerDrainOnDisposeTimeout"/> to be effective.
+    /// </para>
+    /// <para>
+    /// Drain mechanism: each subscription sends UNSUB, then a PING and waits
+    /// for the PONG (bounded by <see cref="DrainPingTimeout"/>). 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 <see cref="ConsumerDrainOnDisposeTimeout"/>
+    /// wait if set, not multiplied by the number of subscriptions.
+    /// </para>
+    /// </remarks>
+    public bool DrainSubscriptionsOnDispose { get; init; } = false;
+
+    /// <summary>
+    /// 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)
+    /// </summary>
+    /// <remarks>
+    /// Only effective when <see cref="DrainSubscriptionsOnDispose"/> 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.
+    /// </remarks>
+    public TimeSpan? ConsumerDrainOnDisposeTimeout { get; init; } = null;
+
+    /// <summary>
+    /// Per-subscription PING/PONG timeout used during drain on dispose.
+    /// (default: 5 seconds)
+    /// </summary>
+    /// <remarks>
+    /// 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
+    /// <see cref="DrainSubscriptionsOnDispose"/> is true.
+    /// </remarks>
+    public TimeSpan DrainPingTimeout { get; init; } = TimeSpan.FromSeconds(5);
+
     /// <summary>
     /// Maximum number of reconnect attempts. (default: -1, unlimited)
     /// </summary>

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;
 
     /// <summary>
     /// Creates a new instance of <see cref="NatsSubBase"/>.
@@ -358,6 +359,82 @@ internal static bool IsHeader503(ReadOnlySequence<byte>? headersBuffer) =>
 
     internal void ClearException() => Interlocked.Exchange(ref _exception, null);
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// No-op unless <see cref="NatsOpts.DrainSubscriptionsOnDispose"/> is
+    /// enabled. The connection-level dispose path also gates entry into
+    /// this method on the same flag (see ordering in
+    /// <see cref="NatsConnection.DisposeAsync"/>).
+    /// </remarks>
+    /// <returns>A <see cref="ValueTask"/> that represents the asynchronous drain operation.</returns>
+    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);
+    }
+
     /// <summary>
     /// 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<byte>? headersBuffer) =>
     /// <returns>ValueTask</returns>
     internal virtual ValueTask WriteReconnectCommandsAsync(CommandWriter commandWriter, int sid) => commandWriter.SubscribeAsync(sid, Subject, QueueGroup, PendingMsgs, CancellationToken.None);
 
+    /// <summary>
+    /// Marks the user-side reader as active so dispose can wait for it to
+    /// drain buffered messages before completing. Pairs with
+    /// <see cref="MarkReaderInactive"/>.
+    /// </summary>
+    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);
+    }
+
+    /// <summary>
+    /// Signals the user-side reader has exited. Pairs with <see cref="MarkReaderActive"/>.
+    /// </summary>
+    internal void MarkReaderInactive()
+    {
+        if (Connection is NatsConnection nc)
+            nc.UnregisterDrainParticipant(this);
+        _readerExited?.TrySetResult();
+    }
+
     /// <summary>
     /// Invoked when a MSG or HMSG arrives for the subscription.
     /// <remarks>
@@ -390,6 +497,21 @@ internal static bool IsHeader503(ReadOnlySequence<byte>? headersBuffer) =>
     /// <returns></returns>
     protected abstract ValueTask ReceiveInternalAsync(string subject, string? replyTo, ReadOnlySequence<byte>? headersBuffer, ReadOnlySequence<byte> payloadBuffer);
 
+    /// <summary>
+    /// 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 <see cref="NatsOpts.ConsumerDrainOnDisposeTimeout"/> is unset.
+    /// </summary>
+    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);
+    }
+
     /// <summary>
     /// 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");
+        }
+    }
 }