Fixed issue with global state when using deferred execution and scoped services. (#7611)

Author: michaelstaib

src/HotChocolate/Core/src/Execution/IRequestContext.cs

@@ -113,9 +113,4 @@ public interface IRequestContext : IHasContextData
     /// Gets or sets an unexpected execution exception.
     /// </summary>
     Exception? Exception { get; set; }
-
-    /// <summary>
-    /// Clones the request context.
-    /// </summary>
-    IRequestContext Clone();
 }

src/HotChocolate/Core/src/Execution/Pipeline/OperationExecutionMiddleware.cs

@@ -89,12 +89,8 @@ private async Task ExecuteOperationRequestAsync(
     {
         if (operation.Definition.Operation is OperationType.Subscription)
         {
-            // since the request context is pooled we need to clone the context for
-            // long-running executions.
-            var cloned = context.Clone();
-
             context.Result = await _subscriptionExecutor
-                .ExecuteAsync(cloned, () => GetQueryRootValue(cloned))
+                .ExecuteAsync(context, () => GetQueryRootValue(context))
                 .ConfigureAwait(false);
         }
         else

src/HotChocolate/Core/src/Execution/RequestContext.cs

@@ -66,38 +66,6 @@ public DocumentValidatorResult? ValidationResult
 
     public Exception? Exception { get; set; }
 
-    public IRequestContext Clone()
-    {
-        var cloned = new RequestContext(
-            Schema,
-            ExecutorVersion,
-            ErrorHandler,
-            DiagnosticEvents)
-        {
-            Request = Request,
-            Services = Services,
-            RequestAborted = RequestAborted,
-            DocumentId = DocumentId,
-            DocumentHash = DocumentHash,
-            IsCachedDocument = IsCachedDocument,
-            Document = Document,
-            ValidationResult = ValidationResult,
-            OperationId = OperationId,
-            Operation = Operation,
-            Variables = Variables,
-            Result = Result,
-            Exception = Exception,
-            RequestIndex = RequestIndex,
-        };
-
-        foreach (var item in _contextData)
-        {
-            cloned._contextData.TryAdd(item.Key, item.Value);
-        }
-
-        return cloned;
-    }
-
     public void Initialize(IOperationRequest request, IServiceProvider services)
     {
         Request = request;
@@ -139,4 +107,3 @@ public void Reset()
         RequestIndex = default;
     }
 }
-

src/HotChocolate/Core/src/Execution/RequestExecutor.cs

@@ -103,7 +103,7 @@ internal async Task<IExecutionResult> ExecuteAsync(
             services.InitializeDataLoaderScope();
         }
 
-        RequestContext? context = _contextPool.Get();
+        var context = _contextPool.Get();
 
         try
         {
@@ -123,7 +123,15 @@ internal async Task<IExecutionResult> ExecuteAsync(
 
             if (scope is null)
             {
-                return context.Result;
+                var localContext = context;
+
+                if (context.Result.IsStreamResult())
+                {
+                    context.Result.RegisterForCleanup(() => _contextPool.Return(localContext));
+                    context = null;
+                }
+
+                return localContext.Result;
             }
 
             if (context.Result.IsStreamResult())

src/HotChocolate/Core/test/Execution.Tests/DeferAndStreamTestSchema.cs

@@ -21,6 +21,22 @@ public static async Task<IRequestExecutor> CreateAsync()
             .BuildRequestExecutorAsync();
     }
 
+    public static IServiceProvider CreateServiceProvider()
+    {
+        return new ServiceCollection()
+            .AddGraphQL()
+            .AddQueryType<Query>()
+            .AddGlobalObjectIdentification()
+            .ModifyOptions(
+                o =>
+                {
+                    o.EnableDefer = true;
+                    o.EnableStream = true;
+                })
+            .Services
+            .BuildServiceProvider();
+    }
+
     public class Query
     {
         private readonly List<Person> _persons =

src/HotChocolate/Core/test/Execution.Tests/DeferTests.cs

@@ -1,4 +1,7 @@
 using CookieCrumble;
+using HotChocolate.AspNetCore;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace HotChocolate.Execution;
 
@@ -109,7 +112,11 @@ ... @defer(if: $defer) {
                         }
                     }
                     """)
-                .SetVariableValues(new Dictionary<string, object?> { { "defer", false }, })
+                .SetVariableValues(
+                    new Dictionary<string, object?>
+                    {
+                        { "defer", false },
+                    })
                 .Build());
 
         Assert.IsType<OperationResult>(result).MatchMarkdownSnapshot();
@@ -232,7 +239,11 @@ fragment Foo on Query {
                         }
                     }
                     """)
-                .SetVariableValues(new Dictionary<string, object?> { { "defer", false }, })
+                .SetVariableValues(
+                    new Dictionary<string, object?>
+                    {
+                        { "defer", false },
+                    })
                 .Build());
 
         Assert.IsType<OperationResult>(result).MatchMarkdownSnapshot();
@@ -301,6 +312,9 @@ ... @defer {
     [Fact]
     public async Task Ensure_GlobalState_Is_Passed_To_DeferContext_Single_Defer()
     {
+        // this test ensures that the request context is not recycled until the
+        // a stream is fully processed when no outer DI scope exists.
+
         // arrange
         var executor = await DeferAndStreamTestSchema.CreateAsync();
 
@@ -323,4 +337,49 @@ ... @defer {
 
         Assert.IsType<ResponseStream>(result).MatchMarkdownSnapshot();
     }
+
+    [Fact]
+    public async Task Ensure_GlobalState_Is_Passed_To_DeferContext_Single_Defer_2()
+    {
+        // this test ensures that the request context is not recycled until the
+        // a stream is fully processed when an outer DI scope exists.
+
+        // arrange
+        var services = DeferAndStreamTestSchema.CreateServiceProvider();
+        var executor = await services.GetRequestExecutorAsync();
+        await using var scope = services.CreateAsyncScope();
+
+        // act
+        var result = await executor.ExecuteAsync(
+            OperationRequestBuilder
+                .New()
+                .SetDocument(
+                    """
+                    {
+                        ... @defer {
+                            ensureState {
+                                state
+                            }
+                        }
+                    }
+                    """)
+                .SetGlobalState("requestState", "state 123")
+                .SetServices(scope.ServiceProvider)
+                .Build());
+
+        Assert.IsType<ResponseStream>(result).MatchMarkdownSnapshot();
+    }
+
+    private class StateRequestInterceptor : DefaultHttpRequestInterceptor
+    {
+        public override ValueTask OnCreateAsync(
+            HttpContext context,
+            IRequestExecutor requestExecutor,
+            OperationRequestBuilder requestBuilder,
+            CancellationToken cancellationToken)
+        {
+            requestBuilder.AddGlobalState("requestState", "bar");
+            return base.OnCreateAsync(context, requestExecutor, requestBuilder, cancellationToken);
+        }
+    }
 }

src/HotChocolate/Core/test/Execution.Tests/__snapshots__/DeferTests.Ensure_GlobalState_Is_Passed_To_DeferContext_Single_Defer_2.md

@@ -0,0 +1,12 @@
+# Ensure_GlobalState_Is_Passed_To_DeferContext_Single_Defer_2
+
+```text
+{
+  "data": {
+    "ensureState": {
+      "state": "state 123"
+    }
+  }
+}
+
+```