diff --git a/src/NATS.Client.Core/Commands/PingCommand.cs b/src/NATS.Client.Core/Commands/PingCommand.cs index 6e6fa6744..966cd5b11 100644 --- a/src/NATS.Client.Core/Commands/PingCommand.cs +++ b/src/NATS.Client.Core/Commands/PingCommand.cs @@ -10,6 +10,7 @@ internal class PingCommand : IValueTaskSource, IObjectPoolNode _core; private PingCommand? _next; + private int _completed; public PingCommand(ObjectPool? pool) { @@ -25,14 +26,27 @@ public PingCommand(ObjectPool? pool) public void Start() => _stopwatch.Restart(); - public void SetResult() => _core.SetResult(_stopwatch.Elapsed); + public void SetResult() + { + if (Interlocked.CompareExchange(ref _completed, 1, 0) == 0) + { + _core.SetResult(_stopwatch.Elapsed); + } + } - public void SetCanceled() => _core.SetException(new OperationCanceledException()); + public void SetCanceled() + { + if (Interlocked.CompareExchange(ref _completed, 1, 0) == 0) + { + _core.SetException(new OperationCanceledException()); + } + } public void Reset() { _stopwatch.Reset(); _core.Reset(); + Volatile.Write(ref _completed, 0); } public ValueTask RunAsync() => new(this, _core.Version); diff --git a/src/NATS.Client.Core/NatsConnection.Ping.cs b/src/NATS.Client.Core/NatsConnection.Ping.cs index 86f31b68d..838233829 100644 --- a/src/NATS.Client.Core/NatsConnection.Ping.cs +++ b/src/NATS.Client.Core/NatsConnection.Ping.cs @@ -24,6 +24,12 @@ public async ValueTask PingAsync(CancellationToken cancellationToken = pingCommand.Start(); +#if NETSTANDARD + using var registration = cancellationToken.Register(static state => ((PingCommand)state!).SetCanceled(), pingCommand); +#else + await using var registration = cancellationToken.UnsafeRegister(static state => ((PingCommand)state!).SetCanceled(), pingCommand); +#endif + await CommandWriter.PingAsync(pingCommand, cancellationToken).ConfigureAwait(false); return await pingCommand.RunAsync().ConfigureAwait(false); diff --git a/tests/NATS.Client.Core2.Tests/PingCancellationTest.cs b/tests/NATS.Client.Core2.Tests/PingCancellationTest.cs new file mode 100644 index 000000000..459c1eeac --- /dev/null +++ b/tests/NATS.Client.Core2.Tests/PingCancellationTest.cs @@ -0,0 +1,150 @@ +using NATS.Client.TestUtilities; +using Synadia.Orbit.Testing.NatsServerProcessManager; + +namespace NATS.Client.Core.Tests; + +public class PingCancellationTest +{ + private readonly ITestOutputHelper _output; + + public PingCancellationTest(ITestOutputHelper output) => _output = output; + + [Fact] + public async Task PingAsync_succeeds_with_mock_server() + { + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + + await using var server = new MockServer( + handler: (_, _) => Task.CompletedTask, + logger: m => _output.WriteLine(m), + cancellationToken: cts.Token); + + await using var nats = new NatsConnection(new NatsOpts { Url = server.Url }); + await nats.ConnectAsync(); + + var rtt = await nats.PingAsync(cts.Token); + rtt.Should().BeGreaterThan(TimeSpan.Zero); + } + + [Fact] + public async Task PingAsync_multiple_sequential_pings_succeed() + { + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + + await using var server = new MockServer( + handler: (_, _) => Task.CompletedTask, + logger: m => _output.WriteLine(m), + cancellationToken: cts.Token); + + await using var nats = new NatsConnection(new NatsOpts { Url = server.Url }); + await nats.ConnectAsync(); + + for (var i = 0; i < 5; i++) + { + var rtt = await nats.PingAsync(cts.Token); + rtt.Should().BeGreaterThan(TimeSpan.Zero); + } + } + + [Fact] + public async Task PingAsync_throws_when_cancelled_waiting_for_pong() + { + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + var pingCount = 0; + + // autoPong: false — handler replies PONG only for the first PING (the connect handshake) + await using var server = new MockServer( + handler: async (client, cmd) => + { + if (cmd.Name == "PING") + { + var n = Interlocked.Increment(ref pingCount); + client.Log($"[S] PING #{n}"); + + if (n == 1) + { + // Reply to the connect-time PING so the connection opens + await client.Writer.WriteAsync("PONG\r\n"); + await client.Writer.FlushAsync(); + } + + // Subsequent PINGs get no PONG + } + }, + logger: m => _output.WriteLine(m), + autoPong: false, + cancellationToken: cts.Token); + + await using var nats = new NatsConnection(new NatsOpts { Url = server.Url }); + await nats.ConnectAsync(); + + // This ping should time out because the server won't reply PONG + using var pingCts = new CancellationTokenSource(TimeSpan.FromMilliseconds(500)); + + var act = () => nats.PingAsync(pingCts.Token).AsTask(); + await act.Should().ThrowAsync(); + } + + [Fact] + public async Task PingAsync_concurrent_pings_all_complete() + { + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + + await using var server = new MockServer( + handler: (_, _) => Task.CompletedTask, + logger: m => _output.WriteLine(m), + cancellationToken: cts.Token); + + await using var nats = new NatsConnection(new NatsOpts { Url = server.Url }); + await nats.ConnectAsync(); + + // Fire multiple pings concurrently — exercises the pool and concurrent SetResult paths + var tasks = new Task[10]; + for (var i = 0; i < tasks.Length; i++) + { + tasks[i] = nats.PingAsync(cts.Token).AsTask(); + } + + var results = await Task.WhenAll(tasks); + foreach (var rtt in results) + { + rtt.Should().BeGreaterOrEqualTo(TimeSpan.Zero); + } + } + + [Fact] + public async Task PingAsync_succeeds_against_real_server() + { + await using var server = await NatsServerProcess.StartAsync(); + + await using var nats = new NatsConnection(new NatsOpts { Url = server.Url }); + await nats.ConnectAsync(); + + var rtt = await nats.PingAsync(); + rtt.Should().BeGreaterThan(TimeSpan.Zero); + } + + [Fact] + public async Task PingAsync_times_out_after_server_stopped() + { + await using var server = await NatsServerProcess.StartAsync(); + + await using var nats = new NatsConnection(new NatsOpts { Url = server.Url }); + await nats.ConnectAsync(); + + // Verify ping works while server is up + var rtt = await nats.PingAsync(); + rtt.Should().BeGreaterThan(TimeSpan.Zero); + _output.WriteLine($"Ping RTT before stop: {rtt}"); + + // Stop the server — no more PONGs + await server.StopAsync(); + _output.WriteLine("Server stopped"); + + // Ping with a timeout should throw since no PONG will arrive + using var pingCts = new CancellationTokenSource(TimeSpan.FromSeconds(2)); + + var act = () => nats.PingAsync(pingCts.Token).AsTask(); + await act.Should().ThrowAsync(); + } +} diff --git a/tests/NATS.Client.CoreUnit.Tests/PingCommandTest.cs b/tests/NATS.Client.CoreUnit.Tests/PingCommandTest.cs new file mode 100644 index 000000000..fd63aab09 --- /dev/null +++ b/tests/NATS.Client.CoreUnit.Tests/PingCommandTest.cs @@ -0,0 +1,94 @@ +namespace NATS.Client.Core.Tests; + +public class PingCommandTest +{ + [Fact] + public async Task SetResult_DeliversElapsedTime() + { + var cmd = new PingCommand(pool: null); + cmd.Start(); + await Task.Delay(50); + cmd.SetResult(); + + var elapsed = await cmd.RunAsync(); + elapsed.TotalMilliseconds.Should().BeGreaterThan(0); + } + + [Fact] + public async Task SetCanceled_ThrowsOperationCanceledException() + { + var cmd = new PingCommand(pool: null); + cmd.Start(); + cmd.SetCanceled(); + + var act = () => cmd.RunAsync().AsTask(); + await act.Should().ThrowAsync(); + } + + [Fact] + public void SetResult_CalledTwice_DoesNotThrow() + { + var cmd = new PingCommand(pool: null); + cmd.Start(); + cmd.SetResult(); + cmd.SetResult(); // second call should be a no-op + } + + [Fact] + public void SetCanceled_CalledTwice_DoesNotThrow() + { + var cmd = new PingCommand(pool: null); + cmd.Start(); + cmd.SetCanceled(); + cmd.SetCanceled(); // second call should be a no-op + } + + [Fact] + public async Task SetResult_AfterSetCanceled_StillThrows() + { + var cmd = new PingCommand(pool: null); + cmd.Start(); + cmd.SetCanceled(); + cmd.SetResult(); // should be ignored + + var act = () => cmd.RunAsync().AsTask(); + await act.Should().ThrowAsync(); + } + + [Fact] + public async Task SetCanceled_AfterSetResult_StillReturnsResult() + { + var cmd = new PingCommand(pool: null); + cmd.Start(); + cmd.SetResult(); + cmd.SetCanceled(); // should be ignored + + var elapsed = await cmd.RunAsync(); + elapsed.TotalMilliseconds.Should().BeGreaterOrEqualTo(0); + } + + [Fact] + public async Task ConcurrentSetResultAndSetCanceled_DoesNotThrow() + { + // Run multiple times to increase chance of hitting the race + for (var i = 0; i < 100; i++) + { + var cmd = new PingCommand(pool: null); + cmd.Start(); + + var t1 = Task.Run(() => cmd.SetResult()); + var t2 = Task.Run(() => cmd.SetCanceled()); + await Task.WhenAll(t1, t2); + + // Should not throw - either result or cancellation wins + try + { + await cmd.RunAsync(); + } + catch (OperationCanceledException) + { + // This is also acceptable + } + } + } +} diff --git a/tests/NATS.Client.TestUtilities/MockServer.cs b/tests/NATS.Client.TestUtilities/MockServer.cs index d4968b632..4bde4af8a 100644 --- a/tests/NATS.Client.TestUtilities/MockServer.cs +++ b/tests/NATS.Client.TestUtilities/MockServer.cs @@ -17,13 +17,16 @@ public class MockServer : IAsyncDisposable private readonly List _clients = new(); private readonly Task _accept; private readonly CancellationTokenSource _cts; + private readonly bool _autoPong; public MockServer( Func handler, Action? logger = null, string info = "{\"max_payload\":1048576}", + bool autoPong = true, CancellationToken cancellationToken = default) { + _autoPong = autoPong; _logger = logger ?? (_ => { }); _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); cancellationToken = _cts.Token; @@ -79,8 +82,13 @@ public MockServer( { // B: PING␍␊ // B: PONG␍␊ - await sw.WriteAsync("PONG\r\n"); - await sw.FlushAsync(); + if (_autoPong) + { + await sw.WriteAsync("PONG\r\n"); + await sw.FlushAsync(); + } + + await handler(client, new Cmd("PING", string.Empty, null, 0, 0, null, string.Empty, client)); } else if (line.StartsWith("SUB")) {