diff --git a/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommand.java b/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommand.java index 5d1e1d90c67..3fb0b2a9f27 100644 --- a/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommand.java +++ b/besu/src/main/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommand.java @@ -38,6 +38,7 @@ import org.hyperledger.besu.plugin.services.storage.SegmentedKeyValueStorage; import org.hyperledger.besu.plugin.services.storage.SegmentedKeyValueStorageTransaction; +import java.util.function.Consumer; import java.util.function.Function; import org.apache.tuweni.bytes.Bytes; @@ -133,10 +134,9 @@ Hash rebuildTrie(final BonsaiWorldStateKeyValueStorage worldStateStorage) { // rebuild the trie by inserting everything into a StoredMerklePatriciaTrie // and incrementally (naively) commit after each account while streaming - // TODO: optimize to incrementally commit tx after a certain threshold final var wss = worldStateStorage.getComposedWorldStateStorage(); - final var accountTrie = + var accountTrie = new StoredMerklePatriciaTrie<>( new StoredNodeFactory<>( // this may be inefficient, and we can read through an incrementally committing tx @@ -147,9 +147,9 @@ Hash rebuildTrie(final BonsaiWorldStateKeyValueStorage worldStateStorage) { MerkleTrie.EMPTY_TRIE_NODE_HASH); final var accountsTx = new WrappedTransaction(worldStateStorage.getComposedWorldStateStorage()); - final Runnable accountTrieCommit = - () -> - accountTrie.commit( + final Consumer> accountTrieCommit = + (trie) -> + trie.commit( (loc, hash, value) -> accountsTx.put( TRIE_BRANCH_STORAGE, loc.toArrayUnsafe(), value.toArrayUnsafe())); @@ -181,13 +181,25 @@ Hash rebuildTrie(final BonsaiWorldStateKeyValueStorage worldStateStorage) { LOG.info("committing account trie at account {}", accountPair.getFirst()); // commit the account trie if we have exceeded the forced commit interval - accountTrieCommit.run(); + accountTrieCommit.accept(accountTrie); accountsTx.commitAndReopen(); + + // new trie with new root, GC trie nodes + accountTrie = + new StoredMerklePatriciaTrie<>( + new StoredNodeFactory<>( + // this may be inefficient, and we can read through an incrementally committing + // tx + // instead + worldStateStorage::getAccountStateTrieNode, + Function.identity(), + Function.identity()), + accountTrie.getRootHash()); } } // final commit - accountTrieCommit.run(); + accountTrieCommit.accept(accountTrie); accountsTx.commit(); // return the new state trie root hash @@ -209,19 +221,22 @@ Hash rebuildAccountTrie( var accountStorageTx = new WrappedTransaction(worldStateStorage.getComposedWorldStateStorage()); var wss = worldStateStorage.getComposedWorldStateStorage(); + Function> newAccountStorageTrie = + (rootHash) -> + new StoredMerklePatriciaTrie<>( + new StoredNodeFactory<>( + (location, hash) -> + worldStateStorage.getAccountStorageTrieNode(accountHash, location, hash), + Function.identity(), + Function.identity()), + rootHash); + // create account storage trie - var accountStorageTrie = - new StoredMerklePatriciaTrie<>( - new StoredNodeFactory<>( - (location, hash) -> - worldStateStorage.getAccountStorageTrieNode(accountHash, location, hash), - Function.identity(), - Function.identity()), - MerkleTrie.EMPTY_TRIE_NODE_HASH); + var accountStorageTrie = newAccountStorageTrie.apply(MerkleTrie.EMPTY_TRIE_NODE_HASH); - Runnable accountStorageCommit = - () -> - accountStorageTrie.commit( + Consumer> accountStorageCommit = + (trie) -> + trie.commit( (location, hash, value) -> accountStorageTx.put( TRIE_BRANCH_STORAGE, @@ -241,11 +256,13 @@ Hash rebuildAccountTrie( // commit the account storage trie if (accountStorageCount++ % FORCED_COMMIT_INTERVAL == 0) { - accountStorageCommit.run(); + accountStorageCommit.accept(accountStorageTrie); accountStorageTx.commitAndReopen(); + // new trie with new root, GC trie nodes + accountStorageTrie = newAccountStorageTrie.apply(accountStorageTrie.getRootHash()); } } - accountStorageCommit.run(); + accountStorageCommit.accept(accountStorageTrie); accountStorageTx.commit(); return Hash.wrap(accountStorageTrie.getRootHash()); diff --git a/besu/src/test/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommandTest.java b/besu/src/test/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommandTest.java index c844884b7fc..3a74e249f1e 100644 --- a/besu/src/test/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommandTest.java +++ b/besu/src/test/java/org/hyperledger/besu/cli/subcommands/storage/RebuildBonsaiStateTrieSubCommandTest.java @@ -55,9 +55,8 @@ void setupBonsaiBlockchain() { blockchainSetupUtil = BlockchainSetupUtil.createForEthashChain( // Mainnet is a more robust test resource, but it takes upwards of 1 minute to generate - //BlockTestUtil.getMainnetResources(), - BlockTestUtil.getSnapTestChainResources(), - DataStorageFormat.BONSAI); + // BlockTestUtil.getMainnetResources(), + BlockTestUtil.getSnapTestChainResources(), DataStorageFormat.BONSAI); blockchainSetupUtil.importAllBlocks(HeaderValidationMode.NONE, HeaderValidationMode.NONE); }