feature: Add optional worldstate move flag to debug_setHead (#7821)

* add optional worldstate move to debug_setHead
* make state rolling occur incrementally so as not to overwhelm memory and resources

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

ethereum/api/src/main/java/org/hyperledger/besu/ethereum/api/jsonrpc/internal/methods/DebugSetHead.java

@@ -21,22 +21,41 @@ import org.hyperledger.besu.ethereum.ProtocolContext;
 import org.hyperledger.besu.ethereum.api.jsonrpc.RpcMethod;
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.JsonRpcRequestContext;
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.exception.InvalidJsonRpcParameters;
-import org.hyperledger.besu.ethereum.api.jsonrpc.internal.parameters.BlockParameter;
+import org.hyperledger.besu.ethereum.api.jsonrpc.internal.parameters.BlockParameterOrBlockHash;
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.parameters.JsonRpcParameter.JsonRpcParameterException;
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.response.JsonRpcErrorResponse;
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.response.JsonRpcSuccessResponse;
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.response.RpcErrorType;
 import org.hyperledger.besu.ethereum.api.query.BlockchainQueries;
+import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
+import org.hyperledger.besu.ethereum.core.BlockHeader;
+import org.hyperledger.besu.ethereum.trie.diffbased.common.DiffBasedWorldStateProvider;
 
 import java.util.Optional;
 
-public class DebugSetHead extends AbstractBlockParameterMethod {
+import graphql.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class DebugSetHead extends AbstractBlockParameterOrBlockHashMethod {
   private final ProtocolContext protocolContext;
+  private static final Logger LOG = LoggerFactory.getLogger(DebugSetHead.class);
+  private static final int DEFAULT_MAX_TRIE_LOGS_TO_ROLL_AT_ONCE = 512;
+
+  private final long maxTrieLogsToRollAtOnce;
 
   public DebugSetHead(final BlockchainQueries blockchain, final ProtocolContext protocolContext) {
-    super(blockchain);
+    this(blockchain, protocolContext, DEFAULT_MAX_TRIE_LOGS_TO_ROLL_AT_ONCE);
+  }
 
+  @VisibleForTesting
+  DebugSetHead(
+      final BlockchainQueries blockchain,
+      final ProtocolContext protocolContext,
+      final long maxTrieLogsToRollAtOnce) {
+    super(blockchain);
     this.protocolContext = protocolContext;
+    this.maxTrieLogsToRollAtOnce = Math.abs(maxTrieLogsToRollAtOnce);
   }
 
   @Override
@@ -45,26 +64,108 @@ public class DebugSetHead extends AbstractBlockParameterMethod {
   }
 
   @Override
-  protected BlockParameter blockParameter(final JsonRpcRequestContext request) {
+  protected BlockParameterOrBlockHash blockParameterOrBlockHash(
+      final JsonRpcRequestContext requestContext) {
     try {
-      return request.getRequiredParameter(0, BlockParameter.class);
+      return requestContext.getRequiredParameter(0, BlockParameterOrBlockHash.class);
     } catch (JsonRpcParameterException e) {
       throw new InvalidJsonRpcParameters(
-          "Invalid block parameter (index 0)", RpcErrorType.INVALID_BLOCK_PARAMS, e);
+          "Invalid block or block hash parameter (index 0)", RpcErrorType.INVALID_BLOCK_PARAMS, e);
     }
   }
 
   @Override
-  protected Object resultByBlockNumber(
+  protected Object resultByBlockHash(final JsonRpcRequestContext request, final Hash blockHash) {
-      final JsonRpcRequestContext request, final long blockNumber) {
+    var blockchainQueries = getBlockchainQueries();
-    final Optional<Hash> maybeBlockHash = getBlockchainQueries().getBlockHashByNumber(blockNumber);
+    var blockchain = protocolContext.getBlockchain();
+    Optional<BlockHeader> maybeBlockHeader = blockchainQueries.getBlockHeaderByHash(blockHash);
+    Optional<Boolean> maybeMoveWorldstate = shouldMoveWorldstate(request);
 
-    if (maybeBlockHash.isEmpty()) {
+    if (maybeBlockHeader.isEmpty()) {
       return new JsonRpcErrorResponse(request.getRequest().getId(), UNKNOWN_BLOCK);
     }
 
-    protocolContext.getBlockchain().rewindToBlock(maybeBlockHash.get());
+    // Optionally move the worldstate to the specified blockhash, if it is present in the chain
+    if (maybeMoveWorldstate.orElse(Boolean.FALSE)) {
+      var archive = blockchainQueries.getWorldStateArchive();
+
+      // Only DiffBasedWorldState's need to be moved:
+      if (archive instanceof DiffBasedWorldStateProvider diffBasedArchive) {
+        if (rollIncrementally(maybeBlockHeader.get(), blockchain, diffBasedArchive)) {
+          return JsonRpcSuccessResponse.SUCCESS_RESULT;
+        }
+      }
+    }
+
+    // If we are not rolling incrementally or if there was an error incrementally rolling,
+    // move the blockchain to the requested hash:
+    blockchain.rewindToBlock(maybeBlockHeader.get().getBlockHash());
 
     return JsonRpcSuccessResponse.SUCCESS_RESULT;
   }
+
+  private boolean rollIncrementally(
+      final BlockHeader target,
+      final MutableBlockchain blockchain,
+      final DiffBasedWorldStateProvider archive) {
+
+    try {
+      if (archive.isWorldStateAvailable(target.getStateRoot(), target.getBlockHash())) {
+        // WARNING, this can be dangerous for a DiffBasedWorldstate if a concurrent
+        //          process attempts to move or modify the head worldstate.
+        //          Ensure no block processing is occuring when using this feature.
+        //          No engine-api, block import, sync, mining or other rpc calls should be running.
+
+        Optional<BlockHeader> currentHead =
+            archive
+                .getWorldStateKeyValueStorage()
+                .getWorldStateBlockHash()
+                .flatMap(blockchain::getBlockHeader);
+
+        while (currentHead.isPresent()
+            && !target.getStateRoot().equals(currentHead.get().getStateRoot())) {
+          long delta = currentHead.get().getNumber() - target.getNumber();
+
+          if (maxTrieLogsToRollAtOnce < Math.abs(delta)) {
+            // do we need to move forward or backward?
+            long distanceToMove = (delta > 0) ? -maxTrieLogsToRollAtOnce : maxTrieLogsToRollAtOnce;
+
+            // Add distanceToMove to the current block number to get the interim target header
+            var interimHead =
+                blockchain.getBlockHeader(currentHead.get().getNumber() + distanceToMove);
+
+            interimHead.ifPresent(
+                it -> {
+                  blockchain.rewindToBlock(it.getBlockHash());
+                  archive.getMutable(it.getStateRoot(), it.getBlockHash());
+                  LOG.info("incrementally rolled worldstate to {}", it.toLogString());
+                });
+            currentHead = interimHead;
+
+          } else {
+            blockchain.rewindToBlock(target.getBlockHash());
+            archive.getMutable(target.getStateRoot(), target.getBlockHash());
+            currentHead = Optional.of(target);
+            LOG.info("finished rolling worldstate to {}", target.toLogString());
+          }
+        }
+      }
+
+      return true;
+    } catch (Exception ex) {
+      LOG.error("Failed to incrementally roll blockchain to " + target.toLogString(), ex);
+      return false;
+    }
+  }
+
+  private Optional<Boolean> shouldMoveWorldstate(final JsonRpcRequestContext request) {
+    try {
+      return request.getOptionalParameter(1, Boolean.class);
+    } catch (JsonRpcParameterException e) {
+      throw new InvalidJsonRpcParameters(
+          "Invalid should move worldstate boolean parameter (index 1)",
+          RpcErrorType.INVALID_PARAMS,
+          e);
+    }
+  }
 }

ethereum/api/src/test/java/org/hyperledger/besu/ethereum/api/jsonrpc/internal/methods/DebugSetHeadTest.java

@@ -0,0 +1,198 @@
+/*
+ * Copyright contributors to Besu.
+ *
+ * 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.api.jsonrpc.internal.methods;
+
+import static java.lang.Boolean.FALSE;
+import static java.lang.Boolean.TRUE;
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+
+import org.hyperledger.besu.crypto.Hash;
+import org.hyperledger.besu.ethereum.ProtocolContext;
+import org.hyperledger.besu.ethereum.api.jsonrpc.AbstractJsonRpcHttpServiceTest;
+import org.hyperledger.besu.ethereum.api.jsonrpc.internal.JsonRpcRequest;
+import org.hyperledger.besu.ethereum.api.jsonrpc.internal.JsonRpcRequestContext;
+import org.hyperledger.besu.ethereum.api.jsonrpc.internal.parameters.BlockParameterOrBlockHash;
+import org.hyperledger.besu.ethereum.api.query.BlockchainQueries;
+import org.hyperledger.besu.ethereum.chain.Blockchain;
+import org.hyperledger.besu.ethereum.core.BlockHeader;
+import org.hyperledger.besu.ethereum.core.MiningConfiguration;
+import org.hyperledger.besu.ethereum.mainnet.ProtocolSchedule;
+import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
+import org.hyperledger.besu.plugin.services.rpc.RpcResponseType;
+
+import java.util.Optional;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import org.apache.tuweni.bytes.Bytes;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+/**
+ * This test only exercises bonsai worldstate since forest is essentially a no-op for moving the
+ * worldstate.
+ */
+public class DebugSetHeadTest extends AbstractJsonRpcHttpServiceTest {
+
+  DebugSetHead debugSetHead;
+  Blockchain blockchain;
+  WorldStateArchive archive;
+  ProtocolContext protocolContext;
+  ProtocolSchedule protocolSchedule;
+
+  @Override
+  @BeforeEach
+  public void setup() throws Exception {
+    setupBonsaiBlockchain();
+    blockchain = blockchainSetupUtil.getBlockchain();
+    protocolContext = blockchainSetupUtil.getProtocolContext();
+    protocolSchedule = blockchainSetupUtil.getProtocolSchedule();
+    ;
+    archive = blockchainSetupUtil.getWorldArchive();
+    debugSetHead =
+        new DebugSetHead(
+            new BlockchainQueries(
+                protocolSchedule, blockchain, archive, MiningConfiguration.MINING_DISABLED),
+            protocolContext,
+            // a value of 2 here exercises all the state rolling code paths
+            2);
+    startService();
+  }
+
+  @ParameterizedTest
+  @ValueSource(
+      strings = {"0x01", "0x4e9a67b663f9abe03e7e9fd5452c9497998337077122f44ee78a466f6a7358de"})
+  public void assertOnlyChainHeadMovesWorldParameterAbsent(final String blockParam) {
+    var chainTip = blockchain.getChainHead().getBlockHeader();
+    var blockOne = getBlockHeaderForHashOrNumber(blockParam).orElse(null);
+
+    assertThat(blockOne).isNotNull();
+    assertThat(blockOne).isNotEqualTo(chainTip);
+
+    // move the head to param val, number or hash
+    debugSetHead.response(debugSetHead(blockParam, Optional.empty()));
+
+    // get the new chainTip:
+    var newChainTip = blockchain.getChainHead().getBlockHeader();
+
+    // assert the chain moved, and the worldstate did not
+    assertThat(newChainTip).isEqualTo(blockOne);
+    assertThat(archive.getMutable().rootHash()).isEqualTo(chainTip.getStateRoot());
+  }
+
+  @ParameterizedTest
+  @ValueSource(
+      strings = {
+        "0x01",
+        "0x02",
+        "0x3d813a0ffc9cd04436e17e3e9c309f1e80df0407078e50355ce0d570b5424812",
+        "0x4e9a67b663f9abe03e7e9fd5452c9497998337077122f44ee78a466f6a7358de"
+      })
+  public void assertOnlyChainHeadMoves(final String blockParam) {
+    var chainTip = blockchain.getChainHead().getBlockHeader();
+    var blockOne = getBlockHeaderForHashOrNumber(blockParam).orElse(null);
+
+    assertThat(blockOne).isNotNull();
+    assertThat(blockOne).isNotEqualTo(chainTip);
+
+    // move the head to param val, number or hash
+    debugSetHead.response(debugSetHead(blockParam, Optional.of(FALSE)));
+
+    // get the new chainTip:
+    var newChainTip = blockchain.getChainHead().getBlockHeader();
+
+    // assert the chain moved, and the worldstate did not
+    assertThat(newChainTip).isEqualTo(blockOne);
+    assertThat(archive.getMutable().rootHash()).isEqualTo(chainTip.getStateRoot());
+  }
+
+  @ParameterizedTest
+  @ValueSource(
+      strings = {
+        "0x01",
+        "0x02",
+        "0x3d813a0ffc9cd04436e17e3e9c309f1e80df0407078e50355ce0d570b5424812",
+        "0x4e9a67b663f9abe03e7e9fd5452c9497998337077122f44ee78a466f6a7358de"
+      })
+  public void assertBothChainHeadAndWorldStatByNumber(final String blockParam) {
+    var chainTip = blockchain.getChainHead().getBlockHeader();
+    var blockOne = getBlockHeaderForHashOrNumber(blockParam).orElse(null);
+
+    assertThat(blockOne).isNotNull();
+    assertThat(blockOne).isNotEqualTo(chainTip);
+
+    // move the head and worldstate to param val number or hash
+    debugSetHead.response(debugSetHead(blockParam, Optional.of(TRUE)));
+
+    // get the new chainTip:
+    var newChainTip = blockchain.getChainHead().getBlockHeader();
+
+    // assert both the chain and worldstate moved to block one
+    assertThat(newChainTip).isEqualTo(blockOne);
+    assertThat(archive.getMutable().rootHash()).isEqualTo(blockOne.getStateRoot());
+  }
+
+  @Test
+  public void assertNotFound() {
+    var chainTip = blockchain.getChainHead().getBlockHeader();
+
+    // move the head to number just after chain head
+    var resp =
+        debugSetHead.response(debugSetHead("" + chainTip.getNumber() + 1, Optional.of(TRUE)));
+    assertThat(resp.getType()).isEqualTo(RpcResponseType.ERROR);
+
+    // move the head to some arbitrary hash
+    var resp2 =
+        debugSetHead.response(
+            debugSetHead(
+                Hash.keccak256(Bytes.fromHexString("0xdeadbeef")).toHexString(),
+                Optional.of(TRUE)));
+    assertThat(resp2.getType()).isEqualTo(RpcResponseType.ERROR);
+
+    // get the new chainTip:
+    var newChainTip = blockchain.getChainHead().getBlockHeader();
+
+    // assert neither the chain nor the worldstate moved
+    assertThat(newChainTip).isEqualTo(chainTip);
+    assertThat(archive.getMutable().rootHash()).isEqualTo(chainTip.getStateRoot());
+  }
+
+  private JsonRpcRequestContext debugSetHead(
+      final String numberOrHash, final Optional<Boolean> moveWorldState) {
+    if (moveWorldState.isPresent()) {
+      return new JsonRpcRequestContext(
+          new JsonRpcRequest(
+              "2.0", "debug_setHead", new Object[] {numberOrHash, moveWorldState.get()}));
+    } else {
+      return new JsonRpcRequestContext(
+          new JsonRpcRequest("2.0", "debug_setHead", new Object[] {numberOrHash}));
+    }
+  }
+
+  private Optional<BlockHeader> getBlockHeaderForHashOrNumber(final String input) {
+    try {
+      var param = new BlockParameterOrBlockHash(input);
+      if (param.getHash().isPresent()) {
+        return blockchain.getBlockHeader(param.getHash().get());
+      } else if (param.getNumber().isPresent()) {
+        return blockchain.getBlockHeader(param.getNumber().getAsLong());
+      }
+    } catch (JsonProcessingException ignored) {
+      // meh
+    }
+    return Optional.empty();
+  }
+}