GraphQL: Make sure the context terminate

Signed-off-by: Phillip Kruger <phillip.kruger@gmail.com>

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;
+    }
+}

extensions/smallrye-graphql/runtime/src/main/java/io/quarkus/smallrye/graphql/runtime/SmallRyeGraphQLAbstractHandler.java

@@ -58,9 +58,16 @@ public void handle(final RoutingContext ctx) {
             currentManagedContext.activate();
             handleWithIdentity(ctx);
 
-            ctx.response().bodyEndHandler((e) -> {
-                currentManagedContext.terminate();
-            });
+            ctx.response()
+                    .endHandler((e) -> {
+                        currentManagedContext.terminate();
+                    })
+                    .exceptionHandler((e) -> {
+                        currentManagedContext.terminate();
+                    })
+                    .closeHandler((e) -> {
+                        currentManagedContext.terminate();
+                    });
         }
     }