add shutdown hook to close thread pool automatically

Change-Id: I876fdab935cd29248fd37b3f2850366dc1983d4f
apache · Feb 2, 2019 · ab656b9 · ab656b9
1 parent 1885132
commit ab656b9
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeGraph.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeGraph.java
@@ -90,6 +90,10 @@ public class HugeGraph implements GremlinGraph {
                                  HugeGraphStepStrategy.instance());
         TraversalStrategies.GlobalCache.registerStrategies(HugeGraph.class,
                                                            strategies);
+
+        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+            HugeGraph.shutdown(30L);
+        }));
     }
 
     private volatile boolean closed;
@@ -568,9 +572,16 @@ public Id[] mapVlName2Id(String[] vertexLabels) {
      * @param timout seconds
      * @throws InterruptedException when be interrupted
      */
-    public static void shutdown(long timout) throws InterruptedException {
-        EventHub.destroy(timout);
-        TaskManager.instance().shutdown(timout);
+    public static void shutdown(long timeout) {
+        try {
+            if (!EventHub.destroy(timeout)) {
+                throw new TimeoutException(timeout + "s");
+            }
+            TaskManager.instance().shutdown(timeout);
+        } catch (Throwable e) {
+            LOG.error("Error while shutdown", e);
+            throw new HugeException("Failed to shutdown", e);
+        }
     }
 
     private class TinkerpopTransaction extends AbstractThreadLocalTransaction {

diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/task/TaskManager.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/task/TaskManager.java
@@ -28,6 +28,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import com.baidu.hugegraph.HugeException;
 import com.baidu.hugegraph.HugeGraph;
@@ -116,30 +117,37 @@ public TaskScheduler getScheduler(HugeGraph graph) {
     }
 
     public void shutdown(long timeout) {
-        Throwable ex = null;
         assert this.schedulers.isEmpty() : this.schedulers.size();
 
+        Throwable ex = null;
+        boolean terminated = this.taskExecutor.isTerminated();
+        final TimeUnit unit = TimeUnit.SECONDS;
+
         if (!this.taskExecutor.isShutdown()) {
             this.taskExecutor.shutdown();
             try {
-                this.taskExecutor.awaitTermination(timeout, TimeUnit.SECONDS);
+                terminated = this.taskExecutor.awaitTermination(timeout, unit);
             } catch (Throwable e) {
                 ex = e;
             }
         }
 
-        if (!this.dbExecutor.isShutdown()) {
+        if (terminated && !this.dbExecutor.isShutdown()) {
             this.dbExecutor.shutdown();
             try {
-                this.dbExecutor.awaitTermination(timeout, TimeUnit.SECONDS);
+                terminated = this.dbExecutor.awaitTermination(timeout, unit);
             } catch (Throwable e) {
                 ex = e;
             }
         }
 
+        if (!terminated) {
+            ex = new TimeoutException(timeout + "s");
+        }
         if (ex != null) {
             throw new HugeException("Failed to wait for TaskScheduler", ex);
         }
+
     }
 
     public int workerPoolSize() {