fix: eventHub notify after shutdown

apache · Mar 31, 2022 · 2eb6c9b · 2eb6c9b
1 parent f621ea7
commit 2eb6c9b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 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
@@ -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
@@ -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();
     }
 }