Fix PingCommand cancellation

src/NATS.Client.Core/Commands/PingCommand.cs

@@ -10,6 +10,7 @@ internal class PingCommand : IValueTaskSource<TimeSpan>, IObjectPoolNode<PingCom
     private Stopwatch _stopwatch;
     private ManualResetValueTaskSourceCore<TimeSpan> _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<TimeSpan> RunAsync() => new(this, _core.Version);

src/NATS.Client.Core/NatsConnection.Ping.cs

@@ -24,6 +24,12 @@ public async ValueTask<TimeSpan> 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);

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<OperationCanceledException>();
+    }
+
+    [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<TimeSpan>[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<OperationCanceledException>();
+    }
+}

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<OperationCanceledException>();
+    }
+
+    [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<OperationCanceledException>();
+    }
+
+    [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
+            }
+        }
+    }
+}

tests/NATS.Client.TestUtilities/MockServer.cs

@@ -17,13 +17,16 @@ public class MockServer : IAsyncDisposable
     private readonly List<Task> _clients = new();
     private readonly Task _accept;
     private readonly CancellationTokenSource _cts;
+    private readonly bool _autoPong;
 
     public MockServer(
         Func<Client, Cmd, Task> handler,
         Action<string>? 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"))
                             {