Fix bonsai trie mining (#2101)

Signed-off-by: Karim TAAM <karim.t2am@gmail.com>

besu/src/test/java/org/hyperledger/besu/services/BesuEventsImplTest.java

@@ -16,6 +16,7 @@ package org.hyperledger.besu.services;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -31,10 +32,10 @@ import org.hyperledger.besu.ethereum.core.BlockBody;
 import org.hyperledger.besu.ethereum.core.BlockDataGenerator;
 import org.hyperledger.besu.ethereum.core.BlockHeaderTestFixture;
 import org.hyperledger.besu.ethereum.core.Difficulty;
+import org.hyperledger.besu.ethereum.core.MutableWorldState;
 import org.hyperledger.besu.ethereum.core.TransactionReceipt;
 import org.hyperledger.besu.ethereum.core.TransactionTestFixture;
 import org.hyperledger.besu.ethereum.core.Wei;
-import org.hyperledger.besu.ethereum.core.WorldState;
 import org.hyperledger.besu.ethereum.eth.manager.EthContext;
 import org.hyperledger.besu.ethereum.eth.manager.EthMessages;
 import org.hyperledger.besu.ethereum.eth.manager.EthPeer;
@@ -98,7 +99,7 @@ public class BesuEventsImplTest {
   @Mock private MainnetTransactionValidator mockTransactionValidator;
   @Mock private ProtocolSpec mockProtocolSpec;
   @Mock private WorldStateArchive mockWorldStateArchive;
-  @Mock private WorldState mockWorldState;
+  @Mock private MutableWorldState mockWorldState;
   private TransactionPool transactionPool;
   private BlockBroadcaster blockBroadcaster;
   private BesuEventsImpl serviceImpl;
@@ -131,7 +132,8 @@ public class BesuEventsImplTest {
         .thenReturn(ValidationResult.valid());
     when(mockTransactionValidator.validateForSender(any(), any(), any()))
         .thenReturn(ValidationResult.valid());
-    when(mockWorldStateArchive.get(any(), any())).thenReturn(Optional.of(mockWorldState));
+    when(mockWorldStateArchive.getMutable(any(), any(), anyBoolean()))
+        .thenReturn(Optional.of(mockWorldState));
 
     blockBroadcaster = new BlockBroadcaster(mockEthContext);
     syncState = new SyncState(blockchain, mockEthPeers);

ethereum/api/src/main/java/org/hyperledger/besu/ethereum/api/query/BlockchainQueries.java

@@ -802,7 +802,7 @@ public class BlockchainQueries {
     final Optional<BlockHeader> header = blockchain.getBlockHeader(blockHash);
     return header.flatMap(
         blockHeader ->
-            worldStateArchive.getMutable(blockHeader.getStateRoot(), blockHeader.getHash()));
+            worldStateArchive.getMutable(blockHeader.getStateRoot(), blockHeader.getHash(), false));
   }
 
   public Optional<Long> gasPrice() {

ethereum/blockcreation/src/main/java/org/hyperledger/besu/ethereum/blockcreation/AbstractBlockCreator.java

@@ -240,11 +240,11 @@ public abstract class AbstractBlockCreator implements AsyncBlockCreator {
     final MutableWorldState worldState =
         protocolContext
             .getWorldStateArchive()
-            .getMutable(parentStateRoot, parentHeader.getHash())
+            .getMutable(parentStateRoot, parentHeader.getHash(), false)
             .orElseThrow(
                 () -> {
                   LOG.info("Unable to create block because world state is not available");
-                  return new IllegalStateException(
+                  return new CancellationException(
                       "World state not available for block "
                           + parentHeader.getNumber()
                           + " with state root "

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiInMemoryWorldState.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.hyperledger.besu.ethereum.bonsai;
+
+import org.hyperledger.besu.ethereum.core.BlockHeader;
+import org.hyperledger.besu.ethereum.core.Hash;
+
+public class BonsaiInMemoryWorldState extends BonsaiPersistedWorldState {
+
+  public BonsaiInMemoryWorldState(
+      final BonsaiWorldStateArchive archive,
+      final BonsaiWorldStateKeyValueStorage worldStateStorage) {
+    super(archive, worldStateStorage);
+  }
+
+  @Override
+  public Hash rootHash() {
+    final BonsaiWorldStateKeyValueStorage.Updater updater = worldStateStorage.updater();
+    final Hash calculatedRootHash = calculateRootHash(updater);
+    updater.rollback();
+    return Hash.wrap(calculatedRootHash);
+  }
+
+  @Override
+  public void persist(final BlockHeader blockHeader) {
+    throw new UnsupportedOperationException("In Memory worldState can not be persisted.");
+  }
+}

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiInMemoryWorldStateKeyValueStorage.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.hyperledger.besu.ethereum.bonsai;
+
+import org.hyperledger.besu.ethereum.worldstate.WorldStateStorage;
+import org.hyperledger.besu.plugin.services.storage.KeyValueStorage;
+import org.hyperledger.besu.plugin.services.storage.KeyValueStorageTransaction;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+public class BonsaiInMemoryWorldStateKeyValueStorage extends BonsaiWorldStateKeyValueStorage
+    implements WorldStateStorage {
+
+  private static final Logger LOG = LogManager.getLogger();
+
+  public BonsaiInMemoryWorldStateKeyValueStorage(
+      final KeyValueStorage accountStorage,
+      final KeyValueStorage codeStorage,
+      final KeyValueStorage storageStorage,
+      final KeyValueStorage trieBranchStorage,
+      final KeyValueStorage trieLogStorage) {
+    super(accountStorage, codeStorage, storageStorage, trieBranchStorage, trieLogStorage);
+  }
+
+  @Override
+  public InMemoryUpdater updater() {
+    return new InMemoryUpdater(
+        accountStorage.startTransaction(),
+        codeStorage.startTransaction(),
+        storageStorage.startTransaction(),
+        trieBranchStorage.startTransaction(),
+        trieLogStorage.startTransaction());
+  }
+
+  public static class InMemoryUpdater extends BonsaiWorldStateKeyValueStorage.Updater
+      implements WorldStateStorage.Updater {
+
+    public InMemoryUpdater(
+        final KeyValueStorageTransaction accountStorageTransaction,
+        final KeyValueStorageTransaction codeStorageTransaction,
+        final KeyValueStorageTransaction storageStorageTransaction,
+        final KeyValueStorageTransaction trieBranchStorageTransaction,
+        final KeyValueStorageTransaction trieLogStorageTransaction) {
+      super(
+          accountStorageTransaction,
+          codeStorageTransaction,
+          storageStorageTransaction,
+          trieBranchStorageTransaction,
+          trieLogStorageTransaction);
+    }
+
+    @Override
+    public void commit() {
+      LOG.trace("Cannot commit using an in memory key value storage");
+    }
+  }
+}

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiLayeredWorldState.java

@@ -258,8 +258,16 @@ public class BonsaiLayeredWorldState implements MutableWorldState, BonsaiWorldVi
 
   @Override
   public MutableWorldState copy() {
-    throw new UnsupportedOperationException(
+    final BonsaiPersistedWorldState bonsaiPersistedWorldState =
-        "Bonsai Tries does not support direct duplication of the persisted tries.");
+        ((BonsaiPersistedWorldState) archive.getMutable());
+    return new BonsaiInMemoryWorldState(
+        archive,
+        new BonsaiInMemoryWorldStateKeyValueStorage(
+            bonsaiPersistedWorldState.getWorldStateStorage().accountStorage,
+            bonsaiPersistedWorldState.getWorldStateStorage().codeStorage,
+            bonsaiPersistedWorldState.getWorldStateStorage().storageStorage,
+            bonsaiPersistedWorldState.getWorldStateStorage().trieBranchStorage,
+            bonsaiPersistedWorldState.getWorldStateStorage().trieLogStorage));
   }
 
   @Override

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiPersistedWorldState.java

@@ -48,7 +48,7 @@ public class BonsaiPersistedWorldState implements MutableWorldState, BonsaiWorld
 
   private static final Logger LOG = LogManager.getLogger();
 
-  private final BonsaiWorldStateKeyValueStorage worldStateStorage;
+  protected final BonsaiWorldStateKeyValueStorage worldStateStorage;
 
   private final BonsaiWorldStateArchive archive;
   private final BonsaiWorldStateUpdater updater;
@@ -75,8 +75,14 @@ public class BonsaiPersistedWorldState implements MutableWorldState, BonsaiWorld
 
   @Override
   public MutableWorldState copy() {
-    throw new UnsupportedOperationException(
+    return new BonsaiInMemoryWorldState(
-        "Bonsai Tries does not support direct duplication of the persisted tries.");
+        archive,
+        new BonsaiInMemoryWorldStateKeyValueStorage(
+            worldStateStorage.accountStorage,
+            worldStateStorage.codeStorage,
+            worldStateStorage.storageStorage,
+            worldStateStorage.trieBranchStorage,
+            worldStateStorage.trieLogStorage));
   }
 
   @Override
@@ -89,7 +95,11 @@ public class BonsaiPersistedWorldState implements MutableWorldState, BonsaiWorld
     worldStateRootHash = blockHeader.getStateRoot();
   }
 
-  private Hash calculateRootHash(final BonsaiWorldStateKeyValueStorage.Updater stateUpdater) {
+  public BonsaiWorldStateKeyValueStorage getWorldStateStorage() {
+    return worldStateStorage;
+  }
+
+  protected Hash calculateRootHash(final BonsaiWorldStateKeyValueStorage.Updater stateUpdater) {
     // first clear storage
     for (final Address address : updater.getStorageToClear()) {
       // because we are clearing persisted values we need the account root as persisted

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateArchive.java

@@ -153,7 +153,7 @@ public class BonsaiWorldStateArchive implements WorldStateArchive {
                   this,
                   Optional.empty(),
                   header.getNumber(),
-                  blockHash,
+                  header.getStateRoot(),
                   trieLogLayer.get()));
         }
       }

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateKeyValueStorage.java

@@ -36,11 +36,11 @@ public class BonsaiWorldStateKeyValueStorage implements WorldStateStorage {
   public static final byte[] WORLD_BLOCK_HASH_KEY =
       "worldBlockHash".getBytes(StandardCharsets.UTF_8);
 
-  private final KeyValueStorage accountStorage;
+  protected final KeyValueStorage accountStorage;
-  private final KeyValueStorage codeStorage;
+  protected final KeyValueStorage codeStorage;
-  private final KeyValueStorage storageStorage;
+  protected final KeyValueStorage storageStorage;
-  private final KeyValueStorage trieBranchStorage;
+  protected final KeyValueStorage trieBranchStorage;
-  private final KeyValueStorage trieLogStorage;
+  protected final KeyValueStorage trieLogStorage;
 
   public BonsaiWorldStateKeyValueStorage(final StorageProvider provider) {
     accountStorage =

ethereum/eth/src/main/java/org/hyperledger/besu/ethereum/eth/transactions/TransactionPool.java

@@ -251,7 +251,7 @@ public class TransactionPool implements BlockAddedObserver {
 
     return protocolContext
         .getWorldStateArchive()
-        .get(chainHeadBlockHeader.getStateRoot(), chainHeadBlockHeader.getHash())
+        .getMutable(chainHeadBlockHeader.getStateRoot(), chainHeadBlockHeader.getHash(), false)
         .map(
             worldState -> {
               final Account senderAccount = worldState.get(transaction.getSender());