fix(control-plane): debounce ingester pool rebalances

Author: fulmicoton-dd

quickwit/quickwit-cluster/src/change.rs

@@ -328,7 +328,6 @@ async fn try_new_node(
 #[cfg(any(test, feature = "testsuite"))]
 pub mod for_test {
     use quickwit_config::GrpcConfig;
-    use tokio::sync::mpsc;
 
     use super::*;
 
@@ -337,21 +336,6 @@ pub mod for_test {
         ChannelFactory::for_grpc(&GrpcConfig::default())
             .expect("plaintext channel factory should build")
     }
-
-    pub struct ClusterChangeStreamForTest {
-        pub cluster_change_rx: ClusterChangeStream,
-        pub cluster_change_tx: mpsc::UnboundedSender<ClusterChange>,
-    }
-
-    impl Default for ClusterChangeStreamForTest {
-        fn default() -> Self {
-            let (cluster_change_rx, cluster_change_tx) = ClusterChangeStream::new_unbounded();
-            Self {
-                cluster_change_rx,
-                cluster_change_tx,
-            }
-        }
-    }
 }
 
 #[cfg(test)]

quickwit/quickwit-common/src/tower/pool.rs

@@ -21,13 +21,15 @@ use std::hash::Hash;
 use std::sync::{Arc, RwLock};
 
 use futures::{Stream, StreamExt};
+use tokio::sync::watch;
 
 use super::Change;
 
 /// A pool of `V` values identified by `K` keys. The pool can be updated manually by calling the
 /// `add/remove` methods or by listening to a stream of changes.
 pub struct Pool<K, V> {
     pool: Arc<RwLock<HashMap<K, V>>>,
+    generation_tx: watch::Sender<u64>,
 }
 
 impl<K, V> fmt::Debug for Pool<K, V>
@@ -44,6 +46,7 @@ impl<K, V> Clone for Pool<K, V> {
     fn clone(&self) -> Self {
         Self {
             pool: self.pool.clone(),
+            generation_tx: self.generation_tx.clone(),
         }
     }
 }
@@ -52,8 +55,10 @@ impl<K, V> Default for Pool<K, V>
 where K: Eq + PartialEq + Hash
 {
     fn default() -> Self {
+        let (generation_tx, _generation_rx) = watch::channel(0);
         Self {
             pool: Arc::new(RwLock::new(HashMap::default())),
+            generation_tx,
         }
     }
 }
@@ -86,6 +91,24 @@ where
         tokio::spawn(future);
     }
 
+    /// Returns the current pool generation.
+    pub fn generation(&self) -> u64 {
+        *self.generation_tx.borrow()
+    }
+
+    /// Subscribes to pool generation changes.
+    pub fn subscribe(&self) -> watch::Receiver<u64> {
+        self.generation_tx.subscribe()
+    }
+
+    fn increment_generation(&self) {
+        self.generation_tx.send_modify(|generation| {
+            *generation = generation
+                .checked_add(1)
+                .expect("pool generation should not overflow")
+        });
+    }
+
     /// Returns whether the pool is empty.
     pub fn is_empty(&self) -> bool {
         self.pool
@@ -180,14 +203,19 @@ where
             .write()
             .expect("lock should not be poisoned")
             .insert(key, service);
+        self.increment_generation();
     }
 
     /// Removes a value from the pool.
     fn remove(&self, key: &K) {
-        self.pool
+        let removed = self
+            .pool
             .write()
             .expect("lock should not be poisoned")
             .remove(key);
+        if removed.is_some() {
+            self.increment_generation();
+        }
     }
 }
 
@@ -196,8 +224,10 @@ where K: Eq + PartialEq + Hash
 {
     fn from_iter<I>(iter: I) -> Self
     where I: IntoIterator<Item = (K, V)> {
+        let (generation_tx, _generation_rx) = watch::channel(0);
         Self {
             pool: Arc::new(RwLock::new(HashMap::from_iter(iter))),
+            generation_tx,
         }
     }
 }
@@ -210,6 +240,32 @@ mod tests {
 
     use super::*;
 
+    #[tokio::test]
+    async fn test_pool_generation() {
+        let pool = Pool::default();
+        let mut generation_rx = pool.subscribe();
+
+        assert_eq!(pool.generation(), 0);
+
+        pool.insert(1, 11);
+        generation_rx.changed().await.unwrap();
+        assert_eq!(*generation_rx.borrow_and_update(), 1);
+        assert_eq!(pool.get(&1), Some(11));
+
+        pool.insert(1, 12);
+        generation_rx.changed().await.unwrap();
+        assert_eq!(*generation_rx.borrow_and_update(), 2);
+        assert_eq!(pool.get(&1), Some(12));
+
+        pool.remove(&2);
+        assert_eq!(pool.generation(), 2);
+
+        pool.remove(&1);
+        generation_rx.changed().await.unwrap();
+        assert_eq!(*generation_rx.borrow_and_update(), 3);
+        assert!(!pool.contains_key(&1));
+    }
+
     #[tokio::test]
     async fn test_pool() {
         let (change_stream_tx, change_stream_rx) = tokio::sync::mpsc::channel(10);