Add explicit subscription and consumer drain

src/NATS.Client.Core/INatsSub.cs

@@ -27,4 +27,19 @@ public interface INatsSub<T> : IAsyncDisposable
     /// </summary>
     /// <returns>A <see cref="ValueTask"/> that represents the asynchronous server UNSUB operation.</returns>
     public ValueTask UnsubscribeAsync();
+
+    /// <summary>
+    /// Drain the subscription: stop receiving new messages while letting messages
+    /// already buffered be read, keeping the connection usable.
+    /// </summary>
+    /// <remarks>
+    /// Sends an unsubscribe to the server, waits for a PING/PONG round-trip so messages
+    /// already in flight are delivered, then completes <see cref="Msgs"/>. Keep reading
+    /// <see cref="Msgs"/> until it completes to process the buffered messages. Unlike
+    /// <see cref="UnsubscribeAsync"/>, drain does not drop messages still in the socket
+    /// buffer. The connection stays open.
+    /// </remarks>
+    /// <param name="cancellationToken">Bounds the best-effort PING/PONG fence wait, which lasts at most <see cref="NatsOpts.DrainPingTimeout"/>; the token can only shorten it, not extend it past that timeout.</param>
+    /// <returns>A <see cref="ValueTask"/> that represents the asynchronous drain operation.</returns>
+    public ValueTask DrainAsync(CancellationToken cancellationToken = default);
 }

src/NATS.Client.Core/NatsConnection.cs

@@ -1254,7 +1254,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));
     }

src/NATS.Client.Core/NatsOpts.cs

@@ -197,14 +197,15 @@ public sealed record NatsOpts
     public TimeSpan? ConsumerDrainOnDisposeTimeout { get; init; } = null;
 
     /// <summary>
-    /// Per-subscription PING/PONG timeout used during drain on dispose.
+    /// Per-subscription PING/PONG timeout used during drain.
     /// (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.
+    /// where a single round-trip can exceed the default. Applies both to drain on
+    /// dispose (when <see cref="DrainSubscriptionsOnDispose"/> is true) and to the
+    /// explicit per-subscription drain API.
     /// </remarks>
     public TimeSpan DrainPingTimeout { get; init; } = TimeSpan.FromSeconds(5);
 

src/NATS.Client.Core/NatsSubBase.cs

@@ -278,6 +278,29 @@ public ValueTask UnsubscribeAsync()
         return _manager.RemoveAsync(this);
     }
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// Unlike <see cref="UnsubscribeAsync"/>, 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 <paramref name="cancellationToken"/> and
+    /// <see cref="NatsOpts.DrainPingTimeout"/>.
+    /// <para>
+    /// Cancellation only shortens the PING/PONG fence wait; it cannot abort the
+    /// drain. Once <see cref="DrainAsync"/> is entered the subscription is always
+    /// committed to the drained state, and cancelling the token does not throw or
+    /// keep the subscription alive.
+    /// </para>
+    /// </remarks>
+    /// <param name="cancellationToken">Bounds the best-effort PING/PONG fence wait, which lasts at most <see cref="NatsOpts.DrainPingTimeout"/>; the token can only shorten it, not extend it past that timeout. Cancellation does not abort the drain.</param>
+    /// <returns>A <see cref="ValueTask"/> that represents the asynchronous drain operation.</returns>
+    public virtual ValueTask DrainAsync(CancellationToken cancellationToken = default) => DrainCoreAsync(cancellationToken);
+
     /// <summary>
     /// Disposes the instance asynchronously.
     /// </summary>
@@ -405,10 +428,9 @@ 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.
+    /// Dispose-path drain. Runs the same UNSUB plus PING/PONG fence as
+    /// <see cref="DrainAsync"/>, then waits for an active user reader to drain
+    /// the channel before returning.
     /// </summary>
     /// <remarks>
     /// No-op unless <see cref="NatsOpts.DrainSubscriptionsOnDispose"/> is
@@ -417,64 +439,12 @@ internal static bool IsHeader503(ReadOnlySequence<byte>? headersBuffer) =>
     /// <see cref="NatsConnection.DisposeAsync"/>).
     /// </remarks>
     /// <returns>A <see cref="ValueTask"/> that represents the asynchronous drain operation.</returns>
-    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
@@ -544,6 +514,18 @@ internal void MarkReaderInactive()
     /// <returns></returns>
     protected abstract ValueTask ReceiveInternalAsync(string subject, string? replyTo, ReadOnlySequence<byte>? headersBuffer, ReadOnlySequence<byte> payloadBuffer);
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    protected virtual void StopDelivery()
+    {
+    }
+
     /// <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
@@ -657,6 +639,87 @@ private void DecrementActiveSubscription()
             Telemetry.ActiveSubscriptions.Add(-1, Telemetry.BuildMetricTags(Connection, Telemetry.Constants.OpSub));
     }
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    private async ValueTask DrainCoreAsync(CancellationToken cancellationToken)
+    {
+        var needsUnsub = false;
+        lock (_gate)
+        {
+            if (!_unsubscribed)
+            {
+                _unsubscribed = true;
+                needsUnsub = true;
+            }
+        }
+
+        if (!needsUnsub)
+            return;
+
+        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.
+            // 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);
+        }
+        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.
+        }
+        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 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)
+                mgr?.RemoveFromRouting(sid);
+
+            TryComplete();
+        }
+    }
+
     private async ValueTask WaitForReaderDrainCoreAsync(Task task, TimeSpan timeout)
     {
         try

src/NATS.Client.JetStream/Internal/NatsJSConsume.cs

@@ -50,6 +50,7 @@ internal class NatsJSConsume<TMsg> : 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<TMsg>)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,7 +206,15 @@ public ValueTask CallMsgNextAsync(string origin, ConsumerGetnextRequest request,
 
     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)
     {
@@ -222,7 +238,7 @@ public override async ValueTask DisposeAsync()
     {
         try
         {
-            await DrainAsync().ConfigureAwait(false);
+            await DrainOnDisposeAsync().ConfigureAwait(false);
             await base.DisposeAsync().ConfigureAwait(false);
         }
         finally
@@ -299,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,
@@ -529,6 +554,12 @@ private void CompleteStop()
             state =>
             {
                 var self = (NatsJSConsume<TMsg>)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);
             },

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

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