From 159032df5b3db816d8333168ae60a016c26fad13 Mon Sep 17 00:00:00 2001 From: Caleb Lloyd Date: Tue, 4 Jun 2024 22:40:13 -0400 Subject: [PATCH] prefer sub channel ReadAllAsync and rm SubAnchor Signed-off-by: Caleb Lloyd --- .../NatsConnection.RequestReply.cs | 21 +++------ .../NatsConnection.Subscribe.cs | 33 -------------- src/NATS.Client.JetStream/NatsJSConsumer.cs | 9 ---- src/NATS.Client.JetStream/NatsJSContext.cs | 38 +++++++--------- .../NatsJSOrderedConsumer.cs | 3 -- src/NATS.Client.KeyValueStore/NatsKVStore.cs | 44 +++++++------------ 6 files changed, 38 insertions(+), 110 deletions(-) diff --git a/src/NATS.Client.Core/NatsConnection.RequestReply.cs b/src/NATS.Client.Core/NatsConnection.RequestReply.cs index 766c93522..cec7504dc 100644 --- a/src/NATS.Client.Core/NatsConnection.RequestReply.cs +++ b/src/NATS.Client.Core/NatsConnection.RequestReply.cs @@ -42,12 +42,9 @@ public async ValueTask> RequestAsync( await using var sub1 = await RequestSubAsync(subject, data, headers, requestSerializer, replySerializer, requestOpts, replyOpts, cancellationToken) .ConfigureAwait(false); - if (await sub1.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var msg in sub1.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - if (sub1.Msgs.TryRead(out var msg)) - { - return msg; - } + return msg; } throw new NatsNoReplyException(); @@ -63,12 +60,9 @@ public async ValueTask> RequestAsync( await using var sub = await RequestSubAsync(subject, data, headers, requestSerializer, replySerializer, requestOpts, replyOpts, cancellationToken) .ConfigureAwait(false); - if (await sub.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - if (sub.Msgs.TryRead(out var msg)) - { - return msg; - } + return msg; } throw new NatsNoReplyException(); @@ -105,12 +99,9 @@ public async IAsyncEnumerable> RequestManyAsync(subject, data, headers, requestSerializer, replySerializer, requestOpts, replyOpts, cancellationToken) .ConfigureAwait(false); - while (await sub.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - while (sub.Msgs.TryRead(out var msg)) - { - yield return msg; - } + yield return msg; } } diff --git a/src/NATS.Client.Core/NatsConnection.Subscribe.cs b/src/NATS.Client.Core/NatsConnection.Subscribe.cs index 7f169d98e..d163db30e 100644 --- a/src/NATS.Client.Core/NatsConnection.Subscribe.cs +++ b/src/NATS.Client.Core/NatsConnection.Subscribe.cs @@ -5,18 +5,11 @@ namespace NATS.Client.Core; public partial class NatsConnection { - // Keep subscription alive until the channel reader completes. - // Otherwise subscription is collected because subscription manager - // only holds a weak reference to it. - private readonly ConcurrentDictionary _subAnchor = new(); - private long _subAnchorId; - /// public async IAsyncEnumerable> SubscribeAsync(string subject, string? queueGroup = default, INatsDeserialize? serializer = default, NatsSubOpts? opts = default, [EnumeratorCancellation] CancellationToken cancellationToken = default) { serializer ??= Opts.SerializerRegistry.GetDeserializer(); - // call to RegisterSubAnchor is no longer needed; sub is kept alive in ActivityEndingMsgReader await using var sub = new NatsSub(this, SubscriptionManager.GetManagerFor(subject), subject, queueGroup, opts, serializer, cancellationToken); await SubAsync(sub, cancellationToken: cancellationToken).ConfigureAwait(false); @@ -36,30 +29,4 @@ public async ValueTask> SubscribeCoreAsync(string subject, string await SubAsync(sub, cancellationToken).ConfigureAwait(false); return sub; } - - /// - /// Make sure subscription is not collected until the end of the scope. - /// - /// Subscription object - /// Disposable - /// - /// 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. - /// - 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 _); - } } diff --git a/src/NATS.Client.JetStream/NatsJSConsumer.cs b/src/NATS.Client.JetStream/NatsJSConsumer.cs index 8f1d0f224..e1c7b400d 100644 --- a/src/NATS.Client.JetStream/NatsJSConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSConsumer.cs @@ -62,9 +62,6 @@ public async IAsyncEnumerable> ConsumeAsync( opts ??= _context.Opts.DefaultConsumeOpts; await using var cc = await ConsumeInternalAsync(serializer, opts, cancellationToken).ConfigureAwait(false); - // Keep subscription alive (since it's a weak ref in subscription manager) until we're done. - using var anchor = _context.Connection.RegisterSubAnchor(cc); - while (!cancellationToken.IsCancellationRequested) { // We have to check calls individually since we can't use yield return in try-catch blocks. @@ -151,9 +148,6 @@ public async IAsyncEnumerable> ConsumeAsync( serializer, cancellationToken: cancellationToken).ConfigureAwait(false); - // Keep subscription alive (since it's a weak ref in subscription manager) until we're done. - using var anchor = _context.Connection.RegisterSubAnchor(f); - await foreach (var natsJSMsg in f.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { return natsJSMsg; @@ -173,9 +167,6 @@ public async IAsyncEnumerable> FetchAsync( await using var fc = await FetchInternalAsync(opts, serializer, cancellationToken).ConfigureAwait(false); - // Keep subscription alive (since it's a weak ref in subscription manager) until we're done. - using var anchor = _context.Connection.RegisterSubAnchor(fc); - while (!cancellationToken.IsCancellationRequested) { // We have to check calls individually since we can't use yield return in try-catch blocks. diff --git a/src/NATS.Client.JetStream/NatsJSContext.cs b/src/NATS.Client.JetStream/NatsJSContext.cs index f16c71166..2097af019 100644 --- a/src/NATS.Client.JetStream/NatsJSContext.cs +++ b/src/NATS.Client.JetStream/NatsJSContext.cs @@ -144,17 +144,14 @@ public async ValueTask PublishAsync( try { - while (await sub.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - while (sub.Msgs.TryRead(out var msg)) + if (msg.Data == null) { - if (msg.Data == null) - { - throw new NatsJSException("No response data received"); - } - - return msg.Data; + throw new NatsJSException("No response data received"); } + + return msg.Data; } } catch (NatsNoRespondersException) @@ -226,27 +223,24 @@ internal async ValueTask> JSRequestAsync(default, jsError.Error); - } - - throw error; + return new NatsJSResponse(default, jsError.Error); } - if (msg.Data == null) - { - throw new NatsJSException("No response data received"); - } + throw error; + } - return new NatsJSResponse(msg.Data, default); + if (msg.Data == null) + { + throw new NatsJSException("No response data received"); } + + return new NatsJSResponse(msg.Data, default); } if (sub is NatsSubBase { EndReason: NatsSubEndReason.Exception, Exception: not null } sb) diff --git a/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs b/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs index cb0a3a862..a52d3a0f5 100644 --- a/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs @@ -80,9 +80,6 @@ public async IAsyncEnumerable> ConsumeAsync( await using (var cc = await consumer.OrderedConsumeInternalAsync(serializer, opts, cancellationToken)) { - // Keep subscription alive (since it's a wek ref in subscription manager) until we're done. - using var anchor = _context.Connection.RegisterSubAnchor(cc); - while (true) { // We have to check every call to WaitToReadAsync and TryRead for diff --git a/src/NATS.Client.KeyValueStore/NatsKVStore.cs b/src/NATS.Client.KeyValueStore/NatsKVStore.cs index ba8e26185..c773b1db7 100644 --- a/src/NATS.Client.KeyValueStore/NatsKVStore.cs +++ b/src/NATS.Client.KeyValueStore/NatsKVStore.cs @@ -331,12 +331,9 @@ public async IAsyncEnumerable> WatchAsync(IEnumerable } } - while (await watcher.Entries.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var entry in watcher.Entries.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - while (watcher.Entries.TryRead(out var entry)) - { - yield return entry; - } + yield return entry; } } @@ -359,14 +356,11 @@ public async IAsyncEnumerable> HistoryAsync(string key, INatsD await using var watcher = await WatchInternalAsync([key], serializer, opts, cancellationToken); - while (await watcher.Entries.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var entry in watcher.Entries.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - while (watcher.Entries.TryRead(out var entry)) - { - yield return entry; - if (entry.Delta == 0) - yield break; - } + yield return entry; + if (entry.Delta == 0) + yield break; } } @@ -398,15 +392,12 @@ public async ValueTask PurgeDeletesAsync(NatsKVPurgeOpts? opts = default, Cancel if (watcher.InitialConsumer.Info.NumPending == 0) return; - while (await watcher.Entries.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var entry in watcher.Entries.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - while (watcher.Entries.TryRead(out var entry)) - { - if (entry.Operation is NatsKVOperation.Purge or NatsKVOperation.Del) - deleted.Add(entry); - if (entry.Delta == 0) - goto PURGE_LOOP_DONE; - } + if (entry.Operation is NatsKVOperation.Purge or NatsKVOperation.Del) + deleted.Add(entry); + if (entry.Delta == 0) + goto PURGE_LOOP_DONE; } } @@ -442,15 +433,12 @@ public async IAsyncEnumerable GetKeysAsync(NatsKVWatchOpts? opts = defau if (watcher.InitialConsumer.Info.NumPending == 0) yield break; - while (await watcher.Entries.WaitToReadAsync(cancellationToken).ConfigureAwait(false)) + await foreach (var entry in watcher.Entries.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - while (watcher.Entries.TryRead(out var entry)) - { - if (entry.Operation is NatsKVOperation.Put) - yield return entry.Key; - if (entry.Delta == 0) - yield break; - } + if (entry.Operation is NatsKVOperation.Put) + yield return entry.Key; + if (entry.Delta == 0) + yield break; } }