Fix the coredump that occurs when calling KeyValueStorageRocksDB.count after rocksdb has been closed (#4581)

* Fix the coredump that occurs when calling KeyValueStorageRocksDB.count() (possibly triggered by Prometheus) after RocksDB has been closed(#4243)

* fix race when count op in process and db gets closed.

---------

Co-authored-by: zhaizhibo <zhaizhibo@kuaishou.com>
(cherry picked from commit 2831ed3)

Author: 2 people

bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/ldb/KeyValueStorageRocksDB.java

@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.Map.Entry;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import org.apache.bookkeeper.bookie.storage.ldb.KeyValueStorageFactory.DbConfigType;
 import org.apache.bookkeeper.conf.ServerConfiguration;
 import org.apache.commons.lang3.StringUtils;
@@ -73,6 +74,8 @@ public class KeyValueStorageRocksDB implements KeyValueStorage {
     static KeyValueStorageFactory factory = (defaultBasePath, subPath, dbConfigType, conf) ->
             new KeyValueStorageRocksDB(defaultBasePath, subPath, dbConfigType, conf);
 
+    private volatile boolean closed = true;
+    private final ReentrantReadWriteLock closedLock = new ReentrantReadWriteLock();
     private final RocksDB db;
     private RocksObject options;
     private List<ColumnFamilyDescriptor> columnFamilyDescriptors;
@@ -144,6 +147,7 @@ public KeyValueStorageRocksDB(String basePath, String subPath, DbConfigType dbCo
         optionDontCache.setFillCache(false);
 
         this.writeBatchMaxSize = conf.getMaxOperationNumbersInSingleRocksDBBatch();
+        this.closed = false;
     }
 
     private RocksDB initializeRocksDBWithConfFile(String basePath, String subPath, DbConfigType dbConfigType,
@@ -284,7 +288,13 @@ private RocksDB initializeRocksDBWithBookieConf(String basePath, String subPath,
 
     @Override
     public void close() throws IOException {
-        db.close();
+        try {
+            closedLock.writeLock().lock();
+            closed = true;
+            db.close();
+        } finally {
+            closedLock.writeLock().unlock();
+        }
         if (cache != null) {
             cache.close();
         }
@@ -475,7 +485,18 @@ public void close() {
     @Override
     public long count() throws IOException {
         try {
-            return db.getLongProperty("rocksdb.estimate-num-keys");
+            if (closed) {
+                throw new IOException("RocksDB is closed");
+            }
+            try {
+                closedLock.readLock().lock();
+                if (!closed) {
+                    return db.getLongProperty("rocksdb.estimate-num-keys");
+                }
+                throw new IOException("RocksDB is closed");
+            } finally {
+                closedLock.readLock().unlock();
+            }
         } catch (RocksDBException e) {
             throw new IOException("Error in getting records count", e);
         }

bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/storage/ldb/KeyValueStorageRocksDBTest.java

@@ -21,9 +21,11 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import java.io.File;
+import java.io.IOException;
 import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Paths;
@@ -118,4 +120,19 @@ public void testReadChecksumTypeFromConfigurationFile() throws Exception {
         // After the PR: https://github.com/facebook/rocksdb/pull/10826 merge, we can turn on this test.
         assertEquals(ChecksumType.kxxHash, ((BlockBasedTableConfig) familyOptions.tableFormatConfig()).checksumType());
     }
+
+    @Test
+    public void testCallCountAfterClose() throws IOException {
+        ServerConfiguration configuration = new ServerConfiguration();
+        URL url = getClass().getClassLoader().getResource("test_entry_location_rocksdb.conf");
+        configuration.setEntryLocationRocksdbConf(url.getPath());
+        File tmpDir = Files.createTempDirectory("bk-kv-rocksdbtest-file").toFile();
+        Files.createDirectory(Paths.get(tmpDir.toString(), "subDir"));
+        KeyValueStorageRocksDB rocksDB = new KeyValueStorageRocksDB(tmpDir.toString(), "subDir",
+                KeyValueStorageFactory.DbConfigType.EntryLocation, configuration);
+        assertNotNull(rocksDB.getColumnFamilyDescriptors());
+        rocksDB.close();
+        IOException exception = assertThrows(IOException.class, rocksDB::count);
+        assertEquals("RocksDB is closed", exception.getMessage());
+    }
 }