From 2eb6c9b3dd505a25a883829490bd4154163cb722 Mon Sep 17 00:00:00 2001
From: zyxxoo <1318247699@qq.com>
Date: Wed, 30 Mar 2022 13:17:25 +0800
Subject: [PATCH] fix: eventHub notify after shutdown

---
 .../java/com/baidu/hugegraph/HugeFactory.java | 21 +++++++++++++++----
 .../baidu/hugegraph/dist/HugeGraphServer.java |  6 ++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeFactory.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeFactory.java
index de73284598..5276c8a742 100644
--- a/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeFactory.java
+++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeFactory.java
@@ -24,6 +24,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.commons.configuration.Configuration;
 import org.apache.commons.configuration.ConfigurationException;
@@ -43,20 +44,24 @@ public class HugeFactory {
 
     private static final Logger LOG = Log.logger(HugeGraph.class);
 
+    private static final Thread shutdownHook = new Thread(() -> {
+        LOG.info("HugeGraph is shutting down");
+        HugeFactory.shutdown(30L);
+    }, "hugegraph-shutdown");
+
     static {
         SerialEnum.registerInternalEnums();
         HugeGraph.registerTraversalStrategies(StandardHugeGraph.class);
 
-        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
-            LOG.info("HugeGraph is shutting down");
-            HugeFactory.shutdown(30L);
-        }, "hugegraph-shutdown"));
+        Runtime.getRuntime().addShutdownHook(shutdownHook);
     }
 
     private static final String NAME_REGEX = "^[A-Za-z][A-Za-z0-9_]{0,47}$";
 
     private static final Map<String, HugeGraph> graphs = new HashMap<>();
 
+    private static final AtomicBoolean shutdown = new AtomicBoolean(false);
+
     public static synchronized HugeGraph open(Configuration config) {
         HugeConfig conf = config instanceof HugeConfig ?
                           (HugeConfig) config : new HugeConfig(config);
@@ -138,6 +143,9 @@ public static PropertiesConfiguration getRemoteConfig(URL url) {
      * @param timeout seconds
      */
     public static void shutdown(long timeout) {
+        if (!shutdown.compareAndSet(false, true)) {
+            return;
+        }
         try {
             if (!EventHub.destroy(timeout)) {
                 throw new TimeoutException(timeout + "s");
@@ -146,7 +154,12 @@ public static void shutdown(long timeout) {
             OltpTraverser.destroy();
         } catch (Throwable e) {
             LOG.error("Error while shutdown", e);
+            shutdown.compareAndSet(true, false);
             throw new HugeException("Failed to shutdown", e);
         }
     }
+
+    public static void removeShutdownHook() {
+        Runtime.getRuntime().removeShutdownHook(shutdownHook);
+    }
 }
diff --git a/hugegraph-dist/src/main/java/com/baidu/hugegraph/dist/HugeGraphServer.java b/hugegraph-dist/src/main/java/com/baidu/hugegraph/dist/HugeGraphServer.java
index 8d5d5d5853..c7d06fd673 100644
--- a/hugegraph-dist/src/main/java/com/baidu/hugegraph/dist/HugeGraphServer.java
+++ b/hugegraph-dist/src/main/java/com/baidu/hugegraph/dist/HugeGraphServer.java
@@ -122,5 +122,11 @@ public static void main(String[] args) throws Exception {
             server.stop();
             LOG.info("HugeGraphServer stopped");
         }, "hugegraph-server-shutdown"));
+        /*
+         * HugeFactory shutdown hook be invoked before server stop, so that
+         * eventHub execute notify event on a shutdown executor, so here
+         * remove the HugeFactory shutdown hook.
+         */
+        HugeFactory.removeShutdownHook();
     }
 }