From 5a851d72ba448fc18af2f05dc33d8e535c1e76bc Mon Sep 17 00:00:00 2001
From: zyxxoo <1318247699@qq.com>
Date: Thu, 31 Mar 2022 16:23:59 +0800
Subject: [PATCH] improve code

---
 .../hugegraph/auth/ContextGremlinServer.java  |  2 -
 .../baidu/hugegraph/core/GraphManager.java    | 41 +++++++++++--------
 .../java/com/baidu/hugegraph/HugeFactory.java |  2 +
 .../baidu/hugegraph/dist/HugeGraphServer.java | 17 +++++---
 4 files changed, 36 insertions(+), 26 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
index 9e319b4fe9..3aaf7a409b 100644
--- a/hugegraph-api/src/main/java/com/baidu/hugegraph/auth/ContextGremlinServer.java
+++ b/hugegraph-api/src/main/java/com/baidu/hugegraph/auth/ContextGremlinServer.java
@@ -98,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-api/src/main/java/com/baidu/hugegraph/core/GraphManager.java b/hugegraph-api/src/main/java/com/baidu/hugegraph/core/GraphManager.java
index be36b1af4d..3e61929007 100644
--- a/hugegraph-api/src/main/java/com/baidu/hugegraph/core/GraphManager.java
+++ b/hugegraph-api/src/main/java/com/baidu/hugegraph/core/GraphManager.java
@@ -337,28 +337,33 @@ private void checkBackendVersionOrExit(HugeConfig config) {
         for (String graph : this.graphs()) {
             // TODO: close tx from main thread
             HugeGraph hugegraph = this.graph(graph);
-            if (!hugegraph.backendStoreFeatures().supportsPersistence()) {
-                hugegraph.initBackend();
-                if (this.requireAuthentication()) {
-                    String token = config.get(ServerOptions.AUTH_ADMIN_TOKEN);
-                    try {
-                        this.authenticator.initAdminUser(token);
-                    } catch (Exception e) {
-                        throw new BackendException(
-                                  "The backend store of '%s' can't " +
-                                  "initialize admin user", hugegraph.name());
+            try {
+                if (!hugegraph.backendStoreFeatures().supportsPersistence()) {
+                    hugegraph.initBackend();
+                    if (this.requireAuthentication()) {
+                        String token = config.get(
+                                              ServerOptions.AUTH_ADMIN_TOKEN);
+                        try {
+                            this.authenticator.initAdminUser(token);
+                        } catch (Exception e) {
+                            throw new BackendException(
+                                      "The backend store of '%s' can't " +
+                                      "initialize admin user", hugegraph.name());
+                        }
                     }
                 }
-            }
-            BackendStoreSystemInfo info = hugegraph.backendStoreSystemInfo();
-            if (!info.exists()) {
-                throw new BackendException(
+                BackendStoreSystemInfo info = hugegraph.backendStoreSystemInfo();
+                if (!info.exists()) {
+                    throw new BackendException(
                           "The backend store of '%s' has not been initialized",
                           hugegraph.name());
-            }
-            if (!info.checkVersion()) {
-                throw new BackendException(
-                          "The backend store version is inconsistent");
+                }
+                if (!info.checkVersion()) {
+                    throw new BackendException(
+                              "The backend store version is inconsistent");
+                }
+            } finally {
+                hugegraph.tx().close();
             }
         }
     }
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 5276c8a742..56ee22e84b 100644
--- a/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeFactory.java
+++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/HugeFactory.java
@@ -157,6 +157,8 @@ public static void shutdown(long timeout) {
             shutdown.compareAndSet(true, false);
             throw new HugeException("Failed to shutdown", e);
         }
+
+        LOG.info("HugeFactory shutdown");
     }
 
     public static void removeShutdownHook() {
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 c7d06fd673..77c5854774 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
@@ -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,16 +119,19 @@ 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"));
-        /*
-         * 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();
+        serverStopped.get();
     }
 }