diff --git a/src/NATS.Client.Core/INatsSub.cs b/src/NATS.Client.Core/INatsSub.cs index a9b7a2c36..b4e939573 100644 --- a/src/NATS.Client.Core/INatsSub.cs +++ b/src/NATS.Client.Core/INatsSub.cs @@ -27,4 +27,19 @@ public interface INatsSub : IAsyncDisposable /// /// A that represents the asynchronous server UNSUB operation. public ValueTask UnsubscribeAsync(); + + /// + /// Drain the subscription: stop receiving new messages while letting messages + /// already buffered be read, keeping the connection usable. + /// + /// + /// Sends an unsubscribe to the server, waits for a PING/PONG round-trip so messages + /// already in flight are delivered, then completes . Keep reading + /// until it completes to process the buffered messages. Unlike + /// , drain does not drop messages still in the socket + /// buffer. The connection stays open. + /// + /// Bounds the best-effort PING/PONG fence wait, which lasts at most ; the token can only shorten it, not extend it past that timeout. + /// A that represents the asynchronous drain operation. + public ValueTask DrainAsync(CancellationToken cancellationToken = default); } diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs index 70f4219ea..743287a31 100644 --- a/src/NATS.Client.Core/NatsConnection.cs +++ b/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)); } diff --git a/src/NATS.Client.Core/NatsOpts.cs b/src/NATS.Client.Core/NatsOpts.cs index e4ecf23a2..7be0acc5c 100644 --- a/src/NATS.Client.Core/NatsOpts.cs +++ b/src/NATS.Client.Core/NatsOpts.cs @@ -197,14 +197,15 @@ public sealed record NatsOpts public TimeSpan? ConsumerDrainOnDisposeTimeout { get; init; } = null; /// - /// Per-subscription PING/PONG timeout used during drain on dispose. + /// Per-subscription PING/PONG timeout used during drain. /// (default: 5 seconds) /// /// /// 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 - /// is true. + /// where a single round-trip can exceed the default. Applies both to drain on + /// dispose (when is true) and to the + /// explicit per-subscription drain API. /// public TimeSpan DrainPingTimeout { get; init; } = TimeSpan.FromSeconds(5); diff --git a/src/NATS.Client.Core/NatsSubBase.cs b/src/NATS.Client.Core/NatsSubBase.cs index 86464973c..b8b655a13 100644 --- a/src/NATS.Client.Core/NatsSubBase.cs +++ b/src/NATS.Client.Core/NatsSubBase.cs @@ -278,6 +278,29 @@ public ValueTask UnsubscribeAsync() return _manager.RemoveAsync(this); } + /// + /// 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. + /// + /// + /// Unlike , 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 and + /// . + /// + /// Cancellation only shortens the PING/PONG fence wait; it cannot abort the + /// drain. Once is entered the subscription is always + /// committed to the drained state, and cancelling the token does not throw or + /// keep the subscription alive. + /// + /// + /// Bounds the best-effort PING/PONG fence wait, which lasts at most ; the token can only shorten it, not extend it past that timeout. Cancellation does not abort the drain. + /// A that represents the asynchronous drain operation. + public virtual ValueTask DrainAsync(CancellationToken cancellationToken = default) => DrainCoreAsync(cancellationToken); + /// /// Disposes the instance asynchronously. /// @@ -405,10 +428,9 @@ internal static bool IsHeader503(ReadOnlySequence? headersBuffer) => internal void ClearException() => Interlocked.Exchange(ref _exception, null); /// - /// 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 + /// , then waits for an active user reader to drain + /// the channel before returning. /// /// /// No-op unless is @@ -417,64 +439,12 @@ internal static bool IsHeader503(ReadOnlySequence? headersBuffer) => /// ). /// /// A that represents the asynchronous drain operation. - 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() /// protected abstract ValueTask ReceiveInternalAsync(string subject, string? replyTo, ReadOnlySequence? headersBuffer, ReadOnlySequence payloadBuffer); + /// + /// 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. + /// + protected virtual void StopDelivery() + { + } + /// /// 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)); } + /// + /// 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. + /// + 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 diff --git a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs index ca276d7c1..ed812c721 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs +++ b/src/NATS.Client.JetStream/Internal/NatsJSConsume.cs @@ -50,6 +50,7 @@ internal class NatsJSConsume : 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)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)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); }, diff --git a/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs b/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs index 5dcf6d429..6abc50b9a 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSFetch.cs +++ b/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 diff --git a/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs b/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs index cc2cacf4a..e595cc348 100644 --- a/src/NATS.Client.JetStream/Internal/NatsJSOrderedConsume.cs +++ b/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 diff --git a/src/NATS.Client.JetStream/NatsJSConsumer.cs b/src/NATS.Client.JetStream/NatsJSConsumer.cs index dd66e0b43..bc08edf43 100644 --- a/src/NATS.Client.JetStream/NatsJSConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSConsumer.cs @@ -66,20 +66,38 @@ public async IAsyncEnumerable> ConsumeAsync( [EnumeratorCancellation] CancellationToken cancellationToken = default) { opts ??= _context.Opts.DefaultConsumeOpts; + var drainOnCancel = opts.DrainOnCancel; await using var cc = await ConsumeInternalAsync(serializer, opts, cancellationToken).ConfigureAwait(false); cc.MarkReaderActive(); + + // When draining on cancel we keep reading buffered messages until the channel + // completes instead of bailing out immediately. The drain stops pulling, fences + // in-flight messages with a PING/PONG and completes the channel, leaving the + // connection usable so handlers can still ACK the messages they receive. + var draining = false; + var drainTask = default(ValueTask); try { - while (!cancellationToken.IsCancellationRequested) + while (true) { // We have to check calls individually since we can't use yield return in try-catch blocks. + // WaitToReadAsync throws if the token is already canceled (even with buffered + // messages), so the catch below is the single place that reacts to cancellation: + // start the drain when DrainOnCancel is set, otherwise stop. bool ready; try { - ready = await cc.Msgs.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + ready = await cc.Msgs.WaitToReadAsync(draining ? CancellationToken.None : cancellationToken).ConfigureAwait(false); } catch (OperationCanceledException) { + if (drainOnCancel && !draining) + { + draining = true; + drainTask = cc.DrainAsync(); + continue; + } + ready = false; } catch (NatsConnectionFailedException) @@ -96,8 +114,13 @@ public async IAsyncEnumerable> ConsumeAsync( if (!ready) yield break; - while (!cancellationToken.IsCancellationRequested) + while (true) { + // Stop promptly on cancel unless we are draining, in which case we + // keep reading the buffered messages until the channel completes. + if (!draining && cancellationToken.IsCancellationRequested) + break; + bool read; NatsJSMsg jsMsg; try @@ -135,6 +158,11 @@ public async IAsyncEnumerable> ConsumeAsync( finally { cc.MarkReaderInactive(); + + // Observe the drain completion. DrainAsync swallows cancellation and + // connection errors internally, so this does not throw on the normal path. + if (draining) + await drainTask.ConfigureAwait(false); } } diff --git a/src/NATS.Client.JetStream/NatsJSOpts.cs b/src/NATS.Client.JetStream/NatsJSOpts.cs index e36868d90..ef7cf614e 100644 --- a/src/NATS.Client.JetStream/NatsJSOpts.cs +++ b/src/NATS.Client.JetStream/NatsJSOpts.cs @@ -160,6 +160,19 @@ public record NatsJSConsumeOpts /// Set to -1 to disable this check. (default: 10) /// public int MaxConsecutive503Errors { get; init; } = 10; + + /// + /// When true, cancelling the token passed to ConsumeAsync drains the consumer + /// gracefully (stop pulling, deliver buffered messages so handlers can ACK, + /// then complete) and leaves the connection usable, instead of stopping + /// immediately. (default: false) + /// + /// + /// Only the standard pull consumer's ConsumeAsync honors this. Ordered consumers + /// (see ) do not ACK and always stop promptly + /// on cancel, so the option has no effect there. + /// + public bool DrainOnCancel { get; init; } } /// diff --git a/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs new file mode 100644 index 000000000..56b8d63d0 --- /dev/null +++ b/tests/NATS.Client.Core2.Tests/SubscriptionDrainTest.cs @@ -0,0 +1,122 @@ +using NATS.Client.Core2.Tests; + +namespace NATS.Client.Core.Tests; + +[Collection("nats-server")] +public class SubscriptionDrainTest +{ + private readonly NatsServerFixture _server; + + public SubscriptionDrainTest(NatsServerFixture server) + { + _server = server; + } + + [Fact] + public async Task Drain_preserves_buffered_messages_and_completes_channel() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + const int count = 10; + for (var i = 0; i < count; i++) + await nats.PublishAsync(subject, i, cancellationToken: cancellationToken); + + // PING/PONG round-trip guarantees the published messages have been + // delivered back to our subscription channel before we drain. + await nats.PingAsync(cancellationToken); + + await sub.DrainAsync(cancellationToken); + + // All buffered messages are still readable and the channel completes. + var received = new List(); + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken)) + received.Add(msg.Data); + + Assert.Equal(Enumerable.Range(0, count), received); + } + + [Fact] + public async Task Drain_delivers_messages_still_in_flight_after_unsub() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + // Kept below the default SubPendingChannelCapacity so nothing is dropped for being + // over capacity; what is under test is the drain fence, not back-pressure. + const int count = 1000; + for (var i = 0; i < count; i++) + await nats.PublishAsync(subject, i, cancellationToken: cancellationToken); + + // Unlike the test above there is no PING/PONG before draining, so messages the + // server sent before it processed our UNSUB are still in flight (not yet in the + // channel). The drain fences inbound delivery with its own PING/PONG so they all + // land before the channel completes; without that fence the tail would be dropped. + await sub.DrainAsync(cancellationToken); + + var received = new List(); + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken)) + received.Add(msg.Data); + + // Check the count first so a dropped message surfaces as a count mismatch rather + // than a confusing order-equality failure. + Assert.Equal(count, received.Count); + Assert.Equal(Enumerable.Range(0, count), received); + } + + [Fact] + public async Task Drain_stops_new_messages_but_connection_stays_usable() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + await nats.PublishAsync(subject, 1, cancellationToken: cancellationToken); + await nats.PingAsync(cancellationToken); + + await sub.DrainAsync(cancellationToken); + + // Published after the UNSUB; must not reach the drained subscription. + await nats.PublishAsync(subject, 2, cancellationToken: cancellationToken); + await nats.PingAsync(cancellationToken); + + var received = new List(); + await foreach (var msg in sub.Msgs.ReadAllAsync(cancellationToken)) + received.Add(msg.Data); + + Assert.Equal(new[] { 1 }, received); + + // The connection is still usable for a fresh subscription and publish. + var sub2 = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + await nats.PublishAsync(subject, 42, cancellationToken: cancellationToken); + var next = await sub2.Msgs.ReadAsync(cancellationToken); + Assert.Equal(42, next.Data); + await sub2.DisposeAsync(); + } + + [Fact] + public async Task Drain_is_idempotent_and_dispose_after_drain_is_safe() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + var cancellationToken = cts.Token; + + var subject = $"foo.{Guid.NewGuid():N}"; + var sub = await nats.SubscribeCoreAsync(subject, cancellationToken: cancellationToken); + + await sub.DrainAsync(cancellationToken); + await sub.DrainAsync(cancellationToken); + await sub.DisposeAsync(); + } +} diff --git a/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs b/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs new file mode 100644 index 000000000..182fa5633 --- /dev/null +++ b/tests/NATS.Client.JetStream.Tests/ConsumeDrainTest.cs @@ -0,0 +1,111 @@ +using NATS.Client.Core2.Tests; +using NATS.Client.TestUtilities2; + +namespace NATS.Client.JetStream.Tests; + +[Collection("nats-server")] +public class ConsumeDrainTest +{ + private readonly NatsServerFixture _server; + + public ConsumeDrainTest(NatsServerFixture server) => _server = server; + + [Fact] + public async Task Consume_drain_on_cancel_delivers_buffered_and_keeps_connection() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + await nats.ConnectRetryAsync(); + var prefix = _server.GetNextId(); + + var js = new NatsJSContext(nats); + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + + await js.CreateStreamAsync($"{prefix}s1", new[] { $"{prefix}s1.*" }, cts.Token); + await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token); + + const int total = 100; + for (var i = 0; i < total; i++) + { + var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + } + + // MaxMsgs above the published count so a single pull buffers all messages. + var consumerOpts = new NatsJSConsumeOpts { MaxMsgs = 100, DrainOnCancel = true }; + var consumer = (NatsJSConsumer)await js.GetConsumerAsync($"{prefix}s1", $"{prefix}c1", cts.Token); + + var received = 0; + await foreach (var msg in consumer.ConsumeAsync(serializer: TestDataJsonSerializer.Default, consumerOpts, cancellationToken: cts.Token)) + { + // Simulate some processing time. + await Task.Delay(10, CancellationToken.None); + + // ACK after cancellation must still succeed since the connection stays open. + await msg.AckAsync(cancellationToken: CancellationToken.None); + received++; + + // Cancel once consumption is under way; drain must still deliver the + // remaining buffered/in-flight messages instead of dropping them. + if (received == 1) +#if NET8_0_OR_GREATER + await cts.CancelAsync(); +#else + cts.Cancel(); +#endif + } + + Assert.Equal(total, received); + + // Connection is left usable after the drain completes. + Assert.Equal(NatsConnectionState.Open, nats.ConnectionState); + await nats.PingAsync(CancellationToken.None); + var ack2 = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = total }, serializer: TestDataJsonSerializer.Default, cancellationToken: CancellationToken.None); + ack2.EnsureSuccess(); + } + + [Fact] + public async Task Consume_without_drain_on_cancel_drops_buffered_messages() + { + await using var nats = new NatsConnection(new NatsOpts { Url = _server.Url }); + await nats.ConnectRetryAsync(); + var prefix = _server.GetNextId(); + + var js = new NatsJSContext(nats); + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30)); + + await js.CreateStreamAsync($"{prefix}s1", new[] { $"{prefix}s1.*" }, cts.Token); + await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token); + + const int total = 100; + for (var i = 0; i < total; i++) + { + var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + } + + // DrainOnCancel defaults to false: cancelling stops promptly and does not + // deliver the messages still buffered in the consumer. + var consumerOpts = new NatsJSConsumeOpts { MaxMsgs = 100 }; + var consumer = (NatsJSConsumer)await js.GetConsumerAsync($"{prefix}s1", $"{prefix}c1", cts.Token); + + var received = 0; + await foreach (var msg in consumer.ConsumeAsync(serializer: TestDataJsonSerializer.Default, consumerOpts, cancellationToken: cts.Token)) + { + await msg.AckAsync(cancellationToken: CancellationToken.None); + received++; + + if (received == 1) +#if NET8_0_OR_GREATER + await cts.CancelAsync(); +#else + cts.Cancel(); +#endif + } + + // Without drain the loop stops promptly on cancel (fired after the first + // message) and abandons the buffered messages rather than delivering them. + Assert.True(received <= 5, $"expected a prompt stop (<= 5 messages), got {received} of {total}"); + } +} diff --git a/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs b/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs index 32ba7645c..a32f6bc11 100644 --- a/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs +++ b/tests/NATS.Net.OpenTelemetry.Tests/OpenTelemetryTest.cs @@ -298,9 +298,22 @@ public async Task Reconnect_counter() await nats.ReconnectAsync(); await openedTwice.Task.WaitAsync(TimeSpan.FromSeconds(30)); - var reconnects = meter.LongMeasurements - .Where(m => m.Name == "nats.client.reconnects") - .ToList(); + // The reconnects counter is incremented just after the ConnectionOpened event is + // queued, so the measurement can land slightly after the event we awaited. Poll + // for it rather than reading once and racing the increment. + List<(string Name, long Value, KeyValuePair[] Tags)> reconnects = new(); + var deadline = DateTime.UtcNow + TimeSpan.FromSeconds(10); + while (DateTime.UtcNow < deadline) + { + reconnects = meter.LongMeasurements + .Where(m => m.Name == "nats.client.reconnects") + .ToList(); + + if (reconnects.Count > 0) + break; + + await Task.Delay(25); + } reconnects.Sum(m => m.Value).Should().Be(1); @@ -539,6 +552,9 @@ public void AssertAllStopped() private sealed class MeterTracker : IDisposable { private readonly MeterListener _listener; + private readonly object _sync = new(); + private readonly List<(string Name, long Value, KeyValuePair[] Tags)> _longMeasurements = new(); + private readonly List<(string Name, double Value, KeyValuePair[] Tags)> _doubleMeasurements = new(); public MeterTracker() { @@ -551,17 +567,39 @@ public MeterTracker() }, }; + // Measurement callbacks fire on background threads (e.g. the reconnect loop), + // so guard the lists and hand out snapshots to readers. _listener.SetMeasurementEventCallback((inst, val, tags, _) => - LongMeasurements.Add((inst.Name, val, tags.ToArray()))); + { + lock (_sync) + _longMeasurements.Add((inst.Name, val, tags.ToArray())); + }); _listener.SetMeasurementEventCallback((inst, val, tags, _) => - DoubleMeasurements.Add((inst.Name, val, tags.ToArray()))); + { + lock (_sync) + _doubleMeasurements.Add((inst.Name, val, tags.ToArray())); + }); _listener.Start(); } - public List<(string Name, long Value, KeyValuePair[] Tags)> LongMeasurements { get; } = new(); + public IReadOnlyList<(string Name, long Value, KeyValuePair[] Tags)> LongMeasurements + { + get + { + lock (_sync) + return _longMeasurements.ToArray(); + } + } - public List<(string Name, double Value, KeyValuePair[] Tags)> DoubleMeasurements { get; } = new(); + public IReadOnlyList<(string Name, double Value, KeyValuePair[] Tags)> DoubleMeasurements + { + get + { + lock (_sync) + return _doubleMeasurements.ToArray(); + } + } public void Dispose() => _listener.Dispose(); } diff --git a/tests/xunit.runsettings b/tests/xunit.runsettings index b64f9a84c..6c4ab22e3 100644 --- a/tests/xunit.runsettings +++ b/tests/xunit.runsettings @@ -4,9 +4,6 @@ 600000 - - false + 1