Allow middleware to register resolver task bound cleanup tasks. (#5294)

src/HotChocolate/AspNetCore/src/Transport.Sockets/MessageProcessor.cs

@@ -26,8 +26,8 @@ public async Task ProcessMessagesAsync(CancellationToken cancellationToken)
             while (true)
             {
                 SequencePosition? position;
-                ReadResult result = await _reader.ReadAsync(cancellationToken);
-                ReadOnlySequence<byte> buffer = result.Buffer;
+                var result = await _reader.ReadAsync(cancellationToken);
+                var buffer = result.Buffer;
 
                 do
                 {

src/HotChocolate/AspNetCore/src/Transport.Sockets/MessageReceiver.cs

@@ -44,7 +44,7 @@ public async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
         static void WriteEndOfMessage(IBufferWriter<byte> writer)
         {
             const int length = 1;
-            Span<byte> span = writer.GetSpan(length);
+            var span = writer.GetSpan(length);
             span[0] = EndOfText;
             writer.Advance(length);
         }

src/HotChocolate/Core/src/Execution/HotChocolate.Execution.csproj

@@ -165,6 +165,18 @@
     <Compile Update="Processing\Tasks\ResolverTask.Pooling.cs">
       <DependentUpon>ResolverTask.cs</DependentUpon>
     </Compile>
+    <Compile Update="Processing\OperationContext.Execution.cs">
+      <DependentUpon>OperationContext.cs</DependentUpon>
+    </Compile>
+    <Compile Update="Processing\OperationContext.Operation.cs">
+      <DependentUpon>OperationContext.cs</DependentUpon>
+    </Compile>
+    <Compile Update="Processing\OperationContext.Services.cs">
+      <DependentUpon>OperationContext.cs</DependentUpon>
+    </Compile>
+    <Compile Update="Processing\OperationContext.Utilities.cs">
+      <DependentUpon>OperationContext.cs</DependentUpon>
+    </Compile>
   </ItemGroup>
 
   <ItemGroup>

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

@@ -109,8 +109,10 @@ public void ReportError(Exception exception, Action<IErrorBuilder>? configure =
     public ValueTask<T> ResolveAsync<T>()
         => _middlewareContext.ResolveAsync<T>();
 
-    public void RegisterForCleanup(Action action)
-        => _middlewareContext.RegisterForCleanup(action);
+    public void RegisterForCleanup(
+        Func<ValueTask> action,
+        CleanAfter cleanAfter = CleanAfter.Resolver)
+        => _middlewareContext.RegisterForCleanup(action, cleanAfter);
 
     public IReadOnlyDictionary<string, ArgumentValue> ReplaceArguments(
         IReadOnlyDictionary<string, ArgumentValue> argumentValues)

src/HotChocolate/Core/src/Execution/Processing/MiddlewareContext.Arguments.cs

@@ -12,29 +12,6 @@ internal partial class MiddlewareContext
 {
     public IReadOnlyDictionary<string, ArgumentValue> Arguments { get; set; } = default!;
 
-    public T Argument<T>(string name)
-    {
-        if (string.IsNullOrEmpty(name))
-        {
-            throw new ArgumentNullException(nameof(name));
-        }
-
-        if (typeof(IValueNode).IsAssignableFrom(typeof(T)))
-        {
-            var literal = ArgumentLiteral<IValueNode>(name);
-
-            if (literal is T casted)
-            {
-                return casted;
-            }
-
-            throw ResolverContext_LiteralNotCompatible(
-                _selection.SyntaxNode, Path, name, typeof(T), literal.GetType());
-        }
-
-        return ArgumentValue<T>(name);
-    }
-
     public T ArgumentValue<T>(string name)
     {
         if (string.IsNullOrEmpty(name))

src/HotChocolate/Core/src/Execution/Processing/MiddlewareContext.Global.cs

@@ -11,6 +11,7 @@ namespace HotChocolate.Execution.Processing;
 
 internal partial class MiddlewareContext : IMiddlewareContext
 {
+    private readonly List<Func<ValueTask>> _cleanupTasks = new();
     private OperationContext _operationContext = default!;
     private IServiceProvider _services = default!;
     private InputParser _parser = default!;
@@ -25,8 +26,6 @@ public IServiceProvider Services
 
     public ISchema Schema => _operationContext.Schema;
 
-    public IObjectType RootType => _operationContext.Operation.RootType;
-
     public IOperation Operation => _operationContext.Operation;
 
     public IDictionary<string, object?> ContextData => _operationContext.ContextData;
@@ -35,6 +34,8 @@ public IServiceProvider Services
 
     public CancellationToken RequestAborted { get; private set; }
 
+    public bool HasCleanupTasks => _cleanupTasks.Count > 0;
+
     public IReadOnlyList<ISelection> GetSelections(
         IObjectType typeContext,
         ISelection? selection = null,
@@ -192,9 +193,33 @@ public object Service(Type service)
         return Services.GetRequiredService(service);
     }
 
-    public void RegisterForCleanup(Action action) =>
-        _operationContext.RegisterForCleanup(action);
+    public void RegisterForCleanup(
+        Func<ValueTask> action,
+        CleanAfter cleanAfter = CleanAfter.Resolver)
+    {
+        if (action is null)
+        {
+            throw new ArgumentNullException(nameof(action));
+        }
+
+        if (cleanAfter is CleanAfter.Request)
+        {
+            _operationContext.Result.RegisterForCleanup(action);
+        }
+        else
+        {
+            _cleanupTasks.Add(action);
+        }
+    }
+
+    public async ValueTask ExecuteCleanupTasksAsync()
+    {
+        foreach (var task in _cleanupTasks)
+        {
+            await task.Invoke().ConfigureAwait(false);
+        }
+    }
 
-    public T GetQueryRoot<T>() =>
-        _operationContext.GetQueryRoot<T>();
+    public T GetQueryRoot<T>() 
+        => _operationContext.GetQueryRoot<T>();
 }

src/HotChocolate/Core/src/Execution/Processing/MiddlewareContext.Pooling.cs

@@ -37,6 +37,7 @@ public MiddlewareContext()
     public void Clean()
     {
         _childContext.Clear();
+        _cleanupTasks.Clear();
         _operationContext = default!;
         _services = default!;
         _selection = default!;

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

@@ -43,18 +43,6 @@ public ResultBuilder Result
         }
     }
 
-    /// <summary>
-    /// Register cleanup tasks that will be executed after resolver execution is finished.
-    /// </summary>
-    /// <param name="action">
-    /// Cleanup action.
-    /// </param>
-    public void RegisterForCleanup(Action action)
-    {
-        AssertInitialized();
-        _cleanupActions.Add(action);
-    }
-
     public ResolverTask CreateResolverTask(
         ISelection selection,
         object? parent,

src/HotChocolate/Core/src/Execution/Processing/OperationContext.Pooling.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using System.Threading;
@@ -14,7 +13,6 @@ namespace HotChocolate.Execution.Processing;
 
 internal sealed partial class OperationContext
 {
-    private readonly ConcurrentBag<Action> _cleanupActions = new();
     private readonly IFactory<ResolverTask> _resolverTaskFactory;
     private readonly WorkScheduler _workScheduler;
     private readonly DeferredWorkScheduler _deferredWorkScheduler;
@@ -106,14 +104,6 @@ public void Clean()
     {
         if (_isInitialized)
         {
-            if (!_cleanupActions.IsEmpty)
-            {
-                while (_cleanupActions.TryTake(out var clean))
-                {
-                    clean();
-                }
-            }
-
             _pathFactory.Clear();
             _workScheduler.Clear();
             _resultHelper.Clear();

src/HotChocolate/Core/src/Execution/Processing/Result/ResultBuilder.Pooling.cs

@@ -31,6 +31,7 @@ public void Clear()
         _nonNullViolations.Clear();
         _extensions.Clear();
         _contextData.Clear();
+        _cleanupTasks.Clear();
 
         InitializeResult();
 

src/HotChocolate/Core/src/Execution/Processing/Result/ResultBuilder.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Threading.Tasks;
 using HotChocolate.Execution.Properties;
 
 namespace HotChocolate.Execution.Processing;
@@ -15,6 +16,7 @@ internal sealed partial class ResultBuilder
     private readonly object _syncExtensions = new();
     private readonly Dictionary<string, object?> _extensions = new();
     private readonly Dictionary<string, object?> _contextData = new();
+    private readonly List<Func<ValueTask>> _cleanupTasks = new();
 
     private ResultMemoryOwner _resultOwner = default!;
     private ObjectResult? _data;
@@ -43,6 +45,20 @@ public void SetContextData(string key, object? value)
         }
     }
 
+    /// <summary>
+    /// Register cleanup tasks that will be executed after resolver execution is finished.
+    /// </summary>
+    /// <param name="action">
+    /// Cleanup action.
+    /// </param>
+    public void RegisterForCleanup(Func<ValueTask> action)
+    {
+        lock (_syncExtensions)
+        {
+            _cleanupTasks.Add(action);
+        }
+    }
+
     public void SetPath(Path? path)
         => _path = path;
 
@@ -100,7 +116,10 @@ public IQueryResult BuildResult()
             CreateExtensionData(_contextData),
             _label,
             _path,
-            _hasNext
+            _hasNext,
+            _cleanupTasks.Count > 0
+                ? _cleanupTasks.ToArray()
+                : Array.Empty<Func<ValueTask>>()
         );
 
         if (_data is not null)

src/HotChocolate/Core/src/Execution/Processing/Tasks/ResolverTask.Execute.cs

@@ -52,6 +52,12 @@ private async Task ExecuteAsync(CancellationToken cancellationToken)
         finally
         {
             _operationContext.Scheduler.Complete(this);
+
+            if (_resolverContext.HasCleanupTasks)
+            {
+                await _resolverContext.ExecuteCleanupTasksAsync().ConfigureAwait(false);
+            }
+
             _objectPool.Return(this);
         }
     }

src/HotChocolate/Core/src/Execution/Processing/Tasks/ResolverTask.cs

@@ -71,6 +71,6 @@ public void BeginExecute(CancellationToken cancellationToken)
     }
 
     /// <inheritdoc />
-    public Task WaitForCompletionAsync(CancellationToken cancellationToken) =>
-        _task ?? Task.CompletedTask;
+    public Task WaitForCompletionAsync(CancellationToken cancellationToken)
+        => _task ?? Task.CompletedTask;
 }

src/HotChocolate/Core/src/Types/Resolvers/CleanAfter.cs

@@ -0,0 +1,19 @@
+#nullable enable
+
+namespace HotChocolate.Resolvers;
+
+/// <summary>
+/// Specifies when the cleanup task shall be applied.
+/// </summary>
+public enum CleanAfter
+{
+    /// <summary>
+    /// The cleanup task shall be applied after the resolver task is completed.
+    /// </summary>
+    Resolver,
+
+    /// <summary>
+    /// The cleanup task shall be when the query result is being disposed.
+    /// </summary>
+    Request
+}

src/HotChocolate/Core/src/Types/Resolvers/Expressions/Parameters/ServiceHelper.cs

@@ -43,19 +43,17 @@ internal static class ServiceHelper
         FieldMiddlewareDefinition serviceMiddleware =
             new(next => async context =>
                 {
-                    var services = context.Services;
-                    var objectPool = services.GetRequiredService<ObjectPool<TService>>();
+                    var objectPool = context.Services.GetRequiredService<ObjectPool<TService>>();
                     var service = objectPool.Get();
 
-                    try
+                    context.RegisterForCleanup(() =>
                     {
-                        context.SetLocalState(scopedServiceName, service);
-                        await next(context).ConfigureAwait(false);
-                    }
-                    finally
-                    {
-                        objectPool.Return(service!);
-                    }
+                        objectPool.Return(service);
+                        return default;
+                    });
+
+                    context.SetLocalState(scopedServiceName, service);
+                    await next(context).ConfigureAwait(false);
                 },
                 isRepeatable: true,
                 key: WellKnownMiddleware.PooledService);
@@ -86,7 +84,12 @@ internal static class ServiceHelper
             middleware = new FieldMiddlewareDefinition(
                 next => async context =>
                 {
-                    using var scope = context.Services.CreateScope();
+                    var scope = context.Services.CreateScope();
+                    context.RegisterForCleanup(() =>
+                    {
+                        scope.Dispose();
+                        return default;
+                    });
                     context.SetLocalState(WellKnownContextData.ResolverServiceScope, scope);
                     await next(context).ConfigureAwait(false);
                 },

src/HotChocolate/Core/src/Types/Resolvers/IMiddlewareContext.cs

@@ -8,10 +8,14 @@
 namespace HotChocolate.Resolvers;
 
 /// <summary>
-/// Encapsulates all resolver-specific information about the execution of an individual field selection.
+/// Encapsulates all resolver-specific information about the execution of
+/// an individual field selection.
 /// </summary>
 public interface IMiddlewareContext : IResolverContext
 {
+    /// <summary>
+    /// Gets or sets the value type hint.
+    /// </summary>
     IType? ValueType { get; set; }
 
     /// <summary>
@@ -40,9 +44,12 @@ public interface IMiddlewareContext : IResolverContext
     /// Register cleanup tasks that will be executed after resolver execution is finished.
     /// </summary>
     /// <param name="action">
-    /// Cleanup action.
+    /// The cleanup action.
+    /// </param>
+    /// <param name="cleanAfter">
+    /// Specifies when the cleanup task shall be applied.
     /// </param>
-    void RegisterForCleanup(Action action);
+    void RegisterForCleanup(Func<ValueTask> action, CleanAfter cleanAfter = CleanAfter.Resolver);
 
     /// <summary>
     /// Replaces the argument values for the current field execution pipeline.

src/HotChocolate/Data/src/EntityFramework/DbContextParameterExpressionBuilder.cs

@@ -52,10 +52,11 @@ public void ApplyConfiguration(ParameterInfo parameter, ObjectFieldDescriptor de
 
             case ServiceKind.Resolver:
                 ServiceExpressionHelper.ApplyConfiguration(parameter, descriptor, _kind);
-                ObjectFieldDefinition definition = descriptor.Extend().Definition;
-                FieldMiddlewareDefinition placeholderMiddleware =
-                    new(_ => _ => throw new NotSupportedException(), key: ToList);
-                FieldMiddlewareDefinition serviceMiddleware =
+                var definition = descriptor.Extend().Definition;
+                var placeholderMiddleware = new FieldMiddlewareDefinition(
+                    _ => _ => throw new NotSupportedException(),
+                    key: ToList);
+                var serviceMiddleware =
                     definition.MiddlewareDefinitions.Last(t => t.Key == PooledService);
                 var index = definition.MiddlewareDefinitions.IndexOf(serviceMiddleware) + 1;
                 definition.MiddlewareDefinitions.Insert(index, placeholderMiddleware);