Add cancelled token handling for consumers

src/NATS.Client.JetStream/NatsJSConsumer.cs

@@ -350,6 +350,7 @@ internal void SetPinId(string? pinId)
     internal async ValueTask<NatsJSConsume<T>> ConsumeInternalAsync<T>(INatsDeserialize<T>? serializer = default, NatsJSConsumeOpts? opts = default, CancellationToken cancellationToken = default)
     {
         ThrowIfDeleted();
+        cancellationToken.ThrowIfCancellationRequested();
 
         opts ??= new NatsJSConsumeOpts();
         serializer ??= _context.Connection.Opts.SerializerRegistry.GetDeserializer<T>();
@@ -457,6 +458,7 @@ internal async ValueTask<NatsJSFetch<T>> FetchInternalAsync<T>(
         CancellationToken cancellationToken = default)
     {
         ThrowIfDeleted();
+        cancellationToken.ThrowIfCancellationRequested();
         serializer ??= _context.Connection.Opts.SerializerRegistry.GetDeserializer<T>();
 
         var inbox = _context.NewBaseInbox();

src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs

@@ -64,6 +64,7 @@ public async IAsyncEnumerable<INatsJSMsg<T>> ConsumeAsync<T>(
         NatsJSConsumeOpts? opts = default,
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
+        cancellationToken.ThrowIfCancellationRequested();
         opts ??= _context.Opts.DefaultConsumeOpts;
         var consumerName = string.Empty;
         var notificationHandler = opts.NotificationHandler;
@@ -198,6 +199,8 @@ public async IAsyncEnumerable<INatsJSMsg<T>> FetchAsync<T>(
         INatsDeserialize<T>? serializer = default,
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
+        cancellationToken.ThrowIfCancellationRequested();
+
         cancellationToken = CancellationTokenSource.CreateLinkedTokenSource(_cancellationToken, cancellationToken).Token;
         var processed = 0;
         var bytesProcessed = 0;
@@ -266,6 +269,8 @@ public async IAsyncEnumerable<INatsJSMsg<T>> FetchNoWaitAsync<T>(
         INatsDeserialize<T>? serializer = default,
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
+        cancellationToken.ThrowIfCancellationRequested();
+
         cancellationToken = CancellationTokenSource.CreateLinkedTokenSource(_cancellationToken, cancellationToken).Token;
         var processed = 0;
         var bytesProcessed = 0;

tests/NATS.Client.JetStream.Tests/CancellationTokenTests.cs

@@ -0,0 +1,212 @@
+using NATS.Client.Core2.Tests;
+using NATS.Client.TestUtilities2;
+
+// ReSharper disable MethodHasAsyncOverload
+namespace NATS.Client.JetStream.Tests;
+
+[Collection("nats-server")]
+public class CancellationTokenTests(NatsServerFixture server)
+{
+    [Fact]
+    public async Task FetchAsync_with_cancelled_token_throws_immediately()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await using var nats = new NatsConnection(new NatsOpts { Url = server.Url });
+        var prefix = server.GetNextId();
+        var js = new NatsJSContext(nats);
+        await js.CreateStreamAsync($"{prefix}s1", [$"{prefix}s1.*"], cts.Token);
+        var consumer = (NatsJSConsumer)await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token);
+
+        for (var i = 0; i < 10; i++)
+        {
+            var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cts.Token);
+            ack.EnsureSuccess();
+        }
+
+        using var cancelledCts = new CancellationTokenSource();
+        cancelledCts.Cancel();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (var unused in consumer.FetchAsync(new NatsJSFetchOpts { MaxMsgs = 10 }, serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cancelledCts.Token))
+            {
+            }
+        });
+
+        // Verify no messages were reserved as pending
+        await consumer.RefreshAsync(cts.Token);
+        Assert.Equal(0, consumer.Info.NumAckPending);
+    }
+
+    [Fact]
+    public async Task FetchNoWaitAsync_with_cancelled_token_throws_immediately()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await using var nats = new NatsConnection(new NatsOpts { Url = server.Url });
+        var prefix = server.GetNextId();
+        var js = new NatsJSContext(nats);
+        await js.CreateStreamAsync($"{prefix}s1", [$"{prefix}s1.*"], cts.Token);
+        var consumer = (NatsJSConsumer)await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token);
+
+        for (var i = 0; i < 5; i++)
+        {
+            var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cts.Token);
+            ack.EnsureSuccess();
+        }
+
+        using var cancelledCts = new CancellationTokenSource();
+        cancelledCts.Cancel();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (var unused in consumer.FetchNoWaitAsync(new NatsJSFetchOpts { MaxMsgs = 5 }, serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cancelledCts.Token))
+            {
+            }
+        });
+
+        // Verify no messages were reserved as pending
+        await consumer.RefreshAsync(cts.Token);
+        Assert.Equal(0, consumer.Info.NumAckPending);
+    }
+
+    [Fact]
+    public async Task NextAsync_with_cancelled_token_throws_immediately()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await using var nats = new NatsConnection(new NatsOpts { Url = server.Url });
+        var prefix = server.GetNextId();
+        var js = new NatsJSContext(nats);
+        await js.CreateStreamAsync($"{prefix}s1", [$"{prefix}s1.*"], cts.Token);
+        var consumer = (NatsJSConsumer)await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token);
+
+        var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = 1 }, serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cts.Token);
+        ack.EnsureSuccess();
+
+        using var cancelledCts = new CancellationTokenSource();
+        cancelledCts.Cancel();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await consumer.NextAsync(serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cancelledCts.Token);
+        });
+
+        // Verify no messages were reserved as pending
+        await consumer.RefreshAsync(cts.Token);
+        Assert.Equal(0, consumer.Info.NumAckPending);
+    }
+
+    [Fact]
+    public async Task ConsumeAsync_with_cancelled_token_throws_immediately()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await using var nats = new NatsConnection(new NatsOpts { Url = server.Url });
+        var prefix = server.GetNextId();
+        var js = new NatsJSContext(nats);
+        await js.CreateStreamAsync($"{prefix}s1", [$"{prefix}s1.*"], cts.Token);
+        var consumer = (NatsJSConsumer)await js.CreateOrUpdateConsumerAsync($"{prefix}s1", $"{prefix}c1", cancellationToken: cts.Token);
+
+        for (var i = 0; i < 5; i++)
+        {
+            var ack = await js.PublishAsync($"{prefix}s1.foo", new TestData { Test = i }, serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cts.Token);
+            ack.EnsureSuccess();
+        }
+
+        using var cancelledCts = new CancellationTokenSource();
+        cancelledCts.Cancel();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (var unused in consumer.ConsumeAsync(serializer: TestDataJsonSerializer<TestData>.Default, cancellationToken: cancelledCts.Token))
+            {
+            }
+        });
+
+        // Verify no messages were reserved as pending
+        await consumer.RefreshAsync(cts.Token);
+        Assert.Equal(0, consumer.Info.NumAckPending);
+    }
+
+    [Fact]
+    public async Task OrderedConsumer_FetchAsync_with_cancelled_token_throws_immediately()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await using var nats = server.CreateNatsConnection();
+        await nats.ConnectRetryAsync();
+        var prefix = server.GetNextId();
+        var js = new NatsJSContext(nats);
+        var stream = await js.CreateStreamAsync($"{prefix}s1", [$"{prefix}s1.*"], cts.Token);
+
+        for (var i = 0; i < 5; i++)
+        {
+            await js.PublishAsync($"{prefix}s1.foo", i, cancellationToken: cts.Token);
+        }
+
+        var consumer = await stream.CreateOrderedConsumerAsync(cancellationToken: cts.Token);
+
+        using var cancelledCts = new CancellationTokenSource();
+        cancelledCts.Cancel();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (var unused in consumer.FetchAsync<int>(new NatsJSFetchOpts { MaxMsgs = 5 }, cancellationToken: cancelledCts.Token))
+            {
+            }
+        });
+    }
+
+    [Fact]
+    public async Task OrderedConsumer_FetchNoWaitAsync_with_cancelled_token_throws_immediately()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await using var nats = server.CreateNatsConnection();
+        await nats.ConnectRetryAsync();
+        var prefix = server.GetNextId();
+        var js = new NatsJSContext(nats);
+        var stream = await js.CreateStreamAsync($"{prefix}s1", [$"{prefix}s1.*"], cts.Token);
+
+        for (var i = 0; i < 5; i++)
+        {
+            await js.PublishAsync($"{prefix}s1.foo", i, cancellationToken: cts.Token);
+        }
+
+        var consumer = await stream.CreateOrderedConsumerAsync(cancellationToken: cts.Token);
+
+        using var cancelledCts = new CancellationTokenSource();
+        cancelledCts.Cancel();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (var unused in consumer.FetchNoWaitAsync<int>(new NatsJSFetchOpts { MaxMsgs = 5 }, cancellationToken: cancelledCts.Token))
+            {
+            }
+        });
+    }
+
+    [Fact]
+    public async Task OrderedConsumer_ConsumeAsync_with_cancelled_token_throws_immediately()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await using var nats = server.CreateNatsConnection();
+        await nats.ConnectRetryAsync();
+        var prefix = server.GetNextId();
+        var js = new NatsJSContext(nats);
+        var stream = await js.CreateStreamAsync($"{prefix}s1", [$"{prefix}s1.*"], cts.Token);
+
+        for (var i = 0; i < 5; i++)
+        {
+            await js.PublishAsync($"{prefix}s1.foo", i, cancellationToken: cts.Token);
+        }
+
+        var consumer = await stream.CreateOrderedConsumerAsync(cancellationToken: cts.Token);
+
+        using var cancelledCts = new CancellationTokenSource();
+        cancelledCts.Cancel();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (var unused in consumer.ConsumeAsync<int>(cancellationToken: cancelledCts.Token))
+            {
+            }
+        });
+    }
+}