sim-rs: apply sharding strategy

Author: SupernaviX

data/simulation/config.d.ts

@@ -58,6 +58,11 @@ export interface Config {
   "tx-generation-distribution": Distribution;
   /** Only supported by Rust simulation. */
   "tx-size-bytes-distribution": Distribution;
+  /**
+   * What percentage of transactions have at least one sharded input?
+   * 
+   * Only supported by Rust simulation. */
+  "tx-sharded-percentage": number;
   /** Only supported by Rust simulation. */
   "tx-validation-cpu-time-ms": number;
   /** Only supported by Rust simulation. */
@@ -76,8 +81,24 @@ export interface Config {
   // Input Block Configuration
   "ib-generation-probability": number;
   "ib-generation-cpu-time-ms": number;
-  /** Only supported by Rust simulation. */
+  /**
+   * The total number of shards available for IBs.
+   * Must be divisible by ib_shard_group_count.
+
+   * Only supported by Rust simulation. */
   "ib-shards": number;
+  /**
+   * The "m" IB sharding parameter.
+   * Controls how many slots in a row will use the same list of shards.
+   * 
+   * Only supported by Rust simulation. */
+  "ib-shard-period-length-slots": number;
+  /**
+   * The "k" IB sharding parameter.
+   * Controls how many groups of shards are assigned at a time.
+
+   * Only supported by Rust simulation. */
+  "ib-shard-group-count": number;
   "ib-head-size-bytes": bigint;
   "ib-head-validation-cpu-time-ms": number;
   "ib-body-validation-cpu-time-ms-constant": number;

data/simulation/config.default.yaml

@@ -43,6 +43,7 @@ tx-size-bytes-distribution:
   distribution: log-normal
   mu: 6.833
   sigma: 1.127
+tx-sharded-percentage: 0.5
 tx-validation-cpu-time-ms: 1.5
 tx-max-size-bytes: 16384
 
@@ -74,7 +75,9 @@ rb-body-legacy-praos-payload-validation-cpu-time-ms-per-byte: 0.0005
 ################################################################################
 
 ib-generation-probability: 5.0
-ib-shards: 1
+ib-shards: 50
+ib-shard-period-length-slots: 5
+ib-shard-group-count: 25
 
 # ProducerId  32
 # SlotNo      64

data/simulation/config.schema.json

@@ -243,8 +243,16 @@
     "ib-head-validation-cpu-time-ms": {
       "type": "number"
     },
+    "ib-shard-group-count": {
+      "description": "The \"k\" IB sharding parameter.\nControls how many groups of shards are assigned at a time.\n\nOnly supported by Rust simulation.",
+      "type": "number"
+    },
+    "ib-shard-period-length-slots": {
+      "description": "The \"m\" IB sharding parameter.\nControls how many slots in a row will use the same list of shards.\n\nOnly supported by Rust simulation.",
+      "type": "number"
+    },
     "ib-shards": {
-      "description": "Only supported by Rust simulation.",
+      "description": "The total number of shards available for IBs.\nMust be divisible by ib_shard_group_count.\n\nOnly supported by Rust simulation.",
       "type": "number"
     },
     "leios-header-diffusion-time-ms": {
@@ -335,6 +343,10 @@
       "properties": {},
       "type": "number"
     },
+    "tx-sharded-percentage": {
+      "description": "What percentage of transactions have at least one sharded input?\n\nOnly supported by Rust simulation.",
+      "type": "number"
+    },
     "tx-size-bytes-distribution": {
       "$ref": "#/definitions/Distribution",
       "description": "Only supported by Rust simulation."

sim-rs/sim-cli/src/events/aggregate.rs

@@ -59,6 +59,7 @@ impl TraceAggregator {
                 id,
                 publisher,
                 size_bytes,
+                ..
             } => {
                 self.transactions.insert(
                     id,

sim-rs/sim-cli/src/main.rs

@@ -90,7 +90,7 @@ fn read_config(args: &Args) -> Result<SimConfiguration> {
     }
 
     let params: RawParameters = raw_params.extract()?;
-    let mut config = SimConfiguration::build(params, topology);
+    let mut config = SimConfiguration::build(params, topology)?;
     if let Some(slots) = args.slots {
         config.slots = Some(slots);
     }

sim-rs/sim-core/src/config.rs

@@ -67,6 +67,7 @@ pub struct RawParameters {
     // Transaction configuration
     pub tx_generation_distribution: DistributionConfig,
     pub tx_size_bytes_distribution: DistributionConfig,
+    pub tx_sharded_percentage: f64,
     pub tx_validation_cpu_time_ms: f64,
     pub tx_max_size_bytes: u64,
 
@@ -85,6 +86,8 @@ pub struct RawParameters {
     pub ib_generation_probability: f64,
     pub ib_generation_cpu_time_ms: f64,
     pub ib_shards: u64,
+    pub ib_shard_period_length_slots: u64,
+    pub ib_shard_group_count: u64,
     pub ib_head_size_bytes: u64,
     pub ib_head_validation_cpu_time_ms: f64,
     pub ib_body_validation_cpu_time_ms_constant: f64,
@@ -371,6 +374,7 @@ impl TransactionConfig {
         if params.simulate_transactions {
             Self::Real(RealTransactionConfig {
                 max_size: params.tx_max_size_bytes,
+                sharded_percentage: params.tx_sharded_percentage,
                 frequency_ms: params.tx_generation_distribution.into(),
                 size_bytes: params.tx_size_bytes_distribution.into(),
             })
@@ -387,6 +391,7 @@ impl TransactionConfig {
 #[derive(Debug, Clone)]
 pub(crate) struct RealTransactionConfig {
     pub max_size: u64,
+    pub sharded_percentage: f64,
     pub frequency_ms: FloatDistribution,
     pub size_bytes: FloatDistribution,
 }
@@ -434,14 +439,23 @@ pub struct SimConfiguration {
     pub(crate) ib_diffusion_strategy: DiffusionStrategy,
     pub(crate) max_ib_requests_per_peer: usize,
     pub(crate) ib_shards: u64,
+    pub(crate) ib_shard_period_slots: u64,
+    pub(crate) ib_shard_groups: u64,
     pub(crate) cpu_times: CpuTimeConfig,
     pub(crate) sizes: BlockSizeConfig,
     pub(crate) transactions: TransactionConfig,
 }
 
 impl SimConfiguration {
-    pub fn build(params: RawParameters, topology: Topology) -> Self {
-        Self {
+    pub fn build(params: RawParameters, topology: Topology) -> Result<Self> {
+        if params.ib_shards % params.ib_shard_group_count != 0 {
+            bail!(
+                "ib-shards ({}) is not divisible by ib-shard-group-count ({})",
+                params.ib_shards,
+                params.ib_shard_group_count
+            );
+        }
+        Ok(Self {
             seed: 0,
             slots: None,
             emit_conformance_events: false,
@@ -467,10 +481,12 @@ impl SimConfiguration {
             ib_diffusion_strategy: params.ib_diffusion_strategy,
             max_ib_requests_per_peer: params.ib_diffusion_max_bodies_to_request as usize,
             ib_shards: params.ib_shards,
+            ib_shard_period_slots: params.ib_shard_period_length_slots,
+            ib_shard_groups: params.ib_shard_group_count,
             cpu_times: CpuTimeConfig::new(&params),
             sizes: BlockSizeConfig::new(&params),
             transactions: TransactionConfig::new(&params),
-        }
+        })
     }
 }
 

sim-rs/sim-core/src/events.rs

@@ -102,6 +102,7 @@ pub enum Event {
         id: TransactionId,
         publisher: Node,
         size_bytes: u64,
+        shard: Option<u64>,
     },
     TXSent {
         id: TransactionId,
@@ -158,6 +159,7 @@ pub enum Event {
         pipeline: u64,
         producer: Node,
         index: u64,
+        shard: u64,
         header_bytes: u64,
         tx_payload_bytes: u64,
         size_bytes: u64,
@@ -417,6 +419,7 @@ impl EventTracker {
             id: transaction.id,
             publisher: self.to_node(publisher),
             size_bytes: transaction.bytes,
+            shard: transaction.shard,
         });
     }
 
@@ -461,6 +464,7 @@ impl EventTracker {
             pipeline: block.header.id.pipeline,
             producer: self.to_node(block.header.id.producer),
             index: block.header.id.index,
+            shard: block.header.shard,
             header_bytes,
             tx_payload_bytes,
             size_bytes: header_bytes + tx_payload_bytes,

sim-rs/sim-core/src/model.rs

@@ -89,7 +89,7 @@ id_wrapper!(TransactionId, u64);
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Transaction {
     pub id: TransactionId,
-    pub shard: u64,
+    pub shard: Option<u64>,
     pub bytes: u64,
 }
 
@@ -125,6 +125,7 @@ impl<Node: Display> Serialize for InputBlockId<Node> {
 pub struct InputBlockHeader {
     pub id: InputBlockId,
     pub vrf: u64,
+    pub shard: u64,
     pub timestamp: Timestamp,
     pub bytes: u64,
 }

sim-rs/sim-core/src/sim/node.rs

@@ -548,6 +548,8 @@ impl Node {
             }
         }
         for (slot, vrfs) in slot_vrfs {
+            let shards = self.find_available_ib_shards(slot);
+            assert!(!shards.is_empty());
             let headers = vrfs
                 .into_iter()
                 .enumerate()
@@ -559,6 +561,7 @@ impl Node {
                         index: index as u64,
                     },
                     vrf,
+                    shard: shards[vrf as usize % shards.len()],
                     timestamp: self.clock.now(),
                     bytes: self.sim_config.sizes.ib_header,
                 })
@@ -567,6 +570,13 @@ impl Node {
         }
     }
 
+    fn find_available_ib_shards(&self, slot: u64) -> Vec<u64> {
+        let period = slot / self.sim_config.ib_shard_period_slots;
+        let group = period % self.sim_config.ib_shard_groups;
+        let shards_per_group = self.sim_config.ib_shards / self.sim_config.ib_shard_groups;
+        (group * shards_per_group..(group + 1) * shards_per_group).collect()
+    }
+
     fn generate_endorser_blocks(&mut self, slot: u64) {
         for next_p in vrf_probabilities(self.sim_config.eb_generation_probability) {
             if self.run_vrf(next_p).is_some() {
@@ -719,7 +729,7 @@ impl Node {
             // Add one transaction, the right size for the extra RB payload
             let tx = Transaction {
                 id: config.next_id(),
-                shard: 0,
+                shard: None,
                 bytes: config.rb_size,
             };
             self.tracker.track_transaction_generated(&tx, self.id);
@@ -1267,21 +1277,20 @@ impl Node {
         if let TransactionConfig::Mock(config) = &self.sim_config.transactions {
             let tx = Transaction {
                 id: config.next_id(),
-                shard: 0,
+                shard: None,
                 bytes: config.ib_size,
             };
             self.tracker.track_transaction_generated(&tx, self.id);
             ib.transactions.push(Arc::new(tx));
             return;
         }
 
-        let shard = ib.header.vrf % self.sim_config.ib_shards;
         let candidate_txs: Vec<_> = self
             .leios
             .mempool
             .values()
             .filter_map(|tx| {
-                if tx.shard == shard {
+                if tx.shard.is_none_or(|shard| shard == ib.header.shard) {
                     Some((tx.id, tx.bytes))
                 } else {
                     None

sim-rs/sim-core/src/sim/tx.rs

@@ -49,7 +49,9 @@ impl TransactionProducer {
         let mut rng = &mut self.rng;
         loop {
             let id = TransactionId::new(next_tx_id);
-            let shard = rng.random_range(0..self.ib_shards);
+            let shard = rng
+                .random_bool(config.sharded_percentage)
+                .then(|| rng.random_range(0..self.ib_shards));
             let bytes = (config.size_bytes.sample(&mut rng) as u64).min(config.max_size);
             let tx = Transaction { id, shard, bytes };