Use BTreeMap to store graph nodes and channels

Author: jjyr

Cargo.lock

@@ -57,7 +57,6 @@ thiserror = "1.0.58"
 tokio-util = {version = "0.7.10", features = ["rt"]}
 tracing = "0.1"
 tracing-subscriber = {version = "0.3", features = ["env-filter"]}
-indexmap = "2.11.0"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 biscuit-auth = "6.0.0-beta.3"

crates/fiber-lib/Cargo.toml

@@ -22,11 +22,11 @@ use crate::fiber::types::PaymentHopData;
 use crate::invoice::CkbInvoice;
 use crate::now_timestamp_as_millis_u64;
 use ckb_types::packed::{OutPoint, Script};
-use indexmap::IndexMap;
 use parking_lot::Mutex;
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
+use std::ops::Bound;
 use std::sync::Arc;
 use tentacle::multiaddr::MultiAddr;
 use tentacle::secio::PeerId;
@@ -448,9 +448,9 @@ pub struct NetworkGraph<S> {
     // The count of private channels
     pub(crate) private_channels_count: usize,
     // All the channels in the network.
-    pub(crate) channels: IndexMap<OutPoint, ChannelInfo>,
+    pub(crate) channels: BTreeMap<OutPoint, ChannelInfo>,
     // All the nodes in the network.
-    pub(crate) nodes: IndexMap<Pubkey, NodeInfo>,
+    pub(crate) nodes: BTreeMap<Pubkey, NodeInfo>,
 
     // Channel stats map, used to track the attempts for each channel,
     // this information is used to HELP the path finding algorithm for better routing in two ways:
@@ -528,9 +528,9 @@ where
             always_process_gossip_message: false,
             source,
             private_channels_count: 0,
-            channels: IndexMap::new(),
+            channels: BTreeMap::new(),
             channel_stats: Default::default(),
-            nodes: IndexMap::new(),
+            nodes: BTreeMap::new(),
             latest_cursor: Cursor::default(),
             store: store.clone(),
             history: PaymentHistory::new(source, None, store),
@@ -620,7 +620,7 @@ where
                     .insert(channel_info.channel_outpoint.clone(), channel_info);
             }
             OwnedChannelUpdateEvent::Down(channel_outpoint) => {
-                if let Some(channel_info) = self.channels.swap_remove(&channel_outpoint) {
+                if let Some(channel_info) = self.channels.remove(&channel_outpoint) {
                     if !channel_info.is_public {
                         self.private_channels_count -= 1;
                     }
@@ -851,14 +851,21 @@ where
         self.nodes.values()
     }
 
-    pub fn get_nodes_with_params(&self, limit: usize, after: Option<u64>) -> Vec<NodeInfo> {
-        let after = after.unwrap_or_default();
-        self.nodes
-            .iter()
-            .skip(after as usize)
-            .take(limit)
-            .map(|(_pubkey, node)| node.to_owned())
-            .collect()
+    pub fn get_nodes_with_params(&self, limit: usize, after: Option<Pubkey>) -> Vec<NodeInfo> {
+        match after {
+            Some(after) => self
+                .nodes
+                .range((Bound::Excluded(after), Bound::Unbounded))
+                .take(limit)
+                .map(|(_pubkey, node)| node.to_owned())
+                .collect(),
+            None => self
+                .nodes
+                .iter()
+                .take(limit)
+                .map(|(_pubkey, node)| node.to_owned())
+                .collect(),
+        }
     }
 
     pub fn get_node(&self, node_id: &Pubkey) -> Option<&NodeInfo> {
@@ -904,20 +911,32 @@ where
         }
     }
 
-    pub fn get_channels_with_params(&self, limit: usize, after: Option<u64>) -> Vec<ChannelInfo> {
-        let after = after.unwrap_or_default();
-        self.channels
-            .iter()
-            .skip(after as usize)
-            .take(limit)
-            .filter_map(|(_out_point, channel_info)| {
-                if channel_info.is_public {
-                    Some(channel_info.to_owned())
-                } else {
-                    None
-                }
-            })
-            .collect()
+    pub fn get_channels_with_params(
+        &self,
+        limit: usize,
+        after: Option<OutPoint>,
+    ) -> Vec<ChannelInfo> {
+        let filter = |(_out_point, channel_info): (&OutPoint, &ChannelInfo)| {
+            if channel_info.is_public {
+                Some(channel_info.to_owned())
+            } else {
+                None
+            }
+        };
+        match after {
+            Some(after) => self
+                .channels
+                .range((Bound::Excluded(after), Bound::Unbounded))
+                .take(limit)
+                .filter_map(filter)
+                .collect(),
+            None => self
+                .channels
+                .iter()
+                .take(limit)
+                .filter_map(filter)
+                .collect(),
+        }
     }
 
     pub fn get_channels_by_peer(&self, node_id: Pubkey) -> impl Iterator<Item = &ChannelInfo> {

crates/fiber-lib/src/fiber/graph.rs

@@ -287,7 +287,8 @@ async fn test_rpc_graph() {
         .all(|n| n.version == *env!("CARGO_PKG_VERSION")));
     assert!(!graph_nodes.nodes[0].features.is_empty());
 
-    let graph_nodes: GraphNodesResult = node_0
+    // query nodes by page
+    let graph_nodes_p1: GraphNodesResult = node_0
         .send_rpc_request(
             "graph_nodes",
             GraphNodesParams {
@@ -298,8 +299,33 @@ async fn test_rpc_graph() {
         .await
         .unwrap();
 
-    assert_eq!(graph_nodes.total_count.value(), 2);
-    assert_eq!(graph_nodes.nodes.len(), 1);
+    assert_eq!(graph_nodes_p1.total_count.value(), 2);
+    assert_eq!(graph_nodes_p1.nodes.len(), 1);
+
+    let graph_nodes_p2: GraphNodesResult = node_0
+        .send_rpc_request(
+            "graph_nodes",
+            GraphNodesParams {
+                limit: Some(1),
+                after: Some(graph_nodes_p1.last_cursor),
+            },
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(graph_nodes_p2.total_count.value(), 2);
+    assert_eq!(graph_nodes_p2.nodes.len(), 1);
+
+    let mut nodes = graph_nodes_p1.nodes.clone();
+    nodes.extend(graph_nodes_p2.nodes.clone());
+    assert_eq!(
+        graph_nodes
+            .nodes
+            .iter()
+            .map(|n| n.node_id)
+            .collect::<Vec<_>>(),
+        nodes.iter().map(|n| n.node_id).collect::<Vec<_>>(),
+    );
 
     let graph_channels: GraphChannelsResult = node_0
         .send_rpc_request(
@@ -315,6 +341,21 @@ async fn test_rpc_graph() {
     // only public channels
     assert_eq!(graph_channels.total_count.value(), 1);
     assert_eq!(graph_channels.channels.len(), 1);
+
+    // next query is empty
+    let graph_channels: GraphChannelsResult = node_0
+        .send_rpc_request(
+            "graph_channels",
+            GraphChannelsParams {
+                limit: Some(1),
+                after: Some(graph_channels.last_cursor),
+            },
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(graph_channels.total_count.value(), 1);
+    assert!(graph_channels.channels.is_empty());
 }
 
 #[tokio::test]

crates/fiber-lib/src/fiber/tests/rpc.rs

@@ -18,6 +18,7 @@ use ckb_types::H256;
 use jsonrpsee::proc_macros::rpc;
 use jsonrpsee::types::ErrorObjectOwned;
 
+use molecule::prelude::Entity;
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 use std::sync::Arc;
@@ -338,15 +339,16 @@ where
         let network_graph = self.network_graph.read().await;
         let default_max_limit = 500;
         let limit = params.limit.unwrap_or(default_max_limit) as usize;
-        let after = params.after.as_ref().map(|after| {
-            let buf: [u8; 8] = after.as_bytes().try_into().unwrap_or_default();
-            u64::from_le_bytes(buf)
-        });
+        let after = params
+            .after
+            .as_ref()
+            .and_then(|after| Pubkey::from_slice(after.as_bytes()).ok());
         let nodes = network_graph.get_nodes_with_params(limit, after);
         let last_cursor = JsonBytes::from_vec(
-            (after.unwrap_or_default() + nodes.len() as u64)
-                .to_le_bytes()
-                .to_vec(),
+            nodes
+                .last()
+                .map(|n| n.node_id.serialize().to_vec())
+                .unwrap_or_default(),
         );
         let nodes = nodes.into_iter().map(Into::into).collect();
         let total_count = (network_graph.nodes.len() as u64).into();
@@ -365,16 +367,20 @@ where
         let default_max_limit = 500;
         let network_graph = self.network_graph.read().await;
         let limit = params.limit.unwrap_or(default_max_limit) as usize;
-        let after = params.after.as_ref().map(|after| {
-            let buf: [u8; 8] = after.as_bytes().try_into().unwrap_or_default();
-            u64::from_le_bytes(buf)
+        let after = params.after.as_ref().and_then(|after| {
+            if after.is_empty() {
+                None
+            } else {
+                OutPoint::from_slice(after.as_bytes()).ok()
+            }
         });
 
         let channels = network_graph.get_channels_with_params(limit, after);
         let last_cursor = JsonBytes::from_vec(
-            (after.unwrap_or_default() + channels.len() as u64)
-                .to_le_bytes()
-                .to_vec(),
+            channels
+                .last()
+                .map(|c| c.channel_outpoint.as_slice().to_vec())
+                .unwrap_or_default(),
         );
 
         let channels = channels.into_iter().map(Into::into).collect();