Expose live nodes watch::Receiver

quickwit-oss · Feb 23, 2024 · 45ef977 · 45ef977
1 parent fcdd0d0
commit 45ef977
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/chitchat/src/lib.rs b/chitchat/src/lib.rs
@@ -257,10 +257,14 @@ impl Chitchat {
     /// - updates its max version
     ///
     /// Heartbeats are not notified.
-    pub fn live_nodes_watcher(&self) -> WatchStream<BTreeMap<ChitchatId, NodeState>> {
+    pub fn live_nodes_watch_stream(&self) -> WatchStream<BTreeMap<ChitchatId, NodeState>> {
         WatchStream::new(self.live_nodes_watcher_rx.clone())
     }
 
+    pub fn live_nodes_watcher(&self) -> watch::Receiver<BTreeMap<ChitchatId, NodeState>> {
+        self.live_nodes_watcher_rx.clone()
+    }
+
     /// Returns the set of nodes considered dead by the failure detector.
     pub fn dead_nodes(&self) -> impl Iterator<Item = &ChitchatId> {
         self.failure_detector.dead_nodes()
@@ -475,7 +479,7 @@ mod tests {
             chitchat
                 .lock()
                 .await
-                .live_nodes_watcher()
+                .live_nodes_watch_stream()
                 .skip_while(|live_nodes| {
                     if live_nodes.len() != expected_nodes.len() {
                         return true;
@@ -530,7 +534,7 @@ mod tests {
         let transport = ChannelTransport::with_mtu(MAX_UDP_DATAGRAM_PAYLOAD_SIZE);
         let nodes = setup_nodes(20001..=20005, &transport).await;
         let node2 = nodes.get(1).unwrap();
-        let mut live_rx = node2.chitchat().lock().await.live_nodes_watcher();
+        let mut live_rx = node2.chitchat().lock().await.live_nodes_watch_stream();
         let live_members = loop {
             let live_members = live_rx.next().await.unwrap();
             if live_members.len() == 5 {
@@ -574,7 +578,7 @@ mod tests {
 
         let mut num_nodes = 0;
         assert!(tokio::time::timeout(Duration::from_secs(1), async {
-            let mut live_rx = nodes[2].chitchat().lock().await.live_nodes_watcher();
+            let mut live_rx = nodes[2].chitchat().lock().await.live_nodes_watch_stream();
             loop {
                 let live_members = live_rx.next().await.unwrap();
                 num_nodes = live_members.len();
@@ -606,7 +610,7 @@ mod tests {
             .self_node_state()
             .set("READY", "true");
 
-        let mut live_rx = nodes[2].chitchat().lock().await.live_nodes_watcher();
+        let mut live_rx = nodes[2].chitchat().lock().await.live_nodes_watch_stream();
         let live_members = loop {
             let live_members = live_rx.next().await.unwrap();
             if live_members.len() == 3 {

diff --git a/chitchat/src/server.rs b/chitchat/src/server.rs
@@ -633,7 +633,7 @@ mod tests {
             .chitchat()
             .lock()
             .await
-            .live_nodes_watcher()
+            .live_nodes_watch_stream()
             .skip_while(|live_nodes| live_nodes.is_empty());
 
         {

diff --git a/chitchat/tests/perf_test.rs b/chitchat/tests/perf_test.rs
@@ -49,7 +49,7 @@ async fn wait_until<P: Fn(&BTreeMap<ChitchatId, NodeState>) -> bool>(
     predicate: P,
 ) -> Duration {
     let start = Instant::now();
-    let mut node_watcher = handle.chitchat().lock().await.live_nodes_watcher();
+    let mut node_watcher = handle.chitchat().lock().await.live_nodes_watch_stream();
     while let Some(nodes) = node_watcher.next().await {
         if predicate(&nodes) {
             break;