Feature/multi version flat db

besu/src/main/java/org/hyperledger/besu/cli/BesuCommand.java

@@ -2111,7 +2111,7 @@ private void issueOptionWarnings() {
 
     CommandLineUtils.failIfOptionDoesntMeetRequirement(
         commandLine,
-        "--Xsnapsync-synchronizer-flat option can only be used when -Xsnapsync-synchronizer-flat-db-healing-enabled is true",
+        "--Xsnapsync-synchronizer-flat option can only be used when --Xsnapsync-synchronizer-flat-db-healing-enabled is true",
         unstableSynchronizerOptions.isSnapsyncFlatDbHealingEnabled(),
         asList(
             "--Xsnapsync-synchronizer-flat-account-healed-count-per-request",

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

@@ -1073,6 +1073,11 @@ WorldStateArchive createWorldStateArchive(
       final CachedMerkleTrieLoader cachedMerkleTrieLoader) {
     switch (dataStorageConfiguration.getDataStorageFormat()) {
       case BONSAI:
+        // TODO, better integrate. Just for PoC, explicitly set our bonsai context chain head:
+        ((BonsaiWorldStateKeyValueStorage) worldStateStorage)
+            .getFlatDbStrategy()
+            .updateBlockContext(blockchain.getChainHeadHeader());
+
         return new BonsaiWorldStateProvider(
             (BonsaiWorldStateKeyValueStorage) worldStateStorage,
             blockchain,

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright Hyperledger Besu Contributors.
+ *
+ * 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.plugin.data.BlockHeader;
+
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+
+/** Context which holds information relevant to a bonsai archive storage query. */
+public class BonsaiContext {
+
+  private final AtomicReference<BlockHeader> blockHeader;
+
+  public BonsaiContext() {
+    blockHeader = new AtomicReference<>();
+  }
+
+  public BonsaiContext copy() {
+    var newCtx = new BonsaiContext();
+    Optional.ofNullable(blockHeader.get()).ifPresent(newCtx::setBlockHeader);
+    return newCtx;
+  }
+
+  public BonsaiContext setBlockHeader(final BlockHeader blockHeader) {
+    this.blockHeader.set(blockHeader);
+    return this;
+  }
+
+  public Optional<BlockHeader> getBlockHeader() {
+    return Optional.ofNullable(blockHeader.get());
+  }
+}

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

@@ -23,6 +23,7 @@
 import org.hyperledger.besu.ethereum.bonsai.cache.CachedMerkleTrieLoader;
 import org.hyperledger.besu.ethereum.bonsai.cache.CachedWorldStorageManager;
 import org.hyperledger.besu.ethereum.bonsai.storage.BonsaiWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.bonsai.storage.flat.ArchiveFlatDbStrategy;
 import org.hyperledger.besu.ethereum.bonsai.trielog.TrieLogManager;
 import org.hyperledger.besu.ethereum.bonsai.worldview.BonsaiWorldState;
 import org.hyperledger.besu.ethereum.bonsai.worldview.BonsaiWorldStateUpdateAccumulator;
@@ -159,6 +160,19 @@ public Optional<MutableWorldState> getMutable(
     if (shouldPersistState) {
       return getMutable(blockHeader.getStateRoot(), blockHeader.getHash());
     } else {
+      // TODO this needs to be better integrated && ensure block is canonical
+      // HACK for kikori PoC, if we have the trielog for this block, we can assume we have it in
+      // flatDB
+      // although, in practice we can only serve canonical chain worldstates and need to fall back
+      // to state rolling if the requested block is a fork.
+      if (this.worldStateStorage.getFlatDbStrategy() instanceof ArchiveFlatDbStrategy
+          && trieLogManager.getTrieLogLayer(blockHeader.getBlockHash()).isPresent()) {
+
+        var contextSafeCopy = worldStateStorage.getContextSafeCopy();
+        contextSafeCopy.getFlatDbStrategy().updateBlockContext(blockHeader);
+        return Optional.of(new BonsaiWorldState(this, contextSafeCopy));
+      }
+
       final BlockHeader chainHeadBlockHeader = blockchain.getChainHeadHeader();
       if (chainHeadBlockHeader.getNumber() - blockHeader.getNumber()
           >= trieLogManager.getMaxLayersToLoad()) {

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

@@ -21,6 +21,8 @@
 
 import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.datatypes.StorageSlotKey;
+import org.hyperledger.besu.ethereum.bonsai.BonsaiContext;
+import org.hyperledger.besu.ethereum.bonsai.storage.flat.ArchiveFlatDbStrategy;
 import org.hyperledger.besu.ethereum.bonsai.storage.flat.FlatDbStrategy;
 import org.hyperledger.besu.ethereum.bonsai.storage.flat.FullFlatDbStrategy;
 import org.hyperledger.besu.ethereum.bonsai.storage.flat.PartialFlatDbStrategy;
@@ -114,21 +116,33 @@ private void loadFlatDbStrategy() {
       this.flatDbMode = newFlatDbMode;
       if (flatDbMode == FlatDbMode.FULL) {
         this.flatDbStrategy = new FullFlatDbStrategy(metricsSystem);
+      } else if (flatDbMode == FlatDbMode.ARCHIVE) {
+        this.flatDbStrategy = new ArchiveFlatDbStrategy(new BonsaiContext(), metricsSystem);
       } else {
         this.flatDbStrategy = new PartialFlatDbStrategy(metricsSystem);
       }
     }
   }
 
+  public BonsaiWorldStateKeyValueStorage getContextSafeCopy() {
+    return new BonsaiWorldStateKeyValueStorage(
+        flatDbMode,
+        flatDbStrategy.contextSafeClone(),
+        composedWorldStateStorage,
+        trieLogStorage,
+        metricsSystem);
+  }
+
   public FlatDbMode deriveFlatDbStrategy() {
     var flatDbMode =
         FlatDbMode.fromVersion(
             composedWorldStateStorage
                 .get(TRIE_BRANCH_STORAGE, FLAT_DB_MODE)
                 .map(Bytes::wrap)
-                .orElse(FlatDbMode.PARTIAL.getVersion()));
+                // .orElse(FlatDbMode.PARTIAL.getVersion()));
+                // TODO: this is for archive testing only, remove me.
+                .orElse(FlatDbMode.ARCHIVE.getVersion()));
     LOG.info("Bonsai flat db mode found {}", flatDbMode);
-
     return flatDbMode;
   }