simln-lib/refactor: fully deterministic produce events

Author: f3r10

simln-lib/src/lib.rs

@@ -352,7 +352,7 @@ pub trait LightningNode: Send {
 pub struct DestinationGenerationError(String);
 
 /// A trait for selecting destination nodes for payments in the Lightning Network.
-pub trait DestinationGenerator: Send {
+pub trait DestinationGenerator: Send + Sync {
     /// choose_destination picks a destination node within the network, returning the node's information and its
     /// capacity (if available).
     fn choose_destination(
@@ -624,6 +624,34 @@ struct ExecutorKit {
     payment_generator: Box<dyn PaymentGenerator>,
 }
 
+struct PaymentEvent {
+    wait_time: Duration,
+    destination: NodeInfo,
+    current_count: u64,
+    capacity: Option<u64>,
+    executor: ExecutorKit,
+}
+
+impl Ord for PaymentEvent {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.wait_time.cmp(&other.wait_time)
+    }
+}
+
+impl PartialOrd for PaymentEvent {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl PartialEq for PaymentEvent {
+    fn eq(&self, other: &Self) -> bool {
+        self.wait_time == other.wait_time
+    }
+}
+
+impl Eq for PaymentEvent {}
+
 impl<C: Clock + 'static> Simulation<C> {
     pub fn new(
         cfg: SimulationCfg,
@@ -1071,32 +1099,6 @@ impl<C: Clock + 'static> Simulation<C> {
         channels
     }
 
-    fn generate_payments<A: PaymentGenerator + ?Sized>(
-        &self,
-        heap: &mut BinaryHeap<Reverse<(Duration, usize)>>,
-        source: &NodeInfo,
-        node_generator: &A,
-        index: usize,
-    ) -> Result<(), SimulationError> {
-        let mut current_count = 0;
-        loop {
-            let wait = get_payment_delay(current_count, source, node_generator)?;
-            heap.push(Reverse((wait, index)));
-            if node_generator.payment_count().is_none() {
-                return Ok(());
-            }
-            if let Some(c) = node_generator.payment_count() {
-                if c == current_count {
-                    log::info!(
-                        "Payment count has been met for {source}: {c} payments. Stopping the activity."
-                    );
-                    return Ok(());
-                }
-            }
-            current_count += 1;
-        }
-    }
-
     /// Responsible for spinning up producers for a set of activities. Requires that a consumer channel is present
     /// for every source node in the set of executors.
     async fn dispatch_producers(
@@ -1105,74 +1107,108 @@ impl<C: Clock + 'static> Simulation<C> {
         producer_channels: HashMap<PublicKey, Sender<SimulationEvent>>,
         tasks: &TaskTracker,
     ) -> Result<(), SimulationError> {
-        let mut heap: BinaryHeap<Reverse<(Duration, usize)>> = BinaryHeap::new();
-        for (index, executor) in executors.iter().enumerate() {
-            self.generate_payments(
-                &mut heap,
-                &executor.source_info,
-                executor.payment_generator.as_ref(),
-                index,
-            )?;
+        let mut heap: BinaryHeap<Reverse<PaymentEvent>> = BinaryHeap::new();
+        for executor in executors {
+            generate_payments(&mut heap, executor, 0)?;
         }
-        while let Some(Reverse((wait, index))) = heap.pop() {
-            let executor = executors
-                .get(index)
-                .ok_or(SimulationError::MissingNodeError(
-                    "Missing executor".to_string(),
-                ))?;
-            let sender = producer_channels.get(&executor.source_info.pubkey).ok_or(
-                SimulationError::RandomActivityError(RandomActivityError::ValueError(format!(
-                    "Activity producer for: {} not found.",
-                    executor.source_info.pubkey,
-                ))),
-            )?;
-            let (destination, capacity) = executor
-                .network_generator
-                .choose_destination(executor.source_info.pubkey)
-                .map_err(SimulationError::DestinationGenerationError)?;
-
-            // pe: produce events
-            let pe_shutdown = self.shutdown_trigger.clone();
-            let pe_listener = self.shutdown_listener.clone();
-            let pe_sender = sender.clone();
-            let clock = self.clock.clone();
-            let source = executor.source_info.clone();
-            // Only proceed with a payment if the amount is non-zero, otherwise skip this round. If we can't get
-            // a payment amount something has gone wrong (because we should have validated that we can always
-            // generate amounts), so we exit.
-            let amount = match executor.payment_generator.payment_amount(capacity) {
-                Ok(amt) => {
-                    if amt == 0 {
-                        log::debug!("Skipping zero amount payment for {source} -> {destination}.");
-                        continue;
-                    }
-                    amt
-                },
-                Err(e) => {
-                    return Err(SimulationError::PaymentGenerationError(e));
+
+        let listener = self.shutdown_listener.clone();
+        let shutdown = self.shutdown_trigger.clone();
+        let clock = self.clock.clone();
+        let t = tasks.clone();
+
+        tasks.spawn(async move {
+        loop {
+            select! {
+                biased;
+                _ = listener.clone()  => {
+                    return Ok(())
                 },
-            };
-            tasks.spawn(async move {
-                log::info!("Starting activity producer for {}.", source);
-
-                if let Err(e) = produce_events(
-                    source.clone(),
-                    clock,
-                    pe_sender,
-                    pe_listener,
-                    wait,
-                    destination,
-                    amount,
-                )
-                .await
-                {
-                    pe_shutdown.trigger();
-                    log::debug!("Activity producer for {source} exited with error {e}.");
-                } else {
-                    log::debug!("Activity producer for {source} completed successfully.");
+                _ = async {} => {
+                    match heap.pop() {
+                        Some(Reverse(PaymentEvent {
+                            wait_time,
+                            destination,
+                            current_count,
+                            capacity,
+                            executor,
+                        })) => {
+
+                            if let Some(c) = executor.payment_generator.payment_count() {
+                                if c == current_count {
+                                    log::info!(
+                                        "Payment count has been met for {}: {c} payments. Stopping the activity."
+                                    , executor.source_info);
+                                    return Ok(());
+                                }
+                            }
+                            let sender = producer_channels.get(&executor.source_info.pubkey).ok_or(
+                                SimulationError::RandomActivityError(RandomActivityError::ValueError(
+                                    format!(
+                                        "Activity producer for: {} not found.",
+                                        executor.source_info.pubkey,
+                                    ),
+                                )),
+                            )?;
+
+                            // pe: produce events
+                            let pe_shutdown = shutdown.clone();
+                            let pe_listener = listener.clone();
+                            let pe_sender = sender.clone();
+                            let pe_clock = clock.clone();
+                            let source = executor.source_info.clone();
+                            // Only proceed with a payment if the amount is non-zero, otherwise skip this round. If we can't get
+                            // a payment amount something has gone wrong (because we should have validated that we can always
+                            // generate amounts), so we exit.
+                            let payment_amount = executor.payment_generator.payment_amount(capacity);
+                            let amount = match payment_amount {
+                                Ok(amt) => {
+                                    if amt == 0 {
+                                        log::debug!(
+                                            "Skipping zero amount payment for {source} -> {destination}."
+                                        );
+                                        generate_payments(&mut heap, executor, current_count)?;
+                                        continue;
+                                    } else {
+                                        generate_payments(&mut heap, executor, current_count + 1)?;
+                                    }
+                                    amt
+                                },
+                                Err(e) => {
+                                    return Err(SimulationError::PaymentGenerationError(e));
+                                },
+                            };
+                            // Wait until our time to next payment has elapsed then execute a random amount payment to a random
+                            // destination.
+                            pe_clock.sleep(wait_time).await;
+                            t.spawn(async move {
+                                log::info!("Starting activity producer for {}.", source);
+
+                                if let Err(e) = produce_events(
+                                    source.clone(),
+                                    pe_sender,
+                                    pe_listener,
+                                    destination,
+                                    amount,
+                                )
+                                .await
+                                {
+                                    pe_shutdown.trigger();
+                                    log::debug!("Activity producer for {source} exited with error {e}.");
+                                } else {
+                                    log::debug!("Activity producer for {source} completed successfully.");
+                                }
+                            });
+                        },
+                        None => {
+                            break
+                        },
+                    }
                 }
-            });
-        }
+            }
+        };
+        Ok(())
+        });
         Ok(())
     }
 }
@@ -1263,10 +1299,8 @@ async fn consume_events(
 /// exit if other threads signal that they have errored out.
 async fn produce_events(
     source: NodeInfo,
-    clock: Arc<dyn Clock>,
     sender: Sender<SimulationEvent>,
     listener: Listener,
-    wait: Duration,
     destination: NodeInfo,
     amount: u64,
 ) -> Result<(), SimulationError> {
@@ -1275,9 +1309,7 @@ async fn produce_events(
         _ = listener.clone() => {
             return Ok(());
         },
-        // Wait until our time to next payment has elapsed then execute a random amount payment to a random
-        // destination.
-        _ = clock.sleep(wait) => {
+        _ = async {} => {
             log::debug!("Generated payment: {source} -> {}: {amount} msat.", destination);
 
             // Send the payment, exiting if we can no longer send to the consumer.
@@ -1578,6 +1610,31 @@ async fn track_payment_result(
     Ok(())
 }
 
+fn generate_payments(
+    heap: &mut BinaryHeap<Reverse<PaymentEvent>>,
+    executor: ExecutorKit,
+    current_count: u64,
+) -> Result<(), SimulationError> {
+    let wait_time = get_payment_delay(
+        current_count,
+        &executor.source_info,
+        executor.payment_generator.as_ref(),
+    )?;
+    let (destination, capacity) = executor
+        .network_generator
+        .choose_destination(executor.source_info.pubkey)
+        .map_err(SimulationError::DestinationGenerationError)?;
+    let payment_event = PaymentEvent {
+        wait_time,
+        destination,
+        current_count,
+        capacity,
+        executor,
+    };
+    heap.push(Reverse(payment_event));
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{