Add cancellation coverage for the internal-end execute path

[mdaigle]

Stacked on #4435. Base is dev/cheena/fix-async-cancel-attention, so the diff here is just the two tests.

Opening as a draft because both new tests fail against the current head of #4435 by design. They are the reproduction for the review feedback on that PR.

Why

#4435 removes lock (_stateObj) from EndExecuteReaderAsync, EndExecuteNonQueryAsync and EndExecuteXmlReaderAsync, and keeps it in CreateLocalCompletionTask. The five tests it adds all exercise the paths where the lock was removed, so they pass with or without its own change. Nothing covers the path that is still locked.

That path is reachable in production whenever column encryption is enabled and the parameter metadata came from the cache.

What the tests show

Cancelling 1s into a 60s WAITFOR, with the metadata cache warm and _internalEndExecuteInitiated confirmed true:

	result
lock retained (head of #4435)	60018ms, no exception, ExecuteReaderAsync returns a reader
lock removed	1016ms, SqlException with Number == 0

Number == 0 is the error TdsParser.TryRun adds when it consumes the attention ack, so it is a direct signal that the round trip happened.

Full DataReaderCancellationTest class:

• head of Fix async cancellation failing to send TDS attention signal #4435: 5 passed, 1 failed, ExecuteReaderAsync returned a reader 60015ms after cancellation was requested
• with the lock removed from CreateLocalCompletionTask: 6 passed, the new one in 12ms

The two tests

CancellationOnInternalEndExecutePath_SendsAttention forces the path with the existing DEBUG _forceInternalEndQuery hook, so it runs without an Always Encrypted setup. It no-ops against a Release driver, where the hook is compiled out.

TestAsyncCancellationSendsAttention_WithAlwaysEncryptedCommand_AfterPartialResults is the Always Encrypted equivalent, verified against SQL Server 2022 with a custom key store provider.

It differs from the existing TestAsyncCancellationSendsAttention_WithAlwaysEncryptedCommand in one way: it flushes a partial result before the delay. Both shapes reach CreateLocalCompletionTask, but only this one catches the bug.

query shape	result
WAITFOR DELAY @Delay; SELECT ... (existing test)	cancels in 27ms, passes
RAISERROR(...) WITH NOWAIT; WAITFOR DELAY @Delay; SELECT ...	60018ms, reader returned, fails

With WAITFOR first, nothing comes back until the delay expires, so the continuation never runs while the server is busy and cancellation is handled out in ExecuteReaderAsync's await. The lock is never held across the blocking read, which is the thing under test.

Suggested resolution on #4435

Remove the lock from CreateLocalCompletionTask as well. The justification comment there says the continuation runs after the data is buffered, but localCompletion is completed by BeginExecuteReaderInternalReadStage → _stateObj.ReadSni(completion), which fires on the first packet.

ac8acd72 restored that lock to fix an AE test. The cancellation exception on the unlocked path is a SqlException carrying two errors, A severe error occurred on the current command. followed by Operation cancelled by user., and TestCancellationToken asserts an exact message match. That is a correct cancellation, so the assertion is worth relaxing to check Number == 0 instead.

• Tests added
• Public API changes documented (none)
• Verified against a live SQL Server 2022 instance
• No breaking changes

[Copilot]

Pull request overview

This PR adds two new manual tests intended to reproduce and prevent a regression where async cancellation fails to send a TDS attention signal on the “internal end” completion path (notably reachable with Always Encrypted + cached parameter metadata, or forced via a DEBUG-only hook). The changes extend the existing cancellation coverage by ensuring the scenario includes an early “partial results” packet before a server-side block.

Changes:

• Added a new DataReaderCancellationTest case that forces the internal-end path via the DEBUG-only _forceInternalEndQuery hook.
• Added a new Always Encrypted cancellation test variant that flushes partial results (RAISERROR ... WITH NOWAIT) before WAITFOR, to exercise the internal-end path while the server is blocked.
• Extended the manual-test reflection helper to access _forceInternalEndQuery.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/DataReaderTest/DataReaderCancellationTest.cs	Adds a DEBUG-hook-based regression test for cancellation on the internal-end execute path.
src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/Common/SystemDataInternals/CommandHelper.cs	Adds reflection access for _forceInternalEndQuery to drive the new regression test.
src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ApiShould.cs	Adds an AE cancellation test variant that flushes partial results before a blocking delay.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (2)

src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ApiShould.cs:2462

• This test starts ExecuteReaderAsync before subscribing to InfoMessage and doesn’t verify the RAISERROR WITH NOWAIT message was actually observed. That makes the setup racy: if the INFO token arrives before the handler is attached (or never arrives), the test can cancel without confirming it reached the intended “partial results received/internal-end continuation running” state, producing false positives/flakiness. Attach the handler first, assert it fires within the timeout, and detach it before continuing so later commands on the same connection aren’t affected.

                        // Wait for the RAISERROR NOWAIT packet so we know the continuation is
                        // inside endFunc before cancelling.
                        var infoMessageReceived = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                        sqlConnection.InfoMessage += (_, __) => infoMessageReceived.TrySetResult(true);
                        await Task.WhenAny(infoMessageReceived.Task, Task.Delay(TimeSpan.FromSeconds(5)));

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/DataReaderTest/DataReaderCancellationTest.cs:382

• The test starts ExecuteReaderAsync before subscribing to InfoMessage and doesn’t assert that the RAISERROR WITH NOWAIT message was actually observed. That creates a race where the INFO token can be missed, and the test can proceed to cancel without confirming it reached the intended “partial results received/internal-end continuation running” state (risking false positives/flakiness). Subscribe before starting the command, assert the InfoMessage arrives within the timeout, and unsubscribe to avoid affecting later commands (see ExceptionTest.cs:39-49 for the add/remove pattern).

                        // Let the RAISERROR NOWAIT packet land so the internal-end continuation is
                        // inside endFunc, and therefore inside lock (_stateObj), before we cancel.
                        var infoMessageReceived = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                        connection.InfoMessage += (_, __) => infoMessageReceived.TrySetResult(true);
                        await Task.WhenAny(infoMessageReceived.Task, Task.Delay(System.TimeSpan.FromSeconds(5)));

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (2)

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/DataReaderTest/DataReaderCancellationTest.cs:387

• The test intends to cancel only after the RAISERROR ... WITH NOWAIT InfoMessage arrives (so cancellation happens in the "partial results received" state), but the current WhenAny(..., Delay(10s)) result is ignored. If the delay wins, the test still cancels and may later see InfoMessageReceived.Task.IsCompleted become true (arriving after cancellation), weakening the assertion and potentially causing false results under slow/loaded environments. Capture the completed task and fail immediately if the InfoMessage was not received before cancellation.

                        // Let the RAISERROR NOWAIT packet land so the internal-end continuation is
                        // inside endFunc, and therefore inside lock (_stateObj), before we cancel.
                        await Task.WhenAny(infoMessageReceived.Task, Task.Delay(System.TimeSpan.FromSeconds(10)));

                        long cancelledAtMs = stopwatch.ElapsedMilliseconds;

src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ApiShould.cs:2467

• This test relies on Task.WhenAny(infoMessageReceived, Delay(10s)) to ensure a partial result was flushed before cancelling, but it doesn't check which task completed. If the delay wins, cancellation may happen before the InfoMessage is processed, and infoMessageReceived.Task.IsCompleted can still become true later (after cancellation), reducing the test’s ability to prove it exercised the internal-end partial-results path. Capture the WhenAny result and assert the InfoMessage arrived before calling Cancel().

                        // Wait for the RAISERROR NOWAIT packet so we know the continuation is
                        // inside endFunc before cancelling.
                        await Task.WhenAny(infoMessageReceived.Task, Task.Delay(TimeSpan.FromSeconds(10)));

                        long cancelledAtMs = stopwatch.ElapsedMilliseconds;

[cheenamalhotra]

Thanks for putting this together — the _forceInternalEndQuery hook is the piece that made this testable without an AE setup, and I've merged this branch into #4435 (commit 120ba67).

To close the loop on the finding: the internal-end path is now fixed there rather than left as a known gap. I removed lock (_stateObj) from CreateLocalCompletionTask, which is what was starving TdsParserStateObject.Cancel()'s Monitor.TryEnter poll. Verified against a live server — CancellationOnInternalEndExecutePath_SendsAttention returned at 60005ms before the change and passes in ~1s after.

The reason the earlier revision reverted this (commit ac8acd72) turned out to be the AE TestCancellationToken assertion, not the fix: on the unlocked path the exception is a SqlException with Number == 0 carrying two errors — A severe error occurred on the current command. followed by Operation cancelled by user. The first is what TryRun adds when it consumes the attention ack, so that is correct behaviour; the test was asserting an exact whole-message match. I relaxed it to Assert.Contains.

Since the tests and the fix now both live in #4435, this PR can be closed — happy to do that, or leave it to you.