From adfb03e9cf0a9158b907f154bb0abfc570899838 Mon Sep 17 00:00:00 2001 From: Nadav Gov-Ari Date: Thu, 13 Aug 2026 14:38:09 -0400 Subject: [PATCH 1/4] Slight refactor improvements --- .../src/indexing_scheduler/mod.rs | 180 +++++++++++++-- .../src/indexing_scheduler/scheduling/mod.rs | 137 +++++++++-- .../scheduling/scheduling_logic.rs | 215 +++++++++++++++++- .../scheduling/scheduling_logic_model.rs | 102 ++++++++- quickwit/quickwit-control-plane/src/lib.rs | 1 + .../quickwit-control-plane/src/metrics.rs | 10 +- quickwit/quickwit-serve/src/lib.rs | 1 + 7 files changed, 597 insertions(+), 49 deletions(-) diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index b3291c4d05b..ec555500b43 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -29,12 +29,14 @@ use quickwit_config::{ FileSourceParams, SourceParams, disable_ingest_v1, indexing_pipeline_params_fingerprint, }; use quickwit_proto::indexing::{ - ApplyIndexingPlanRequest, CpuCapacity, IndexingService, IndexingTask, PIPELINE_FULL_CAPACITY, + ApplyIndexingPlanRequest, IndexingService, IndexingTask, PIPELINE_FULL_CAPACITY, PIPELINE_THROUGHPUT, }; use quickwit_proto::ingest::ingester::IngesterStatus; use quickwit_proto::types::NodeId; -use scheduling::{SourceToSchedule, SourceToScheduleType}; +use scheduling::{ + Eligibility, IndexerInfo, SourceToSchedule, SourceToScheduleType, is_shard_nearby, +}; use serde::Serialize; use tracing::{debug, info, warn}; use ulid::Ulid; @@ -48,6 +50,8 @@ use crate::{IndexerNodeInfo, IndexerPool}; const DEFAULT_ENABLE_VARIABLE_SHARD_LOAD: bool = false; +const DEFAULT_ENABLE_AZ_AWARE_SCHEDULING: bool = false; + pub(crate) const MIN_DURATION_BETWEEN_SCHEDULING: Duration = if cfg!(any(test, feature = "testsuite")) { Duration::from_millis(50) @@ -67,6 +71,8 @@ pub struct IndexingSchedulerState { pub num_schedule_indexing_plan: usize, pub last_applied_physical_plan: Option, #[serde(skip)] + pub last_applied_indexer_statuses: FnvHashMap, + #[serde(skip)] pub last_applied_plan_timestamp: Option, } @@ -160,6 +166,16 @@ fn enable_variable_shard_load() -> bool { *IS_SHARD_LOAD_CP_ENABLED } +fn enable_az_aware_scheduling() -> bool { + static IS_AZ_AWARE_SCHEDULING_ENABLED: LazyLock = LazyLock::new(|| { + quickwit_common::get_bool_from_env( + "QW_ENABLE_AZ_AWARE_SCHEDULING", + DEFAULT_ENABLE_AZ_AWARE_SCHEDULING, + ) + }); + *IS_AZ_AWARE_SCHEDULING_ENABLED +} + /// Computes the CPU load associated to a single shard of a given index. /// /// The array passed contains all of data we have about the shard of the index. @@ -290,6 +306,72 @@ fn get_sources_to_schedule( } sources } +/// In the case where there are only draining indexers left, they are eligible to index any shards, +/// including shards hosted on other indexers. Otherwise, they can only index their own shards. +fn determine_draining_indexer_eligibility(indexers: &[IndexerNodeInfo]) -> Eligibility { + let has_ready_indexer = indexers + .iter() + .any(|indexer| indexer.ingester_status == IngesterStatus::Ready); + if has_ready_indexer { + return Eligibility::SelfHostedOnly; + } + warn!("no ready indexer available, letting draining indexers index any shard"); + Eligibility::Any +} + +fn build_indexer_info( + indexer: &IndexerNodeInfo, + draining_eligibility: Eligibility, + locality_aware: bool, +) -> IndexerInfo { + let eligibility = match indexer.ingester_status { + IngesterStatus::Ready => Eligibility::Any, + // For draining indexers, if they're the last ones left in the cluster, they need to be + // able to drain all remaining shards (Eligibility::Any). Otherwise, they just drain + // their own (Eligibility::SelfHostedOnly). + _ => draining_eligibility, + }; + let availability_zone = if locality_aware { + indexer.availability_zone.clone() + } else { + None + }; + IndexerInfo { + cpu_capacity: indexer.indexing_capacity, + availability_zone, + eligibility, + } +} + +fn build_indexer_infos( + indexers: &[IndexerNodeInfo], + locality_aware: bool, +) -> FnvHashMap { + let draining_eligibility = determine_draining_indexer_eligibility(indexers); + let mut indexer_infos: FnvHashMap = FnvHashMap::default(); + for indexer in indexers { + if indexer.indexing_capacity.cpu_millis() == 0 { + continue; + } + let indexer_info = build_indexer_info(indexer, draining_eligibility, locality_aware); + indexer_infos.insert(indexer.node_id.to_string(), indexer_info); + } + indexer_infos +} + +fn build_indexer_statuses(indexers: &[IndexerNodeInfo]) -> FnvHashMap { + indexers + .iter() + .map(|indexer| (indexer.node_id.to_string(), indexer.ingester_status)) + .collect() +} + +fn build_indexer_tasks(indexers: &[IndexerNodeInfo]) -> FnvHashMap> { + indexers + .iter() + .map(|indexer| (indexer.node_id.to_string(), indexer.indexing_tasks.clone())) + .collect() +} impl IndexingScheduler { pub fn new(cluster_id: String, self_node_id: NodeId, indexer_pool: IndexerPool) -> Self { @@ -320,18 +402,12 @@ impl IndexingScheduler { let indexers: Vec = self.select_available_indexers_for_scheduling(); - let indexer_id_to_cpu_capacities: FnvHashMap = indexers - .iter() - .filter_map(|indexer| { - if indexer.indexing_capacity.cpu_millis() > 0 { - Some((indexer.node_id.to_string(), indexer.indexing_capacity)) - } else { - None - } - }) - .collect(); + let locality_aware = enable_az_aware_scheduling(); + + let indexer_infos: FnvHashMap = + build_indexer_infos(&indexers, locality_aware); - if indexer_id_to_cpu_capacities.is_empty() { + if indexer_infos.is_empty() { if !sources.is_empty() { warn!("no indexing capacity available, cannot schedule an indexing plan"); } @@ -341,23 +417,29 @@ impl IndexingScheduler { let shard_locations = model.shard_locations(); let new_physical_plan = build_physical_indexing_plan( &sources, - &indexer_id_to_cpu_capacities, + &indexer_infos, + locality_aware, self.state.last_applied_physical_plan.as_ref(), &shard_locations, ); let shard_locality_metrics = - get_shard_locality_metrics(&new_physical_plan, &shard_locations); + get_shard_locality_metrics(&new_physical_plan, &shard_locations, &indexer_infos); shard_locality_metrics.publish(); + + let indexer_statuses = build_indexer_statuses(&indexers); if let Some(last_applied_plan) = &self.state.last_applied_physical_plan { let plans_diff = get_indexing_plans_diff( last_applied_plan.indexing_tasks_per_indexer(), new_physical_plan.indexing_tasks_per_indexer(), + &self.state.last_applied_indexer_statuses, + &indexer_statuses, ); // No need to apply the new plan as it is the same as the old one. if plans_diff.is_empty() { return; } } + self.state.last_applied_indexer_statuses = indexer_statuses; self.apply_physical_indexing_plan(new_physical_plan, Some(notify_on_drop)); self.state.num_schedule_indexing_plan += 1; } @@ -384,17 +466,17 @@ impl IndexingScheduler { return; } let indexers: Vec = self.select_available_indexers_for_scheduling(); - let running_indexing_tasks_by_node_id: FnvHashMap> = indexers - .iter() - .map(|indexer| (indexer.node_id.to_string(), indexer.indexing_tasks.clone())) - .collect(); + let running_indexer_tasks = build_indexer_tasks(&indexers); + let running_indexer_statuses = build_indexer_statuses(&indexers); let indexing_plans_diff = get_indexing_plans_diff( - &running_indexing_tasks_by_node_id, + &running_indexer_tasks, last_applied_plan.indexing_tasks_per_indexer(), + &running_indexer_statuses, + &self.state.last_applied_indexer_statuses, ); if !indexing_plans_diff.has_same_nodes() { - info!(plans_diff=?indexing_plans_diff, "running plan and last applied plan node IDs differ: schedule an indexing plan"); + info!(plans_diff=?indexing_plans_diff, "running plan and last applied plan indexers differ: schedule an indexing plan"); self.rebuild_plan(model); } else if !indexing_plans_diff.has_same_tasks() { // Some nodes may have not received their tasks, apply it again. @@ -404,6 +486,28 @@ impl IndexingScheduler { } fn select_available_indexers_for_scheduling(&self) -> Vec { + if enable_az_aware_scheduling() { + return self.select_ready_and_draining_indexers(); + } + self.select_ready_or_retiring_indexers() + } + + fn select_ready_and_draining_indexers(&self) -> Vec { + self.indexer_pool + .values() + .into_iter() + .filter(|indexer| { + matches!( + indexer.ingester_status, + IngesterStatus::Ready + | IngesterStatus::Retiring + | IngesterStatus::Decommissioning + ) + }) + .collect() + } + + fn select_ready_or_retiring_indexers(&self) -> Vec { let (ready, retiring): (Vec, Vec) = self .indexer_pool .values() @@ -504,13 +608,16 @@ impl IndexingScheduler { struct IndexingPlansDiff<'a> { pub missing_node_ids: FnvHashSet<&'a str>, pub unplanned_node_ids: FnvHashSet<&'a str>, + pub nodes_with_changed_ingester_status: FnvHashSet<&'a str>, pub missing_tasks_by_node_id: FnvHashMap<&'a str, Vec<&'a IndexingTask>>, pub unplanned_tasks_by_node_id: FnvHashMap<&'a str, Vec<&'a IndexingTask>>, } impl IndexingPlansDiff<'_> { pub fn has_same_nodes(&self) -> bool { - self.missing_node_ids.is_empty() && self.unplanned_node_ids.is_empty() + self.missing_node_ids.is_empty() + && self.unplanned_node_ids.is_empty() + && self.nodes_with_changed_ingester_status.is_empty() } pub fn has_same_tasks(&self) -> bool { @@ -535,8 +642,10 @@ impl IndexingPlansDiff<'_> { fn get_shard_locality_metrics( physical_plan: &PhysicalIndexingPlan, shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, ) -> ShardLocalityMetrics { let mut num_local_shards = 0; + let mut num_nearby_shards = 0; let mut num_remote_shards = 0; for (indexer, tasks) in physical_plan.indexing_tasks_per_indexer() { for task in tasks { @@ -547,6 +656,8 @@ fn get_shard_locality_metrics( .any(|node| node.as_str() == indexer) { num_local_shards += 1; + } else if is_shard_nearby(indexer, shard_id, shard_locations, indexer_infos) { + num_nearby_shards += 1; } else { num_remote_shards += 1; } @@ -555,6 +666,7 @@ fn get_shard_locality_metrics( } ShardLocalityMetrics { num_remote_shards, + num_nearby_shards, num_local_shards, } } @@ -582,6 +694,14 @@ impl fmt::Debug for IndexingPlansDiff<'_> { )?; separator = ", " } + if !self.nodes_with_changed_ingester_status.is_empty() { + write!( + formatter, + "{separator}nodes_with_changed_ingester_status={:?}", + PrettySample::new(&self.nodes_with_changed_ingester_status, 10) + )?; + separator = ", " + } if !self.missing_tasks_by_node_id.is_empty() { write!(formatter, "{separator}missing_tasks_by_node_id=",)?; format_indexing_task_map(formatter, &self.missing_tasks_by_node_id)?; @@ -676,6 +796,8 @@ fn format_indexing_task_map( fn get_indexing_plans_diff<'a>( running_plan: &'a FnvHashMap>, last_applied_plan: &'a FnvHashMap>, + running_ingester_statuses: &'a FnvHashMap, + last_applied_ingester_statuses: &'a FnvHashMap, ) -> IndexingPlansDiff<'a> { // Nodes diff. let running_node_ids: FnvHashSet<&str> = running_plan @@ -694,6 +816,19 @@ fn get_indexing_plans_diff<'a>( .difference(&planned_node_ids) .copied() .collect(); + // Ingester status diff. + let running_node_states: FnvHashSet<(&str, IngesterStatus)> = running_ingester_statuses + .iter() + .map(|(node_id, ingester_status)| (node_id.as_str(), *ingester_status)) + .collect(); + let planned_node_states: FnvHashSet<(&str, IngesterStatus)> = last_applied_ingester_statuses + .iter() + .map(|(node_id, ingester_status)| (node_id.as_str(), *ingester_status)) + .collect(); + let nodes_with_changed_ingester_status: FnvHashSet<&str> = running_node_states + .difference(&planned_node_states) + .map(|(node_id, _)| *node_id) + .collect(); // Tasks diff. let mut missing_tasks_by_node_id: FnvHashMap<&str, Vec<&IndexingTask>> = FnvHashMap::default(); let mut unplanned_tasks_by_node_id: FnvHashMap<&str, Vec<&IndexingTask>> = @@ -715,6 +850,7 @@ fn get_indexing_plans_diff<'a>( IndexingPlansDiff { missing_node_ids, unplanned_node_ids, + nodes_with_changed_ingester_status, missing_tasks_by_node_id, unplanned_tasks_by_node_id, } diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs index f2b32f4e784..fa8fbfab6fc 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs @@ -22,7 +22,8 @@ use fnv::{FnvHashMap, FnvHashSet}; use quickwit_common::rate_limited_debug; use quickwit_proto::indexing::{CpuCapacity, IndexingTask}; use quickwit_proto::types::{PipelineUid, ShardId, SourceUid}; -use scheduling_logic_model::{IndexerOrd, SourceOrd}; +pub use scheduling_logic_model::Eligibility; +use scheduling_logic_model::{IndexerLocality, IndexerOrd, LocalityGroup, SourceOrd}; use tracing::{error, warn}; use crate::indexing_plan::PhysicalIndexingPlan; @@ -150,6 +151,13 @@ fn convert_physical_plan_to_solution( } } +#[derive(Debug)] +pub struct IndexerInfo { + pub cpu_capacity: CpuCapacity, + pub availability_zone: Option, + pub eligibility: Eligibility, +} + #[derive(Debug)] pub struct SourceToSchedule { pub source_uid: SourceUid, @@ -357,6 +365,7 @@ fn convert_scheduling_solution_to_physical_plan( sources: &[SourceToSchedule], previous_plan_opt: Option<&PhysicalIndexingPlan>, shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, ) -> PhysicalIndexingPlan { let mut indexer_assignments = solution.indexer_assignments.clone(); let mut new_physical_plan = PhysicalIndexingPlan::with_indexer_ids(&id_to_ord_map.indexer_ids); @@ -435,6 +444,7 @@ fn convert_scheduling_solution_to_physical_plan( missing_shards, remaining_num_shards_per_node, shard_locations, + indexer_infos, ); for (shard_id, indexer) in shard_to_indexer { add_shard_to_indexer( @@ -453,6 +463,48 @@ fn convert_scheduling_solution_to_physical_plan( new_physical_plan } +fn indexer_availability_zone<'a>( + node_id: &str, + indexer_infos: &'a FnvHashMap, +) -> Option<&'a str> { + indexer_infos.get(node_id)?.availability_zone.as_deref() +} + +fn shard_availability_zone<'a>( + shard_id: &ShardId, + shard_locations: &ShardLocations, + indexer_infos: &'a FnvHashMap, +) -> Option<&'a str> { + let hosting_node_id = shard_locations.get_shard_locations(shard_id).first()?; + indexer_availability_zone(hosting_node_id.as_str(), indexer_infos) +} + +pub(crate) fn is_shard_nearby( + indexer: &str, + shard_id: &ShardId, + shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, +) -> bool { + let availability_zone = + shard_availability_zone(shard_id, shard_locations, indexer_infos); + availability_zone.is_some() + && availability_zone == indexer_availability_zone(indexer, indexer_infos) +} + +fn find_nearby_indexer( + shard_id: &ShardId, + remaining_num_shards_per_node: &HashMap, + shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, +) -> Option { + remaining_num_shards_per_node + .iter() + .filter(|(node_id, _)| is_shard_nearby(node_id, shard_id, shard_locations, indexer_infos)) + // Fill up the nearly-full indexers first. Ties break on node id, for determinism. + .min_by_key(|(node_id, num_remaining_shards)| (**num_remaining_shards, *node_id)) + .map(|(node_id, _)| node_id.clone()) +} + /// This function is meant to be called after we have solved the scheduling /// problem, so we already know the number of shards to be assigned on each indexer node. /// We now need to precisely where each shard should be assigned. @@ -470,6 +522,7 @@ fn assign_shards( missing_shards: Vec, mut remaining_num_shards_per_node: HashMap, shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, ) -> HashMap { let mut shard_to_indexer: HashMap = HashMap::with_capacity(missing_shards.len()); @@ -496,19 +549,44 @@ fn assign_shards( } for shard_id in remaining_missing_shards { - let indexer = remaining_num_shards_per_node - .keys() - .next() - .expect("failed to assign all shards. please report") - .to_string(); + let indexer = find_indexer_for_shard( + &shard_id, + &remaining_num_shards_per_node, + shard_locations, + indexer_infos, + ); decrement_num_shards(&indexer, &mut remaining_num_shards_per_node); - shard_to_indexer.insert(shard_id, indexer.to_string()); + shard_to_indexer.insert(shard_id, indexer); } assert!(remaining_num_shards_per_node.is_empty()); shard_to_indexer } +// Try to place the shard on an indexer in the same availability zone, if enabled; otherwise, pick +// the next available one. +fn find_indexer_for_shard( + shard_id: &ShardId, + remaining_num_shards_per_node: &HashMap, + shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, +) -> String { + if let Some(indexer) = find_nearby_indexer( + shard_id, + remaining_num_shards_per_node, + shard_locations, + indexer_infos, + ) { + indexer + } else { + remaining_num_shards_per_node + .keys() + .next() + .expect("failed to assign all shards. please report") + .to_string() + } +} + fn decrement_num_shards( node_id: &str, remaining_num_shards_to_schedule_per_indexers: &mut HashMap, @@ -643,7 +721,8 @@ fn inflate_node_capacities_if_necessary(problem: &mut SchedulingProblem) { /// Panics if any sources has no shards. pub fn build_physical_indexing_plan( sources: &[SourceToSchedule], - indexer_id_to_cpu_capacities: &FnvHashMap, + indexer_infos: &FnvHashMap, + locality_aware: bool, previous_plan_opt: Option<&PhysicalIndexingPlan>, shard_locations: &ShardLocations, ) -> PhysicalIndexingPlan { @@ -656,7 +735,7 @@ pub fn build_physical_indexing_plan( // Similarly, instead of accurate locality, we just keep the number of shards local // to an indexer. let (id_to_ord_map, problem) = - convert_to_simplified_problem(indexer_id_to_cpu_capacities, sources, shard_locations); + convert_to_simplified_problem(indexer_infos, locality_aware, sources, shard_locations); // Populate the previous solution, if any. let mut previous_solution = problem.new_solution(); @@ -674,6 +753,7 @@ pub fn build_physical_indexing_plan( sources, previous_plan_opt, shard_locations, + indexer_infos, ); assert_post_condition_physical_plan_match_solution( @@ -695,8 +775,21 @@ fn check_sources(sources: &[SourceToSchedule]) { } } +fn intern_locality_group<'a>( + availability_zone: &'a Option, + locality_groups: &mut FnvHashMap<&'a str, LocalityGroup>, +) -> Option { + let availability_zone = availability_zone.as_deref()?; + let next_group_ord = locality_groups.len(); + let locality_group = locality_groups + .entry(availability_zone) + .or_insert_with(|| LocalityGroup::from_ord(next_group_ord)); + Some(*locality_group) +} + fn convert_to_simplified_problem<'a>( - indexer_id_to_cpu_capacities: &'a FnvHashMap, + indexer_infos: &'a FnvHashMap, + locality_aware: bool, sources: &'a [SourceToSchedule], shard_locations: &ShardLocations, ) -> (IdToOrdMap<'a>, SchedulingProblem) { @@ -704,15 +797,29 @@ fn convert_to_simplified_problem<'a>( let mut id_to_ord_map: IdToOrdMap<'a> = IdToOrdMap::default(); // We use a Vec as a `IndexOrd` -> Max load map. - let mut indexer_cpu_capacities: Vec = - Vec::with_capacity(indexer_id_to_cpu_capacities.len()); - for (indexer_id, &cpu_capacity) in indexer_id_to_cpu_capacities { + let mut indexer_cpu_capacities: Vec = Vec::with_capacity(indexer_infos.len()); + let mut indexer_localities: Vec = Vec::with_capacity(indexer_infos.len()); + let mut locality_groups: FnvHashMap<&str, LocalityGroup> = FnvHashMap::default(); + for (indexer_id, indexer_info) in indexer_infos { let indexer_ord = id_to_ord_map.add_indexer_id(indexer_id.clone()); assert_eq!(indexer_ord, indexer_cpu_capacities.len() as IndexerOrd); - indexer_cpu_capacities.push(cpu_capacity); + indexer_cpu_capacities.push(indexer_info.cpu_capacity); + if !locality_aware { + continue; + } + let locality_group = + intern_locality_group(&indexer_info.availability_zone, &mut locality_groups); + indexer_localities.push(IndexerLocality { + group: locality_group, + eligibility: indexer_info.eligibility, + }); } - let mut problem = SchedulingProblem::with_indexer_cpu_capacities(indexer_cpu_capacities); + let mut problem = if locality_aware { + SchedulingProblem::with_indexer_localities(indexer_cpu_capacities, indexer_localities) + } else { + SchedulingProblem::with_indexer_cpu_capacities(indexer_cpu_capacities) + }; for source in sources { if let Some(source_ord) = populate_problem(source, &mut problem) { diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs index 86f992774f4..e702006a7ae 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs @@ -131,11 +131,20 @@ fn attempt_solve( enforce_indexers_cpu_capacity(problem, &mut solution); // The solution now meets the constraint, but it does not necessarily // contains all of the shards that we need to assign. - // - // We first assign sources to indexers that have some affinity with them - // (provided they have the capacity.) - place_unassigned_shards_with_affinity(problem, &mut solution); - // Finally we assign the remaining shards, regardess of whether they have affinity + if problem.is_locality_aware() { + // First, we remove remote shards from indexers that are only eligible + // to index their own shards. + strip_self_hosted_only_indexers(problem, &mut solution); + // Then, we place shards that are hosted on the indexers. + let leftover_shards_per_source = place_self_hosted_shards(problem, &mut solution); + // After that, we place remaining shards on remote indexers, but in the same locality. + place_nearby_shards(problem, &leftover_shards_per_source, &mut solution); + } else { + // If locality awareness is disabled, we directly assign sources to indexers that have some + // affinity with them (provided they have the capacity.) + place_unassigned_shards_with_affinity(problem, &mut solution); + } + // Finally we assign the remaining shards, regardess of whether they have affinity or locality // or not. place_unassigned_shards_ignoring_affinity(problem, &mut solution)?; Ok(solution) @@ -281,6 +290,31 @@ fn assert_enforce_nodes_cpu_capacity_post_condition( ); } +/// Strips non-local shards off retiring/decommissioning indexers, which are only eligible to index +/// their own shards. +fn strip_self_hosted_only_indexers( + problem: &SchedulingProblem, + solution: &mut SchedulingSolution, +) { + for indexer_assignment in &mut solution.indexer_assignments { + let indexer_ord = indexer_assignment.indexer_ord; + if problem.is_eligible_for_foreign_shards(indexer_ord) { + continue; + } + let mut num_foreign_shards_per_source: Vec<(SourceOrd, u32)> = Vec::new(); + for (&source_ord, &num_shards) in &indexer_assignment.num_shards_per_source { + let num_self_hosted_shards = problem.source_affinity(source_ord, indexer_ord); + if num_shards > num_self_hosted_shards { + num_foreign_shards_per_source + .push((source_ord, num_shards - num_self_hosted_shards)); + } + } + for (source_ord, num_foreign_shards) in num_foreign_shards_per_source { + indexer_assignment.remove_shards(source_ord, num_foreign_shards); + } + } +} + // ---------------------------------------------------- // Phase 3 // Place unassigned sources. @@ -311,6 +345,7 @@ fn attempt_place_unassigned_shards( for source in unassigned_shards { let indexers_with_most_available_capacity = compute_indexer_available_capacity(problem, &solution) + .filter(|&(indexer_ord, _)| problem.is_eligible_for_foreign_shards(indexer_ord)) .sorted_by_key(|(indexer_ord, capacity)| Reverse((*capacity, *indexer_ord))); place_unassigned_shards_single_source( source, @@ -356,6 +391,134 @@ fn place_unassigned_shards_with_affinity( } } +fn unassigned_sources_by_decreasing_load( + problem: &SchedulingProblem, + solution: &SchedulingSolution, +) -> Vec { + let mut unassigned_sources: Vec = compute_unassigned_sources(problem, solution); + unassigned_sources.sort_by_key(|source| { + let load = source.num_shards * source.load_per_shard.get(); + Reverse(load) + }); + unassigned_sources +} + +fn available_cpu_capacity( + indexer_ord: IndexerOrd, + problem: &SchedulingProblem, + solution: &SchedulingSolution, +) -> CpuCapacity { + let available_cpu_millis = + solution.indexer_assignments[indexer_ord].indexer_available_capacity(problem); + CpuCapacity::from_cpu_millis(available_cpu_millis as u32) +} + +fn place_self_hosted_shards_on_indexer( + source: &Source, + indexer_ord: IndexerOrd, + num_self_hosted_shards: u32, + problem: &SchedulingProblem, + solution: &mut SchedulingSolution, +) -> u32 { + let available_capacity = available_cpu_capacity(indexer_ord, problem, solution); + let num_placable_shards = available_capacity.cpu_millis() / source.load_per_shard; + let num_shards_placed = num_placable_shards.min(num_self_hosted_shards); + solution.indexer_assignments[indexer_ord].add_shards(source.source_ord, num_shards_placed); + num_shards_placed +} + +fn place_self_hosted_shards( + problem: &SchedulingProblem, + solution: &mut SchedulingSolution, +) -> Vec { + let num_locality_groups = problem.num_locality_groups(); + let mut leftover_shards_per_source: Vec = (0..problem.num_sources()) + .map(|_| LeftoverShards::none(num_locality_groups)) + .collect(); + let unassigned_sources: Vec = unassigned_sources_by_decreasing_load(problem, solution); + for unassigned_source in &unassigned_sources { + let source_ord = unassigned_source.source_ord as usize; + let leftover_shards = &mut leftover_shards_per_source[source_ord]; + for (&indexer_ord, &num_self_hosted_shards) in &unassigned_source.affinities { + let num_shards_placed = place_self_hosted_shards_on_indexer( + unassigned_source, + indexer_ord, + num_self_hosted_shards, + problem, + solution, + ); + let Some(locality_group) = problem.indexer_locality_group(indexer_ord) else { + continue; + }; + leftover_shards.add(locality_group, num_self_hosted_shards - num_shards_placed); + } + } + leftover_shards_per_source +} +/// Gets a list of the indexers in the locality group that are eligible to receive other indexers' +/// shards. +fn nearby_eligible_indexer_ords( + locality_group: LocalityGroup, + problem: &SchedulingProblem, +) -> Vec { + (0..problem.num_indexers()) + .filter(|&indexer_ord| problem.is_eligible_for_foreign_shards(indexer_ord)) + .filter(|&indexer_ord| problem.indexer_locality_group(indexer_ord) == Some(locality_group)) + .collect() +} + +fn place_shards_on_emptiest_indexers( + source: &Source, + num_shards_to_place: u32, + indexer_ords: &[IndexerOrd], + problem: &SchedulingProblem, + solution: &mut SchedulingSolution, +) { + if num_shards_to_place == 0 { + return; + } + let emptiest_indexers: Vec<(IndexerOrd, CpuCapacity)> = indexer_ords + .iter() + .map(|&indexer_ord| { + let available_capacity = available_cpu_capacity(indexer_ord, problem, solution); + (indexer_ord, available_capacity) + }) + .sorted_by_key(|(indexer_ord, available_capacity)| { + Reverse((*available_capacity, *indexer_ord)) + }) + .collect(); + place_shards_capped( + source, + num_shards_to_place, + emptiest_indexers.into_iter(), + solution, + ); +} + +fn place_nearby_shards( + problem: &SchedulingProblem, + leftover_shards_per_source: &[LeftoverShards], + solution: &mut SchedulingSolution, +) { + let unassigned_sources: Vec = unassigned_sources_by_decreasing_load(problem, solution); + for group_ord in 0..problem.num_locality_groups() { + let locality_group = LocalityGroup::from_ord(group_ord); + let indexer_ords = nearby_eligible_indexer_ords(locality_group, problem); + for unassigned_source in &unassigned_sources { + let source_ord = unassigned_source.source_ord as usize; + let num_leftover_shards = + leftover_shards_per_source[source_ord].in_locality_group(locality_group); + place_shards_on_emptiest_indexers( + unassigned_source, + num_leftover_shards, + &indexer_ords, + problem, + solution, + ); + } + } +} + /// Places the still-unassigned shards onto the indexers with the most available /// capacity, ignoring affinity. /// @@ -423,6 +586,48 @@ fn place_unassigned_shards_single_source( Ok(()) } +/// By placing a capped number of shards on indexers one pass at a time, we can better control +/// which shards end up where. AZ-aware routing is accomplished this way- place the shards one pass +/// at a time, first self-hosted, then same-AZ, and finally cross-AZ shards. +fn place_shards_capped( + source: &Source, + num_shards_to_place: u32, + indexers: impl Iterator, + solution: &mut SchedulingSolution, +) { + let mut num_shards_remaining = num_shards_to_place; + for (indexer_ord, available_capacity) in indexers { + if num_shards_remaining == 0 { + break; + } + let num_placable_shards = available_capacity.cpu_millis() / source.load_per_shard; + let num_shards_placed_on_indexer = num_placable_shards.min(num_shards_remaining); + solution.indexer_assignments[indexer_ord] + .add_shards(source.source_ord, num_shards_placed_on_indexer); + num_shards_remaining -= num_shards_placed_on_indexer; + } +} + +struct LeftoverShards { + num_shards_per_locality_group: Vec, +} + +impl LeftoverShards { + fn none(num_locality_groups: usize) -> LeftoverShards { + LeftoverShards { + num_shards_per_locality_group: vec![0u32; num_locality_groups], + } + } + + fn add(&mut self, locality_group: LocalityGroup, num_shards: u32) { + self.num_shards_per_locality_group[locality_group.ord()] += num_shards; + } + + fn in_locality_group(&self, locality_group: LocalityGroup) -> u32 { + self.num_shards_per_locality_group[locality_group.ord()] + } +} + /// Compute the sources/shards that have not been assigned to any indexer yet. /// Affinity are also updated, with the limitation described in `Source`. fn compute_unassigned_sources( diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic_model.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic_model.rs index 629b89043f4..a8d916d4d40 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic_model.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic_model.rs @@ -21,6 +21,34 @@ use quickwit_proto::indexing::CpuCapacity; pub type SourceOrd = u32; pub type IndexerOrd = usize; +/// Identifies a set of indexers that are local to each other, ie. availability zone. +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct LocalityGroup(usize); + +impl LocalityGroup { + pub fn from_ord(group_ord: usize) -> LocalityGroup { + LocalityGroup(group_ord) + } + + pub fn ord(self) -> usize { + self.0 + } +} + +/// Whether an indexer may be assigned shards it does not host. +/// An indexer in a not-ready state (Retiring, Decommissioning) can only be assigned its own shards. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Eligibility { + Any, + SelfHostedOnly, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct IndexerLocality { + pub group: Option, + pub eligibility: Eligibility, +} + #[derive(Clone, Debug, Eq, PartialEq)] pub struct Source { pub source_ord: SourceOrd, @@ -78,6 +106,17 @@ impl Source { pub struct SchedulingProblem { sources: Vec, indexer_cpu_capacities: Vec, + indexer_localities: Option>, +} + +/// Panics if the list of indexers is empty or if one of the indexer has a null capacity. +fn assert_valid_indexer_cpu_capacities(indexer_cpu_capacities: &[CpuCapacity]) { + assert!(!indexer_cpu_capacities.is_empty()); + assert!( + indexer_cpu_capacities + .iter() + .all(|cpu_capacity| cpu_capacity.cpu_millis() > 0) + ); } impl SchedulingProblem { @@ -88,16 +127,29 @@ impl SchedulingProblem { pub fn with_indexer_cpu_capacities( indexer_cpu_capacities: Vec, ) -> SchedulingProblem { - assert!(!indexer_cpu_capacities.is_empty()); - assert!( - indexer_cpu_capacities - .iter() - .all(|cpu_capacity| cpu_capacity.cpu_millis() > 0) - ); + assert_valid_indexer_cpu_capacities(&indexer_cpu_capacities); // TODO assert for affinity. SchedulingProblem { sources: Vec::new(), indexer_cpu_capacities, + indexer_localities: None, + } + } + + /// Same as `with_indexer_cpu_capacities`, but the problem also carries the locality group and + /// the eligibility of every indexer. + /// + /// Panics if the two vectors do not have the same length. + pub fn with_indexer_localities( + indexer_cpu_capacities: Vec, + indexer_localities: Vec, + ) -> SchedulingProblem { + assert_valid_indexer_cpu_capacities(&indexer_cpu_capacities); + assert_eq!(indexer_cpu_capacities.len(), indexer_localities.len()); + SchedulingProblem { + sources: Vec::new(), + indexer_cpu_capacities, + indexer_localities: Some(indexer_localities), } } @@ -169,6 +221,44 @@ impl SchedulingProblem { pub fn num_indexers(&self) -> usize { self.indexer_cpu_capacities.len() } + + pub fn is_locality_aware(&self) -> bool { + self.indexer_localities.is_some() + } + + /// Retiring/Decommissioning indexers are only eligible for shards they host. + pub fn is_eligible_for_foreign_shards(&self, indexer_ord: IndexerOrd) -> bool { + let Some(indexer_localities) = &self.indexer_localities else { + return true; + }; + indexer_localities[indexer_ord].eligibility == Eligibility::Any + } + + pub fn indexer_locality_group(&self, indexer_ord: IndexerOrd) -> Option { + let indexer_localities = self.indexer_localities.as_ref()?; + indexer_localities[indexer_ord].group + } + + pub fn num_locality_groups(&self) -> usize { + let Some(indexer_localities) = &self.indexer_localities else { + return 0; + }; + indexer_localities + .iter() + .filter_map(|indexer_locality| indexer_locality.group) + .map(|group| group.ord() + 1) + .max() + .unwrap_or(0) + } + + /// Number of shards of the source hosted on the indexer. + pub fn source_affinity(&self, source_ord: SourceOrd, indexer_ord: IndexerOrd) -> u32 { + self.sources[source_ord as usize] + .affinities + .get(&indexer_ord) + .copied() + .unwrap_or(0u32) + } } #[derive(Clone, Eq, PartialEq)] diff --git a/quickwit/quickwit-control-plane/src/lib.rs b/quickwit/quickwit-control-plane/src/lib.rs index e4d0b956110..3eccf4d1e1b 100644 --- a/quickwit/quickwit-control-plane/src/lib.rs +++ b/quickwit/quickwit-control-plane/src/lib.rs @@ -33,6 +33,7 @@ pub struct IndexerNodeInfo { pub indexing_tasks: Vec, pub indexing_capacity: CpuCapacity, pub ingester_status: IngesterStatus, + pub availability_zone: Option, } pub type IndexerPool = Pool; diff --git a/quickwit/quickwit-control-plane/src/metrics.rs b/quickwit/quickwit-control-plane/src/metrics.rs index aee4c1c16b7..aef28daa98d 100644 --- a/quickwit/quickwit-control-plane/src/metrics.rs +++ b/quickwit/quickwit-control-plane/src/metrics.rs @@ -16,13 +16,18 @@ use quickwit_metrics::{LabelNames, LazyCounter, LazyGauge, label_names, lazy_cou #[derive(Debug, Clone, Copy)] pub struct ShardLocalityMetrics { + // shards on other indexers if az-awareness is off; cross-az if its on pub num_remote_shards: usize, + // not used if az-awareness is off; same-az if its on + pub num_nearby_shards: usize, + // shards hosted on this indexer pub num_local_shards: usize, } impl ShardLocalityMetrics { pub fn publish(self) { LOCAL_SHARDS.set(self.num_local_shards as f64); + NEARBY_SHARDS.set(self.num_nearby_shards as f64); REMOTE_SHARDS.set(self.num_remote_shards as f64); } } @@ -47,13 +52,16 @@ pub(crate) const INDEX_ID_LABEL_NAMES: LabelNames<1> = label_names!("index_id"); static INDEXED_SHARDS: LazyGauge = lazy_gauge!( name: "indexed_shards", - description: "Number of (remote/local) shards in the indexing plan", + description: "Number of (remote/nearby/local) shards in the indexing plan", subsystem: "control_plane", ); pub(crate) static LOCAL_SHARDS: LazyGauge = lazy_gauge!(parent: INDEXED_SHARDS, "locality" => "local"); +pub(crate) static NEARBY_SHARDS: LazyGauge = + lazy_gauge!(parent: INDEXED_SHARDS, "locality" => "nearby"); + pub(crate) static REMOTE_SHARDS: LazyGauge = lazy_gauge!(parent: INDEXED_SHARDS, "locality" => "remote"); diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index cacafe31423..4b52cd6fee7 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -1515,6 +1515,7 @@ fn build_indexer_insert_change( indexing_tasks: node.indexing_tasks.to_vec(), indexing_capacity: node.indexing_cpu_capacity, ingester_status: node.ingester_status, + availability_zone: node.availability_zone().map(|az| az.to_string()), }, ) } From fb5add222e25e62118ce6515961d75a9ba503d4f Mon Sep 17 00:00:00 2001 From: Nadav Gov-Ari Date: Fri, 14 Aug 2026 10:32:02 -0400 Subject: [PATCH 2/4] WIP undo --- .../src/control_plane.rs | 5 + .../src/indexing_scheduler/mod.rs | 194 ++++++- .../src/indexing_scheduler/scheduling/mod.rs | 497 ++++++++++++++++-- .../scheduling/scheduling_logic.rs | 164 +++++- quickwit/quickwit-control-plane/src/tests.rs | 1 + 5 files changed, 810 insertions(+), 51 deletions(-) diff --git a/quickwit/quickwit-control-plane/src/control_plane.rs b/quickwit/quickwit-control-plane/src/control_plane.rs index c637a0869b2..73a2a811e12 100644 --- a/quickwit/quickwit-control-plane/src/control_plane.rs +++ b/quickwit/quickwit-control-plane/src/control_plane.rs @@ -1275,6 +1275,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(1_000), ingester_status: IngesterStatus::Ready, + availability_zone: None, }; indexer_pool.insert(self_node_id.clone(), indexer_info); @@ -1851,6 +1852,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: IngesterStatus::Ready, + availability_zone: None, }; indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info); let ingester_pool = IngesterPool::default(); @@ -1999,6 +2001,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: IngesterStatus::Ready, + availability_zone: None, }; indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info); let ingester_pool = IngesterPool::default(); @@ -2076,6 +2079,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: IngesterStatus::Ready, + availability_zone: None, }; indexer_pool.insert(indexer_node_info.node_id.clone(), indexer_node_info); let ingester_pool = IngesterPool::default(); @@ -2701,6 +2705,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(1_000), ingester_status: IngesterStatus::Ready, + availability_zone: None, }; 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 ec555500b43..bef4a04bba3 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -909,15 +909,22 @@ mod tests { use quickwit_proto::types::{IndexUid, PipelineUid, ShardId, SourceUid}; use super::*; + use crate::indexing_scheduler::scheduling::build_physical_indexing_plan_without_locality; use crate::model::ShardLocations; #[test] fn test_indexing_plans_diff() { let index_uid = IndexUid::from_str("index-1:11111111111111111111111111").unwrap(); let index_uid2 = IndexUid::from_str("index-2:11111111111111111111111111").unwrap(); + let indexer_statuses: FnvHashMap = FnvHashMap::default(); { let running_plan = FnvHashMap::default(); let desired_plan = FnvHashMap::default(); - let indexing_plans_diff = get_indexing_plans_diff(&running_plan, &desired_plan); + let indexing_plans_diff = get_indexing_plans_diff( + &running_plan, + &desired_plan, + &indexer_statuses, + &indexer_statuses, + ); assert!(indexing_plans_diff.is_empty()); } { @@ -952,7 +959,12 @@ mod tests { "indexer-1".to_string(), vec![task_2, task_1.clone(), task_1b.clone()], ); - let indexing_plans_diff = get_indexing_plans_diff(&running_plan, &desired_plan); + let indexing_plans_diff = get_indexing_plans_diff( + &running_plan, + &desired_plan, + &indexer_statuses, + &indexer_statuses, + ); assert!(indexing_plans_diff.is_empty()); } { @@ -975,7 +987,12 @@ mod tests { running_plan.insert("indexer-1".to_string(), vec![task_1.clone()]); desired_plan.insert("indexer-1".to_string(), vec![task_2.clone()]); - let indexing_plans_diff = get_indexing_plans_diff(&running_plan, &desired_plan); + let indexing_plans_diff = get_indexing_plans_diff( + &running_plan, + &desired_plan, + &indexer_statuses, + &indexer_statuses, + ); assert!(!indexing_plans_diff.is_empty()); assert!(indexing_plans_diff.has_same_nodes()); assert!(!indexing_plans_diff.has_same_tasks()); @@ -1009,7 +1026,12 @@ mod tests { running_plan.insert("indexer-2".to_string(), vec![task_2.clone()]); desired_plan.insert("indexer-1".to_string(), vec![task_1.clone()]); - let indexing_plans_diff = get_indexing_plans_diff(&running_plan, &desired_plan); + let indexing_plans_diff = get_indexing_plans_diff( + &running_plan, + &desired_plan, + &indexer_statuses, + &indexer_statuses, + ); assert!(!indexing_plans_diff.is_empty()); assert!(!indexing_plans_diff.has_same_nodes()); assert!(!indexing_plans_diff.has_same_tasks()); @@ -1061,7 +1083,12 @@ mod tests { vec![task_1a.clone(), task_1b.clone(), task_1c.clone()], ); - let indexing_plans_diff = get_indexing_plans_diff(&running_plan, &desired_plan); + let indexing_plans_diff = get_indexing_plans_diff( + &running_plan, + &desired_plan, + &indexer_statuses, + &indexer_statuses, + ); assert!(!indexing_plans_diff.is_empty()); assert!(indexing_plans_diff.has_same_nodes()); assert!(!indexing_plans_diff.has_same_tasks()); @@ -1070,6 +1097,38 @@ mod tests { FnvHashMap::from_iter([("indexer-1", vec![&task_1b, &task_1c])]) ); } + { + let mut running_plan = FnvHashMap::default(); + let mut desired_plan = FnvHashMap::default(); + let task_1 = IndexingTask { + pipeline_uid: Some(PipelineUid::for_test(1u128)), + index_uid: Some(index_uid.clone()), + source_id: "source-1".to_string(), + shard_ids: Vec::new(), + params_fingerprint: 0, + }; + running_plan.insert("indexer-1".to_string(), vec![task_1.clone()]); + desired_plan.insert("indexer-1".to_string(), vec![task_1.clone()]); + + let mut running_statuses = FnvHashMap::default(); + running_statuses.insert("indexer-1".to_string(), IngesterStatus::Retiring); + let mut last_applied_statuses = FnvHashMap::default(); + last_applied_statuses.insert("indexer-1".to_string(), IngesterStatus::Ready); + + let indexing_plans_diff = get_indexing_plans_diff( + &running_plan, + &desired_plan, + &running_statuses, + &last_applied_statuses, + ); + assert!(!indexing_plans_diff.is_empty()); + assert!(indexing_plans_diff.has_same_tasks()); + assert!(!indexing_plans_diff.has_same_nodes()); + assert_eq!( + indexing_plans_diff.nodes_with_changed_ingester_status, + FnvHashSet::from_iter(["indexer-1"]) + ); + } } #[test] @@ -1218,12 +1277,16 @@ mod tests { params_fingerprint: 0, }, ]; - let mut indexer_max_loads = FnvHashMap::default(); - indexer_max_loads.insert("indexer1".to_string(), mcpu(3_000)); - indexer_max_loads.insert("indexer2".to_string(), mcpu(3_000)); + let mut indexer_infos = FnvHashMap::default(); + indexer_infos.insert("indexer1".to_string(), IndexerInfo::for_test(mcpu(3_000))); + indexer_infos.insert("indexer2".to_string(), IndexerInfo::for_test(mcpu(3_000))); let shard_locations = ShardLocations::default(); - let physical_plan = - build_physical_indexing_plan(&sources[..], &indexer_max_loads, None, &shard_locations); + let physical_plan = build_physical_indexing_plan_without_locality( + &sources[..], + &indexer_infos, + None, + &shard_locations, + ); assert_eq!(physical_plan.indexing_tasks_per_indexer().len(), 2); let indexing_tasks_1 = physical_plan.indexer("indexer1").unwrap(); assert_eq!(indexing_tasks_1.len(), 2); @@ -1265,6 +1328,7 @@ mod tests { let plan = IndexingPlansDiff { missing_node_ids: FnvHashSet::default(), unplanned_node_ids: FnvHashSet::default(), + nodes_with_changed_ingester_status: FnvHashSet::default(), missing_tasks_by_node_id: map, unplanned_tasks_by_node_id: FnvHashMap::default(), }; @@ -1293,13 +1357,13 @@ mod tests { } let sources: Vec = get_sources_to_schedule(&model, false); - let mut indexer_max_loads = FnvHashMap::default(); + let mut indexer_infos = FnvHashMap::default(); for i in 0..num_indexers { let indexer_id = format!("indexer-{i}"); - indexer_max_loads.insert(indexer_id, mcpu(4_000)); + indexer_infos.insert(indexer_id, IndexerInfo::for_test(mcpu(4_000))); } let shard_locations = ShardLocations::default(); - let _physical_indexing_plan = build_physical_indexing_plan(&sources, &indexer_max_loads, None, &shard_locations); + let _physical_indexing_plan = build_physical_indexing_plan_without_locality(&sources, &indexer_infos, None, &shard_locations); } } @@ -1319,6 +1383,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, + availability_zone: None, } } @@ -1337,7 +1402,7 @@ mod tests { NodeId::from_str("control-plane"), indexer_pool, ); - let selected = scheduler.select_available_indexers_for_scheduling(); + let selected = scheduler.select_ready_or_retiring_indexers(); assert_eq!(selected.len(), 2); assert!( @@ -1370,7 +1435,7 @@ mod tests { NodeId::from_str("control-plane"), indexer_pool, ); - let selected = scheduler.select_available_indexers_for_scheduling(); + let selected = scheduler.select_ready_or_retiring_indexers(); assert_eq!(selected.len(), 2); assert!( @@ -1388,10 +1453,104 @@ mod tests { NodeId::from_str("control-plane"), indexer_pool, ); - let selected = scheduler.select_available_indexers_for_scheduling(); + let selected = scheduler.select_ready_or_retiring_indexers(); assert!(selected.is_empty()); } + #[test] + fn test_select_ready_and_draining_indexers() { + let indexer_pool = IndexerPool::default(); + let statuses = [ + IngesterStatus::Unspecified, + IngesterStatus::Initializing, + IngesterStatus::Ready, + IngesterStatus::Retiring, + IngesterStatus::Decommissioning, + IngesterStatus::Decommissioned, + IngesterStatus::Failed, + ]; + for status in statuses { + let node_id = format!("indexer-{status:?}"); + let indexer = mock_indexer_node_info(&node_id, status); + indexer_pool.insert(indexer.node_id.clone(), indexer); + } + + let scheduler = IndexingScheduler::new( + "test-cluster".to_string(), + NodeId::from_str("control-plane"), + indexer_pool, + ); + let selected = scheduler.select_ready_and_draining_indexers(); + + let selected_statuses: FnvHashSet = selected + .iter() + .map(|indexer| indexer.ingester_status) + .collect(); + let expected_statuses = FnvHashSet::from_iter([ + IngesterStatus::Ready, + IngesterStatus::Retiring, + IngesterStatus::Decommissioning, + ]); + assert_eq!(selected.len(), 3); + assert_eq!(selected_statuses, expected_statuses); + } + + #[test] + fn test_build_indexer_infos_assigns_draining_eligibility() { + let locality_aware = true; + { + let mut ready = mock_indexer_node_info("indexer-ready", IngesterStatus::Ready); + ready.availability_zone = Some("az-a".to_string()); + let retiring = mock_indexer_node_info("indexer-retiring", IngesterStatus::Retiring); + let decommissioning = + mock_indexer_node_info("indexer-decommissioning", IngesterStatus::Decommissioning); + let indexers = vec![ready, retiring, decommissioning]; + + let indexer_infos = build_indexer_infos(&indexers, locality_aware); + + assert_eq!(indexer_infos["indexer-ready"].eligibility, Eligibility::Any); + assert_eq!( + indexer_infos["indexer-ready"].availability_zone, + Some("az-a".to_string()) + ); + assert_eq!( + indexer_infos["indexer-retiring"].eligibility, + Eligibility::SelfHostedOnly + ); + assert_eq!( + indexer_infos["indexer-decommissioning"].eligibility, + Eligibility::SelfHostedOnly + ); + } + { + let retiring = mock_indexer_node_info("indexer-retiring", IngesterStatus::Retiring); + let decommissioning = + mock_indexer_node_info("indexer-decommissioning", IngesterStatus::Decommissioning); + let indexers = vec![retiring, decommissioning]; + + let indexer_infos = build_indexer_infos(&indexers, locality_aware); + + assert_eq!( + indexer_infos["indexer-retiring"].eligibility, + Eligibility::Any + ); + assert_eq!( + indexer_infos["indexer-decommissioning"].eligibility, + Eligibility::Any + ); + } + { + let mut ready = mock_indexer_node_info("indexer-ready", IngesterStatus::Ready); + ready.availability_zone = Some("az-a".to_string()); + let indexers = vec![ready]; + let locality_unaware = false; + + let indexer_infos = build_indexer_infos(&indexers, locality_unaware); + + assert_eq!(indexer_infos["indexer-ready"].availability_zone, None); + } + } + // Only ready, retiring, and decommissioning indexers receive a plan; indexers in any other // state must be skipped entirely. See `apply_physical_indexing_plan`. #[tokio::test] @@ -1492,6 +1651,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, + availability_zone: None, } } @@ -1509,6 +1669,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, + availability_zone: None, } } @@ -1536,6 +1697,7 @@ mod tests { indexing_tasks: Vec::new(), indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: status, + availability_zone: None, } } diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs index fa8fbfab6fc..712dc53cf0b 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs @@ -158,6 +158,17 @@ pub struct IndexerInfo { pub eligibility: Eligibility, } +#[cfg(test)] +impl IndexerInfo { + pub(crate) fn for_test(cpu_capacity: CpuCapacity) -> IndexerInfo { + IndexerInfo { + cpu_capacity, + availability_zone: None, + eligibility: Eligibility::Any, + } + } +} + #[derive(Debug)] pub struct SourceToSchedule { pub source_uid: SourceUid, @@ -765,6 +776,23 @@ pub fn build_physical_indexing_plan( new_physical_plan } +#[cfg(test)] +pub(crate) fn build_physical_indexing_plan_without_locality( + sources: &[SourceToSchedule], + indexer_infos: &FnvHashMap, + previous_plan_opt: Option<&PhysicalIndexingPlan>, + shard_locations: &ShardLocations, +) -> PhysicalIndexingPlan { + let locality_aware = false; + build_physical_indexing_plan( + sources, + indexer_infos, + locality_aware, + previous_plan_opt, + shard_locations, + ) +} + /// Makes any checks on the sources. /// Sharded sources are not allowed to have no shards. fn check_sources(sources: &[SourceToSchedule]) { @@ -848,7 +876,7 @@ fn convert_to_simplified_problem<'a>( #[cfg(test)] mod tests { - use std::collections::{HashMap, HashSet}; + use std::collections::{BTreeMap, HashMap, HashSet}; use std::num::NonZeroU32; use std::str::FromStr; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -857,17 +885,95 @@ mod tests { use itertools::Itertools; use quickwit_proto::indexing::{CpuCapacity, IndexingTask, mcpu}; use quickwit_proto::types::{IndexUid, NodeId, PipelineUid, ShardId, SourceUid}; + use rand::SeedableRng; use rand::prelude::IndexedRandom; + use rand::rngs::StdRng; + use super::scheduling_logic::solve; use super::{ - SourceToSchedule, SourceToScheduleType, build_physical_indexing_plan, + Eligibility, IndexerInfo, SourceToSchedule, SourceToScheduleType, + build_physical_indexing_plan, build_physical_indexing_plan_without_locality, convert_scheduling_solution_to_physical_plan_single_node_single_source, + convert_to_simplified_problem, }; use crate::indexing_plan::PhysicalIndexingPlan; use crate::indexing_scheduler::get_shard_locality_metrics; use crate::indexing_scheduler::scheduling::assign_shards; use crate::model::ShardLocations; + fn indexer_info_in_az( + cpu_capacity: CpuCapacity, + availability_zone: &str, + eligibility: Eligibility, + ) -> IndexerInfo { + IndexerInfo { + cpu_capacity, + availability_zone: Some(availability_zone.to_string()), + eligibility, + } + } + + struct IndexerSpec { + node_id: NodeId, + cpu_capacity: CpuCapacity, + availability_zone: Option, + } + + impl IndexerSpec { + fn new( + node_id: &str, + cpu_capacity: CpuCapacity, + availability_zone: Option<&str>, + ) -> IndexerSpec { + IndexerSpec { + node_id: NodeId::from_str(node_id), + cpu_capacity, + availability_zone: availability_zone.map(|az| az.to_string()), + } + } + + fn to_indexer_info(&self) -> IndexerInfo { + IndexerInfo { + cpu_capacity: self.cpu_capacity, + availability_zone: self.availability_zone.clone(), + eligibility: Eligibility::Any, + } + } + } + + fn shard_ids_for_indexer(plan: &PhysicalIndexingPlan, indexer: &str) -> Vec { + let mut shard_ids: Vec = plan + .indexer(indexer) + .unwrap() + .iter() + .flat_map(|task| task.shard_ids.iter().cloned()) + .collect(); + shard_ids.sort(); + shard_ids + } + + fn shard_counts_per_az( + plan: &PhysicalIndexingPlan, + indexer_infos: &FnvHashMap, + ) -> BTreeMap, Vec> { + let mut counts_per_az: BTreeMap, Vec> = BTreeMap::default(); + for (indexer, tasks) in plan.indexing_tasks_per_indexer() { + let num_shards: usize = tasks.iter().map(|task| task.shard_ids.len()).sum(); + if num_shards == 0 { + continue; + } + let availability_zone = indexer_infos[indexer].availability_zone.clone(); + counts_per_az + .entry(availability_zone) + .or_default() + .push(num_shards); + } + for counts in counts_per_az.values_mut() { + counts.sort(); + } + counts_per_az + } + fn source_id() -> SourceUid { static COUNTER: AtomicUsize = AtomicUsize::new(0); let index = IndexUid::for_test("test_index", 0); @@ -915,13 +1021,13 @@ mod tests { source_type: SourceToScheduleType::IngestV1, params_fingerprint: 0, }; - let mut indexer_id_to_cpu_capacities = FnvHashMap::default(); - indexer_id_to_cpu_capacities.insert(indexer1.clone(), mcpu(16_000)); - indexer_id_to_cpu_capacities.insert(indexer2.clone(), mcpu(16_000)); + let mut indexer_infos = FnvHashMap::default(); + indexer_infos.insert(indexer1.clone(), IndexerInfo::for_test(mcpu(16_000))); + indexer_infos.insert(indexer2.clone(), IndexerInfo::for_test(mcpu(16_000))); let shard_locations = ShardLocations::default(); - let indexing_plan = build_physical_indexing_plan( + let indexing_plan = build_physical_indexing_plan_without_locality( &[source_0, source_1, source_2], - &indexer_id_to_cpu_capacities, + &indexer_infos, None, &shard_locations, ); @@ -978,11 +1084,12 @@ mod tests { }) .collect(); - let mut indexer_id_to_cpu_capacities = FnvHashMap::default(); + let mut indexer_infos = FnvHashMap::default(); for indexer in &indexers { - indexer_id_to_cpu_capacities.insert(indexer.as_str().to_string(), mcpu(16_000)); + let indexer_info = IndexerInfo::for_test(mcpu(16_000)); + indexer_infos.insert(indexer.as_str().to_string(), indexer_info); } - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(42); let mut shard_locations = ShardLocations::default(); for shard_id in &shard_ids { @@ -990,14 +1097,14 @@ mod tests { shard_locations.add_location(shard_id, indexer); } - let plan = build_physical_indexing_plan( + let plan = build_physical_indexing_plan_without_locality( &sources, - &indexer_id_to_cpu_capacities, + &indexer_infos, None, &shard_locations, ); assert_eq!(plan.indexing_tasks_per_indexer().len(), num_indexers); - let metrics = get_shard_locality_metrics(&plan, &shard_locations); + let metrics = get_shard_locality_metrics(&plan, &shard_locations, &indexer_infos); assert_eq!( metrics.num_remote_shards + metrics.num_local_shards, num_shards @@ -1005,6 +1112,309 @@ mod tests { assert!(metrics.num_remote_shards < 10); } + // asserts that if the flag is off, even if az/eligibiltiy are applied, the indexing plan + // falls back to the old logic. + #[test] + fn test_flag_off_ignores_locality_and_eligibility() { + let shard0 = ShardId::from(0); + let shard1 = ShardId::from(1); + let shard2 = ShardId::from(2); + let sharded_source = SourceToSchedule { + source_uid: source_id(), + source_type: SourceToScheduleType::Sharded { + shard_ids: vec![shard0.clone(), shard1.clone(), shard2.clone()], + load_per_shard: NonZeroU32::new(1_000).unwrap(), + }, + params_fingerprint: 0, + }; + let source_without_affinity = SourceToSchedule { + source_uid: source_id(), + source_type: SourceToScheduleType::NonSharded { + num_pipelines: 1, + load_per_pipeline: NonZeroU32::new(1_000).unwrap(), + }, + params_fingerprint: 0, + }; + let sources = vec![sharded_source, source_without_affinity]; + + let indexer1 = NodeId::from_str("indexer1"); + let indexer2 = NodeId::from_str("indexer2"); + let mut shard_locations = ShardLocations::default(); + shard_locations.add_location(&shard0, &indexer1); + shard_locations.add_location(&shard1, &indexer1); + shard_locations.add_location(&shard2, &indexer2); + + let mut plain_indexer_infos = FnvHashMap::default(); + plain_indexer_infos.insert(indexer1.to_string(), IndexerInfo::for_test(mcpu(4_000))); + plain_indexer_infos.insert(indexer2.to_string(), IndexerInfo::for_test(mcpu(4_000))); + + let indexer1_in_az = indexer_info_in_az(mcpu(4_000), "az-a", Eligibility::SelfHostedOnly); + let indexer2_in_az = indexer_info_in_az(mcpu(4_000), "az-b", Eligibility::SelfHostedOnly); + let mut az_indexer_infos = FnvHashMap::default(); + az_indexer_infos.insert(indexer1.to_string(), indexer1_in_az); + az_indexer_infos.insert(indexer2.to_string(), indexer2_in_az); + + let locality_aware = false; + let (_, plain_problem) = convert_to_simplified_problem( + &plain_indexer_infos, + locality_aware, + &sources, + &shard_locations, + ); + let (_, az_problem) = convert_to_simplified_problem( + &az_indexer_infos, + locality_aware, + &sources, + &shard_locations, + ); + + assert!(!plain_problem.is_locality_aware()); + assert!(!az_problem.is_locality_aware()); + assert_eq!(plain_problem.num_locality_groups(), 0); + assert_eq!(az_problem.num_locality_groups(), 0); + for indexer_ord in 0..az_problem.num_indexers() { + assert!(az_problem.is_eligible_for_foreign_shards(indexer_ord)); + } + + let plain_previous_solution = plain_problem.new_solution(); + let az_previous_solution = az_problem.new_solution(); + let plain_solution = solve(plain_problem, plain_previous_solution); + let az_solution = solve(az_problem, az_previous_solution); + assert_eq!( + plain_solution.indexer_assignments, + az_solution.indexer_assignments + ); + } + + struct TopologyOutcome { + num_local_shards: usize, + num_nearby_shards: usize, + num_remote_shards: usize, + shard_counts_per_az: BTreeMap, Vec>, + } + + fn assert_stable_locality_topology( + indexer_specs: &[IndexerSpec], + shard_ids: &[ShardId], + shard_hosts: &[&NodeId], + ) -> TopologyOutcome { + let source = SourceToSchedule { + source_uid: source_id(), + source_type: SourceToScheduleType::Sharded { + shard_ids: shard_ids.to_vec(), + load_per_shard: NonZeroU32::new(1_000).unwrap(), + }, + params_fingerprint: 0, + }; + let sources = vec![source]; + + let mut shard_locations = ShardLocations::default(); + for (shard_id, shard_host) in shard_ids.iter().zip(shard_hosts) { + shard_locations.add_location(shard_id, shard_host); + } + + let mut indexer_infos = FnvHashMap::default(); + for indexer_spec in indexer_specs { + let indexer_info = indexer_spec.to_indexer_info(); + indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); + } + + let locality_aware = true; + let plan = build_physical_indexing_plan( + &sources, + &indexer_infos, + locality_aware, + None, + &shard_locations, + ); + + let scheduled_shard_ids: Vec = plan + .indexing_tasks_per_indexer() + .values() + .flatten() + .flat_map(|task| task.shard_ids.iter().cloned()) + .collect(); + let unique_shard_ids: HashSet<&ShardId> = scheduled_shard_ids.iter().collect(); + assert_eq!(scheduled_shard_ids.len(), shard_ids.len()); + assert_eq!(unique_shard_ids.len(), shard_ids.len()); + + let replanned = build_physical_indexing_plan( + &sources, + &indexer_infos, + locality_aware, + Some(&plan), + &shard_locations, + ); + assert_eq!(plan, replanned); + + let mut reversed_indexer_infos = FnvHashMap::default(); + for indexer_spec in indexer_specs.iter().rev() { + let indexer_info = indexer_spec.to_indexer_info(); + reversed_indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); + } + let reversed_plan = build_physical_indexing_plan( + &sources, + &reversed_indexer_infos, + locality_aware, + None, + &shard_locations, + ); + let counts_per_az = shard_counts_per_az(&plan, &indexer_infos); + let reversed_counts_per_az = shard_counts_per_az(&reversed_plan, &reversed_indexer_infos); + assert_eq!(counts_per_az, reversed_counts_per_az); + + let metrics = get_shard_locality_metrics(&plan, &shard_locations, &indexer_infos); + TopologyOutcome { + num_local_shards: metrics.num_local_shards, + num_nearby_shards: metrics.num_nearby_shards, + num_remote_shards: metrics.num_remote_shards, + shard_counts_per_az: counts_per_az, + } + } + + #[test] + fn test_locality_aware_topology_matrix_is_stable() { + let shard_ids: Vec = (0..5u64).map(ShardId::from).collect(); + let unschedulable_host = NodeId::from_str("unschedulable"); + { + let indexer_specs = vec![IndexerSpec::new("indexer1", mcpu(4_000), None)]; + let shard_hosts = vec![&indexer_specs[0].node_id, &unschedulable_host]; + let outcome = + assert_stable_locality_topology(&indexer_specs, &shard_ids[..2], &shard_hosts); + assert_eq!(outcome.num_local_shards, 1); + assert_eq!(outcome.num_nearby_shards, 0); + assert_eq!(outcome.num_remote_shards, 1); + let expected_counts_per_az = BTreeMap::from_iter([(None, vec![2])]); + assert_eq!(outcome.shard_counts_per_az, expected_counts_per_az); + } + { + let indexer_specs = vec![ + IndexerSpec::new("indexer1", mcpu(4_000), Some("az-a")), + IndexerSpec::new("indexer2", mcpu(4_000), Some("az-a")), + ]; + let shard_hosts = vec![&indexer_specs[0].node_id, &indexer_specs[0].node_id]; + let outcome = + assert_stable_locality_topology(&indexer_specs, &shard_ids[..2], &shard_hosts); + assert_eq!(outcome.num_local_shards, 1); + assert_eq!(outcome.num_nearby_shards, 1); + assert_eq!(outcome.num_remote_shards, 0); + let expected_counts_per_az = + BTreeMap::from_iter([(Some("az-a".to_string()), vec![1, 1])]); + assert_eq!(outcome.shard_counts_per_az, expected_counts_per_az); + } + { + let indexer_specs = vec![ + IndexerSpec::new("indexer1", mcpu(4_000), Some("az-a")), + IndexerSpec::new("indexer2", mcpu(4_000), Some("az-b")), + ]; + let shard_hosts = vec![&indexer_specs[0].node_id, &indexer_specs[0].node_id]; + let outcome = + assert_stable_locality_topology(&indexer_specs, &shard_ids[..2], &shard_hosts); + assert_eq!(outcome.num_local_shards, 1); + assert_eq!(outcome.num_nearby_shards, 0); + assert_eq!(outcome.num_remote_shards, 1); + let expected_counts_per_az = BTreeMap::from_iter([ + (Some("az-a".to_string()), vec![1]), + (Some("az-b".to_string()), vec![1]), + ]); + assert_eq!(outcome.shard_counts_per_az, expected_counts_per_az); + } + { + let indexer_specs = vec![ + IndexerSpec::new("indexer1", mcpu(8_000), Some("az-a")), + IndexerSpec::new("indexer2", mcpu(4_000), Some("az-a")), + IndexerSpec::new("indexer3", mcpu(4_000), Some("az-b")), + IndexerSpec::new("indexer4", mcpu(4_000), Some("az-c")), + ]; + let shard_hosts = vec![&indexer_specs[0].node_id; 5]; + let outcome = assert_stable_locality_topology(&indexer_specs, &shard_ids, &shard_hosts); + assert_eq!(outcome.num_local_shards, 2); + assert_eq!(outcome.num_nearby_shards, 1); + assert_eq!(outcome.num_remote_shards, 2); + let expected_counts_per_az = BTreeMap::from_iter([ + (Some("az-a".to_string()), vec![1, 2]), + (Some("az-b".to_string()), vec![1]), + (Some("az-c".to_string()), vec![1]), + ]); + assert_eq!(outcome.shard_counts_per_az, expected_counts_per_az); + } + } + + #[test] + fn test_draining_indexer_keeps_only_hosted_shard_ids() { + let shard0 = ShardId::from(0); + let shard1 = ShardId::from(1); + let source_uid = source_id(); + let source = SourceToSchedule { + source_uid: source_uid.clone(), + source_type: SourceToScheduleType::Sharded { + shard_ids: vec![shard0.clone(), shard1.clone()], + load_per_shard: NonZeroU32::new(1_000).unwrap(), + }, + params_fingerprint: 0, + }; + let sources = vec![source]; + + let draining_indexer = NodeId::from_str("indexer1"); + let ready_indexer = NodeId::from_str("indexer2"); + let mut shard_locations = ShardLocations::default(); + shard_locations.add_location(&shard0, &draining_indexer); + shard_locations.add_location(&shard1, &ready_indexer); + + let draining_info = indexer_info_in_az(mcpu(4_000), "az-a", Eligibility::SelfHostedOnly); + let ready_info = indexer_info_in_az(mcpu(4_000), "az-a", Eligibility::Any); + let mut indexer_infos = FnvHashMap::default(); + indexer_infos.insert(draining_indexer.to_string(), draining_info); + indexer_infos.insert(ready_indexer.to_string(), ready_info); + + let indexer_ids = vec![draining_indexer.to_string(), ready_indexer.to_string()]; + let mut swapped_plan = PhysicalIndexingPlan::with_indexer_ids(&indexer_ids); + let draining_task = IndexingTask { + index_uid: Some(source_uid.index_uid.clone()), + source_id: source_uid.source_id.clone(), + pipeline_uid: Some(PipelineUid::for_test(1u128)), + shard_ids: vec![shard1.clone()], + params_fingerprint: 0, + }; + let ready_task = IndexingTask { + index_uid: Some(source_uid.index_uid.clone()), + source_id: source_uid.source_id.clone(), + pipeline_uid: Some(PipelineUid::for_test(2u128)), + shard_ids: vec![shard0.clone()], + params_fingerprint: 0, + }; + swapped_plan.add_indexing_task(draining_indexer.as_str(), draining_task); + swapped_plan.add_indexing_task(ready_indexer.as_str(), ready_task); + + let locality_aware = true; + let plan = build_physical_indexing_plan( + &sources, + &indexer_infos, + locality_aware, + Some(&swapped_plan), + &shard_locations, + ); + + let draining_shard_ids = shard_ids_for_indexer(&plan, draining_indexer.as_str()); + let ready_shard_ids = shard_ids_for_indexer(&plan, ready_indexer.as_str()); + assert_eq!(draining_shard_ids, vec![shard0.clone()]); + assert_eq!(ready_shard_ids, vec![shard1.clone()]); + + let metrics = get_shard_locality_metrics(&plan, &shard_locations, &indexer_infos); + assert_eq!(metrics.num_local_shards, 2); + assert_eq!(metrics.num_nearby_shards, 0); + assert_eq!(metrics.num_remote_shards, 0); + + let replanned = build_physical_indexing_plan( + &sources, + &indexer_infos, + locality_aware, + Some(&plan), + &shard_locations, + ); + assert_eq!(plan, replanned); + } + #[tokio::test] async fn test_build_physical_indexing_plan_with_not_enough_indexers() { let source_uid1 = source_id(); @@ -1019,23 +1429,31 @@ mod tests { let sources = vec![source_1]; let indexer1 = "indexer1".to_string(); - let mut indexer_max_loads = FnvHashMap::default(); + let mut indexer_infos = FnvHashMap::default(); let shard_locations = ShardLocations::default(); { - indexer_max_loads.insert(indexer1.clone(), mcpu(1_999)); + indexer_infos.insert(indexer1.clone(), IndexerInfo::for_test(mcpu(1_999))); // This test what happens when there isn't enough capacity on the cluster. - let physical_plan = - build_physical_indexing_plan(&sources, &indexer_max_loads, None, &shard_locations); + let physical_plan = build_physical_indexing_plan_without_locality( + &sources, + &indexer_infos, + None, + &shard_locations, + ); assert_eq!(physical_plan.indexing_tasks_per_indexer().len(), 1); let expected_tasks = physical_plan.indexer(&indexer1).unwrap(); assert_eq!(expected_tasks.len(), 2); assert_eq!(&expected_tasks[0].source_id, &source_uid1.source_id); } { - indexer_max_loads.insert(indexer1.clone(), mcpu(2_000)); + indexer_infos.insert(indexer1.clone(), IndexerInfo::for_test(mcpu(2_000))); // This test what happens when there isn't enough capacity on the cluster. - let physical_plan = - build_physical_indexing_plan(&sources, &indexer_max_loads, None, &shard_locations); + let physical_plan = build_physical_indexing_plan_without_locality( + &sources, + &indexer_infos, + None, + &shard_locations, + ); assert_eq!(physical_plan.indexing_tasks_per_indexer().len(), 1); let expected_tasks = physical_plan.indexer(&indexer1).unwrap(); assert_eq!(expected_tasks.len(), 2); @@ -1093,16 +1511,16 @@ mod tests { }, params_fingerprint: 0, }]; - let mut indexer_id_to_cpu_capacities = FnvHashMap::default(); - indexer_id_to_cpu_capacities.insert("node1".to_string(), mcpu(10_000)); + let mut indexer_infos = FnvHashMap::default(); + indexer_infos.insert("node1".to_string(), IndexerInfo::for_test(mcpu(10_000))); let mut indexing_plan = PhysicalIndexingPlan::with_indexer_ids(&["node1".to_string()]); for indexing_task in indexing_tasks { indexing_plan.add_indexing_task("node1", indexing_task); } let shard_locations = ShardLocations::default(); - let new_plan = build_physical_indexing_plan( + let new_plan = build_physical_indexing_plan_without_locality( &sources, - &indexer_id_to_cpu_capacities, + &indexer_infos, Some(&indexing_plan), &shard_locations, ); @@ -1134,16 +1552,16 @@ mod tests { params_fingerprint: 0, }]; const NODE: &str = "node1"; - let mut indexer_id_to_cpu_capacities = FnvHashMap::default(); - indexer_id_to_cpu_capacities.insert(NODE.to_string(), mcpu(10_000)); + let mut indexer_infos = FnvHashMap::default(); + indexer_infos.insert(NODE.to_string(), IndexerInfo::for_test(mcpu(10_000))); let mut indexing_plan = PhysicalIndexingPlan::with_indexer_ids(&["node1".to_string()]); for indexing_task in indexing_tasks { indexing_plan.add_indexing_task(NODE, indexing_task); } let shard_locations = ShardLocations::default(); - let new_plan = build_physical_indexing_plan( + let new_plan = build_physical_indexing_plan_without_locality( &sources, - &indexer_id_to_cpu_capacities, + &indexer_infos, Some(&indexing_plan), &shard_locations, ); @@ -1278,6 +1696,16 @@ mod tests { remaining_num_shards_per_node .insert(node2.as_str().to_string(), NonZeroU32::new(1).unwrap()); + let mut indexer_infos = FnvHashMap::default(); + indexer_infos.insert( + node1.as_str().to_string(), + IndexerInfo::for_test(mcpu(4_000)), + ); + indexer_infos.insert( + node2.as_str().to_string(), + IndexerInfo::for_test(mcpu(4_000)), + ); + let mut shard_locations: ShardLocations = ShardLocations::default(); // shard1 on 1 shard_locations.add_location(&shard1, &node1); @@ -1292,6 +1720,7 @@ mod tests { missing_shards, remaining_num_shards_per_node, &shard_locations, + &indexer_infos, ); assert_eq!(shard_to_indexer.len(), 4); assert_eq!(shard_to_indexer.get(&shard1).unwrap(), "node1"); @@ -1327,10 +1756,16 @@ mod tests { params_fingerprint: 0, }, ]; - let mut capacities = FnvHashMap::default(); - capacities.insert("indexer-1".to_string(), CpuCapacity::from_cpu_millis(8000)); + let mut indexer_infos = FnvHashMap::default(); + let indexer_info = IndexerInfo::for_test(CpuCapacity::from_cpu_millis(8000)); + indexer_infos.insert("indexer-1".to_string(), indexer_info); let shard_locations = ShardLocations::default(); - build_physical_indexing_plan(&sources_to_schedule, &capacities, None, &shard_locations); + build_physical_indexing_plan_without_locality( + &sources_to_schedule, + &indexer_infos, + None, + &shard_locations, + ); } #[test] diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs index e702006a7ae..8f6ed37eab6 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs @@ -416,13 +416,13 @@ fn available_cpu_capacity( fn place_self_hosted_shards_on_indexer( source: &Source, indexer_ord: IndexerOrd, - num_self_hosted_shards: u32, + num_shards_to_place: u32, problem: &SchedulingProblem, solution: &mut SchedulingSolution, ) -> u32 { let available_capacity = available_cpu_capacity(indexer_ord, problem, solution); let num_placable_shards = available_capacity.cpu_millis() / source.load_per_shard; - let num_shards_placed = num_placable_shards.min(num_self_hosted_shards); + let num_shards_placed = num_placable_shards.min(num_shards_to_place); solution.indexer_assignments[indexer_ord].add_shards(source.source_ord, num_shards_placed); num_shards_placed } @@ -439,18 +439,24 @@ fn place_self_hosted_shards( for unassigned_source in &unassigned_sources { let source_ord = unassigned_source.source_ord as usize; let leftover_shards = &mut leftover_shards_per_source[source_ord]; + let mut num_unaccounted_shards = unassigned_source.num_shards; for (&indexer_ord, &num_self_hosted_shards) in &unassigned_source.affinities { + let num_shards_to_place = num_self_hosted_shards.min(num_unaccounted_shards); + if num_shards_to_place == 0 { + break; + } let num_shards_placed = place_self_hosted_shards_on_indexer( unassigned_source, indexer_ord, - num_self_hosted_shards, + num_shards_to_place, problem, solution, ); + num_unaccounted_shards -= num_shards_to_place; let Some(locality_group) = problem.indexer_locality_group(indexer_ord) else { continue; }; - leftover_shards.add(locality_group, num_self_hosted_shards - num_shards_placed); + leftover_shards.add(locality_group, num_shards_to_place - num_shards_placed); } } leftover_shards_per_source @@ -788,6 +794,36 @@ mod tests { assert_eq!(solution.indexer_assignments[4].num_shards(2), 2); } + #[test] + fn test_self_hosted_only_indexer_removes_foreign_work() { + let draining_locality = IndexerLocality { + group: Some(LocalityGroup::from_ord(0)), + eligibility: Eligibility::SelfHostedOnly, + }; + let ready_locality = IndexerLocality { + group: Some(LocalityGroup::from_ord(0)), + eligibility: Eligibility::Any, + }; + let mut problem = SchedulingProblem::with_indexer_localities( + vec![mcpu(3_000), mcpu(4_000)], + vec![draining_locality, ready_locality], + ); + problem.add_source(3, NonZeroU32::new(1_000).unwrap()); + problem.add_source(1, NonZeroU32::new(1_000).unwrap()); + problem.inc_affinity(0, 0); + + let mut previous_solution = problem.new_solution(); + previous_solution.indexer_assignments[0].add_shards(0, 2); + previous_solution.indexer_assignments[0].add_shards(1, 1); + + let solution = attempt_solve(&problem, previous_solution).unwrap(); + + assert_eq!(solution.indexer_assignments[0].num_shards(0), 1); + assert_eq!(solution.indexer_assignments[0].num_shards(1), 0); + assert_eq!(solution.indexer_assignments[1].num_shards(0), 2); + assert_eq!(solution.indexer_assignments[1].num_shards(1), 1); + } + #[test] fn test_compute_unassigned_shards_simple() { let mut problem = SchedulingProblem::with_indexer_cpu_capacities(vec![mcpu(4_000)]); @@ -997,6 +1033,126 @@ mod tests { }) } + fn locality_group_strat(num_groups: usize) -> impl Strategy> { + prop_oneof![ + 1 => Just(None), + 4 => (0..num_groups).prop_map(|group_ord| Some(LocalityGroup::from_ord(group_ord))), + ] + } + + fn eligibility_strat() -> impl Strategy { + prop_oneof![ + 3 => Just(Eligibility::Any), + 1 => Just(Eligibility::SelfHostedOnly), + ] + } + + fn locality_source_strat( + num_indexers: usize, + ) -> impl Strategy)> { + let load_strat = prop_oneof![ + Just(1u32), + Just(250u32), + Just(1_000u32), + Just(1_200u32), + Just(3_200u32), + 1u32..1_000u32, + ]; + (0u32..12u32, load_strat).prop_flat_map(move |(num_shards, load)| { + let host_strat = prop_oneof![ + 3 => (0..num_indexers).prop_map(Some), + 1 => Just(None), + ]; + let shard_hosts_strat = proptest::collection::vec(host_strat, num_shards as usize); + let load_per_shard = NonZeroU32::new(load).unwrap(); + shard_hosts_strat.prop_map(move |shard_hosts| { + let hosting_indexer_ords: Vec = + shard_hosts.into_iter().flatten().collect(); + (num_shards, load_per_shard, hosting_indexer_ords) + }) + }) + } + + fn locality_problem_strategy( + num_indexers: usize, + num_sources: usize, + num_groups: usize, + ) -> impl Strategy { + let cpu_capacities_strat = + proptest::collection::vec(indexer_cpu_capacity_strat(), num_indexers); + let groups_strat = proptest::collection::vec(locality_group_strat(num_groups), num_indexers); + let eligibilities_strat = proptest::collection::vec(eligibility_strat(), num_indexers); + let sources_strat = + proptest::collection::vec(locality_source_strat(num_indexers), num_sources); + ( + cpu_capacities_strat, + groups_strat, + eligibilities_strat, + sources_strat, + ) + .prop_map( + |(cpu_capacities, groups, mut eligibilities, sources)| { + eligibilities[0] = Eligibility::Any; + let indexer_localities: Vec = groups + .into_iter() + .zip(eligibilities) + .map(|(group, eligibility)| IndexerLocality { group, eligibility }) + .collect(); + let mut problem = SchedulingProblem::with_indexer_localities( + cpu_capacities, + indexer_localities, + ); + for (num_shards, load_per_shard, hosting_indexer_ords) in sources { + let source_ord = problem.add_source(num_shards, load_per_shard); + for hosting_indexer_ord in hosting_indexer_ords { + problem.inc_affinity(source_ord, hosting_indexer_ord); + } + } + problem + }, + ) + } + + fn locality_problem_solution_strategy() + -> impl Strategy { + (1usize..8, 0usize..8, 1usize..4).prop_flat_map( + |(num_indexers, num_sources, num_groups)| { + ( + locality_problem_strategy(num_indexers, num_sources, num_groups), + initial_solution_strategy(num_indexers, num_sources), + ) + }, + ) + } + + proptest! { + #[test] + fn test_proptest_locality_aware_idempotence((problem, solution) in locality_problem_solution_strategy()) { + let solution_1 = solve(problem.clone(), solution); + let solution_2 = solve(problem.clone(), solution_1.clone()); + assert_eq!( + solution_1.indexer_assignments, solution_2.indexer_assignments, + "solution unstable!\nSolution 1: {solution_1:?}\nSolution 2: {solution_2:?}" + ); + for indexer_assignment in &solution_1.indexer_assignments { + let indexer_ord = indexer_assignment.indexer_ord; + if problem.is_eligible_for_foreign_shards(indexer_ord) { + continue; + } + for source in problem.sources() { + let num_shards = indexer_assignment.num_shards(source.source_ord); + let num_self_hosted_shards = problem.source_affinity(source.source_ord, indexer_ord); + assert!( + num_shards <= num_self_hosted_shards, + "self-hosted-only indexer {indexer_ord} holds {num_shards} shards of source \ + {} but hosts {num_self_hosted_shards}", + source.source_ord + ); + } + } + } + } + #[test] fn test_problem_missing_capacities() { let mut problem = diff --git a/quickwit/quickwit-control-plane/src/tests.rs b/quickwit/quickwit-control-plane/src/tests.rs index bdfecd0701a..581844951ec 100644 --- a/quickwit/quickwit-control-plane/src/tests.rs +++ b/quickwit/quickwit-control-plane/src/tests.rs @@ -86,6 +86,7 @@ pub fn test_indexer_change_stream( indexing_tasks, indexing_capacity: CpuCapacity::from_cpu_millis(4_000), ingester_status: node.ingester_status, + availability_zone: None, }, ); Some(change) From b79f2318e39d17816d2cd00ef3de9a8d53b8367c Mon Sep 17 00:00:00 2001 From: Nadav Gov-Ari Date: Fri, 14 Aug 2026 16:55:06 -0400 Subject: [PATCH 3/4] Add a bunch of tests, including all sorts of crazy regression tests --- .../src/indexing_scheduler/mod.rs | 37 +- .../scheduling/churn_tests.rs | 352 ++++++++++ .../src/indexing_scheduler/scheduling/mod.rs | 241 +++++-- .../scheduling/scale_tests.rs | 630 ++++++++++++++++++ .../scheduling/scheduling_logic.rs | 290 ++++++-- 5 files changed, 1450 insertions(+), 100 deletions(-) create mode 100644 quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/churn_tests.rs create mode 100644 quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scale_tests.rs diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index bef4a04bba3..30e3669ff78 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -324,6 +324,13 @@ fn build_indexer_info( draining_eligibility: Eligibility, locality_aware: bool, ) -> IndexerInfo { + if !locality_aware { + return IndexerInfo { + cpu_capacity: indexer.indexing_capacity, + availability_zone: None, + eligibility: Eligibility::Any, + }; + } let eligibility = match indexer.ingester_status { IngesterStatus::Ready => Eligibility::Any, // For draining indexers, if they're the last ones left in the cluster, they need to be @@ -331,14 +338,9 @@ fn build_indexer_info( // their own (Eligibility::SelfHostedOnly). _ => draining_eligibility, }; - let availability_zone = if locality_aware { - indexer.availability_zone.clone() - } else { - None - }; IndexerInfo { cpu_capacity: indexer.indexing_capacity, - availability_zone, + availability_zone: indexer.availability_zone.clone(), eligibility, } } @@ -1128,6 +1130,18 @@ mod tests { indexing_plans_diff.nodes_with_changed_ingester_status, FnvHashSet::from_iter(["indexer-1"]) ); + + let mirrored_plans_diff = get_indexing_plans_diff( + &running_plan, + &desired_plan, + &last_applied_statuses, + &running_statuses, + ); + assert!(!mirrored_plans_diff.has_same_nodes()); + assert_eq!( + mirrored_plans_diff.nodes_with_changed_ingester_status, + FnvHashSet::from_iter(["indexer-1"]) + ); } } @@ -1542,12 +1556,21 @@ mod tests { { let mut ready = mock_indexer_node_info("indexer-ready", IngesterStatus::Ready); ready.availability_zone = Some("az-a".to_string()); - let indexers = vec![ready]; + let mut retiring = + mock_indexer_node_info("indexer-retiring", IngesterStatus::Retiring); + retiring.availability_zone = Some("az-b".to_string()); + let indexers = vec![ready, retiring]; let locality_unaware = false; let indexer_infos = build_indexer_infos(&indexers, locality_unaware); assert_eq!(indexer_infos["indexer-ready"].availability_zone, None); + assert_eq!(indexer_infos["indexer-retiring"].availability_zone, None); + assert_eq!(indexer_infos["indexer-ready"].eligibility, Eligibility::Any); + assert_eq!( + indexer_infos["indexer-retiring"].eligibility, + Eligibility::Any + ); } } diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/churn_tests.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/churn_tests.rs new file mode 100644 index 00000000000..91b8959dd0b --- /dev/null +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/churn_tests.rs @@ -0,0 +1,352 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::num::NonZeroU32; + +use fnv::{FnvHashMap, FnvHashSet}; +use quickwit_proto::indexing::CpuCapacity; +use quickwit_proto::types::{IndexUid, ShardId, SourceUid}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +use super::{ + Eligibility, IndexerInfo, IndexerSpec, SourceToSchedule, SourceToScheduleType, + build_physical_indexing_plan, shard_ids_of_source, total_num_shards, +}; +use crate::indexing_plan::PhysicalIndexingPlan; +use crate::model::ShardLocations; + +const NUM_INDEXERS: usize = 500; +const NUM_AZS: usize = 3; +const INDEXER_CPU_CAPACITY: CpuCapacity = CpuCapacity::from_cpu_millis(4_000); +const SEED: u64 = 20_260_814; +const LOADS_PER_SHARD: [u32; 8] = [50, 250, 500, 1_000, 1_200, 1_600, 2_000, 3_200]; + +const CHURN_NUM_SOURCES: usize = 100; +const CHURN_ITERATIONS: usize = 200; +const CHURN_PHASE_LENGTH: usize = 15; +const CHURN_MAX_DRAINING_INDEXERS: usize = NUM_INDEXERS / 10; +const CHURN_SCALE_UP_FACTOR: f32 = 1.5; +const CHURN_MAX_SHARDS: usize = 4_000; + +struct ChurnWorld { + indexer_specs: Vec, + sources: Vec, + host_ord_per_shard: FnvHashMap, + draining_indexer_ords: FnvHashSet, + next_shard_id: u64, +} + +impl ChurnWorld { + fn new(rng: &mut StdRng) -> ChurnWorld { + let indexer_specs = (0..NUM_INDEXERS) + .map(|indexer_ord| { + let availability_zone = format!("az-{}", indexer_ord % NUM_AZS); + let node_id = format!("indexer-{indexer_ord}"); + IndexerSpec::new(&node_id, INDEXER_CPU_CAPACITY, Some(&availability_zone)) + }) + .collect(); + let index_uid = IndexUid::for_test("churn-index", 0); + let mut world = ChurnWorld { + indexer_specs, + sources: Vec::new(), + host_ord_per_shard: FnvHashMap::default(), + draining_indexer_ords: FnvHashSet::default(), + next_shard_id: 0, + }; + for source_ord in 0..CHURN_NUM_SOURCES { + let load_per_shard = LOADS_PER_SHARD[rng.random_range(0..LOADS_PER_SHARD.len())]; + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: format!("source-{source_ord}"), + }; + world.sources.push(SourceToSchedule { + source_uid, + source_type: SourceToScheduleType::Sharded { + shard_ids: Vec::new(), + load_per_shard: NonZeroU32::new(load_per_shard).unwrap(), + }, + params_fingerprint: 0, + }); + world.open_shards(source_ord, 1 + rng.random_range(0..8), rng); + } + world + } + + fn open_shards(&mut self, source_ord: usize, num_shards_to_open: usize, rng: &mut StdRng) { + let mut host_ords = Vec::with_capacity(num_shards_to_open); + for _ in 0..num_shards_to_open { + host_ords.push(self.pick_routable_indexer_ord(rng)); + } + let SourceToScheduleType::Sharded { shard_ids, .. } = + &mut self.sources[source_ord].source_type + else { + return; + }; + for host_ord in host_ords { + let shard_id = ShardId::from(self.next_shard_id); + self.next_shard_id += 1; + shard_ids.push(shard_id.clone()); + self.host_ord_per_shard.insert(shard_id, host_ord); + } + } + + fn pick_routable_indexer_ord(&self, rng: &mut StdRng) -> usize { + loop { + let indexer_ord = rng.random_range(0..self.indexer_specs.len()); + if !self.draining_indexer_ords.contains(&indexer_ord) { + return indexer_ord; + } + } + } + + fn close_shards(&mut self, source_ord: usize, num_shards_to_close: usize) { + let SourceToScheduleType::Sharded { shard_ids, .. } = + &mut self.sources[source_ord].source_type + else { + return; + }; + for _ in 0..num_shards_to_close { + if shard_ids.len() <= 1 { + return; + } + let shard_id = shard_ids.remove(0); + self.host_ord_per_shard.remove(&shard_id); + } + } + + fn close_shards_on_draining_indexers(&mut self, rng: &mut StdRng) { + for source_ord in 0..self.sources.len() { + let drained_shard_ids: Vec = shard_ids_of_source(&self.sources[source_ord]) + .iter() + .filter(|shard_id| { + let host_ord = self.host_ord_per_shard[*shard_id]; + self.draining_indexer_ords.contains(&host_ord) + }) + .filter(|_| rng.random_range(0..100) < 40) + .cloned() + .collect(); + for shard_id in drained_shard_ids { + let SourceToScheduleType::Sharded { shard_ids, .. } = + &mut self.sources[source_ord].source_type + else { + continue; + }; + if shard_ids.len() <= 1 { + break; + } + shard_ids.retain(|candidate| candidate != &shard_id); + self.host_ord_per_shard.remove(&shard_id); + } + } + } + + fn num_hosted_shards(&self, indexer_ord: usize) -> usize { + self.host_ord_per_shard + .values() + .filter(|host_ord| **host_ord == indexer_ord) + .count() + } + + fn advance_drain_lifecycle(&mut self, rng: &mut StdRng) { + let fully_drained: Vec = self + .draining_indexer_ords + .iter() + .copied() + .filter(|indexer_ord| self.num_hosted_shards(*indexer_ord) == 0) + .collect(); + for indexer_ord in fully_drained { + self.draining_indexer_ords.remove(&indexer_ord); + } + while self.draining_indexer_ords.len() < CHURN_MAX_DRAINING_INDEXERS { + let indexer_ord = rng.random_range(0..self.indexer_specs.len()); + if !self.draining_indexer_ords.insert(indexer_ord) { + break; + } + } + } + + fn scale_shards(&mut self, growing: bool, rng: &mut StdRng) { + let num_shards_in_cluster = total_num_shards(&self.sources); + for source_ord in 0..self.sources.len() { + if rng.random_range(0..100) >= 25 { + continue; + } + let num_shards = shard_ids_of_source(&self.sources[source_ord]).len(); + if growing && num_shards_in_cluster < CHURN_MAX_SHARDS { + let target_num_shards = + (num_shards as f32 * CHURN_SCALE_UP_FACTOR).ceil() as usize; + self.open_shards(source_ord, target_num_shards - num_shards, rng); + } else if !growing { + let target_num_shards = + (num_shards as f32 / CHURN_SCALE_UP_FACTOR).floor().max(1.0) as usize; + self.close_shards(source_ord, num_shards - target_num_shards); + } + } + } + + fn shard_locations(&self) -> ShardLocations<'_> { + let mut shard_locations = ShardLocations::default(); + for source in &self.sources { + for shard_id in shard_ids_of_source(source) { + let host_ord = self.host_ord_per_shard[shard_id]; + shard_locations.add_location(shard_id, &self.indexer_specs[host_ord].node_id); + } + } + shard_locations + } + + fn indexer_infos(&self, locality_aware: bool) -> FnvHashMap { + let mut indexer_infos = FnvHashMap::default(); + for (indexer_ord, indexer_spec) in self.indexer_specs.iter().enumerate() { + let draining = self.draining_indexer_ords.contains(&indexer_ord); + if !locality_aware { + if draining { + continue; + } + let indexer_info = IndexerInfo::for_test(INDEXER_CPU_CAPACITY); + indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); + continue; + } + let eligibility = if draining { + Eligibility::SelfHostedOnly + } else { + Eligibility::Any + }; + let indexer_info = indexer_spec.to_indexer_info(eligibility); + indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); + } + indexer_infos + } +} + +fn indexer_per_shard(plan: &PhysicalIndexingPlan) -> FnvHashMap<&ShardId, &String> { + let mut indexer_per_shard = FnvHashMap::default(); + for (indexer, tasks) in plan.indexing_tasks_per_indexer() { + for task in tasks { + for shard_id in &task.shard_ids { + indexer_per_shard.insert(shard_id, indexer); + } + } + } + indexer_per_shard +} + +fn count_shards_that_moved( + plan: &PhysicalIndexingPlan, + replanned: &PhysicalIndexingPlan, +) -> (usize, usize) { + let indexer_per_shard_before = indexer_per_shard(plan); + let indexer_per_shard_after = indexer_per_shard(replanned); + let mut num_surviving_shards = 0; + let mut num_moved_shards = 0; + for (shard_id, indexer_before) in &indexer_per_shard_before { + let Some(indexer_after) = indexer_per_shard_after.get(*shard_id) else { + continue; + }; + num_surviving_shards += 1; + if indexer_before != indexer_after { + num_moved_shards += 1; + } + } + (num_moved_shards, num_surviving_shards) +} + +#[derive(Default)] +struct ChurnTally { + num_moved_shards: usize, + num_surviving_shards: usize, + max_moved_percent: f32, +} + +impl ChurnTally { + fn record(&mut self, num_moved_shards: usize, num_surviving_shards: usize) { + self.num_moved_shards += num_moved_shards; + self.num_surviving_shards += num_surviving_shards; + let moved_percent = num_moved_shards as f32 * 100.0 / num_surviving_shards.max(1) as f32; + self.max_moved_percent = self.max_moved_percent.max(moved_percent); + } + + fn moved_percent(&self) -> f32 { + self.num_moved_shards as f32 * 100.0 / self.num_surviving_shards.max(1) as f32 + } +} + +/// Churn across successive plans, which is what the cluster actually pays for: every shard that +/// changes indexer is a pipeline restart and a re-read over the network. +/// +/// The simulation is a cluster under ingest pressure. Shards scale up by 1.5x per event on a +/// quarter of the sources and drain away again, indexers enter and leave decommissioning, and the +/// plan is rebuilt after every change from the plan before it. +/// +/// The same events are replayed with locality awareness disabled, which also means draining +/// indexers are left out of planning entirely, as that path selects only ready ones. +#[test] +fn test_churn_across_successive_plans() { + let mut rng = StdRng::seed_from_u64(SEED); + let mut world = ChurnWorld::new(&mut rng); + let mut previous_plans: [Option; 2] = [None, None]; + let mut tallies: [ChurnTally; 2] = [ChurnTally::default(), ChurnTally::default()]; + let mut min_num_shards = usize::MAX; + let mut max_num_shards = 0; + + for iteration in 0..CHURN_ITERATIONS { + let growing = (iteration / CHURN_PHASE_LENGTH) % 2 == 0; + world.scale_shards(growing, &mut rng); + world.advance_drain_lifecycle(&mut rng); + world.close_shards_on_draining_indexers(&mut rng); + + let shard_locations = world.shard_locations(); + let num_shards = total_num_shards(&world.sources); + min_num_shards = min_num_shards.min(num_shards); + max_num_shards = max_num_shards.max(num_shards); + + for (path_ord, locality_aware) in [true, false].into_iter().enumerate() { + let indexer_infos = world.indexer_infos(locality_aware); + let plan = build_physical_indexing_plan( + &world.sources, + &indexer_infos, + locality_aware, + previous_plans[path_ord].as_ref(), + &shard_locations, + ); + if let Some(previous_plan) = &previous_plans[path_ord] { + let (num_moved_shards, num_surviving_shards) = + count_shards_that_moved(previous_plan, &plan); + tallies[path_ord].record(num_moved_shards, num_surviving_shards); + } + previous_plans[path_ord] = Some(plan); + } + } + + println!( + "{CHURN_ITERATIONS} rebuilds, shards {min_num_shards}..{max_num_shards}\n new moved \ + {} of {} ({:.2}%), worst rebuild {:.1}%\n legacy moved {} of {} ({:.2}%), worst rebuild \ + {:.1}%", + tallies[0].num_moved_shards, + tallies[0].num_surviving_shards, + tallies[0].moved_percent(), + tallies[0].max_moved_percent, + tallies[1].num_moved_shards, + tallies[1].num_surviving_shards, + tallies[1].moved_percent(), + tallies[1].max_moved_percent, + ); + assert!( + tallies[0].moved_percent() <= tallies[1].moved_percent() * 1.2, + "moved {:.2}% of shards against legacy's {:.2}% over the same events", + tallies[0].moved_percent(), + tallies[1].moved_percent() + ); +} diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs index 712dc53cf0b..a3d901055df 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs @@ -12,6 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(test)] +mod churn_tests; +#[cfg(test)] +mod scale_tests; pub mod scheduling_logic; pub mod scheduling_logic_model; @@ -21,6 +25,8 @@ use std::num::NonZeroU32; use fnv::{FnvHashMap, FnvHashSet}; use quickwit_common::rate_limited_debug; use quickwit_proto::indexing::{CpuCapacity, IndexingTask}; +#[cfg(test)] +use quickwit_proto::types::NodeId; use quickwit_proto::types::{PipelineUid, ShardId, SourceUid}; pub use scheduling_logic_model::Eligibility; use scheduling_logic_model::{IndexerLocality, IndexerOrd, LocalityGroup, SourceOrd}; @@ -169,6 +175,64 @@ impl IndexerInfo { } } +#[cfg(test)] +pub(crate) struct IndexerSpec { + pub node_id: NodeId, + pub cpu_capacity: CpuCapacity, + pub availability_zone: Option, +} + +#[cfg(test)] +impl IndexerSpec { + pub(crate) fn new( + node_id: &str, + cpu_capacity: CpuCapacity, + availability_zone: Option<&str>, + ) -> IndexerSpec { + IndexerSpec { + node_id: NodeId::from_str(node_id), + cpu_capacity, + availability_zone: availability_zone.map(|az| az.to_string()), + } + } + + pub(crate) fn to_indexer_info(&self, eligibility: Eligibility) -> IndexerInfo { + IndexerInfo { + cpu_capacity: self.cpu_capacity, + availability_zone: self.availability_zone.clone(), + eligibility, + } + } +} + +#[cfg(test)] +pub(crate) fn shard_ids_of_source(source: &SourceToSchedule) -> &[ShardId] { + let SourceToScheduleType::Sharded { shard_ids, .. } = &source.source_type else { + return &[]; + }; + shard_ids +} + +#[cfg(test)] +pub(crate) fn total_num_shards(sources: &[SourceToSchedule]) -> usize { + sources + .iter() + .map(|source| shard_ids_of_source(source).len()) + .sum() +} + +#[cfg(test)] +pub(crate) fn shard_ids_for_indexer(plan: &PhysicalIndexingPlan, indexer: &str) -> Vec { + let mut shard_ids: Vec = plan + .indexer(indexer) + .unwrap() + .iter() + .flat_map(|task| task.shard_ids.iter().cloned()) + .collect(); + shard_ids.sort(); + shard_ids +} + #[derive(Debug)] pub struct SourceToSchedule { pub source_uid: SourceUid, @@ -413,6 +477,7 @@ fn convert_scheduling_solution_to_physical_plan( HashMap::with_capacity(new_physical_plan.num_indexers()); for (indexer, indexing_tasks) in new_physical_plan.indexing_tasks_per_indexer_mut() { let indexer_ord = id_to_ord_map.indexer_ord(indexer).unwrap(); + let eligibility = indexer_infos[indexer].eligibility; let mut num_shards_for_indexer_source: u32 = indexer_assignments[indexer_ord].num_shards(source_ord); for indexing_task in indexing_tasks { @@ -420,6 +485,15 @@ fn convert_scheduling_solution_to_physical_plan( && indexing_task.source_id == source.source_uid.source_id { indexing_task.shard_ids.retain(|shard_id| { + if !may_keep_shard_in_previous_pipeline( + indexer, + eligibility, + shard_id, + shard_locations, + indexer_infos, + ) { + return false; + } let shard_added = scheduled_shards.insert(shard_id.clone()); if shard_added { true @@ -474,6 +548,45 @@ fn convert_scheduling_solution_to_physical_plan( new_physical_plan } +fn is_shard_local(indexer: &str, shard_id: &ShardId, shard_locations: &ShardLocations) -> bool { + shard_locations + .get_shard_locations(shard_id) + .iter() + .any(|node_id| node_id.as_str() == indexer) +} + +fn is_shard_hosted_on_draining_indexer( + shard_id: &ShardId, + shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, +) -> bool { + shard_locations + .get_shard_locations(shard_id) + .iter() + .any(|node_id| { + let Some(indexer_info) = indexer_infos.get(node_id.as_str()) else { + return false; + }; + indexer_info.eligibility == Eligibility::SelfHostedOnly + }) +} + +fn may_keep_shard_in_previous_pipeline( + indexer: &str, + eligibility: Eligibility, + shard_id: &ShardId, + shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, +) -> bool { + if is_shard_local(indexer, shard_id, shard_locations) { + return true; + } + if eligibility == Eligibility::SelfHostedOnly { + return false; + } + !is_shard_hosted_on_draining_indexer(shard_id, shard_locations, indexer_infos) +} + fn indexer_availability_zone<'a>( node_id: &str, indexer_infos: &'a FnvHashMap, @@ -891,10 +1004,10 @@ mod tests { use super::scheduling_logic::solve; use super::{ - Eligibility, IndexerInfo, SourceToSchedule, SourceToScheduleType, + Eligibility, IndexerInfo, IndexerSpec, SourceToSchedule, SourceToScheduleType, build_physical_indexing_plan, build_physical_indexing_plan_without_locality, convert_scheduling_solution_to_physical_plan_single_node_single_source, - convert_to_simplified_problem, + convert_to_simplified_problem, shard_ids_for_indexer, }; use crate::indexing_plan::PhysicalIndexingPlan; use crate::indexing_scheduler::get_shard_locality_metrics; @@ -913,45 +1026,6 @@ mod tests { } } - struct IndexerSpec { - node_id: NodeId, - cpu_capacity: CpuCapacity, - availability_zone: Option, - } - - impl IndexerSpec { - fn new( - node_id: &str, - cpu_capacity: CpuCapacity, - availability_zone: Option<&str>, - ) -> IndexerSpec { - IndexerSpec { - node_id: NodeId::from_str(node_id), - cpu_capacity, - availability_zone: availability_zone.map(|az| az.to_string()), - } - } - - fn to_indexer_info(&self) -> IndexerInfo { - IndexerInfo { - cpu_capacity: self.cpu_capacity, - availability_zone: self.availability_zone.clone(), - eligibility: Eligibility::Any, - } - } - } - - fn shard_ids_for_indexer(plan: &PhysicalIndexingPlan, indexer: &str) -> Vec { - let mut shard_ids: Vec = plan - .indexer(indexer) - .unwrap() - .iter() - .flat_map(|task| task.shard_ids.iter().cloned()) - .collect(); - shard_ids.sort(); - shard_ids - } - fn shard_counts_per_az( plan: &PhysicalIndexingPlan, indexer_infos: &FnvHashMap, @@ -1215,7 +1289,7 @@ mod tests { let mut indexer_infos = FnvHashMap::default(); for indexer_spec in indexer_specs { - let indexer_info = indexer_spec.to_indexer_info(); + let indexer_info = indexer_spec.to_indexer_info(Eligibility::Any); indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); } @@ -1249,7 +1323,7 @@ mod tests { let mut reversed_indexer_infos = FnvHashMap::default(); for indexer_spec in indexer_specs.iter().rev() { - let indexer_info = indexer_spec.to_indexer_info(); + let indexer_info = indexer_spec.to_indexer_info(Eligibility::Any); reversed_indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); } let reversed_plan = build_physical_indexing_plan( @@ -1340,6 +1414,87 @@ mod tests { } } + #[test] + fn test_small_clusters_are_not_unbalanced() { + const NUM_SHARDS: usize = 12; + for num_indexers in [2usize, 4, 5] { + let indexer_specs: Vec = (0..num_indexers) + .map(|indexer_ord| { + let node_id = format!("indexer{indexer_ord}"); + IndexerSpec::new(&node_id, mcpu(4_000), Some("az-a")) + }) + .collect(); + let shard_ids: Vec = (0..NUM_SHARDS as u64).map(ShardId::from).collect(); + let source = SourceToSchedule { + source_uid: source_id(), + source_type: SourceToScheduleType::Sharded { + shard_ids: shard_ids.clone(), + load_per_shard: NonZeroU32::new(1_000).unwrap(), + }, + params_fingerprint: 0, + }; + let sources = vec![source]; + + let mut shard_locations = ShardLocations::default(); + for shard_id in &shard_ids { + shard_locations.add_location(shard_id, &indexer_specs[0].node_id); + } + + let mut indexer_infos = FnvHashMap::default(); + for indexer_spec in &indexer_specs { + let indexer_info = indexer_spec.to_indexer_info(Eligibility::Any); + indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); + } + + let locality_aware = true; + let plan = build_physical_indexing_plan( + &sources, + &indexer_infos, + locality_aware, + None, + &shard_locations, + ); + + let mut legacy_indexer_infos = FnvHashMap::default(); + for indexer_spec in &indexer_specs { + let indexer_info = IndexerInfo::for_test(mcpu(4_000)); + legacy_indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); + } + let legacy_plan = build_physical_indexing_plan_without_locality( + &sources, + &legacy_indexer_infos, + None, + &shard_locations, + ); + + let mut shard_counts: Vec = indexer_specs + .iter() + .map(|indexer_spec| { + shard_ids_for_indexer(&plan, indexer_spec.node_id.as_str()).len() + }) + .collect(); + shard_counts.sort(); + let mut legacy_shard_counts: Vec = indexer_specs + .iter() + .map(|indexer_spec| { + shard_ids_for_indexer(&legacy_plan, indexer_spec.node_id.as_str()).len() + }) + .collect(); + legacy_shard_counts.sort(); + let fair_share = NUM_SHARDS.div_ceil(num_indexers); + let max_shard_count = *shard_counts.last().unwrap(); + println!( + "{num_indexers} indexers hold {shard_counts:?}, legacy holds {legacy_shard_counts:?}" + ); + assert_eq!(shard_counts.iter().sum::(), NUM_SHARDS); + assert_eq!(shard_counts, legacy_shard_counts); + assert!( + max_shard_count <= fair_share + 1, + "{num_indexers} indexers hold {shard_counts:?}, fair share is {fair_share}" + ); + } + } + #[test] fn test_draining_indexer_keeps_only_hosted_shard_ids() { let shard0 = ShardId::from(0); diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scale_tests.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scale_tests.rs new file mode 100644 index 00000000000..1c8c60a041e --- /dev/null +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scale_tests.rs @@ -0,0 +1,630 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::num::NonZeroU32; + +use fnv::{FnvHashMap, FnvHashSet}; +use quickwit_proto::indexing::CpuCapacity; +use quickwit_proto::types::{IndexUid, PipelineUid, ShardId, SourceUid}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +use super::{ + Eligibility, IndexerInfo, IndexerSpec, SourceToSchedule, SourceToScheduleType, + build_physical_indexing_plan, compute_max_num_shards_per_pipeline, shard_ids_for_indexer, + shard_ids_of_source, total_num_shards, +}; +use crate::indexing_plan::PhysicalIndexingPlan; +use crate::indexing_scheduler::get_shard_locality_metrics; +use crate::model::ShardLocations; + +const NUM_INDEXERS: usize = 500; +const NUM_SOURCES: usize = 1_000; +const NUM_AZS: usize = 3; +const NUM_DRAINING_INDEXERS: usize = 250; +const MAX_SHARDS_PER_SOURCE: usize = 4_000; +const MAX_TOTAL_SHARDS: usize = 20_000; +const INDEXER_CPU_CAPACITY: CpuCapacity = CpuCapacity::from_cpu_millis(4_000); +const SEED: u64 = 20_260_814; +const LOADS_PER_SHARD: [u32; 8] = [50, 250, 500, 1_000, 1_200, 1_600, 2_000, 3_200]; + +struct SourceSizeTier { + num_sources: usize, + max_shards_per_source: usize, +} + +const SOURCE_SIZE_TIERS: [SourceSizeTier; 5] = [ + SourceSizeTier { + num_sources: 2, + max_shards_per_source: MAX_SHARDS_PER_SOURCE, + }, + SourceSizeTier { + num_sources: 10, + max_shards_per_source: 500, + }, + SourceSizeTier { + num_sources: 50, + max_shards_per_source: 60, + }, + SourceSizeTier { + num_sources: 200, + max_shards_per_source: 12, + }, + SourceSizeTier { + num_sources: 738, + max_shards_per_source: 1, + }, +]; + +/// Scheduling at scale: 500 indexers, 3 zones, 1000 sources, ~10k shards, fixed seed. +/// +/// These tests attempt to assert that indexing planning is resilient to "scale" conditions, +/// where the assumptions we make for solving the "normal" problem might break down. They attempt +/// to answer the questions +/// * "Does this work well at petabyte+ scale?" +/// * "Does a large drain of indexers across multiple AZs result in a balanced plan?" +/// * "Does the removal of a whole AZ result in a balanced plan?" +/// +/// Assertions are on <1% drift overall. +#[test] +fn test_scale_all_indexers_ready() { + let mut rng = StdRng::seed_from_u64(SEED); + let indexer_specs = build_indexer_specs(); + let sources = build_sources(&mut rng); + let shard_locations = build_shard_locations(&sources, &indexer_specs, &mut rng); + let no_draining_indexers = FnvHashSet::default(); + let indexer_infos = build_indexer_infos(&indexer_specs, &no_draining_indexers); + let num_shards = total_num_shards(&sources); + + let locality_aware = true; + let plan = build_physical_indexing_plan( + &sources, + &indexer_infos, + locality_aware, + None, + &shard_locations, + ); + + assert_every_shard_scheduled_once(&plan, &sources); + assert_pipelines_within_limits(&plan, &sources); + assert_load_is_balanced(&plan, &sources); + assert_eq!(num_idle_indexers(&plan, &indexer_specs), 0); + + let metrics = get_shard_locality_metrics(&plan, &shard_locations, &indexer_infos); + println!( + "{num_shards} shards: {} local, {} nearby, {} remote", + metrics.num_local_shards, metrics.num_nearby_shards, metrics.num_remote_shards + ); + assert_eq!( + metrics.num_local_shards + metrics.num_nearby_shards + metrics.num_remote_shards, + num_shards + ); + assert!( + metrics.num_nearby_shards > 0, + "no shard overflowed to a same-az peer, so pass two never engaged" + ); + assert_eq!( + metrics.num_remote_shards, 0, + "{} shards crossed an availability zone despite every zone having room", + metrics.num_remote_shards + ); + + let replanned = build_physical_indexing_plan( + &sources, + &indexer_infos, + locality_aware, + Some(&plan), + &shard_locations, + ); + assert_eq!(plan, replanned); +} + +#[test] +fn test_scale_drain_spread_across_azs() { + let mut rng = StdRng::seed_from_u64(SEED); + let indexer_specs = build_indexer_specs(); + let sources = build_sources(&mut rng); + let shard_locations = build_shard_locations(&sources, &indexer_specs, &mut rng); + let host_per_shard = build_host_per_shard(&sources, &shard_locations); + let num_shards = total_num_shards(&sources); + + let no_draining_indexers = FnvHashSet::default(); + let ready_indexer_infos = build_indexer_infos(&indexer_specs, &no_draining_indexers); + let locality_aware = true; + let baseline_plan = build_physical_indexing_plan( + &sources, + &ready_indexer_infos, + locality_aware, + None, + &shard_locations, + ); + + let draining_indexer_ords = spread_draining_indexer_ords(); + assert_eq!(draining_indexer_ords.len(), NUM_DRAINING_INDEXERS); + let draining_indexer_infos = build_indexer_infos(&indexer_specs, &draining_indexer_ords); + let drained_plan = build_physical_indexing_plan( + &sources, + &draining_indexer_infos, + locality_aware, + Some(&baseline_plan), + &shard_locations, + ); + + assert_every_shard_scheduled_once(&drained_plan, &sources); + assert_pipelines_within_limits(&drained_plan, &sources); + assert_load_is_balanced(&drained_plan, &sources); + assert_draining_indexers_index_only_own_shards( + &drained_plan, + &host_per_shard, + &draining_indexer_infos, + ); + + let num_hosted_on_draining = + num_shards_hosted_on_draining_indexers(&host_per_shard, &draining_indexer_infos); + let num_indexed_by_draining = + num_shards_indexed_by_draining_indexers(&drained_plan, &draining_indexer_infos); + let metrics = get_shard_locality_metrics(&drained_plan, &shard_locations, &ready_indexer_infos); + println!( + "{num_shards} shards: draining indexers host {num_hosted_on_draining} and index \ + {num_indexed_by_draining}; {} local, {} nearby, {} remote", + metrics.num_local_shards, metrics.num_nearby_shards, metrics.num_remote_shards + ); + assert!( + num_indexed_by_draining * 100 >= num_hosted_on_draining * 85, + "draining indexers index {num_indexed_by_draining} of the {num_hosted_on_draining} shards \ + they host" + ); + + let replanned = build_physical_indexing_plan( + &sources, + &draining_indexer_infos, + locality_aware, + Some(&drained_plan), + &shard_locations, + ); + let num_released_shards = num_hosted_on_draining - num_indexed_by_draining; + let num_churned_shards = num_shards_with_changed_pipeline(&drained_plan, &replanned); + println!( + "{num_churned_shards} of {num_shards} shards changed pipeline on replan, out of \ + {num_released_shards} released by the draining indexers" + ); + assert!( + num_churned_shards <= num_released_shards, + "{num_churned_shards} shards changed pipeline but only {num_released_shards} were released" + ); + assert_every_shard_scheduled_once(&replanned, &sources); + assert_locality_of_hosted_shards_is_stable( + &drained_plan, + &replanned, + &shard_locations, + &ready_indexer_infos, + ); + assert_draining_indexers_index_only_own_shards( + &replanned, + &host_per_shard, + &draining_indexer_infos, + ); +} + +#[test] +fn test_scale_drain_whole_az() { + let mut rng = StdRng::seed_from_u64(SEED); + let indexer_specs = build_indexer_specs(); + let sources = build_sources(&mut rng); + let shard_locations = build_shard_locations(&sources, &indexer_specs, &mut rng); + let host_per_shard = build_host_per_shard(&sources, &shard_locations); + let num_shards = total_num_shards(&sources); + + let no_draining_indexers = FnvHashSet::default(); + let ready_indexer_infos = build_indexer_infos(&indexer_specs, &no_draining_indexers); + let locality_aware = true; + let baseline_plan = build_physical_indexing_plan( + &sources, + &ready_indexer_infos, + locality_aware, + None, + &shard_locations, + ); + + let draining_indexer_ords = whole_az_draining_indexer_ords(); + assert_eq!(draining_indexer_ords.len(), NUM_DRAINING_INDEXERS); + let draining_indexer_infos = build_indexer_infos(&indexer_specs, &draining_indexer_ords); + let drained_plan = build_physical_indexing_plan( + &sources, + &draining_indexer_infos, + locality_aware, + Some(&baseline_plan), + &shard_locations, + ); + + assert_every_shard_scheduled_once(&drained_plan, &sources); + assert_pipelines_within_limits(&drained_plan, &sources); + assert_load_is_balanced(&drained_plan, &sources); + assert_draining_indexers_index_only_own_shards( + &drained_plan, + &host_per_shard, + &draining_indexer_infos, + ); + assert_drained_az_spills_across_zones( + &drained_plan, + &host_per_shard, + &draining_indexer_infos, + "az-0", + ); + + let num_hosted_on_draining = + num_shards_hosted_on_draining_indexers(&host_per_shard, &draining_indexer_infos); + let num_indexed_by_draining = + num_shards_indexed_by_draining_indexers(&drained_plan, &draining_indexer_infos); + let metrics = get_shard_locality_metrics(&drained_plan, &shard_locations, &ready_indexer_infos); + println!( + "{num_shards} shards: draining indexers host {num_hosted_on_draining} and index \ + {num_indexed_by_draining}; {} local, {} nearby, {} remote", + metrics.num_local_shards, metrics.num_nearby_shards, metrics.num_remote_shards + ); + assert!( + num_indexed_by_draining * 100 >= num_hosted_on_draining * 85, + "draining indexers index {num_indexed_by_draining} of the {num_hosted_on_draining} shards \ + they host" + ); + assert!( + metrics.num_remote_shards > 0, + "the fully drained az has no eligible peer, so its overflow must cross a zone" + ); + + let replanned = build_physical_indexing_plan( + &sources, + &draining_indexer_infos, + locality_aware, + Some(&drained_plan), + &shard_locations, + ); + let num_released_shards = num_hosted_on_draining - num_indexed_by_draining; + let num_churned_shards = num_shards_with_changed_pipeline(&drained_plan, &replanned); + println!( + "{num_churned_shards} of {num_shards} shards changed pipeline on replan, out of \ + {num_released_shards} released by the draining indexers" + ); + assert!( + num_churned_shards <= num_released_shards, + "{num_churned_shards} shards changed pipeline but only {num_released_shards} were released" + ); + assert_every_shard_scheduled_once(&replanned, &sources); + assert_locality_of_hosted_shards_is_stable( + &drained_plan, + &replanned, + &shard_locations, + &ready_indexer_infos, + ); + assert_draining_indexers_index_only_own_shards( + &replanned, + &host_per_shard, + &draining_indexer_infos, + ); +} + +fn build_indexer_specs() -> Vec { + (0..NUM_INDEXERS) + .map(|indexer_ord| { + let node_id = format!("indexer-{indexer_ord}"); + let availability_zone = format!("az-{}", indexer_ord % NUM_AZS); + IndexerSpec::new(&node_id, INDEXER_CPU_CAPACITY, Some(&availability_zone)) + }) + .collect() +} + +fn build_sources(rng: &mut StdRng) -> Vec { + let index_uid = IndexUid::for_test("scale-test-index", 0); + let mut next_shard_id = 0u64; + let mut sources = Vec::with_capacity(NUM_SOURCES); + for source_size_tier in &SOURCE_SIZE_TIERS { + for _ in 0..source_size_tier.num_sources { + let num_shards = 1 + rng.random_range(0..source_size_tier.max_shards_per_source); + let load_per_shard = LOADS_PER_SHARD[rng.random_range(0..LOADS_PER_SHARD.len())]; + let shard_ids: Vec = (0..num_shards) + .map(|_| { + let shard_id = ShardId::from(next_shard_id); + next_shard_id += 1; + shard_id + }) + .collect(); + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: format!("source-{}", sources.len()), + }; + sources.push(SourceToSchedule { + source_uid, + source_type: SourceToScheduleType::Sharded { + shard_ids, + load_per_shard: NonZeroU32::new(load_per_shard).unwrap(), + }, + params_fingerprint: 0, + }); + } + } + assert_eq!(sources.len(), NUM_SOURCES); + assert!(total_num_shards(&sources) <= MAX_TOTAL_SHARDS); + sources +} + +fn build_shard_locations<'a>( + sources: &'a [SourceToSchedule], + indexer_specs: &'a [IndexerSpec], + rng: &mut StdRng, +) -> ShardLocations<'a> { + let mut shard_locations = ShardLocations::default(); + for source in sources { + for shard_id in shard_ids_of_source(source) { + let host_ord = rng.random_range(0..indexer_specs.len()); + shard_locations.add_location(shard_id, &indexer_specs[host_ord].node_id); + } + } + shard_locations +} + +fn build_indexer_infos( + indexer_specs: &[IndexerSpec], + draining_indexer_ords: &FnvHashSet, +) -> FnvHashMap { + let mut indexer_infos = FnvHashMap::default(); + for (indexer_ord, indexer_spec) in indexer_specs.iter().enumerate() { + let eligibility = if draining_indexer_ords.contains(&indexer_ord) { + Eligibility::SelfHostedOnly + } else { + Eligibility::Any + }; + let indexer_info = indexer_spec.to_indexer_info(eligibility); + indexer_infos.insert(indexer_spec.node_id.to_string(), indexer_info); + } + indexer_infos +} + +fn spread_draining_indexer_ords() -> FnvHashSet { + (0..NUM_INDEXERS).step_by(2).collect() +} + +fn whole_az_draining_indexer_ords() -> FnvHashSet { + let mut draining_indexer_ords: FnvHashSet = (0..NUM_INDEXERS) + .filter(|indexer_ord| indexer_ord % NUM_AZS == 0) + .collect(); + let num_remaining_to_drain = NUM_DRAINING_INDEXERS - draining_indexer_ords.len(); + let next_az_indexer_ords = (0..NUM_INDEXERS) + .filter(|indexer_ord| indexer_ord % NUM_AZS == 1) + .take(num_remaining_to_drain); + draining_indexer_ords.extend(next_az_indexer_ords); + draining_indexer_ords +} + +fn assert_every_shard_scheduled_once(plan: &PhysicalIndexingPlan, sources: &[SourceToSchedule]) { + let mut scheduled_shard_ids: Vec = Vec::new(); + for tasks in plan.indexing_tasks_per_indexer().values() { + for task in tasks { + scheduled_shard_ids.extend(task.shard_ids.iter().cloned()); + } + } + let unique_scheduled_shard_ids: HashSet<&ShardId> = scheduled_shard_ids.iter().collect(); + let expected_shard_ids: HashSet<&ShardId> = + sources.iter().flat_map(shard_ids_of_source).collect(); + assert_eq!( + scheduled_shard_ids.len(), + expected_shard_ids.len(), + "scheduled {} shard slots for {} shards", + scheduled_shard_ids.len(), + expected_shard_ids.len() + ); + assert_eq!(unique_scheduled_shard_ids, expected_shard_ids); +} + +fn assert_pipelines_within_limits(plan: &PhysicalIndexingPlan, sources: &[SourceToSchedule]) { + let mut max_shards_per_source: FnvHashMap<&str, usize> = FnvHashMap::default(); + for source in sources { + let max_num_shards = compute_max_num_shards_per_pipeline(&source.source_type); + max_shards_per_source.insert(&source.source_uid.source_id, max_num_shards.get() as usize); + } + for tasks in plan.indexing_tasks_per_indexer().values() { + for task in tasks { + let max_num_shards = max_shards_per_source[task.source_id.as_str()]; + assert!( + !task.shard_ids.is_empty(), + "pipeline for {} holds no shards at all", + task.source_id + ); + assert!( + task.shard_ids.len() <= max_num_shards, + "pipeline for {} holds {} shards, limit is {max_num_shards}", + task.source_id, + task.shard_ids.len() + ); + } + } +} + +fn load_per_indexer(plan: &PhysicalIndexingPlan, sources: &[SourceToSchedule]) -> Vec { + let mut load_per_source: FnvHashMap<&str, u64> = FnvHashMap::default(); + for source in sources { + let SourceToScheduleType::Sharded { load_per_shard, .. } = &source.source_type else { + continue; + }; + load_per_source.insert(&source.source_uid.source_id, load_per_shard.get() as u64); + } + plan.indexing_tasks_per_indexer() + .values() + .map(|tasks| { + tasks + .iter() + .map(|task| { + let load_per_shard = load_per_source[task.source_id.as_str()]; + load_per_shard * task.shard_ids.len() as u64 + }) + .sum() + }) + .collect() +} + +fn num_idle_indexers(plan: &PhysicalIndexingPlan, indexer_specs: &[IndexerSpec]) -> usize { + indexer_specs + .iter() + .filter(|indexer_spec| { + let node_id = indexer_spec.node_id.as_str(); + shard_ids_for_indexer(plan, node_id).is_empty() + }) + .count() +} + +fn assert_load_is_balanced(plan: &PhysicalIndexingPlan, sources: &[SourceToSchedule]) { + let loads = load_per_indexer(plan, sources); + let total_load: u64 = loads.iter().sum(); + let mean_load = total_load / loads.len() as u64; + let max_load = *loads.iter().max().unwrap(); + assert!( + max_load * 2 <= mean_load * 3, + "most loaded indexer holds {max_load} mcpu against a mean of {mean_load} mcpu" + ); +} + +fn build_host_per_shard<'a>( + sources: &'a [SourceToSchedule], + shard_locations: &ShardLocations, +) -> FnvHashMap<&'a ShardId, String> { + let mut host_per_shard = FnvHashMap::default(); + for source in sources { + for shard_id in shard_ids_of_source(source) { + let Some(host) = shard_locations.get_shard_locations(shard_id).first() else { + continue; + }; + host_per_shard.insert(shard_id, host.to_string()); + } + } + host_per_shard +} + +fn pipeline_per_shard( + plan: &PhysicalIndexingPlan, +) -> FnvHashMap<&ShardId, (&String, Option)> { + let mut pipeline_per_shard = FnvHashMap::default(); + for (indexer, tasks) in plan.indexing_tasks_per_indexer() { + for task in tasks { + for shard_id in &task.shard_ids { + pipeline_per_shard.insert(shard_id, (indexer, task.pipeline_uid)); + } + } + } + pipeline_per_shard +} + +fn num_shards_with_changed_pipeline( + plan: &PhysicalIndexingPlan, + replanned: &PhysicalIndexingPlan, +) -> usize { + let pipeline_per_shard_before = pipeline_per_shard(plan); + let pipeline_per_shard_after = pipeline_per_shard(replanned); + pipeline_per_shard_before + .iter() + .filter(|(shard_id, pipeline_before)| { + pipeline_per_shard_after[**shard_id] != **pipeline_before + }) + .count() +} + +fn assert_locality_of_hosted_shards_is_stable( + plan: &PhysicalIndexingPlan, + replanned: &PhysicalIndexingPlan, + shard_locations: &ShardLocations, + indexer_infos: &FnvHashMap, +) { + let metrics_before = get_shard_locality_metrics(plan, shard_locations, indexer_infos); + let metrics_after = get_shard_locality_metrics(replanned, shard_locations, indexer_infos); + assert_eq!( + metrics_before.num_local_shards, metrics_after.num_local_shards, + "{} shards were indexed on their host before the replan and {} after", + metrics_before.num_local_shards, metrics_after.num_local_shards + ); + let num_displaced_before = + metrics_before.num_nearby_shards + metrics_before.num_remote_shards; + let num_displaced_after = metrics_after.num_nearby_shards + metrics_after.num_remote_shards; + assert_eq!(num_displaced_before, num_displaced_after); +} + +fn is_draining(indexer: &str, indexer_infos: &FnvHashMap) -> bool { + indexer_infos[indexer].eligibility == Eligibility::SelfHostedOnly +} + +fn assert_draining_indexers_index_only_own_shards( + plan: &PhysicalIndexingPlan, + host_per_shard: &FnvHashMap<&ShardId, String>, + indexer_infos: &FnvHashMap, +) { + for (indexer, tasks) in plan.indexing_tasks_per_indexer() { + if !is_draining(indexer, indexer_infos) { + continue; + } + for task in tasks { + for shard_id in &task.shard_ids { + let host = &host_per_shard[shard_id]; + assert_eq!( + host, indexer, + "draining indexer {indexer} indexes shard {shard_id:?} hosted on {host}" + ); + } + } + } +} + +fn assert_drained_az_spills_across_zones( + plan: &PhysicalIndexingPlan, + host_per_shard: &FnvHashMap<&ShardId, String>, + indexer_infos: &FnvHashMap, + drained_az: &str, +) { + for (indexer, tasks) in plan.indexing_tasks_per_indexer() { + for task in tasks { + for shard_id in &task.shard_ids { + let host = &host_per_shard[shard_id]; + let host_az = indexer_infos[host.as_str()].availability_zone.as_deref(); + if host_az != Some(drained_az) || host == indexer { + continue; + } + let indexer_az = indexer_infos[indexer].availability_zone.as_deref(); + assert_ne!( + indexer_az, + Some(drained_az), + "shard {shard_id:?} hosted on {host} moved to {indexer}, still inside the \ + fully drained {drained_az}" + ); + } + } + } +} + +fn num_shards_hosted_on_draining_indexers( + host_per_shard: &FnvHashMap<&ShardId, String>, + indexer_infos: &FnvHashMap, +) -> usize { + host_per_shard + .values() + .filter(|host| is_draining(host, indexer_infos)) + .count() +} + +fn num_shards_indexed_by_draining_indexers( + plan: &PhysicalIndexingPlan, + indexer_infos: &FnvHashMap, +) -> usize { + plan.indexing_tasks_per_indexer() + .iter() + .filter(|(indexer, _)| is_draining(indexer, indexer_infos)) + .map(|(_, tasks)| tasks.iter().map(|task| task.shard_ids.len()).sum::()) + .sum() +} diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs index 8f6ed37eab6..a384608652b 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs @@ -50,52 +50,40 @@ pub fn solve( check_contract_conditions(&problem, &previous_solution); let base_problem = problem; - // Due to the inherent nature of bin-packing, it is possible that the first inflation - // is not sufficient to solve the problem. - // - // In that case, we inflate the capacity iteratively until we find a solution. - let mut best_solution: Option = None; - let mut best_attempt: u32 = 0; - for attempt in 0..MAX_INFLATION_ATTEMPTS { - let scaled_problem = problem_at_inflation_level(&base_problem, attempt); - if let Ok(solution) = attempt_solve(&scaled_problem, previous_solution.clone()) { - best_solution = Some(solution); - best_attempt = attempt; - break; - } - } - let mut best_solution = - best_solution.expect("failed to assign all of the sources (logical bug)"); - - // Just stopping here would not offer any stability guarantee. - // - // We descend: we re-feed the candidate solution to the algorithm at lower - // inflation levels to find the true minimal feasible level for *this* solution. - // This is what guarantees stability: the returned solution succeeds at - // `best_attempt` but fails at `best_attempt - 1`. On the next call to `solve` - // starting from this solution, the ascending search will fail at every level - // below `best_attempt` (less capacity than a level that already failed) and - // succeed at `best_attempt`, where the pipeline is a no-op. Hence `solve` is - // idempotent. - while let Some(lower_attempt) = best_attempt.checked_sub(1) { - let scaled_problem = problem_at_inflation_level(&base_problem, lower_attempt); - match attempt_solve(&scaled_problem, best_solution.clone()) { - Ok(solution) => { - best_solution = solution; - best_attempt = lower_attempt; - } - Err(NotEnoughCapacity) => break, - } - } - - if best_attempt > 0 { + let inflation_attempt = minimal_feasible_inflation_attempt(&base_problem); + if inflation_attempt > 0 { // the higher the attempt number, the more unbalanced the solution tracing::warn!( - attempt_number = best_attempt, + attempt_number = inflation_attempt, "capacity re-scaled, scheduling solution likely unbalanced" ); } - best_solution + let scaled_problem = problem_at_inflation_level(&base_problem, inflation_attempt); + if let Ok(solution) = attempt_solve(&scaled_problem, previous_solution) { + return solution; + } + let empty_solution = scaled_problem.new_solution(); + attempt_solve(&scaled_problem, empty_solution) + .expect("failed to assign all of the sources (logical bug)") +} + +/// Smallest inflation attempt at which the placement succeeds from an empty solution. +/// +/// The attempt is deliberately derived from the problem alone, never from the previous solution. +/// A search seeded with the previous solution is not reproducible: shedding whole sources off +/// over-capacity indexers is what makes a tight problem placeable, so raising the capacities can +/// stop that shedding and turn a feasible attempt infeasible. Feasibility is therefore not +/// monotonic in the attempt number, and any rule of the form "stop at the first attempt that +/// works" would pick a different attempt depending on where it started. +fn minimal_feasible_inflation_attempt(base_problem: &SchedulingProblem) -> u32 { + for attempt in 0..MAX_INFLATION_ATTEMPTS { + let scaled_problem = problem_at_inflation_level(base_problem, attempt); + let empty_solution = scaled_problem.new_solution(); + if attempt_solve(&scaled_problem, empty_solution).is_ok() { + return attempt; + } + } + panic!("failed to assign all of the sources (logical bug)") } /// Returns a clone of `base_problem` with its node capacities scaled by `1.2^inflation_attempt`. @@ -135,9 +123,10 @@ fn attempt_solve( // First, we remove remote shards from indexers that are only eligible // to index their own shards. strip_self_hosted_only_indexers(problem, &mut solution); - // Then, we place shards that are hosted on the indexers. + // Then, we place shards that are hosted on the indexers, reclaiming a draining + // indexer's own shards from its peers so that it gets priority to them. let leftover_shards_per_source = place_self_hosted_shards(problem, &mut solution); - // After that, we place remaining shards on remote indexers, but in the same locality. + // Next we place remaining shards on remote indexers, but in the same locality. place_nearby_shards(problem, &leftover_shards_per_source, &mut solution); } else { // If locality awareness is disabled, we directly assign sources to indexers that have some @@ -315,6 +304,44 @@ fn strip_self_hosted_only_indexers( } } +fn num_foreign_shards_on_indexer( + source: &Source, + indexer_ord: IndexerOrd, + problem: &SchedulingProblem, + solution: &SchedulingSolution, +) -> u32 { + let num_assigned_shards = + solution.indexer_assignments[indexer_ord].num_shards(source.source_ord); + let num_self_hosted_shards = problem.source_affinity(source.source_ord, indexer_ord); + num_assigned_shards.saturating_sub(num_self_hosted_shards) +} + +fn release_foreign_shards_from_peers( + source: &Source, + num_shards_to_release: u32, + problem: &SchedulingProblem, + solution: &mut SchedulingSolution, +) -> u32 { + let mut peer_ords: Vec = (0..problem.num_indexers()) + .filter(|&peer_ord| problem.is_eligible_for_foreign_shards(peer_ord)) + .collect(); + peer_ords.sort_by_key(|&peer_ord| { + solution.indexer_assignments[peer_ord].indexer_available_capacity(problem) + }); + let mut num_shards_remaining = num_shards_to_release; + for peer_ord in peer_ords { + if num_shards_remaining == 0 { + break; + } + let num_foreign_shards = num_foreign_shards_on_indexer(source, peer_ord, problem, solution); + let num_shards_released = num_foreign_shards.min(num_shards_remaining); + solution.indexer_assignments[peer_ord] + .remove_shards(source.source_ord, num_shards_released); + num_shards_remaining -= num_shards_released; + } + num_shards_to_release - num_shards_remaining +} + // ---------------------------------------------------- // Phase 3 // Place unassigned sources. @@ -427,6 +454,24 @@ fn place_self_hosted_shards_on_indexer( num_shards_placed } +fn reclaim_self_hosted_shards_from_peers( + source: &Source, + indexer_ord: IndexerOrd, + num_self_hosted_shards: u32, + num_unaccounted_shards: u32, + problem: &SchedulingProblem, + solution: &mut SchedulingSolution, +) -> u32 { + let available_capacity = available_cpu_capacity(indexer_ord, problem, solution); + let num_placable_shards = available_capacity.cpu_millis() / source.load_per_shard; + let num_shards_wanted = num_self_hosted_shards.min(num_placable_shards); + let num_shards_short = num_shards_wanted.saturating_sub(num_unaccounted_shards); + if num_shards_short == 0 { + return 0; + } + release_foreign_shards_from_peers(source, num_shards_short, problem, solution) +} + fn place_self_hosted_shards( problem: &SchedulingProblem, solution: &mut SchedulingSolution, @@ -440,10 +485,29 @@ fn place_self_hosted_shards( let source_ord = unassigned_source.source_ord as usize; let leftover_shards = &mut leftover_shards_per_source[source_ord]; let mut num_unaccounted_shards = unassigned_source.num_shards; - for (&indexer_ord, &num_self_hosted_shards) in &unassigned_source.affinities { + // A draining indexer claims its own shards before any other host of the same source. + let (draining_indexer_ords, ready_indexer_ords): (Vec, Vec) = + unassigned_source + .affinities + .keys() + .copied() + .partition(|&indexer_ord| !problem.is_eligible_for_foreign_shards(indexer_ord)); + for indexer_ord in draining_indexer_ords.into_iter().chain(ready_indexer_ords) { + let num_self_hosted_shards = unassigned_source.affinities[&indexer_ord]; + if !problem.is_eligible_for_foreign_shards(indexer_ord) { + let num_shards_reclaimed = reclaim_self_hosted_shards_from_peers( + unassigned_source, + indexer_ord, + num_self_hosted_shards, + num_unaccounted_shards, + problem, + solution, + ); + num_unaccounted_shards += num_shards_reclaimed; + } let num_shards_to_place = num_self_hosted_shards.min(num_unaccounted_shards); if num_shards_to_place == 0 { - break; + continue; } let num_shards_placed = place_self_hosted_shards_on_indexer( unassigned_source, @@ -824,6 +888,71 @@ mod tests { assert_eq!(solution.indexer_assignments[1].num_shards(1), 1); } + fn locality_in_az(group_ord: usize, eligibility: Eligibility) -> IndexerLocality { + IndexerLocality { + group: Some(LocalityGroup::from_ord(group_ord)), + eligibility, + } + } + + #[test] + fn test_draining_indexer_reclaims_own_shards_from_peers() { + let draining_locality = locality_in_az(0, Eligibility::SelfHostedOnly); + let same_az_peer_locality = locality_in_az(0, Eligibility::Any); + let other_az_peer_locality = locality_in_az(1, Eligibility::Any); + { + let mut problem = SchedulingProblem::with_indexer_localities( + vec![mcpu(4_000), mcpu(4_000), mcpu(4_000)], + vec![ + draining_locality, + same_az_peer_locality, + other_az_peer_locality, + ], + ); + problem.add_source(4, NonZeroU32::new(1_000).unwrap()); + problem.inc_affinity(0, 0); + problem.inc_affinity(0, 0); + problem.inc_affinity(0, 2); + problem.inc_affinity(0, 2); + + let mut previous_solution = problem.new_solution(); + previous_solution.indexer_assignments[1].add_shards(0, 3); + + let solution = attempt_solve(&problem, previous_solution).unwrap(); + + assert_eq!(solution.indexer_assignments[0].num_shards(0), 2); + assert_eq!(solution.indexer_assignments[1].num_shards(0), 2); + assert_eq!(solution.indexer_assignments[2].num_shards(0), 0); + + let settled = attempt_solve(&problem, solution.clone()).unwrap(); + assert_eq!(settled.indexer_assignments, solution.indexer_assignments); + } + { + let mut problem = SchedulingProblem::with_indexer_localities( + vec![mcpu(1_000), mcpu(4_000), mcpu(4_000)], + vec![ + draining_locality, + same_az_peer_locality, + other_az_peer_locality, + ], + ); + problem.add_source(4, NonZeroU32::new(1_000).unwrap()); + problem.inc_affinity(0, 0); + problem.inc_affinity(0, 0); + problem.inc_affinity(0, 2); + problem.inc_affinity(0, 2); + + let mut previous_solution = problem.new_solution(); + previous_solution.indexer_assignments[1].add_shards(0, 3); + + let solution = attempt_solve(&problem, previous_solution).unwrap(); + + assert_eq!(solution.indexer_assignments[0].num_shards(0), 1); + assert_eq!(solution.indexer_assignments[1].num_shards(0), 3); + assert_eq!(solution.indexer_assignments[2].num_shards(0), 0); + } + } + #[test] fn test_compute_unassigned_shards_simple() { let mut problem = SchedulingProblem::with_indexer_cpu_capacities(vec![mcpu(4_000)]); @@ -1033,10 +1162,16 @@ mod tests { }) } - fn locality_group_strat(num_groups: usize) -> impl Strategy> { + fn locality_groups_strat( + num_indexers: usize, + num_groups: usize, + ) -> impl Strategy>> { + let group_strat = (0..num_groups).prop_map(|group_ord| Some(LocalityGroup::from_ord(group_ord))); + let every_indexer_in_a_zone = proptest::collection::vec(group_strat, num_indexers); + let no_indexer_in_a_zone = Just(vec![None; num_indexers]); prop_oneof![ - 1 => Just(None), - 4 => (0..num_groups).prop_map(|group_ord| Some(LocalityGroup::from_ord(group_ord))), + 4 => every_indexer_in_a_zone, + 1 => no_indexer_in_a_zone, ] } @@ -1080,7 +1215,7 @@ mod tests { ) -> impl Strategy { let cpu_capacities_strat = proptest::collection::vec(indexer_cpu_capacity_strat(), num_indexers); - let groups_strat = proptest::collection::vec(locality_group_strat(num_groups), num_indexers); + let groups_strat = locality_groups_strat(num_indexers, num_groups); let eligibilities_strat = proptest::collection::vec(eligibility_strat(), num_indexers); let sources_strat = proptest::collection::vec(locality_source_strat(num_indexers), num_sources); @@ -1132,7 +1267,8 @@ mod tests { let solution_2 = solve(problem.clone(), solution_1.clone()); assert_eq!( solution_1.indexer_assignments, solution_2.indexer_assignments, - "solution unstable!\nSolution 1: {solution_1:?}\nSolution 2: {solution_2:?}" + "solution unstable!\nProblem: {problem:?}\nSolution 1: {solution_1:?}\nSolution \ + 2: {solution_2:?}" ); for indexer_assignment in &solution_1.indexer_assignments { let indexer_ord = indexer_assignment.indexer_ord; @@ -1231,6 +1367,60 @@ mod tests { ); } + #[test] + fn test_reproduce_non_monotonic_inflation() { + let localities = vec![ + locality_in_az(0, Eligibility::Any), + locality_in_az(0, Eligibility::SelfHostedOnly), + locality_in_az(0, Eligibility::Any), + locality_in_az(0, Eligibility::SelfHostedOnly), + ]; + let mut problem = SchedulingProblem::with_indexer_localities( + vec![mcpu(4_237), mcpu(4_146), mcpu(1_953), mcpu(1_964)], + localities, + ); + problem.add_source(2, NonZeroU32::new(3_200).unwrap()); + problem.add_source(3, NonZeroU32::new(250).unwrap()); + problem.add_source(10, NonZeroU32::new(1_000).unwrap()); + let affinities = [ + (0u32, 1usize, 1), + (0, 2, 1), + (1, 0, 2), + (1, 3, 1), + (2, 0, 3), + (2, 1, 1), + (2, 2, 3), + (2, 3, 1), + ]; + for (source_ord, indexer_ord, num_shards) in affinities { + for _ in 0..num_shards { + problem.inc_affinity(source_ord, indexer_ord); + } + } + + let mut previous_solution = problem.new_solution(); + let seeded_assignments = [ + (0usize, 0u32, 1), + (0, 1, 2), + (0, 2, 6), + (1, 2, 1), + (2, 0, 1), + (2, 2, 2), + (3, 1, 1), + (3, 2, 1), + ]; + for (indexer_ord, source_ord, num_shards) in seeded_assignments { + previous_solution.indexer_assignments[indexer_ord].add_shards(source_ord, num_shards); + } + + let solution_1 = solve(problem.clone(), previous_solution); + let solution_2 = solve(problem, solution_1.clone()); + assert_eq!( + solution_1.indexer_assignments, solution_2.indexer_assignments, + "solution unstable!\nSolution 1: {solution_1:?}\nSolution 2: {solution_2:?}" + ); + } + #[test] fn test_capacity_scaling_iteration_required() { // Sources 0 and 1 each need 2500 mcpu; source 2 needs 1500 mcpu. From 43f7a4624f2893725b0a749d0908a631fc15eddd Mon Sep 17 00:00:00 2001 From: Nadav Gov-Ari Date: Fri, 14 Aug 2026 17:38:48 -0400 Subject: [PATCH 4/4] With all tests and the locality threshold rebuild --- .../src/indexing_scheduler/mod.rs | 167 +++++++++++++++++- .../scheduling/scheduling_logic.rs | 10 +- .../quickwit-control-plane/src/metrics.rs | 11 ++ 3 files changed, 176 insertions(+), 12 deletions(-) diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index 30e3669ff78..8d96a95ef91 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -52,6 +52,11 @@ const DEFAULT_ENABLE_VARIABLE_SHARD_LOAD: bool = false; const DEFAULT_ENABLE_AZ_AWARE_SCHEDULING: bool = false; +const DEFAULT_MIN_SHARD_LOCALITY_PERCENT: u32 = 30; + +/// Minimum period before being able to rebuild the plan from scratch. +const PLAN_FROM_SCRATCH_COOLDOWN_PERIOD: Duration = Duration::from_mins(30); + pub(crate) const MIN_DURATION_BETWEEN_SCHEDULING: Duration = if cfg!(any(test, feature = "testsuite")) { Duration::from_millis(50) @@ -74,6 +79,8 @@ pub struct IndexingSchedulerState { pub last_applied_indexer_statuses: FnvHashMap, #[serde(skip)] pub last_applied_plan_timestamp: Option, + #[serde(skip)] + pub next_plan_from_scratch_timestamp: Option, } /// The [`IndexingScheduler`] is responsible for listing indexing tasks and assigning them to @@ -417,15 +424,23 @@ impl IndexingScheduler { }; let shard_locations = model.shard_locations(); - let new_physical_plan = build_physical_indexing_plan( + let plan_from_previous = build_physical_indexing_plan( &sources, &indexer_infos, locality_aware, self.state.last_applied_physical_plan.as_ref(), &shard_locations, ); - let shard_locality_metrics = - get_shard_locality_metrics(&new_physical_plan, &shard_locations, &indexer_infos); + let locality_metrics_from_previous = + get_shard_locality_metrics(&plan_from_previous, &shard_locations, &indexer_infos); + let (new_physical_plan, shard_locality_metrics) = self.maybe_build_plan_from_scratch( + plan_from_previous, + locality_metrics_from_previous, + &sources, + &indexer_infos, + locality_aware, + &shard_locations, + ); shard_locality_metrics.publish(); let indexer_statuses = build_indexer_statuses(&indexers); @@ -446,6 +461,47 @@ impl IndexingScheduler { self.state.num_schedule_indexing_plan += 1; } + /// A plan built from the previous one can lose locality over time but never regain it. Below + /// the threshold we try one built from scratch (equivalent to restarting the control plane). + fn maybe_build_plan_from_scratch( + &mut self, + plan: PhysicalIndexingPlan, + locality_metrics: ShardLocalityMetrics, + sources: &[SourceToSchedule], + indexer_infos: &FnvHashMap, + locality_aware: bool, + shard_locations: &ShardLocations, + ) -> (PhysicalIndexingPlan, ShardLocalityMetrics) { + if locality_metrics.locality_percent() >= min_shard_locality_percent() { + return (plan, locality_metrics); + } + let now = Instant::now(); + if let Some(next_plan_from_scratch_timestamp) = self.state.next_plan_from_scratch_timestamp + && now < next_plan_from_scratch_timestamp + { + return (plan, locality_metrics); + } + let plan_from_scratch = build_physical_indexing_plan( + sources, + indexer_infos, + locality_aware, + None, + shard_locations, + ); + let locality_metrics_from_scratch = + get_shard_locality_metrics(&plan_from_scratch, shard_locations, indexer_infos); + if locality_metrics_from_scratch.locality_percent() <= locality_metrics.locality_percent() { + return (plan, locality_metrics); + } + info!( + locality_percent = locality_metrics.locality_percent(), + locality_percent_from_scratch = locality_metrics_from_scratch.locality_percent(), + "rebuilding the indexing plan from scratch to restore shard locality" + ); + self.state.next_plan_from_scratch_timestamp = Some(now + PLAN_FROM_SCRATCH_COOLDOWN_PERIOD); + (plan_from_scratch, locality_metrics_from_scratch) + } + /// Checks if the last applied plan corresponds to the running indexing tasks present in the /// chitchat cluster state. If true, do nothing. /// - If node IDs differ, schedule a new indexing plan. @@ -641,6 +697,15 @@ impl IndexingPlansDiff<'_> { } } +fn min_shard_locality_percent() -> u32 { + quickwit_common::get_from_env_cached!( + u32, + "QW_MIN_SHARD_LOCALITY_PERCENT", + DEFAULT_MIN_SHARD_LOCALITY_PERCENT, + false + ) +} + fn get_shard_locality_metrics( physical_plan: &PhysicalIndexingPlan, shard_locations: &ShardLocations, @@ -911,7 +976,9 @@ mod tests { use quickwit_proto::types::{IndexUid, PipelineUid, ShardId, SourceUid}; use super::*; - use crate::indexing_scheduler::scheduling::build_physical_indexing_plan_without_locality; + use crate::indexing_scheduler::scheduling::{ + build_physical_indexing_plan_without_locality, shard_ids_for_indexer, + }; use crate::model::ShardLocations; #[test] fn test_indexing_plans_diff() { @@ -1145,6 +1212,98 @@ mod tests { } } + #[test] + fn test_maybe_build_plan_from_scratch() { + let indexer1 = NodeId::from_str("indexer1"); + let indexer2 = NodeId::from_str("indexer2"); + let shard1 = ShardId::from(1); + let shard2 = ShardId::from(2); + let source_uid = SourceUid { + index_uid: IndexUid::for_test("test-index", 0), + source_id: "test-source".to_string(), + }; + let sources = vec![SourceToSchedule { + source_uid: source_uid.clone(), + source_type: SourceToScheduleType::Sharded { + shard_ids: vec![shard1.clone(), shard2.clone()], + load_per_shard: NonZeroU32::new(1_000).unwrap(), + }, + params_fingerprint: 0, + }]; + let mut shard_locations = ShardLocations::default(); + shard_locations.add_location(&shard1, &indexer1); + shard_locations.add_location(&shard2, &indexer2); + + let mut indexer_infos = FnvHashMap::default(); + indexer_infos.insert(indexer1.to_string(), IndexerInfo::for_test(mcpu(4_000))); + indexer_infos.insert(indexer2.to_string(), IndexerInfo::for_test(mcpu(4_000))); + + // Each indexer indexes the shard the other one hosts, so nothing is local. + let swapped_plan = || { + let indexer_ids = vec![indexer1.to_string(), indexer2.to_string()]; + let mut plan = PhysicalIndexingPlan::with_indexer_ids(&indexer_ids); + for (indexer, shard_id) in [(&indexer1, &shard2), (&indexer2, &shard1)] { + plan.add_indexing_task( + indexer.as_str(), + IndexingTask { + index_uid: Some(source_uid.index_uid.clone()), + source_id: source_uid.source_id.clone(), + pipeline_uid: Some(PipelineUid::random()), + shard_ids: vec![shard_id.clone()], + params_fingerprint: 0, + }, + ); + } + plan + }; + + let mut scheduler = IndexingScheduler::new( + "test-cluster".to_string(), + NodeId::from_str("control-plane"), + IndexerPool::default(), + ); + let locality_aware = false; + + let plan = swapped_plan(); + let metrics = get_shard_locality_metrics(&plan, &shard_locations, &indexer_infos); + assert_eq!(metrics.locality_percent(), 0); + let (plan, metrics) = scheduler.maybe_build_plan_from_scratch( + plan, + metrics, + &sources, + &indexer_infos, + locality_aware, + &shard_locations, + ); + assert_eq!(metrics.locality_percent(), 100); + assert_eq!(shard_ids_for_indexer(&plan, indexer1.as_str()), vec![shard1.clone()]); + + let plan_in_cooldown = swapped_plan(); + let metrics_in_cooldown = + get_shard_locality_metrics(&plan_in_cooldown, &shard_locations, &indexer_infos); + let (_, metrics_in_cooldown) = scheduler.maybe_build_plan_from_scratch( + plan_in_cooldown, + metrics_in_cooldown, + &sources, + &indexer_infos, + locality_aware, + &shard_locations, + ); + assert_eq!(metrics_in_cooldown.locality_percent(), 0); + + scheduler.state.next_plan_from_scratch_timestamp = None; + let (_, metrics_above_threshold) = scheduler.maybe_build_plan_from_scratch( + plan, + metrics, + &sources, + &indexer_infos, + locality_aware, + &shard_locations, + ); + assert_eq!(metrics_above_threshold.locality_percent(), 100); + assert!(scheduler.state.next_plan_from_scratch_timestamp.is_none()); + } + #[test] fn test_get_sources_to_schedule() { let mut model = ControlPlaneModel::default(); diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs index a384608652b..ef21c57f594 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs @@ -67,14 +67,8 @@ pub fn solve( .expect("failed to assign all of the sources (logical bug)") } -/// Smallest inflation attempt at which the placement succeeds from an empty solution. -/// -/// The attempt is deliberately derived from the problem alone, never from the previous solution. -/// A search seeded with the previous solution is not reproducible: shedding whole sources off -/// over-capacity indexers is what makes a tight problem placeable, so raising the capacities can -/// stop that shedding and turn a feasible attempt infeasible. Feasibility is therefore not -/// monotonic in the attempt number, and any rule of the form "stop at the first attempt that -/// works" would pick a different attempt depending on where it started. +/// Don't derive this from the previous plan. More capacity means fewer sources get evicted off +/// overloaded indexers, and that eviction is sometimes the only reason a plan fits at all. fn minimal_feasible_inflation_attempt(base_problem: &SchedulingProblem) -> u32 { for attempt in 0..MAX_INFLATION_ATTEMPTS { let scaled_problem = problem_at_inflation_level(base_problem, attempt); diff --git a/quickwit/quickwit-control-plane/src/metrics.rs b/quickwit/quickwit-control-plane/src/metrics.rs index aef28daa98d..c26aea3ad79 100644 --- a/quickwit/quickwit-control-plane/src/metrics.rs +++ b/quickwit/quickwit-control-plane/src/metrics.rs @@ -25,6 +25,17 @@ pub struct ShardLocalityMetrics { } impl ShardLocalityMetrics { + /// Share of shards indexed without crossing an availability zone, as a percentage. + pub fn locality_percent(self) -> u32 { + let num_shards = + self.num_local_shards + self.num_nearby_shards + self.num_remote_shards; + if num_shards == 0 { + return 100; + } + let num_local_or_nearby_shards = self.num_local_shards + self.num_nearby_shards; + (num_local_or_nearby_shards * 100 / num_shards) as u32 + } + pub fn publish(self) { LOCAL_SHARDS.set(self.num_local_shards as f64); NEARBY_SHARDS.set(self.num_nearby_shards as f64);