From 20704bb236bb529e828029a104a24e3abba846db Mon Sep 17 00:00:00 2001
From: "vaughn.zhang" <vaughn.zhang@zoom.us>
Date: Sat, 16 Dec 2023 09:50:40 +0800
Subject: [PATCH] fix: rocksdb lock avaliable

---
 .../apache/hugegraph/StandardHugeGraph.java   |  2 +-
 .../store/rocksdb/RocksDBStdSessions.java     |  7 +++++
 .../backend/store/rocksdb/RocksDBStore.java   | 10 +++++++
 .../store/rocksdb/RocksDBStoreProvider.java   | 28 +++++++++++++++++++
 4 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java
index cb2e369ea0..d55c474c59 100644
--- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java
+++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java
@@ -971,8 +971,8 @@ public synchronized void close() throws Exception {
             this.closeTx();
         } finally {
             this.closed = true;
-            this.storeProvider.close();
             LockUtil.destroy(this.name);
+            this.storeProvider.close();
         }
         // Make sure that all transactions are closed in all threads
         if (!tx.closed()) {
diff --git a/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStdSessions.java b/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStdSessions.java
index bcbe37b7c3..456976605b 100644
--- a/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStdSessions.java
+++ b/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStdSessions.java
@@ -343,6 +343,13 @@ protected synchronized void doClose() {
         this.rocksdb.close();
     }
 
+    public synchronized void forceClose() {
+        if (!this.rocksdb.isOwningHandle()) {
+            return;
+        }
+        this.rocksdb.close();
+    }
+
     private void checkValid() {
         E.checkState(this.rocksdb.isOwningHandle(),
                      "It seems RocksDB has been closed");
diff --git a/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStore.java b/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStore.java
index 283baa622a..e7e28f00cd 100644
--- a/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStore.java
+++ b/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStore.java
@@ -423,6 +423,16 @@ public void close() {
         this.closeSessions();
     }
 
+    public void forceClose() {
+        try {
+            this.checkOpened();
+            this.closeSessions();
+        } catch (Throwable ignore) {
+            return;
+        }
+        ((RocksDBStdSessions)this.sessions).forceClose();
+    }
+
     @Override
     public boolean opened() {
         this.checkDbOpened();
diff --git a/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStoreProvider.java b/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStoreProvider.java
index 2e5006a54c..44041125f1 100644
--- a/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStoreProvider.java
+++ b/hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBStoreProvider.java
@@ -19,13 +19,18 @@
 
 import java.io.File;
 
+import org.apache.hugegraph.backend.BackendException;
 import org.apache.hugegraph.backend.store.AbstractBackendStoreProvider;
 import org.apache.hugegraph.backend.store.BackendStore;
 import org.apache.hugegraph.config.HugeConfig;
 import org.apache.hugegraph.util.ConfigUtil;
+import org.apache.hugegraph.util.Log;
+import org.slf4j.Logger;
 
 public class RocksDBStoreProvider extends AbstractBackendStoreProvider {
 
+    private static final Logger LOG = Log.logger(RocksDBStoreProvider.class);
+
     protected String database() {
         return this.graph().toLowerCase();
     }
@@ -68,6 +73,29 @@ protected BackendStore newSystemStore(HugeConfig config, String store) {
         return new RocksDBStore.RocksDBSystemStore(this, this.database(), store);
     }
 
+    @Override
+    public void close() throws BackendException {
+        super.close();
+        if (this.stores == null) {
+            return;
+        }
+        this.stores.values().forEach(store -> {
+            try {
+                if (!store.opened()) {
+                    return;
+                }
+            } catch (Throwable ignore) {
+                return;
+            }
+            try {
+                ((RocksDBStore)store).forceClose();
+            } catch (Exception e) {
+                LOG.error("Failed to close store '%s'", store.store(), e);
+                throw new BackendException("Failed to close store '%s'", store.store(), e);
+            }
+        });
+    }
+
     @Override
     public String type() {
         return "rocksdb";