JetStream ordered consume stability fixes

src/NATS.Client.Core/NatsConnection.Subscribe.cs

@@ -14,31 +14,22 @@ public partial class NatsConnection
     /// <inheritdoc />
     public async IAsyncEnumerable<NatsMsg<T>> SubscribeAsync<T>(string subject, string? queueGroup = default, INatsDeserialize<T>? serializer = default, NatsSubOpts? opts = default, [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
-        var anchor = Interlocked.Increment(ref _subAnchorId);
-        try
-        {
-            serializer ??= Opts.SerializerRegistry.GetDeserializer<T>();
-
-            await using var sub = new NatsSub<T>(this, SubscriptionManager.GetManagerFor(subject), subject, queueGroup, opts, serializer, cancellationToken);
+        serializer ??= Opts.SerializerRegistry.GetDeserializer<T>();
 
-            _subAnchor[anchor] = sub;
+        await using var sub = new NatsSub<T>(this, SubscriptionManager.GetManagerFor(subject), subject, queueGroup, opts, serializer, cancellationToken);
+        using var anchor = RegisterSubAnchor(sub);
 
-            await SubAsync(sub, cancellationToken: cancellationToken).ConfigureAwait(false);
+        await SubAsync(sub, cancellationToken: cancellationToken).ConfigureAwait(false);
 
-            // We don't cancel the channel reader here because we want to keep reading until the subscription
-            // channel writer completes so that messages left in the channel can be consumed before exit the loop.
-            while (await sub.Msgs.WaitToReadAsync(CancellationToken.None).ConfigureAwait(false))
+        // We don't cancel the channel reader here because we want to keep reading until the subscription
+        // channel writer completes so that messages left in the channel can be consumed before exit the loop.
+        while (await sub.Msgs.WaitToReadAsync(CancellationToken.None).ConfigureAwait(false))
+        {
+            while (sub.Msgs.TryRead(out var msg))
             {
-                while (sub.Msgs.TryRead(out var msg))
-                {
-                    yield return msg;
-                }
+                yield return msg;
             }
         }
-        finally
-        {
-            _subAnchor.TryRemove(anchor, out _);
-        }
     }
 
     internal async ValueTask<NatsSub<T>> SubscribeInternalAsync<T>(string subject, string? queueGroup = default, INatsDeserialize<T>? serializer = default, NatsSubOpts? opts = default, CancellationToken cancellationToken = default)
@@ -48,4 +39,30 @@ internal async ValueTask<NatsSub<T>> SubscribeInternalAsync<T>(string subject, s
         await SubAsync(sub, cancellationToken).ConfigureAwait(false);
         return sub;
     }
+
+    /// <summary>
+    /// Make sure subscription is not collected until the end of the scope.
+    /// </summary>
+    /// <param name="sub">Subscription object</param>
+    /// <returns>Disposable</returns>
+    /// <remarks>
+    /// We must keep subscription alive until the end of its scope especially in async iterators.
+    /// Otherwise subscription is collected because subscription manager only holds a weak reference to it.
+    /// </remarks>
+    internal IDisposable RegisterSubAnchor(NatsSubBase sub) => new SubAnchor(this, sub);
+
+    internal class SubAnchor : IDisposable
+    {
+        private readonly NatsConnection _nats;
+        private readonly long _anchor;
+
+        public SubAnchor(NatsConnection nats, NatsSubBase sub)
+        {
+            _nats = nats;
+            _anchor = Interlocked.Increment(ref _nats._subAnchorId);
+            _nats._subAnchor[_anchor] = sub;
+        }
+
+        public void Dispose() => _nats._subAnchor.TryRemove(_anchor, out _);
+    }
 }