Merge pull request #306 from elnosh/fix-excluded-nodes

return all `SimNode`s from `ln_node_from_graph`

Author: carlaKC

sim-cli/src/parsing.rs

@@ -301,7 +301,7 @@ pub async fn create_simulation_with_network(
     ));
 
     // Copy all simulated channels into a read-only routing graph, allowing to pathfind for
-    // individual payments without locking th simulation graph (this is a duplication of the channels,
+    // individual payments without locking the simulation graph (this is a duplication of the channels,
     // but the performance tradeoff is worthwhile for concurrent pathfinding).
     let routing_graph = Arc::new(
         populate_network_graph(channels, clock.clone())
@@ -312,7 +312,7 @@ pub async fn create_simulation_with_network(
     // custom actions on the simulated network. For the nodes we'll pass our simulation, cast them
     // to a dyn trait and exclude any nodes that shouldn't be included in random activity
     // generation.
-    let nodes = ln_node_from_graph(simulation_graph.clone(), routing_graph, clock.clone()).await?;
+    let nodes = ln_node_from_graph(simulation_graph, routing_graph, clock.clone()).await?;
     let mut nodes_dyn: HashMap<_, Arc<Mutex<dyn LightningNode>>> = nodes
         .iter()
         .map(|(pk, node)| (*pk, Arc::clone(node) as Arc<Mutex<dyn LightningNode>>))

simln-lib/src/cln.rs

@@ -264,10 +264,10 @@ impl LightningNode for ClnNode {
         }
     }
 
-    async fn list_channels(&self) -> Result<Vec<u64>, LightningError> {
+    async fn channel_capacities(&self) -> Result<u64, LightningError> {
         let mut node_channels = self.node_channels(true).await?;
         node_channels.extend(self.node_channels(false).await?);
-        Ok(node_channels)
+        Ok(node_channels.iter().sum())
     }
 
     async fn get_graph(&self) -> Result<Graph, LightningError> {

simln-lib/src/eclair.rs

@@ -222,7 +222,7 @@ impl LightningNode for EclairNode {
         })
     }
 
-    async fn list_channels(&self) -> Result<Vec<u64>, LightningError> {
+    async fn channel_capacities(&self) -> Result<u64, LightningError> {
         let client = self.client.lock().await;
         let channels: ChannelsResponse = client
             .request("channels", None)
@@ -242,7 +242,7 @@ impl LightningNode for EclairNode {
             })
             .collect();
 
-        Ok(capacities_msat)
+        Ok(capacities_msat.iter().sum())
     }
 
     async fn get_graph(&self) -> Result<Graph, LightningError> {

simln-lib/src/lib.rs

@@ -343,9 +343,9 @@ pub trait LightningNode: Send {
     ) -> Result<PaymentResult, LightningError>;
     /// Gets information on a specific node.
     async fn get_node_info(&self, node_id: &PublicKey) -> Result<NodeInfo, LightningError>;
-    /// Lists all channels, at present only returns a vector of channel capacities in msat because no further
-    /// information is required.
-    async fn list_channels(&self) -> Result<Vec<u64>, LightningError>;
+    /// Sum of channel capacities. This is used when running with random activity generation to
+    /// determine how much the node will send per month.
+    async fn channel_capacities(&self) -> Result<u64, LightningError>;
     /// Get the network graph from the point of view of a given node.
     async fn get_graph(&self) -> Result<Graph, LightningError>;
 }
@@ -1003,7 +1003,7 @@ impl<C: Clock + 'static> Simulation<C> {
         // While we're at it, we get the node info and store it with capacity to create activity generators in our
         // second pass.
         for (pk, node) in self.nodes.iter() {
-            let chan_capacity = node.lock().await.list_channels().await?.iter().sum::<u64>();
+            let chan_capacity = node.lock().await.channel_capacities().await?;
 
             if let Err(e) = RandomPaymentActivity::validate_capacity(
                 chan_capacity,
@@ -1935,8 +1935,8 @@ mod tests {
             .expect_get_network()
             .returning(|| Network::Regtest);
         mock_node
-            .expect_list_channels()
-            .returning(|| Ok(vec![100_000_000]));
+            .expect_channel_capacities()
+            .returning(|| Ok(100_000_000));
         mock_node
             .expect_get_node_info()
             .returning(move |_| Ok(node_info.clone()));

simln-lib/src/lnd.rs

@@ -252,7 +252,7 @@ impl LightningNode for LndNode {
         }
     }
 
-    async fn list_channels(&self) -> Result<Vec<u64>, LightningError> {
+    async fn channel_capacities(&self) -> Result<u64, LightningError> {
         let mut client = self.client.lock().await;
         let channels = client
             .lightning()
@@ -268,7 +268,7 @@ impl LightningNode for LndNode {
             .channels
             .iter()
             .map(|channel| 1000 * channel.capacity as u64)
-            .collect())
+            .sum())
     }
 
     async fn get_graph(&self) -> Result<Graph, LightningError> {

simln-lib/src/sim_node.rs

@@ -477,6 +477,7 @@ impl SimulatedChannel {
 
 /// SimNetwork represents a high level network coordinator that is responsible for the task of actually propagating
 /// payments through the simulated network.
+#[async_trait]
 pub trait SimNetwork: Send + Sync {
     /// Sends payments over the route provided through the network, reporting the final payment outcome to the sender
     /// channel provided.
@@ -490,7 +491,7 @@ pub trait SimNetwork: Send + Sync {
     );
 
     /// Looks up a node in the simulated network and a list of its channel capacities.
-    fn lookup_node(&self, node: &PublicKey) -> Result<(NodeInfo, Vec<u64>), LightningError>;
+    async fn lookup_node(&self, node: &PublicKey) -> Result<(NodeInfo, Vec<u64>), LightningError>;
     /// Lists all nodes in the simulated network.
     fn list_nodes(&self) -> Vec<NodeInfo>;
 }
@@ -794,11 +795,17 @@ impl<T: SimNetwork, C: Clock> LightningNode for SimNode<T, C> {
     }
 
     async fn get_node_info(&self, node_id: &PublicKey) -> Result<NodeInfo, LightningError> {
-        Ok(self.network.lock().await.lookup_node(node_id)?.0)
+        Ok(self.network.lock().await.lookup_node(node_id).await?.0)
     }
 
-    async fn list_channels(&self) -> Result<Vec<u64>, LightningError> {
-        Ok(self.network.lock().await.lookup_node(&self.info.pubkey)?.1)
+    async fn channel_capacities(&self) -> Result<u64, LightningError> {
+        let channels = self
+            .network
+            .lock()
+            .await
+            .lookup_node(&self.info.pubkey)
+            .await?;
+        Ok(channels.1.iter().sum())
     }
 
     async fn get_graph(&self) -> Result<Graph, LightningError> {
@@ -1017,9 +1024,9 @@ async fn handle_intercepted_htlc(
 
 /// Graph is the top level struct that is used to coordinate simulation of lightning nodes.
 pub struct SimGraph {
-    /// nodes caches the list of nodes in the network with a vector of their channel capacities, only used for quick
+    /// nodes caches the list of nodes in the network with a vector of their channel ids, only used for quick
     /// lookup.
-    nodes: HashMap<PublicKey, (NodeInfo, Vec<u64>)>,
+    nodes: HashMap<PublicKey, (NodeInfo, Vec<ShortChannelID>)>,
 
     /// channels maps the scid of a channel to its current simulation state.
     channels: Arc<Mutex<HashMap<ShortChannelID, SimulatedChannel>>>,
@@ -1051,7 +1058,7 @@ impl SimGraph {
         default_custom_records: CustomRecords,
         shutdown_signal: (Trigger, Listener),
     ) -> Result<Self, SimulationError> {
-        let mut nodes: HashMap<PublicKey, (NodeInfo, Vec<u64>)> = HashMap::new();
+        let mut nodes: HashMap<PublicKey, (NodeInfo, Vec<ShortChannelID>)> = HashMap::new();
         let mut channels = HashMap::new();
 
         for channel in graph_channels.iter() {
@@ -1068,18 +1075,16 @@ impl SimGraph {
                 Entry::Vacant(v) => v.insert(channel.clone()),
             };
 
-            if !channel.exclude_capacity {
-                // It's okay to have duplicate pubkeys because one node can have many channels.
-                for info in [&channel.node_1.policy, &channel.node_2.policy] {
-                    match nodes.entry(info.pubkey) {
-                        Entry::Occupied(o) => o.into_mut().1.push(channel.capacity_msat),
-                        Entry::Vacant(v) => {
-                            v.insert((
-                                node_info(info.pubkey, info.alias.clone()),
-                                vec![channel.capacity_msat],
-                            ));
-                        },
-                    }
+            // It's okay to have duplicate pubkeys because one node can have many channels.
+            for info in [&channel.node_1.policy, &channel.node_2.policy] {
+                match nodes.entry(info.pubkey) {
+                    Entry::Occupied(o) => o.into_mut().1.push(channel.short_channel_id),
+                    Entry::Vacant(v) => {
+                        v.insert((
+                            node_info(info.pubkey, info.alias.clone()),
+                            vec![channel.short_channel_id],
+                        ));
+                    },
                 }
             }
         }
@@ -1101,9 +1106,11 @@ pub async fn ln_node_from_graph<C: Clock>(
     routing_graph: Arc<LdkNetworkGraph>,
     clock: Arc<C>,
 ) -> Result<HashMap<PublicKey, Arc<Mutex<SimNode<SimGraph, C>>>>, LightningError> {
-    let mut nodes: HashMap<PublicKey, Arc<Mutex<SimNode<SimGraph, C>>>> = HashMap::new();
+    let sim_graph = graph.lock().await;
+    let mut nodes: HashMap<PublicKey, Arc<Mutex<SimNode<SimGraph, C>>>> =
+        HashMap::with_capacity(sim_graph.nodes.len());
 
-    for node in graph.lock().await.nodes.iter() {
+    for node in sim_graph.nodes.iter() {
         nodes.insert(
             *node.0,
             Arc::new(Mutex::new(SimNode::new(
@@ -1182,6 +1189,7 @@ pub fn populate_network_graph<C: Clock>(
     Ok(graph)
 }
 
+#[async_trait]
 impl SimNetwork for SimGraph {
     /// dispatch_payment asynchronously propagates a payment through the simulated network, returning a tracking
     /// channel that can be used to obtain the result of the payment. At present, MPP payments are not supported.
@@ -1231,13 +1239,27 @@ impl SimNetwork for SimGraph {
     }
 
     /// lookup_node fetches a node's information and channel capacities.
-    fn lookup_node(&self, node: &PublicKey) -> Result<(NodeInfo, Vec<u64>), LightningError> {
-        match self.nodes.get(node) {
-            Some(node) => Ok(node.clone()),
-            None => Err(LightningError::GetNodeInfoError(
-                "Node not found".to_string(),
-            )),
-        }
+    async fn lookup_node(&self, node: &PublicKey) -> Result<(NodeInfo, Vec<u64>), LightningError> {
+        let node_info = match self.nodes.get(node) {
+            Some(node) => node.clone(),
+            None => {
+                return Err(LightningError::GetNodeInfoError(format!(
+                    "Node {} not found",
+                    node
+                )))
+            },
+        };
+
+        let channels = self.channels.lock().await;
+        let capacities: Vec<u64> = node_info
+            .1
+            .iter()
+            .filter_map(|scid| channels.get(scid))
+            .filter(|channel| !channel.exclude_capacity)
+            .map(|channel| channel.capacity_msat)
+            .collect();
+
+        Ok((node_info.0, capacities))
     }
 
     fn list_nodes(&self) -> Vec<NodeInfo> {
@@ -1965,34 +1987,25 @@ mod tests {
             .await
             .unwrap();
 
-        let node_1_channels = nodes
-            .get(&pk1)
-            .unwrap()
-            .lock()
-            .await
-            .list_channels()
-            .await
-            .unwrap();
+        assert!(nodes.len() == 3);
 
-        // Node 1 has 2 channels but one was excluded so here we should only have the one that was
-        // not excluded.
-        assert!(node_1_channels.len() == 1);
-        assert!(node_1_channels[0] == capacity_1);
+        let node_1 = nodes.get(&pk1).unwrap().lock().await;
+        let node_1_capacity = node_1.channel_capacities().await.unwrap();
 
-        let node_2_channels = nodes
-            .get(&pk2)
-            .unwrap()
-            .lock()
-            .await
-            .list_channels()
-            .await
-            .unwrap();
+        // Node 1 has 2 channels but one was excluded so here we should only have the capacity of
+        // the channel that was not excluded.
+        assert!(node_1_capacity == capacity_1);
 
-        assert!(node_2_channels.len() == 1);
-        assert!(node_2_channels[0] == capacity_1);
+        let node_2 = nodes.get(&pk2).unwrap().lock().await;
+        let node_2_capacity = node_2.channel_capacities().await.unwrap();
+        assert!(node_2_capacity == capacity_1);
 
-        // Node 3's only channel was excluded so it won't be present here.
-        assert!(!nodes.contains_key(&pk3));
+        // Node 3 should be returned from ln_node_from_graph but it won't have any channel capacity
+        // present because its only channel was excluded.
+        let node_3 = nodes.get(&pk3);
+        assert!(node_3.is_some());
+        let node_3 = node_3.unwrap().lock().await;
+        assert!(node_3.channel_capacities().await.unwrap() == 0);
     }
 
     /// Tests basic functionality of a `SimulatedChannel` but does no endeavor to test the underlying
@@ -2062,6 +2075,7 @@ mod tests {
     mock! {
         Network{}
 
+        #[async_trait]
         impl SimNetwork for Network{
             fn dispatch_payment(
                 &mut self,
@@ -2072,7 +2086,7 @@ mod tests {
                 sender: Sender<Result<PaymentResult, LightningError>>,
             );
 
-            fn lookup_node(&self, node: &PublicKey) -> Result<(NodeInfo, Vec<u64>), LightningError>;
+            async fn lookup_node(&self, node: &PublicKey) -> Result<(NodeInfo, Vec<u64>), LightningError>;
             fn list_nodes(&self) -> Vec<NodeInfo>;
         }
     }
@@ -2103,12 +2117,17 @@ mod tests {
             .lock()
             .await
             .expect_lookup_node()
-            .returning(move |_| Ok((node_info(lookup_pk, String::default()), vec![1, 2, 3])));
+            .returning(move |_| {
+                Ok((
+                    node_info(lookup_pk, String::default()),
+                    vec![10_000, 20_000, 10_000],
+                ))
+            });
 
         // Assert that we get three channels from the mock.
         let node_info = node.get_node_info(&lookup_pk).await.unwrap();
         assert_eq!(lookup_pk, node_info.pubkey);
-        assert_eq!(node.list_channels().await.unwrap().len(), 3);
+        assert_eq!(node.channel_capacities().await.unwrap(), 40_000);
 
         // Next, we're going to test handling of in-flight payments. To do this, we'll mock out calls to our dispatch
         // function to send different results depending on the destination.

simln-lib/src/test_utils.rs

@@ -89,7 +89,7 @@ mock! {
                 shutdown: triggered::Listener,
             ) -> Result<crate::PaymentResult, LightningError>;
         async fn get_node_info(&self, node_id: &PublicKey) -> Result<NodeInfo, LightningError>;
-        async fn list_channels(&self) -> Result<Vec<u64>, LightningError>;
+        async fn channel_capacities(&self) -> Result<u64, LightningError>;
         async fn get_graph(&self) -> Result<Graph, LightningError>;
     }
 }