Fixes ReadAsync() behavior to register Cancellation token action before streaming results (#1781)

Author: cheenamalhotra

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlDataReader.cs

@@ -4733,6 +4733,13 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return Task.FromException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed()));
                 }
 
+                // Register first to catch any already expired tokens to be able to trigger cancellation event.
+                IDisposable registration = null;
+                if (cancellationToken.CanBeCanceled)
+                {
+                    registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
+                }
+
                 // If user's token is canceled, return a canceled task
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -4831,12 +4838,6 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return source.Task;
                 }
 
-                IDisposable registration = null;
-                if (cancellationToken.CanBeCanceled)
-                {
-                    registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
-                }
-
                 ReadAsyncCallContext context = null;
                 if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
                 {

src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlDataReader.cs

@@ -5326,6 +5326,13 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return ADP.CreatedTaskWithException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed("ReadAsync")));
                 }
 
+                // Register first to catch any already expired tokens to be able to trigger cancellation event.
+                IDisposable registration = null;
+                if (cancellationToken.CanBeCanceled)
+                {
+                    registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
+                }
+
                 // If user's token is canceled, return a canceled task
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -5425,12 +5432,6 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return source.Task;
                 }
 
-                IDisposable registration = null;
-                if (cancellationToken.CanBeCanceled)
-                {
-                    registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
-                }
-
                 var context = Interlocked.Exchange(ref _cachedReadAsyncContext, null) ?? new ReadAsyncCallContext();
 
                 Debug.Assert(context._reader == null && context._source == null && context._disposable == null, "cached ReadAsyncCallContext was not properly disposed");

src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj

@@ -86,6 +86,7 @@
     <Compile Include="SQL\ConnectivityTests\TcpDefaultForAzureTest.cs" />
     <Compile Include="SQL\DataBaseSchemaTest\ConnectionSchemaTest.cs" />
     <Compile Include="SQL\DataClassificationTest\DataClassificationTest.cs" />
+    <Compile Include="SQL\DataReaderTest\DataReaderCancellationTest.cs" />
     <Compile Include="SQL\DataReaderTest\DataReaderStreamsTest.cs" />
     <Compile Include="SQL\DataReaderTest\DataReaderTest.cs" />
     <Compile Include="SQL\DataStreamTest\DataStreamTest.cs" />

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

@@ -0,0 +1,83 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.ManualTesting.Tests
+{
+    public class DataReaderCancellationTest
+    {
+        /// <summary>
+        /// Test ensures cancellation token is registered before ReadAsync starts processing results from TDS Stream,
+        /// such that when Cancel is triggered, the token is capable of canceling reading further results.
+        /// </summary>
+        /// <returns>Async Task</returns>
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        public static async Task CancellationTokenIsRespected_ReadAsync()
+        {
+            const string longRunningQuery = @"
+with TenRows as (select Value from (values (1), (2), (3), (4), (5), (6), (7), (8), (9), (10)) as TenRows (Value)),
+    ThousandRows as (select A.Value as A, B.Value as B, C.Value as C from TenRows as A, TenRows as B, TenRows as C)
+select *
+from ThousandRows as A, ThousandRows as B, ThousandRows as C;";
+
+            using (var source = new CancellationTokenSource())
+            using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
+            {
+                await connection.OpenAsync(source.Token);
+
+                Stopwatch stopwatch = Stopwatch.StartNew();
+                await Assert.ThrowsAsync<TaskCanceledException>(async () =>
+                {
+                    using (var command = new SqlCommand(longRunningQuery, connection))
+                    using (var reader = await command.ExecuteReaderAsync(source.Token))
+                    {
+                        while (await reader.ReadAsync(source.Token))
+                        {
+                            source.Cancel();
+                        }
+                    }
+                });
+                Assert.True(stopwatch.ElapsedMilliseconds < 10000, "Cancellation did not trigger on time.");
+            }
+        }
+
+        /// <summary>
+        /// Test ensures cancellation token is registered before ReadAsync starts processing results from TDS Stream,
+        /// such that when Cancel is triggered, the token is capable of canceling reading further results.
+        /// </summary>
+        /// <returns>Async Task</returns>
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        public static async Task CancelledCancellationTokenIsRespected_ReadAsync()
+        {
+            const string longRunningQuery = @"
+with TenRows as (select Value from (values (1), (2), (3), (4), (5), (6), (7), (8), (9), (10)) as TenRows (Value)),
+    ThousandRows as (select A.Value as A, B.Value as B, C.Value as C from TenRows as A, TenRows as B, TenRows as C)
+select *
+from ThousandRows as A, ThousandRows as B, ThousandRows as C;";
+
+            using (var source = new CancellationTokenSource())
+            using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
+            {
+                await connection.OpenAsync(source.Token);
+
+                Stopwatch stopwatch = Stopwatch.StartNew();
+                await Assert.ThrowsAsync<TaskCanceledException>(async () =>
+                {
+                    using (var command = new SqlCommand(longRunningQuery, connection))
+                    using (var reader = await command.ExecuteReaderAsync(source.Token))
+                    {
+                        source.Cancel();
+                        while (await reader.ReadAsync(source.Token))
+                        { }
+                    }
+                });
+                Assert.True(stopwatch.ElapsedMilliseconds < 10000, "Cancellation did not trigger on time.");
+            }
+        }
+    }
+}