From 650ec93a5990c1e1820c4cb49ba0e39b4ac1a28e Mon Sep 17 00:00:00 2001
From: Ladislav Thon <lthon@redhat.com>
Date: Tue, 20 Feb 2024 15:07:34 +0100
Subject: [PATCH] SmallRye Health: terminate request context properly

When the SmallRye Health route handler activates request context on its own,
it also needs to terminate it. However, that termination needs to happen
only after all health checks complete. That's what this commit does.
---
 .../runtime/SmallRyeHealthHandlerBase.java       | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/extensions/smallrye-health/runtime/src/main/java/io/quarkus/smallrye/health/runtime/SmallRyeHealthHandlerBase.java b/extensions/smallrye-health/runtime/src/main/java/io/quarkus/smallrye/health/runtime/SmallRyeHealthHandlerBase.java
index e999375418769..c24dcfb60cc2b 100644
--- a/extensions/smallrye-health/runtime/src/main/java/io/quarkus/smallrye/health/runtime/SmallRyeHealthHandlerBase.java
+++ b/extensions/smallrye-health/runtime/src/main/java/io/quarkus/smallrye/health/runtime/SmallRyeHealthHandlerBase.java
@@ -28,18 +28,19 @@ abstract class SmallRyeHealthHandlerBase implements Handler<RoutingContext> {
     public void handle(RoutingContext ctx) {
         ManagedContext requestContext = Arc.container().requestContext();
         if (requestContext.isActive()) {
-            doHandle(ctx);
+            doHandle(ctx, null);
         } else {
             requestContext.activate();
             try {
-                doHandle(ctx);
-            } finally {
+                doHandle(ctx, requestContext);
+            } catch (Exception e) {
                 requestContext.terminate();
+                throw e;
             }
         }
     }
 
-    private void doHandle(RoutingContext ctx) {
+    private void doHandle(RoutingContext ctx, ManagedContext requestContext) {
         QuarkusHttpUser user = (QuarkusHttpUser) ctx.user();
         if (user != null) {
             Arc.container().instance(CurrentIdentityAssociation.class).get().setIdentity(user.getSecurityIdentity());
@@ -48,6 +49,9 @@ private void doHandle(RoutingContext ctx) {
         Context context = Vertx.currentContext();
         getHealth(reporter, ctx).emitOn(MutinyHelper.executor(context))
                 .subscribe().with(health -> {
+                    if (requestContext != null) {
+                        requestContext.terminate();
+                    }
                     HttpServerResponse resp = ctx.response();
                     if (health.isDown()) {
                         resp.setStatusCode(503);
@@ -60,6 +64,10 @@ private void doHandle(RoutingContext ctx) {
                     } catch (IOException e) {
                         throw new UncheckedIOException(e);
                     }
+                }, failure -> {
+                    if (requestContext != null) {
+                        requestContext.terminate();
+                    }
                 });
     }
 }