added healthcheckregistry instead of rotation

flipkart-incubator · Nov 15, 2024 · 570621b · 570621b
1 parent d7735a5
commit 570621b
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/core/src/main/java/com/flipkart/gjex/core/healthcheck/GrpcHealthCheckService.java b/core/src/main/java/com/flipkart/gjex/core/healthcheck/GrpcHealthCheckService.java
@@ -1,13 +1,17 @@
 package com.flipkart.gjex.core.healthcheck;
 
 import com.flipkart.gjex.core.filter.grpc.MethodFilters;
+import io.dropwizard.metrics5.health.HealthCheck;
 import io.grpc.health.v1.HealthCheckRequest;
 import io.grpc.health.v1.HealthCheckResponse;
 import io.grpc.health.v1.HealthGrpc;
 import io.grpc.stub.StreamObserver;
 
 import javax.inject.Named;
 import javax.inject.Singleton;
+import javax.servlet.ServletContext;
+import javax.ws.rs.core.Context;
+import java.util.SortedMap;
 
 /**
  * Default GRPC HealthCheck Service
@@ -18,6 +22,9 @@
 @Named("GrpcHealthCheckService")
 public class GrpcHealthCheckService extends HealthGrpc.HealthImplBase {
 
+  @Context
+  private ServletContext servletContext;
+
   public GrpcHealthCheckService(){
 
   }
@@ -27,7 +34,10 @@ public GrpcHealthCheckService(){
   public void check(HealthCheckRequest request,
                     StreamObserver<HealthCheckResponse> responseObserver) {
     HealthCheckResponse.Builder builder = HealthCheckResponse.newBuilder();
-    if (RotationManagementBasedHealthCheck.inRotation()) {
+    HealthCheckRegistry registry = (HealthCheckRegistry) servletContext
+        .getAttribute(HealthCheckRegistry.HEALTHCHECK_REGISTRY_NAME);
+    SortedMap<String, HealthCheck.Result> results = registry.runHealthChecks();
+    if (results.values().stream().anyMatch(result -> !result.isHealthy())){
       builder.setStatus(HealthCheckResponse.ServingStatus.SERVING);
     } else {
       builder.setStatus(HealthCheckResponse.ServingStatus.NOT_SERVING);

diff --git a/core/src/main/java/com/flipkart/gjex/core/web/DashboardHealthCheckResource.java b/core/src/main/java/com/flipkart/gjex/core/web/DashboardHealthCheckResource.java
@@ -46,10 +46,12 @@ public class DashboardHealthCheckResource {
 	@GET
 	@Produces(MediaType.APPLICATION_JSON)
 	public Response performHealthChecks() {
-		HealthCheckRegistry registry = (HealthCheckRegistry) servletContext
-				.getAttribute(HealthCheckRegistry.HEALTHCHECK_REGISTRY_NAME);
-		SortedMap<String, HealthCheck.Result> results = registry.runHealthChecks();
-		return Response.status(Response.Status.OK).entity(results).build();
+    HealthCheckRegistry registry = (HealthCheckRegistry) servletContext
+      .getAttribute(HealthCheckRegistry.HEALTHCHECK_REGISTRY_NAME);
+    SortedMap<String, HealthCheck.Result> results = registry.runHealthChecks();
+    if (results.values().stream().anyMatch(result -> !result.isHealthy())){
+      return Response.status(Response.Status.SERVICE_UNAVAILABLE).entity(results).build();
+    }
+    return Response.status(Response.Status.OK).entity(results).build();
 	}
-
 }