Improve context based healthchecks

implementation/src/main/java/io/smallrye/health/SmallRyeHealthReporter.java

@@ -24,7 +24,6 @@
 import jakarta.enterprise.inject.Instance;
 import jakarta.enterprise.inject.spi.Bean;
 import jakarta.enterprise.inject.spi.BeanManager;
-import jakarta.enterprise.inject.spi.CDI;
 import jakarta.inject.Inject;
 import jakarta.json.JsonArray;
 import jakarta.json.JsonArrayBuilder;
@@ -409,14 +408,12 @@ private <T> List<T> getHealthChecksWithoutHealthGroup(Class<T> checkClass) {
     }
 
     private Uni<SmallRyeHealth> getHealthAsync(Uni<SmallRyeHealth> cachedHealth, HealthType... types) {
-        if (!checksInitialized) {
-            initChecks();
-        }
-
         if (contextPropagated) {
-            recreateCheckUnis();
-            return computeHealth(types);
+            return computeHealthWithContext(types);
         } else {
+            if (!checksInitialized) {
+                initChecks();
+            }
             if (additionalListsChanged(types) || additionalListsChanged || cachedHealth == null) {
                 additionalListsChanged = false;
                 cachedHealth = computeHealth(types);
@@ -426,14 +423,29 @@ private Uni<SmallRyeHealth> getHealthAsync(Uni<SmallRyeHealth> cachedHealth, Hea
         }
     }
 
-    private void recreateCheckUnis() {
-        livenessUnis.clear();
-        readinessUnis.clear();
-        wellnessUnis.clear();
-        startupUnis.clear();
-        checksInitialized = false;
-
-        initChecks();
+    private Uni<SmallRyeHealth> computeHealthWithContext(HealthType... types) {
+        List<Uni<HealthCheckResponse>> checks = new ArrayList<>();
+        for (HealthType type : types) {
+            switch (type) {
+                case LIVENESS:
+                    initUnis(checks, livenessChecks, asyncLivenessChecks);
+                    checks.addAll(livenessHealthRegistry.getChecks(healthChecksConfigs));
+                    break;
+                case READINESS:
+                    initUnis(checks, readinessChecks, asyncReadinessChecks);
+                    checks.addAll(readinessHealthRegistry.getChecks(healthChecksConfigs));
+                    break;
+                case WELLNESS:
+                    initUnis(checks, wellnessChecks, asyncWellnessChecks);
+                    checks.addAll(wellnessHealthRegistry.getChecks(healthChecksConfigs));
+                    break;
+                case STARTUP:
+                    initUnis(checks, startupChecks, asyncStartupChecks);
+                    checks.addAll(startupHealthRegistry.getChecks(healthChecksConfigs));
+                    break;
+            }
+        }
+        return getHealthAsync(checks);
     }
 
     private boolean additionalListsChanged(HealthType... types) {
@@ -467,7 +479,6 @@ private boolean additionalListsChanged(HealthType... types) {
 
     private Uni<SmallRyeHealth> computeHealth(HealthType[] types) {
         List<Uni<HealthCheckResponse>> checks = new ArrayList<>();
-
         for (HealthType type : types) {
             switch (type) {
                 case LIVENESS:
@@ -488,7 +499,6 @@ private Uni<SmallRyeHealth> computeHealth(HealthType[] types) {
                     break;
             }
         }
-
         return getHealthAsync(checks);
     }
 
@@ -508,7 +518,7 @@ private Uni<SmallRyeHealth> getHealthAsync(Collection<Uni<HealthCheckResponse>>
         }
 
         return Uni.combine().all().unis(healthCheckUnis)
-                .combinedWith(responses -> {
+                .with(responses -> {
                     JsonArrayBuilder results = jsonProvider.createArrayBuilder();
                     HealthCheckResponse.Status status = HealthCheckResponse.Status.UP;
 

implementation/src/main/java/io/smallrye/health/registry/HealthRegistryImpl.java

@@ -1,10 +1,10 @@
 package io.smallrye.health.registry;
 
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.eclipse.microprofile.health.HealthCheck;
 import org.eclipse.microprofile.health.HealthCheckResponse;
@@ -16,31 +16,32 @@
 
 public class HealthRegistryImpl implements HealthRegistry {
 
-    Map<String, Uni<HealthCheckResponse>> checks = new HashMap<>();
+    Map<String, HealthCheck> checks = new HashMap<>();
+    Map<String, AsyncHealthCheck> asyncChecks = new HashMap<>();
     private boolean checksChanged = false;
 
     AsyncHealthCheckFactory asyncHealthCheckFactory = new AsyncHealthCheckFactory();
 
     @Override
     public HealthRegistry register(String id, HealthCheck healthCheck) {
-        return register(id, asyncHealthCheckFactory.callSync(healthCheck));
+        return register(id, healthCheck, checks);
     }
 
     @Override
     public HealthRegistry register(String id, AsyncHealthCheck asyncHealthCheck) {
-        return register(id, asyncHealthCheckFactory.callAsync(asyncHealthCheck));
+        return register(id, asyncHealthCheck, asyncChecks);
     }
 
-    private HealthRegistry register(String id, Uni<HealthCheckResponse> healthCheckResponseUni) {
-        checks.put(id, healthCheckResponseUni);
+    private <T> HealthRegistry register(String id, T check, Map<String, T> checks) {
+        checks.put(id, check);
         checksChanged = true;
         return this;
     }
 
     @Override
     public HealthRegistry remove(String id) {
         try {
-            if (checks.remove(id) == null) {
+            if (checks.remove(id) == null && asyncChecks.remove(id) == null) {
                 throw new IllegalStateException(String.format("ID '%s' not found", id));
             }
             checksChanged = true;
@@ -51,9 +52,13 @@ public HealthRegistry remove(String id) {
     }
 
     public Collection<Uni<HealthCheckResponse>> getChecks(Map<String, Boolean> healthChecksConfigs) {
-        return Collections.unmodifiableCollection(checks.entrySet().stream()
+        var enabledChecks = checks.entrySet().stream()
                 .filter(e -> healthChecksConfigs.getOrDefault(e.getKey(), true))
-                .map(Map.Entry::getValue).collect(Collectors.toList()));
+                .map(e -> asyncHealthCheckFactory.callSync(e.getValue()));
+        var enabledAsyncChecks = asyncChecks.entrySet().stream()
+                .filter(e -> healthChecksConfigs.getOrDefault(e.getKey(), true))
+                .map(e -> asyncHealthCheckFactory.callAsync(e.getValue()));
+        return Stream.concat(enabledChecks, enabledAsyncChecks).collect(Collectors.toList());
     }
 
     public boolean checksChanged() {

testsuite/experimental/pom.xml

@@ -32,6 +32,12 @@
       <scope>compile</scope>
     </dependency>
 
+    <dependency>
+      <groupId>jakarta.servlet</groupId>
+      <artifactId>jakarta.servlet-api</artifactId>
+      <scope>provided</scope>
+    </dependency>
+
   </dependencies>
 
   <profiles>

testsuite/experimental/src/test/java/io/smallrye/health/deployment/ContextHealthCheck.java

@@ -0,0 +1,23 @@
+package io.smallrye.health.deployment;
+
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Inject;
+
+import org.eclipse.microprofile.health.HealthCheck;
+import org.eclipse.microprofile.health.HealthCheckResponse;
+
+import io.smallrye.health.api.Wellness;
+
+@Wellness
+@ApplicationScoped
+public class ContextHealthCheck implements HealthCheck {
+
+    @Inject
+    ContextInfo context;
+
+    @Override
+    public HealthCheckResponse call() {
+        return HealthCheckResponse.up(context.getValue());
+    }
+
+}

testsuite/experimental/src/test/java/io/smallrye/health/deployment/ContextInfo.java

@@ -0,0 +1,17 @@
+package io.smallrye.health.deployment;
+
+import jakarta.enterprise.context.RequestScoped;
+
+@RequestScoped
+public class ContextInfo {
+
+    private String value;
+
+    public String getValue() {
+        return value;
+    }
+
+    public void setValue(String value) {
+        this.value = value;
+    }
+}

testsuite/experimental/src/test/java/io/smallrye/health/deployment/RequestFilter.java

@@ -0,0 +1,29 @@
+package io.smallrye.health.deployment;
+
+import java.io.IOException;
+
+import jakarta.inject.Inject;
+import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.annotation.WebFilter;
+import jakarta.servlet.http.HttpFilter;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+
+@WebFilter(urlPatterns = "*", filterName = "ContextInfoFilter")
+public class RequestFilter extends HttpFilter {
+
+    public static final String TEST_CONTEXT_HEADER = "X-CONTEXT";
+
+    @Inject
+    ContextInfo contextInfo;
+
+    @Override
+    protected void doFilter(HttpServletRequest req, HttpServletResponse res, FilterChain chain)
+            throws IOException, ServletException {
+        var value = req.getHeader(TEST_CONTEXT_HEADER);
+        contextInfo.setValue(value);
+        super.doFilter(req, res, chain);
+    }
+
+}

testsuite/experimental/src/test/java/io/smallrye/health/test/ContextPropagatedHealthTest.java

@@ -0,0 +1,51 @@
+package io.smallrye.health.test;
+
+import java.util.List;
+import java.util.Map;
+
+import org.jboss.arquillian.container.test.api.Deployment;
+import org.jboss.arquillian.container.test.api.RunAsClient;
+import org.jboss.shrinkwrap.api.Archive;
+import org.jboss.shrinkwrap.api.asset.StringAsset;
+import org.jboss.shrinkwrap.api.spec.WebArchive;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import io.smallrye.health.deployment.ContextHealthCheck;
+import io.smallrye.health.deployment.ContextInfo;
+import io.smallrye.health.deployment.RequestFilter;
+import io.smallrye.mutiny.Uni;
+import io.smallrye.mutiny.infrastructure.Infrastructure;
+
+@RunAsClient
+class ContextPropagatedHealthTest extends TCKBase {
+
+    @Deployment
+    public static Archive<WebArchive> getDeployment() {
+        return DeploymentUtils
+                .createWarFileWithClasses(ContextPropagatedHealthTest.class.getSimpleName(),
+                        ContextHealthCheck.class, ContextInfo.class, RequestFilter.class, TCKBase.class)
+                .addAsManifestResource(
+                        new StringAsset("io.smallrye.health.context.propagation=true"),
+                        "microprofile-config.properties");
+    }
+
+    // got to repeat this test a couple of times to trigger potential race condition issue at least once.
+    @Test(invocationCount = 20)
+    void twoParallelRequests_shouldNotMixupContext() {
+        String contextValue = "test-value";
+        String otherContextValue = "other-test-value";
+        Uni<Response> response = Uni.createFrom().voidItem().emitOn(Infrastructure.getDefaultExecutor()).onItem()
+                .transform(v -> getUrlHealthContents(Map.of(RequestFilter.TEST_CONTEXT_HEADER, contextValue)));
+        Uni<Response> otherResponse = Uni.createFrom().voidItem().emitOn(Infrastructure.getDefaultExecutor()).onItem()
+                .transform(v -> getUrlHealthContents(Map.of(RequestFilter.TEST_CONTEXT_HEADER, otherContextValue)));
+
+        List<Response> responses = Uni.join().all(response, otherResponse).andCollectFailures().await().indefinitely();
+
+        String responseBody = responses.get(0).getBody().orElseThrow();
+        String otherResponseBody = responses.get(1).getBody().orElseThrow();
+        Assert.assertNotEquals(responseBody, otherResponseBody);
+        Assert.assertTrue(responseBody.contains(contextValue), responseBody);
+        Assert.assertTrue(otherResponseBody.contains(otherContextValue), otherResponseBody);
+    }
+}

testsuite/experimental/src/test/java/io/smallrye/health/test/TCKBase.java

@@ -27,6 +27,7 @@
 import java.io.StringReader;
 import java.lang.reflect.Method;
 import java.net.URI;
+import java.util.Map;
 import java.util.Optional;
 import java.util.logging.Logger;
 
@@ -65,6 +66,10 @@ public void beforeMethod(Method method) {
         LOG.info(String.format("Running test: %s#%s", method.getDeclaringClass().getSimpleName(), method.getName()));
     }
 
+    Response getUrlHealthContents(Map<String, String> headers) {
+        return getUrlContents(this.uri + "/health", false, true, headers);
+    }
+
     Response getUrlHealthContents() {
         return getUrlContents(this.uri + "/health", false);
     }
@@ -98,6 +103,10 @@ private Response getUrlContents(String theUrl, boolean useAuth) {
     }
 
     private Response getUrlContents(String theUrl, boolean useAuth, boolean followRedirects) {
+        return getUrlContents(theUrl, useAuth, followRedirects, Map.of());
+    }
+
+    private Response getUrlContents(String theUrl, boolean useAuth, boolean followRedirects, Map<String, String> headers) {
 
         StringBuilder content = new StringBuilder();
         int code;
@@ -117,8 +126,9 @@ private Response getUrlContents(String theUrl, boolean useAuth, boolean followRe
             }
 
             HttpClient client = builder.build();
-
-            HttpResponse response = client.execute(new HttpGet(theUrl));
+            var httpGet = new HttpGet(theUrl);
+            headers.forEach(httpGet::addHeader);
+            HttpResponse response = client.execute(httpGet);
             code = response.getStatusLine().getStatusCode();
 
             if (response.getEntity() != null) {