bugfix: dont schedule indexing tasks on retiring indexers (#6427)

* Track ingester status in control-plane indexer pool

* filter out retiring indexers when scheduling indexing tasks

* make fix

* Unit test to update the indexer_pool when changing indexer readiness status

* Unit tests for select_available_indexers_for_scheduling + renaming

* Fix merge conflicts with commit 710aa11

Author: ncoiffier-celonis

quickwit/quickwit-control-plane/src/control_plane.rs

@@ -1395,6 +1395,7 @@ mod tests {
             client: indexer,
             indexing_tasks: Vec::new(),
             indexing_capacity: CpuCapacity::from_cpu_millis(1_000),
+            ingester_status: IngesterStatus::Ready,
         };
         indexer_pool.insert(self_node_id.clone(), indexer_info);
 
@@ -1824,6 +1825,7 @@ mod tests {
             client,
             indexing_tasks: Vec::new(),
             indexing_capacity: CpuCapacity::from_cpu_millis(4_000),
+            ingester_status: IngesterStatus::Ready,
         };
         indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info);
         let ingester_pool = IngesterPool::default();
@@ -1973,6 +1975,7 @@ mod tests {
             client,
             indexing_tasks: Vec::new(),
             indexing_capacity: CpuCapacity::from_cpu_millis(4_000),
+            ingester_status: IngesterStatus::Ready,
         };
         indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info);
         let ingester_pool = IngesterPool::default();
@@ -2051,6 +2054,7 @@ mod tests {
             client,
             indexing_tasks: Vec::new(),
             indexing_capacity: CpuCapacity::from_cpu_millis(4_000),
+            ingester_status: IngesterStatus::Ready,
         };
         indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info);
         let ingester_pool = IngesterPool::default();
@@ -2777,6 +2781,7 @@ mod tests {
             client: indexer,
             indexing_tasks: Vec::new(),
             indexing_capacity: CpuCapacity::from_cpu_millis(1_000),
+            ingester_status: IngesterStatus::Ready,
         };
         indexer_pool.insert(ingester_id.clone(), indexer_info);
 

quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs

@@ -30,6 +30,7 @@ use quickwit_proto::indexing::{
     ApplyIndexingPlanRequest, CpuCapacity, IndexingService, IndexingTask, PIPELINE_FULL_CAPACITY,
     PIPELINE_THROUGHPUT,
 };
+use quickwit_proto::ingest::ingester::IngesterStatus;
 use quickwit_proto::types::NodeId;
 use scheduling::{SourceToSchedule, SourceToScheduleType};
 use serde::Serialize;
@@ -301,7 +302,7 @@ impl IndexingScheduler {
 
         let sources = get_sources_to_schedule(model);
 
-        let indexers: Vec<IndexerNodeInfo> = self.get_indexers_from_indexer_pool();
+        let indexers: Vec<IndexerNodeInfo> = self.select_available_indexers_for_scheduling();
 
         let indexer_id_to_cpu_capacities: FnvHashMap<String, CpuCapacity> = indexers
             .iter()
@@ -366,7 +367,7 @@ impl IndexingScheduler {
         {
             return;
         }
-        let indexers: Vec<IndexerNodeInfo> = self.get_indexers_from_indexer_pool();
+        let indexers: Vec<IndexerNodeInfo> = self.select_available_indexers_for_scheduling();
         let running_indexing_tasks_by_node_id: FnvHashMap<String, Vec<IndexingTask>> = indexers
             .iter()
             .map(|indexer| (indexer.node_id.to_string(), indexer.indexing_tasks.clone()))
@@ -386,8 +387,23 @@ impl IndexingScheduler {
         }
     }
 
-    fn get_indexers_from_indexer_pool(&self) -> Vec<IndexerNodeInfo> {
-        self.indexer_pool.values()
+    fn select_available_indexers_for_scheduling(&self) -> Vec<IndexerNodeInfo> {
+        let (ready, not_ready): (Vec<IndexerNodeInfo>, Vec<IndexerNodeInfo>) = self
+            .indexer_pool
+            .values()
+            .into_iter()
+            .partition(|indexer| indexer.ingester_status == IngesterStatus::Ready);
+
+        if ready.is_empty() {
+            // Allow scheduling on retiring indexers to drain shards
+            // and avoid decommission timeouts (e.g. single-node cluster).
+            warn!(
+                "no ready indexer available, falling back to retiring indexers for shard draining"
+            );
+            not_ready
+        } else {
+            ready
+        }
     }
 
     fn apply_physical_indexing_plan(
@@ -1090,9 +1106,82 @@ mod tests {
     }
 
     use quickwit_config::SourceInputFormat;
-    use quickwit_proto::indexing::mcpu;
+    use quickwit_proto::indexing::{CpuCapacity, IndexingServiceClient, MockIndexingService, mcpu};
     use quickwit_proto::ingest::{Shard, ShardState};
 
+    fn mock_indexer_node_info(node_id: &str, status: IngesterStatus) -> IndexerNodeInfo {
+        let mock_indexer = MockIndexingService::new();
+        let client = IndexingServiceClient::from_mock(mock_indexer);
+        IndexerNodeInfo {
+            node_id: NodeId::from_str(node_id),
+            generation_id: 0,
+            client,
+            indexing_tasks: Vec::new(),
+            indexing_capacity: CpuCapacity::from_cpu_millis(4_000),
+            ingester_status: status,
+        }
+    }
+
+    #[test]
+    fn test_select_available_indexers_returns_only_ready_when_available() {
+        let indexer_pool = IndexerPool::default();
+        let ready_indexer = mock_indexer_node_info("indexer-ready-1", IngesterStatus::Ready);
+        let ready_indexer_2 = mock_indexer_node_info("indexer-ready-2", IngesterStatus::Ready);
+        let retiring_indexer = mock_indexer_node_info("indexer-retiring", IngesterStatus::Retiring);
+        indexer_pool.insert(ready_indexer.node_id.clone(), ready_indexer);
+        indexer_pool.insert(ready_indexer_2.node_id.clone(), ready_indexer_2);
+        indexer_pool.insert(retiring_indexer.node_id.clone(), retiring_indexer);
+
+        let scheduler = IndexingScheduler::new(
+            "test-cluster".to_string(),
+            NodeId::from_str("control-plane"),
+            indexer_pool,
+        );
+        let selected = scheduler.select_available_indexers_for_scheduling();
+
+        assert_eq!(selected.len(), 2);
+        assert!(
+            selected
+                .iter()
+                .all(|i| i.ingester_status == IngesterStatus::Ready)
+        );
+    }
+
+    #[test]
+    fn test_select_available_indexers_falls_back_to_retiring_when_no_ready() {
+        let indexer_pool = IndexerPool::default();
+        let retiring_1 = mock_indexer_node_info("indexer-retiring-1", IngesterStatus::Retiring);
+        let retiring_2 = mock_indexer_node_info("indexer-retiring-2", IngesterStatus::Retiring);
+        indexer_pool.insert(retiring_1.node_id.clone(), retiring_1);
+        indexer_pool.insert(retiring_2.node_id.clone(), retiring_2);
+
+        let scheduler = IndexingScheduler::new(
+            "test-cluster".to_string(),
+            NodeId::from_str("control-plane"),
+            indexer_pool,
+        );
+        let selected = scheduler.select_available_indexers_for_scheduling();
+
+        assert_eq!(selected.len(), 2);
+        assert!(
+            selected
+                .iter()
+                .all(|i| i.ingester_status == IngesterStatus::Retiring)
+        );
+    }
+
+    #[test]
+    fn test_select_available_indexers_returns_empty_when_pool_is_empty() {
+        let indexer_pool = IndexerPool::default();
+        let scheduler = IndexingScheduler::new(
+            "test-cluster".to_string(),
+            NodeId::from_str("control-plane"),
+            indexer_pool,
+        );
+        let selected = scheduler.select_available_indexers_for_scheduling();
+        assert!(selected.is_empty());
+    }
+
     fn kafka_source_params_for_test() -> SourceParams {
         SourceParams::Kafka(KafkaSourceParams {
             topic: "topic".to_string(),

quickwit/quickwit-control-plane/src/lib.rs

@@ -21,6 +21,7 @@ pub(crate) mod model;
 
 use quickwit_common::tower::Pool;
 use quickwit_proto::indexing::{CpuCapacity, IndexingServiceClient, IndexingTask};
+use quickwit_proto::ingest::ingester::IngesterStatus;
 use quickwit_proto::types::NodeId;
 
 /// Indexer-node specific information stored in the pool of available indexer nodes
@@ -31,6 +32,7 @@ pub struct IndexerNodeInfo {
     pub client: IndexingServiceClient,
     pub indexing_tasks: Vec<IndexingTask>,
     pub indexing_capacity: CpuCapacity,
+    pub ingester_status: IngesterStatus,
 }
 
 pub type IndexerPool = Pool<NodeId, IndexerNodeInfo>;

quickwit/quickwit-control-plane/src/tests.rs

@@ -85,6 +85,7 @@ pub fn test_indexer_change_stream(
                             client,
                             indexing_tasks,
                             indexing_capacity: CpuCapacity::from_cpu_millis(4_000),
+                            ingester_status: node.ingester_status,
                         },
                     );
                     Some(change)

quickwit/quickwit-serve/src/lib.rs

@@ -1266,6 +1266,17 @@ fn setup_indexer_pool(
                     );
                     Some(change)
                 }
+                ClusterChange::Update { previous, updated }
+                    if updated.is_indexer()
+                        && previous.ingester_status != updated.ingester_status =>
+                {
+                    let change = build_indexer_insert_change(
+                        &updated,
+                        indexing_service_clone_opt,
+                        grpc_max_message_size,
+                    );
+                    Some(change)
+                }
                 ClusterChange::Remove(node) if node.is_indexer() => {
                     let change = build_indexer_remove_change(&node);
                     Some(change)
@@ -1300,6 +1311,7 @@ fn build_indexer_insert_change(
             client,
             indexing_tasks: node.indexing_tasks.to_vec(),
             indexing_capacity: node.indexing_cpu_capacity,
+            ingester_status: node.ingester_status,
         },
     )
 }
@@ -1654,6 +1666,33 @@ mod tests {
 
         assert_eq!(indexer_pool.len(), 1);
 
+        // changing the ingester status of an indexer node refreshes the indexer pool
+        let updated_indexer_node = ClusterNode::for_test(
+            "test-indexer-node",
+            1,
+            true,
+            &["indexer"],
+            &[],
+            IngesterStatus::Retiring,
+        )
+        .await;
+        cluster_change_stream_tx
+            .send(ClusterChange::Update {
+                previous: new_indexer_node.clone(),
+                updated: updated_indexer_node.clone(),
+            })
+            .unwrap();
+        tokio::time::sleep(Duration::from_millis(1)).await;
+
+        assert_eq!(indexer_pool.len(), 1);
+        assert_eq!(
+            indexer_pool
+                .get(&NodeId::from_str("test-indexer-node"))
+                .expect("indexer node should be in the pool")
+                .ingester_status,
+            IngesterStatus::Retiring
+        );
+
         // removing an indexer node refreshes the indexer pool
         cluster_change_stream_tx
             .send(ClusterChange::Remove(new_indexer_node))