SnapServer using FlatDB Strategy

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

@@ -744,9 +744,6 @@ public BesuController build() {
             peerValidators,
             Optional.empty());
 
-    final Optional<SnapProtocolManager> maybeSnapProtocolManager =
-        createSnapProtocolManager(peerValidators, ethPeers, snapMessages, worldStateArchive);
-
     final PivotBlockSelector pivotBlockSelector =
         createPivotSelector(
             protocolSchedule, protocolContext, ethContext, syncState, metricsSystem);
@@ -764,6 +761,9 @@ public BesuController build() {
 
     protocolContext.setSynchronizer(Optional.of(synchronizer));
 
+    final Optional<SnapProtocolManager> maybeSnapProtocolManager =
+        createSnapProtocolManager(protocolContext, peerValidators, ethPeers, snapMessages);
+
     final MiningCoordinator miningCoordinator =
         createMiningCoordinator(
             protocolSchedule,
@@ -1086,12 +1086,12 @@ protected ProtocolContext createProtocolContext(
   }
 
   private Optional<SnapProtocolManager> createSnapProtocolManager(
+      final ProtocolContext protocolContext,
       final List<PeerValidator> peerValidators,
       final EthPeers ethPeers,
-      final EthMessages snapMessages,
-      final WorldStateArchive worldStateArchive) {
+      final EthMessages snapMessages) {
     return Optional.of(
-        new SnapProtocolManager(peerValidators, ethPeers, snapMessages, worldStateArchive));
+        new SnapProtocolManager(peerValidators, ethPeers, snapMessages, protocolContext));
   }
 
   WorldStateArchive createWorldStateArchive(

ethereum/core/build.gradle

@@ -47,6 +47,7 @@ dependencies {
 
   implementation 'com.fasterxml.jackson.core:jackson-databind'
   implementation 'com.google.guava:guava'
+  implementation 'com.github.ben-manes.caffeine:caffeine'
   implementation 'com.google.dagger:dagger'
   annotationProcessor 'com.google.dagger:dagger-compiler'
   implementation 'io.opentelemetry:opentelemetry-api'

ethereum/core/src/integration-test/java/org/hyperledger/besu/ethereum/trie/forest/pruner/PrunerIntegrationTest.java

@@ -42,12 +42,13 @@
 import org.hyperledger.besu.testutil.MockExecutorService;
 
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
+import java.util.NavigableMap;
 import java.util.Optional;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -59,7 +60,8 @@ public class PrunerIntegrationTest {
 
   private final BlockDataGenerator gen = new BlockDataGenerator();
   private final NoOpMetricsSystem metricsSystem = new NoOpMetricsSystem();
-  private final Map<Bytes, Optional<byte[]>> hashValueStore = new HashMap<>();
+  private final NavigableMap<Bytes, Optional<byte[]>> hashValueStore =
+      new TreeMap<>(Comparator.comparing(Bytes::toHexString));
   private final InMemoryKeyValueStorage stateStorage = new TestInMemoryStorage(hashValueStore);
   private final WorldStateStorage worldStateStorage =
       new ForestWorldStateKeyValueStorage(stateStorage);
@@ -260,7 +262,7 @@ private MerkleTrie<Bytes32, Bytes> createStorageTrie(final Bytes32 rootHash) {
   // Proxy class so that we have access to the constructor that takes our own map
   private static class TestInMemoryStorage extends InMemoryKeyValueStorage {
 
-    public TestInMemoryStorage(final Map<Bytes, Optional<byte[]>> hashValueStore) {
+    public TestInMemoryStorage(final NavigableMap<Bytes, Optional<byte[]>> hashValueStore) {
       super(hashValueStore);
     }
   }

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/proof/WorldStateProofProvider.java

@@ -28,6 +28,7 @@
 import org.hyperledger.besu.ethereum.worldstate.StateTrieAccountValue;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateStorage;
 
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -41,6 +42,8 @@
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
 import org.apache.tuweni.units.bigints.UInt256;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * The WorldStateProofProvider class is responsible for providing proofs for world state entries. It
@@ -49,6 +52,7 @@
 public class WorldStateProofProvider {
 
   private final WorldStateStorage worldStateStorage;
+  private static final Logger LOG = LoggerFactory.getLogger(WorldStateProofProvider.class);
 
   public WorldStateProofProvider(final WorldStateStorage worldStateStorage) {
     this.worldStateStorage = worldStateStorage;
@@ -85,7 +89,8 @@ private SortedMap<UInt256, Proof<Bytes>> getStorageProofs(
       final List<UInt256> accountStorageKeys) {
     final MerkleTrie<Bytes32, Bytes> storageTrie =
         newAccountStorageTrie(accountHash, account.getStorageRoot());
-    final NavigableMap<UInt256, Proof<Bytes>> storageProofs = new TreeMap<>();
+    final NavigableMap<UInt256, Proof<Bytes>> storageProofs =
+        new TreeMap<>(Comparator.comparing(Bytes32::toHexString));
     accountStorageKeys.forEach(
         key -> storageProofs.put(key, storageTrie.getValueWithProof(Hash.hash(key))));
     return storageProofs;
@@ -153,19 +158,26 @@ public boolean isValidRangeProof(
       final SortedMap<Bytes32, Bytes> keys) {
 
     // check if it's monotonic increasing
-    if (!Ordering.natural().isOrdered(keys.keySet())) {
+    if (keys.size() > 1 && !Ordering.natural().isOrdered(keys.keySet())) {
       return false;
     }
 
-    // when proof is empty we need to have all the keys to reconstruct the trie
+    // when proof is empty and we requested the full range, we should
+    // have all the keys to reconstruct the trie
     if (proofs.isEmpty()) {
-      final MerkleTrie<Bytes, Bytes> trie = new SimpleMerklePatriciaTrie<>(Function.identity());
-      // add the received keys in the trie
-      for (Map.Entry<Bytes32, Bytes> key : keys.entrySet()) {
-        trie.put(key.getKey(), key.getValue());
+      if (startKeyHash.equals(Bytes32.ZERO)) {
+        final MerkleTrie<Bytes, Bytes> trie = new SimpleMerklePatriciaTrie<>(Function.identity());
+        // add the received keys in the trie
+        for (Map.Entry<Bytes32, Bytes> key : keys.entrySet()) {
+          trie.put(key.getKey(), key.getValue());
+        }
+        return rootHash.equals(trie.getRootHash());
+      } else {
+        // TODO: possibly accept a node loader so we can verify this with already
+        //  completed partial storage requests
+        LOG.info("failing proof due to incomplete range without proofs");
+        return false;
       }
-
-      return rootHash.equals(trie.getRootHash());
     }
 
     // reconstruct a part of the trie with the proof

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

@@ -30,12 +30,13 @@
 import org.hyperledger.besu.ethereum.trie.bonsai.cache.CachedMerkleTrieLoader;
 import org.hyperledger.besu.ethereum.trie.bonsai.cache.CachedWorldStorageManager;
 import org.hyperledger.besu.ethereum.trie.bonsai.storage.BonsaiWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.trie.bonsai.storage.flat.FullFlatDbStrategy;
 import org.hyperledger.besu.ethereum.trie.bonsai.trielog.TrieLogManager;
 import org.hyperledger.besu.ethereum.trie.bonsai.worldview.BonsaiWorldState;
 import org.hyperledger.besu.ethereum.trie.bonsai.worldview.BonsaiWorldStateUpdateAccumulator;
 import org.hyperledger.besu.ethereum.trie.patricia.StoredMerklePatriciaTrie;
+import org.hyperledger.besu.ethereum.worldstate.FlatWorldStateArchive;
 import org.hyperledger.besu.ethereum.worldstate.StateTrieAccountValue;
-import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
 import org.hyperledger.besu.evm.internal.EvmConfiguration;
 import org.hyperledger.besu.evm.worldstate.WorldState;
 import org.hyperledger.besu.plugin.BesuContext;
@@ -54,7 +55,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class BonsaiWorldStateProvider implements WorldStateArchive {
+public class BonsaiWorldStateProvider implements FlatWorldStateArchive {
 
   private static final Logger LOG = LoggerFactory.getLogger(BonsaiWorldStateProvider.class);
 
@@ -130,7 +131,7 @@ public Optional<WorldState> get(final Hash rootHash, final Hash blockHash) {
 
   @Override
   public boolean isWorldStateAvailable(final Hash rootHash, final Hash blockHash) {
-    return cachedWorldStorageManager.containWorldStateStorage(blockHash)
+    return cachedWorldStorageManager.contains(blockHash)
         || persistedState.blockHash().equals(blockHash)
         || worldStateStorage.isWorldStateAvailable(rootHash, blockHash);
   }
@@ -336,8 +337,9 @@ public TrieLogManager getTrieLogManager() {
     return trieLogManager;
   }
 
-  public CachedWorldStorageManager getCachedWorldStorageManager() {
-    return cachedWorldStorageManager;
+  @Override
+  public Optional<CachedWorldStorageManager> getCachedWorldStorageManager() {
+    return Optional.of(cachedWorldStorageManager);
   }
 
   @Override
@@ -348,6 +350,11 @@ public void resetArchiveStateTo(final BlockHeader blockHeader) {
         blockHeader, persistedState.getWorldStateRootHash(), persistedState);
   }
 
+  @Override
+  public boolean isFullFlat() {
+    return worldStateStorage.getFlatDbStrategy() instanceof FullFlatDbStrategy;
+  }
+
   @Override
   public <U> Optional<U> getAccountProof(
       final BlockHeader blockHeader,

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/trie/bonsai/cache/CachedWorldStorageManager.java

@@ -15,6 +15,7 @@
 package org.hyperledger.besu.ethereum.trie.bonsai.cache;
 
 import org.hyperledger.besu.datatypes.Hash;
+import org.hyperledger.besu.ethereum.chain.Blockchain;
 import org.hyperledger.besu.ethereum.core.BlockHeader;
 import org.hyperledger.besu.ethereum.trie.bonsai.BonsaiWorldStateProvider;
 import org.hyperledger.besu.ethereum.trie.bonsai.storage.BonsaiSnapshotWorldStateKeyValueStorage;
@@ -29,8 +30,11 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import org.apache.tuweni.bytes.Bytes32;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -41,6 +45,11 @@ public class CachedWorldStorageManager
   private static final Logger LOG = LoggerFactory.getLogger(CachedWorldStorageManager.class);
   private final BonsaiWorldStateProvider archive;
   private final EvmConfiguration evmConfiguration;
+  private final Cache<Hash, BlockHeader> stateRootToBlockHeaderCache =
+      Caffeine.newBuilder()
+          .maximumSize(RETAINED_LAYERS)
+          .expireAfterWrite(100, TimeUnit.MINUTES)
+          .build();
 
   private final BonsaiWorldStateKeyValueStorage rootWorldStateStorage;
   private final Map<Bytes32, CachedBonsaiWorldView> cachedWorldStatesByHash;
@@ -104,6 +113,8 @@ public synchronized void addCachedLayer(
                 blockHeader,
                 ((BonsaiWorldStateLayerStorage) forWorldState.getWorldStateStorage()).clone()));
       }
+      // add stateroot -> blockHeader cache entry
+      stateRootToBlockHeaderCache.put(blockHeader.getStateRoot(), blockHeader);
     }
     scrubCachedLayers(blockHeader.getNumber());
   }
@@ -192,14 +203,59 @@ public Optional<BonsaiWorldState> getHeadWorldState(
             });
   }
 
-  public boolean containWorldStateStorage(final Hash blockHash) {
+  public boolean contains(final Hash blockHash) {
     return cachedWorldStatesByHash.containsKey(blockHash);
   }
 
   public void reset() {
     this.cachedWorldStatesByHash.clear();
   }
 
+  public void primeRootToBlockHashCache(final Blockchain blockchain, final int numEntries) {
+    // prime the stateroot-to-blockhash cache
+    long head = blockchain.getChainHeadHeader().getNumber();
+    for (long i = head; i > Math.max(0, head - numEntries); i--) {
+      blockchain
+          .getBlockHeader(i)
+          .ifPresent(header -> stateRootToBlockHeaderCache.put(header.getStateRoot(), header));
+    }
+  }
+
+  /**
+   * Returns the worldstate for the supplied root hash, or the head worldstate if no root hash is
+   * supplied. synchronized to prevent concurrent adds to the cache of the same root hash.
+   *
+   * @param rootHash optional rootHash to supply worldstate storage for
+   * @return Optional worldstate storage
+   */
+  public synchronized Optional<BonsaiWorldStateKeyValueStorage> getStorageByRootHash(
+      final Optional<Hash> rootHash) {
+    if (rootHash.isPresent()) {
+      // if we supplied a hash, return the worldstate for that hash if it is available:
+      return rootHash
+          .map(stateRootToBlockHeaderCache::getIfPresent)
+          .flatMap(
+              header ->
+                  Optional.ofNullable(cachedWorldStatesByHash.get(header.getHash()))
+                      .map(CachedBonsaiWorldView::getWorldStateStorage)
+                      .or(
+                          () -> {
+                            // if not cached already, maybe fetch and cache this worldstate
+                            var maybeWorldState =
+                                archive.getMutable(header, false).map(BonsaiWorldState.class::cast);
+                            maybeWorldState.ifPresent(
+                                ws -> addCachedLayer(header, header.getStateRoot(), ws));
+                            return maybeWorldState.map(BonsaiWorldState::getWorldStateStorage);
+                          }));
+    } else {
+      // if we did not supply a hash, return the head worldstate from cachedWorldStates
+      return rootWorldStateStorage
+          .getWorldStateBlockHash()
+          .map(cachedWorldStatesByHash::get)
+          .map(CachedBonsaiWorldView::getWorldStateStorage);
+    }
+  }
+
   @Override
   public void onClearStorage() {
     this.cachedWorldStatesByHash.clear();

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/trie/bonsai/cache/NoOpCachedWorldStorageManager.java

@@ -38,7 +38,7 @@ public synchronized void addCachedLayer(
   }
 
   @Override
-  public boolean containWorldStateStorage(final Hash blockHash) {
+  public boolean contains(final Hash blockHash) {
     return false;
   }
 

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

@@ -41,14 +41,14 @@
 
 import java.nio.charset.StandardCharsets;
 import java.util.List;
-import java.util.Map;
 import java.util.NavigableMap;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 
+import kotlin.Pair;
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
 import org.slf4j.Logger;
@@ -157,9 +157,7 @@ public Optional<Bytes> getAccountStorageTrieNode(
 
   @Override
   public Optional<Bytes> getTrieNodeUnsafe(final Bytes key) {
-    return composedWorldStateStorage
-        .get(TRIE_BRANCH_STORAGE, Bytes.concatenate(key).toArrayUnsafe())
-        .map(Bytes::wrap);
+    return composedWorldStateStorage.get(TRIE_BRANCH_STORAGE, key.toArrayUnsafe()).map(Bytes::wrap);
   }
 
   public Optional<byte[]> getTrieLog(final Hash blockHash) {
@@ -217,22 +215,42 @@ public Optional<Bytes> getStorageValueByStorageSlotKey(
   }
 
   @Override
-  public Map<Bytes32, Bytes> streamFlatAccounts(
+  public NavigableMap<Bytes32, Bytes> streamFlatAccounts(
       final Bytes startKeyHash, final Bytes32 endKeyHash, final long max) {
     return flatDbStrategyProvider
         .getFlatDbStrategy(composedWorldStateStorage)
         .streamAccountFlatDatabase(composedWorldStateStorage, startKeyHash, endKeyHash, max);
   }
 
   @Override
-  public Map<Bytes32, Bytes> streamFlatStorages(
+  public NavigableMap<Bytes32, Bytes> streamFlatAccounts(
+      final Bytes startKeyHash,
+      final Bytes32 endKeyHash,
+      final Predicate<Pair<Bytes32, Bytes>> takeWhile) {
+    return getFlatDbStrategy()
+        .streamAccountFlatDatabase(composedWorldStateStorage, startKeyHash, endKeyHash, takeWhile);
+  }
+
+  @Override
+  public NavigableMap<Bytes32, Bytes> streamFlatStorages(
       final Hash accountHash, final Bytes startKeyHash, final Bytes32 endKeyHash, final long max) {
     return flatDbStrategyProvider
         .getFlatDbStrategy(composedWorldStateStorage)
         .streamStorageFlatDatabase(
             composedWorldStateStorage, accountHash, startKeyHash, endKeyHash, max);
   }
 
+  @Override
+  public NavigableMap<Bytes32, Bytes> streamFlatStorages(
+      final Hash accountHash,
+      final Bytes startKeyHash,
+      final Bytes32 endKeyHash,
+      final Predicate<Pair<Bytes32, Bytes>> takeWhile) {
+    return getFlatDbStrategy()
+        .streamStorageFlatDatabase(
+            composedWorldStateStorage, accountHash, startKeyHash, endKeyHash, takeWhile);
+  }
+
   public NavigableMap<Bytes32, AccountStorageEntry> storageEntriesFrom(
       final Hash addressHash, final Bytes32 startKeyHash, final int limit) {
     throw new RuntimeException("Bonsai Tries does not currently support enumerating storage");