Fix async cancellation failing to send TDS attention signal

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.NonQuery.cs

@@ -360,17 +360,9 @@ private int EndExecuteNonQueryAsync(IAsyncResult asyncResult)
             }
 
             ThrowIfReconnectionHasBeenCanceled();
-            // lock on _stateObj prevents races with close/cancel.
-            // If we have already initiated the End call internally, we have already done that, so
-            // no point doing it again.
-            if (!_internalEndExecuteInitiated)
-            {
-                lock (_stateObj)
-                {
-                    return EndExecuteNonQueryInternal(asyncResult);
-                }
-            }
 
+            // Note: We intentionally do NOT lock on _stateObj here.
+            // See comment in EndExecuteReaderAsync and GitHub issue #4424 for details.
             return EndExecuteNonQueryInternal(asyncResult);
         }
 

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs

@@ -671,15 +671,16 @@ private SqlDataReader EndExecuteReaderAsync(IAsyncResult asyncResult)
 
             ThrowIfReconnectionHasBeenCanceled();
 
-            // Lock on _stateObj prevents race with close/cancel
-            if (!_internalEndExecuteInitiated)
-            {
-                lock (_stateObj)
-                {
-                    return EndExecuteReaderInternal(asyncResult);
-                }
-            }
-
+            // Note: We intentionally do NOT lock on _stateObj here.
+            // Taking lock(_stateObj) would prevent Cancel() from acquiring the stateObj
+            // monitor to send a TDS attention signal while FinishExecuteReader may be
+            // blocked on a synchronous network read (e.g., waiting for metadata after
+            // partial results like RAISERROR WITH NOWAIT). This caused cancellation to
+            // hang until the full query completed. See GitHub issue #4424.
+            //
+            // Concurrent close is handled by parser state checks within TryRun
+            // (detects Broken/Closed state). Cancel() uses Monitor.TryEnter with polling
+            // and checks parser state in its loop, so it handles concurrency safely.
             return EndExecuteReaderInternal(asyncResult);
         }
 

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Xml.cs

@@ -393,17 +393,8 @@ private XmlReader EndExecuteXmlReaderAsync(IAsyncResult asyncResult)
 
             ThrowIfReconnectionHasBeenCanceled();
 
-            // Locking _stateObj prevents races with close/cancel.
-            // If we have already initiated the End call internally, we have already done that, so
-            // no point doing it again.
-            if (!_internalEndExecuteInitiated)
-            {
-                lock (_stateObj)
-                {
-                    return EndExecuteXmlReaderInternal(asyncResult);
-                }
-            }
-
+            // Note: We intentionally do NOT lock on _stateObj here.
+            // See comment in EndExecuteReaderAsync and GitHub issue #4424 for details.
             return EndExecuteXmlReaderInternal(asyncResult);
         }
 

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs

@@ -2451,11 +2451,9 @@ private void CreateLocalCompletionTask(
                         Debug.Assert(!_internalEndExecuteInitiated);
                         _internalEndExecuteInitiated = true;
 
-                        // Lock on _stateObj prevents races with close/cancel
-                        lock (_stateObj)
-                        {
-                            endFunc(this, task, /*isInternal:*/ true, endMethod);
-                        }
+                        // Note: We intentionally do NOT lock on _stateObj here.
+                        // See comment in EndExecuteReaderAsync and GitHub issue #4424.
+                        endFunc(this, task, /*isInternal:*/ true, endMethod);
 
                         globalCompletion.TrySetResult(task.Result);
                     }

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

@@ -2276,6 +2276,108 @@ public void TestSqlCommandCancellationToken(string connection, int initalValue,
         }
 
 
+        /// <summary>
+        /// Validates that async cancellation via CancellationToken sends a TDS attention signal
+        /// when an AE-enabled command is blocked on a server-side wait (e.g., WAITFOR DELAY).
+        /// This covers the internal-end path in CreateLocalCompletionTask where the lock on
+        /// _stateObj previously prevented Cancel() from sending attention.
+        /// See GitHub issue #4424.
+        /// Synapse: Incompatible query.
+        /// </summary>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsTargetReadyForAeWithKeyStore))]
+        [ClassData(typeof(AEConnectionStringProvider))]
+        public async Task TestAsyncCancellationSendsAttention_WithAlwaysEncryptedCommand(string connection)
+        {
+            CleanUpTable(connection, _tableName);
+
+            IList<object> values = GetValues(dataHint: 60);
+            int numberOfRows = 10;
+            int rowsAffected = InsertRows(tableName: _tableName, numberofRows: numberOfRows, values: values, connection: connection);
+            Assert.True(rowsAffected == numberOfRows, "number of rows affected is unexpected.");
+
+            using (SqlConnection sqlConnection = new SqlConnection(connection))
+            {
+                await sqlConnection.OpenAsync();
+
+                // Use the same CommandText for both warmup and cancellation test so the
+                // query metadata cache key matches and the second execution goes through
+                // the CreateLocalCompletionTask internal-end path.
+                // WAITFOR is placed BEFORE SELECT so that EndExecuteReaderInternal blocks
+                // waiting for result-set metadata — this is the exact CreateLocalCompletionTask
+                // code path we need to exercise.
+                string commandText = $"WAITFOR DELAY @Delay; SELECT CustomerId, FirstName, LastName FROM [{_tableName}] WHERE FirstName = @FirstName AND CustomerId = @CustomerId;";
+
+                // Warmup: execute with zero delay to populate the metadata cache.
+                using (SqlCommand warmupCmd = new SqlCommand(
+                    commandText, sqlConnection, null, SqlCommandColumnEncryptionSetting.Enabled))
+                {
+                    warmupCmd.Parameters.AddWithValue("@CustomerId", values[0]);
+                    warmupCmd.Parameters.AddWithValue("@FirstName", values[1]);
+                    warmupCmd.Parameters.AddWithValue("@Delay", "00:00:00");
+
+                    using (var reader = await warmupCmd.ExecuteReaderAsync())
+                    {
+                        while (await reader.ReadAsync()) { }
+                        while (await reader.NextResultAsync()) { }
+                    }
+                }
+
+                // Now execute the same command with a long delay and cancel it.
+                // With cached metadata, this goes through CreateLocalCompletionTask's internal-end path.
+                // The WAITFOR blocks before any result-set metadata is returned, so
+                // ExecuteReaderAsync should be cancelled before a reader is produced.
+                using (SqlCommand sqlCommand = new SqlCommand(
+                    commandText, sqlConnection, null, SqlCommandColumnEncryptionSetting.Enabled))
+                {
+                    sqlCommand.Parameters.AddWithValue("@CustomerId", values[0]);
+                    sqlCommand.Parameters.AddWithValue("@FirstName", values[1]);
+                    sqlCommand.Parameters.AddWithValue("@Delay", "00:01:00");
+                    sqlCommand.CommandTimeout = 90;
+
+                    using (CancellationTokenSource cts = new CancellationTokenSource(TimeSpan.FromSeconds(3)))
+                    {
+                        System.Diagnostics.Stopwatch stopwatch = System.Diagnostics.Stopwatch.StartNew();
+
+                        Exception caughtException = null;
+                        try
+                        {
+                            using (SqlDataReader reader = await sqlCommand.ExecuteReaderAsync(cts.Token))
+                            {
+                                // If we reach here, cancellation failed during EndExecuteReaderInternal.
+                                Assert.Fail("ExecuteReaderAsync should have been cancelled before returning a reader.");
+                            }
+                        }
+                        catch (OperationCanceledException ex)
+                        {
+                            caughtException = ex;
+                        }
+                        catch (SqlException ex)
+                        {
+                            // Attention ack from server manifests as SqlException
+                            caughtException = ex;
+                        }
+
+                        stopwatch.Stop();
+
+                        Assert.NotNull(caughtException);
+                        Assert.True(cts.IsCancellationRequested,
+                            "CancellationTokenSource was not cancelled; exception may be unrelated to cancellation.");
+                        Assert.True(stopwatch.ElapsedMilliseconds < 30000,
+                            $"Cancellation took {stopwatch.ElapsedMilliseconds}ms, expected < 30000ms. " +
+                            "Attention signal may not have been sent during AE async execution.");
+                    }
+                }
+
+                // Verify the connection is still usable after cancellation.
+                using (SqlCommand verifyCmd = new SqlCommand(
+                    $"SELECT COUNT(*) FROM [{_tableName}]", sqlConnection))
+                {
+                    object result = await verifyCmd.ExecuteScalarAsync();
+                    Assert.Equal(numberOfRows, (int)result);
+                }
+            }
+        }
+
         [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.IsSGXEnclaveConnStringSetup))]
         public void TestNoneAttestationProtocolWithSGXEnclave()
         {

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

@@ -82,5 +82,132 @@ await Assert.ThrowsAsync<TaskCanceledException>(async () =>
                 Assert.True(stopwatch.ElapsedMilliseconds < 10000, "Cancellation did not trigger on time.");
             }
         }
+        /// <summary>
+        /// Validates that async cancellation sends a TDS attention signal to SQL Server
+        /// when the server has sent partial results (RAISERROR WITH NOWAIT at severity 10)
+        /// followed by a blocking operation (WAITFOR). Without the fix for GitHub issue #4424,
+        /// cancellation would hang until WAITFOR completed naturally.
+        /// Synapse: Incompatible query.
+        /// </summary>
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureSynapse))]
+        public static async Task CancellationSendsAttention_WhenPartialResultsReceived()
+        {
+            // Severity 10 informational message flushed via NOWAIT sends a partial TDS
+            // response, then WAITFOR blocks for 60s. Cancellation should send attention
+            // and abort within seconds.
+            const string query = @"
+RAISERROR('partial result', 10, 1) WITH NOWAIT;
+WAITFOR DELAY '00:01:00';
+SELECT 1 AS Result;";
+
+            using (var cts = new CancellationTokenSource())
+            using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
+            {
+                await connection.OpenAsync();
+
+                using (var command = new SqlCommand(query, connection))
+                {
+                    command.CommandTimeout = 90;
+
+                    // Schedule cancellation BEFORE ExecuteReaderAsync so it fires while
+                    // the async completion path may still be consuming metadata or while
+                    // ReadAsync/NextResultAsync is blocked waiting for the server.
+                    cts.CancelAfter(System.TimeSpan.FromSeconds(2));
+
+                    Stopwatch stopwatch = Stopwatch.StartNew();
+
+                    // Cancellation during async read may surface as either
+                    // OperationCanceledException or SqlException (attention ack).
+                    System.Exception caughtException = null;
+                    try
+                    {
+                        using (var reader = await command.ExecuteReaderAsync(cts.Token))
+                        {
+                            // If we reach here, cancellation failed to abort ExecuteReaderAsync while it was waiting
+                            // for metadata after a partial response (e.g., RAISERROR WITH NOWAIT).
+                            Assert.Fail("ExecuteReaderAsync should have been cancelled before returning a reader.");
+                        }
+                    }
+                    catch (System.OperationCanceledException ex)
+                    {
+                        caughtException = ex;
+                    }
+                    catch (SqlException ex)
+                    {
+                        // Attention acknowledgment from server manifests as SqlException
+                        caughtException = ex;
+                    }
+
+                    stopwatch.Stop();
+
+                    Assert.NotNull(caughtException);
+                    // Ensure the CTS actually fired — guards against false positives
+                    // from unrelated SqlExceptions.
+                    Assert.True(cts.IsCancellationRequested,
+                        "CancellationTokenSource was not cancelled; exception may be unrelated to cancellation.");
+                    // The key assertion: cancellation should complete well before the
+                    // 60-second WAITFOR. Allow up to 30 seconds for CI variability.
+                    Assert.True(stopwatch.ElapsedMilliseconds < 30000,
+                        $"Cancellation took {stopwatch.ElapsedMilliseconds}ms, expected < 30000ms. " +
+                        "Attention signal may not have been sent to the server.");
+                }
+            }
+        }
+
+        /// <summary>
+        /// Validates that cancellation during ExecuteReaderAsync itself sends a TDS attention
+        /// signal when the server is blocked before returning any result set metadata.
+        /// With WAITFOR as the first statement, ExecuteReaderAsync should never return a
+        /// reader — cancellation must abort the operation during the await.
+        /// Synapse: Incompatible query.
+        /// </summary>
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureSynapse))]
+        public static async Task CancellationDuringExecuteReaderAsync_SendsAttention()
+        {
+            // WAITFOR as the first statement means no metadata is returned until it
+            // completes. ExecuteReaderAsync will be blocked in the async completion path.
+            const string query = "WAITFOR DELAY '00:01:00'; SELECT 1 AS Result;";
+
+            using (var cts = new CancellationTokenSource(System.TimeSpan.FromSeconds(2)))
+            using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
+            {
+                await connection.OpenAsync();
+
+                using (var command = new SqlCommand(query, connection))
+                {
+                    command.CommandTimeout = 90;
+                    Stopwatch stopwatch = Stopwatch.StartNew();
+
+                    System.Exception caughtException = null;
+                    try
+                    {
+                        // ExecuteReaderAsync should be cancelled via attention before
+                        // a reader is ever returned.
+                        using (var reader = await command.ExecuteReaderAsync(cts.Token))
+                        {
+                            // If we reach here, cancellation failed to abort ExecuteReaderAsync.
+                            Assert.Fail("ExecuteReaderAsync should have been cancelled before returning a reader.");
+                        }
+                    }
+                    catch (System.OperationCanceledException ex)
+                    {
+                        caughtException = ex;
+                    }
+                    catch (SqlException ex)
+                    {
+                        caughtException = ex;
+                    }
+
+                    stopwatch.Stop();
+
+                    Assert.NotNull(caughtException);
+                    Assert.True(cts.IsCancellationRequested,
+                        "CancellationTokenSource was not cancelled; exception may be unrelated to cancellation.");
+                    Assert.True(stopwatch.ElapsedMilliseconds < 30000,
+                        $"Cancellation took {stopwatch.ElapsedMilliseconds}ms, expected < 30000ms. " +
+                        "Attention signal may not have been sent during ExecuteReaderAsync.");
+                }
+            }
+        }
     }
 }