simln-lib/bugfix: return all nodes from ln_node_from_graph

39257da introduced bug where we
would not return all `SimNode`s from `ln_node_from_graph` because
`SimGraph` did not have the excluded nodes. Here we store all
the nodes but exclude them from the capacity based on the
exclude field in the `SimulatedChannel`.

Author: elnosh

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/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,16 @@ 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 channel_capacities(&self) -> Result<u64, LightningError> {
-        let channels = self.network.lock().await.lookup_node(&self.info.pubkey)?;
+        let channels = self
+            .network
+            .lock()
+            .await
+            .lookup_node(&self.info.pubkey)
+            .await?;
         Ok(channels.1.iter().sum())
     }
 
@@ -1018,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>>>,
@@ -1052,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() {
@@ -1069,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],
+                        ));
+                    },
                 }
             }
         }
@@ -1102,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(
@@ -1183,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.
@@ -1232,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> {
@@ -1966,19 +1987,25 @@ mod tests {
             .await
             .unwrap();
 
+        assert!(nodes.len() == 3);
+
         let node_1 = nodes.get(&pk1).unwrap().lock().await;
         let node_1_capacity = node_1.channel_capacities().await.unwrap();
 
-        // Node 1 has 2 channels but one was excluded so here we should only have the one that was
-        // not excluded.
+        // 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);
 
         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
@@ -2048,6 +2075,7 @@ mod tests {
     mock! {
         Network{}
 
+        #[async_trait]
         impl SimNetwork for Network{
             fn dispatch_payment(
                 &mut self,
@@ -2058,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>;
         }
     }