Abandon OperationContext when cancellations happen.

ChilliCream · Jan 23, 2023 · 8ad9a73 · 8ad9a73
1 parent ee27855
commit 8ad9a73
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/src/HotChocolate/Core/src/Execution/Pipeline/OperationExecutionMiddleware.cs b/src/HotChocolate/Core/src/Execution/Pipeline/OperationExecutionMiddleware.cs
@@ -97,7 +97,10 @@ await ExecuteOperationAsync(context, batchDispatcher, context.Operation)
             try
             {
                 await ExecuteQueryOrMutationAsync(
-                    context, batchDispatcher, operation, operationContext)
+                        context,
+                        batchDispatcher,
+                        operation,
+                        operationContext)
                     .ConfigureAwait(false);
 
                 if (operationContext.DeferredScheduler.HasResults &&
@@ -115,6 +118,12 @@ await ExecuteOperationAsync(context, batchDispatcher, context.Operation)
 
                 await _next(context).ConfigureAwait(false);
             }
+            catch (OperationCanceledException)
+            {
+                // if an operation is canceled we will abandon the the rented operation context
+                // to ensure that that abandoned tasks to not leak execution into new operations.
+                operationContextOwner = null;
+            }
             finally
             {
                 operationContextOwner?.Dispose();

diff --git a/src/HotChocolate/Core/src/Execution/Processing/SubscriptionExecutor.Subscription.cs b/src/HotChocolate/Core/src/Execution/Processing/SubscriptionExecutor.Subscription.cs
@@ -153,7 +153,7 @@ private async Task<IQueryResult> OnEvent(object payload)
 
                 // we store the event payload on the scoped context so that it is accessible
                 // in the resolvers.
-                var scopedContext =
+                var scopedContextData =
                     _scopedContextData.SetItem(WellKnownContextData.EventMessage, payload);
 
                 // next we resolve the subscription instance.
@@ -179,13 +179,21 @@ private async Task<IQueryResult> OnEvent(object payload)
                     payload);
 
                 var result = await _queryExecutor
-                    .ExecuteAsync(operationContext, scopedContext)
+                    .ExecuteAsync(operationContext, scopedContextData)
                     .ConfigureAwait(false);
 
                 _diagnosticEvents.SubscriptionEventResult(new(this, payload), result);
 
                 return result;
             }
+            catch (OperationCanceledException ex)
+            {
+                operationContext = null;
+                _diagnosticEvents.SubscriptionEventError(
+                    new SubscriptionEventContext(this, payload),
+                    ex);
+                throw;
+            }
             catch (Exception ex)
             {
                 _diagnosticEvents.SubscriptionEventError(
@@ -195,7 +203,13 @@ private async Task<IQueryResult> OnEvent(object payload)
             }
             finally
             {
-                _operationContextPool.Return(operationContext);
+                // if the operation context is null a cancellation has happened and we will
+                // abandon the operation context in order to not have leakage into the
+                // new operations.
+                if (operationContext is not null)
+                {
+                    _operationContextPool.Return(operationContext);
+                }
             }
         }