sim-rs: support weighted TX production, stake pools don't generate TXs

Author: SupernaviX

data/simulation/topology.d.ts

@@ -27,6 +27,12 @@ export interface Node<Location> {
    * What fraction of TXs (from 0 to 1) should introduce conflicts with transactions which were produced before?
    * Only supported by Rust simulation.  */
   "tx-conflict-fraction"?: number | null;
+  /**
+   * How likely is this node to generate transactions, compared to its peers?
+   * Default is 0 for nodes with stake, 1 for nodes with no stake.
+   * Only supported by Rust simulation.
+   */
+  "tx-generation-weight"?: number | null;
   /** If not null, the node will behave according to the given Behaviour.
   *
   * Only supported by Haskell simulation.

data/simulation/topology.schema.json

@@ -52,6 +52,10 @@
         "tx-conflict-fraction": {
           "description": "What fraction of TXs (from 0 to 1) should introduce conflicts with transactions which were produced before?\nOnly supported by Rust simulation.",
           "type": "number"
+        },
+        "tx-generation-weight": {
+          "description": "How likely is this node to generate transactions, compared to its peers?\nDefault is 0 for nodes with stake, 1 for nodes with no stake.\nOnly supported by Rust simulation.",
+          "type": "number"
         }
       },
       "type": "object"
@@ -95,6 +99,10 @@
         "tx-conflict-fraction": {
           "description": "What fraction of TXs (from 0 to 1) should introduce conflicts with transactions which were produced before?\nOnly supported by Rust simulation.",
           "type": "number"
+        },
+        "tx-generation-weight": {
+          "description": "How likely is this node to generate transactions, compared to its peers?\nDefault is 0 for nodes with stake, 1 for nodes with no stake.\nOnly supported by Rust simulation.",
+          "type": "number"
         }
       },
       "type": "object"

sim-rs/sim-cli/src/bin/gen-test-data/strategy/utils.rs

@@ -171,6 +171,7 @@ impl GraphBuilder {
                     location: RawNodeLocation::Coords(n.location),
                     cpu_core_count: n.cores,
                     tx_conflict_fraction: None,
+                    tx_generation_weight: None,
                     producers: BTreeMap::new(),
                 };
                 (name, node)

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

@@ -182,6 +182,8 @@ pub struct RawNode {
     pub cpu_core_count: Option<u64>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub tx_conflict_fraction: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tx_generation_weight: Option<u64>,
     pub producers: BTreeMap<String, RawLinkInfo>,
 }
 
@@ -260,6 +262,7 @@ impl From<RawTopology> for Topology {
                     cpu_multiplier: 1.0,
                     cores: node.cpu_core_count,
                     tx_conflict_fraction: node.tx_conflict_fraction,
+                    tx_generation_weight: node.tx_generation_weight,
                     consumers: vec![],
                 },
             );
@@ -563,6 +566,7 @@ pub struct NodeConfiguration {
     pub cpu_multiplier: f64,
     pub cores: Option<u64>,
     pub tx_conflict_fraction: Option<f64>,
+    pub tx_generation_weight: Option<u64>,
     pub consumers: Vec<NodeId>,
 }
 

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

@@ -14,6 +14,7 @@ use crate::{
 struct NodeState {
     sink: mpsc::UnboundedSender<Arc<Transaction>>,
     tx_conflict_fraction: Option<f64>,
+    tx_generation_weight: u64,
 }
 
 pub struct TransactionProducer {
@@ -36,10 +37,14 @@ impl TransactionProducer {
             .iter()
             .map(|node| {
                 let sink = node_tx_sinks.remove(&node.id).unwrap();
-                let state = NodeState {
-                    sink,
-                    tx_conflict_fraction: node.tx_conflict_fraction,
-                };
+                let state =
+                    NodeState {
+                        sink,
+                        tx_conflict_fraction: node.tx_conflict_fraction,
+                        tx_generation_weight: node
+                            .tx_generation_weight
+                            .unwrap_or(if node.stake > 0 { 0 } else { 1 }),
+                    };
                 (node.id, state)
             })
             .collect();
@@ -60,7 +65,6 @@ impl TransactionProducer {
             self.clock.wait_forever().await;
             return Ok(());
         };
-        let node_count = self.nodes.len();
         let mut next_tx_id = 0;
         let mut next_tx_at = Timestamp::zero();
         let mut next_input_id = 0;
@@ -71,10 +75,15 @@ impl TransactionProducer {
             next_tx_at = start;
         };
 
+        let node_weights = self.nodes.iter().filter_map(|(id, node)| {
+            let weight = node.tx_generation_weight;
+            (weight != 0).then_some((*id, weight))
+        });
+        let node_lookup = WeightedLookup::new(node_weights);
+
         loop {
-            let node_index = rng.random_range(0..node_count);
-            let node_id = NodeId::new(node_index);
-            let node = self.nodes.get(&node_id).unwrap();
+            let node_id = node_lookup.sample(rng).unwrap();
+            let node = self.nodes.get(node_id).unwrap();
 
             let conflict_fraction = node
                 .tx_conflict_fraction
@@ -116,3 +125,39 @@ impl TransactionProducer {
         }
     }
 }
+
+struct WeightedLookup<T> {
+    elements: Vec<(T, u64)>,
+    total_weight: u64,
+}
+
+impl<T> WeightedLookup<T> {
+    pub fn new(weights: impl IntoIterator<Item = (T, u64)>) -> Self {
+        let elements: Vec<_> = weights
+            .into_iter()
+            .scan(0, |cum_weight, (element, weight)| {
+                *cum_weight += weight;
+                Some((element, *cum_weight))
+            })
+            .collect();
+        let total_weight = elements
+            .last()
+            .map(|(_, weight)| *weight)
+            .unwrap_or_default();
+        Self {
+            elements,
+            total_weight,
+        }
+    }
+
+    pub fn sample<R: Rng>(&self, rng: &mut R) -> Option<&T> {
+        let choice = rng.random_range(0..self.total_weight);
+        match self
+            .elements
+            .binary_search_by_key(&choice, |(_, weight)| *weight)
+        {
+            Ok(index) => self.elements.get(index).map(|(el, _)| el),
+            Err(index) => self.elements.get(index).map(|(el, _)| el),
+        }
+    }
+}