Fixes cancelation token registration on ReadAsync

Author: cheenamalhotra

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

@@ -4739,6 +4739,12 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return Task.FromCanceled<bool>(cancellationToken);
                 }
 
+                IDisposable registration = null;
+                if (cancellationToken.CanBeCanceled)
+                {
+                    registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
+                }
+
                 // Check for existing async
                 if (_currentTask != null)
                 {
@@ -4831,12 +4837,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

@@ -5332,6 +5332,12 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return ADP.CreatedTaskWithCancellation<bool>();
                 }
 
+                IDisposable registration = null;
+                if (cancellationToken.CanBeCanceled)
+                {
+                    registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
+                }
+
                 // Check for existing async
                 if (_currentTask != null)
                 {
@@ -5425,12 +5431,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

@@ -272,6 +272,7 @@
     <Compile Include="SQL\Common\SystemDataInternals\DataReaderHelper.cs" />
     <Compile Include="SQL\Common\SystemDataInternals\TdsParserHelper.cs" />
     <Compile Include="SQL\Common\SystemDataInternals\TdsParserStateObjectHelper.cs" />
+    <Compile Include="SQL\DataReaderTest\DataReaderCancellationTest.cs" />
     <Compile Include="SQL\SqlCommand\SqlCommandStoredProcTest.cs" />
     <Compile Include="TracingTests\TestTdsServer.cs" />
     <Compile Include="XUnitAssemblyAttributes.cs" />

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

@@ -0,0 +1,49 @@
+// 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 < 1000, "Cancellation did not trigger on time.");
+            }
+        }
+    }
+}