[AMORO-3056] Support graceful shutdown for AMS (#3057)

apache · Jul 24, 2024 · f5d2e85 · f5d2e85
1 parent fafb3f4
commit f5d2e85
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/amoro-ams/amoro-ams-server/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java b/amoro-ams/amoro-ams-server/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java
@@ -102,6 +102,14 @@ public AmoroServiceContainer() throws Exception {
   public static void main(String[] args) {
     try {
       AmoroServiceContainer service = new AmoroServiceContainer();
+      Runtime.getRuntime()
+          .addShutdownHook(
+              new Thread(
+                  () -> {
+                    LOG.info("AMS service is shutting down...");
+                    service.dispose();
+                    LOG.info("AMS service has been shut down");
+                  }));
       while (true) {
         try {
           service.waitLeaderShip();
@@ -165,24 +173,29 @@ private void addHandlerChain(RuntimeHandlerChain chain) {
   }
 
   public void dispose() {
-    if (tableManagementServer != null) {
+    if (tableManagementServer != null && tableManagementServer.isServing()) {
+      LOG.info("Stopping table management server...");
       tableManagementServer.stop();
     }
-    if (optimizingServiceServer != null) {
+    if (optimizingServiceServer != null && optimizingServiceServer.isServing()) {
+      LOG.info("Stopping optimizing server...");
       optimizingServiceServer.stop();
     }
     if (httpServer != null) {
+      LOG.info("Stopping http server...");
       try {
         httpServer.close();
       } catch (Exception e) {
         LOG.error("Error stopping http server", e);
       }
     }
     if (tableService != null) {
+      LOG.info("Stopping table service...");
       tableService.dispose();
       tableService = null;
     }
     if (terminalManager != null) {
+      LOG.info("Stopping terminal manager...");
       terminalManager.dispose();
       terminalManager = null;
     }

diff --git a/...-ams/amoro-ams-server/src/main/java/org/apache/amoro/server/terminal/TerminalManager.java b/...-ams/amoro-ams-server/src/main/java/org/apache/amoro/server/terminal/TerminalManager.java
@@ -69,7 +69,7 @@ public class TerminalManager {
   private final Object sessionMapLock = new Object();
   private final Map<String, TerminalSessionContext> sessionMap = Maps.newHashMap();
   private final Thread gcThread;
-  private boolean stop = false;
+  private volatile boolean running = true;
 
   private final ThreadPoolExecutor executionPool =
       new ThreadPoolExecutor(
@@ -237,7 +237,10 @@ public LatestSessionInfo getLastSessionInfo(String terminalId) {
   }
 
   public void dispose() {
-    stop = true;
+    if (!running) {
+      return;
+    }
+    running = false;
     if (gcThread != null) {
       gcThread.interrupt();
     }
@@ -399,7 +402,7 @@ public void run() {
       LOG.info(
           "Terminal Session Clean Task, check interval: " + SESSION_TIMEOUT_CHECK_INTERVAL + " ms");
       LOG.info("Terminal Session Timeout: {} minutes", sessionTimeout);
-      while (!stop) {
+      while (running) {
         try {
           List<TerminalSessionContext> sessionToRelease = checkIdleSession();
           sessionToRelease.forEach(this::releaseSession);