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,7 +2451,14 @@ private void CreateLocalCompletionTask(
                         Debug.Assert(!_internalEndExecuteInitiated);
                         _internalEndExecuteInitiated = true;
 
-                        // Lock on _stateObj prevents races with close/cancel
+                        // Lock on _stateObj serializes this internal-end path with
+                        // close/cancel. Note: stateObj.Cancel() also acquires this monitor,
+                        // so this lock CAN block Cancel() while endFunc is executing.
+                        // This is retained here (unlike the user-facing EndExecute* methods)
+                        // because this continuation runs after the initial async I/O has
+                        // already completed — the blocking metadata read that caused #4424
+                        // in the user-facing path does not apply here since the data is
+                        // already buffered by the time this continuation fires.
                         lock (_stateObj)
                         {
                             endFunc(this, task, /*isInternal:*/ true, endMethod);

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

@@ -2276,6 +2276,124 @@ 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 exercises the EndExecuteReaderAsync path 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())
+                    {
+                        System.Diagnostics.Stopwatch stopwatch = System.Diagnostics.Stopwatch.StartNew();
+
+                        // Start ExecuteReaderAsync without awaiting so we can trigger
+                        // cancellation from a separate thread once the query is in flight.
+                        // The 60-second WAITFOR gives a wide window during which
+                        // cancellation must send a TDS attention signal to the server.
+                        Task<SqlDataReader> execTask = sqlCommand.ExecuteReaderAsync(cts.Token);
+
+                        // Cancel from another thread after briefly yielding to ensure the
+                        // async operation has been dispatched and reached the server-side
+                        // WAITFOR. This avoids the flakiness of a preemptive timer that
+                        // could fire before the query is actually in flight.
+                        Task cancelTask = Task.Run(async () =>
+                        {
+                            await Task.Delay(TimeSpan.FromMilliseconds(500));
+                            cts.Cancel();
+                        });
+
+                        Exception caughtException = null;
+                        try
+                        {
+                            using (SqlDataReader reader = await execTask)
+                            {
+                                // 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;
+                        }
+
+                        await cancelTask;
+                        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()
         {