diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeGraph.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeGraph.java index 25640bbf9b..7a4d96ada5 100644 --- 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 index 9b2c7a4c40..df8a3906a3 100644 --- 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() {