make RocksDBSnapshotTransaction.copy() *also* copy the transaction st…

…ate, not just restart a new transaction from the original snapshot

Signed-off-by: garyschulte <garyschulte@gmail.com>

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

@@ -75,7 +75,6 @@ public BonsaiWorldStateArchive getArchive() {
 
   @Override
   public MutableWorldState copy() {
-    // TODO: consider returning a snapshot rather than a copy here.
     BonsaiInMemoryWorldStateKeyValueStorage bonsaiInMemoryWorldStateKeyValueStorage =
         new BonsaiInMemoryWorldStateKeyValueStorage(
             worldStateStorage.accountStorage,

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

@@ -34,6 +34,7 @@ public class BonsaiSnapshotPersistedWorldState extends BonsaiPersistedWorldState
   private final SnappedKeyValueStorage codeSnap;
   private final SnappedKeyValueStorage storageSnap;
   private final SnappedKeyValueStorage trieBranchSnap;
+  private boolean isPersisted = false;
 
   private BonsaiSnapshotPersistedWorldState(
       final BonsaiWorldStateArchive archive,
@@ -60,17 +61,11 @@ public static BonsaiSnapshotPersistedWorldState create(
 
   @Override
   public Hash rootHash() {
-    return rootHash(updater.copy());
-  }
-
-  public Hash rootHash(final BonsaiWorldStateUpdater localUpdater) {
-    final BonsaiWorldStateKeyValueStorage.BonsaiUpdater updater = worldStateStorage.updater();
-    try {
-      final Hash calculatedRootHash = calculateRootHash(updater, localUpdater);
-      return Hash.wrap(calculatedRootHash);
-    } finally {
-      updater.rollback();
+    if (!isPersisted) {
+      this.worldStateRootHash = calculateRootHash(worldStateStorage.updater(), updater);
+      isPersisted = true;
     }
+    return this.worldStateRootHash;
   }
 
   @Override

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

@@ -61,7 +61,7 @@ public static BonsaiSnapshotWorldState create(
   @Override
   public void persist(final BlockHeader blockHeader) {
     super.persist(blockHeader);
-    // persist roothash to snapshot tx
+    // persist roothash to trie branch snapshot tx
     trieBranchSnap
         .getSnapshotTransaction()
         .put(
@@ -72,6 +72,9 @@ public void persist(final BlockHeader blockHeader) {
   @Override
   public MutableWorldState copy() {
     // return a clone-based copy of worldstate storage
+    // TODO: this is currently broken.  We need to clone the in-memory updater in addition to
+    // the storage transactions in order to get a true copy.
+
     return new BonsaiSnapshotWorldState(
         archive,
         new BonsaiSnapshotWorldStateKeyValueStorage(

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

@@ -140,8 +140,8 @@ public Optional<MutableWorldState> getMutableSnapshot(final Hash blockHash) {
     // TODO: decide whether we want to use BonsaiSnapshotWorldState or
     // BonsaiSnapshotPersistedWorldState
     return rollMutableStateToBlockHash(
-            BonsaiSnapshotWorldState.create(this, worldStateStorage), blockHash)
-        //            BonsaiSnapshotPersistedWorldState.create(this, worldStateStorage), blockHash)
+            //            BonsaiSnapshotWorldState.create(this, worldStateStorage), blockHash)
+            BonsaiSnapshotPersistedWorldState.create(this, worldStateStorage), blockHash)
         .map(SnapshotMutableWorldState.class::cast);
   }
 

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

@@ -24,6 +24,7 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 
 import org.apache.tuweni.bytes.Bytes32;
 import org.slf4j.Logger;
@@ -70,6 +71,14 @@ public void addCachedLayer(
                     new CachedSnapshotWorldState(snapshot, 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.empty();
+  }
+
   @Override
   public void updateCachedLayers(final Hash blockParentHash, final Hash blockHash) {
     // no-op.  Snapshots are independent and do not need to update 'next' worldstates

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

@@ -62,6 +62,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -246,7 +247,38 @@ public void testSnapshotCloneIsolation() {
   @Test
   public void assertSnapshotDoesNotClose() {
     // TODO: add unit test to assert snapshot does not close on clone if parent tx is closed
+    Address testAddress = Address.fromHexString("0xdeadbeef");
+
+    // create a snapshot worldstate, and then clone it:
+    var isolated = archive.getMutableSnapshot(genesisState.getBlock().getHash()).get();
+
+    // 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(res.isSuccessful()).isTrue();
+    Consumer<MutableWorldState> checkIsolatedState =
+        (ws) -> {
+          assertThat(ws.rootHash()).isEqualTo(firstBlock.getHeader().getStateRoot());
+          assertThat(ws.get(testAddress)).isNotNull();
+          assertThat(ws.get(testAddress).getBalance())
+              .isEqualTo(Wei.of(1_000_000_000_000_000_000L));
+        };
+    checkIsolatedState.accept(isolated);
+
+    var isolatedClone = isolated.copy();
+    checkIsolatedState.accept(isolatedClone);
+
+    try {
+      // close the first snapshot worldstate.  The second worldstate should still be able to read
+      // through its snapshot
+      isolated.close();
+    } catch (Exception ex) {
+      // meh
+    }
 
+    // copy of closed isolated worldstate should still pass check
+    checkIsolatedState.accept(isolatedClone);
   }
 
   @Test

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

@@ -22,14 +22,15 @@
 
 /**
  * Wraps and reference counts a Snapshot object from an OptimisticTransactionDB such that it can be
- * used as the basis of multiple SnapshotTransactions, and released once it is no longer in use.
+ * used as the basis of multiple RocksDBSnapshotTransaction's, and released once it is no longer in
+ * use.
  */
-public class RocksDBSnapshot {
+class RocksDBSnapshot {
   private final OptimisticTransactionDB db;
   private final Snapshot dbSnapshot;
   private final AtomicInteger usages = new AtomicInteger(0);
 
-  public RocksDBSnapshot(final OptimisticTransactionDB db) {
+  RocksDBSnapshot(final OptimisticTransactionDB db) {
     this.db = db;
     this.dbSnapshot = db.getSnapshot();
   }

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

@@ -63,14 +63,16 @@ private RocksDBSnapshotTransaction(
       final OptimisticTransactionDB db,
       final ColumnFamilyHandle columnFamilyHandle,
       final RocksDBMetrics metrics,
-      final RocksDBSnapshot snapshot) {
+      final RocksDBSnapshot snapshot,
+      final Transaction snapTx,
+      final ReadOptions readOptions) {
     this.metrics = metrics;
     this.db = db;
     this.columnFamilyHandle = columnFamilyHandle;
     this.snapshot = snapshot;
     this.writeOptions = new WriteOptions();
-    this.snapTx = db.beginTransaction(writeOptions);
-    this.readOptions = new ReadOptions().setSnapshot(snapshot.markAndUseSnapshot());
+    this.readOptions = readOptions;
+    this.snapTx = snapTx;
   }
 
   public Optional<byte[]> get(final byte[] key) {
@@ -136,7 +138,17 @@ public void rollback() {
   }
 
   public RocksDBSnapshotTransaction copy() {
-    return new RocksDBSnapshotTransaction(db, columnFamilyHandle, metrics, snapshot);
+    try {
+      var copyReadOptions = new ReadOptions().setSnapshot(snapshot.markAndUseSnapshot());
+      var copySnapTx = db.beginTransaction(writeOptions);
+      copySnapTx.rebuildFromWriteBatch(snapTx.getWriteBatch().getWriteBatch());
+      return new RocksDBSnapshotTransaction(
+          db, columnFamilyHandle, metrics, snapshot, copySnapTx, copyReadOptions);
+    } catch (Exception ex) {
+      LOG.error("Failed to copy snapshot transaction", ex);
+      snapshot.unMarkSnapshot();
+      throw new StorageException(ex);
+    }
   }
 
   @Override