From 711f5791f8b8499a6a4467c2c5e1cb5a8648607f Mon Sep 17 00:00:00 2001 From: Phillip Kruger Date: Mon, 8 Aug 2022 12:23:38 +1000 Subject: [PATCH] Also cleanup GraphQL Context Signed-off-by: Phillip Kruger --- .../SmallRyeGraphQLAbstractHandler.java | 17 +++++++++++++--- .../datafetcher/AbstractAsyncDataFetcher.java | 20 +++++++++++++++++-- .../QuarkusDefaultDataFetcher.java | 10 ++++++++-- 3 files changed, 40 insertions(+), 7 deletions(-) diff --git a/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/SmallRyeGraphQLAbstractHandler.java b/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/SmallRyeGraphQLAbstractHandler.java index 1bacbca3126c1..3e011fe3a5cb9 100644 --- a/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/SmallRyeGraphQLAbstractHandler.java +++ b/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/SmallRyeGraphQLAbstractHandler.java @@ -19,6 +19,7 @@ import io.quarkus.vertx.http.runtime.CurrentVertxRequest; import io.quarkus.vertx.http.runtime.security.QuarkusHttpUser; import io.smallrye.graphql.execution.ExecutionService; +import io.smallrye.graphql.execution.context.SmallRyeContextManager; import io.vertx.core.Handler; import io.vertx.core.MultiMap; import io.vertx.ext.web.RoutingContext; @@ -56,7 +57,7 @@ private Handler createCurrentManagedContextTerminationHandler() { return new Handler<>() { @Override public void handle(Object e) { - currentManagedContext.terminate(); + terminate(); } }; } @@ -73,9 +74,9 @@ public void handle(final RoutingContext ctx) { } try { handleWithIdentity(ctx); - currentManagedContext.deactivate(); + deactivate(); } catch (Throwable t) { - currentManagedContext.terminate(); + terminate(); throw t; } } @@ -128,4 +129,14 @@ private Map> getHeaders(RoutingContext ctx) { } return h; } + + private void deactivate() { + SmallRyeContextManager.clearCurrentSmallRyeContext(); + currentManagedContext.deactivate(); + } + + private void terminate() { + SmallRyeContextManager.clearCurrentSmallRyeContext(); + currentManagedContext.terminate(); + } } diff --git a/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/AbstractAsyncDataFetcher.java b/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/AbstractAsyncDataFetcher.java index b96e890b7ee5e..037e3efec9815 100644 --- a/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/AbstractAsyncDataFetcher.java +++ b/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/AbstractAsyncDataFetcher.java @@ -10,6 +10,7 @@ import io.quarkus.arc.Arc; import io.quarkus.arc.ManagedContext; import io.smallrye.graphql.SmallRyeGraphQLServerMessages; +import io.smallrye.graphql.execution.context.SmallRyeContextManager; import io.smallrye.graphql.execution.datafetcher.AbstractDataFetcher; import io.smallrye.graphql.schema.model.Operation; import io.smallrye.graphql.schema.model.Type; @@ -36,6 +37,11 @@ protected O invokeAndTransform( return (O) uni .onItemOrFailure() .transformToUni((result, throwable, emitter) -> { + + emitter.onTermination(() -> { + deactivate(requestContext); + }); + if (throwable != null) { eventEmitter.fireOnDataFetchError(dfe.getExecutionId().toString(), throwable); if (throwable instanceof GraphQLException) { @@ -55,16 +61,26 @@ protected O invokeAndTransform( te.appendDataFetcherResult(resultBuilder, dfe); } } - emitter.complete(resultBuilder.build()); }) + .onCancellation().invoke(() -> { + deactivate(requestContext); + }) + .onTermination().invoke(() -> { + deactivate(requestContext); + }) .subscribe() .asCompletionStage(); } finally { - requestContext.deactivate(); + deactivate(requestContext); } } + private void deactivate(ManagedContext requestContext) { + SmallRyeContextManager.clearCurrentSmallRyeContext(); + requestContext.deactivate(); + } + protected abstract Uni handleUserMethodCall(DataFetchingEnvironment dfe, final Object[] transformedArguments) throws Exception; diff --git a/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/QuarkusDefaultDataFetcher.java b/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/QuarkusDefaultDataFetcher.java index 0abc666a65e18..b08595baaa2ef 100644 --- a/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/QuarkusDefaultDataFetcher.java +++ b/extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/spi/datafetcher/QuarkusDefaultDataFetcher.java @@ -12,6 +12,7 @@ import io.quarkus.arc.Arc; import io.quarkus.arc.ManagedContext; import io.smallrye.context.SmallRyeThreadContext; +import io.smallrye.graphql.execution.context.SmallRyeContextManager; import io.smallrye.graphql.execution.datafetcher.DefaultDataFetcher; import io.smallrye.graphql.schema.model.Operation; import io.smallrye.graphql.schema.model.Type; @@ -40,7 +41,7 @@ public T invokeAndTransform(DataFetchingEnvironment dfe, DataFetcherResult.B return invokeAndTransformBlocking(dfe, resultBuilder, transformedArguments, vc); } } finally { - requestContext.deactivate(); + deactivate(requestContext); } } @@ -57,7 +58,7 @@ public CompletionStage> invokeBatch(DataFetchingEnvironment dfe, Object[ return invokeBatchBlocking(dfe, arguments, vc); } } finally { - requestContext.deactivate(); + deactivate(requestContext); } } @@ -113,4 +114,9 @@ private CompletionStage> invokeBatchBlocking(DataFetchingEnvironment dfe private boolean runBlocking(DataFetchingEnvironment dfe) { return dfe.getGraphQlContext().get("runBlocking"); } + + private void deactivate(ManagedContext requestContext) { + SmallRyeContextManager.clearCurrentSmallRyeContext(); + requestContext.deactivate(); + } }