Import export trie log

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

@@ -62,23 +62,28 @@ public class DataStorageOptions implements CLIOptions<DataStorageConfiguration>
   private final DataStorageOptions.Unstable unstableOptions = new Unstable();
 
   static class Unstable {
+    private static final String BONSAI_LIMIT_TRIE_LOGS_ENABLED =
+        "--Xbonsai-limit-trie-logs-enabled";
+    private static final String BONSAI_TRIE_LOGS_RETENTION_THRESHOLD =
+        "--Xbonsai-trie-logs-retention-threshold";
+    private static final String BONSAI_TRIE_LOG_PRUNING_LIMIT = "--Xbonsai-trie-logs-pruning-limit";
 
     @CommandLine.Option(
         hidden = true,
-        names = {"--Xbonsai-trie-log-pruning-enabled"},
+        names = {BONSAI_LIMIT_TRIE_LOGS_ENABLED},
         description = "Enable trie log pruning. (default: ${DEFAULT-VALUE})")
     private boolean bonsaiTrieLogPruningEnabled = DEFAULT_BONSAI_TRIE_LOG_PRUNING_ENABLED;
 
     @CommandLine.Option(
         hidden = true,
-        names = {"--Xbonsai-trie-log-retention-threshold"},
+        names = {BONSAI_TRIE_LOGS_RETENTION_THRESHOLD},
         description =
             "The number of blocks for which to retain trie logs. (default: ${DEFAULT-VALUE})")
     private long bonsaiTrieLogRetentionThreshold = DEFAULT_BONSAI_TRIE_LOG_RETENTION_THRESHOLD;
 
     @CommandLine.Option(
         hidden = true,
-        names = {"--Xbonsai-trie-log-pruning-limit"},
+        names = {BONSAI_TRIE_LOG_PRUNING_LIMIT},
         description =
             "The max number of blocks to load and prune trie logs for at startup. (default: ${DEFAULT-VALUE})")
     private int bonsaiTrieLogPruningLimit = DEFAULT_BONSAI_TRIE_LOG_PRUNING_LIMIT;

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

@@ -22,7 +22,11 @@
 import org.hyperledger.besu.ethereum.chain.Blockchain;
 import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
 import org.hyperledger.besu.ethereum.core.BlockHeader;
+import org.hyperledger.besu.ethereum.rlp.BytesValueRLPInput;
+import org.hyperledger.besu.ethereum.rlp.RLP;
 import org.hyperledger.besu.ethereum.trie.bonsai.storage.BonsaiWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.trie.bonsai.trielog.TrieLogFactoryImpl;
+import org.hyperledger.besu.ethereum.trie.bonsai.trielog.TrieLogLayer;
 import org.hyperledger.besu.ethereum.worldstate.DataStorageConfiguration;
 
 import java.io.File;
@@ -32,13 +36,15 @@
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.PrintWriter;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -97,16 +103,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 +123,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 +212,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 +266,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,24 +285,67 @@ 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();
-
     } catch (IOException | ClassNotFoundException e) {
-
       LOG.error(e.getMessage());
       throw new RuntimeException(e);
     }
 
     return trieLogs;
   }
 
+  private static void saveTrieLogsAsRlpInFile(
+      final List<Hash> trieLogsKeys,
+      final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
+      final String batchFileName) {
+    File file = new File(batchFileName);
+    if (file.exists()) {
+      LOG.error("File already exists, skipping file creation");
+      return;
+    }
+
+    final IdentityHashMap<byte[], byte[]> trieLogs =
+        getTrieLogs(trieLogsKeys, rootWorldStateStorage);
+    final Bytes rlp =
+        RLP.encode(
+            o ->
+                o.writeList(
+                    trieLogs.entrySet(), (val, out) -> out.writeRaw(Bytes.wrap(val.getValue()))));
+    try {
+      Files.write(file.toPath(), rlp.toArrayUnsafe());
+    } catch (IOException e) {
+      LOG.error(e.getMessage());
+      throw new RuntimeException(e);
+    }
+  }
+
+  static IdentityHashMap<byte[], byte[]> readTrieLogsAsRlpFromFile(final String batchFileName) {
+    try {
+      final Bytes file = Bytes.wrap(Files.readAllBytes(Path.of(batchFileName)));
+      final BytesValueRLPInput input = new BytesValueRLPInput(file, false);
+
+      input.enterList();
+      final IdentityHashMap<byte[], byte[]> trieLogs = new IdentityHashMap<>();
+      while (!input.isEndOfCurrentList()) {
+        final Bytes trieLogBytes = input.currentListAsBytes();
+        TrieLogLayer trieLogLayer =
+            TrieLogFactoryImpl.readFrom(new BytesValueRLPInput(Bytes.wrap(trieLogBytes), false));
+        trieLogs.put(trieLogLayer.getBlockHash().toArrayUnsafe(), trieLogBytes.toArrayUnsafe());
+      }
+      input.leaveList();
+
+      return trieLogs;
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   private static IdentityHashMap<byte[], byte[]> getTrieLogs(
       final List<Hash> trieLogKeys, final BonsaiWorldStateKeyValueStorage rootWorldStateStorage) {
     IdentityHashMap<byte[], byte[]> trieLogsToRetain = new IdentityHashMap<>();
@@ -357,5 +400,25 @@ 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 trieLogFilePath) {
+
+    var trieLog = readTrieLogsAsRlpFromFile(trieLogFilePath.toString());
+
+    var updater = rootWorldStateStorage.updater();
+    trieLog.forEach((key, value) -> updater.getTrieLogStorageTransaction().put(key, value));
+    updater.getTrieLogStorageTransaction().commit();
+  }
+
+  static void exportTrieLog(
+      final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
+      final List<Hash> trieLogHash,
+      final Path directoryPath)
+      throws IOException {
+    final String trieLogFile = directoryPath.toString();
+
+    saveTrieLogsAsRlpInFile(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

@@ -19,16 +19,19 @@
 
 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;
+import java.util.List;
 
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.core.config.Configurator;
@@ -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,102 @@ public void run() {
     }
   }
 
+  @Command(
+      name = "export",
+      description = "This command exports the trie log of a determined block to a binary file",
+      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-block-hash",
+        description =
+            "Comma separated list of hashes from the blocks you want to export the trie logs of",
+        split = " {0,1}, {0,1}",
+        arity = "1..*")
+    private List<String> trieLogBlockHashList;
+
+    @CommandLine.Option(
+        names = "--trie-log-file-path",
+        description = "The file you want to export the trie logs to",
+        arity = "1..1")
+    private Path trieLogFilePath = null;
+
+    @Override
+    public void run() {
+      if (trieLogFilePath == null) {
+        trieLogFilePath =
+            Paths.get(
+                TrieLogSubCommand.parentCommand
+                    .parentCommand
+                    .dataDir()
+                    .resolve("trie-logs.bin")
+                    .toAbsolutePath()
+                    .toString());
+      }
+
+      TrieLogContext context = getTrieLogContext();
+
+      final List<Hash> listOfBlockHashes =
+          trieLogBlockHashList.stream().map(Hash::fromHexString).toList();
+
+      try {
+        TrieLogHelper.exportTrieLog(
+            context.rootWorldStateStorage(), listOfBlockHashes, trieLogFilePath);
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
+
+  @Command(
+      name = "import",
+      description = "This command imports a trie log exported by another besu node",
+      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-file-path",
+        description = "The file you want to import the trie logs from",
+        arity = "1..1")
+    private Path trieLogFilePath = null;
+
+    @Override
+    public void run() {
+      if (trieLogFilePath == null) {
+        trieLogFilePath =
+            Paths.get(
+                TrieLogSubCommand.parentCommand
+                    .parentCommand
+                    .dataDir()
+                    .resolve("trie-logs.bin")
+                    .toAbsolutePath()
+                    .toString());
+      }
+
+      TrieLogContext context = getTrieLogContext();
+
+      TrieLogHelper.importTrieLog(context.rootWorldStateStorage(), trieLogFilePath);
+    }
+  }
+
   record TrieLogContext(
       DataStorageConfiguration config,
       BonsaiWorldStateKeyValueStorage rootWorldStateStorage,