Skip to content

Commit

Permalink
Import export trie log (#6363)
Browse files Browse the repository at this point in the history
* Import and export trie log subcommands
* change option name and fix descriptions

Signed-off-by: Gabriel Fukushima <[email protected]>
Co-authored-by: Jason Frame <[email protected]>
  • Loading branch information
gfukushima and jframe authored Jan 22, 2024
1 parent 24718e3 commit cfea3ab
Show file tree
Hide file tree
Showing 5 changed files with 349 additions and 83 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -62,23 +62,28 @@ public class DataStorageOptions implements CLIOptions<DataStorageConfiguration>
private final DataStorageOptions.Unstable unstableOptions = new Unstable();

static class Unstable {
private static final String BONSAI_LIMIT_TRIE_LOGS_ENABLED =
"--Xbonsai-limit-trie-logs-enabled";
private static final String BONSAI_TRIE_LOGS_RETENTION_THRESHOLD =
"--Xbonsai-trie-logs-retention-threshold";
private static final String BONSAI_TRIE_LOG_PRUNING_LIMIT = "--Xbonsai-trie-logs-pruning-limit";

@CommandLine.Option(
hidden = true,
names = {"--Xbonsai-trie-log-pruning-enabled"},
names = {BONSAI_LIMIT_TRIE_LOGS_ENABLED},
description = "Enable trie log pruning. (default: ${DEFAULT-VALUE})")
private boolean bonsaiTrieLogPruningEnabled = DEFAULT_BONSAI_TRIE_LOG_PRUNING_ENABLED;

@CommandLine.Option(
hidden = true,
names = {"--Xbonsai-trie-log-retention-threshold"},
names = {BONSAI_TRIE_LOGS_RETENTION_THRESHOLD},
description =
"The number of blocks for which to retain trie logs. (default: ${DEFAULT-VALUE})")
private long bonsaiTrieLogRetentionThreshold = DEFAULT_BONSAI_TRIE_LOG_RETENTION_THRESHOLD;

@CommandLine.Option(
hidden = true,
names = {"--Xbonsai-trie-log-pruning-limit"},
names = {BONSAI_TRIE_LOG_PRUNING_LIMIT},
description =
"The max number of blocks to load and prune trie logs for at startup. (default: ${DEFAULT-VALUE})")
private int bonsaiTrieLogPruningLimit = DEFAULT_BONSAI_TRIE_LOG_PRUNING_LIMIT;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,11 @@
import org.hyperledger.besu.ethereum.chain.Blockchain;
import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
import org.hyperledger.besu.ethereum.core.BlockHeader;
import org.hyperledger.besu.ethereum.rlp.BytesValueRLPInput;
import org.hyperledger.besu.ethereum.rlp.RLP;
import org.hyperledger.besu.ethereum.trie.bonsai.storage.BonsaiWorldStateKeyValueStorage;
import org.hyperledger.besu.ethereum.trie.bonsai.trielog.TrieLogFactoryImpl;
import org.hyperledger.besu.ethereum.trie.bonsai.trielog.TrieLogLayer;
import org.hyperledger.besu.ethereum.worldstate.DataStorageConfiguration;

import java.io.File;
Expand All @@ -32,13 +36,15 @@
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.PrintWriter;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.tuweni.bytes.Bytes;
import org.apache.tuweni.bytes.Bytes32;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand Down Expand Up @@ -97,16 +103,15 @@ private static void processTrieLogBatches(
final String batchFileNameBase) {

for (long batchNumber = 1; batchNumber <= numberOfBatches; batchNumber++) {

final String batchFileName = batchFileNameBase + "-" + batchNumber;
final long firstBlockOfBatch = chainHeight - ((batchNumber - 1) * BATCH_SIZE);

final long lastBlockOfBatch =
Math.max(chainHeight - (batchNumber * BATCH_SIZE), lastBlockNumberToRetainTrieLogsFor);

final List<Hash> trieLogKeys =
getTrieLogKeysForBlocks(blockchain, firstBlockOfBatch, lastBlockOfBatch);

saveTrieLogBatches(batchFileNameBase, rootWorldStateStorage, batchNumber, trieLogKeys);
LOG.info("Saving trie logs to retain in file (batch {})...", batchNumber);
saveTrieLogBatches(batchFileName, rootWorldStateStorage, trieLogKeys);
}

LOG.info("Clear trie logs...");
Expand All @@ -118,15 +123,12 @@ private static void processTrieLogBatches(
}

private static void saveTrieLogBatches(
final String batchFileNameBase,
final String batchFileName,
final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
final long batchNumber,
final List<Hash> trieLogKeys) {

LOG.info("Saving trie logs to retain in file (batch {})...", batchNumber);

try {
saveTrieLogsInFile(trieLogKeys, rootWorldStateStorage, batchNumber, batchFileNameBase);
saveTrieLogsInFile(trieLogKeys, rootWorldStateStorage, batchFileName);
} catch (IOException e) {
LOG.error("Error saving trie logs to file: {}", e.getMessage());
throw new RuntimeException(e);
Expand Down Expand Up @@ -210,9 +212,8 @@ private static void recreateTrieLogs(
final String batchFileNameBase)
throws IOException {
// process in chunk to avoid OOM

IdentityHashMap<byte[], byte[]> trieLogsToRetain =
readTrieLogsFromFile(batchFileNameBase, batchNumber);
final String batchFileName = batchFileNameBase + "-" + batchNumber;
IdentityHashMap<byte[], byte[]> trieLogsToRetain = readTrieLogsFromFile(batchFileName);
final int chunkSize = ROCKSDB_MAX_INSERTS_PER_TRANSACTION;
List<byte[]> keys = new ArrayList<>(trieLogsToRetain.keySet());

Expand Down Expand Up @@ -265,11 +266,10 @@ private static void validatePruneConfiguration(final DataStorageConfiguration co
private static void saveTrieLogsInFile(
final List<Hash> trieLogsKeys,
final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
final long batchNumber,
final String batchFileNameBase)
final String batchFileName)
throws IOException {

File file = new File(batchFileNameBase + "-" + batchNumber);
File file = new File(batchFileName);
if (file.exists()) {
LOG.error("File already exists, skipping file creation");
return;
Expand All @@ -285,24 +285,67 @@ private static void saveTrieLogsInFile(
}

@SuppressWarnings("unchecked")
private static IdentityHashMap<byte[], byte[]> readTrieLogsFromFile(
final String batchFileNameBase, final long batchNumber) {
static IdentityHashMap<byte[], byte[]> readTrieLogsFromFile(final String batchFileName) {

IdentityHashMap<byte[], byte[]> trieLogs;
try (FileInputStream fis = new FileInputStream(batchFileNameBase + "-" + batchNumber);
try (FileInputStream fis = new FileInputStream(batchFileName);
ObjectInputStream ois = new ObjectInputStream(fis)) {

trieLogs = (IdentityHashMap<byte[], byte[]>) ois.readObject();

} catch (IOException | ClassNotFoundException e) {

LOG.error(e.getMessage());
throw new RuntimeException(e);
}

return trieLogs;
}

private static void saveTrieLogsAsRlpInFile(
final List<Hash> trieLogsKeys,
final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
final String batchFileName) {
File file = new File(batchFileName);
if (file.exists()) {
LOG.error("File already exists, skipping file creation");
return;
}

final IdentityHashMap<byte[], byte[]> trieLogs =
getTrieLogs(trieLogsKeys, rootWorldStateStorage);
final Bytes rlp =
RLP.encode(
o ->
o.writeList(
trieLogs.entrySet(), (val, out) -> out.writeRaw(Bytes.wrap(val.getValue()))));
try {
Files.write(file.toPath(), rlp.toArrayUnsafe());
} catch (IOException e) {
LOG.error(e.getMessage());
throw new RuntimeException(e);
}
}

static IdentityHashMap<byte[], byte[]> readTrieLogsAsRlpFromFile(final String batchFileName) {
try {
final Bytes file = Bytes.wrap(Files.readAllBytes(Path.of(batchFileName)));
final BytesValueRLPInput input = new BytesValueRLPInput(file, false);

input.enterList();
final IdentityHashMap<byte[], byte[]> trieLogs = new IdentityHashMap<>();
while (!input.isEndOfCurrentList()) {
final Bytes trieLogBytes = input.currentListAsBytes();
TrieLogLayer trieLogLayer =
TrieLogFactoryImpl.readFrom(new BytesValueRLPInput(Bytes.wrap(trieLogBytes), false));
trieLogs.put(trieLogLayer.getBlockHash().toArrayUnsafe(), trieLogBytes.toArrayUnsafe());
}
input.leaveList();

return trieLogs;
} catch (IOException e) {
throw new RuntimeException(e);
}
}

private static IdentityHashMap<byte[], byte[]> getTrieLogs(
final List<Hash> trieLogKeys, final BonsaiWorldStateKeyValueStorage rootWorldStateStorage) {
IdentityHashMap<byte[], byte[]> trieLogsToRetain = new IdentityHashMap<>();
Expand Down Expand Up @@ -357,5 +400,25 @@ static void printCount(final PrintWriter out, final TrieLogCount count) {
count.total, count.canonicalCount, count.forkCount, count.orphanCount);
}

static void importTrieLog(
final BonsaiWorldStateKeyValueStorage rootWorldStateStorage, final Path trieLogFilePath) {

var trieLog = readTrieLogsAsRlpFromFile(trieLogFilePath.toString());

var updater = rootWorldStateStorage.updater();
trieLog.forEach((key, value) -> updater.getTrieLogStorageTransaction().put(key, value));
updater.getTrieLogStorageTransaction().commit();
}

static void exportTrieLog(
final BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
final List<Hash> trieLogHash,
final Path directoryPath)
throws IOException {
final String trieLogFile = directoryPath.toString();

saveTrieLogsAsRlpInFile(trieLogHash, rootWorldStateStorage, trieLogFile);
}

record TrieLogCount(int total, int canonicalCount, int forkCount, int orphanCount) {}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,16 +19,19 @@

import org.hyperledger.besu.cli.util.VersionProvider;
import org.hyperledger.besu.controller.BesuController;
import org.hyperledger.besu.datatypes.Hash;
import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
import org.hyperledger.besu.ethereum.storage.StorageProvider;
import org.hyperledger.besu.ethereum.trie.bonsai.storage.BonsaiWorldStateKeyValueStorage;
import org.hyperledger.besu.ethereum.trie.bonsai.trielog.TrieLogPruner;
import org.hyperledger.besu.ethereum.worldstate.DataStorageConfiguration;
import org.hyperledger.besu.ethereum.worldstate.DataStorageFormat;

import java.io.IOException;
import java.io.PrintWriter;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

import org.apache.logging.log4j.Level;
import org.apache.logging.log4j.core.config.Configurator;
Expand All @@ -43,7 +46,12 @@
description = "Manipulate trie logs",
mixinStandardHelpOptions = true,
versionProvider = VersionProvider.class,
subcommands = {TrieLogSubCommand.CountTrieLog.class, TrieLogSubCommand.PruneTrieLog.class})
subcommands = {
TrieLogSubCommand.CountTrieLog.class,
TrieLogSubCommand.PruneTrieLog.class,
TrieLogSubCommand.ExportTrieLog.class,
TrieLogSubCommand.ImportTrieLog.class
})
public class TrieLogSubCommand implements Runnable {

@SuppressWarnings("UnusedVariable")
Expand Down Expand Up @@ -123,6 +131,102 @@ public void run() {
}
}

@Command(
name = "export",
description = "This command exports the trie log of a determined block to a binary file",
mixinStandardHelpOptions = true,
versionProvider = VersionProvider.class)
static class ExportTrieLog implements Runnable {

@SuppressWarnings("unused")
@ParentCommand
private TrieLogSubCommand parentCommand;

@SuppressWarnings("unused")
@CommandLine.Spec
private CommandLine.Model.CommandSpec spec; // Picocli injects reference to command spec

@CommandLine.Option(
names = "--trie-log-block-hash",
description =
"Comma separated list of hashes from the blocks you want to export the trie logs of",
split = " {0,1}, {0,1}",
arity = "1..*")
private List<String> trieLogBlockHashList;

@CommandLine.Option(
names = "--trie-log-file-path",
description = "The file you want to export the trie logs to",
arity = "1..1")
private Path trieLogFilePath = null;

@Override
public void run() {
if (trieLogFilePath == null) {
trieLogFilePath =
Paths.get(
TrieLogSubCommand.parentCommand
.parentCommand
.dataDir()
.resolve("trie-logs.bin")
.toAbsolutePath()
.toString());
}

TrieLogContext context = getTrieLogContext();

final List<Hash> listOfBlockHashes =
trieLogBlockHashList.stream().map(Hash::fromHexString).toList();

try {
TrieLogHelper.exportTrieLog(
context.rootWorldStateStorage(), listOfBlockHashes, trieLogFilePath);
} catch (IOException e) {
throw new RuntimeException(e);
}
}
}

@Command(
name = "import",
description = "This command imports a trie log exported by another besu node",
mixinStandardHelpOptions = true,
versionProvider = VersionProvider.class)
static class ImportTrieLog implements Runnable {

@SuppressWarnings("unused")
@ParentCommand
private TrieLogSubCommand parentCommand;

@SuppressWarnings("unused")
@CommandLine.Spec
private CommandLine.Model.CommandSpec spec; // Picocli injects reference to command spec

@CommandLine.Option(
names = "--trie-log-file-path",
description = "The file you want to import the trie logs from",
arity = "1..1")
private Path trieLogFilePath = null;

@Override
public void run() {
if (trieLogFilePath == null) {
trieLogFilePath =
Paths.get(
TrieLogSubCommand.parentCommand
.parentCommand
.dataDir()
.resolve("trie-logs.bin")
.toAbsolutePath()
.toString());
}

TrieLogContext context = getTrieLogContext();

TrieLogHelper.importTrieLog(context.rootWorldStateStorage(), trieLogFilePath);
}
}

record TrieLogContext(
DataStorageConfiguration config,
BonsaiWorldStateKeyValueStorage rootWorldStateStorage,
Expand Down
Loading

0 comments on commit cfea3ab

Please sign in to comment.