Fix close exception and server-info EXPIRED_INTERVAL (#1804)

apache · Apr 7, 2022 · d778eb1 · d778eb1
1 parent 4a2c2d2
commit d778eb1
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 7 deletions.
diff --git a/hugegraph-api/src/main/java/com/baidu/hugegraph/auth/ContextGremlinServer.java b/hugegraph-api/src/main/java/com/baidu/hugegraph/auth/ContextGremlinServer.java
@@ -53,6 +53,10 @@ public class ContextGremlinServer extends GremlinServer {
 
     private final EventHub eventHub;
 
+    static {
+        HugeGraphAuthProxy.setContext(Context.admin());
+    }
+
     public ContextGremlinServer(final Settings settings, EventHub eventHub) {
         /*
          * pass custom Executor https://github.com/apache/tinkerpop/pull/813
@@ -94,8 +98,6 @@ public synchronized CompletableFuture<Void> stop() {
     }
 
     public void injectAuthGraph() {
-        HugeGraphAuthProxy.setContext(Context.admin());
-
         GraphManager manager = this.getServerGremlinExecutor()
                                    .getGraphManager();
         for (String name : manager.getGraphNames()) {

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,14 @@ 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);
         }
+
+        LOG.info("HugeFactory shutdown");
+    }
+
+    public static void removeShutdownHook() {
+        Runtime.getRuntime().removeShutdownHook(shutdownHook);
     }
 }
diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/task/HugeServerInfo.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/task/HugeServerInfo.java
@@ -51,7 +51,7 @@ public class HugeServerInfo {
 
     // Unit millisecond
     private static final long EXPIRED_INTERVAL =
-                              TaskManager.SCHEDULE_PERIOD * 1000L * 3;
+                              TaskManager.SCHEDULE_PERIOD * 10;
 
     private Id id;
     private NodeRole role;

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
@@ -31,6 +31,8 @@
 import com.baidu.hugegraph.util.ConfigUtil;
 import com.baidu.hugegraph.util.Log;
 
+import java.util.concurrent.CompletableFuture;
+
 public class HugeGraphServer {
 
     private static final Logger LOG = Log.logger(HugeGraphServer.class);
@@ -117,10 +119,21 @@ public static void main(String[] args) throws Exception {
         HugeGraphServer.register();
 
         HugeGraphServer server = new HugeGraphServer(args[0], args[1]);
+
+        /*
+         * HugeFactory.shutdown hook may be invoked before server stop,
+         * causes event-hub can't execute notification events for another
+         * shutdown executor such as gremling-stop-shutdown
+         */
+        HugeFactory.removeShutdownHook();
+
+        CompletableFuture<?> serverStopped = new CompletableFuture<>();
         Runtime.getRuntime().addShutdownHook(new Thread(() -> {
             LOG.info("HugeGraphServer stopping");
             server.stop();
             LOG.info("HugeGraphServer stopped");
+            serverStopped.complete(null);
         }, "hugegraph-server-shutdown"));
+        serverStopped.get();
     }
 }