Fix for block unavailability during SnapshotTrieLog caching (#4666)

* supplier workaround for Cached snapshots to defer snapshots until the block is added to the chain
* handle cache update when worldstate is fast-syncing
* add additional coverage for SnapshotTrieLogManager caching

Signed-off-by: garyschulte <garyschulte@gmail.com>
Co-authored-by: Sally MacFarlane <macfarla.github@gmail.com>

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

@@ -131,7 +131,8 @@ public abstract class AbstractTrieLogManager<T extends CachedWorldState> impleme
   @Override
   public Optional<MutableWorldState> getBonsaiCachedWorldState(final Hash blockHash) {
     if (cachedWorldStatesByHash.containsKey(blockHash)) {
-      return Optional.of(cachedWorldStatesByHash.get(blockHash).getMutableWorldState());
+      return Optional.ofNullable(cachedWorldStatesByHash.get(blockHash))
+          .map(T::getMutableWorldState);
     }
     return Optional.empty();
   }

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

@@ -25,7 +25,9 @@ import org.hyperledger.besu.ethereum.core.MutableWorldState;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Supplier;
 
+import com.google.common.base.Suppliers;
 import org.apache.tuweni.bytes.Bytes32;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -61,38 +63,45 @@ public class SnapshotTrieLogManager
         "adding snapshot world state for block {}, state root hash {}",
         blockHeader::toLogString,
         worldStateRootHash::toHexString);
-    worldStateArchive
-        .getMutableSnapshot(worldStateRootHash)
-        .map(BonsaiSnapshotWorldState.class::cast)
-        .ifPresent(
-            snapshot ->
-                cachedWorldStatesByHash.put(
-                    blockHeader.getHash(),
-                    new CachedSnapshotWorldState(snapshot, trieLog, blockHeader.getNumber())));
+    cachedWorldStatesByHash.put(
+        blockHeader.getHash(),
+        new CachedSnapshotWorldState(
+            () ->
+                worldStateArchive
+                    .getMutableSnapshot(blockHeader.getHash())
+                    .map(BonsaiSnapshotWorldState.class::cast)
+                    .orElse(null),
+            trieLog,
+            blockHeader.getNumber()));
   }
 
   @Override
   public Optional<MutableWorldState> getBonsaiCachedWorldState(final Hash blockHash) {
     if (cachedWorldStatesByHash.containsKey(blockHash)) {
-      return Optional.of(cachedWorldStatesByHash.get(blockHash).getMutableWorldState().copy());
+      return Optional.ofNullable(cachedWorldStatesByHash.get(blockHash).getMutableWorldState())
+          .map(MutableWorldState::copy);
     }
     return Optional.empty();
   }
 
   @Override
   public void updateCachedLayers(final Hash blockParentHash, final Hash blockHash) {
-    // no-op.  Snapshots are independent and do not need to update 'next' worldstates
+    // fetch the snapshot supplier as soon as its block has been added:
+    Optional.ofNullable(cachedWorldStatesByHash.get(blockHash))
+        .ifPresent(CachedSnapshotWorldState::getMutableWorldState);
   }
 
   public static class CachedSnapshotWorldState implements CachedWorldState {
 
-    final BonsaiSnapshotWorldState snapshot;
+    final Supplier<BonsaiSnapshotWorldState> snapshot;
     final TrieLogLayer trieLog;
     final long height;
 
     public CachedSnapshotWorldState(
-        final BonsaiSnapshotWorldState snapshot, final TrieLogLayer trieLog, final long height) {
-      this.snapshot = snapshot;
+        final Supplier<BonsaiSnapshotWorldState> snapshotSupplier,
+        final TrieLogLayer trieLog,
+        final long height) {
+      this.snapshot = Suppliers.memoize(snapshotSupplier::get);
       this.trieLog = trieLog;
       this.height = height;
     }
@@ -109,7 +118,7 @@ public class SnapshotTrieLogManager
 
     @Override
     public MutableWorldState getMutableWorldState() {
-      return snapshot;
+      return snapshot.get();
     }
   }
 }

ethereum/core/src/test/java/org/hyperledger/besu/ethereum/bonsai/BonsaiSnapshotIsolationTests.java

@@ -134,6 +134,11 @@ public class BonsaiSnapshotIsolationTests {
     assertThat(res.isSuccessful()).isTrue();
     assertThat(res2.isSuccessful()).isTrue();
 
+    assertThat(archive.getTrieLogManager().getBonsaiCachedWorldState(firstBlock.getHash()))
+        .isNotEmpty();
+    assertThat(archive.getTrieLogManager().getBonsaiCachedWorldState(secondBlock.getHash()))
+        .isNotEmpty();
+
     assertThat(archive.getMutable().get(testAddress)).isNotNull();
     assertThat(archive.getMutable().get(testAddress).getBalance())
         .isEqualTo(Wei.of(2_000_000_000_000_000_000L));
@@ -164,6 +169,9 @@ public class BonsaiSnapshotIsolationTests {
     var firstBlock = forTransactions(List.of(burnTransaction(sender1, 0L, testAddress)));
     var res = executeBlock(isolated.get(), firstBlock);
 
+    assertThat(archive.getTrieLogManager().getBonsaiCachedWorldState(firstBlock.getHash()))
+        .isNotEmpty();
+
     assertThat(res.isSuccessful()).isTrue();
     assertThat(isolated.get().get(testAddress)).isNotNull();
     assertThat(isolated.get().get(testAddress).getBalance())
@@ -249,6 +257,9 @@ public class BonsaiSnapshotIsolationTests {
     assertThat(firstBlockTrieLog).isNotEmpty();
     assertThat(firstBlockTrieLog.get().getAccount(testAddress)).isNotEmpty();
     assertThat(firstBlockTrieLog.get().getAccount(altTestAddress)).isEmpty();
+    assertThat(archive.getTrieLogManager().getBonsaiCachedWorldState(firstBlock.getHash()))
+        .isNotEmpty();
+
     var cloneForkTrieLog = archive.getTrieLogManager().getTrieLogLayer(cloneForkBlock.getHash());
     assertThat(cloneForkTrieLog.get().getAccount(testAddress)).isEmpty();
     assertThat(cloneForkTrieLog.get().getAccount(altTestAddress)).isNotEmpty();
@@ -271,6 +282,8 @@ public class BonsaiSnapshotIsolationTests {
     // execute a block with a single transaction on the first snapshot:
     var firstBlock = forTransactions(List.of(burnTransaction(sender1, 0L, testAddress)));
     var res = executeBlock(isolated, firstBlock);
+    assertThat(archive.getTrieLogManager().getBonsaiCachedWorldState(firstBlock.getHash()))
+        .isNotEmpty();
 
     assertThat(res.isSuccessful()).isTrue();
     Consumer<MutableWorldState> checkIsolatedState =

plugins/rocksdb/src/main/java/org/hyperledger/besu/plugin/services/storage/rocksdb/segmented/RocksDBSnapshot.java

@@ -47,13 +47,4 @@ class RocksDBSnapshot {
       dbSnapshot.close();
     }
   }
-
-  // TODO:  this is a stopgap measure, revisit with https://github.com/hyperledger/besu/issues/4641
-  @Override
-  protected void finalize() {
-    if (usages.decrementAndGet() > 0) {
-      db.releaseSnapshot(dbSnapshot);
-      dbSnapshot.close();
-    }
-  }
 }

plugins/rocksdb/src/main/java/org/hyperledger/besu/plugin/services/storage/rocksdb/segmented/RocksDBSnapshotTransaction.java

@@ -15,17 +15,13 @@
  */
 package org.hyperledger.besu.plugin.services.storage.rocksdb.segmented;
 
-import static org.hyperledger.besu.util.Slf4jLambdaHelper.debugLambda;
-
 import org.hyperledger.besu.plugin.services.exception.StorageException;
 import org.hyperledger.besu.plugin.services.metrics.OperationTimer;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorageTransaction;
 import org.hyperledger.besu.plugin.services.storage.rocksdb.RocksDBMetrics;
 import org.hyperledger.besu.plugin.services.storage.rocksdb.RocksDbIterator;
 
-import java.util.Arrays;
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Stream;
 
 import org.apache.commons.lang3.tuple.Pair;
@@ -50,9 +46,6 @@ public class RocksDBSnapshotTransaction implements KeyValueStorageTransaction, A
   private final RocksDBSnapshot snapshot;
   private final WriteOptions writeOptions;
   private final ReadOptions readOptions;
-  // TODO:  this is a stopgap measure, revisit with https://github.com/hyperledger/besu/issues/4641
-  private final AtomicBoolean isClosed = new AtomicBoolean(false);
-  private final StackTraceElement[] stackTrace;
 
   RocksDBSnapshotTransaction(
       final OptimisticTransactionDB db,
@@ -65,7 +58,6 @@ public class RocksDBSnapshotTransaction implements KeyValueStorageTransaction, A
     this.writeOptions = new WriteOptions();
     this.snapTx = db.beginTransaction(writeOptions);
     this.readOptions = new ReadOptions().setSnapshot(snapshot.markAndUseSnapshot());
-    this.stackTrace = Thread.currentThread().getStackTrace();
   }
 
   private RocksDBSnapshotTransaction(
@@ -82,7 +74,6 @@ public class RocksDBSnapshotTransaction implements KeyValueStorageTransaction, A
     this.writeOptions = new WriteOptions();
     this.readOptions = readOptions;
     this.snapTx = snapTx;
-    this.stackTrace = Thread.currentThread().getStackTrace();
   }
 
   public Optional<byte[]> get(final byte[] key) {
@@ -173,18 +164,5 @@ public class RocksDBSnapshotTransaction implements KeyValueStorageTransaction, A
     writeOptions.close();
     readOptions.close();
     snapshot.unMarkSnapshot();
-    isClosed.set(true);
-  }
-
-  // TODO:  this is a stopgap measure, revisit with https://github.com/hyperledger/besu/issues/4641
-  @Override
-  protected void finalize() {
-    if (!isClosed.get()) {
-      debugLambda(
-          LOG,
-          "RocksDBSnapshotTransaction was not closed.  This is a memory leak.  Stack trace: {}",
-          () -> Arrays.toString(stackTrace));
-      close();
-    }
   }
 }