Add HealthStatusManager and correctly enterTerminalState (deephaven#2848

) Add HealthStatusManager and correctly enterTerminalState
jmao-denver · Oct 3, 2022 · d8f0c11 · d8f0c11
1 parent 6a040d2
commit d8f0c11
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 8 deletions.
diff --git a/server/src/main/java/io/deephaven/server/healthcheck/HealthCheckModule.java b/server/src/main/java/io/deephaven/server/healthcheck/HealthCheckModule.java
@@ -19,10 +19,6 @@ public class HealthCheckModule {
     @Singleton
     public HealthStatusManager bindHealthStatusManager() {
         HealthStatusManager healthStatusManager = new HealthStatusManager();
-        ProcessEnvironment.getGlobalShutdownManager().registerTask(
-                ShutdownManager.OrderingCategory.FIRST,
-                healthStatusManager::enterTerminalState);
-
         return healthStatusManager;
     }
 

diff --git a/server/src/main/java/io/deephaven/server/runner/DeephavenApiServer.java b/server/src/main/java/io/deephaven/server/runner/DeephavenApiServer.java
@@ -25,6 +25,8 @@
 import io.deephaven.util.annotations.VisibleForTesting;
 import io.deephaven.util.process.ProcessEnvironment;
 import io.deephaven.util.process.ShutdownManager;
+import io.grpc.health.v1.HealthCheckResponse;
+import io.grpc.protobuf.services.HealthStatusManager;
 
 import javax.inject.Inject;
 import javax.inject.Provider;
@@ -51,6 +53,7 @@ public class DeephavenApiServer {
     private final Map<String, AuthenticationRequestHandler> authenticationHandlers;
     private final Provider<ExecutionContext> executionContextProvider;
     private final ServerConfig serverConfig;
+    private final HealthStatusManager healthStatusManager;
 
     @Inject
     public DeephavenApiServer(
@@ -64,7 +67,8 @@ public DeephavenApiServer(
             final SessionService sessionService,
             final Map<String, AuthenticationRequestHandler> authenticationHandlers,
             final Provider<ExecutionContext> executionContextProvider,
-            final ServerConfig serverConfig) {
+            final ServerConfig serverConfig,
+            final HealthStatusManager healthStatusManager) {
         this.server = server;
         this.ugp = ugp;
         this.logInit = logInit;
@@ -76,6 +80,7 @@ public DeephavenApiServer(
         this.authenticationHandlers = authenticationHandlers;
         this.executionContextProvider = executionContextProvider;
         this.serverConfig = serverConfig;
+        this.healthStatusManager = healthStatusManager;
     }
 
     @VisibleForTesting
@@ -100,7 +105,12 @@ SessionService sessionService() {
     public DeephavenApiServer run() throws IOException, ClassNotFoundException, TimeoutException {
         // Stop accepting new gRPC requests.
         ProcessEnvironment.getGlobalShutdownManager().registerTask(ShutdownManager.OrderingCategory.FIRST,
-                () -> server.stopWithTimeout(10, TimeUnit.SECONDS));
+                () -> {
+                    // healthStatusManager.enterTerminalState() must be called before server.stopWithTimeout().
+                    // If we add multiple `OrderingCategory.FIRST` callbacks, they'll execute in the wrong order.
+                    healthStatusManager.enterTerminalState();
+                    server.stopWithTimeout(10, TimeUnit.SECONDS);
+                });
 
         // Close outstanding sessions to give any gRPCs closure.
         ProcessEnvironment.getGlobalShutdownManager().registerTask(ShutdownManager.OrderingCategory.MIDDLE,
@@ -151,6 +161,7 @@ public DeephavenApiServer run() throws IOException, ClassNotFoundException, Time
         log.info().append("Starting server...").endl();
         server.start();
         log.info().append("Server started on port ").append(server.getPort()).endl();
+        healthStatusManager.setStatus("", HealthCheckResponse.ServingStatus.SERVING);
         return this;
     }
 

diff --git a/server/src/main/java/io/deephaven/server/runner/DeephavenApiServerModule.java b/server/src/main/java/io/deephaven/server/runner/DeephavenApiServerModule.java
@@ -11,6 +11,7 @@
 import io.deephaven.engine.updategraph.UpdateGraphProcessor;
 import io.deephaven.engine.util.ScriptSession;
 import io.deephaven.server.notebook.FilesystemStorageServiceModule;
+import io.deephaven.server.healthcheck.HealthCheckModule;
 import io.deephaven.server.object.ObjectServiceModule;
 import io.deephaven.server.partitionedtable.PartitionedTableServiceModule;
 import io.deephaven.server.plugin.PluginsModule;
@@ -27,6 +28,7 @@
 import io.deephaven.util.thread.NamingThreadFactory;
 import io.grpc.BindableService;
 import io.grpc.ServerInterceptor;
+import io.grpc.protobuf.services.HealthStatusManager;
 import org.jetbrains.annotations.NotNull;
 
 import javax.inject.Named;
@@ -58,13 +60,14 @@
         PluginsModule.class,
         PartitionedTableServiceModule.class,
         FilesystemStorageServiceModule.class,
+        HealthCheckModule.class,
 })
 public class DeephavenApiServerModule {
 
     @Provides
     @ElementsIntoSet
-    static Set<BindableService> primeServices() {
-        return Collections.emptySet();
+    static Set<BindableService> primeServices(HealthStatusManager healthStatusManager) {
+        return Collections.singleton(healthStatusManager.getHealthService());
     }
 
     @Provides