feat(db/rocksdb): improve resource management with try-with-resources

Author: halibobo1205

chainbase/src/main/java/org/tron/common/storage/rocksdb/RocksDbDataSourceImpl.java

@@ -308,7 +308,9 @@ public void updateByBatch(Map<byte[], byte[]> rows, WriteOptionsWrapper optionsW
 
   @Override
   public void updateByBatch(Map<byte[], byte[]> rows) {
-    this.updateByBatch(rows, new WriteOptions());
+    try (WriteOptions writeOptions = new WriteOptions()) {
+      this.updateByBatch(rows, writeOptions);
+    }
   }
 
   private void updateByBatch(Map<byte[], byte[]> rows, WriteOptions options) {
@@ -419,6 +421,22 @@ public void backup(String dir) throws RocksDBException {
     }
   }
 
+  /**
+   * Returns an iterator over the database.
+   *
+   * <p><b>CRITICAL:</b> The returned iterator holds native resources and <b>MUST</b> be closed
+   * after use to prevent memory leaks. It is strongly recommended to use a try-with-resources
+   * statement.
+   *
+   * <p>Example of correct usage:
+   * <pre>{@code
+   * try (RocksIterator iterator = getRocksIterator()) {
+   *  // do something
+   * }
+   * }</pre>
+   *
+   * @return a new database iterator that must be closed.
+   */
   private RocksIterator getRocksIterator() {
     try (ReadOptions readOptions = new ReadOptions().setFillCache(false)) {
       throwIfNotAlive();

common/src/main/java/org/tron/common/setting/RocksDbSettings.java

@@ -163,6 +163,21 @@ public static LRUCache getCache() {
     return cache;
   }
 
+  /**
+   * Creates a new RocksDB Options.
+   *
+   * <p><b>CRITICAL:</b> Must be closed after use to prevent native memory leaks.
+   * Use try-with-resources.
+   *
+   * <pre>{@code
+   * try (Options options = getOptionsByDbName(dbName)) {
+   *     // do something
+   * }
+   * }</pre>
+   *
+   * @param dbName  db name
+   * @return a new Options instance that must be closed
+   */
   public static Options getOptionsByDbName(String dbName) {
     RocksDbSettings settings = getSettings();
 

plugins/src/main/java/common/org/tron/plugins/DbConvert.java

@@ -223,8 +223,8 @@ private void batchInsert(RocksDB rocks, List<byte[]> keys, List<byte[]> values)
      * @throws Exception RocksDBException
      */
     private void write(RocksDB rocks, org.rocksdb.WriteBatch batch) throws Exception {
-      try {
-        rocks.write(new org.rocksdb.WriteOptions(), batch);
+      try (org.rocksdb.WriteOptions writeOptions = new org.rocksdb.WriteOptions()) {
+        rocks.write(writeOptions, batch);
       } catch (RocksDBException e) {
         // retry
         if (maybeRetry(e)) {

plugins/src/main/java/common/org/tron/plugins/DbLite.java

@@ -656,12 +656,13 @@ private boolean isLite(String databaseDir) throws RocksDBException, IOException
   private long getSecondBlock(String databaseDir) throws RocksDBException, IOException {
     long num = 0;
     DBInterface sourceBlockIndexDb = DbTool.getDB(databaseDir, BLOCK_INDEX_DB_NAME);
-    DBIterator iterator = sourceBlockIndexDb.iterator();
-    iterator.seek(ByteArray.fromLong(1));
-    if (iterator.hasNext()) {
-      num =  Longs.fromByteArray(iterator.getKey());
+    try (DBIterator iterator = sourceBlockIndexDb.iterator()) {
+      iterator.seek(ByteArray.fromLong(1));
+      if (iterator.hasNext()) {
+        num = Longs.fromByteArray(iterator.getKey());
+      }
+      return num;
     }
-    return num;
   }
 
   private DBInterface getCheckpointDb(String sourceDir) throws IOException, RocksDBException {

plugins/src/main/java/common/org/tron/plugins/utils/DBUtils.java

@@ -88,6 +88,21 @@ public static org.iq80.leveldb.Options newDefaultLevelDbOptions() {
     return dbOptions;
   }
 
+  /**
+   * Creates a new RocksDB Options.
+   *
+   * <p><b>CRITICAL:</b> Must be closed after use to prevent native memory leaks.
+   * Use try-with-resources.
+   *
+   * <pre>{@code
+   * try (Options options = newDefaultRocksDbOptions(false)) {
+   *     // do something
+   * }
+   * }</pre>
+   *
+   * @param forBulkLoad if true, optimizes for bulk loading
+   * @return a new Options instance that must be closed
+   */
   private static Options newDefaultRocksDbOptions(boolean forBulkLoad) {
     Options options = new Options();
     options.setCreateIfMissing(true);

plugins/src/main/java/common/org/tron/plugins/utils/db/DBInterface.java

@@ -12,6 +12,22 @@ public interface DBInterface extends Closeable {
 
   void delete(byte[] key);
 
+  /**
+   * Returns an iterator over the database.
+   *
+   * <p><b>CRITICAL:</b> The returned iterator holds native resources and <b>MUST</b> be closed
+   * after use to prevent memory leaks. It is strongly recommended to use a try-with-resources
+   * statement.
+   *
+   * <p>Example of correct usage:
+   * <pre>{@code
+   * try (DBIterator iterator = db.iterator()) {
+   *  // do something
+   * }
+   * }</pre>
+   *
+   * @return a new database iterator that must be closed.
+   */
   DBIterator iterator();
 
   long size();

plugins/src/main/java/common/org/tron/plugins/utils/db/LevelDBImpl.java

@@ -41,7 +41,7 @@ public DBIterator iterator() {
 
   @Override
   public long size() {
-    return Streams.stream(leveldb.iterator()).count();
+    return Streams.stream(this.iterator()).count();
   }
 
   @Override

plugins/src/main/java/common/org/tron/plugins/utils/db/RockDBIterator.java

@@ -2,14 +2,18 @@
 
 import java.io.IOException;
 import java.util.Map;
+import org.rocksdb.ReadOptions;
+import org.rocksdb.RocksDB;
 import org.rocksdb.RocksIterator;
 
 public class RockDBIterator implements DBIterator {
 
   private final RocksIterator iterator;
+  private final ReadOptions readOptions;
 
-  public RockDBIterator(RocksIterator iterator) {
-    this.iterator = iterator;
+  public RockDBIterator(RocksDB rocksdb) {
+    this.readOptions = new ReadOptions().setFillCache(false);
+    this.iterator = rocksdb.newIterator(readOptions);
   }
 
   @Override
@@ -73,5 +77,6 @@ public byte[] setValue(byte[] value) {
   @Override
   public void close() throws IOException {
     iterator.close();
+    readOptions.close();
   }
 }

plugins/src/main/java/common/org/tron/plugins/utils/db/RocksDBImpl.java

@@ -1,9 +1,9 @@
 package org.tron.plugins.utils.db;
 
+import com.google.common.collect.Streams;
 import java.io.IOException;
 import lombok.Getter;
 import org.rocksdb.RocksDBException;
-import org.rocksdb.RocksIterator;
 
 public class RocksDBImpl implements DBInterface {
 
@@ -47,19 +47,12 @@ public void delete(byte[] key) {
 
   @Override
   public DBIterator iterator() {
-    return new RockDBIterator(rocksDB.newIterator(
-        new org.rocksdb.ReadOptions().setFillCache(false)));
+    return new RockDBIterator(rocksDB);
   }
 
   @Override
   public long size() {
-    RocksIterator iterator = rocksDB.newIterator();
-    long size = 0;
-    for (iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
-      size++;
-    }
-    iterator.close();
-    return size;
+    return Streams.stream(this.iterator()).count();
   }
 
   @Override

plugins/src/test/java/org/tron/plugins/DbTest.java

@@ -5,6 +5,7 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Paths;
 import java.util.UUID;
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.rules.TemporaryFolder;
 import org.rocksdb.RocksDBException;
@@ -83,6 +84,7 @@ private static void initDB(String sourceDir, String dbName, DbTool.DbType dbType
           db.put(bytes, bytes);
         }
       }
+      Assert.assertTrue(db.size() > 0);
     }
   }