Remove departed ingesters from routing table

quickwit/quickwit-ingest/src/ingest_v2/mod.rs

@@ -45,6 +45,7 @@ pub use broadcast::{
 use bytes::buf::Writer;
 use bytes::{BufMut, BytesMut};
 use bytesize::ByteSize;
+use quickwit_common::pubsub::Event;
 use quickwit_common::tower::Pool;
 use quickwit_proto::ingest::ingester::{IngesterServiceClient, IngesterStatus};
 use quickwit_proto::ingest::router::{IngestRequestV2, IngestSubrequest};
@@ -96,6 +97,14 @@ impl IngesterPoolEntry {
 
 pub type IngesterPool = Pool<NodeId, IngesterPoolEntry>;
 
+/// Published when an ingester leaves the cluster.
+#[derive(Debug, Clone)]
+pub struct IngesterLeft {
+    pub node_id: NodeId,
+}
+
+impl Event for IngesterLeft {}
+
 /// Identifies an ingester client, typically a source, for logging and debugging purposes.
 pub type ClientId = String;
 

quickwit/quickwit-ingest/src/ingest_v2/router.rs

@@ -48,7 +48,7 @@ use super::debouncing::{
 use super::ingester::PERSIST_REQUEST_TIMEOUT;
 use super::routing_table::RoutingTable;
 use super::workbench::IngestWorkbench;
-use super::{IngesterPool, pending_subrequests};
+use super::{IngesterLeft, IngesterPool, pending_subrequests};
 use crate::get_ingest_router_buffer_size;
 use crate::ingest_v2::metrics::{
     INGEST_ATTEMPTS, INGEST_RESULT_CIRCUIT_BREAKER, INGEST_RESULT_INDEX_NOT_FOUND,
@@ -151,7 +151,10 @@ impl IngestRouter {
     pub fn subscribe(&self) {
         let weak_router_state = WeakRouterState(Arc::downgrade(&self.state));
         self.event_broker
-            .subscribe::<IngesterCapacityScoreUpdate>(weak_router_state)
+            .subscribe::<IngesterCapacityScoreUpdate>(weak_router_state.clone())
+            .forever();
+        self.event_broker
+            .subscribe::<IngesterLeft>(weak_router_state)
             .forever();
     }
 
@@ -613,6 +616,26 @@ impl EventSubscriber<IngesterCapacityScoreUpdate> for WeakRouterState {
     }
 }
 
+/// Clears a departed ingester's routing entries.
+#[async_trait]
+impl EventSubscriber<IngesterLeft> for WeakRouterState {
+    async fn handle_event(&mut self, departed_ingester: IngesterLeft) {
+        let Some(state) = self.0.upgrade() else {
+            return;
+        };
+        let mut state_guard = state.lock().await;
+        let num_entries = state_guard.routing_table.remove_node(&departed_ingester.node_id);
+        drop(state_guard);
+
+        if num_entries > 0 {
+            info!(
+                node_id=%departed_ingester.node_id,
+                "removed ingester from routing table after it left the cluster"
+            );
+        }
+    }
+}
+
 pub(super) struct PersistRequestSummary {
     pub leader_id: NodeId,
     pub subrequest_ids: Vec<SubrequestId>,

quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs

@@ -323,6 +323,18 @@ impl RoutingTable {
         }
         entry.seeded_from_cp = true;
     }
+
+    /// Removes a node that left the cluster from every routing entry.
+    pub fn remove_node(&mut self, node_id: &NodeId) -> usize {
+        let mut num_entries = 0;
+
+        for entry in self.table.values_mut() {
+            if entry.nodes.remove(node_id).is_some() {
+                num_entries += 1;
+            }
+        }
+        num_entries
+    }
 }
 
 #[cfg(test)]
@@ -763,6 +775,35 @@ mod tests {
         assert!(entry.nodes.contains_key("node-4"));
     }
 
+    #[test]
+    fn test_remove_node() {
+        let mut table = RoutingTable::default();
+        let index_uid = IndexUid::for_test("test-index", 0);
+
+        table.apply_capacity_update(
+            NodeId::from_str("node-1"),
+            index_uid.clone(),
+            "test-source".into(),
+            8,
+            3,
+        );
+        table.apply_capacity_update(
+            NodeId::from_str("node-2"),
+            index_uid,
+            "test-source".into(),
+            6,
+            2,
+        );
+        assert_eq!(table.remove_node(&NodeId::from_str("node-1")), 1);
+
+        let entry = table
+            .table
+            .get(&("test-index".to_string(), "test-source".to_string()))
+            .unwrap();
+        assert!(!entry.nodes.contains_key("node-1"));
+        assert!(entry.nodes.contains_key("node-2"));
+    }
+
     #[test]
     fn test_classify_az_locality() {
         let table = RoutingTable::new(Some("az-1".to_string()));

quickwit/quickwit-serve/src/lib.rs

@@ -88,11 +88,11 @@ use quickwit_indexing::actors::{IndexingService, MergeSchedulerService};
 use quickwit_indexing::models::ShardPositionsService;
 use quickwit_indexing::{IndexingSplitCache, start_indexing_service};
 use quickwit_ingest::{
-    GetMemoryCapacity, IngestRequest, IngestRouter, IngestServiceClient, Ingester, IngesterPool,
-    IngesterPoolEntry, LocalShardsUpdate, get_idle_shard_timeout, notify_ingester_decommission,
-    setup_ingester_capacity_update_listener, setup_local_shards_update_listener,
-    start_ingest_api_service, try_get_ingester_status, wait_for_ingester_decommission,
-    wait_for_ingester_status,
+    GetMemoryCapacity, IngestRequest, IngestRouter, IngestServiceClient, Ingester,
+    IngesterLeft, IngesterPool, IngesterPoolEntry, LocalShardsUpdate, get_idle_shard_timeout,
+    notify_ingester_decommission, setup_ingester_capacity_update_listener,
+    setup_local_shards_update_listener, start_ingest_api_service, try_get_ingester_status,
+    wait_for_ingester_decommission, wait_for_ingester_status,
 };
 use quickwit_jaeger::JaegerService;
 use quickwit_janitor::{JanitorService, start_janitor_service};
@@ -1205,6 +1205,7 @@ async fn setup_ingest_v2(
         cluster.change_stream(),
         ingester_opt.clone(),
         ingester_pool,
+        event_broker.clone(),
         grpc_compression_encoding_opt,
         node_config.grpc_config.max_message_size,
     );
@@ -1215,11 +1216,13 @@ fn setup_ingester_pool(
     cluster_change_stream: ClusterChangeStream,
     ingester_opt: Option<Ingester>,
     ingester_pool: IngesterPool,
+    event_broker: EventBroker,
     grpc_compression_encoding_opt: Option<CompressionEncoding>,
     grpc_max_message_size: ByteSize,
 ) {
     let ingester_change_stream = cluster_change_stream.filter_map(move |cluster_change| {
         let ingester_opt_clone = ingester_opt.clone();
+        let event_broker_clone = event_broker.clone();
         Box::pin(async move {
             match cluster_change {
                 ClusterChange::Add(node) if node.is_indexer() => {
@@ -1254,6 +1257,9 @@ fn setup_ingester_pool(
                     Some(change)
                 }
                 ClusterChange::Remove(node) if node.is_indexer() => {
+                    event_broker_clone.publish(IngesterLeft {
+                        node_id: node.node_id.clone(),
+                    });
                     let change = build_ingester_remove_change(&node);
                     Some(change)
                 }
@@ -2164,10 +2170,17 @@ mod tests {
         let (cluster_change_stream, cluster_change_stream_tx) =
             ClusterChangeStream::new_unbounded();
         let ingester_pool = IngesterPool::default();
+        let event_broker = EventBroker::default();
+        let departures: Arc<Mutex<Vec<NodeId>>> = Arc::new(Mutex::new(Vec::new()));
+        let departures_clone = departures.clone();
+        let _subscription = event_broker.subscribe(move |departure: IngesterLeft| {
+            departures_clone.lock().unwrap().push(departure.node_id);
+        });
         setup_ingester_pool(
             cluster_change_stream,
             None::<Ingester>,
             ingester_pool.clone(),
+            event_broker,
             None,
             ByteSize::mib(20),
         );
@@ -2249,6 +2262,11 @@ mod tests {
         tokio::time::sleep(Duration::from_millis(1)).await;
 
         assert!(ingester_pool.is_empty());
+        // Routers are told to drop the departed node's routing entries.
+        assert_eq!(
+            *departures.lock().unwrap(),
+            vec![NodeId::from_str("test-ingester-node")]
+        );
     }
 
     #[tokio::test]