diff --git a/src/NATS.Client.Core/NatsConnection.Publish.cs b/src/NATS.Client.Core/NatsConnection.Publish.cs index 1f5312413..0b6a91161 100644 --- a/src/NATS.Client.Core/NatsConnection.Publish.cs +++ b/src/NATS.Client.Core/NatsConnection.Publish.cs @@ -67,7 +67,17 @@ public ValueTask PublishAsync(in NatsMsg msg, INatsSerialize? serialize private async ValueTask ConnectAndPublishAsync(string subject, T? data, NatsHeaders? headers, string? replyTo, INatsSerialize serializer, CancellationToken cancellationToken) { - await ConnectAsync().AsTask().WaitAsync(cancellationToken).ConfigureAwait(false); + if (Opts.PublishTimeoutOnDisconnected) + { + using var cts1 = new CancellationTokenSource(Opts.CommandTimeout); + using var cts2 = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, cts1.Token); + await ConnectAsync().AsTask().WaitAsync(cts2.Token).ConfigureAwait(false); + } + else + { + await ConnectAsync().AsTask().WaitAsync(cancellationToken).ConfigureAwait(false); + } + await CommandWriter.PublishAsync(subject, data, headers, replyTo, serializer, cancellationToken).ConfigureAwait(false); } } diff --git a/src/NATS.Client.Core/NatsOpts.cs b/src/NATS.Client.Core/NatsOpts.cs index f9f755c86..910ed4afc 100644 --- a/src/NATS.Client.Core/NatsOpts.cs +++ b/src/NATS.Client.Core/NatsOpts.cs @@ -206,6 +206,14 @@ public sealed record NatsOpts /// public bool RetryOnInitialConnect { get; init; } + /// + /// Gets or sets a value indicating whether publish would throw an exception + /// when the connection is disconnected and the is reached. + /// The default is false, meaning publish will not throw on disconnected state + /// and will wait to publish the message until reconnected. + /// + public bool PublishTimeoutOnDisconnected { get; init; } = false; + internal NatsUri[] GetSeedUris(bool suppressRandomization = false) { var urls = Url.Split(','); diff --git a/tests/NATS.Client.Core2.Tests/CancellationTest.cs b/tests/NATS.Client.Core2.Tests/CancellationTest.cs index 7b2673c13..90786a73b 100644 --- a/tests/NATS.Client.Core2.Tests/CancellationTest.cs +++ b/tests/NATS.Client.Core2.Tests/CancellationTest.cs @@ -89,4 +89,93 @@ await Assert.ThrowsAsync(async () => } }); } + + [Fact] + public async Task PublishTimeoutOnDisconnected_WhenEnabled_ThrowsOnTimeout() + { + var server = await NatsServerProcess.StartAsync(); + + await using var conn = new NatsConnection(new NatsOpts + { + Url = server.Url, + CommandTimeout = TimeSpan.FromMilliseconds(500), + PublishTimeoutOnDisconnected = true, + }); + await conn.ConnectAsync(); + + // Kill the server + await server.DisposeAsync(); + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + + // Wait for reconnect loop to kick in + while (conn.ConnectionState != NatsConnectionState.Reconnecting) + { + await Task.Delay(1, cts.Token); + } + + // PublishAsync should throw TaskCanceledException due to CommandTimeout + await Assert.ThrowsAsync(() => conn.PublishAsync("test", "data").AsTask()); + } + + [Fact] + public async Task PublishTimeoutOnDisconnected_WhenDisabled_WaitsIndefinitely() + { + var server = await NatsServerProcess.StartAsync(); + + await using var conn = new NatsConnection(new NatsOpts + { + Url = server.Url, + PublishTimeoutOnDisconnected = false, // default + }); + await conn.ConnectAsync(); + + // Kill the server + await server.DisposeAsync(); + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + + // Wait for reconnect loop to kick in + while (conn.ConnectionState != NatsConnectionState.Reconnecting) + { + await Task.Delay(1, cts.Token); + } + + // With option disabled, publish waits indefinitely - user cancellation should work + var userCts = new CancellationTokenSource(TimeSpan.FromMilliseconds(500)); + + // Should throw TaskCanceledException (user-initiated cancellation) + await Assert.ThrowsAsync(() => conn.PublishAsync("test", "data", cancellationToken: userCts.Token).AsTask()); + } + + [Fact] + public async Task PublishTimeoutOnDisconnected_UserCancellation_StillWorks() + { + var server = await NatsServerProcess.StartAsync(); + + await using var conn = new NatsConnection(new NatsOpts + { + Url = server.Url, + CommandTimeout = TimeSpan.FromSeconds(30), // Long timeout + PublishTimeoutOnDisconnected = true, + }); + await conn.ConnectAsync(); + + // Kill the server + await server.DisposeAsync(); + + var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + + // Wait for reconnect loop to kick in + while (conn.ConnectionState != NatsConnectionState.Reconnecting) + { + await Task.Delay(1, cts.Token); + } + + // User cancellation should still work even with PublishTimeoutOnDisconnected enabled + var userCts = new CancellationTokenSource(TimeSpan.FromMilliseconds(500)); + + // Should throw TaskCanceledException (user-initiated cancellation) before CommandTimeout + await Assert.ThrowsAsync(() => conn.PublishAsync("test", "data", cancellationToken: userCts.Token).AsTask()); + } }