diff --git a/src/NATS.Client.JetStream/NatsJSConsumer.cs b/src/NATS.Client.JetStream/NatsJSConsumer.cs index f6b5220d7..f59a86357 100644 --- a/src/NATS.Client.JetStream/NatsJSConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSConsumer.cs @@ -350,6 +350,7 @@ internal void SetPinId(string? pinId) internal async ValueTask> ConsumeInternalAsync(INatsDeserialize? serializer = default, NatsJSConsumeOpts? opts = default, CancellationToken cancellationToken = default) { ThrowIfDeleted(); + cancellationToken.ThrowIfCancellationRequested(); opts ??= new NatsJSConsumeOpts(); serializer ??= _context.Connection.Opts.SerializerRegistry.GetDeserializer(); @@ -457,6 +458,7 @@ internal async ValueTask> FetchInternalAsync( CancellationToken cancellationToken = default) { ThrowIfDeleted(); + cancellationToken.ThrowIfCancellationRequested(); serializer ??= _context.Connection.Opts.SerializerRegistry.GetDeserializer(); var inbox = _context.NewBaseInbox(); diff --git a/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs b/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs index 5dc06e599..e8d51e3d5 100644 --- a/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs +++ b/src/NATS.Client.JetStream/NatsJSOrderedConsumer.cs @@ -64,6 +64,7 @@ public async IAsyncEnumerable> ConsumeAsync( 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> FetchAsync( INatsDeserialize? 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> FetchNoWaitAsync( INatsDeserialize? serializer = default, [EnumeratorCancellation] CancellationToken cancellationToken = default) { + cancellationToken.ThrowIfCancellationRequested(); + cancellationToken = CancellationTokenSource.CreateLinkedTokenSource(_cancellationToken, cancellationToken).Token; var processed = 0; var bytesProcessed = 0; diff --git a/tests/NATS.Client.JetStream.Tests/CancellationTokenTests.cs b/tests/NATS.Client.JetStream.Tests/CancellationTokenTests.cs new file mode 100644 index 000000000..771ad7dc1 --- /dev/null +++ b/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.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + } + + using var cancelledCts = new CancellationTokenSource(); + cancelledCts.Cancel(); + + await Assert.ThrowsAsync(async () => + { + await foreach (var unused in consumer.FetchAsync(new NatsJSFetchOpts { MaxMsgs = 10 }, serializer: TestDataJsonSerializer.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.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + } + + using var cancelledCts = new CancellationTokenSource(); + cancelledCts.Cancel(); + + await Assert.ThrowsAsync(async () => + { + await foreach (var unused in consumer.FetchNoWaitAsync(new NatsJSFetchOpts { MaxMsgs = 5 }, serializer: TestDataJsonSerializer.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.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + + using var cancelledCts = new CancellationTokenSource(); + cancelledCts.Cancel(); + + await Assert.ThrowsAsync(async () => + { + await consumer.NextAsync(serializer: TestDataJsonSerializer.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.Default, cancellationToken: cts.Token); + ack.EnsureSuccess(); + } + + using var cancelledCts = new CancellationTokenSource(); + cancelledCts.Cancel(); + + await Assert.ThrowsAsync(async () => + { + await foreach (var unused in consumer.ConsumeAsync(serializer: TestDataJsonSerializer.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(async () => + { + await foreach (var unused in consumer.FetchAsync(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(async () => + { + await foreach (var unused in consumer.FetchNoWaitAsync(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(async () => + { + await foreach (var unused in consumer.ConsumeAsync(cancellationToken: cancelledCts.Token)) + { + } + }); + } +}