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

Author: halibobo1205

actuator/src/main/java/org/tron/core/vm/repository/WriteOptionsWrapper.java

@@ -1,6 +1,8 @@
 package org.tron.common.storage;
 
-public class WriteOptionsWrapper {
+import java.io.Closeable;
+
+public class WriteOptionsWrapper implements Closeable {
 
   public org.rocksdb.WriteOptions rocks = null;
   public org.iq80.leveldb.WriteOptions level = null;
@@ -9,6 +11,23 @@ private WriteOptionsWrapper() {
 
   }
 
+  /**
+   * Returns an WriteOptionsWrapper.
+   *
+   * <p><b>CRITICAL:</b> The returned WriteOptionsWrapper 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 ( WriteOptionsWrapper readOptions = WriteOptionsWrapper.getInstance()) {
+   *  // do something
+   * }
+   * }</pre>
+   *
+   * @return a new WriteOptionsWrapper that must be closed.
+   */
   public static WriteOptionsWrapper getInstance() {
     WriteOptionsWrapper wrapper = new WriteOptionsWrapper();
     wrapper.level = new org.iq80.leveldb.WriteOptions();
@@ -23,4 +42,12 @@ public WriteOptionsWrapper sync(boolean bool) {
     this.rocks.setSync(bool);
     return this;
   }
+
+  @Override
+  public void close() {
+    if (rocks != null) {
+      rocks.close();
+    }
+    // leveldb WriteOptions has no close method, and does not need to be closed
+  }
 }

chainbase/src/main/java/org/tron/common/storage/WriteOptionsWrapper.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();

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

@@ -27,7 +27,7 @@ public abstract class TronDatabase<T> implements ITronChainBase<T> {
   protected DbSourceInter<byte[]> dbSource;
   @Getter
   private String dbName;
-  private WriteOptionsWrapper writeOptions = WriteOptionsWrapper.getInstance()
+  private final WriteOptionsWrapper writeOptions = WriteOptionsWrapper.getInstance()
           .sync(CommonParameter.getInstance().getStorage().isDbSync());
 
   @Autowired
@@ -77,6 +77,7 @@ public void reset() {
   public void close() {
     logger.info("******** Begin to close {}. ********", getName());
     try {
+      writeOptions.close();
       dbSource.closeDB();
     } catch (Exception e) {
       logger.warn("Failed to close {}.", getName(), e);

chainbase/src/main/java/org/tron/core/db/TronDatabase.java

@@ -13,7 +13,7 @@ public class LevelDB implements DB<byte[], byte[]>, Flusher {
 
   @Getter
   private LevelDbDataSourceImpl db;
-  private WriteOptionsWrapper writeOptions = WriteOptionsWrapper.getInstance()
+  private final WriteOptionsWrapper writeOptions = WriteOptionsWrapper.getInstance()
       .sync(CommonParameter.getInstance().getStorage().isDbSync());
 
   public LevelDB(LevelDbDataSourceImpl db) {
@@ -65,6 +65,7 @@ public void flush(Map<WrappedByteArray, WrappedByteArray> batch) {
 
   @Override
   public void close() {
+    this.writeOptions.close();
     db.closeDB();
   }
 

chainbase/src/main/java/org/tron/core/db2/common/LevelDB.java

@@ -14,7 +14,7 @@ public class RocksDB implements DB<byte[], byte[]>, Flusher {
   @Getter
   private RocksDbDataSourceImpl db;
 
-  private WriteOptionsWrapper optionsWrapper = WriteOptionsWrapper.getInstance()
+  private final WriteOptionsWrapper writeOptions = WriteOptionsWrapper.getInstance()
       .sync(CommonParameter.getInstance().getStorage().isDbSync());
 
   public RocksDB(RocksDbDataSourceImpl db) {
@@ -61,11 +61,12 @@ public void flush(Map<WrappedByteArray, WrappedByteArray> batch) {
     Map<byte[], byte[]> rows = batch.entrySet().stream()
         .map(e -> Maps.immutableEntry(e.getKey().getBytes(), e.getValue().getBytes()))
         .collect(HashMap::new, (m, k) -> m.put(k.getKey(), k.getValue()), HashMap::putAll);
-    db.updateByBatch(rows, optionsWrapper);
+    db.updateByBatch(rows, writeOptions);
   }
 
   @Override
   public void close() {
+    writeOptions.close();
     db.closeDB();
   }
 

chainbase/src/main/java/org/tron/core/db2/common/RocksDB.java

@@ -29,7 +29,6 @@
 import org.tron.common.error.TronDBException;
 import org.tron.common.es.ExecutorServiceManager;
 import org.tron.common.parameter.CommonParameter;
-import org.tron.common.storage.WriteOptionsWrapper;
 import org.tron.common.utils.FileUtil;
 import org.tron.common.utils.StorageUtils;
 import org.tron.core.db.RevokingDatabase;
@@ -357,7 +356,6 @@ public void flush() {
 
   public void createCheckpoint() {
     TronDatabase<byte[]> checkPointStore = null;
-    boolean syncFlag;
     try {
       Map<WrappedByteArray, WrappedByteArray> batch = new HashMap<>();
       for (Chainbase db : dbs) {
@@ -389,16 +387,13 @@ public void createCheckpoint() {
       if (isV2Open()) {
         String dbName = String.valueOf(System.currentTimeMillis());
         checkPointStore = getCheckpointDB(dbName);
-        syncFlag = CommonParameter.getInstance().getStorage().isCheckpointSync();
       } else {
         checkPointStore = checkTmpStore;
-        syncFlag = CommonParameter.getInstance().getStorage().isDbSync();
       }
 
-      checkPointStore.getDbSource().updateByBatch(batch.entrySet().stream()
+      checkPointStore.updateByBatch(batch.entrySet().stream()
               .map(e -> Maps.immutableEntry(e.getKey().getBytes(), e.getValue().getBytes()))
-              .collect(HashMap::new, (m, k) -> m.put(k.getKey(), k.getValue()), HashMap::putAll),
-          WriteOptionsWrapper.getInstance().sync(syncFlag));
+              .collect(HashMap::new, (m, k) -> m.put(k.getKey(), k.getValue()), HashMap::putAll));
 
     } catch (Exception e) {
       throw new TronDBException(e);

chainbase/src/main/java/org/tron/core/db2/core/SnapshotManager.java

@@ -1,18 +1,23 @@
 package org.tron.core.store;
 
 import com.google.protobuf.InvalidProtocolBufferException;
+import java.util.Map;
+import java.util.Spliterator;
+import java.util.function.Consumer;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.tron.common.parameter.CommonParameter;
+import org.tron.common.storage.WriteOptionsWrapper;
 import org.tron.core.db.TronDatabase;
 import org.tron.core.exception.BadItemException;
 import org.tron.core.exception.ItemNotFoundException;
 
-import java.util.Spliterator;
-import java.util.function.Consumer;
-
 @Slf4j(topic = "DB")
 public class CheckPointV2Store extends TronDatabase<byte[]> {
 
+  private final WriteOptionsWrapper writeOptions = WriteOptionsWrapper.getInstance()
+      .sync(CommonParameter.getInstance().getStorage().isCheckpointSync());
+
   @Autowired
   public CheckPointV2Store(String dbPath) {
     super(dbPath);
@@ -52,13 +57,19 @@ public Spliterator spliterator() {
   protected void init() {
   }
 
+  @Override
+  public void updateByBatch(Map<byte[], byte[]> rows) {
+    this.dbSource.updateByBatch(rows, writeOptions);
+  }
+
   /**
    * close the database.
    */
   @Override
   public void close() {
     logger.debug("******** Begin to close {}. ********", getName());
     try {
+      writeOptions.close();
       dbSource.closeDB();
     } catch (Exception e) {
       logger.warn("Failed to close {}.", getName(), e);

chainbase/src/main/java/org/tron/core/store/CheckPointV2Store.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();
 

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

@@ -172,7 +172,9 @@ public void testupdateByBatchInner() {
     rows.clear();
     rows.put(key1.getBytes(), null);
     rows.put(key2.getBytes(), null);
-    dataSource.updateByBatch(rows, WriteOptionsWrapper.getInstance());
+    try (WriteOptionsWrapper options = WriteOptionsWrapper.getInstance()) {
+      dataSource.updateByBatch(rows, options);
+    }
     assertEquals(0, dataSource.allKeys().size());
 
     rows.clear();