Merge pull request #2978 from ljedrz/feat/testchain-generator-txs

[Feat] Allow the testchain-generator to be fed with pregenerated txs

Author: kaimast

.circleci/config.yml

@@ -684,21 +684,21 @@ jobs:
 
   ledger-with-rocksdb-partition1:
     executor: rust-docker
-    resource_class: << pipeline.parameters.large >>
+    resource_class: << pipeline.parameters.xlarge >>
     steps:
       - run_test:
           workspace_member: snarkvm-ledger
-          flags: --release --features=rocks --partition count:1/2 -- --test-threads=8
+          flags: --release --features=rocks --partition count:1/2 -- --test-threads=10
           no_output_timeout: 20m
           timeout: 30m
 
   ledger-with-rocksdb-partition2:
     executor: rust-docker
-    resource_class: << pipeline.parameters.large >>
+    resource_class: << pipeline.parameters.xlarge >>
     steps:
       - run_test:
           workspace_member: snarkvm-ledger
-          flags: --release --features=rocks --partition count:2/2 -- --test-threads=8
+          flags: --release --features=rocks --partition count:2/2 -- --test-threads=10
           no_output_timeout: 20m
           timeout: 30m
 

Cargo.lock

@@ -68,12 +68,31 @@ pub struct TestChainBuilder<N: Network> {
 }
 
 /// Additional options you can pass to the builder when generating a set of blocks.
-#[derive(Clone, Default)]
-pub struct GenerateBlocksOptions {
+#[derive(Clone)]
+pub struct GenerateBlocksOptions<N: Network> {
     /// Do not include votes to the previous leader certificate
     pub skip_votes: bool,
     /// Do not generate certificates for the specific node indices (to simulate a partition).
     pub skip_nodes: Vec<usize>,
+    /// A flag indicating that a number of initial "placeholder blocks" should be baked
+    /// wthout transactions in order to skip to the latest version of consensus.
+    pub skip_to_current_version: bool,
+    /// The number of validators.
+    pub num_validators: usize,
+    /// Preloaded transactions to populate the blocks with.
+    pub transactions: Vec<Transaction<N>>,
+}
+
+impl<N: Network> Default for GenerateBlocksOptions<N> {
+    fn default() -> Self {
+        Self {
+            skip_votes: false,
+            skip_nodes: Default::default(),
+            skip_to_current_version: false,
+            num_validators: 0,
+            transactions: Default::default(),
+        }
+    }
 }
 
 /// Additional options you can pass to the builder when generating a single block.
@@ -178,7 +197,6 @@ impl<N: Network> TestChainBuilder<N> {
         Ok(Self {
             private_keys,
             ledger,
-
             genesis_block,
             last_batch_round: Default::default(),
             last_committed_batch_round: Default::default(),
@@ -190,7 +208,9 @@ impl<N: Network> TestChainBuilder<N> {
 
     /// Create multiple blocks, with fully-connected DAGs.
     pub fn generate_blocks(&mut self, num_blocks: usize, rng: &mut TestRng) -> Result<Vec<Block<N>>> {
-        self.generate_blocks_with_opts(num_blocks, GenerateBlocksOptions::default(), rng)
+        let num_validators = self.private_keys.len();
+
+        self.generate_blocks_with_opts(num_blocks, GenerateBlocksOptions { num_validators, ..Default::default() }, rng)
     }
 
     /// Create multiple blocks, with additional parameters.
@@ -200,20 +220,53 @@ impl<N: Network> TestChainBuilder<N> {
     pub fn generate_blocks_with_opts(
         &mut self,
         num_blocks: usize,
-        options: GenerateBlocksOptions,
+        mut options: GenerateBlocksOptions<N>,
         rng: &mut TestRng,
     ) -> Result<Vec<Block<N>>> {
         assert!(num_blocks > 0, "Need to build at least one block");
 
-        let options = GenerateBlockOptions {
-            skip_votes: options.skip_votes,
-            skip_nodes: options.skip_nodes,
-            ..Default::default()
-        };
-
         let mut result = vec![];
-        for _ in 0..num_blocks {
-            let block = self.generate_block_with_opts(options.clone(), rng)?;
+
+        // If configured, skip enough blocks to reach the current consensus version.
+        if options.skip_to_current_version {
+            let (version, target_height) = TEST_CONSENSUS_VERSION_HEIGHTS.last().unwrap();
+            let mut current_height = self.ledger.latest_height();
+
+            let diff = target_height.saturating_sub(current_height);
+
+            if diff > 0 {
+                println!("Skipping {diff} blocks to reach {version}");
+
+                while current_height < *target_height && result.len() < num_blocks {
+                    let options = GenerateBlockOptions {
+                        skip_votes: options.skip_votes,
+                        skip_nodes: options.skip_nodes.clone(),
+                        ..Default::default()
+                    };
+
+                    let block = self.generate_block_with_opts(options, rng)?;
+                    current_height = block.height();
+                    result.push(block);
+                }
+
+                println!("Advanced to the current consensus version at height {target_height}");
+            } else {
+                debug!("Already at the current consensus version. No blocks to skip.");
+            }
+        }
+
+        while result.len() < num_blocks {
+            let num_txs = (BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH * options.num_validators)
+                .min(options.transactions.len());
+
+            let options = GenerateBlockOptions {
+                skip_votes: options.skip_votes,
+                skip_nodes: options.skip_nodes.clone(),
+                transactions: options.transactions.drain(..num_txs).collect(),
+                ..Default::default()
+            };
+
+            let block = self.generate_block_with_opts(options, rng)?;
             result.push(block);
         }
 
@@ -262,8 +315,6 @@ impl<N: Network> TestChainBuilder<N> {
             transmissions.insert(transmission_id, transmission);
         }
 
-        let transmission_ids: IndexSet<_> = transmissions.keys().copied().collect();
-
         // =======================================
         // Create certificates for the new block.
         // =======================================
@@ -303,6 +354,13 @@ impl<N: Network> TestChainBuilder<N> {
                     continue;
                 }
 
+                let transmission_ids: IndexSet<_> = transmissions
+                    .keys()
+                    .skip(key1_idx * BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH)
+                    .take(BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH)
+                    .copied()
+                    .collect();
+
                 let batch_header = BatchHeader::new(
                     private_key_1,
                     round,
@@ -404,6 +462,7 @@ impl<N: Network> TestChainBuilder<N> {
 
         // Construct the block.
         let subdag = Subdag::from(subdag_map).unwrap();
+
         let block = self.ledger.prepare_advance_to_next_quorum_block(subdag, transmissions, rng)?;
 
         // Skip to increase performance.

ledger/src/test_helpers/chain_builder.rs

@@ -3390,7 +3390,7 @@ fn test_subdag_with_long_branch() -> Result<()> {
 #[test]
 fn test_subdag_with_gc_length() -> Result<()> {
     let rng = &mut TestRng::default();
-    let mut chain_builder = TestChainBuilder::new(rng)?;
+    let mut chain_builder = TestChainBuilder::new(rng).unwrap();
 
     let blocks = chain_builder.generate_blocks_with_opts(
         BatchHeader::<CurrentNetwork>::MAX_GC_ROUNDS / 2,

ledger/src/tests.rs

@@ -27,6 +27,7 @@ workspace = true
 
 [dependencies.snarkvm-console]
 workspace = true
+features = [ "test_consensus_heights" ]
 
 [dependencies.snarkvm-ledger]
 workspace = true
@@ -35,6 +36,9 @@ features = [ "test-helpers", "rocks" ]
 [dependencies.anyhow]
 workspace = true
 
+[dependencies.tracing]
+workspace = true
+
 [dependencies.tracing-subscriber]
 version = "0.3"
 

ledger/testchain-generator/Cargo.toml

@@ -13,13 +13,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use snarkvm_console::prelude::{CanaryV0, MainnetV0, Network, TestRng, TestnetV0, ToBytes};
-use snarkvm_ledger::{Ledger, store::helpers::rocksdb::ConsensusDB, test_helpers::TestChainBuilder};
+use snarkvm_console::prelude::{
+    CanaryV0,
+    MainnetV0,
+    Network,
+    TEST_CONSENSUS_VERSION_HEIGHTS,
+    TestRng,
+    TestnetV0,
+    ToBytes,
+};
+use snarkvm_ledger::{
+    Ledger,
+    Transaction,
+    store::helpers::rocksdb::ConsensusDB,
+    test_helpers::{TestChainBuilder, chain_builder::GenerateBlocksOptions},
+};
 
 use aleo_std::StorageMode;
 use anyhow::{Context, Result, bail};
 use clap::{Parser, builder::PossibleValuesParser};
-use std::fs;
+use std::{fs, path::Path, str::FromStr};
+use tracing::debug;
 
 #[derive(Parser)]
 struct Args {
@@ -38,6 +52,10 @@ struct Args {
     /// Remove existing ledger if it already exists.
     #[clap(long, short = 'f')]
     force: bool,
+    /// Load the transactions to be used with the generated blocks. They are expected to be
+    /// stored in a JSON-encoded format.
+    #[clap(long)]
+    txs_path: Option<String>,
     /// The name of the network to generate the chain for.
     #[clap(long, value_parser=PossibleValuesParser::new(vec![CanaryV0::SHORT_NAME, TestnetV0::SHORT_NAME, MainnetV0::SHORT_NAME]), default_value=TestnetV0::SHORT_NAME)]
     network: String,
@@ -95,6 +113,32 @@ fn generate_testchain<N: Network>(args: Args) -> Result<()> {
 
     remove_ledger(N::ID, &storage_mode, args.force)?;
 
+    let mut txs = if let Some(path) = args.txs_path {
+        let path = Path::new(&path);
+
+        if path.is_dir() {
+            println!("Attempting to load transactions from \"{}\"", path.display());
+        } else {
+            bail!("Cannot load transactions from \"{}\": not a valid directory", path.display());
+        }
+
+        let mut txs = Vec::new();
+        for entry in fs::read_dir(path)? {
+            let entry = entry?;
+            let path = entry.path();
+
+            let buffer = fs::read_to_string(path)?;
+            let tx = Transaction::<N>::from_str(&buffer)?;
+
+            txs.push(tx);
+        }
+
+        println!("Loaded {} tranactionss from \"{}\"", txs.len(), path.display());
+        txs
+    } else {
+        Default::default()
+    };
+
     let num_validators = args.num_validators;
     let num_blocks = args.num_blocks;
 
@@ -110,9 +154,31 @@ fn generate_testchain<N: Network>(args: Args) -> Result<()> {
     let mut pos = 0;
     let mut blocks = vec![];
 
+    // How many blocks to generate in a single batch.
+    const BATCH_SIZE: usize = 100;
+
+    // How many transactions to insert per block.
+    let latest_consensus_height = TEST_CONSENSUS_VERSION_HEIGHTS.last().unwrap().1 as usize;
+    let num_txn_blocks = num_blocks.saturating_sub(latest_consensus_height);
+    let txns_per_block = txs.len().div_ceil(num_txn_blocks);
+
     while blocks.len() < num_blocks {
-        let batch_size = (num_blocks - blocks.len()).min(100);
-        let mut batch = builder.generate_blocks(batch_size, &mut rng).with_context(|| "Failed to generate blocks")?;
+        let current_height = blocks.len();
+        // How many blocks to generate in this batch.
+        let batch_size = (num_blocks - current_height).min(BATCH_SIZE);
+        // Generate set of transactions to insert in this batch.
+        let num_empty_blocks = latest_consensus_height.saturating_sub(current_height);
+        let num_txns = (batch_size.saturating_sub(num_empty_blocks)) * txns_per_block;
+        let transactions = txs.drain(..num_txns).collect();
+
+        debug!("Generating next batch with {batch_size} blocks and {num_txns} transactions");
+        let mut batch = builder
+            .generate_blocks_with_opts(
+                batch_size,
+                GenerateBlocksOptions { transactions, skip_to_current_version: true, ..Default::default() },
+                &mut rng,
+            )
+            .with_context(|| "Failed to generate blocks")?;
 
         pos += batch_size;
         println!("Generated {pos} of {num_blocks} blocks");