diff --git a/quickwit/quickwit-control-plane/src/control_plane.rs b/quickwit/quickwit-control-plane/src/control_plane.rs index 2e8dd9f9f48..5828acf3a93 100644 --- a/quickwit/quickwit-control-plane/src/control_plane.rs +++ b/quickwit/quickwit-control-plane/src/control_plane.rs @@ -954,12 +954,20 @@ impl DeferableReplyHandler for ControlPlane { async fn handle_message( &mut self, - _request: ReportIndexerStateRequest, + request: ReportIndexerStateRequest, reply: impl FnOnce(Self::Reply) + Send + Sync + 'static, _ctx: &ActorContext, ) -> Result<(), ActorExitStatus> { - // TODO: implement me reply(Ok(ReportIndexerStateResponse {})); + // TODO: implement shard update handling + + if let Some(indexing_tasks_update) = request.indexing_tasks_update { + self.indexing_scheduler.record_reported_tasks( + &request.node_id, + request.generation_id, + indexing_tasks_update.indexing_tasks, + ); + } Ok(()) } } @@ -1298,6 +1306,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(1_000), ingester_status: IngesterStatus::Ready, availability_zone: None, + enable_shard_scaling_v2: false, }; indexer_pool.insert(self_node_id.clone(), indexer_info); @@ -1875,6 +1884,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: IngesterStatus::Ready, availability_zone: None, + enable_shard_scaling_v2: false, }; indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info); let ingester_pool = IngesterPool::default(); @@ -2024,6 +2034,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: IngesterStatus::Ready, availability_zone: None, + enable_shard_scaling_v2: false, }; indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info); let ingester_pool = IngesterPool::default(); @@ -2102,6 +2113,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: IngesterStatus::Ready, availability_zone: None, + enable_shard_scaling_v2: false, }; indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info); let ingester_pool = IngesterPool::default(); @@ -2794,6 +2806,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(1_000), ingester_status: IngesterStatus::Ready, availability_zone: None, + enable_shard_scaling_v2: false, }; indexer_pool.insert(ingester_id.clone(), indexer_info); diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index 6d445eb1a8d..06314340906 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -126,9 +126,15 @@ pub struct IndexingScheduler { self_node_id: NodeId, indexer_pool: IndexerPool, state: IndexingSchedulerState, + reported_tasks: FnvHashMap, pub(crate) next_rebuild_tracker: RebuildNotifier, } +struct ReportedIndexingTasks { + generation_id: u64, + indexing_tasks: Vec, +} + impl fmt::Debug for IndexingScheduler { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_struct("IndexingScheduler") @@ -358,10 +364,34 @@ fn build_indexer_statuses(indexers: &[IndexerPoolEntry]) -> FnvHashMap FnvHashMap> { +/// Once all indexers are on enable_shard_scaling_v2, the gRPC path becomes the source of truth for +/// indexing tasks. Until then, we fall back to the legacy gossip path. +fn build_indexer_tasks( + indexers: &[IndexerPoolEntry], + reported_tasks: &FnvHashMap, +) -> FnvHashMap> { + if !indexers + .iter() + .all(|indexer| indexer.enable_shard_scaling_v2) { + return indexers + .iter() + .map(|indexer| (indexer.node_id.clone(), indexer.indexing_tasks.clone())) + .collect(); + } indexers .iter() - .map(|indexer| (indexer.node_id.clone(), indexer.indexing_tasks.clone())) + .map(|indexer| { + let indexing_tasks: &[IndexingTask] = match reported_tasks.get(&indexer.node_id) { + Some(reported) if reported.generation_id == indexer.generation_id => { + &reported.indexing_tasks + } + // In the case that the reported indexer isn't the one seen in the live ingester + // pool, don't take the reported tasks. The next iteration of the control plane + // loop will clean up the stale tasks. + _ => &[], + }; + (indexer.node_id.clone(), indexing_tasks.to_vec()) + }) .collect() } @@ -380,6 +410,7 @@ impl IndexingScheduler { self_node_id, indexer_pool, state: IndexingSchedulerState::default(), + reported_tasks: FnvHashMap::default(), next_rebuild_tracker: RebuildNotifier::default(), } } @@ -388,6 +419,27 @@ impl IndexingScheduler { self.state.clone() } + pub(crate) fn record_reported_tasks( + &mut self, + node_id: &str, + generation_id: u64, + indexing_tasks: Vec, + ) { + let node_id = NodeId::from_str(node_id); + if let Some(indexer) = self.indexer_pool.get(&node_id) + && generation_id < indexer.generation_id + { + return; + } + self.reported_tasks.insert( + node_id, + ReportedIndexingTasks { + generation_id, + indexing_tasks, + }, + ); + } + // Should be called whenever a change in the list of index/shard // has happened. // @@ -424,8 +476,9 @@ impl IndexingScheduler { }; let shard_locations = model.shard_locations(); + let running_indexer_tasks = build_indexer_tasks(&indexers, &self.reported_tasks); let can_optimize_plan = is_plan_eligible_for_optimization( - &indexers, + &running_indexer_tasks, &indexer_statuses, is_locality_aware, &mut self.state, @@ -482,6 +535,11 @@ impl IndexingScheduler { /// - If node IDs differ, schedule a new indexing plan. /// - If indexing tasks differ, apply again the last plan. pub(crate) fn control_running_plan(&mut self, model: &ControlPlaneModel) { + self.reported_tasks + .retain(|node_id, report| match self.indexer_pool.get(node_id) { + Some(indexer) => report.generation_id >= indexer.generation_id, + None => false, + }); let last_applied_plan = if let Some(last_applied_plan) = &self.state.last_applied_physical_plan { last_applied_plan @@ -499,7 +557,7 @@ impl IndexingScheduler { return; } let indexers: Vec = self.select_available_indexers_for_scheduling(); - let running_indexer_tasks = build_indexer_tasks(&indexers); + let running_indexer_tasks = build_indexer_tasks(&indexers, &self.reported_tasks); let running_indexer_statuses = build_indexer_statuses(&indexers); let indexing_plans_diff = get_indexing_plans_diff( @@ -1572,7 +1630,165 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, availability_zone: None, + enable_shard_scaling_v2: false, + } + } + + #[test] + fn test_build_indexer_tasks_uses_reported_tasks_with_matching_generation() { + let pool = IndexerPool::default(); + let mut indexer = mock_indexer_node_info("indexer", IngesterStatus::Ready); + indexer.generation_id = 10; + let gossip_task = IndexingTask { + pipeline_uid: Some(PipelineUid::for_test(1)), + ..Default::default() + }; + let rpc_task = IndexingTask { + pipeline_uid: Some(PipelineUid::for_test(2)), + ..Default::default() + }; + indexer.indexing_tasks = vec![gossip_task.clone()]; + indexer.enable_shard_scaling_v2 = true; + pool.insert(indexer.node_id.clone(), indexer.clone()); + let mut legacy_indexer = mock_indexer_node_info("legacy-indexer", IngesterStatus::Ready); + pool.insert(legacy_indexer.node_id.clone(), legacy_indexer.clone()); + let mut scheduler = + IndexingScheduler::new("cluster".to_string(), NodeId::from_str("cp"), pool.clone()); + let tasks = |scheduler: &IndexingScheduler| { + build_indexer_tasks(&pool.values(), &scheduler.reported_tasks)[&indexer.node_id].clone() + }; + // Gossip is used until all indexers enable shard scaling v2. + scheduler.record_reported_tasks("indexer", 10, vec![rpc_task.clone()]); + assert_eq!(tasks(&scheduler), vec![gossip_task]); + + legacy_indexer.enable_shard_scaling_v2 = true; + pool.insert(legacy_indexer.node_id.clone(), legacy_indexer); + assert_eq!(tasks(&scheduler), vec![rpc_task.clone()]); + + // A report from a newer generation is not used until the pool catches up. + scheduler.record_reported_tasks("indexer", 11, vec![rpc_task.clone()]); + assert!(tasks(&scheduler).is_empty()); + indexer.generation_id = 11; + pool.insert(indexer.node_id.clone(), indexer.clone()); + assert_eq!(tasks(&scheduler), vec![rpc_task.clone()]); + + // Reports older than the indexer's generation are ignored. + scheduler.record_reported_tasks("indexer", 10, Vec::new()); + assert_eq!(tasks(&scheduler), vec![rpc_task]); + } + + #[test] + fn test_control_running_plan_prunes_reported_tasks() { + let pool = IndexerPool::default(); + let mut scheduler = + IndexingScheduler::new("cluster".to_string(), NodeId::from_str("cp"), pool.clone()); + for (node_id, reported_generation_id, generation_id) in + [("stale", 1, 2), ("current", 2, 2), ("future", 3, 2)] + { + scheduler.record_reported_tasks(node_id, reported_generation_id, Vec::new()); + let mut indexer = mock_indexer_node_info(node_id, IngesterStatus::Ready); + indexer.generation_id = generation_id; + pool.insert(indexer.node_id.clone(), indexer); } + scheduler.record_reported_tasks("departed", 1, Vec::new()); + scheduler.state.last_applied_physical_plan = + Some(PhysicalIndexingPlan::with_indexer_ids(&pool.keys())); + scheduler.state.last_applied_plan_timestamp = Some(Instant::now()); + + scheduler.control_running_plan(&ControlPlaneModel::default()); + + let mut reported_node_ids: Vec<&str> = scheduler + .reported_tasks + .keys() + .map(|node_id| node_id.as_str()) + .collect(); + reported_node_ids.sort_unstable(); + assert_eq!(reported_node_ids, ["current", "future"]); + } + + #[tokio::test] + async fn test_control_running_plan_uses_reported_tasks() { + let pool = IndexerPool::default(); + let mut indexer = mock_indexer_node_info("indexer", IngesterStatus::Ready); + indexer.enable_shard_scaling_v2 = true; + let task = IndexingTask { + index_uid: Some(IndexUid::for_test("index", 0)), + source_id: "source".to_string(), + pipeline_uid: Some(PipelineUid::for_test(1)), + ..Default::default() + }; + let (applied_tx, mut applied_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut mock_indexer = MockIndexingService::new(); + mock_indexer + .expect_apply_indexing_plan() + .once() + .returning(move |request| { + applied_tx.send(request).unwrap(); + Ok(ApplyIndexingPlanResponse {}) + }); + indexer.client = IndexingServiceClient::from_mock(mock_indexer); + pool.insert(indexer.node_id.clone(), indexer.clone()); + let mut scheduler = + IndexingScheduler::new("cluster".to_string(), NodeId::from_str("cp"), pool.clone()); + let mut plan = + PhysicalIndexingPlan::with_indexer_ids(std::slice::from_ref(&indexer.node_id)); + plan.add_indexing_task(&indexer.node_id, task.clone()); + scheduler.state.last_applied_physical_plan = Some(plan); + scheduler.state.last_applied_indexer_statuses = build_indexer_statuses(&pool.values()); + + // Gossip reports no tasks, but the reported tasks match the plan. + scheduler.record_reported_tasks("indexer", 0, vec![task.clone()]); + scheduler.control_running_plan(&ControlPlaneModel::default()); + assert_eq!(scheduler.state.num_applied_physical_indexing_plan, 0); + + // Gossip matches the plan, but the reported tasks do not. + indexer.indexing_tasks = vec![task.clone()]; + pool.insert(indexer.node_id.clone(), indexer); + scheduler.record_reported_tasks("indexer", 0, Vec::new()); + scheduler.control_running_plan(&ControlPlaneModel::default()); + assert_eq!(scheduler.state.num_applied_physical_indexing_plan, 1); + let applied = tokio::time::timeout(Duration::from_secs(1), applied_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(applied.indexing_tasks, vec![task]); + } + + #[test] + fn test_optimization_eligibility_uses_reported_tasks() { + let pool = IndexerPool::default(); + let mut indexer = mock_indexer_node_info("indexer", IngesterStatus::Ready); + indexer.enable_shard_scaling_v2 = true; + pool.insert(indexer.node_id.clone(), indexer.clone()); + let mut scheduler = + IndexingScheduler::new("cluster".to_string(), NodeId::from_str("cp"), pool.clone()); + let task = IndexingTask { + index_uid: Some(IndexUid::for_test("index", 0)), + source_id: "source".to_string(), + pipeline_uid: Some(PipelineUid::for_test(1)), + ..Default::default() + }; + let mut plan = + PhysicalIndexingPlan::with_indexer_ids(std::slice::from_ref(&indexer.node_id)); + plan.add_indexing_task(&indexer.node_id, task.clone()); + let statuses = build_indexer_statuses(&pool.values()); + scheduler.state.last_applied_physical_plan = Some(plan); + scheduler.state.last_applied_indexer_statuses = statuses.clone(); + scheduler.state.last_applied_plan_timestamp = + Some(Instant::now() - MIN_DURATION_BETWEEN_SCHEDULING * 2); + assert!(!is_plan_eligible_for_optimization( + &build_indexer_tasks(&pool.values(), &scheduler.reported_tasks), + &statuses, + true, + &mut scheduler.state + )); + scheduler.record_reported_tasks("indexer", 0, vec![task]); + assert!(is_plan_eligible_for_optimization( + &build_indexer_tasks(&pool.values(), &scheduler.reported_tasks), + &statuses, + true, + &mut scheduler.state + )); } #[test] @@ -1899,6 +2115,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, availability_zone: None, + enable_shard_scaling_v2: false, } } @@ -1917,6 +2134,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, availability_zone: None, + enable_shard_scaling_v2: false, } } @@ -1945,6 +2163,7 @@ mod tests { indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, availability_zone: None, + enable_shard_scaling_v2: false, } } diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs index adc436df37b..a4947271951 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs @@ -25,17 +25,15 @@ use super::{ Eligibility, IndexerInfo, SourceToSchedule, compute_max_num_shards_per_pipeline, shard_availability_zone, }; -use crate::IndexerPoolEntry; use crate::indexing_plan::PhysicalIndexingPlan; use crate::indexing_scheduler::{ - IndexingSchedulerState, MIN_DURATION_BETWEEN_SCHEDULING, build_indexer_tasks, - get_indexing_plans_diff, + IndexingSchedulerState, MIN_DURATION_BETWEEN_SCHEDULING, get_indexing_plans_diff, }; use crate::metrics::LOCALITY_REPAIRS_TOTAL; use crate::model::ShardLocations; pub(in crate::indexing_scheduler) fn is_plan_eligible_for_optimization( - indexers: &[IndexerPoolEntry], + running_indexer_tasks: &FnvHashMap>, indexer_statuses: &FnvHashMap, locality_aware: bool, state: &mut IndexingSchedulerState, @@ -43,8 +41,7 @@ pub(in crate::indexing_scheduler) fn is_plan_eligible_for_optimization( if !locality_aware || !is_plan_improvement_due(state) { return false; } - let running_indexer_tasks = build_indexer_tasks(indexers); - if !is_running_plan_stable(&running_indexer_tasks, indexer_statuses, state) { + if !is_running_plan_stable(running_indexer_tasks, indexer_statuses, state) { return false; } state.last_plan_improvement_attempt_timestamp = Some(Instant::now()); diff --git a/quickwit/quickwit-control-plane/src/lib.rs b/quickwit/quickwit-control-plane/src/lib.rs index 45ee0275d21..f3e27c27b8b 100644 --- a/quickwit/quickwit-control-plane/src/lib.rs +++ b/quickwit/quickwit-control-plane/src/lib.rs @@ -34,6 +34,7 @@ pub struct IndexerPoolEntry { pub indexing_capacity: CpuCapacity, pub ingester_status: IngesterStatus, pub availability_zone: Option, + pub enable_shard_scaling_v2: bool, } pub type IndexerPool = Pool; diff --git a/quickwit/quickwit-control-plane/src/tests.rs b/quickwit/quickwit-control-plane/src/tests.rs index df0b5dec4c8..70e5dc17cea 100644 --- a/quickwit/quickwit-control-plane/src/tests.rs +++ b/quickwit/quickwit-control-plane/src/tests.rs @@ -87,6 +87,7 @@ pub fn test_indexer_change_stream( indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: node.ingester_status, availability_zone: None, + enable_shard_scaling_v2: false, }, ); Some(change) diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index dc12eac8d3e..86139931618 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -1514,6 +1514,7 @@ fn build_indexer_insert_change( indexing_capacity: node.indexing_cpu_capacity, ingester_status: node.ingester_status, availability_zone: node.availability_zone(), + enable_shard_scaling_v2: node.enable_shard_scaling_v2(), }, ) }