Make healing an action of the WorldStateArchive

besu/src/main/java/org/hyperledger/besu/controller/BesuControllerBuilder.java

@@ -93,6 +93,7 @@
 import org.hyperledger.besu.ethereum.worldstate.DataStorageConfiguration;
 import org.hyperledger.besu.ethereum.worldstate.DiffBasedSubStorageConfiguration;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
+import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive.WorldStateHealer;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateKeyValueStorage;
 import org.hyperledger.besu.ethereum.worldstate.WorldStatePreimageStorage;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateStorageCoordinator;
@@ -113,6 +114,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
 import org.slf4j.Logger;
@@ -589,9 +591,14 @@ public BesuController build() {
             .map(BesuComponent::getCachedMerkleTrieLoader)
             .orElseGet(() -> new BonsaiCachedMerkleTrieLoader(metricsSystem));
 
+    final var worldStateHealerSupplier = new AtomicReference<WorldStateHealer>();
+
     final WorldStateArchive worldStateArchive =
         createWorldStateArchive(
-            worldStateStorageCoordinator, blockchain, bonsaiCachedMerkleTrieLoader);
+            worldStateStorageCoordinator,
+            blockchain,
+            bonsaiCachedMerkleTrieLoader,
+            worldStateHealerSupplier::get);
 
     if (maybeStoredGenesisBlockHash.isEmpty()) {
       genesisState.writeStateTo(worldStateArchive.getMutable());
@@ -713,6 +720,8 @@ public BesuController build() {
             ethProtocolManager,
             pivotBlockSelector);
 
+    worldStateHealerSupplier.set(synchronizer::healWorldState);
+
     ethPeers.setTrailingPeerRequirementsSupplier(synchronizer::calculateTrailingPeerRequirements);
 
     if (syncConfig.getSyncMode() == SyncMode.SNAP
@@ -1101,7 +1110,8 @@ private Optional<SnapProtocolManager> createSnapProtocolManager(
   WorldStateArchive createWorldStateArchive(
       final WorldStateStorageCoordinator worldStateStorageCoordinator,
       final Blockchain blockchain,
-      final BonsaiCachedMerkleTrieLoader bonsaiCachedMerkleTrieLoader) {
+      final BonsaiCachedMerkleTrieLoader bonsaiCachedMerkleTrieLoader,
+      final Supplier<WorldStateHealer> worldStateHealerSupplier) {
     return switch (dataStorageConfiguration.getDataStorageFormat()) {
       case BONSAI -> {
         final BonsaiWorldStateKeyValueStorage worldStateKeyValueStorage =
@@ -1116,7 +1126,8 @@ yield new BonsaiWorldStateProvider(
                     .getMaxLayersToLoad()),
             bonsaiCachedMerkleTrieLoader,
             besuComponent.map(BesuComponent::getBesuPluginContext).orElse(null),
-            evmConfiguration);
+            evmConfiguration,
+            worldStateHealerSupplier);
       }
       case FOREST -> {
         final WorldStatePreimageStorage preimageStorage =

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/MainnetBlockValidator.java

@@ -181,7 +181,7 @@ public BlockProcessingResult validateAndProcessBlock(
             Optional.of(new BlockProcessingOutputs(worldState, receipts, maybeRequests)));
       }
     } catch (MerkleTrieException ex) {
-      context.getSynchronizer().healWorldState(ex.getMaybeAddress(), ex.getLocation());
+      context.getWorldStateArchive().heal(ex.getMaybeAddress(), ex.getLocation());
       return new BlockProcessingResult(Optional.empty(), ex);
     } catch (StorageException ex) {
       var retval = new BlockProcessingResult(Optional.empty(), ex);

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/BonsaiWorldStateProvider.java

@@ -34,6 +34,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.apache.tuweni.bytes.Bytes;
@@ -44,16 +45,19 @@ public class BonsaiWorldStateProvider extends DiffBasedWorldStateProvider {
 
   private static final Logger LOG = LoggerFactory.getLogger(BonsaiWorldStateProvider.class);
   private final BonsaiCachedMerkleTrieLoader bonsaiCachedMerkleTrieLoader;
+  private final Supplier<WorldStateHealer> worldStateHealerSupplier;
 
   public BonsaiWorldStateProvider(
       final BonsaiWorldStateKeyValueStorage worldStateKeyValueStorage,
       final Blockchain blockchain,
       final Optional<Long> maxLayersToLoad,
       final BonsaiCachedMerkleTrieLoader bonsaiCachedMerkleTrieLoader,
       final BesuContext pluginContext,
-      final EvmConfiguration evmConfiguration) {
+      final EvmConfiguration evmConfiguration,
+      final Supplier<WorldStateHealer> worldStateHealerSupplier) {
     super(worldStateKeyValueStorage, blockchain, maxLayersToLoad, pluginContext);
     this.bonsaiCachedMerkleTrieLoader = bonsaiCachedMerkleTrieLoader;
+    this.worldStateHealerSupplier = worldStateHealerSupplier;
     provideCachedWorldStorageManager(
         new BonsaiCachedWorldStorageManager(
             this, worldStateKeyValueStorage, this::cloneBonsaiWorldStateConfig));
@@ -69,9 +73,11 @@ public BonsaiWorldStateProvider(
       final BonsaiWorldStateKeyValueStorage worldStateKeyValueStorage,
       final Blockchain blockchain,
       final BonsaiCachedMerkleTrieLoader bonsaiCachedMerkleTrieLoader,
-      final EvmConfiguration evmConfiguration) {
+      final EvmConfiguration evmConfiguration,
+      final Supplier<WorldStateHealer> worldStateHealerSupplier) {
     super(worldStateKeyValueStorage, blockchain, trieLogManager);
     this.bonsaiCachedMerkleTrieLoader = bonsaiCachedMerkleTrieLoader;
+    this.worldStateHealerSupplier = worldStateHealerSupplier;
     provideCachedWorldStorageManager(bonsaiCachedWorldStorageManager);
     loadPersistedState(
         new BonsaiWorldState(
@@ -151,4 +157,9 @@ public void prepareStateHealing(final Address address, final Bytes location) {
   private DiffBasedWorldStateConfig cloneBonsaiWorldStateConfig() {
     return new DiffBasedWorldStateConfig(defaultWorldStateConfig);
   }
+
+  @Override
+  public void heal(final Optional<Address> maybeAccountToRepair, final Bytes location) {
+    worldStateHealerSupplier.get().heal(maybeAccountToRepair, location);
+  }
 }

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/trie/forest/ForestWorldStateArchive.java

@@ -112,6 +112,11 @@ public <U> Optional<U> getAccountProof(
             blockHeader.getStateRoot(), accountAddress, accountStorageKeys));
   }
 
+  @Override
+  public void heal(final Optional<Address> maybeAccountToRepair, final Bytes location) {
+    // no heal needed for Forest
+  }
+
   @Override
   public void close() {
     // no op

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/worldstate/WorldStateArchive.java

@@ -64,4 +64,24 @@ <U> Optional<U> getAccountProof(
       final Address accountAddress,
       final List<UInt256> accountStorageKeys,
       final Function<Optional<WorldStateProof>, ? extends Optional<U>> mapper);
+
+  /**
+   * Heal the world state to fix inconsistency
+   *
+   * @param maybeAccountToRepair the optional account to repair
+   * @param location the location of the inconsistency
+   */
+  void heal(Optional<Address> maybeAccountToRepair, Bytes location);
+
+  /** A world state healer */
+  @FunctionalInterface
+  interface WorldStateHealer {
+    /**
+     * Heal the world state to fix inconsistency
+     *
+     * @param maybeAccountToRepair the optional account to repair
+     * @param location the location of the inconsistency
+     */
+    void heal(Optional<Address> maybeAccountToRepair, Bytes location);
+  }
 }

ethereum/core/src/test-support/java/org/hyperledger/besu/ethereum/core/InMemoryKeyValueStorageProvider.java

@@ -14,6 +14,8 @@
  */
 package org.hyperledger.besu.ethereum.core;
 
+import static org.hyperledger.besu.ethereum.core.WorldStateHealerHelper.throwingWorldStateHealerSupplier;
+
 import org.hyperledger.besu.ethereum.chain.Blockchain;
 import org.hyperledger.besu.ethereum.chain.DefaultBlockchain;
 import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
@@ -105,7 +107,8 @@ public static BonsaiWorldStateProvider createBonsaiInMemoryWorldStateArchive(
         Optional.empty(),
         bonsaiCachedMerkleTrieLoader,
         null,
-        evmConfiguration);
+        evmConfiguration,
+        throwingWorldStateHealerSupplier());
   }
 
   public static MutableWorldState createInMemoryWorldState() {

ethereum/core/src/test-support/java/org/hyperledger/besu/ethereum/core/WorldStateHealerHelper.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright contributors to Besu.
+ *
+ * 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.core;
+
+import org.hyperledger.besu.datatypes.Address;
+import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive.WorldStateHealer;
+
+import java.util.Optional;
+import java.util.function.Supplier;
+
+import org.apache.tuweni.bytes.Bytes;
+
+public class WorldStateHealerHelper {
+
+  public static WorldStateHealer throwingHealer(
+      final Optional<Address> maybeAccountToRepair, final Bytes location) {
+    throw new RuntimeException(
+        "World state needs to be healed: "
+            + maybeAccountToRepair.map(address -> "account to repair: " + address).orElse("")
+            + " location: "
+            + location.toHexString());
+  }
+
+  public static Supplier<WorldStateHealer> throwingWorldStateHealerSupplier() {
+    return () -> WorldStateHealerHelper::throwingHealer;
+  }
+}

ethereum/core/src/test/java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/AbstractIsolationTests.java

@@ -15,6 +15,7 @@
 package org.hyperledger.besu.ethereum.trie.diffbased.bonsai;
 
 import static org.hyperledger.besu.ethereum.core.InMemoryKeyValueStorageProvider.createInMemoryBlockchain;
+import static org.hyperledger.besu.ethereum.core.WorldStateHealerHelper.throwingWorldStateHealerSupplier;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.mock;
@@ -167,7 +168,8 @@ public void createStorage() {
             Optional.of(16L),
             new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
             null,
-            EvmConfiguration.DEFAULT);
+            EvmConfiguration.DEFAULT,
+            throwingWorldStateHealerSupplier());
     var ws = archive.getMutable();
     genesisState.writeStateTo(ws);
     protocolContext = new ProtocolContext(blockchain, archive, null, new BadBlockManager());

ethereum/core/src/test/java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/BonsaiWorldStateProviderTest.java

@@ -15,6 +15,7 @@
 package org.hyperledger.besu.ethereum.trie.diffbased.bonsai;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.hyperledger.besu.ethereum.core.WorldStateHealerHelper.throwingWorldStateHealerSupplier;
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.BLOCKCHAIN;
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.TRIE_BRANCH_STORAGE;
 import static org.hyperledger.besu.ethereum.trie.diffbased.common.storage.DiffBasedWorldStateKeyValueStorage.WORLD_BLOCK_HASH_KEY;
@@ -111,7 +112,8 @@ void testGetMutableReturnPersistedStateWhenNeeded() {
                 DataStorageConfiguration.DEFAULT_BONSAI_CONFIG),
             blockchain,
             new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
-            EvmConfiguration.DEFAULT);
+            EvmConfiguration.DEFAULT,
+            throwingWorldStateHealerSupplier());
 
     assertThat(bonsaiWorldStateArchive.getMutable(chainHead, true))
         .containsInstanceOf(BonsaiWorldState.class);
@@ -129,7 +131,8 @@ void testGetMutableReturnEmptyWhenLoadMoreThanLimitLayersBack() {
             Optional.of(512L),
             new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
             null,
-            EvmConfiguration.DEFAULT);
+            EvmConfiguration.DEFAULT,
+            throwingWorldStateHealerSupplier());
     final BlockHeader blockHeader = blockBuilder.number(0).buildHeader();
     final BlockHeader chainHead = blockBuilder.number(512).buildHeader();
     when(blockchain.getChainHeadHeader()).thenReturn(chainHead);
@@ -150,7 +153,8 @@ void testGetMutableWhenLoadLessThanLimitLayersBack() {
                 DataStorageConfiguration.DEFAULT_BONSAI_CONFIG),
             blockchain,
             new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
-            EvmConfiguration.DEFAULT);
+            EvmConfiguration.DEFAULT,
+            throwingWorldStateHealerSupplier());
     final BlockHeader blockHeader = blockBuilder.number(0).buildHeader();
     final BlockHeader chainHead = blockBuilder.number(511).buildHeader();
     final BonsaiWorldState mockWorldState = mock(BonsaiWorldState.class);
@@ -185,7 +189,8 @@ void testGetMutableWithStorageInconsistencyRollbackTheState() {
                 worldStateKeyValueStorage,
                 blockchain,
                 new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
-                EvmConfiguration.DEFAULT));
+                EvmConfiguration.DEFAULT,
+                throwingWorldStateHealerSupplier()));
     final BlockHeader blockHeader = blockBuilder.number(0).buildHeader();
 
     when(blockchain.getBlockHeader(blockHeader.getHash())).thenReturn(Optional.of(blockHeader));
@@ -214,7 +219,8 @@ void testGetMutableWithStorageConsistencyNotRollbackTheState() {
                 worldStateKeyValueStorage,
                 blockchain,
                 new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
-                EvmConfiguration.DEFAULT));
+                EvmConfiguration.DEFAULT,
+                throwingWorldStateHealerSupplier()));
 
     final BlockHeader blockHeader = blockBuilder.number(0).buildHeader();
 
@@ -254,7 +260,8 @@ void testGetMutableWithStorageConsistencyToRollbackAndRollForwardTheState() {
                 worldStateKeyValueStorage,
                 blockchain,
                 new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
-                EvmConfiguration.DEFAULT));
+                EvmConfiguration.DEFAULT,
+                throwingWorldStateHealerSupplier()));
 
     // initial persisted state hash key
     when(blockchain.getBlockHeader(Hash.ZERO)).thenReturn(Optional.of(blockHeaderChainA));
@@ -297,7 +304,8 @@ void testGetMutableWithRollbackNotOverrideTrieLogLayer() {
                     DataStorageConfiguration.DEFAULT_BONSAI_CONFIG),
                 blockchain,
                 new BonsaiCachedMerkleTrieLoader(new NoOpMetricsSystem()),
-                EvmConfiguration.DEFAULT));
+                EvmConfiguration.DEFAULT,
+                throwingWorldStateHealerSupplier()));
 
     // initial persisted state hash key
     when(blockchain.getBlockHeader(Hash.ZERO)).thenReturn(Optional.of(blockHeaderChainA));