Clear _modificationCommandGraph when returning the context to pool (#…

…32871)

Reset StateManager and DatabaseFacade more consistently
Fix functional tests so they actually return the context to pool.

Fixes #32652

src/EFCore.Relational/Infrastructure/EntityFrameworkRelationalServicesBuilder.cs

@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using Microsoft.EntityFrameworkCore.Diagnostics.Internal;
-using Microsoft.EntityFrameworkCore.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Migrations.Internal;
 using Microsoft.EntityFrameworkCore.Query.Internal;
 using Microsoft.EntityFrameworkCore.Storage.Internal;
@@ -152,6 +151,7 @@ public override EntityFrameworkServicesBuilder TryAddCoreServices()
         TryAdd<IRelationalCommandBuilderFactory, RelationalCommandBuilderFactory>();
         TryAdd<IRawSqlCommandBuilder, RawSqlCommandBuilder>();
         TryAdd<ICommandBatchPreparer, CommandBatchPreparer>();
+        TryAdd<IResettableService, ICommandBatchPreparer>(p => p.GetRequiredService<ICommandBatchPreparer>());
         TryAdd<IModificationCommandFactory, ModificationCommandFactory>();
         TryAdd<IMigrationsModelDiffer, MigrationsModelDiffer>();
         TryAdd<IMigrationsSqlGenerator, MigrationsSqlGenerator>();

src/EFCore.Relational/Update/ICommandBatchPreparer.cs

@@ -24,7 +24,7 @@ namespace Microsoft.EntityFrameworkCore.Update;
 ///         for more information and examples.
 ///     </para>
 /// </remarks>
-public interface ICommandBatchPreparer
+public interface ICommandBatchPreparer : IResettableService
 {
     /// <summary>
     ///     Creates the command batches needed to insert/update/delete the entities represented by the given

src/EFCore.Relational/Update/Internal/CommandBatchPreparer.cs

@@ -1334,6 +1334,17 @@ private void AddSameTableEdges()
         }
     }
 
+    /// <inheritdoc/>
+    void IResettableService.ResetState() => _modificationCommandGraph.Clear();
+
+    /// <inheritdoc/>
+    Task IResettableService.ResetStateAsync(CancellationToken cancellationToken)
+    {
+        ((IResettableService)this).ResetState();
+
+        return Task.CompletedTask;
+    }
+
     private sealed record class CommandDependency
     {
         public CommandDependency(IAnnotatable metadata, bool breakable = false)

src/EFCore/DbContext.cs

@@ -4,6 +4,7 @@
 using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
+using Microsoft.EntityFrameworkCore.ChangeTracking;
 using Microsoft.EntityFrameworkCore.ChangeTracking.Internal;
 using Microsoft.EntityFrameworkCore.Internal;
 using Microsoft.EntityFrameworkCore.Metadata.Internal;
@@ -54,7 +55,7 @@ public class DbContext :
 {
     private readonly DbContextOptions _options;
 
-    private IDictionary<(Type Type, string? Name), object>? _sets;
+    private Dictionary<(Type Type, string? Name), object>? _sets;
     private IDbContextServices? _contextServices;
     private IDbContextDependencies? _dbContextDependencies;
     private DatabaseFacade? _database;
@@ -141,7 +142,13 @@ public virtual DatabaseFacade Database
         {
             CheckDisposed();
 
-            return _database ??= new DatabaseFacade(this);
+            if (_database == null)
+            {
+                _database = new DatabaseFacade(this);
+                _cachedResettableServices?.Add(_database);
+            }
+
+            return _database;
         }
     }
 
@@ -152,7 +159,18 @@ public virtual DatabaseFacade Database
     ///     See <see href="https://aka.ms/efcore-docs-change-tracking">EF Core change tracking</see> for more information and examples.
     /// </remarks>
     public virtual ChangeTracker ChangeTracker
-        => _changeTracker ??= InternalServiceProvider.GetRequiredService<IChangeTrackerFactory>().Create();
+    {
+        get
+        {
+            if (_changeTracker == null)
+            {
+                _changeTracker = InternalServiceProvider.GetRequiredService<IChangeTrackerFactory>().Create();
+                _cachedResettableServices?.Add(_changeTracker);
+            }
+
+            return _changeTracker;
+        }
+    }
 
     /// <summary>
     ///     The metadata about the shape of entities, the relationships between them, and how they map to the database.
@@ -882,7 +900,6 @@ private void SetLeaseInternal(DbContextLease lease)
             || _configurationSnapshot.HasChangeTrackerConfiguration)
         {
             var changeTracker = ChangeTracker;
-            ((IResettableService)changeTracker).ResetState();
             changeTracker.AutoDetectChangesEnabled = _configurationSnapshot.AutoDetectChangesEnabled;
             if (_configurationSnapshot.QueryTrackingBehavior.HasValue)
             {
@@ -1018,11 +1035,20 @@ private IEnumerable<IResettableService> GetResettableServices()
             _cachedResettableServices = resettableServices;
         }
 
-        if (_sets is not null)
+        if (_changeTracker != null)
+        {
+            resettableServices.Add(_changeTracker);
+        }
+        else if (_sets is not null)
         {
             resettableServices.AddRange(_sets.Values.OfType<IResettableService>());
         }
 
+        if (_database != null)
+        {
+            resettableServices.Add(_database);
+        }
+
         return resettableServices;
     }
 
@@ -1081,6 +1107,12 @@ private bool DisposeSync(bool leaseActive, bool contextShouldBeDisposed)
         {
             if (contextShouldBeDisposed)
             {
+                if (_contextServices != null)
+                {
+                    // Make sure to create the model before the context is marked as disposed
+                    // This is necessary for the corner case where a pooled context is used only for design-time operations
+                    var _ = Model;
+                }
                 _disposed = true;
                 _lease = DbContextLease.InactiveLease;
             }

src/EFCore/Internal/DbContextLease.cs

@@ -60,7 +60,7 @@ public DbContextLease(IDbContextPool contextPool, bool standalone)
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public bool IsActive
+    public readonly bool IsActive
         => _contextPool != null;
 
     /// <summary>

test/EFCore.Specification.Tests/NonSharedModelTestBase.cs

@@ -8,7 +8,7 @@ namespace Microsoft.EntityFrameworkCore;
 
 public abstract class NonSharedModelTestBase : IDisposable, IAsyncLifetime
 {
-    public static IEnumerable<object[]> IsAsyncData = new object[][] { [false], [true] };
+    public static IEnumerable<object[]> IsAsyncData = [[false], [true]];
 
     protected abstract string StoreName { get; }
     protected abstract ITestStoreFactory TestStoreFactory { get; }
@@ -106,7 +106,7 @@ protected ContextFactory<TContext> CreateContextFactory<TContext>(
         addServices?.Invoke(services);
 
         services = usePooling
-            ? services.AddDbContextPool(typeof(TContext), (s, b) => ConfigureOptions(useServiceProvider ? s : null, b, onConfiguring))
+            ? services.AddPooledDbContextFactory(typeof(TContext), (s, b) => ConfigureOptions(useServiceProvider ? s : null, b, onConfiguring))
             : services.AddDbContext(
                 typeof(TContext),
                 (s, b) => ConfigureOptions(useServiceProvider ? s : null, b, onConfiguring),
@@ -180,21 +180,23 @@ public ContextFactory(IServiceProvider serviceProvider, bool usePooling, TestSto
             UsePooling = usePooling;
             if (usePooling)
             {
-                ContextPool ??= (IDbContextPool)ServiceProvider
-                    .GetRequiredService(typeof(IDbContextPool<>).MakeGenericType(typeof(TContext)));
+                PooledContextFactory = (IDbContextFactory<TContext>)ServiceProvider
+                    .GetRequiredService(typeof(IDbContextFactory<TContext>));
             }
 
             TestStore = testStore;
         }
 
         public IServiceProvider ServiceProvider { get; }
         protected virtual bool UsePooling { get; }
-        private IDbContextPool ContextPool { get; }
+
+        private IDbContextFactory<TContext> PooledContextFactory { get; }
+
         public TestStore TestStore { get; }
 
         public virtual TContext CreateContext()
             => UsePooling
-                ? (TContext)new DbContextLease(ContextPool, standalone: true).Context
+                ? PooledContextFactory.CreateDbContext()
                 : (TContext)ServiceProvider.GetRequiredService(typeof(TContext));
     }
 }

test/EFCore.Specification.Tests/SharedStoreFixtureBase.cs

@@ -1,8 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using Microsoft.EntityFrameworkCore.Internal;
-
 // ReSharper disable VirtualMemberCallInConstructor
 namespace Microsoft.EntityFrameworkCore;
 
@@ -31,33 +29,40 @@ public TestStore TestStore
     protected virtual bool UsePooling
         => true;
 
-    private IDbContextPool _contextPool;
+    private object _contextFactory;
 
-    private IDbContextPool ContextPool
-        => _contextPool ??= (IDbContextPool)ServiceProvider
-            .GetRequiredService(typeof(IDbContextPool<>).MakeGenericType(ContextType));
+    private object ContextFactory
+        => _contextFactory ??= ServiceProvider
+            .GetRequiredService(typeof(IDbContextFactory<>).MakeGenericType(ContextType));
 
     private ListLoggerFactory _listLoggerFactory;
 
     public ListLoggerFactory ListLoggerFactory
         => _listLoggerFactory ??= (ListLoggerFactory)ServiceProvider.GetRequiredService<ILoggerFactory>();
 
+    private MethodInfo _createDbContext;
+
     public virtual Task InitializeAsync()
     {
         _testStore = TestStoreFactory.GetOrCreate(StoreName);
 
         var services = AddServices(TestStoreFactory.AddProviderServices(new ServiceCollection()));
-        if (UsePooling)
-        {
-            services = services.AddDbContextPool(ContextType, (s, b) => ConfigureOptions(s, b));
-        }
-        else
-        {
-            services = services.AddDbContext(
+        services = UsePooling
+            ? services.AddPooledDbContextFactory(ContextType, (s, b) => ConfigureOptions(s, b))
+            : services.AddDbContext(
                 ContextType,
                 (s, b) => ConfigureOptions(s, b),
                 ServiceLifetime.Transient,
                 ServiceLifetime.Singleton);
+
+        if (UsePooling)
+        {
+            _createDbContext
+                = typeof(IDbContextFactory<>).MakeGenericType(ContextType)
+                    .GetTypeInfo().GetDeclaredMethods(nameof(IDbContextFactory<TContext>.CreateDbContext))
+                    .Single(
+                        mi => mi.GetParameters().Length == 0
+                            && mi.GetGenericArguments().Length == 0);
         }
 
         _serviceProvider = services.BuildServiceProvider(validateScopes: true);
@@ -69,7 +74,7 @@ public virtual Task InitializeAsync()
 
     public virtual TContext CreateContext()
         => UsePooling
-            ? (TContext)new DbContextLease(ContextPool, standalone: true).Context
+            ? (TContext)_createDbContext.Invoke(ContextFactory, null)
             : (TContext)ServiceProvider.GetRequiredService(ContextType);
 
     public DbContextOptions CreateOptions()

test/EFCore.Specification.Tests/TestUtilities/ServiceCollectionExtensions.cs

@@ -30,7 +30,7 @@ private static readonly MethodInfo _addDbContextPool
                     && mi.GetParameters()[1].ParameterType == typeof(Action<IServiceProvider, DbContextOptionsBuilder>)
                     && mi.GetGenericArguments().Length == 1);
 
-    public static IServiceCollection AddDbContextPool(
+    public static IServiceCollection AddPooledDbContextFactory(
         this IServiceCollection serviceCollection,
         Type contextType,
         Action<IServiceProvider, DbContextOptionsBuilder> optionsAction)