feat: add server shutdown util (#20)

Author: aaron-steinfeld

grpc-server-utils/build.gradle.kts

@@ -24,5 +24,9 @@ dependencies {
     }
   }
 
+  annotationProcessor("org.projectlombok:lombok:1.18.20")
+  compileOnly("org.projectlombok:lombok:1.18.20")
+
   testImplementation("org.junit.jupiter:junit-jupiter:5.7.0")
+  testImplementation("org.mockito:mockito-core:3.12.1")
 }

grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/ServerManagementUtil.java

@@ -0,0 +1,49 @@
+package org.hypertrace.core.grpcutils.server;
+
+import io.grpc.Server;
+import java.time.Duration;
+import java.util.concurrent.TimeUnit;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class ServerManagementUtil {
+
+  public static void shutdownServer(Server grpcServer, String name, Duration timeout) {
+    log.info("Starting shutdown for service [{}]", name);
+    grpcServer.shutdown();
+    boolean gracefullyShutdown = waitForGracefulShutdown(grpcServer, name, timeout);
+    if (!gracefullyShutdown) {
+      forceShutdown(grpcServer, name);
+    }
+  }
+
+  private static boolean waitForGracefulShutdown(Server grpcServer, String name, Duration timeout) {
+    boolean successfullyShutdown = waitForTermination(grpcServer, name, timeout);
+    if (successfullyShutdown) {
+      log.info("Shutdown service successfully [{}]", name);
+    }
+    return successfullyShutdown;
+  }
+
+  private static void forceShutdown(Server grpcServer, String name) {
+    log.error("Shutting down service [{}] forcefully", name);
+    grpcServer.shutdownNow();
+    if (waitForTermination(grpcServer, name, Duration.ofSeconds(5))) {
+      log.error("Forced service [{}] shutdown successful", name);
+    } else {
+      log.error("Unable to force service [{}] shutdown in 5s - giving up!", name);
+    }
+  }
+
+  private static boolean waitForTermination(Server grpcServer, String name, Duration deadline) {
+    try {
+      if (!grpcServer.awaitTermination(deadline.toMillis(), TimeUnit.MILLISECONDS)) {
+        log.error("Service [{}] did not shut down after waiting", name);
+      }
+    } catch (InterruptedException ex) {
+      log.error("There has been an interruption while waiting for service [{}] to shutdown", name);
+      Thread.currentThread().interrupt();
+    }
+    return grpcServer.isTerminated();
+  }
+}

grpc-server-utils/src/test/java/org/hypertrace/core/grpcutils/server/ServerManagementUtilTest.java

@@ -0,0 +1,46 @@
+package org.hypertrace.core.grpcutils.server;
+
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import io.grpc.Server;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
+import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.Test;
+import org.mockito.InOrder;
+
+class ServerManagementUtilTest {
+
+  @Test
+  void canShutdownGracefully() throws InterruptedException {
+    Server mockServer = mock(Server.class);
+    when(mockServer.isTerminated()).thenReturn(true);
+    ServerManagementUtil.shutdownServer(
+        mockServer, "mockServer", Duration.of(10, ChronoUnit.MILLIS));
+
+    InOrder serverVerifier = inOrder(mockServer);
+    serverVerifier.verify(mockServer).shutdown();
+    serverVerifier.verify(mockServer).awaitTermination(10, TimeUnit.MILLISECONDS);
+    serverVerifier.verify(mockServer).isTerminated();
+    serverVerifier.verifyNoMoreInteractions();
+  }
+
+  @Test
+  void canShutdownForcefully() throws InterruptedException {
+    Server mockServer = mock(Server.class);
+    when(mockServer.isTerminated()).thenReturn(false);
+    ServerManagementUtil.shutdownServer(
+        mockServer, "mockServer", Duration.of(10, ChronoUnit.MILLIS));
+
+    InOrder serverVerifier = inOrder(mockServer);
+    serverVerifier.verify(mockServer).shutdown();
+    serverVerifier.verify(mockServer).awaitTermination(10, TimeUnit.MILLISECONDS);
+    serverVerifier.verify(mockServer).isTerminated();
+    serverVerifier.verify(mockServer).shutdownNow();
+    serverVerifier.verify(mockServer).awaitTermination(5000, TimeUnit.MILLISECONDS);
+    serverVerifier.verify(mockServer).isTerminated();
+    serverVerifier.verifyNoMoreInteractions();
+  }
+}