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,98 @@ 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();
+
+                // First query to warm the metadata cache (so subsequent calls use the internal-end path)
+                using (SqlCommand warmupCmd = new SqlCommand(
+                    $"SELECT CustomerId, FirstName, LastName FROM [{_tableName}] WHERE FirstName = @FirstName AND CustomerId = @CustomerId",
+                    sqlConnection, null, SqlCommandColumnEncryptionSetting.Enabled))
+                {
+                    warmupCmd.Parameters.AddWithValue("@CustomerId", values[0]);
+                    warmupCmd.Parameters.AddWithValue("@FirstName", values[1]);
+
+                    using (var reader = await warmupCmd.ExecuteReaderAsync())
+                    {
+                        while (await reader.ReadAsync()) { }
+                    }
+                }
+
+                // Now execute a long-running query with AE enabled and cancel it.
+                // The WAITFOR ensures the server blocks after initial metadata/results are sent.
+                // With cached metadata, this goes through CreateLocalCompletionTask's internal-end path.
+                using (SqlCommand sqlCommand = new SqlCommand(
+                    $"SELECT CustomerId, FirstName, LastName FROM [{_tableName}] WHERE FirstName = @FirstName AND CustomerId = @CustomerId; WAITFOR DELAY '00:01:00';",
+                    sqlConnection, null, SqlCommandColumnEncryptionSetting.Enabled))
+                {
+                    sqlCommand.Parameters.AddWithValue("@CustomerId", values[0]);
+                    sqlCommand.Parameters.AddWithValue("@FirstName", values[1]);
+                    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))
+                            {
+                                while (await reader.ReadAsync(cts.Token)) { }
+                                // NextResultAsync will block on WAITFOR — cancellation should abort it
+                                await reader.NextResultAsync(cts.Token);
+                            }
+                        }
+                        catch (OperationCanceledException ex)
+                        {
+                            caughtException = ex;
+                        }
+                        catch (SqlException ex)
+                        {
+                            // Attention ack from server manifests as SqlException
+                            caughtException = ex;
+                        }
+
+                        stopwatch.Stop();
+
+                        Assert.NotNull(caughtException);
+                        // Cancellation should complete well before the 60-second WAITFOR.
+                        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,125 @@ 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) 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()
+        {
+            // This query sends a partial response (RAISERROR WITH NOWAIT), then blocks
+            // for 60 seconds. Cancellation should send attention and abort within seconds.
+            const string query = @"
+RAISERROR('partial result', 0, 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))
+                        {
+                            while (await reader.ReadAsync(cts.Token))
+                            { }
+                            // Advance to next result set (blocked by WAITFOR)
+                            await reader.NextResultAsync(cts.Token);
+                        }
+                    }
+                    catch (System.OperationCanceledException ex)
+                    {
+                        caughtException = ex;
+                    }
+                    catch (SqlException ex)
+                    {
+                        // Attention acknowledgment from server manifests as SqlException
+                        caughtException = ex;
+                    }
+
+                    stopwatch.Stop();
+
+                    Assert.NotNull(caughtException);
+                    // 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 (not just ReadAsync)
+        /// sends a TDS attention signal. This covers the case where the async completion path
+        /// in EndExecuteReaderInternal is blocked on synchronous metadata consumption while
+        /// the server is still processing (e.g., a long-running batch where initial metadata
+        /// is delayed). This also covers the internal-end path used by Always Encrypted retry.
+        /// Synapse: Incompatible query.
+        /// </summary>
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureSynapse))]
+        public static async Task CancellationDuringExecuteReaderAsync_SendsAttention()
+        {
+            // Use a query that blocks immediately so cancellation must fire during
+            // ExecuteReaderAsync's async completion (before any reader is returned).
+            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 itself should be cancelled via attention
+                        using (var reader = await command.ExecuteReaderAsync(cts.Token))
+                        {
+                            await reader.ReadAsync(cts.Token);
+                        }
+                    }
+                    catch (System.OperationCanceledException ex)
+                    {
+                        caughtException = ex;
+                    }
+                    catch (SqlException ex)
+                    {
+                        caughtException = ex;
+                    }
+
+                    stopwatch.Stop();
+
+                    Assert.NotNull(caughtException);
+                    Assert.True(stopwatch.ElapsedMilliseconds < 30000,
+                        $"Cancellation took {stopwatch.ElapsedMilliseconds}ms, expected < 30000ms. " +
+                        "Attention signal may not have been sent during ExecuteReaderAsync.");
+                }
+            }
+        }
     }
 }