From 6f247c0fdc8ae5f1c34816020f5c38e5698eb15d Mon Sep 17 00:00:00 2001
From: contrueCT <contrue_CT@outlook.com>
Date: Fri, 9 Oct 2026 20:05:08 +0800
Subject: [PATCH] fix(rocksdb): synchronize session close and reset

Serialize native batch disposal with cross-thread reset.
Skip reset after the session has closed.
Cover the disposed-but-registered window with real JNI.
Keep retained closed references safe after pool removal.
---
 .../store/rocksdb/RocksDBStdSessions.java     |  9 ++-
 .../unit/rocksdb/RocksDBSessionsTest.java     | 81 +++++++++++++++++++
 2 files changed, 87 insertions(+), 3 deletions(-)

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 7dca2fc0ec..a6d6615cf2 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
@@ -732,7 +732,7 @@ public void open() {
         }
 
         @Override
-        public void close() {
+        public synchronized void close() {
             assert this.closeable();
             try {
                 this.batch.close();
@@ -751,8 +751,11 @@ public boolean closed() {
         }
 
         @Override
-        public void reset() {
-            this.batch.clear();
+        public synchronized void reset() {
+            // A pool iterator may retain this session after its final detach.
+            if (this.opened) {
+                this.batch.clear();
+            }
         }
 
         /**
diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/rocksdb/RocksDBSessionsTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/rocksdb/RocksDBSessionsTest.java
index 5c8b14c71f..a4897533e7 100644
--- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/rocksdb/RocksDBSessionsTest.java
+++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/rocksdb/RocksDBSessionsTest.java
@@ -55,6 +55,87 @@
 
 public class RocksDBSessionsTest extends BaseRocksDBUnitTest {
 
+    @Test
+    public void testResetWaitsForNativeCloseAndSkipsRetainedClosedSession() throws Exception {
+        CountDownLatch disposed = new CountDownLatch(1);
+        CountDownLatch finishClose = new CountDownLatch(1);
+        CountDownLatch resetStarted = new CountDownLatch(1);
+        CountDownLatch resetFinished = new CountDownLatch(1);
+        AtomicReference<Throwable> failure = new AtomicReference<>();
+        AtomicReference<RocksDBSessions.Session> closing = new AtomicReference<>();
+        Thread closer = new Thread(() -> {
+            try {
+                RocksDBSessions.Session session = this.rocks.session();
+                closing.set(session);
+                WriteBatch previous = Whitebox.getInternalState(session, "batch");
+                previous.close();
+                WriteBatch guarded = new WriteBatch() {
+                    @Override
+                    public void close() {
+                        super.close();
+                        disposed.countDown();
+                        try {
+                            Assert.assertTrue(finishClose.await(30, TimeUnit.SECONDS));
+                        } catch (InterruptedException e) {
+                            Thread.currentThread().interrupt();
+                            throw new AssertionError(e);
+                        }
+                    }
+
+                    @Override
+                    public void clear() {
+                        // Trap the bad interleaving before calling a freed JNI address.
+                        Assert.assertTrue("Reset touched a disposed native batch", this.isOwningHandle());
+                        super.clear();
+                    }
+                };
+                Whitebox.setInternalState(session, "batch", guarded);
+                this.rocks.close();
+            } catch (Throwable e) {
+                failure.compareAndSet(null, e);
+                disposed.countDown();
+            }
+        }, "session-close");
+        Thread resetter = new Thread(() -> {
+            resetStarted.countDown();
+            try {
+                this.rocks.forceResetSessions();
+            } catch (Throwable e) {
+                failure.compareAndSet(null, e);
+            } finally {
+                resetFinished.countDown();
+            }
+        }, "session-reset");
+        closer.start();
+        try {
+            Assert.assertTrue(disposed.await(30, TimeUnit.SECONDS));
+            Assert.assertNull(failure.get());
+            // Close is still running and the pool still exposes this session.
+            Assert.assertTrue(closing.get().opened());
+            resetter.start();
+            Assert.assertTrue(resetStarted.await(30, TimeUnit.SECONDS));
+            long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(30);
+            while (resetter.getState() != Thread.State.BLOCKED && resetFinished.getCount() != 0 &&
+                   System.nanoTime() < deadline) {
+                Thread.sleep(1L);
+            }
+            Assert.assertNull(failure.get());
+            Assert.assertEquals(Thread.State.BLOCKED, resetter.getState());
+        } finally {
+            finishClose.countDown();
+            closer.join(TimeUnit.SECONDS.toMillis(30));
+            resetter.join(TimeUnit.SECONDS.toMillis(30));
+        }
+        Assert.assertFalse(closer.isAlive());
+        Assert.assertFalse(resetter.isAlive());
+        Assert.assertNull(failure.get());
+        Assert.assertTrue(closing.get().closed());
+        // A reference retained before pool removal must remain safe to reset.
+        closing.get().reset();
+        this.put("after-close", "usable");
+        Assert.assertEquals("usable", this.get("after-close"));
+    }
+
     @Test
     public void testResetDiscardsPendingWritesWithoutLeakingNativeBatch() throws Exception {
         RocksDBSessions.Session session = this.rocks.session();