Move context management out of flat DB strategy

Signed-off-by: Matthew Whitehead <[email protected]>
hyperledger · Jan 20, 2025 · 633fad4 · 633fad4
1 parent 25c5bc6
commit 633fad4
Show file tree

Hide file tree

Showing 20 changed files with 121 additions and 99 deletions.
diff --git a/.../org/hyperledger/besu/ethereum/trie/diffbased/bonsai/BonsaiArchiveWorldStateProvider.java b/.../org/hyperledger/besu/ethereum/trie/diffbased/bonsai/BonsaiArchiveWorldStateProvider.java
@@ -97,15 +97,14 @@ public Optional<MutableWorldState> getMutable(
             trieLogManager.getMaxLayersToLoad());
         Optional<MutableWorldState> cachedWorldState =
             cachedWorldStorageManager
-                .getHeadWorldState(blockchain::getBlockHeader)
+                .getWorldState(chainHeadBlockHeader.getHash())
                 .map(MutableWorldState::disableTrie)
                 .flatMap(
                     worldState ->
                         rollMutableArchiveStateToBlockHash( // This is a tiny action for archive
                             // state
                             (DiffBasedWorldState) worldState, blockHeader.getHash()))
                 .map(MutableWorldState::freeze);
-
         return cachedWorldState;
       }
       return super.getMutable(blockHeader, false);

diff --git a/...in/java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/BonsaiWorldStateProvider.java b/...in/java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/BonsaiWorldStateProvider.java
@@ -105,7 +105,7 @@ public Optional<MutableWorldState> getMutable(
       return cachedWorldStorageManager
           .getWorldState(blockHeader.getHash())
           .or(() -> cachedWorldStorageManager.getNearestWorldState(blockHeader))
-          .or(() -> cachedWorldStorageManager.getHeadWorldState(blockchain::getBlockHeader))
+          .or(() -> cachedWorldStorageManager.getWorldState(chainHeadBlockHeader.getHash()))
           .flatMap(worldState -> rollMutableStateToBlockHash(worldState, blockHeader.getHash()))
           .map(MutableWorldState::freeze);
     }

diff --git a/.../besu/ethereum/trie/diffbased/bonsai/storage/BonsaiSnapshotWorldStateKeyValueStorage.java b/.../besu/ethereum/trie/diffbased/bonsai/storage/BonsaiSnapshotWorldStateKeyValueStorage.java
@@ -71,7 +71,8 @@ public Updater updater() {
     return new Updater(
         ((SnappedKeyValueStorage) composedWorldStateStorage).getSnapshotTransaction(),
         trieLogStorage.startTransaction(),
-        getFlatDbStrategy());
+        getFlatDbStrategy(),
+        composedWorldStateStorage);
   }
 
   @Override

diff --git a/...erledger/besu/ethereum/trie/diffbased/bonsai/storage/BonsaiWorldStateKeyValueStorage.java b/...erledger/besu/ethereum/trie/diffbased/bonsai/storage/BonsaiWorldStateKeyValueStorage.java
@@ -187,27 +187,34 @@ public Updater updater() {
     return new Updater(
         composedWorldStateStorage.startTransaction(),
         trieLogStorage.startTransaction(),
-        getFlatDbStrategy());
+        getFlatDbStrategy(),
+        composedWorldStateStorage);
   }
 
   public static class Updater implements DiffBasedWorldStateKeyValueStorage.Updater {
 
     private final SegmentedKeyValueStorageTransaction composedWorldStateTransaction;
     private final KeyValueStorageTransaction trieLogStorageTransaction;
     private final FlatDbStrategy flatDbStrategy;
+    private final SegmentedKeyValueStorage worldStorage;
 
     public Updater(
         final SegmentedKeyValueStorageTransaction composedWorldStateTransaction,
         final KeyValueStorageTransaction trieLogStorageTransaction,
-        final FlatDbStrategy flatDbStrategy) {
+        final FlatDbStrategy flatDbStrategy,
+        final SegmentedKeyValueStorage worldStorage) {
 
       this.composedWorldStateTransaction = composedWorldStateTransaction;
       this.trieLogStorageTransaction = trieLogStorageTransaction;
       this.flatDbStrategy = flatDbStrategy;
+      this.worldStorage =
+          worldStorage; // An update could need to read from world storage to decide how to PUT to
+      // it (i.e. Bonsai archive)
     }
 
     public Updater removeCode(final Hash accountHash, final Hash codeHash) {
-      flatDbStrategy.removeFlatCode(composedWorldStateTransaction, accountHash, codeHash);
+      flatDbStrategy.removeFlatCode(
+          worldStorage, composedWorldStateTransaction, accountHash, codeHash);
       return this;
     }
 
@@ -222,12 +229,13 @@ public Updater putCode(final Hash accountHash, final Hash codeHash, final Bytes
         // Don't save empty values
         return this;
       }
-      flatDbStrategy.putFlatCode(composedWorldStateTransaction, accountHash, codeHash, code);
+      flatDbStrategy.putFlatCode(
+          worldStorage, composedWorldStateTransaction, accountHash, codeHash, code);
       return this;
     }
 
     public Updater removeAccountInfoState(final Hash accountHash) {
-      flatDbStrategy.removeFlatAccount(composedWorldStateTransaction, accountHash);
+      flatDbStrategy.removeFlatAccount(worldStorage, composedWorldStateTransaction, accountHash);
       return this;
     }
 
@@ -236,7 +244,8 @@ public Updater putAccountInfoState(final Hash accountHash, final Bytes accountVa
         // Don't save empty values
         return this;
       }
-      flatDbStrategy.putFlatAccount(composedWorldStateTransaction, accountHash, accountValue);
+      flatDbStrategy.putFlatAccount(
+          worldStorage, composedWorldStateTransaction, accountHash, accountValue);
       return this;
     }
 
@@ -283,14 +292,14 @@ public synchronized Updater putAccountStorageTrieNode(
     public synchronized Updater putStorageValueBySlotHash(
         final Hash accountHash, final Hash slotHash, final Bytes storageValue) {
       flatDbStrategy.putFlatAccountStorageValueByStorageSlotHash(
-          composedWorldStateTransaction, accountHash, slotHash, storageValue);
+          worldStorage, composedWorldStateTransaction, accountHash, slotHash, storageValue);
       return this;
     }
 
     public synchronized void removeStorageValueBySlotHash(
         final Hash accountHash, final Hash slotHash) {
       flatDbStrategy.removeFlatAccountStorageValueByStorageSlotHash(
-          composedWorldStateTransaction, accountHash, slotHash);
+          worldStorage, composedWorldStateTransaction, accountHash, slotHash);
     }
 
     @Override

diff --git a/...rledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiArchiveFlatDbStrategy.java b/...rledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiArchiveFlatDbStrategy.java
@@ -18,10 +18,11 @@
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.ACCOUNT_INFO_STATE_ARCHIVE;
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.ACCOUNT_STORAGE_ARCHIVE;
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.ACCOUNT_STORAGE_STORAGE;
+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_NUMBER_KEY;
 
 import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.datatypes.StorageSlotKey;
-import org.hyperledger.besu.ethereum.core.BlockHeader;
 import org.hyperledger.besu.ethereum.trie.NodeLoader;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.BonsaiContext;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.storage.flat.CodeStorageStrategy;
@@ -31,6 +32,7 @@
 import org.hyperledger.besu.plugin.services.storage.SegmentedKeyValueStorage;
 import org.hyperledger.besu.plugin.services.storage.SegmentedKeyValueStorageTransaction;
 
+import java.nio.charset.StandardCharsets;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -44,18 +46,14 @@
 import org.slf4j.LoggerFactory;
 
 public class BonsaiArchiveFlatDbStrategy extends BonsaiFullFlatDbStrategy {
-  private final BonsaiContext context;
   private static final Logger LOG = LoggerFactory.getLogger(BonsaiArchiveFlatDbStrategy.class);
 
   protected final Counter getAccountFromArchiveCounter;
   protected final Counter getStorageFromArchiveCounter;
 
   public BonsaiArchiveFlatDbStrategy(
-      final BonsaiContext context,
-      final MetricsSystem metricsSystem,
-      final CodeStorageStrategy codeStorageStrategy) {
+      final MetricsSystem metricsSystem, final CodeStorageStrategy codeStorageStrategy) {
     super(metricsSystem, codeStorageStrategy);
-    this.context = context;
 
     getAccountFromArchiveCounter =
         metricsSystem.createCounter(
@@ -76,6 +74,26 @@ public BonsaiArchiveFlatDbStrategy(
   public static final byte[] DELETED_CODE_VALUE = new byte[0];
   public static final byte[] DELETED_STORAGE_VALUE = new byte[0];
 
+  public Optional<BonsaiContext> getStateArchiveContext(final SegmentedKeyValueStorage storage) {
+    // For Bonsai archive update the flat DB strategy context to match the block the state/storage
+    // represents
+    Optional<byte[]> archiveContext = storage.get(TRIE_BRANCH_STORAGE, WORLD_BLOCK_NUMBER_KEY);
+    if (archiveContext.isPresent()) {
+      try {
+        // keyNearest, use MAX_BLOCK_SUFFIX in the absence of a block context:
+        return Optional.of(
+            new BonsaiContext(
+                Long.decode("0x" + (new String(archiveContext.get(), StandardCharsets.UTF_8)))));
+      } catch (NumberFormatException e) {
+        throw new IllegalStateException(
+            "World state archive context invalid format: "
+                + new String(archiveContext.get(), StandardCharsets.UTF_8));
+      }
+    }
+
+    throw new IllegalStateException("World state missing archive context");
+  }
+
   @Override
   public Optional<Bytes> getFlatAccount(
       final Supplier<Optional<Bytes>> worldStateRootHashSupplier,
@@ -87,7 +105,9 @@ public Optional<Bytes> getFlatAccount(
     Optional<Bytes> accountFound;
 
     // keyNearest, use MAX_BLOCK_SUFFIX in the absence of a block context:
-    Bytes keyNearest = calculateArchiveKeyWithMaxSuffix(context, accountHash.toArrayUnsafe());
+    Bytes keyNearest =
+        calculateArchiveKeyWithMaxSuffix(
+            getStateArchiveContext(storage).get(), accountHash.toArrayUnsafe());
 
     // Find the nearest account state for this address and block context
     Optional<SegmentedKeyValueStorage.NearestKeyValue> nearestAccount =
@@ -229,22 +249,29 @@ protected Stream<Pair<Bytes32, Bytes>> storageToPairStream(
    */
   @Override
   public void putFlatAccount(
+      final SegmentedKeyValueStorage storage,
       final SegmentedKeyValueStorageTransaction transaction,
       final Hash accountHash,
       final Bytes accountValue) {
 
     // key suffixed with block context, or MIN_BLOCK_SUFFIX if we have no context:
-    byte[] keySuffixed = calculateArchiveKeyWithMinSuffix(context, accountHash.toArrayUnsafe());
+    byte[] keySuffixed =
+        calculateArchiveKeyWithMinSuffix(
+            getStateArchiveContext(storage).get(), accountHash.toArrayUnsafe());
 
     transaction.put(ACCOUNT_INFO_STATE, keySuffixed, accountValue.toArrayUnsafe());
   }
 
   @Override
   public void removeFlatAccount(
-      final SegmentedKeyValueStorageTransaction transaction, final Hash accountHash) {
+      final SegmentedKeyValueStorage storage,
+      final SegmentedKeyValueStorageTransaction transaction,
+      final Hash accountHash) {
 
     // insert a key suffixed with block context, with 'deleted account' value
-    byte[] keySuffixed = calculateArchiveKeyWithMinSuffix(context, accountHash.toArrayUnsafe());
+    byte[] keySuffixed =
+        calculateArchiveKeyWithMinSuffix(
+            getStateArchiveContext(storage).get(), accountHash.toArrayUnsafe());
 
     transaction.put(ACCOUNT_INFO_STATE, keySuffixed, DELETED_ACCOUNT_VALUE);
   }
@@ -271,7 +298,8 @@ public Optional<Bytes> getFlatStorageValueByStorageSlotKey(
     // get natural key from account hash and slot key
     byte[] naturalKey = calculateNaturalSlotKey(accountHash, storageSlotKey.getSlotHash());
     // keyNearest, use MAX_BLOCK_SUFFIX in the absence of a block context:
-    Bytes keyNearest = calculateArchiveKeyWithMaxSuffix(context, naturalKey);
+    Bytes keyNearest =
+        calculateArchiveKeyWithMaxSuffix(getStateArchiveContext(storage).get(), naturalKey);
 
     // Find the nearest storage for this address, slot key hash, and block context
     Optional<SegmentedKeyValueStorage.NearestKeyValue> nearestStorage =
@@ -321,6 +349,7 @@ public Optional<Bytes> getFlatStorageValueByStorageSlotKey(
    */
   @Override
   public void putFlatAccountStorageValueByStorageSlotHash(
+      final SegmentedKeyValueStorage storage,
       final SegmentedKeyValueStorageTransaction transaction,
       final Hash accountHash,
       final Hash slotHash,
@@ -329,7 +358,8 @@ public void putFlatAccountStorageValueByStorageSlotHash(
     // get natural key from account hash and slot key
     byte[] naturalKey = calculateNaturalSlotKey(accountHash, slotHash);
     // keyNearest, use MIN_BLOCK_SUFFIX in the absence of a block context:
-    byte[] keyNearest = calculateArchiveKeyWithMinSuffix(context, naturalKey);
+    byte[] keyNearest =
+        calculateArchiveKeyWithMinSuffix(getStateArchiveContext(storage).get(), naturalKey);
 
     transaction.put(ACCOUNT_STORAGE_STORAGE, keyNearest, storageValue.toArrayUnsafe());
   }
@@ -339,14 +369,16 @@ public void putFlatAccountStorageValueByStorageSlotHash(
    */
   @Override
   public void removeFlatAccountStorageValueByStorageSlotHash(
+      final SegmentedKeyValueStorage storage,
       final SegmentedKeyValueStorageTransaction transaction,
       final Hash accountHash,
       final Hash slotHash) {
 
     // get natural key from account hash and slot key
     byte[] naturalKey = calculateNaturalSlotKey(accountHash, slotHash);
     // insert a key suffixed with block context, with 'deleted account' value
-    byte[] keySuffixed = calculateArchiveKeyWithMinSuffix(context, naturalKey);
+    byte[] keySuffixed =
+        calculateArchiveKeyWithMinSuffix(getStateArchiveContext(storage).get(), naturalKey);
 
     transaction.put(ACCOUNT_STORAGE_STORAGE, keySuffixed, DELETED_STORAGE_VALUE);
   }
@@ -391,14 +423,4 @@ public static byte[] calculateArchiveKeyWithSuffix(
                   return orElseSuffix;
                 }));
   }
-
-  @Override
-  public void updateBlockContext(final Long blockNumber) {
-    context.setBlockNumber(blockNumber);
-  }
-
-  @Override
-  public void updateBlockContext(final BlockHeader blockHeader) {
-    context.setBlockNumber(blockHeader.getNumber());
-  }
 }
diff --git a/...rg/hyperledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiFlatDbStrategy.java b/...rg/hyperledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiFlatDbStrategy.java
@@ -71,6 +71,7 @@ public abstract Optional<Bytes> getFlatStorageValueByStorageSlotKey(
 
   @Override
   public void putFlatAccount(
+      final SegmentedKeyValueStorage storage,
       final SegmentedKeyValueStorageTransaction transaction,
       final Hash accountHash,
       final Bytes accountValue) {
@@ -79,12 +80,15 @@ public void putFlatAccount(
 
   @Override
   public void removeFlatAccount(
-      final SegmentedKeyValueStorageTransaction transaction, final Hash accountHash) {
+      final SegmentedKeyValueStorage storage,
+      final SegmentedKeyValueStorageTransaction transaction,
+      final Hash accountHash) {
     transaction.remove(ACCOUNT_INFO_STATE, accountHash.toArrayUnsafe());
   }
 
   @Override
   public void putFlatAccountStorageValueByStorageSlotHash(
+      final SegmentedKeyValueStorage storage,
       final SegmentedKeyValueStorageTransaction transaction,
       final Hash accountHash,
       final Hash slotHash,
@@ -97,6 +101,7 @@ public void putFlatAccountStorageValueByStorageSlotHash(
 
   @Override
   public void removeFlatAccountStorageValueByStorageSlotHash(
+      final SegmentedKeyValueStorage storage,
       final SegmentedKeyValueStorageTransaction transaction,
       final Hash accountHash,
       final Hash slotHash) {

diff --git a/...ledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiFlatDbStrategyProvider.java b/...ledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiFlatDbStrategyProvider.java
@@ -16,7 +16,6 @@
 
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.TRIE_BRANCH_STORAGE;
 
-import org.hyperledger.besu.ethereum.trie.diffbased.common.BonsaiContext;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.storage.flat.CodeStorageStrategy;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.storage.flat.FlatDbStrategy;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.storage.flat.FlatDbStrategyProvider;
@@ -91,8 +90,7 @@ protected FlatDbStrategy createFlatDbStrategy(
     if (flatDbMode == FlatDbMode.FULL) {
       return new BonsaiFullFlatDbStrategy(metricsSystem, codeStorageStrategy);
     } else if (flatDbMode == FlatDbMode.ARCHIVE) {
-      return new BonsaiArchiveFlatDbStrategy(
-          new BonsaiContext(), metricsSystem, codeStorageStrategy);
+      return new BonsaiArchiveFlatDbStrategy(metricsSystem, codeStorageStrategy);
     } else {
       return new BonsaiPartialFlatDbStrategy(metricsSystem, codeStorageStrategy);
     }

diff --git a/...yperledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiFullFlatDbStrategy.java b/...yperledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiFullFlatDbStrategy.java
@@ -19,7 +19,6 @@
 
 import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.datatypes.StorageSlotKey;
-import org.hyperledger.besu.ethereum.core.BlockHeader;
 import org.hyperledger.besu.ethereum.trie.NodeLoader;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.storage.flat.CodeStorageStrategy;
 import org.hyperledger.besu.metrics.BesuMetricCategory;
@@ -101,14 +100,4 @@ public void resetOnResync(final SegmentedKeyValueStorage storage) {
     // NOOP
     // not need to reset anything in full mode
   }
-
-  @Override
-  public void updateBlockContext(final Long blockNumber) {
-    // default no-op for strategies that do not care about bonsai context
-  }
-
-  @Override
-  public void updateBlockContext(final BlockHeader blockHeader) {
-    // default no-op for strategies that do not care about bonsai context
-  }
 }
diff --git a/...rledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiPartialFlatDbStrategy.java b/...rledger/besu/ethereum/trie/diffbased/bonsai/storage/flat/BonsaiPartialFlatDbStrategy.java
@@ -19,7 +19,6 @@
 
 import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.datatypes.StorageSlotKey;
-import org.hyperledger.besu.ethereum.core.BlockHeader;
 import org.hyperledger.besu.ethereum.trie.NodeLoader;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.storage.flat.CodeStorageStrategy;
 import org.hyperledger.besu.ethereum.trie.patricia.StoredMerklePatriciaTrie;
@@ -146,6 +145,7 @@ public Optional<Bytes> getFlatStorageValueByStorageSlotKey(
     return response;
   }
 
+  /*
   @Override
   public void updateBlockContext(final Long blockNumber) {
     // Not implemented
@@ -154,5 +154,5 @@ public void updateBlockContext(final Long blockNumber) {
   @Override
   public void updateBlockContext(final BlockHeader blockHeader) {
     // Not implemented
-  }
+  }*/
 }
diff --git a/.../java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/worldview/BonsaiWorldState.java b/.../java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/worldview/BonsaiWorldState.java
@@ -384,7 +384,10 @@ public Hash frontierRootHash() {
     return calculateRootHash(
         Optional.of(
             new BonsaiWorldStateKeyValueStorage.Updater(
-                noOpSegmentedTx, noOpTx, worldStateKeyValueStorage.getFlatDbStrategy())),
+                noOpSegmentedTx,
+                noOpTx,
+                worldStateKeyValueStorage.getFlatDbStrategy(),
+                worldStateKeyValueStorage.getComposedWorldStateStorage())),
         accumulator.copy());
   }