Fixed races in cancellation tests by switching from timer-based to flag-based with retries.

Author: RyanLamansky

docs/claude/tds-endpoint.md

@@ -429,6 +429,11 @@ Batch-scoped variables go with the ended batch either way; connection-scoped tem
 `SimulatedDbCommand.Cancel()` routes to the same machinery, so an in-process `Cancel()` from another thread interrupts a running `WAITFOR` / long batch identically (the reader then drains already-materialized rows, nothing left in flight — the documented in-process reaction bound).
 Oracle: `AttentionTests` (Tests.SqlClient), `WaitForDelayTests.Delay_InterruptedByInProcessCancel_ReturnsPromptly` (Tests).
 
+**Partial response packets don't flush early (probed 2026-07-20).**
+The server accumulates response tokens into a TDS packet and sends it only when the packet fills or the response ends — real SQL Server behaves identically: for `select @p; waitfor delay '00:00:30'` the one-row first result set sits in the send buffer for the full 30 seconds and the client's first `ReadAsync` blocks until the batch ends, on real and sim alike.
+Two consequences: a client can't observe an early small result set to sequence a cancel after (the cancel must land while the client read is parked — `AttentionTests` retries via `CancelUntilComplete` for exactly this reason), and the attention path is the only way a blocked mid-batch client wakes early.
+A `SqlCommand.Cancel()` racing execution *start* is different: a cancel before the batch begins executing is a documented client-side no-op, and one landing inside `Execute*Async`'s setup surfaces SqlClient's own `InvalidOperationException` instead of the cancel `SqlException` — client-side races independent of the server, so tests never fire a single timer-based cancel.
+
 ## MARS (Multiple Active Result Sets)
 
 `MultipleActiveResultSets=True` lets a client run a second command while a reader is still open — the EF-lazy-load shape (iterate a parent query, touch a navigation per row).

tests/SqlServerSimulator.Tests.SqlClient/AttentionTests.cs

@@ -21,13 +21,22 @@ public sealed class AttentionTests
 
     private const string LongWait = "waitfor delay '00:00:30'";
 
-    /// <summary>Cancels the command shortly after execution starts, off the calling thread.</summary>
-    private static void CancelAfter(SqlCommand command, int millis) =>
-        _ = Task.Run(async () =>
+    /// <summary>
+    /// Cancels the command repeatedly until the in-flight execution completes.
+    /// A cancel landing before the batch starts executing is a documented
+    /// <see cref="SqlCommand.Cancel"/> no-op, so a single timer-fired cancel
+    /// can miss on a stalled runner and let the 30-second WAITFOR run to
+    /// natural completion; retrying until the task transitions guarantees an
+    /// attention lands mid-execution.
+    /// </summary>
+    private static async Task CancelUntilComplete(SqlCommand command, Task execution, CancellationToken cancellationToken)
+    {
+        while (!execution.IsCompleted)
         {
-            await Task.Delay(millis);
+            await Task.Delay(100, cancellationToken);
             command.Cancel();
-        });
+        }
+    }
 
     private static async Task AssertSessionReusableAsync(SqlConnection connection, CancellationToken cancellationToken)
     {
@@ -61,9 +70,9 @@ public async Task Cancel_DuringWaitfor_RaisesCancelAndSessionStaysReusable()
 
         await using (var command = new SqlCommand(LongWait, connection) { CommandTimeout = 0 })
         {
-            CancelAfter(command, 200);
-            var error = await ThrowsExactlyAsync<SqlException>(
-                async () => await command.ExecuteNonQueryAsync(TestContext.CancellationToken));
+            var execution = command.ExecuteNonQueryAsync(TestContext.CancellationToken);
+            await CancelUntilComplete(command, execution, TestContext.CancellationToken);
+            var error = await ThrowsExactlyAsync<SqlException>(async () => await execution);
             AreEqual(0, error.Number);
         }
 
@@ -80,10 +89,18 @@ public async Task Cancel_MidLargeResult_AbortsAndSessionStaysReusable()
         const string bigResult = "select value, replicate('x', 200) from generate_series(1, 500000)";
         await using (var command = new SqlCommand(bigResult, connection) { CommandTimeout = 0 })
         {
-            CancelAfter(command, 50);
+            // Cancel only after the first row arrives, so the command is
+            // provably mid-execution. A timer-based cancel races two ways on a
+            // slow runner: landing before execution starts is a documented
+            // Cancel() no-op (the drain then completes with no exception), and
+            // landing during ExecuteReaderAsync's setup trips a client-side
+            // SqlClient race that surfaces InvalidOperationException instead
+            // of the cancel SqlException.
             _ = await ThrowsExactlyAsync<SqlException>(async () =>
             {
                 await using var reader = await command.ExecuteReaderAsync(TestContext.CancellationToken);
+                IsTrue(await reader.ReadAsync(TestContext.CancellationToken));
+                command.Cancel();
                 while (await reader.ReadAsync(TestContext.CancellationToken))
                 {
                 }
@@ -106,9 +123,9 @@ public async Task Cancel_WithOpenTransaction_DefaultXactAbort_LeavesTransactionO
 
         await using (var command = new SqlCommand(LongWait, connection) { CommandTimeout = 0 })
         {
-            CancelAfter(command, 200);
-            _ = await ThrowsExactlyAsync<SqlException>(
-                async () => await command.ExecuteNonQueryAsync(TestContext.CancellationToken));
+            var execution = command.ExecuteNonQueryAsync(TestContext.CancellationToken);
+            await CancelUntilComplete(command, execution, TestContext.CancellationToken);
+            _ = await ThrowsExactlyAsync<SqlException>(async () => await execution);
         }
 
         // XACT_ABORT OFF (the default): the transaction survives the cancel.
@@ -136,9 +153,9 @@ public async Task Cancel_WithOpenTransaction_XactAbortOn_RollsTransactionBack()
 
         await using (var command = new SqlCommand(LongWait, connection) { CommandTimeout = 0 })
         {
-            CancelAfter(command, 200);
-            _ = await ThrowsExactlyAsync<SqlException>(
-                async () => await command.ExecuteNonQueryAsync(TestContext.CancellationToken));
+            var execution = command.ExecuteNonQueryAsync(TestContext.CancellationToken);
+            await CancelUntilComplete(command, execution, TestContext.CancellationToken);
+            _ = await ThrowsExactlyAsync<SqlException>(async () => await execution);
         }
 
         // XACT_ABORT ON: the cancel rolls the transaction back.
@@ -162,9 +179,9 @@ public async Task Cancel_MidBatch_DiscardsRemainingStatements()
         var batchText = $"insert t values (1); {LongWait}; insert t values (2)";
         await using (var batch = new SqlCommand(batchText, connection) { CommandTimeout = 0 })
         {
-            CancelAfter(batch, 200);
-            _ = await ThrowsExactlyAsync<SqlException>(
-                async () => await batch.ExecuteNonQueryAsync(TestContext.CancellationToken));
+            var execution = batch.ExecuteNonQueryAsync(TestContext.CancellationToken);
+            await CancelUntilComplete(batch, execution, TestContext.CancellationToken);
+            _ = await ThrowsExactlyAsync<SqlException>(async () => await execution);
         }
 
         await using var rows = new SqlCommand("select count(*) from t", connection);
@@ -181,18 +198,26 @@ public async Task Cancel_DuringParameterizedRpc_AbortsAndSessionStaysReusable()
         await using (var command = new SqlCommand($"select @p; {LongWait}", connection) { CommandTimeout = 0 })
         {
             _ = command.Parameters.AddWithValue("@p", 7);
-            CancelAfter(command, 200);
-            _ = await ThrowsExactlyAsync<SqlException>(async () =>
+            // The drain blocks inside the first ReadAsync: the server holds
+            // the tiny first result set in its TDS send buffer until the
+            // batch ends (probe-confirmed against real SQL Server), so a
+            // cancel can't be sequenced after an observed row — it has to
+            // land while the read is parked, retried until it takes.
+            var drain = DrainAsync(command, TestContext.CancellationToken);
+            await CancelUntilComplete(command, drain, TestContext.CancellationToken);
+            _ = await ThrowsExactlyAsync<SqlException>(async () => await drain);
+
+            static async Task DrainAsync(SqlCommand command, CancellationToken cancellationToken)
             {
-                await using var reader = await command.ExecuteReaderAsync(TestContext.CancellationToken);
+                await using var reader = await command.ExecuteReaderAsync(cancellationToken);
                 do
                 {
-                    while (await reader.ReadAsync(TestContext.CancellationToken))
+                    while (await reader.ReadAsync(cancellationToken))
                     {
                     }
                 }
-                while (await reader.NextResultAsync(TestContext.CancellationToken));
-            });
+                while (await reader.NextResultAsync(cancellationToken));
+            }
         }
 
         await AssertSessionReusableAsync(connection, TestContext.CancellationToken);
@@ -211,7 +236,7 @@ public async Task Cancel_RacingNaturalCompletion_DoesNotHang()
         for (var i = 0; i < 30; i++)
         {
             await using var command = new SqlCommand("select 1", connection) { CommandTimeout = 5 };
-            CancelAfter(command, 0);
+            _ = Task.Run(command.Cancel, TestContext.CancellationToken);
             try
             {
                 await using var reader = await command.ExecuteReaderAsync(TestContext.CancellationToken);
@@ -243,9 +268,9 @@ public async Task SequentialCommands_AfterCancel_AllSucceed()
 
         await using (var command = new SqlCommand(LongWait, connection) { CommandTimeout = 0 })
         {
-            CancelAfter(command, 200);
-            _ = await ThrowsExactlyAsync<SqlException>(
-                async () => await command.ExecuteNonQueryAsync(TestContext.CancellationToken));
+            var execution = command.ExecuteNonQueryAsync(TestContext.CancellationToken);
+            await CancelUntilComplete(command, execution, TestContext.CancellationToken);
+            _ = await ThrowsExactlyAsync<SqlException>(async () => await execution);
         }
 
         for (var i = 1; i <= 5; i++)

tests/SqlServerSimulator.Tests.SqlClient/MarsTests.cs

@@ -21,12 +21,22 @@ public sealed class MarsTests
 
     private const string MarsExtra = ";MultipleActiveResultSets=True";
 
-    private static void CancelAfter(SqlCommand command, int millis) =>
-        _ = Task.Run(async () =>
+    /// <summary>
+    /// Cancels the command repeatedly until the in-flight execution completes.
+    /// A cancel landing before the batch starts executing is a documented
+    /// <see cref="SqlCommand.Cancel"/> no-op, so a single timer-fired cancel
+    /// can miss on a stalled runner and let the 30-second WAITFOR run to
+    /// natural completion; retrying until the task transitions guarantees an
+    /// attention lands mid-execution.
+    /// </summary>
+    private static async Task CancelUntilComplete(SqlCommand command, Task execution, CancellationToken cancellationToken)
+    {
+        while (!execution.IsCompleted)
         {
-            await Task.Delay(millis);
+            await Task.Delay(100, cancellationToken);
             command.Cancel();
-        });
+        }
+    }
 
     private static Simulation Seeded(int rows = 5)
     {
@@ -243,9 +253,9 @@ public async Task Cancel_OneOfTwoActiveCommands_LeavesOtherReaderUsable()
 
         await using (var b = new SqlCommand("waitfor delay '00:00:30'", connection) { CommandTimeout = 0 })
         {
-            CancelAfter(b, 200);
-            _ = await ThrowsExactlyAsync<SqlException>(
-                async () => await b.ExecuteNonQueryAsync(TestContext.CancellationToken));
+            var execution = b.ExecuteNonQueryAsync(TestContext.CancellationToken);
+            await CancelUntilComplete(b, execution, TestContext.CancellationToken);
+            _ = await ThrowsExactlyAsync<SqlException>(async () => await execution);
         }
 
         var remaining = new List<int> { readerA.GetInt32(0) };

tests/SqlServerSimulator.Tests/WaitForDelayTests.cs

@@ -74,15 +74,24 @@ public void Delay_InterruptedByInProcessCancel_ReturnsPromptly()
         connection.Open();
         using var command = connection.CreateCommand();
         command.CommandText = "waitfor delay '00:00:30'";
+        // Retry the cancel until ExecuteNonQuery returns: a cancel landing
+        // before the execution scope opens targets the prior scope's token
+        // and is dropped, so a single timer-fired cancel can miss on a
+        // stalled runner and let the 30-second wait run to completion.
+        var done = false;
         var canceller = new Thread(() =>
         {
-            Thread.Sleep(200);
-            command.Cancel();
+            while (!Volatile.Read(ref done))
+            {
+                Thread.Sleep(100);
+                command.Cancel();
+            }
         });
 
         var start = Stopwatch.GetTimestamp();
         canceller.Start();
         _ = command.ExecuteNonQuery();
+        Volatile.Write(ref done, true);
         canceller.Join();
         var elapsed = Stopwatch.GetElapsedTime(start);