Refactor

Signed-off-by: Simon Dudley <[email protected]>
siladu · Jan 29, 2024 · 7a55c97 · 7a55c97
1 parent 47e9f43
commit 7a55c97
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 67 deletions.
diff --git a/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RocksDbSubCommand.java b/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RocksDbSubCommand.java
@@ -19,13 +19,7 @@
 import org.hyperledger.besu.cli.util.VersionProvider;
 
 import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.List;
 
-import org.rocksdb.ColumnFamilyDescriptor;
-import org.rocksdb.ColumnFamilyHandle;
-import org.rocksdb.Options;
-import org.rocksdb.RocksDB;
 import org.rocksdb.RocksDBException;
 import picocli.CommandLine;
 import picocli.CommandLine.Command;
@@ -82,34 +76,18 @@ public void run() {
               .concat("/")
               .concat(DATABASE_PATH);
 
-      RocksDB.loadLibrary();
-      Options options = new Options();
-      options.setCreateIfMissing(true);
-
-      // Open the RocksDB database with multiple column families
-      List<byte[]> cfNames;
-      try {
-        cfNames = RocksDB.listColumnFamilies(options, dbPath);
-      } catch (RocksDBException e) {
-        throw new RuntimeException(e);
-      }
-      final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
-      final List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>();
-      for (byte[] cfName : cfNames) {
-        cfDescriptors.add(new ColumnFamilyDescriptor(cfName));
-      }
       RocksDbUsageHelper.printTableHeader(out);
-      try (final RocksDB rocksdb = RocksDB.openReadOnly(dbPath, cfDescriptors, cfHandles)) {
-        for (ColumnFamilyHandle cfHandle : cfHandles) {
-          RocksDbUsageHelper.printUsageForColumnFamily(rocksdb, cfHandle, out);
-        }
-      } catch (RocksDBException e) {
-        throw new RuntimeException(e);
-      } finally {
-        for (ColumnFamilyHandle cfHandle : cfHandles) {
-          cfHandle.close();
-        }
-      }
+
+      RocksDbUsageHelper.forEachColumnFamily(
+          dbPath,
+          (rocksdb, cfHandle) -> {
+            try {
+              RocksDbUsageHelper.printUsageForColumnFamily(rocksdb, cfHandle, out);
+            } catch (RocksDBException e) {
+              throw new RuntimeException(e);
+            }
+            return null;
+          });
     }
   }
 }
diff --git a/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RocksDbUsageHelper.java b/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RocksDbUsageHelper.java
@@ -18,9 +18,14 @@
 import org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier;
 
 import java.io.PrintWriter;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.BiFunction;
 
 import org.bouncycastle.util.Arrays;
+import org.rocksdb.ColumnFamilyDescriptor;
 import org.rocksdb.ColumnFamilyHandle;
+import org.rocksdb.Options;
 import org.rocksdb.RocksDB;
 import org.rocksdb.RocksDBException;
 import org.slf4j.Logger;
@@ -30,6 +35,37 @@
 public class RocksDbUsageHelper {
   private static final Logger LOG = LoggerFactory.getLogger(RocksDbUsageHelper.class);
 
+  static void forEachColumnFamily(
+      final String dbPath, final BiFunction<RocksDB, ColumnFamilyHandle, Void> task) {
+    RocksDB.loadLibrary();
+    Options options = new Options();
+    options.setCreateIfMissing(true);
+
+    // Open the RocksDB database with multiple column families
+    List<byte[]> cfNames;
+    try {
+      cfNames = RocksDB.listColumnFamilies(options, dbPath);
+    } catch (RocksDBException e) {
+      throw new RuntimeException(e);
+    }
+    final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
+    final List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>();
+    for (byte[] cfName : cfNames) {
+      cfDescriptors.add(new ColumnFamilyDescriptor(cfName));
+    }
+    try (final RocksDB rocksdb = RocksDB.openReadOnly(dbPath, cfDescriptors, cfHandles)) {
+      for (ColumnFamilyHandle cfHandle : cfHandles) {
+        task.apply(rocksdb, cfHandle);
+      }
+    } catch (RocksDBException e) {
+      throw new RuntimeException(e);
+    } finally {
+      for (ColumnFamilyHandle cfHandle : cfHandles) {
+        cfHandle.close();
+      }
+    }
+  }
+
   static void printUsageForColumnFamily(
       final RocksDB rocksdb, final ColumnFamilyHandle cfHandle, final PrintWriter out)
       throws RocksDBException, NumberFormatException {

diff --git a/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/TrieLogSubCommand.java b/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/TrieLogSubCommand.java
@@ -33,16 +33,12 @@
 import java.io.PrintWriter;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.core.config.Configurator;
-import org.rocksdb.ColumnFamilyDescriptor;
-import org.rocksdb.ColumnFamilyHandle;
-import org.rocksdb.Options;
-import org.rocksdb.RocksDB;
 import org.rocksdb.RocksDBException;
 import org.slf4j.LoggerFactory;
 import picocli.CommandLine;
@@ -161,35 +157,24 @@ private long getEstimatedSaving() {
               .toString()
               .concat("/")
               .concat(DATABASE_PATH);
-      RocksDB.loadLibrary();
-      Options options = new Options();
-      options.setCreateIfMissing(true);
-      final List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>();
-      List<byte[]> cfNames;
-      try {
-        cfNames = RocksDB.listColumnFamilies(options, dbPath);
-      } catch (RocksDBException e) {
-        throw new RuntimeException(e);
-      }
-      cfDescriptors.add(new ColumnFamilyDescriptor(cfNames.get(0)));
-      cfDescriptors.add(new ColumnFamilyDescriptor(TRIE_LOG_STORAGE.getId()));
-      final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
-      long estimatedSaving = 0L;
-      try (final RocksDB rocksdb = RocksDB.openReadOnly(dbPath, cfDescriptors, cfHandles)) {
-        for (ColumnFamilyHandle cfHandle : cfHandles) {
-          if (Arrays.equals(cfHandle.getName(), TRIE_LOG_STORAGE.getId())) {
-            estimatedSaving =
-                Long.parseLong(rocksdb.getProperty(cfHandle, "rocksdb.estimate-live-data-size"));
-          }
-        }
-      } catch (RocksDBException e) {
-        throw new RuntimeException(e);
-      } finally {
-        for (ColumnFamilyHandle cfHandle : cfHandles) {
-          cfHandle.close();
-        }
-      }
-      return estimatedSaving;
+
+      AtomicLong estimatedSaving = new AtomicLong(0L);
+      RocksDbUsageHelper.forEachColumnFamily(
+          dbPath,
+          (rocksdb, cfHandle) -> {
+            try {
+              if (Arrays.equals(cfHandle.getName(), TRIE_LOG_STORAGE.getId())) {
+                estimatedSaving.set(
+                    Long.parseLong(
+                        rocksdb.getProperty(cfHandle, "rocksdb.estimate-live-data-size")));
+              }
+            } catch (RocksDBException e) {
+              throw new RuntimeException(e);
+            }
+            return null;
+          });
+
+      return estimatedSaving.get();
     }
   }