Added cli tool for pruning blockchain states

rskj-core/src/main/java/co/rsk/cli/tools/PruneState.java

@@ -0,0 +1,85 @@
+/*
+ * This file is part of RskJ
+ * Copyright (C) 2023 RSK Labs Ltd.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package co.rsk.cli.tools;
+
+import co.rsk.cli.PicoCliToolRskContextAware;
+import co.rsk.config.RskSystemProperties;
+import co.rsk.util.BlockStateHandler;
+import org.ethereum.datasource.KeyValueDataSource;
+import org.ethereum.datasource.KeyValueDataSourceUtils;
+import org.ethereum.util.FileUtil;
+import picocli.CommandLine;
+
+import java.lang.invoke.MethodHandles;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+@CommandLine.Command(name = "prune-state", mixinStandardHelpOptions = true, version = "prune-state 1.0",
+        description = "The entry point for pruning blockchain state")
+public class PruneState extends PicoCliToolRskContextAware {
+
+    @CommandLine.Option(names = {"-b", "--numOfBlocks"}, description = "Number of top blocks to preserve from pruning (128 by default)")
+    private Long numOfBlocks;
+
+    public static void main(String[] args) {
+        create(MethodHandles.lookup().lookupClass()).execute(args);
+    }
+
+    @Override
+    public Integer call() throws Exception {
+        if (numOfBlocks == null) {
+            numOfBlocks = 128L; // by default
+        } else if (numOfBlocks <= 0) {
+            throw new IllegalArgumentException("'numOfBlocks' cannot be less than or equal to zero");
+        }
+
+        RskSystemProperties systemProperties = ctx.getRskSystemProperties();
+        Path databasePath = Paths.get(systemProperties.databaseDir());
+        Path trieStorePath = databasePath.resolve("unitrie");
+        if (!trieStorePath.toFile().exists()) {
+            throw new IllegalStateException("'unitrie' db folder not found");
+        }
+
+        Path tmpTrieStorePath = databasePath.resolve("tmp_unitrie");
+        BlockStateHandler stateHandler = new BlockStateHandler(ctx.getBlockchain(), ctx.getTrieStore(), ctx.getStateRootHandler());
+
+        try {
+            KeyValueDataSource destDataSource = KeyValueDataSourceUtils.makeDataSource(tmpTrieStorePath, systemProperties.databaseKind());
+            try {
+                stateHandler.copyState(numOfBlocks, destDataSource);
+            } finally {
+                destDataSource.close();
+            }
+        } catch (Exception e) {
+            printError("State copying failed", e);
+
+            FileUtil.recursiveDelete(tmpTrieStorePath.toString());
+
+            throw e;
+        } finally {
+            ctx.close();
+        }
+
+        FileUtil.recursiveDelete(trieStorePath.toString());
+        Files.move(tmpTrieStorePath, trieStorePath);
+
+        return 0;
+    }
+}

rskj-core/src/main/java/co/rsk/trie/NodeReference.java

@@ -90,6 +90,27 @@ public Optional<Trie> getNode() {
         return node;
     }
 
+    public Optional<Trie> getNodeDetached() {
+        if (lazyNode != null) {
+            return Optional.of(lazyNode);
+        }
+
+        if (lazyHash == null) {
+            return Optional.empty();
+        }
+
+        Optional<Trie> node = store.retrieve(lazyHash.getBytes());
+
+        // Broken database, can't continue
+        if (!node.isPresent()) {
+            logger.error("Broken database, execution can't continue");
+            nodeStopper.stop(1);
+            return Optional.empty();
+        }
+
+        return node;
+    }
+
     /**
      * The hash or empty if this is an empty reference.
      * If the hash is not present but its node is known, it will be calculated.

rskj-core/src/main/java/co/rsk/trie/Trie.java

@@ -753,6 +753,14 @@ public NodeReference getRight() {
         return right;
     }
 
+    public Optional<Trie> getLeftNodeDetached() {
+        return left.getNodeDetached();
+    }
+
+    public Optional<Trie> getRightNodeDetached() {
+        return right.getNodeDetached();
+    }
+
     /**
      * put key with associated value, returning a new NewTrie
      *

rskj-core/src/main/java/co/rsk/util/BlockStateHandler.java

@@ -0,0 +1,117 @@
+/*
+ * This file is part of RskJ
+ * Copyright (C) 2023 RSK Labs Ltd.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package co.rsk.util;
+
+import co.rsk.crypto.Keccak256;
+import co.rsk.db.StateRootHandler;
+import co.rsk.trie.Trie;
+import co.rsk.trie.TrieStore;
+import org.ethereum.core.Block;
+import org.ethereum.core.Blockchain;
+import org.ethereum.datasource.KeyValueDataSource;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nonnull;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+
+public class BlockStateHandler {
+
+    private static final Logger logger = LoggerFactory.getLogger(BlockStateHandler.class);
+
+    private static final int CACHE_SIZE = 10_000_000;
+    private static final int BATCH_SIZE = 1_000_000;
+
+    private final Blockchain blockchain;
+    private final TrieStore trieStore;
+    private final StateRootHandler stateRootHandler;
+
+    public BlockStateHandler(@Nonnull Blockchain blockchain, @Nonnull TrieStore trieStore,
+                             @Nonnull StateRootHandler stateRootHandler) {
+        this.blockchain = Objects.requireNonNull(blockchain);
+        this.trieStore = Objects.requireNonNull(trieStore);
+        this.stateRootHandler = Objects.requireNonNull(stateRootHandler);
+    }
+
+    public void copyState(long lastNumOfBlocksToPreserve, @Nonnull KeyValueDataSource destDataSource) {
+        if (lastNumOfBlocksToPreserve <= 0) {
+            throw new IllegalArgumentException("lastNumOfBlocksToPreserve: " + lastNumOfBlocksToPreserve);
+        }
+
+        Block bestBlock = Objects.requireNonNull(blockchain.getBestBlock());
+        long toBlock = bestBlock.getNumber();
+        long fromBlock = Math.max(0, toBlock - lastNumOfBlocksToPreserve + 1);
+        long curBlock = fromBlock;
+        long blockCount = toBlock - fromBlock + 1;
+        List<Block> blocks;
+
+        TrieHandler trieHandler = new TrieHandler(destDataSource, CACHE_SIZE, BATCH_SIZE);
+
+        logger.info("Start copying states for {} block(s) of range: [{}; {}]", blockCount, fromBlock, toBlock);
+
+        long startTime = System.currentTimeMillis();
+
+        do {
+            blocks = Objects.requireNonNull(blockchain.getBlocksByNumber(curBlock));
+            for (Block block : blocks) {
+                Keccak256 stateRoot = stateRootHandler.translate(block.getHeader());
+
+                Optional<Trie> trieOpt = trieStore.retrieve(stateRoot.getBytes());
+                if (trieOpt.isPresent()) {
+                    Trie trie = trieOpt.get();
+                    trieHandler.copyTrie(curBlock, trie);
+                } else {
+                    logger.info("No trie found at block height: {}. Moving to next one", curBlock);
+                }
+            }
+
+            curBlock++;
+        } while (curBlock <= toBlock || !blocks.isEmpty());
+
+        long endTime = System.currentTimeMillis();
+        Duration duration = Duration.of(endTime - startTime, ChronoUnit.MILLIS);
+
+        logger.info("Finished copying states. Processed {} block(s). Duration: {}", blockCount, duration);
+    }
+
+    public static class MissingBlockException extends RuntimeException {
+
+        public MissingBlockException(String message) {
+            super(message);
+        }
+    }
+
+    public static class InvalidBlockException extends RuntimeException {
+
+        public InvalidBlockException(String message) {
+            super(message);
+        }
+    }
+
+    public static class MissingTrieException extends RuntimeException {
+
+        public MissingTrieException(String message) {
+            super(message);
+        }
+    }
+}