From 44f84fcb16a18f11f24b506db0b5eaf2d5688fd9 Mon Sep 17 00:00:00 2001
From: Smit Patel <smitpatel@users.noreply.github.com>
Date: Mon, 29 Jun 2020 09:05:49 -0700
Subject: [PATCH] Query: Dispose inner data readers for split query

Part of #21441
---
 .../Query/Internal/SplitQueryingEnumerable.cs | 23 +++++--
 .../Query/QueryBugsTest.cs                    | 65 +++++++++++++++++--
 2 files changed, 78 insertions(+), 10 deletions(-)

diff --git a/src/EFCore.Relational/Query/Internal/SplitQueryingEnumerable.cs b/src/EFCore.Relational/Query/Internal/SplitQueryingEnumerable.cs
index 88681646931..29c375aa574 100644
--- a/src/EFCore.Relational/Query/Internal/SplitQueryingEnumerable.cs
+++ b/src/EFCore.Relational/Query/Internal/SplitQueryingEnumerable.cs
@@ -203,6 +203,12 @@ private bool InitializeReader(DbContext _, bool result)
             public void Dispose()
             {
                 _dataReader?.Dispose();
+                foreach (var dataReader in _resultCoordinator.DataReaders)
+                {
+                    dataReader?.DataReader.Dispose();
+                }
+                _resultCoordinator.DataReaders.Clear();
+
                 _dataReader = null;
             }
 
@@ -300,17 +306,22 @@ private async Task<bool> InitializeReaderAsync(DbContext _, bool result, Cancell
                 return result;
             }
 
-            public ValueTask DisposeAsync()
+            public async ValueTask DisposeAsync()
             {
                 if (_dataReader != null)
                 {
-                    var dataReader = _dataReader;
-                    _dataReader = null;
+                    await _dataReader.DisposeAsync().ConfigureAwait(false);
+                    foreach (var dataReader in _resultCoordinator.DataReaders)
+                    {
+                        if (dataReader != null)
+                        {
+                            await dataReader.DataReader.DisposeAsync().ConfigureAwait(false);
+                        }
+                    }
+                    _resultCoordinator.DataReaders.Clear();
 
-                    return dataReader.DisposeAsync();
+                    _dataReader = null;
                 }
-
-                return default;
             }
         }
     }
diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs
index 35fd3817aab..ca835e1f991 100644
--- a/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs
+++ b/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs
@@ -18,6 +18,7 @@
 using Microsoft.EntityFrameworkCore.Infrastructure;
 using Microsoft.EntityFrameworkCore.Query.Internal;
 using Microsoft.EntityFrameworkCore.SqlServer.Diagnostics.Internal;
+using Microsoft.EntityFrameworkCore.Storage;
 using Microsoft.EntityFrameworkCore.Storage.ValueConversion;
 using Microsoft.EntityFrameworkCore.TestUtilities;
 using Microsoft.Extensions.Caching.Memory;
@@ -7467,6 +7468,62 @@ ORDER BY [p].[Id]"
                             });
         }
 
+        [ConditionalFact]
+        public virtual void SplitQuery_disposes_inner_data_readers()
+        {
+            var (options, testSqlLoggerFactory) = CreateOptions21355(null);
+            using var context = new BugContext21355(options);
+            var dbConnection = context.Database.GetDbConnection();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+
+            context.Parents.Include(p => p.Children1).Include(p => p.Children2).AsSplitQuery().ToList();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+        }
+
+        [ConditionalFact]
+        public virtual async Task SplitQuery_disposes_inner_data_readers_async()
+        {
+            var (options, testSqlLoggerFactory) = CreateOptions21355(null);
+            using var context = new BugContext21355(options);
+            var dbConnection = context.Database.GetDbConnection();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+
+            await context.Parents.Include(p => p.Children1).Include(p => p.Children2).AsSplitQuery().ToListAsync();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+        }
+
+        [ConditionalFact]
+        public virtual void SplitQuery_disposes_inner_data_readers_single()
+        {
+            var (options, testSqlLoggerFactory) = CreateOptions21355(null);
+            using var context = new BugContext21355(options);
+            var dbConnection = context.Database.GetDbConnection();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+
+            context.Parents.Include(p => p.Children1).Include(p => p.Children2).AsSplitQuery().Single();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+        }
+
+        [ConditionalFact]
+        public virtual async Task SplitQuery_disposes_inner_data_readers_single_async()
+        {
+            var (options, testSqlLoggerFactory) = CreateOptions21355(null);
+            using var context = new BugContext21355(options);
+            var dbConnection = context.Database.GetDbConnection();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+
+            await context.Parents.Include(p => p.Children1).Include(p => p.Children2).AsSplitQuery().SingleAsync();
+
+            Assert.Equal(ConnectionState.Closed, dbConnection.State);
+        }
+
         private class Parent21355
         {
             public string Id { get; set; }
@@ -7511,10 +7568,10 @@ private class AnotherChild21355
                 {
                     Id = "Parent1",
                     Children1 = new List<Child21355>
-                {
-                    new Child21355(),
-                    new Child21355()
-                }
+                    {
+                        new Child21355(),
+                        new Child21355()
+                    }
                 });
                 context.SaveChanges();
             }