Add ServerConfig option to specify shutdown timeout. (deephaven#5360)

Fixes deephaven#5359
kosak · Apr 14, 2024 · c2017ce · c2017ce
1 parent bf6fcdb
commit c2017ce
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/server/src/main/java/io/deephaven/server/config/ServerConfig.java b/server/src/main/java/io/deephaven/server/config/ServerConfig.java
@@ -21,6 +21,8 @@ public interface ServerConfig {
 
     int DEFAULT_MAX_INBOUND_MESSAGE_SIZE_MiB = 100;
 
+    int DEFAULT_SHUTDOWN_TIMEOUT_MILLIS = 10_000;
+
     String HTTP_SESSION_DURATION_MS = "http.session.durationMs";
 
     String HTTP_HOST = "http.host";
@@ -34,6 +36,7 @@ public interface ServerConfig {
     String GRPC_MAX_INBOUND_MESSAGE_SIZE = "grpc.maxInboundMessageSize";
 
     String PROXY_HINT = "proxy.hint";
+    String SHUTDOWN_TIMEOUT_MILLIS = "shutdown.timeoutMs";
 
     /**
      * Parses the configuration values into the appropriate builder methods.
@@ -92,6 +95,7 @@ static <B extends Builder<?, B>> B buildFromConfig(B builder, Configuration conf
         int schedulerPoolSize = config.getIntegerWithDefault(SCHEDULER_POOL_SIZE, -1);
         int maxInboundMessageSize = config.getIntegerWithDefault(GRPC_MAX_INBOUND_MESSAGE_SIZE, -1);
         String proxyHint = config.getStringWithDefault(PROXY_HINT, null);
+        int shutdownTimeoutMillis = config.getIntegerWithDefault(SHUTDOWN_TIMEOUT_MILLIS, -1);
         if (httpSessionExpireMs > -1) {
             builder.tokenExpire(Duration.ofMillis(httpSessionExpireMs));
         }
@@ -113,6 +117,9 @@ static <B extends Builder<?, B>> B buildFromConfig(B builder, Configuration conf
         if (proxyHint != null) {
             builder.proxyHint(Boolean.parseBoolean(proxyHint));
         }
+        if (shutdownTimeoutMillis != -1) {
+            builder.shutdownTimeout(Duration.ofMillis(shutdownTimeoutMillis));
+        }
         MainHelper.parseSSLConfig(config).ifPresent(builder::ssl);
         MainHelper.parseOutboundSSLConfig(config).ifPresent(builder::outboundSsl);
         return builder;
@@ -165,6 +172,14 @@ default int maxInboundMessageSize() {
         return DEFAULT_MAX_INBOUND_MESSAGE_SIZE_MiB * 1024 * 1024;
     }
 
+    /**
+     * How many do we wait to shut down the server. Defaults to {@value DEFAULT_SHUTDOWN_TIMEOUT_MILLIS}.
+     */
+    @Default
+    default Duration shutdownTimeout() {
+        return Duration.ofMillis(DEFAULT_SHUTDOWN_TIMEOUT_MILLIS);
+    }
+
     /**
      * A hint that the server is running behind a proxy. This may allow consumers of the configuration to make more
      * appropriate default choices.
@@ -219,6 +234,8 @@ interface Builder<T, B extends Builder<T, B>> {
 
         B proxyHint(Boolean proxyHint);
 
+        B shutdownTimeout(Duration shutdownTimeout);
+
         T build();
     }
 }
diff --git a/server/src/main/java/io/deephaven/server/runner/DeephavenApiServer.java b/server/src/main/java/io/deephaven/server/runner/DeephavenApiServer.java
@@ -37,6 +37,7 @@
 import javax.inject.Named;
 import javax.inject.Provider;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -129,7 +130,8 @@ public DeephavenApiServer run() throws IOException, ClassNotFoundException, Time
         // Finally, wait for the http server to be finished stopping
         ProcessEnvironment.getGlobalShutdownManager().registerTask(ShutdownManager.OrderingCategory.LAST, () -> {
             try {
-                server.stopWithTimeout(10, TimeUnit.SECONDS);
+                final Duration duration = serverConfig.shutdownTimeout();
+                server.stopWithTimeout(duration.toNanos(), TimeUnit.NANOSECONDS);
                 server.join();
             } catch (final InterruptedException ignored) {
             }