Skip to content

Prefer sub channel ReadAllAsync and remove SubAnchor #507

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mtmk merged 1 commit into from
Jun 5, 2024
Merged

Prefer sub channel ReadAllAsync and remove SubAnchor #507

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 6 additions & 15 deletions src/NATS.Client.Core/NatsConnection.RequestReply.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,12 +42,9 @@ public async ValueTask<NatsMsg<TReply>> RequestAsync<TRequest, TReply>(
await using var sub1 = await RequestSubAsync<TRequest, TReply>(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();
Expand All @@ -63,12 +60,9 @@ public async ValueTask<NatsMsg<TReply>> RequestAsync<TRequest, TReply>(
await using var sub = await RequestSubAsync<TRequest, TReply>(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();
Expand Down Expand Up @@ -105,12 +99,9 @@ public async IAsyncEnumerable<NatsMsg<TReply>> RequestManyAsync<TRequest, TReply
await using var sub = await RequestSubAsync<TRequest, TReply>(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;
}
}

Expand Down
33 changes: 0 additions & 33 deletions src/NATS.Client.Core/NatsConnection.Subscribe.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<long, NatsSubBase> _subAnchor = new();
private long _subAnchorId;

/// <inheritdoc />
public async IAsyncEnumerable<NatsMsg<T>> SubscribeAsync<T>(string subject, string? queueGroup = default, INatsDeserialize<T>? serializer = default, NatsSubOpts? opts = default, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
serializer ??= Opts.SerializerRegistry.GetDeserializer<T>();

// call to RegisterSubAnchor is no longer needed; sub is kept alive in ActivityEndingMsgReader
await using var sub = new NatsSub<T>(this, SubscriptionManager.GetManagerFor(subject), subject, queueGroup, opts, serializer, cancellationToken);
await SubAsync(sub, cancellationToken: cancellationToken).ConfigureAwait(false);

Expand All @@ -36,30 +29,4 @@ public async ValueTask<INatsSub<T>> SubscribeCoreAsync<T>(string subject, string
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 _);
}
}
9 changes: 0 additions & 9 deletions src/NATS.Client.JetStream/NatsJSConsumer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,9 +62,6 @@ public async IAsyncEnumerable<NatsJSMsg<T>> ConsumeAsync<T>(
opts ??= _context.Opts.DefaultConsumeOpts;
await using var cc = await ConsumeInternalAsync<T>(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.
Expand Down Expand Up @@ -151,9 +148,6 @@ public async IAsyncEnumerable<NatsJSMsg<T>> ConsumeAsync<T>(
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;
Expand All @@ -173,9 +167,6 @@ public async IAsyncEnumerable<NatsJSMsg<T>> FetchAsync<T>(

await using var fc = await FetchInternalAsync<T>(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.
Expand Down
38 changes: 16 additions & 22 deletions src/NATS.Client.JetStream/NatsJSContext.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,17 +144,14 @@ public async ValueTask<PubAckResponse> PublishAsync<T>(

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)
Expand Down Expand Up @@ -226,27 +223,24 @@ internal async ValueTask<NatsJSResponse<TResponse>> JSRequestAsync<TRequest, TRe
cancellationToken: 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))
if (msg.Error is { } error)
{
if (msg.Error is { } error)
if (error.InnerException is NatsJSApiErrorException jsError)
{
if (error.InnerException is NatsJSApiErrorException jsError)
{
return new NatsJSResponse<TResponse>(default, jsError.Error);
}

throw error;
return new NatsJSResponse<TResponse>(default, jsError.Error);
}

if (msg.Data == null)
{
throw new NatsJSException("No response data received");
}
throw error;
}

return new NatsJSResponse<TResponse>(msg.Data, default);
if (msg.Data == null)
{
throw new NatsJSException("No response data received");
}

return new NatsJSResponse<TResponse>(msg.Data, default);
}

if (sub is NatsSubBase { EndReason: NatsSubEndReason.Exception, Exception: not null } sb)
Expand Down
3 changes: 0 additions & 3 deletions src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,9 +80,6 @@ public async IAsyncEnumerable<NatsJSMsg<T>> ConsumeAsync<T>(

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
Expand Down
44 changes: 16 additions & 28 deletions src/NATS.Client.KeyValueStore/NatsKVStore.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,12 +331,9 @@ public async IAsyncEnumerable<NatsKVEntry<T>> WatchAsync<T>(IEnumerable<string>
}
}

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;
}
}

Expand All @@ -359,14 +356,11 @@ public async IAsyncEnumerable<NatsKVEntry<T>> HistoryAsync<T>(string key, INatsD

await using var watcher = await WatchInternalAsync<T>([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;
}
}

Expand Down Expand Up @@ -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;
}
}

Expand Down Expand Up @@ -442,15 +433,12 @@ public async IAsyncEnumerable<string> 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;
}
}

Expand Down