Address usecase for cancelled cancellation token

Author: cheenamalhotra

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

@@ -4733,18 +4733,19 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return Task.FromException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed()));
                 }
 
-                // If user's token is canceled, return a canceled task
-                if (cancellationToken.IsCancellationRequested)
-                {
-                    return Task.FromCanceled<bool>(cancellationToken);
-                }
-
+                // 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)
+                {
+                    return Task.FromCanceled<bool>(cancellationToken);
+                }
+
                 // Check for existing async
                 if (_currentTask != null)
                 {

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

@@ -5326,18 +5326,19 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return ADP.CreatedTaskWithException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed("ReadAsync")));
                 }
 
-                // If user's token is canceled, return a canceled task
-                if (cancellationToken.IsCancellationRequested)
-                {
-                    return ADP.CreatedTaskWithCancellation<bool>();
-                }
-
+                // 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)
+                {
+                    return ADP.CreatedTaskWithCancellation<bool>();
+                }
+
                 // Check for existing async
                 if (_currentTask != null)
                 {

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

@@ -42,7 +42,41 @@ await Assert.ThrowsAsync<TaskCanceledException>(async () =>
                         }
                     }
                 });
-                Assert.True(stopwatch.ElapsedMilliseconds < 1000, "Cancellation did not trigger on time.");
+                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.");
             }
         }
     }