Import export trie log

besu/src/main/java/org/hyperledger/besu/cli/options/stable/DataStorageOptions.java

@@ -42,6 +42,11 @@ public class DataStorageOptions implements CLIOptions<DataStorageConfiguration>
   private static final String BONSAI_STORAGE_FORMAT_MAX_LAYERS_TO_LOAD =
       "--bonsai-historical-block-limit";
 
+  private static final String BONSAI_TRIE_LOG_PRUNING_ENABLED =
+      "--Xbonsai-trie-log-pruning-enabled";
+
+  private static final String BONSAI_LIMIT_TRIE_LOGS_ENABLED = "--Xbonsai-limit-trie-logs-enabled";
+
   // Use Bonsai DB
   @Option(
       names = {DATA_STORAGE_FORMAT},
@@ -65,7 +70,7 @@ static class Unstable {
 
     @CommandLine.Option(
         hidden = true,
-        names = {"--Xbonsai-trie-log-pruning-enabled"},
+        names = {BONSAI_LIMIT_TRIE_LOGS_ENABLED, BONSAI_TRIE_LOG_PRUNING_ENABLED},
         description = "Enable trie log pruning. (default: ${DEFAULT-VALUE})")
     private boolean bonsaiTrieLogPruningEnabled = DEFAULT_BONSAI_TRIE_LOG_PRUNING_ENABLED;
 

besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/TrieLogHelper.java

@@ -16,6 +16,7 @@
 package org.hyperledger.besu.cli.subcommands.storage;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Collections.singletonList;
 import static org.hyperledger.besu.controller.BesuController.DATABASE_PATH;
 
 import org.hyperledger.besu.datatypes.Hash;
@@ -97,16 +98,15 @@ private static void processTrieLogBatches(
       final String batchFileNameBase) {
 
     for (long batchNumber = 1; batchNumber <= numberOfBatches; batchNumber++) {
-
+      final String batchFileName = batchFileNameBase + "-" + batchNumber;
       final long firstBlockOfBatch = chainHeight - ((batchNumber - 1) * BATCH_SIZE);
-
       final long lastBlockOfBatch =
           Math.max(chainHeight - (batchNumber * BATCH_SIZE), lastBlockNumberToRetainTrieLogsFor);
-
       final List<Hash> trieLogKeys =
           getTrieLogKeysForBlocks(blockchain, firstBlockOfBatch, lastBlockOfBatch);
 
-      saveTrieLogBatches(batchFileNameBase, rootWorldStateStorage, batchNumber, trieLogKeys);
+      LOG.info("Saving trie logs to retain in file (batch {})...", batchNumber);
+      saveTrieLogBatches(batchFileName, rootWorldStateStorage, trieLogKeys);
     }
 
     LOG.info("Clear trie logs...");
@@ -118,15 +118,12 @@ private static void processTrieLogBatches(
   }
 
   private static void saveTrieLogBatches(
-      final String batchFileNameBase,
+      final String batchFileName,
       final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
-      final long batchNumber,
       final List<Hash> trieLogKeys) {
 
-    LOG.info("Saving trie logs to retain in file (batch {})...", batchNumber);
-
     try {
-      saveTrieLogsInFile(trieLogKeys, rootWorldStateStorage, batchNumber, batchFileNameBase);
+      saveTrieLogsInFile(trieLogKeys, rootWorldStateStorage, batchFileName);
     } catch (IOException e) {
       LOG.error("Error saving trie logs to file: {}", e.getMessage());
       throw new RuntimeException(e);
@@ -210,9 +207,8 @@ private static void recreateTrieLogs(
       final String batchFileNameBase)
       throws IOException {
     // process in chunk to avoid OOM
-
-    IdentityHashMap<byte[], byte[]> trieLogsToRetain =
-        readTrieLogsFromFile(batchFileNameBase, batchNumber);
+    final String batchFileName = batchFileNameBase + "-" + batchNumber;
+    IdentityHashMap<byte[], byte[]> trieLogsToRetain = readTrieLogsFromFile(batchFileName);
     final int chunkSize = ROCKSDB_MAX_INSERTS_PER_TRANSACTION;
     List<byte[]> keys = new ArrayList<>(trieLogsToRetain.keySet());
 
@@ -265,11 +261,10 @@ private static void validatePruneConfiguration(final DataStorageConfiguration co
   private static void saveTrieLogsInFile(
       final List<Hash> trieLogsKeys,
       final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
-      final long batchNumber,
-      final String batchFileNameBase)
+      final String batchFileName)
       throws IOException {
 
-    File file = new File(batchFileNameBase + "-" + batchNumber);
+    File file = new File(batchFileName);
     if (file.exists()) {
       LOG.error("File already exists, skipping file creation");
       return;
@@ -285,11 +280,10 @@ private static void saveTrieLogsInFile(
   }
 
   @SuppressWarnings("unchecked")
-  private static IdentityHashMap<byte[], byte[]> readTrieLogsFromFile(
-      final String batchFileNameBase, final long batchNumber) {
+  static IdentityHashMap<byte[], byte[]> readTrieLogsFromFile(final String batchFileName) {
 
     IdentityHashMap<byte[], byte[]> trieLogs;
-    try (FileInputStream fis = new FileInputStream(batchFileNameBase + "-" + batchNumber);
+    try (FileInputStream fis = new FileInputStream(batchFileName);
         ObjectInputStream ois = new ObjectInputStream(fis)) {
 
       trieLogs = (IdentityHashMap<byte[], byte[]>) ois.readObject();
@@ -357,5 +351,30 @@ static void printCount(final PrintWriter out, final TrieLogCount count) {
         count.total, count.canonicalCount, count.forkCount, count.orphanCount);
   }
 
+  static void importTrieLog(
+      final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
+      final Path dataDirectoryPath,
+      final Hash trieLogHash) {
+    final String trieLogFile =
+        dataDirectoryPath.resolve(DATABASE_PATH).resolve(trieLogHash.toString()).toString();
+
+    var trieLog = readTrieLogsFromFile(trieLogFile);
+
+    var updater = rootWorldStateStorage.updater();
+    trieLog.forEach((key, value) -> updater.getTrieLogStorageTransaction().put(key, value));
+    updater.getTrieLogStorageTransaction().commit();
+  }
+
+  static void exportTrieLog(
+      final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
+      final Path dataDirectoryPath,
+      final Hash trieLogHash)
+      throws IOException {
+    final String trieLogFile =
+        dataDirectoryPath.resolve(DATABASE_PATH).resolve(trieLogHash.toString()).toString();
+
+    saveTrieLogsInFile(singletonList(trieLogHash), rootWorldStateStorage, trieLogFile);
+  }
+
   record TrieLogCount(int total, int canonicalCount, int forkCount, int orphanCount) {}
 }

besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/TrieLogSubCommand.java

@@ -17,15 +17,18 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import org.hyperledger.besu.cli.DefaultCommandValues;
 import org.hyperledger.besu.cli.util.VersionProvider;
 import org.hyperledger.besu.controller.BesuController;
+import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
 import org.hyperledger.besu.ethereum.storage.StorageProvider;
 import org.hyperledger.besu.ethereum.trie.bonsai.storage.BonsaiWorldStateKeyValueStorage;
 import org.hyperledger.besu.ethereum.trie.bonsai.trielog.TrieLogPruner;
 import org.hyperledger.besu.ethereum.worldstate.DataStorageConfiguration;
 import org.hyperledger.besu.ethereum.worldstate.DataStorageFormat;
 
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -43,7 +46,12 @@
     description = "Manipulate trie logs",
     mixinStandardHelpOptions = true,
     versionProvider = VersionProvider.class,
-    subcommands = {TrieLogSubCommand.CountTrieLog.class, TrieLogSubCommand.PruneTrieLog.class})
+    subcommands = {
+      TrieLogSubCommand.CountTrieLog.class,
+      TrieLogSubCommand.PruneTrieLog.class,
+      TrieLogSubCommand.ExportTrieLog.class,
+      TrieLogSubCommand.ImportTrieLog.class
+    })
 public class TrieLogSubCommand implements Runnable {
 
   @SuppressWarnings("UnusedVariable")
@@ -123,6 +131,78 @@ public void run() {
     }
   }
 
+  @Command(
+      name = "export",
+      description =
+          "This command prunes all trie log layers below the retention threshold, including orphaned trie logs.",
+      mixinStandardHelpOptions = true,
+      versionProvider = VersionProvider.class)
+  static class ExportTrieLog implements Runnable {
+
+    @SuppressWarnings("unused")
+    @ParentCommand
+    private TrieLogSubCommand parentCommand;
+
+    @SuppressWarnings("unused")
+    @CommandLine.Spec
+    private CommandLine.Model.CommandSpec spec; // Picocli injects reference to command spec
+
+    @CommandLine.Option(
+        names = "--trie-log-hash",
+        paramLabel = DefaultCommandValues.MANDATORY_LONG_FORMAT_HELP,
+        description = "The hash of the block you want to export the trie log.",
+        arity = "1..1")
+    private String trieLogHash;
+
+    @Override
+    public void run() {
+      TrieLogContext context = getTrieLogContext();
+      final Path dataDirectoryPath =
+          Paths.get(
+              TrieLogSubCommand.parentCommand.parentCommand.dataDir().toAbsolutePath().toString());
+      try {
+        TrieLogHelper.exportTrieLog(
+            context.rootWorldStateStorage(), dataDirectoryPath, Hash.fromHexString(trieLogHash));
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
+
+  @Command(
+      name = "import",
+      description =
+          "This command prunes all trie log layers below the retention threshold, including orphaned trie logs.",
+      mixinStandardHelpOptions = true,
+      versionProvider = VersionProvider.class)
+  static class ImportTrieLog implements Runnable {
+
+    @SuppressWarnings("unused")
+    @ParentCommand
+    private TrieLogSubCommand parentCommand;
+
+    @SuppressWarnings("unused")
+    @CommandLine.Spec
+    private CommandLine.Model.CommandSpec spec; // Picocli injects reference to command spec
+
+    @CommandLine.Option(
+        names = "--trie-log-hash",
+        paramLabel = DefaultCommandValues.MANDATORY_LONG_FORMAT_HELP,
+        description = "The hash of the block you want to import the trie log.",
+        arity = "1..1")
+    private String trieLogHash;
+
+    @Override
+    public void run() {
+      TrieLogContext context = getTrieLogContext();
+      final Path dataDirectoryPath =
+          Paths.get(
+              TrieLogSubCommand.parentCommand.parentCommand.dataDir().toAbsolutePath().toString());
+      TrieLogHelper.importTrieLog(
+          context.rootWorldStateStorage(), dataDirectoryPath, Hash.fromHexString(trieLogHash));
+    }
+  }
+
   record TrieLogContext(
       DataStorageConfiguration config,
       BonsaiWorldStateKeyValueStorage rootWorldStateStorage,

besu/src/test/java/org/hyperledger/besu/cli/subcommands/storage/TrieLogHelperTest.java

@@ -40,7 +40,6 @@
 
 import org.apache.tuweni.bytes.Bytes;
 import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -65,8 +64,8 @@ class TrieLogHelperTest {
   static BlockHeader blockHeader4;
   static BlockHeader blockHeader5;
 
-  @BeforeAll
-  public static void setup() throws IOException {
+  @BeforeEach
+  public void setup() throws IOException {
 
     blockHeader1 = new BlockHeaderTestFixture().number(1).buildHeader();
     blockHeader2 = new BlockHeaderTestFixture().number(2).buildHeader();
@@ -94,10 +93,11 @@ public static void setup() throws IOException {
         .getTrieLogStorageTransaction()
         .put(blockHeader5.getHash().toArrayUnsafe(), Bytes.fromHexString("0x05").toArrayUnsafe());
     updater.getTrieLogStorageTransaction().commit();
+
+    createDirectory();
   }
 
-  @BeforeEach
-  void createDirectory() throws IOException {
+  static void createDirectory() throws IOException {
     Files.createDirectories(dataDir.resolve("database"));
   }
 
@@ -227,7 +227,6 @@ public void cantPruneIfUserRequiredFurtherThanFinalized() {
 
   @Test
   public void exceptionWhileSavingFileStopsPruneProcess() throws IOException {
-    Files.delete(dataDir.resolve("database"));
 
     DataStorageConfiguration dataStorageConfiguration =
         ImmutableDataStorageConfiguration.builder()
@@ -243,7 +242,11 @@ public void exceptionWhileSavingFileStopsPruneProcess() throws IOException {
     assertThrows(
         RuntimeException.class,
         () ->
-            TrieLogHelper.prune(dataStorageConfiguration, inMemoryWorldState, blockchain, dataDir));
+            TrieLogHelper.prune(
+                dataStorageConfiguration,
+                inMemoryWorldState,
+                blockchain,
+                dataDir.resolve("unknownPath")));
 
     // assert all trie logs are still in the DB
     assertArrayEquals(
@@ -262,4 +265,43 @@ public void exceptionWhileSavingFileStopsPruneProcess() throws IOException {
         inMemoryWorldState.getTrieLog(blockHeader5.getHash()).get(),
         Bytes.fromHexString("0x05").toArrayUnsafe());
   }
+
+  @Test
+  public void exportedTrieMatchesDbTrieLog() throws IOException {
+    TrieLogHelper.exportTrieLog(inMemoryWorldState, dataDir, blockHeader1.getHash());
+    Path trieLogFile = dataDir.resolve("database").resolve(blockHeader1.getHash().toString());
+
+    var trieLog =
+        TrieLogHelper.readTrieLogsFromFile(trieLogFile.toString()).entrySet().stream()
+            .findFirst()
+            .get();
+
+    assertArrayEquals(trieLog.getKey(), blockHeader1.getHash().toArrayUnsafe());
+    assertArrayEquals(trieLog.getValue(), Bytes.fromHexString("0x01").toArrayUnsafe());
+
+    Files.delete(trieLogFile);
+  }
+
+  @Test
+  public void importedTrieLogMatchesDbTrieLog() throws IOException {
+    StorageProvider tempStorageProvider = new InMemoryKeyValueStorageProvider();
+    BonsaiWorldStateKeyValueStorage inMemoryWorldState2 =
+        new BonsaiWorldStateKeyValueStorage(tempStorageProvider, new NoOpMetricsSystem());
+
+    TrieLogHelper.exportTrieLog(inMemoryWorldState, dataDir, blockHeader1.getHash());
+    Path trieLogFile = dataDir.resolve("database").resolve(blockHeader1.getHash().toString());
+
+    var trieLog = TrieLogHelper.readTrieLogsFromFile(trieLogFile.toString());
+    var updater = inMemoryWorldState2.updater();
+
+    trieLog.forEach((k, v) -> updater.getTrieLogStorageTransaction().put(k, v));
+
+    updater.getTrieLogStorageTransaction().commit();
+
+    assertArrayEquals(
+        inMemoryWorldState2.getTrieLog(blockHeader1.getHash()).get(),
+        Bytes.fromHexString("0x01").toArrayUnsafe());
+
+    Files.delete(trieLogFile);
+  }
 }