Fix transaction pool issue (#4964)

* fix transaction pool issue
* add block replay
* support in-memory snapshots

Signed-off-by: Karim TAAM <karim.t2am@gmail.com>
Signed-off-by: garyschulte <garyschulte@gmail.com>

besu/src/test/java/org/hyperledger/besu/services/BesuEventsImplTest.java

@@ -132,6 +132,7 @@ public class BesuEventsImplTest {
         .thenReturn(ValidationResult.valid());
     when(mockTransactionValidator.validateForSender(any(), any(), any()))
         .thenReturn(ValidationResult.valid());
+    when(mockWorldState.copy()).thenReturn(mockWorldState);
     when(mockWorldStateArchive.getMutable(any(), any(), anyBoolean()))
         .thenReturn(Optional.of(mockWorldState));
 

ethereum/api/src/main/java/org/hyperledger/besu/ethereum/api/jsonrpc/internal/processor/BlockReplay.java

@@ -145,15 +145,22 @@ public class BlockReplay {
     if (previous == null) {
       return Optional.empty();
     }
-    try (final MutableWorldState mutableWorldState =
+    try (final var worldState =
         worldStateArchive
-            .getMutable(previous.getStateRoot(), previous.getHash(), false)
+            .getMutable(previous.getStateRoot(), previous.getBlockHash(), false)
+            .map(
+                ws -> {
+                  if (!ws.isPersistable()) {
+                    return ws.copy();
+                  }
+                  return ws;
+                })
             .orElseThrow(
                 () ->
                     new IllegalArgumentException(
                         "Missing worldstate for stateroot "
                             + previous.getStateRoot().toShortHexString()))) {
-      return action.perform(body, header, blockchain, mutableWorldState, transactionProcessor);
+      return action.perform(body, header, blockchain, worldState, transactionProcessor);
     } catch (Exception ex) {
       return Optional.empty();
     }

ethereum/api/src/test/java/org/hyperledger/besu/ethereum/api/jsonrpc/internal/processor/TransactionTracerTest.java

@@ -118,6 +118,7 @@ public class TransactionTracerTest {
     when(protocolSpec.getMiningBeneficiaryCalculator()).thenReturn(BlockHeader::getCoinbase);
     when(blockchain.getChainHeadHeader()).thenReturn(blockHeader);
     when(protocolSpec.getBadBlocksManager()).thenReturn(new BadBlockManager());
+    when(mutableWorldState.copy()).thenReturn(mutableWorldState);
   }
 
   @Test

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

@@ -180,10 +180,10 @@ public class BonsaiSnapshotWorldStateKeyValueStorage extends BonsaiWorldStateKey
 
   public static class SnapshotUpdater implements BonsaiWorldStateKeyValueStorage.BonsaiUpdater {
 
-    private final SnappedKeyValueStorage accountStorage;
-    private final SnappedKeyValueStorage codeStorage;
-    private final SnappedKeyValueStorage storageStorage;
-    private final SnappedKeyValueStorage trieBranchStorage;
+    private final KeyValueStorageTransaction accountStorageTransaction;
+    private final KeyValueStorageTransaction codeStorageTransaction;
+    private final KeyValueStorageTransaction storageStorageTransaction;
+    private final KeyValueStorageTransaction trieBranchStorageTransaction;
     private final KeyValueStorageTransaction trieLogStorageTransaction;
 
     public SnapshotUpdater(
@@ -192,16 +192,16 @@ public class BonsaiSnapshotWorldStateKeyValueStorage extends BonsaiWorldStateKey
         final SnappedKeyValueStorage storageStorage,
         final SnappedKeyValueStorage trieBranchStorage,
         final KeyValueStorage trieLogStorage) {
-      this.accountStorage = accountStorage;
-      this.codeStorage = codeStorage;
-      this.storageStorage = storageStorage;
-      this.trieBranchStorage = trieBranchStorage;
+      this.accountStorageTransaction = accountStorage.getSnapshotTransaction();
+      this.codeStorageTransaction = codeStorage.getSnapshotTransaction();
+      this.storageStorageTransaction = storageStorage.getSnapshotTransaction();
+      this.trieBranchStorageTransaction = trieBranchStorage.getSnapshotTransaction();
       this.trieLogStorageTransaction = trieLogStorage.startTransaction();
     }
 
     @Override
     public BonsaiUpdater removeCode(final Hash accountHash) {
-      codeStorage.getSnapshotTransaction().remove(accountHash.toArrayUnsafe());
+      codeStorageTransaction.remove(accountHash.toArrayUnsafe());
       return this;
     }
 
@@ -212,13 +212,13 @@ public class BonsaiSnapshotWorldStateKeyValueStorage extends BonsaiWorldStateKey
         // Don't save empty values
         return this;
       }
-      codeStorage.getSnapshotTransaction().put(accountHash.toArrayUnsafe(), code.toArrayUnsafe());
+      codeStorageTransaction.put(accountHash.toArrayUnsafe(), code.toArrayUnsafe());
       return this;
     }
 
     @Override
     public BonsaiUpdater removeAccountInfoState(final Hash accountHash) {
-      accountStorage.getSnapshotTransaction().remove(accountHash.toArrayUnsafe());
+      accountStorageTransaction.remove(accountHash.toArrayUnsafe());
       return this;
     }
 
@@ -228,31 +228,26 @@ public class BonsaiSnapshotWorldStateKeyValueStorage extends BonsaiWorldStateKey
         // Don't save empty values
         return this;
       }
-      accountStorage
-          .getSnapshotTransaction()
-          .put(accountHash.toArrayUnsafe(), accountValue.toArrayUnsafe());
+      accountStorageTransaction.put(accountHash.toArrayUnsafe(), accountValue.toArrayUnsafe());
       return this;
     }
 
     @Override
     public BonsaiUpdater putStorageValueBySlotHash(
         final Hash accountHash, final Hash slotHash, final Bytes storage) {
-      storageStorage
-          .getSnapshotTransaction()
-          .put(Bytes.concatenate(accountHash, slotHash).toArrayUnsafe(), storage.toArrayUnsafe());
+      storageStorageTransaction.put(
+          Bytes.concatenate(accountHash, slotHash).toArrayUnsafe(), storage.toArrayUnsafe());
       return this;
     }
 
     @Override
     public void removeStorageValueBySlotHash(final Hash accountHash, final Hash slotHash) {
-      storageStorage
-          .getSnapshotTransaction()
-          .remove(Bytes.concatenate(accountHash, slotHash).toArrayUnsafe());
+      storageStorageTransaction.remove(Bytes.concatenate(accountHash, slotHash).toArrayUnsafe());
     }
 
     @Override
     public KeyValueStorageTransaction getTrieBranchStorageTransaction() {
-      return trieBranchStorage.getSnapshotTransaction();
+      return trieBranchStorageTransaction;
     }
 
     @Override
@@ -263,13 +258,9 @@ public class BonsaiSnapshotWorldStateKeyValueStorage extends BonsaiWorldStateKey
     @Override
     public WorldStateStorage.Updater saveWorldState(
         final Bytes blockHash, final Bytes32 nodeHash, final Bytes node) {
-      trieBranchStorage
-          .getSnapshotTransaction()
-          .put(Bytes.EMPTY.toArrayUnsafe(), node.toArrayUnsafe());
-      trieBranchStorage.getSnapshotTransaction().put(WORLD_ROOT_HASH_KEY, nodeHash.toArrayUnsafe());
-      trieBranchStorage
-          .getSnapshotTransaction()
-          .put(WORLD_BLOCK_HASH_KEY, blockHash.toArrayUnsafe());
+      trieBranchStorageTransaction.put(Bytes.EMPTY.toArrayUnsafe(), node.toArrayUnsafe());
+      trieBranchStorageTransaction.put(WORLD_ROOT_HASH_KEY, nodeHash.toArrayUnsafe());
+      trieBranchStorageTransaction.put(WORLD_BLOCK_HASH_KEY, blockHash.toArrayUnsafe());
       return this;
     }
 
@@ -280,16 +271,14 @@ public class BonsaiSnapshotWorldStateKeyValueStorage extends BonsaiWorldStateKey
         // Don't save empty nodes
         return this;
       }
-      trieBranchStorage
-          .getSnapshotTransaction()
-          .put(location.toArrayUnsafe(), node.toArrayUnsafe());
+      trieBranchStorageTransaction.put(location.toArrayUnsafe(), node.toArrayUnsafe());
       return this;
     }
 
     @Override
     public WorldStateStorage.Updater removeAccountStateTrieNode(
         final Bytes location, final Bytes32 nodeHash) {
-      trieBranchStorage.getSnapshotTransaction().remove(location.toArrayUnsafe());
+      trieBranchStorageTransaction.remove(location.toArrayUnsafe());
       return this;
     }
 
@@ -300,15 +289,17 @@ public class BonsaiSnapshotWorldStateKeyValueStorage extends BonsaiWorldStateKey
         // Don't save empty nodes
         return this;
       }
-      trieBranchStorage
-          .getSnapshotTransaction()
-          .put(Bytes.concatenate(accountHash, location).toArrayUnsafe(), node.toArrayUnsafe());
+      trieBranchStorageTransaction.put(
+          Bytes.concatenate(accountHash, location).toArrayUnsafe(), node.toArrayUnsafe());
       return this;
     }
 
     @Override
     public void commit() {
-      // only commit the trielog layer transaction, leave the snapshot transactions open:
+      accountStorageTransaction.commit();
+      codeStorageTransaction.commit();
+      storageStorageTransaction.commit();
+      trieBranchStorageTransaction.commit();
       trieLogStorageTransaction.commit();
     }
 

ethereum/eth/src/main/java/org/hyperledger/besu/ethereum/eth/transactions/TransactionPool.java

@@ -323,10 +323,18 @@ public class TransactionPool implements BlockAddedObserver {
           "EIP-1559 transaction are not allowed yet");
     }
 
-    try (var worldState =
+    try (final var worldState =
         protocolContext
             .getWorldStateArchive()
-            .getMutable(chainHeadBlockHeader.getStateRoot(), chainHeadBlockHeader.getHash(), false)
+            .getMutable(
+                chainHeadBlockHeader.getStateRoot(), chainHeadBlockHeader.getBlockHash(), false)
+            .map(
+                ws -> {
+                  if (!ws.isPersistable()) {
+                    return ws.copy();
+                  }
+                  return ws;
+                })
             .orElseThrow()) {
       final Account senderAccount = worldState.get(transaction.getSender());
       return new ValidationResultAndAccount(

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

@@ -164,8 +164,7 @@ public class RocksDBSnapshotTransaction implements KeyValueStorageTransaction, A
 
   @Override
   public void commit() throws StorageException {
-    // no-op or throw?
-    throw new UnsupportedOperationException("RocksDBSnapshotTransaction does not support commit");
+    // no-op
   }
 
   @Override

services/kvstore/src/main/java/org/hyperledger/besu/services/kvstore/InMemoryKeyValueStorage.java

@@ -19,6 +19,8 @@ import static java.util.stream.Collectors.toUnmodifiableSet;
 import org.hyperledger.besu.plugin.services.exception.StorageException;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorage;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorageTransaction;
+import org.hyperledger.besu.plugin.services.storage.SnappableKeyValueStorage;
+import org.hyperledger.besu.plugin.services.storage.SnappedKeyValueStorage;
 
 import java.io.PrintStream;
 import java.util.HashMap;
@@ -37,7 +39,8 @@ import org.apache.commons.lang3.tuple.Pair;
 import org.apache.tuweni.bytes.Bytes;
 
 /** The In memory key value storage. */
-public class InMemoryKeyValueStorage implements KeyValueStorage {
+public class InMemoryKeyValueStorage
+    implements SnappedKeyValueStorage, SnappableKeyValueStorage, KeyValueStorage {
 
   private final Map<Bytes, byte[]> hashValueStore;
   private final ReadWriteLock rwLock = new ReentrantReadWriteLock();
@@ -154,6 +157,21 @@ public class InMemoryKeyValueStorage implements KeyValueStorage {
     return Set.copyOf(hashValueStore.keySet());
   }
 
+  @Override
+  public SnappedKeyValueStorage takeSnapshot() {
+    return new InMemoryKeyValueStorage(new HashMap<>(hashValueStore));
+  }
+
+  @Override
+  public KeyValueStorageTransaction getSnapshotTransaction() {
+    return startTransaction();
+  }
+
+  @Override
+  public SnappedKeyValueStorage cloneFromSnapshot() {
+    return takeSnapshot();
+  }
+
   private class InMemoryTransaction implements KeyValueStorageTransaction {
 
     private Map<Bytes, byte[]> updatedValues = new HashMap<>();