Dispose service properties when context is disposed or returned to the pool

ajcvickers · ajcvickers · commit fd96b9488d9f · 2022-12-24T12:23:48.000Z
Fixes #25486
diff --git a/src/EFCore/ChangeTracking/ChangeTracker.cs b/src/EFCore/ChangeTracking/ChangeTracker.cs
@@ -609,7 +609,7 @@ Task IResettableService.ResetStateAsync(CancellationToken cancellationToken)
     ///     </para>
     /// </remarks>
     public virtual void Clear()
-        => StateManager.Clear();
+        => StateManager.Clear(resetting: false);
 
     /// <summary>
     ///     <para>
diff --git a/src/EFCore/ChangeTracking/Internal/IStateManager.cs b/src/EFCore/ChangeTracking/Internal/IStateManager.cs
@@ -457,7 +457,7 @@ IEnumerable<IUpdateEntry> GetDependentsUsingRelationshipSnapshot(
     ///     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>
-    void Unsubscribe();
+    void Unsubscribe(bool resetting);
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -584,5 +584,5 @@ void SetEvents(
     ///     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>
-    void Clear();
+    void Clear(bool resetting);
 }
diff --git a/src/EFCore/ChangeTracking/Internal/InternalEntityEntry.cs b/src/EFCore/ChangeTracking/Internal/InternalEntityEntry.cs
@@ -470,26 +470,29 @@ private void FireStateChanged(EntityState oldState)
 
     private void SetServiceProperties(EntityState oldState, EntityState newState)
     {
-        if (oldState == EntityState.Detached)
+        if (EntityType.HasServiceProperties())
         {
-            foreach (var serviceProperty in EntityType.GetServiceProperties())
+            if (oldState == EntityState.Detached)
             {
-                this[serviceProperty] = (this[serviceProperty] is IInjectableService injectableService
-                        ? injectableService.Attaching(Context, Entity, injectableService)
-                        : null)
-                    ?? serviceProperty
-                        .ParameterBinding
-                        .ServiceDelegate(new MaterializationContext(ValueBuffer.Empty, Context), EntityType, Entity);
+                foreach (var serviceProperty in EntityType.GetServiceProperties())
+                {
+                    this[serviceProperty] = (this[serviceProperty] is IInjectableService injectableService
+                            ? injectableService.Attaching(Context, Entity, injectableService)
+                            : null)
+                        ?? serviceProperty
+                            .ParameterBinding
+                            .ServiceDelegate(new MaterializationContext(ValueBuffer.Empty, Context), EntityType, Entity);
+                }
             }
-        }
-        else if (newState == EntityState.Detached)
-        {
-            foreach (var serviceProperty in EntityType.GetServiceProperties())
+            else if (newState == EntityState.Detached)
             {
-                if (!(this[serviceProperty] is IInjectableService detachable)
-                    || detachable.Detaching(Context, Entity))
+                foreach (var serviceProperty in EntityType.GetServiceProperties())
                 {
-                    this[serviceProperty] = null;
+                    if (!(this[serviceProperty] is IInjectableService detachable)
+                        || detachable.Detaching(Context, Entity))
+                    {
+                        this[serviceProperty] = null;
+                    }
                 }
             }
         }
@@ -2030,6 +2033,11 @@ public bool IsLoaded(INavigationBase navigation)
 
     private ILazyLoader? GetLazyLoader()
     {
+        if (!EntityType.HasServiceProperties())
+        {
+            return null;
+        }
+
         var lazyLoaderProperty = EntityType.GetServiceProperties().FirstOrDefault(p => p.ClrType == typeof(ILazyLoader));
         return lazyLoaderProperty != null ? (ILazyLoader?)this[lazyLoaderProperty] : null;
     }
diff --git a/src/EFCore/ChangeTracking/Internal/StateManager.cs b/src/EFCore/ChangeTracking/Internal/StateManager.cs
@@ -22,6 +22,7 @@ public class StateManager : IStateManager
     private IIdentityMap? _identityMap1;
     private Dictionary<IKey, IIdentityMap>? _identityMaps;
     private bool _needsUnsubscribe;
+    private bool _hasServiceProperties;
     private IChangeDetector? _changeDetector;
 
     private readonly IDiagnosticsLogger<DbLoggerCategory.ChangeTracking> _changeTrackingLogger;
@@ -361,6 +362,11 @@ public virtual InternalEntityEntry StartTrackingFromQuery(
             _needsUnsubscribe = true;
         }
 
+        if (newEntry.EntityType.HasServiceProperties())
+        {
+            _hasServiceProperties = true;
+        }
+
         return newEntry;
     }
 
@@ -600,6 +606,11 @@ public virtual InternalEntityEntry StartTracking(InternalEntityEntry entry)
             _needsUnsubscribe = true;
         }
 
+        if (entry.EntityType.HasServiceProperties())
+        {
+            _hasServiceProperties = true;
+        }
+
         return entry;
     }
 
@@ -661,7 +672,7 @@ public virtual void StopTracking(InternalEntityEntry entry, EntityState oldState
     ///     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 virtual void Unsubscribe()
+    public virtual void Unsubscribe(bool resetting)
     {
         if (_needsUnsubscribe)
         {
@@ -670,6 +681,27 @@ public virtual void Unsubscribe()
                 _internalEntityEntrySubscriber.Unsubscribe(entry);
             }
         }
+
+        if (_hasServiceProperties)
+        {
+            foreach (var entry in Entries)
+            {
+                foreach (var serviceProperty in entry.EntityType.GetServiceProperties())
+                {
+                    var service = entry[serviceProperty];
+                    if (resetting
+                        && service is IDisposable disposable)
+                    {
+                        disposable.Dispose();
+                    }
+                    else if (!(service is IInjectableService detachable)
+                             || detachable.Detaching(Context, entry.Entity))
+                    {
+                        entry[serviceProperty] = null;
+                    }
+                }
+            }
+        }
     }
 
     /// <summary>
@@ -680,7 +712,7 @@ public virtual void Unsubscribe()
     /// </summary>
     public virtual void ResetState()
     {
-        Clear();
+        Clear(resetting: true);
         Dependencies.NavigationFixer.AbortDelayedFixup();
         _changeDetector?.ResetState();
 
@@ -696,9 +728,9 @@ public virtual void ResetState()
     ///     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 virtual void Clear()
+    public virtual void Clear(bool resetting)
     {
-        Unsubscribe();
+        Unsubscribe(resetting);
         ChangedCount = 0;
         _entityReferenceMap.Clear();
         _referencedUntrackedEntities = null;
@@ -708,6 +740,7 @@ public virtual void Clear()
         _identityMap1?.Clear();
 
         _needsUnsubscribe = false;
+        _hasServiceProperties = false;
 
         SavingChanges = false;
 
diff --git a/src/EFCore/DbContext.cs b/src/EFCore/DbContext.cs
@@ -1092,7 +1092,7 @@ private bool DisposeSync(bool leaseActive, bool contextShouldBeDisposed)
 
             _disposed = true;
 
-            _dbContextDependencies?.StateManager.Unsubscribe();
+            _dbContextDependencies?.StateManager.Unsubscribe(resetting: true);
 
             _dbContextDependencies = null;
             _changeTracker = null;
diff --git a/src/EFCore/Internal/EntityFinder.cs b/src/EFCore/Internal/EntityFinder.cs
@@ -447,7 +447,7 @@ public virtual void Load(INavigation navigation, InternalEntityEntry entry, bool
                 {
                     if (stateManager != _stateManager)
                     {
-                        stateManager.ResetState();
+                        stateManager.Clear(resetting: false);
                     }
                 }
             }
@@ -509,7 +509,7 @@ public virtual async Task LoadAsync(
                 {
                     if (stateManager != _stateManager)
                     {
-                        await stateManager.ResetStateAsync(cancellationToken).ConfigureAwait(continueOnCapturedContext: false);
+                        stateManager.Clear(resetting: false);
                     }
                 }
             }
diff --git a/src/EFCore/Internal/ManyToManyLoader.cs b/src/EFCore/Internal/ManyToManyLoader.cs
@@ -61,7 +61,7 @@ public virtual void Load(InternalEntityEntry entry, bool forceIdentityResolution
                 {
                     if (stateManager != entry.StateManager)
                     {
-                        stateManager.ResetState();
+                        stateManager.Clear(resetting: false);
                     }
                 }
             }
@@ -106,7 +106,7 @@ public virtual async Task LoadAsync(
                 {
                     if (stateManager != entry.StateManager)
                     {
-                        await stateManager.ResetStateAsync(cancellationToken).ConfigureAwait(false);
+                        stateManager.Clear(resetting: false);
                     }
                 }
             }
diff --git a/src/EFCore/Metadata/Conventions/ServicePropertyDiscoveryConvention.cs b/src/EFCore/Metadata/Conventions/ServicePropertyDiscoveryConvention.cs
@@ -81,7 +81,8 @@ void DiscoverServiceProperties(IEnumerable<MemberInfo> members)
                 }
 
                 var memberType = memberInfo.GetMemberType();
-                if (entityType.GetServiceProperties().Any(p => p.ClrType == memberType))
+                if (entityType.HasServiceProperties()
+                    && entityType.GetServiceProperties().Any(p => p.ClrType == memberType))
                 {
                     continue;
                 }
diff --git a/src/EFCore/Metadata/IReadOnlyEntityType.cs b/src/EFCore/Metadata/IReadOnlyEntityType.cs
@@ -741,6 +741,12 @@ IReadOnlyProperty GetProperty(string name)
     /// <returns>Derived service properties.</returns>
     IEnumerable<IReadOnlyServiceProperty> GetDerivedServiceProperties();
 
+    /// <summary>
+    ///     Checks whether or not this entity type has any <see cref="IServiceProperty" /> defined.
+    /// </summary>
+    /// <returns><see langword="true"/> if there are any service properties defined on this entity type or base types.</returns>
+    bool HasServiceProperties();
+
     /// <summary>
     ///     Gets all the <see cref="IReadOnlyServiceProperty" /> defined on this entity type.
     /// </summary>
diff --git a/src/EFCore/Metadata/Internal/EntityType.cs b/src/EFCore/Metadata/Internal/EntityType.cs
@@ -2939,6 +2939,15 @@ public virtual ServiceProperty RemoveServiceProperty(ServiceProperty property)
         return property;
     }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     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 virtual bool HasServiceProperties()
+        => _serviceProperties.Count != 0 || _baseType != null && _baseType.HasServiceProperties();
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -5252,6 +5261,16 @@ IEnumerable<IServiceProperty> IEntityType.GetDeclaredServiceProperties()
     IEnumerable<IReadOnlyServiceProperty> IReadOnlyEntityType.GetDerivedServiceProperties()
         => GetDerivedServiceProperties();
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     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>
+    [DebuggerStepThrough]
+    bool IReadOnlyEntityType.HasServiceProperties()
+        => HasServiceProperties();
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
diff --git a/src/EFCore/Metadata/RuntimeEntityType.cs b/src/EFCore/Metadata/RuntimeEntityType.cs
@@ -763,6 +763,9 @@ public virtual RuntimeServiceProperty AddServiceProperty(
             ? property
             : null;
 
+    private bool HasServiceProperties()
+        => _serviceProperties.Count != 0 || _baseType != null && _baseType.HasServiceProperties();
+
     private IEnumerable<RuntimeServiceProperty> GetServiceProperties()
         => _baseType != null
             ? _serviceProperties.Count == 0
@@ -1395,6 +1398,11 @@ IEnumerable<IServiceProperty> IEntityType.GetDeclaredServiceProperties()
     IEnumerable<IReadOnlyServiceProperty> IReadOnlyEntityType.GetDerivedServiceProperties()
         => GetDerivedServiceProperties();
 
+    /// <inheritdoc />
+    [DebuggerStepThrough]
+    bool IReadOnlyEntityType.HasServiceProperties()
+        => HasServiceProperties();
+
     /// <inheritdoc />
     [DebuggerStepThrough]
     IEnumerable<IReadOnlyServiceProperty> IReadOnlyEntityType.GetServiceProperties()
diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs
@@ -512,6 +512,9 @@ public IEnumerable<IReadOnlyProperty> GetDerivedProperties()
         public IEnumerable<IReadOnlyServiceProperty> GetDerivedServiceProperties()
             => throw new NotImplementedException();
 
+        public bool HasServiceProperties()
+            => throw new NotImplementedException();
+
         public IEnumerable<IReadOnlySkipNavigation> GetDerivedSkipNavigations()
             => throw new NotImplementedException();
 
diff --git a/test/EFCore.SqlServer.FunctionalTests/DbContextPoolingTest.cs b/test/EFCore.SqlServer.FunctionalTests/DbContextPoolingTest.cs
diff --git a/test/EFCore.Tests/TestUtilities/FakeStateManager.cs b/test/EFCore.Tests/TestUtilities/FakeStateManager.cs

Original file line number	Diff line number	Diff line change
`@@ -470,26 +470,29 @@ private void FireStateChanged(EntityState oldState)`
`470`	`470`
`471`	`471`	`private void SetServiceProperties(EntityState oldState, EntityState newState)`
`472`	`472`	`{`
`473`		`- if (oldState == EntityState.Detached)`
	`473`	`+ if (EntityType.HasServiceProperties())`
`474`	`474`	`{`
`475`		`- foreach (var serviceProperty in EntityType.GetServiceProperties())`
	`475`	`+ if (oldState == EntityState.Detached)`
`476`	`476`	`{`
`477`		`- this[serviceProperty] = (this[serviceProperty] is IInjectableService injectableService`
`478`		`- ? injectableService.Attaching(Context, Entity, injectableService)`
`479`		`- : null)`
`480`		`- ?? serviceProperty`
`481`		`- .ParameterBinding`
`482`		`- .ServiceDelegate(new MaterializationContext(ValueBuffer.Empty, Context), EntityType, Entity);`
	`477`	`+ foreach (var serviceProperty in EntityType.GetServiceProperties())`
	`478`	`+ {`
	`479`	`+ this[serviceProperty] = (this[serviceProperty] is IInjectableService injectableService`
	`480`	`+ ? injectableService.Attaching(Context, Entity, injectableService)`
	`481`	`+ : null)`
	`482`	`+ ?? serviceProperty`
	`483`	`+ .ParameterBinding`
	`484`	`+ .ServiceDelegate(new MaterializationContext(ValueBuffer.Empty, Context), EntityType, Entity);`
	`485`	`+ }`
`483`	`486`	`}`
`484`		`- }`
`485`		`- else if (newState == EntityState.Detached)`
`486`		`- {`
`487`		`- foreach (var serviceProperty in EntityType.GetServiceProperties())`
	`487`	`+ else if (newState == EntityState.Detached)`
`488`	`488`	`{`
`489`		`- if (!(this[serviceProperty] is IInjectableService detachable)`
`490`		`- \|\| detachable.Detaching(Context, Entity))`
	`489`	`+ foreach (var serviceProperty in EntityType.GetServiceProperties())`
`491`	`490`	`{`
`492`		`- this[serviceProperty] = null;`
	`491`	`+ if (!(this[serviceProperty] is IInjectableService detachable)`
	`492`	`+ \|\| detachable.Detaching(Context, Entity))`
	`493`	`+ {`
	`494`	`+ this[serviceProperty] = null;`
	`495`	`+ }`
`493`	`496`	`}`
`494`	`497`	`}`
`495`	`498`	`}`
`@@ -2030,6 +2033,11 @@ public bool IsLoaded(INavigationBase navigation)`
`2030`	`2033`
`2031`	`2034`	`private ILazyLoader? GetLazyLoader()`
`2032`	`2035`	`{`
	`2036`	`+ if (!EntityType.HasServiceProperties())`
	`2037`	`+ {`
	`2038`	`+ return null;`
	`2039`	`+ }`
	`2040`	`+`
`2033`	`2041`	`var lazyLoaderProperty = EntityType.GetServiceProperties().FirstOrDefault(p => p.ClrType == typeof(ILazyLoader));`
`2034`	`2042`	`return lazyLoaderProperty != null ? (ILazyLoader?)this[lazyLoaderProperty] : null;`
`2035`	`2043`	`}`
Original file line number	Diff line number	Diff line change
`@@ -447,7 +447,7 @@ public virtual void Load(INavigation navigation, InternalEntityEntry entry, bool`
`447`	`447`	`{`
`448`	`448`	`if (stateManager != _stateManager)`
`449`	`449`	`{`
`450`		`- stateManager.ResetState();`
	`450`	`+ stateManager.Clear(resetting: false);`
`451`	`451`	`}`
`452`	`452`	`}`
`453`	`453`	`}`
`@@ -509,7 +509,7 @@ public virtual async Task LoadAsync(`
`509`	`509`	`{`
`510`	`510`	`if (stateManager != _stateManager)`
`511`	`511`	`{`
`512`		`- await stateManager.ResetStateAsync(cancellationToken).ConfigureAwait(continueOnCapturedContext: false);`
	`512`	`+ stateManager.Clear(resetting: false);`
`513`	`513`	`}`
`514`	`514`	`}`
`515`	`515`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ public virtual void Load(InternalEntityEntry entry, bool forceIdentityResolution`
`61`	`61`	`{`
`62`	`62`	`if (stateManager != entry.StateManager)`
`63`	`63`	`{`
`64`		`- stateManager.ResetState();`
	`64`	`+ stateManager.Clear(resetting: false);`
`65`	`65`	`}`
`66`	`66`	`}`
`67`	`67`	`}`
`@@ -106,7 +106,7 @@ public virtual async Task LoadAsync(`
`106`	`106`	`{`
`107`	`107`	`if (stateManager != entry.StateManager)`
`108`	`108`	`{`
`109`		`- await stateManager.ResetStateAsync(cancellationToken).ConfigureAwait(false);`
	`109`	`+ stateManager.Clear(resetting: false);`
`110`	`110`	`}`
`111`	`111`	`}`
`112`	`112`	`}`
Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,8 @@ void DiscoverServiceProperties(IEnumerable<MemberInfo> members)`
`81`	`81`	`}`
`82`	`82`
`83`	`83`	`var memberType = memberInfo.GetMemberType();`
`84`		`- if (entityType.GetServiceProperties().Any(p => p.ClrType == memberType))`
	`84`	`+ if (entityType.HasServiceProperties()`
	`85`	`+ && entityType.GetServiceProperties().Any(p => p.ClrType == memberType))`
`85`	`86`	`{`
`86`	`87`	`continue;`
`87`	`88`	`}`