From b3eb4f635031dc8fcb4011f761b9eb9f7a2348f1 Mon Sep 17 00:00:00 2001
From: Smit Patel <smitpatel@users.noreply.github.com>
Date: Tue, 7 Sep 2021 17:42:39 -0700
Subject: [PATCH] Query: Use cancellationToken appropriately

---
 .../Query/Internal/BufferedDataReader.cs      | 26 ++++++++++++++-----
 ...sitor.ShaperProcessingExpressionVisitor.cs | 10 ++++---
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/src/EFCore.Relational/Query/Internal/BufferedDataReader.cs b/src/EFCore.Relational/Query/Internal/BufferedDataReader.cs
index 745468c6897..4e417c0e94b 100644
--- a/src/EFCore.Relational/Query/Internal/BufferedDataReader.cs
+++ b/src/EFCore.Relational/Query/Internal/BufferedDataReader.cs
@@ -624,7 +624,11 @@ public override bool NextResult()
         ///     doing so can result in application failures when updating to a new Entity Framework Core release.
         /// </summary>
         public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
-            => Task.FromResult(NextResult());
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            return Task.FromResult(NextResult());
+        }
 
         /// <summary>
         ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -884,16 +888,26 @@ public bool IsDBNull(int ordinal)
             public bool Read()
                 => IsDataReady = ++_currentRowNumber < _rowCount;
 
-#pragma warning disable IDE0060 // Remove unused parameter
             public Task<T> GetFieldValueAsync<T>(int ordinal, CancellationToken cancellationToken)
-                => Task.FromResult(GetFieldValue<T>(ordinal));
+            {
+                cancellationToken.ThrowIfCancellationRequested();
+
+                return Task.FromResult(GetFieldValue<T>(ordinal));
+            }
 
             public Task<bool> IsDBNullAsync(int ordinal, CancellationToken cancellationToken)
-                => Task.FromResult(IsDBNull(ordinal));
+            {
+                cancellationToken.ThrowIfCancellationRequested();
+
+                return Task.FromResult(IsDBNull(ordinal));
+            }
 
             public Task<bool> ReadAsync(CancellationToken cancellationToken)
-                => Task.FromResult(Read());
-#pragma warning restore IDE0060 // Remove unused parameter
+            {
+                cancellationToken.ThrowIfCancellationRequested();
+
+                return Task.FromResult(Read());
+            }
 
             public BufferedDataRecord Initialize(DbDataReader reader, IReadOnlyList<ReaderColumn> columns)
             {
diff --git a/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs b/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs
index 314f2e75471..a67cdf1c8f2 100644
--- a/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs
+++ b/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs
@@ -1482,10 +1482,11 @@ private static async Task PopulateSplitIncludeCollectionAsync<TIncludingEntity,
                             (queryContext, relationalCommandCache, detailedErrorsEnabled),
                             ((RelationalQueryContext, RelationalCommandCache, bool) tup, CancellationToken cancellationToken)
                                 => InitializeReaderAsync(tup.Item1, tup.Item2, tup.Item3, cancellationToken),
-                            verifySucceeded: null)
+                            verifySucceeded: null,
+                            queryContext.CancellationToken)
                         .ConfigureAwait(false);
 
-                    async Task<RelationalDataReader> InitializeReaderAsync(
+                    static async Task<RelationalDataReader> InitializeReaderAsync(
                         RelationalQueryContext queryContext,
                         RelationalCommandCache relationalCommandCache,
                         bool detailedErrorsEnabled,
@@ -1801,10 +1802,11 @@ private static async Task PopulateSplitCollectionAsync<TCollection, TElement, TR
                             (queryContext, relationalCommandCache, detailedErrorsEnabled),
                             ((RelationalQueryContext, RelationalCommandCache, bool) tup, CancellationToken cancellationToken)
                                 => InitializeReaderAsync(tup.Item1, tup.Item2, tup.Item3, cancellationToken),
-                            verifySucceeded: null)
+                            verifySucceeded: null,
+                            queryContext.CancellationToken)
                         .ConfigureAwait(false);
 
-                    async Task<RelationalDataReader> InitializeReaderAsync(
+                    static async Task<RelationalDataReader> InitializeReaderAsync(
                         RelationalQueryContext queryContext,
                         RelationalCommandCache relationalCommandCache,
                         bool detailedErrorsEnabled,