backup-state and restore-state subcommands to use bonsai

CHANGELOG.md

@@ -45,6 +45,7 @@
 ### Bug fixes
 - Fix callTracer handling of failed CREATE operations, including correct input field extraction and proper error reporting for both soft failures and revert reasons
 - Upgrade netty to 4.2.10-Final - Fixes `setsockopt() failed: Protocol not available` [#9783](https://github.com/hyperledger/besu/pull/9783)
+- Backup and Restore State Subcommands now support bonsai database format [#9812](https://github.com/hyperledger/besu/pull/9812) 
 
 ## 26.1.0-RC1
 

app/src/main/java/org/hyperledger/besu/cli/subcommands/operator/BackupState.java

@@ -16,15 +16,17 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static org.hyperledger.besu.cli.DefaultCommandValues.MANDATORY_LONG_FORMAT_HELP;
+import static org.hyperledger.besu.cli.DefaultCommandValues.MANDATORY_PATH_FORMAT_HELP;
 
 import org.hyperledger.besu.cli.util.VersionProvider;
 import org.hyperledger.besu.controller.BesuController;
 import org.hyperledger.besu.ethereum.api.query.StateBackupService;
 import org.hyperledger.besu.ethereum.api.query.StateBackupService.BackupStatus;
 import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
 import org.hyperledger.besu.ethereum.eth.manager.EthScheduler;
-import org.hyperledger.besu.ethereum.trie.forest.ForestWorldStateArchive;
-import org.hyperledger.besu.ethereum.trie.forest.storage.ForestWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.trie.pathbased.bonsai.BonsaiWorldStateProvider;
+import org.hyperledger.besu.ethereum.trie.pathbased.bonsai.storage.BonsaiWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
 import org.hyperledger.besu.metrics.noop.NoOpMetricsSystem;
 import org.hyperledger.besu.util.BesuVersionUtils;
 
@@ -60,7 +62,7 @@ public BackupState() {}
   @Option(
       names = "--backup-path",
       required = true,
-      paramLabel = MANDATORY_LONG_FORMAT_HELP,
+      paramLabel = MANDATORY_PATH_FORMAT_HELP,
       description = "The path to store the backup files.",
       arity = "1..1")
   private final File backupDir = null;
@@ -84,9 +86,14 @@ public void run() {
 
     final BesuController besuController = createBesuController();
     final MutableBlockchain blockchain = besuController.getProtocolContext().getBlockchain();
-    final ForestWorldStateKeyValueStorage forestWorldStateKeyValueStorage =
-        ((ForestWorldStateArchive) besuController.getProtocolContext().getWorldStateArchive())
-            .getWorldStateStorage();
+    final WorldStateArchive worldStateArchive =
+        besuController.getProtocolContext().getWorldStateArchive();
+    checkArgument(
+        worldStateArchive instanceof BonsaiWorldStateProvider,
+        "x-backup-state only supports Bonsai world state format.");
+    final BonsaiWorldStateKeyValueStorage bonsaiWorldStateKeyValueStorage =
+        (BonsaiWorldStateKeyValueStorage)
+            ((BonsaiWorldStateProvider) worldStateArchive).getWorldStateKeyValueStorage();
     final EthScheduler scheduler = new EthScheduler(1, 1, 1, 1, new NoOpMetricsSystem());
     try {
       final long targetBlock = Math.min(blockchain.getChainHeadBlockNumber(), this.block);
@@ -96,7 +103,7 @@ public void run() {
               blockchain,
               backupDir.toPath(),
               scheduler,
-              forestWorldStateKeyValueStorage);
+              bonsaiWorldStateKeyValueStorage);
       final BackupStatus status = backup.requestBackup(targetBlock, compress, Optional.empty());
 
       final double refValue = Math.pow(2, 256) / 100.0d;

app/src/main/java/org/hyperledger/besu/cli/subcommands/operator/RestoreState.java

@@ -15,7 +15,6 @@
 package org.hyperledger.besu.cli.subcommands.operator;
 
 import static org.hyperledger.besu.cli.DefaultCommandValues.MANDATORY_LONG_FORMAT_HELP;
-import static org.hyperledger.besu.ethereum.trie.CompactEncoding.bytesToPath;
 
 import org.hyperledger.besu.cli.util.VersionProvider;
 import org.hyperledger.besu.config.JsonUtil;
@@ -32,12 +31,10 @@
 import org.hyperledger.besu.ethereum.mainnet.MainnetBlockHeaderFunctions;
 import org.hyperledger.besu.ethereum.rlp.BytesValueRLPInput;
 import org.hyperledger.besu.ethereum.rlp.RLPInput;
-import org.hyperledger.besu.ethereum.trie.Node;
-import org.hyperledger.besu.ethereum.trie.PersistVisitor;
-import org.hyperledger.besu.ethereum.trie.RestoreVisitor;
 import org.hyperledger.besu.ethereum.trie.common.PmtStateTrieAccountValue;
-import org.hyperledger.besu.ethereum.trie.forest.ForestWorldStateArchive;
-import org.hyperledger.besu.ethereum.trie.forest.storage.ForestWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.trie.pathbased.bonsai.BonsaiWorldStateProvider;
+import org.hyperledger.besu.ethereum.trie.pathbased.bonsai.storage.BonsaiWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
 import org.hyperledger.besu.util.io.RollingFileReader;
 
 import java.io.IOException;
@@ -77,14 +74,14 @@ public class RestoreState implements Runnable {
 
   @ParentCommand private OperatorSubCommand parentCommand;
 
-  private static final int TRIE_NODE_COMMIT_BATCH_SIZE = 100;
+  private static final int COMMIT_BATCH_SIZE = 10_000;
 
   private long targetBlock;
   private long accountCount;
-  private long trieNodeCount;
   private boolean compressed;
   private BesuController besuController;
-  private ForestWorldStateKeyValueStorage.Updater updater;
+  private BonsaiWorldStateKeyValueStorage.Updater updater;
+  private long entryCount;
 
   /** Default Constructor. */
   RestoreState() {}
@@ -165,16 +162,8 @@ private void restoreBlocks() throws IOException {
     LOG.info("Chain data loaded");
   }
 
-  @SuppressWarnings("UnusedVariable")
   private void restoreAccounts() throws IOException {
     newWorldStateUpdater();
-    int storageBranchCount = 0;
-    int storageExtensionCount = 0;
-    int storageLeafCount = 0;
-
-    final PersistVisitor<Bytes> accountPersistVisitor =
-        new PersistVisitor<>(this::updateAccountState);
-    Node<Bytes> root = accountPersistVisitor.initialRoot();
 
     try (final RollingFileReader reader =
         new RollingFileReader(this::accountFileName, compressed)) {
@@ -189,7 +178,7 @@ private void restoreAccounts() throws IOException {
         if (length != 3) {
           throw new RuntimeException("Unexpected account length " + length);
         }
-        final Bytes32 trieKey = accountInput.readBytes32();
+        final Bytes32 accountHash = accountInput.readBytes32();
         final Bytes accountRlp = accountInput.readBytes();
         final Bytes code = accountInput.readBytes();
 
@@ -198,19 +187,20 @@ private void restoreAccounts() throws IOException {
         if (!trieAccount.getCodeHash().equals(Hash.hash(code))) {
           throw new RuntimeException("Code hash doesn't match");
         }
-        if (code.size() > 0) {
-          updateCode(code);
-        }
-
-        final RestoreVisitor<Bytes> accountTrieWriteVisitor =
-            new RestoreVisitor<>(t -> t, accountRlp, accountPersistVisitor);
 
-        root = root.accept(accountTrieWriteVisitor, bytesToPath(trieKey));
+        // write account to flat DB
+        updater.putAccountInfoState(Hash.wrap(accountHash), accountRlp);
+        entryCount++;
+        maybeCommitUpdater();
 
-        final PersistVisitor<Bytes> storagePersistVisitor =
-            new PersistVisitor<>(this::updateAccountStorage);
-        Node<Bytes> storageRoot = storagePersistVisitor.initialRoot();
+        // write code
+        if (code.size() > 0) {
+          updater.putCode(Hash.wrap(accountHash), trieAccount.getCodeHash(), code);
+          entryCount++;
+          maybeCommitUpdater();
+        }
 
+        // read and write storage entries
         while (true) {
           final byte[] trieEntry = reader.readBytes();
           final BytesValueRLPInput trieInput =
@@ -222,68 +212,58 @@ private void restoreAccounts() throws IOException {
           if (len != 2) {
             throw new RuntimeException("Unexpected storage trie entry length " + len);
           }
-          final Bytes32 storageTrieKey = Bytes32.wrap(trieInput.readBytes());
-          final Bytes storageTrieValue = Bytes.wrap(trieInput.readBytes());
-          final RestoreVisitor<Bytes> storageTrieWriteVisitor =
-              new RestoreVisitor<>(t -> t, storageTrieValue, storagePersistVisitor);
-          storageRoot = storageRoot.accept(storageTrieWriteVisitor, bytesToPath(storageTrieKey));
+          final Bytes32 slotHash = Bytes32.wrap(trieInput.readBytes());
+          final Bytes storageValue = Bytes.wrap(trieInput.readBytes());
+
+          updater.putStorageValueBySlotHash(
+              Hash.wrap(accountHash), Hash.wrap(slotHash), storageValue);
+          entryCount++;
+          maybeCommitUpdater();
 
           trieInput.leaveList();
         }
-        storagePersistVisitor.persist(storageRoot);
-        storageBranchCount += storagePersistVisitor.getBranchNodeCount();
-        storageExtensionCount += storagePersistVisitor.getExtensionNodeCount();
-        storageLeafCount += storagePersistVisitor.getLeafNodeCount();
 
         accountInput.leaveList();
       }
     }
-    accountPersistVisitor.persist(root);
+
+    // save world state root from target block header
+    final MutableBlockchain blockchain = besuController.getProtocolContext().getBlockchain();
+    final BlockHeader targetHeader =
+        blockchain
+            .getBlockHeader(targetBlock)
+            .orElseThrow(() -> new RuntimeException("Target block header not found"));
+    updater.saveWorldState(
+        targetHeader.getBlockHash().getBytes(),
+        Bytes32.wrap(targetHeader.getStateRoot().getBytes()),
+        Bytes.EMPTY);
+
     updater.commit();
-    LOG.info("Account BranchNodes: {} ", accountPersistVisitor.getBranchNodeCount());
-    LOG.info("Account ExtensionNodes: {} ", accountPersistVisitor.getExtensionNodeCount());
-    LOG.info("Account LeafNodes: {} ", accountPersistVisitor.getLeafNodeCount());
-    LOG.info("Storage BranchNodes: {} ", storageBranchCount);
-    LOG.info("Storage LeafNodes: {} ", storageExtensionCount);
-    LOG.info("Storage ExtensionNodes: {} ", storageLeafCount);
+    LOG.info("Flat DB entries written: {}", entryCount);
     LOG.info("Account data loaded");
   }
 
   private void newWorldStateUpdater() {
     if (updater != null) {
       updater.commit();
     }
-    final ForestWorldStateKeyValueStorage worldStateKeyValueStorage =
-        ((ForestWorldStateArchive) besuController.getProtocolContext().getWorldStateArchive())
-            .getWorldStateStorage();
+    final WorldStateArchive worldStateArchive =
+        besuController.getProtocolContext().getWorldStateArchive();
+    if (!(worldStateArchive instanceof BonsaiWorldStateProvider)) {
+      throw new IllegalStateException("x-restore-state only supports Bonsai world state format.");
+    }
+    final BonsaiWorldStateKeyValueStorage worldStateKeyValueStorage =
+        (BonsaiWorldStateKeyValueStorage)
+            ((BonsaiWorldStateProvider) worldStateArchive).getWorldStateKeyValueStorage();
     updater = worldStateKeyValueStorage.updater();
   }
 
   private void maybeCommitUpdater() {
-    if (trieNodeCount % TRIE_NODE_COMMIT_BATCH_SIZE == 0) {
+    if (entryCount % COMMIT_BATCH_SIZE == 0) {
       newWorldStateUpdater();
     }
   }
 
-  private void updateCode(final Bytes code) {
-    maybeCommitUpdater();
-    updater.putCode(code);
-  }
-
-  private void updateAccountState(final Bytes32 key, final Bytes value) {
-    maybeCommitUpdater();
-    // restore by path not supported
-    updater.putAccountStateTrieNode(key, value);
-    trieNodeCount++;
-  }
-
-  private void updateAccountStorage(final Bytes32 key, final Bytes value) {
-    maybeCommitUpdater();
-    // restore by path not supported
-    updater.putAccountStorageTrieNode(key, value);
-    trieNodeCount++;
-  }
-
   private BesuController createBesuController() {
     return parentCommand.parentCommand.buildController();
   }