Fixed OperationContext Memory Leak

src/HotChocolate/Core/src/Execution/DependencyInjection/Factories/DeferredWorkStateOwnerFactory.cs

@@ -0,0 +1,42 @@
+using HotChocolate.Execution.DependencyInjection;
+using HotChocolate.Execution.Processing;
+using Microsoft.Extensions.ObjectPool;
+
+namespace Microsoft.Extensions.DependencyInjection;
+
+/// <summary>
+/// The deferred <see cref="DeferredWorkStateOwnerFactory"/> is injected as a scoped services and
+/// preserves the <see cref="DeferredWorkStateOwner"/> instance it creates.
+///
+/// This is done so that the executions running on one service scope share the deferred execution
+/// state between each other.
+///
+/// <see cref="DeferredWorkStateOwner"/> is disposable and will be disposed with the request scope.
+/// </summary>
+internal sealed class DeferredWorkStateOwnerFactory : IFactory<DeferredWorkStateOwner>
+{
+    private readonly object _sync = new();
+    private readonly ObjectPool<DeferredWorkState> _pool;
+    private DeferredWorkStateOwner? _owner;
+
+    public DeferredWorkStateOwnerFactory(ObjectPool<DeferredWorkState> pool)
+    {
+        _pool = pool;
+    }
+
+    public DeferredWorkStateOwner Create()
+    {
+        if (_owner is null)
+        {
+            lock (_sync)
+            {
+                if (_owner is null)
+                {
+                    _owner = new DeferredWorkStateOwner(_pool);
+                }
+            }
+        }
+
+        return _owner;
+    }
+}

src/HotChocolate/Core/src/Execution/DependencyInjection/Factories/OperationContextFactory.cs

@@ -0,0 +1,56 @@
+using System;
+using Microsoft.Extensions.ObjectPool;
+using HotChocolate;
+using HotChocolate.Execution;
+using HotChocolate.Execution.DependencyInjection;
+using HotChocolate.Execution.Processing;
+using HotChocolate.Execution.Processing.Tasks;
+using HotChocolate.Utilities;
+
+namespace Microsoft.Extensions.DependencyInjection;
+
+/// <summary>
+/// The <see cref="OperationContextFactory"/> creates new instances of
+/// <see cref="OperationContext"/>.
+///
+/// Operation context lifetime is managed by the OperationContext pool and
+/// the execution pipeline.
+///
+/// The lifetime MUST NOT be managed or tracked by the DI container.
+///
+/// The <see cref="OperationContextFactory"/> MUST be a singleton.
+/// </summary>
+internal sealed class OperationContextFactory : IFactory<OperationContext>
+{
+    private readonly IFactory<ResolverTask> _resolverTaskFactory;
+    private readonly ObjectPool<PathSegmentBuffer<IndexerPathSegment>> _indexerPathPool;
+    private readonly ObjectPool<PathSegmentBuffer<NamePathSegment>> _namePathPool;
+    private readonly ResultPool _resultPool;
+    private readonly ITypeConverter _typeConverter;
+
+    public OperationContextFactory(
+        IFactory<ResolverTask> resolverTaskFactory,
+        ObjectPool<PathSegmentBuffer<IndexerPathSegment>> indexerPathPool,
+        ObjectPool<PathSegmentBuffer<NamePathSegment>> namePathPool,
+        ResultPool resultPool,
+        ITypeConverter typeConverter)
+    {
+        _resolverTaskFactory = resolverTaskFactory ??
+            throw new ArgumentNullException(nameof(resolverTaskFactory));
+        _indexerPathPool = indexerPathPool ??
+            throw new ArgumentNullException(nameof(indexerPathPool));
+        _namePathPool = namePathPool ??
+            throw new ArgumentNullException(nameof(namePathPool));
+        _resultPool = resultPool ??
+            throw new ArgumentNullException(nameof(resultPool));
+        _typeConverter = typeConverter ??
+            throw new ArgumentNullException(nameof(typeConverter));
+    }
+
+    public OperationContext Create()
+        => new OperationContext(
+            _resolverTaskFactory,
+            new PooledPathFactory(_indexerPathPool, _namePathPool),
+            new ResultBuilder(_resultPool),
+            _typeConverter);
+}

src/HotChocolate/Core/src/Execution/DependencyInjection/Factories/OperationContextOwnerFactory.cs

@@ -0,0 +1,25 @@
+using HotChocolate.Execution.DependencyInjection;
+using HotChocolate.Execution.Processing;
+using Microsoft.Extensions.ObjectPool;
+
+namespace Microsoft.Extensions.DependencyInjection;
+
+/// <summary>
+/// The <see cref="OperationContextOwnerFactory"/> creates new instances of
+/// <see cref="OperationContextOwner"/>.
+///
+/// Each create will create a new instance that MUST NOT be managed or tracked by the DI container.
+///
+/// The <see cref="OperationContextOwnerFactory"/> MUST be a singleton.
+/// </summary>
+internal sealed class OperationContextOwnerFactory : IFactory<OperationContextOwner>
+{
+    private readonly ObjectPool<OperationContext> _pool;
+
+    public OperationContextOwnerFactory(ObjectPool<OperationContext> pool)
+    {
+        _pool = pool;
+    }
+
+    public OperationContextOwner Create() => new(_pool);
+}

src/HotChocolate/Core/src/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs

@@ -50,7 +50,6 @@ internal static IServiceCollection TryAddResultPool(
         services.TryAddSingleton(_ => new ObjectResultPool(maximumRetained, maximumArrayCapacity));
         services.TryAddSingleton(_ => new ListResultPool(maximumRetained, maximumArrayCapacity));
         services.TryAddSingleton<ResultPool>();
-        services.TryAddTransient<ResultBuilder>();
         return services;
     }
 
@@ -76,10 +75,6 @@ internal static IServiceCollection TryAddPathSegmentPool(
             _ => new IndexerPathSegmentPool(maximumRetained));
         services.TryAddSingleton<ObjectPool<PathSegmentBuffer<NamePathSegment>>>(
             _ => new NamePathSegmentPool(maximumRetained));
-        services.TryAddTransient(
-            sp => new PooledPathFactory(
-                sp.GetRequiredService<ObjectPool<PathSegmentBuffer<IndexerPathSegment>>>(),
-                sp.GetRequiredService<ObjectPool<PathSegmentBuffer<NamePathSegment>>>()));
         return services;
     }
 
@@ -94,47 +89,33 @@ internal static IServiceCollection TryAddOperationCompilerPool(
     internal static IServiceCollection TryAddOperationContextPool(
         this IServiceCollection services)
     {
-        services.TryAddSingleton(sp =>
-        {
-            var provider = sp.GetRequiredService<ObjectPoolProvider>();
-            var policy = new OperationContextPooledObjectPolicy(
-                sp.GetRequiredService<IFactory<OperationContext>>());
-            return provider.Create(policy);
-        });
-
-        services.TryAddTransient<OperationContext>();
-
-        services.TryAddTransient(
-            sp => sp.GetRequiredService<ObjectPool<OperationContext>>().GetOwner());
-
-        services.TryAddSingleton<IFactory<OperationContextOwner>>(
-            sp => new ServiceFactory<OperationContextOwner>(sp));
+        services.TryAddSingleton<IFactory<OperationContext>, OperationContextFactory>();
+        services.TryAddSingleton<IFactory<OperationContextOwner>, OperationContextOwnerFactory>();
 
-        services.TryAddSingleton<IFactory<OperationContext>>(
-            sp => new ServiceFactory<OperationContext>(sp));
+        services.TryAddSingleton(
+            sp =>
+            {
+                var provider = sp.GetRequiredService<ObjectPoolProvider>();
+                var policy = new OperationContextPooledObjectPolicy(
+                    sp.GetRequiredService<IFactory<OperationContext>>());
+                return provider.Create(policy);
+            });
 
         return services;
     }
 
     internal static IServiceCollection TryAddDeferredWorkStatePool(
         this IServiceCollection services)
     {
-        services.TryAddSingleton(sp =>
-        {
-            var provider = sp.GetRequiredService<ObjectPoolProvider>();
-            var policy = new DeferredWorkStatePooledObjectPolicy();
-            return provider.Create(policy);
-        });
-
-        services.TryAddScoped(sp =>
-        {
-            var pool = sp.GetRequiredService<ObjectPool<DeferredWorkState>>();
-            var state = pool.Get();
-            return new DeferredWorkStateOwner(state, pool);
-        });
+        services.TryAddSingleton(
+            sp =>
+            {
+                var provider = sp.GetRequiredService<ObjectPoolProvider>();
+                var policy = new DeferredWorkStatePooledObjectPolicy();
+                return provider.Create(policy);
+            });
 
-        services.TryAddScoped<IFactory<DeferredWorkStateOwner>>(
-            sp => new ServiceFactory<DeferredWorkStateOwner>(sp));
+        services.TryAddScoped<IFactory<DeferredWorkStateOwner>, DeferredWorkStateOwnerFactory>();
 
         return services;
     }

src/HotChocolate/Core/src/Execution/Processing/DeferredWorkStateOwner.cs

@@ -6,13 +6,13 @@ namespace HotChocolate.Execution.Processing;
 
 internal sealed class DeferredWorkStateOwner : IDisposable
 {
-    private readonly ObjectPool<DeferredWorkState> _statePool;
+    private readonly ObjectPool<DeferredWorkState> _pool;
     private int _disposed;
 
-    public DeferredWorkStateOwner(DeferredWorkState state, ObjectPool<DeferredWorkState> statePool)
+    public DeferredWorkStateOwner( ObjectPool<DeferredWorkState> pool)
     {
-        State = state ?? throw new ArgumentNullException(nameof(state));
-        _statePool = statePool ?? throw new ArgumentNullException(nameof(statePool));
+        _pool = pool ?? throw new ArgumentNullException(nameof(pool));
+        State = pool.Get();
     }
 
     public DeferredWorkState State { get; }
@@ -21,7 +21,7 @@ public void Dispose()
     {
         if (_disposed == 0 && Interlocked.CompareExchange(ref _disposed, 1, 0) == 0)
         {
-            _statePool.Return(State);
+            _pool.Return(State);
         }
     }
 }

src/HotChocolate/Core/src/Execution/Processing/OperationContextOwner.cs

@@ -1,23 +1,35 @@
 using System;
-using System.Threading;
 using Microsoft.Extensions.ObjectPool;
+using static System.Threading.Interlocked;
 
 namespace HotChocolate.Execution.Processing;
 
+/// <summary>
+/// The operation context owner abstracts the interaction of resolving of
+/// an <see cref="OperationContext"/> instance from its pool and returning to to
+/// the pool through the implementation of <see cref="IDisposable"/>.
+///
+/// In some cases its desirable to not call dispose and abandon a pooled
+/// <see cref="OperationContext"/>.
+/// </summary>
 internal sealed class OperationContextOwner : IDisposable
 {
-    private readonly OperationContext _operationContext;
-    private readonly ObjectPool<OperationContext> _operationContextPool;
+    private readonly ObjectPool<OperationContext> _pool;
+    private readonly OperationContext _context;
     private int _disposed;
 
-    public OperationContextOwner(
-        OperationContext operationContext,
-        ObjectPool<OperationContext> operationContextPool)
+    public OperationContextOwner(ObjectPool<OperationContext> operationContextPool)
     {
-        _operationContext = operationContext;
-        _operationContextPool = operationContextPool;
+        _pool = operationContextPool;
+        _context = operationContextPool.Get();
     }
 
+    /// <summary>
+    /// Gets the pooled <see cref="OperationContext"/>.
+    /// </summary>
+    /// <exception cref="ObjectDisposedException">
+    /// The operation context was already return to the pool.
+    /// </exception>
     public OperationContext OperationContext
     {
         get
@@ -27,15 +39,18 @@ public OperationContext OperationContext
                 throw new ObjectDisposedException(nameof(OperationContextOwner));
             }
 
-            return _operationContext;
+            return _context;
         }
     }
 
+    /// <summary>
+    /// Returns the <see cref="OperationContext"/> back to its pool.
+    /// </summary>
     public void Dispose()
     {
-        if (_disposed is 0 && Interlocked.CompareExchange(ref _disposed, 1, 0) == 0)
+        if (_disposed == 0 && CompareExchange(ref _disposed, 1, 0) == 0)
         {
-            _operationContextPool.Return(_operationContext);
+            _pool.Return(_context);
         }
     }
 }