From f23ca08e47a9731401f1c4472da0e6fe2bc1c413 Mon Sep 17 00:00:00 2001
From: Phillip Kruger <phillip.kruger@gmail.com>
Date: Mon, 18 Jul 2022 21:41:31 +1000
Subject: [PATCH] GraphQL: Make sure the context terminate

Signed-off-by: Phillip Kruger <phillip.kruger@gmail.com>
(cherry picked from commit 48787dec704466fc9bd53cd2cb134d40c77dc880)
---
 .../GraphQLTerminateContextTest.java          | 78 +++++++++++++++++++
 .../SmallRyeGraphQLAbstractHandler.java       | 22 ++++--
 2 files changed, 95 insertions(+), 5 deletions(-)
 create mode 100644 extensions/smallrye-graphql/deployment/src/test/java/io/quarkus/smallrye/graphql/deployment/GraphQLTerminateContextTest.java

diff --git a/extensions/smallrye-graphql/deployment/src/test/java/io/quarkus/smallrye/graphql/deployment/GraphQLTerminateContextTest.java b/extensions/smallrye-graphql/deployment/src/test/java/io/quarkus/smallrye/graphql/deployment/GraphQLTerminateContextTest.java
new file mode 100644
index 0000000000000..ce6c0123bbdf7
--- /dev/null
+++ b/extensions/smallrye-graphql/deployment/src/test/java/io/quarkus/smallrye/graphql/deployment/GraphQLTerminateContextTest.java
@@ -0,0 +1,78 @@
+package io.quarkus.smallrye.graphql.deployment;
+
+import static io.quarkus.smallrye.graphql.deployment.AbstractGraphQLTest.MEDIATYPE_JSON;
+
+import javax.inject.Inject;
+
+import org.eclipse.microprofile.graphql.GraphQLApi;
+import org.eclipse.microprofile.graphql.Query;
+import org.hamcrest.Matchers;
+import org.jboss.shrinkwrap.api.asset.EmptyAsset;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.test.QuarkusUnitTest;
+import io.restassured.RestAssured;
+import io.smallrye.mutiny.Uni;
+import io.vertx.ext.web.RoutingContext;
+
+/**
+ * Testing that the context terminate
+ */
+public class GraphQLTerminateContextTest extends AbstractGraphQLTest {
+
+    @RegisterExtension
+    static QuarkusUnitTest test = new QuarkusUnitTest()
+            .withApplicationRoot((JavaArchive jar) -> jar
+                    .addClasses(TestTerminateContextResource.class)
+                    .addAsManifestResource(EmptyAsset.INSTANCE, "beans.xml"));
+
+    @Test
+    public void testWhoAmI() {
+        runTestWith("user1");
+        runTestWith("user2");
+        runTestWith("user2");
+    }
+
+    public void runTestWith(String expectedUser) {
+        String fooRequest = getPayload("{\n" +
+                "  whoami {\n" +
+                "    name\n" +
+                "  }\n" +
+                "}");
+
+        RestAssured.given().when()
+                .accept(MEDIATYPE_JSON)
+                .contentType(MEDIATYPE_JSON)
+                .body(fooRequest)
+                .with().header("X-Test", expectedUser)
+                .post("/graphql")
+                .then()
+                .assertThat()
+                .statusCode(200)
+                .and()
+                .log().body().and()
+                .body("data.whoami.name", Matchers.equalTo(expectedUser));
+    }
+
+    @GraphQLApi
+    public static class TestTerminateContextResource {
+
+        @Inject
+        RoutingContext ctx;
+
+        @Query
+        public Uni<YouAre> whoami() {
+            return Uni.createFrom().item(() -> {
+                YouAre youAre = new YouAre();
+                youAre.name = ctx.request().headers().get("X-Test");
+                return youAre;
+            });
+        }
+    }
+
+    public static class YouAre {
+        public String name;
+    }
+}
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 dd57cef207078..3a09077350f90 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
@@ -30,7 +30,7 @@ public abstract class SmallRyeGraphQLAbstractHandler implements Handler<RoutingC
     private final CurrentIdentityAssociation currentIdentityAssociation;
     private final CurrentVertxRequest currentVertxRequest;
     private final ManagedContext currentManagedContext;
-
+    private final Handler currentManagedContextTerminationHandler;
     private final boolean runBlocking;
 
     private volatile ExecutionService executionService;
@@ -46,6 +46,13 @@ public SmallRyeGraphQLAbstractHandler(
         this.currentVertxRequest = currentVertxRequest;
         this.currentManagedContext = Arc.container().requestContext();
         this.runBlocking = runBlocking;
+        this.currentManagedContextTerminationHandler = new Handler() {
+            @Override
+            public void handle(Object e) {
+                currentManagedContext.terminate();
+            }
+
+        };
     }
 
     @Override
@@ -56,11 +63,16 @@ public void handle(final RoutingContext ctx) {
         } else {
 
             currentManagedContext.activate();
-            handleWithIdentity(ctx);
-
-            ctx.response().bodyEndHandler((e) -> {
+            ctx.response()
+                    .endHandler(currentManagedContextTerminationHandler)
+                    .exceptionHandler(currentManagedContextTerminationHandler)
+                    .closeHandler(currentManagedContextTerminationHandler);
+
+            try {
+                handleWithIdentity(ctx);
+            } catch (Throwable t) {
                 currentManagedContext.terminate();
-            });
+            }
         }
     }