Reworked RequestContext pooling (#4141)

michaelstaib · web-flow · commit cee5226f5ccb · 2021-08-26T13:33:48.000+02:00
diff --git a/src/HotChocolate/Core/src/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs b/src/HotChocolate/Core/src/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs
@@ -5,7 +5,6 @@
 using Microsoft.Extensions.Options;
 using GreenDonut;
 using HotChocolate.Execution;
-using HotChocolate.Execution.Batching;
 using HotChocolate.Execution.Caching;
 using HotChocolate.Execution.Configuration;
 using HotChocolate.Execution.Internal;
@@ -66,19 +65,17 @@ internal static IServiceCollection TryAddResolverTaskPool(
         }
 
         internal static IServiceCollection TryAddOperationContextPool(
-            this IServiceCollection services,
-            int maximumRetained = -1)
+            this IServiceCollection services)
         {
-            if (maximumRetained < 1)
+            services.TryAddSingleton<ObjectPool<OperationContext>>(sp =>
             {
-                maximumRetained = Environment.ProcessorCount * 2;
-            }
+                ObjectPoolProvider provider = sp.GetRequiredService<ObjectPoolProvider>();
+                var policy = new OperationContextPooledObjectPolicy(
+                    sp.GetRequiredService<ObjectPool<ResolverTask>>(),
+                    sp.GetRequiredService<ResultPool>());
+                return provider.Create(policy);
+            });
 
-            services.TryAddTransient<OperationContext>();
-            services.TryAddSingleton<ObjectPool<OperationContext>>(
-                sp => new DefaultObjectPool<OperationContext>(
-                    new OperationContextPoolPolicy(sp.GetRequiredService<OperationContext>),
-                    maximumRetained));
             return services;
         }
 
@@ -192,5 +189,30 @@ internal static IServiceCollection AddParameterExpressionBuilder<T>(
             services.AddSingleton<IParameterExpressionBuilder>(factory);
             return services;
         }
+
+        private class OperationContextPooledObjectPolicy : PooledObjectPolicy<OperationContext>
+        {
+            private readonly ObjectPool<ResolverTask> _resolverTaskPool;
+            private readonly ResultPool _resultPool;
+
+            public OperationContextPooledObjectPolicy(
+                ObjectPool<ResolverTask> resolverTaskPool,
+                ResultPool resultPool)
+            {
+                _resolverTaskPool = resolverTaskPool ??
+                    throw new ArgumentNullException(nameof(resolverTaskPool));
+                _resultPool = resultPool ??
+                    throw new ArgumentNullException(nameof(resultPool));
+            }
+
+            public override OperationContext Create()
+                => new(_resolverTaskPool, _resultPool);
+
+            public override bool Return(OperationContext obj)
+            {
+                obj.Clean();
+                return true;
+            }
+        }
     }
 }
diff --git a/src/HotChocolate/Core/src/Execution/RequestExecutor.cs b/src/HotChocolate/Core/src/Execution/RequestExecutor.cs
@@ -3,36 +3,27 @@
 using System.Threading;
 using System.Threading.Tasks;
 using HotChocolate.Execution.Batching;
-using HotChocolate.Execution.Instrumentation;
-using HotChocolate.Execution.Processing;
-using HotChocolate.Utilities;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.ObjectPool;
 
 namespace HotChocolate.Execution
 {
     internal sealed class RequestExecutor : IRequestExecutor
     {
         private readonly DefaultRequestContextAccessor _requestContextAccessor;
         private readonly IServiceProvider _applicationServices;
-        private readonly IErrorHandler _errorHandler;
-        private readonly ITypeConverter _converter;
-        private readonly IActivator _activator;
-        private readonly IDiagnosticEvents _diagnosticEvents;
         private readonly RequestDelegate _requestDelegate;
         private readonly BatchExecutor _batchExecutor;
-        private RequestContext? _pooledContext;
+        private readonly ObjectPool<RequestContext> _contextPool;
 
         public RequestExecutor(
             ISchema schema,
             DefaultRequestContextAccessor requestContextAccessor,
             IServiceProvider applicationServices,
             IServiceProvider executorServices,
-            IErrorHandler errorHandler,
-            ITypeConverter converter,
-            IActivator activator,
-            IDiagnosticEvents diagnosticEvents,
             RequestDelegate requestDelegate,
             BatchExecutor batchExecutor,
+            ObjectPool<RequestContext> contextPool,
             ulong version)
         {
             Schema = schema ??
@@ -43,18 +34,12 @@ public RequestExecutor(
                 throw new ArgumentNullException(nameof(applicationServices));
             Services = executorServices ??
                 throw new ArgumentNullException(nameof(executorServices));
-            _errorHandler = errorHandler ??
-                throw new ArgumentNullException(nameof(errorHandler));
-            _converter = converter ??
-                throw new ArgumentNullException(nameof(converter));
-            _activator = activator ??
-                throw new ArgumentNullException(nameof(activator));
-            _diagnosticEvents = diagnosticEvents ??
-                throw new ArgumentNullException(nameof(diagnosticEvents));
             _requestDelegate = requestDelegate ??
                 throw new ArgumentNullException(nameof(requestDelegate));
             _batchExecutor = batchExecutor ??
                 throw new ArgumentNullException(nameof(batchExecutor));
+            _contextPool = contextPool ??
+                throw new ArgumentNullException(nameof(contextPool));
             Version = version;
         }
 
@@ -81,18 +66,11 @@ public async Task<IExecutionResult> ExecuteAsync(
                 ? request.Services!
                 : scope.ServiceProvider;
 
-            RequestContext? context = Interlocked.Exchange(ref _pooledContext, null);
+            RequestContext context = _contextPool.Get();
 
             try
             {
-                context ??= new RequestContext(
-                    Schema,
-                    Version,
-                    _errorHandler,
-                    _converter,
-                    _activator,
-                    _diagnosticEvents) { RequestAborted = cancellationToken };
-
+                context.RequestAborted = cancellationToken;
                 context.Initialize(request, services);
 
                 _requestContextAccessor.RequestContext = context;
@@ -124,8 +102,7 @@ public async Task<IExecutionResult> ExecuteAsync(
             }
             finally
             {
-                context!.Reset();
-                Interlocked.Exchange(ref _pooledContext, context);
+                _contextPool.Return(context);
                 scope?.Dispose();
             }
         }
diff --git a/src/HotChocolate/Core/src/Execution/RequestExecutorResolver.cs b/src/HotChocolate/Core/src/Execution/RequestExecutorResolver.cs
@@ -16,14 +16,16 @@
 using HotChocolate.Types.Descriptors;
 using HotChocolate.Types.Descriptors.Definitions;
 using HotChocolate.Utilities;
+using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Extensions.ObjectPool;
 using static HotChocolate.Execution.ThrowHelper;
 
 namespace HotChocolate.Execution
 {
     internal sealed class RequestExecutorResolver
         : IRequestExecutorResolver
-        , IInternalRequestExecutorResolver
-        , IDisposable
+            , IInternalRequestExecutorResolver
+            , IDisposable
     {
         private readonly SemaphoreSlim _semaphore = new(1, 1);
         private readonly ConcurrentDictionary<string, RegisteredExecutor> _executors = new();
@@ -39,9 +41,9 @@ public RequestExecutorResolver(
             IServiceProvider serviceProvider)
         {
             _optionsMonitor = optionsMonitor ??
-                throw new ArgumentNullException(nameof(optionsMonitor));
+                              throw new ArgumentNullException(nameof(optionsMonitor));
             _applicationServices = serviceProvider ??
-                throw new ArgumentNullException(nameof(serviceProvider));
+                                   throw new ArgumentNullException(nameof(serviceProvider));
             _optionsMonitor.OnChange(EvictRequestExecutor);
         }
 
@@ -145,8 +147,8 @@ private void BeginRunEvictionEvents(RegisteredExecutor registeredExecutor)
                         if (action.AsyncAction is { } task)
                         {
                             await task.Invoke(
-                                registeredExecutor.Executor,
-                                CancellationToken.None)
+                                    registeredExecutor.Executor,
+                                    CancellationToken.None)
                                 .ConfigureAwait(false);
                         }
                     }
@@ -241,18 +243,30 @@ await CreateExecutorOptionsAsync(options, cancellationToken)
                     _applicationServices.GetRequiredService<ITypeConverter>(),
                     _applicationServices.GetRequiredService<InputFormatter>()));
 
+            serviceCollection.TryAddSingleton<ObjectPoolProvider, DefaultObjectPoolProvider>();
+
+            serviceCollection.TryAddSingleton<ObjectPool<RequestContext>>(sp =>
+            {
+                ObjectPoolProvider provider = sp.GetRequiredService<ObjectPoolProvider>();
+                var policy = new RequestContextPooledObjectPolicy(
+                    sp.GetRequiredService<ISchema>(),
+                    sp.GetRequiredService<IErrorHandler>(),
+                    _applicationServices.GetRequiredService<ITypeConverter>(),
+                    sp.GetRequiredService<IActivator>(),
+                    sp.GetRequiredService<IDiagnosticEvents>(),
+                    version);
+                return provider.Create(policy);
+            });
+
             serviceCollection.AddSingleton<IRequestExecutor>(
                 sp => new RequestExecutor(
                     sp.GetRequiredService<ISchema>(),
                     _applicationServices.GetRequiredService<DefaultRequestContextAccessor>(),
                     _applicationServices,
                     sp,
-                    sp.GetRequiredService<IErrorHandler>(),
-                    _applicationServices.GetRequiredService<ITypeConverter>(),
-                    sp.GetRequiredService<IActivator>(),
-                    sp.GetRequiredService<IDiagnosticEvents>(),
                     sp.GetRequiredService<RequestDelegate>(),
                     sp.GetRequiredService<BatchExecutor>(),
+                    sp.GetRequiredService<ObjectPool<RequestContext>>(),
                     version));
 
             foreach (Action<IServiceCollection> configureServices in options.SchemaServices)
@@ -265,12 +279,12 @@ await CreateExecutorOptionsAsync(options, cancellationToken)
 
             lazy.Schema =
                 await CreateSchemaAsync(
-                    schemaName,
-                    options,
-                    executorOptions,
-                    combinedServices,
-                    typeModuleChangeMonitor,
-                    cancellationToken)
+                        schemaName,
+                        options,
+                        executorOptions,
+                        combinedServices,
+                        typeModuleChangeMonitor,
+                        cancellationToken)
                     .ConfigureAwait(false);
 
             return schemaServices;
@@ -342,7 +356,7 @@ private static async ValueTask<RequestExecutorOptions> CreateExecutorOptionsAsyn
             CancellationToken cancellationToken)
         {
             RequestExecutorOptions executorOptions = options.RequestExecutorOptions ??
-                new RequestExecutorOptions();
+                                                     new RequestExecutorOptions();
 
             foreach (RequestExecutorOptionsAction action in options.RequestExecutorOptionsActions)
             {
@@ -531,5 +545,51 @@ await typeModule.CreateTypesAsync(_context, cancellationToken)
                 }
             }
         }
+
+        private class RequestContextPooledObjectPolicy : PooledObjectPolicy<RequestContext>
+        {
+            private readonly ISchema _schema;
+            private readonly ulong _executorVersion;
+            private readonly IErrorHandler _errorHandler;
+            private readonly ITypeConverter _converter;
+            private readonly IActivator _activator;
+            private readonly IDiagnosticEvents _diagnosticEvents;
+
+            public RequestContextPooledObjectPolicy(
+                ISchema schema,
+                IErrorHandler errorHandler,
+                ITypeConverter converter,
+                IActivator activator,
+                IDiagnosticEvents diagnosticEvents,
+                ulong executorVersion)
+            {
+                _schema = schema ??
+                    throw new ArgumentNullException(nameof(schema));
+                _errorHandler = errorHandler ??
+                    throw new ArgumentNullException(nameof(errorHandler));
+                _converter = converter ??
+                    throw new ArgumentNullException(nameof(converter));
+                _activator = activator ??
+                    throw new ArgumentNullException(nameof(activator));
+                _diagnosticEvents = diagnosticEvents ??
+                    throw new ArgumentNullException(nameof(diagnosticEvents));
+                _executorVersion = executorVersion;
+            }
+
+
+            public override RequestContext Create()
+                => new(_schema,
+                    _executorVersion,
+                    _errorHandler,
+                    _converter,
+                    _activator,
+                    _diagnosticEvents);
+
+            public override bool Return(RequestContext obj)
+            {
+                obj.Reset();
+                return true;
+            }
+        }
     }
 }
