diff --git a/quickwit/quickwit-cluster/src/cluster.rs b/quickwit/quickwit-cluster/src/cluster.rs index 0675bf5269f..daea8aed6bd 100644 --- a/quickwit/quickwit-cluster/src/cluster.rs +++ b/quickwit/quickwit-cluster/src/cluster.rs @@ -228,7 +228,7 @@ impl Cluster { ]; if let Some(az) = &self_node.availability_zone { - initial_key_values.push((AVAILABILITY_ZONE_KEY.to_string(), az.clone())); + initial_key_values.push((AVAILABILITY_ZONE_KEY.to_string(), az.to_string())); } initial_key_values.push(( STANDALONE_COMPACTORS_KEY.to_string(), diff --git a/quickwit/quickwit-cluster/src/member.rs b/quickwit/quickwit-cluster/src/member.rs index 9e7a75e5805..510b502ccfe 100644 --- a/quickwit/quickwit-cluster/src/member.rs +++ b/quickwit/quickwit-cluster/src/member.rs @@ -22,7 +22,7 @@ use chitchat::{ChitchatId, NodeState, Version}; use quickwit_common::shared_consts::INGESTER_STATUS_KEY; use quickwit_proto::indexing::{CpuCapacity, IndexingTask}; use quickwit_proto::ingest::ingester::IngesterStatus; -use quickwit_proto::types::NodeId; +use quickwit_proto::types::{AvailabilityZone, NodeId}; use tracing::{error, warn}; use crate::cluster::parse_indexing_tasks; @@ -53,7 +53,7 @@ pub(crate) trait NodeStateExt { fn ingester_status(&self) -> IngesterStatus; - fn availability_zone(&self) -> Option; + fn availability_zone(&self) -> Option; fn enable_standalone_compactors(&self) -> bool; } @@ -93,8 +93,8 @@ impl NodeStateExt for NodeState { .unwrap_or(IngesterStatus::Ready) } - fn availability_zone(&self) -> Option { - self.get(AVAILABILITY_ZONE_KEY).map(|az| az.to_string()) + fn availability_zone(&self) -> Option { + self.get(AVAILABILITY_ZONE_KEY).map(AvailabilityZone::from) } fn enable_standalone_compactors(&self) -> bool { @@ -132,7 +132,7 @@ pub struct ClusterMember { /// Whether the node is ready to serve requests. pub is_ready: bool, /// Availability zone the node is running in, if enabled. - pub availability_zone: Option, + pub availability_zone: Option, /// Whether the node was started with standalone compactors enabled. pub enable_standalone_compactors: bool, } diff --git a/quickwit/quickwit-cluster/src/node.rs b/quickwit/quickwit-cluster/src/node.rs index 27ca6271b72..5189e91776d 100644 --- a/quickwit/quickwit-cluster/src/node.rs +++ b/quickwit/quickwit-cluster/src/node.rs @@ -21,6 +21,7 @@ use quickwit_config::service::QuickwitService; use quickwit_proto::indexing::IndexingTask; #[cfg(any(test, feature = "testsuite"))] use quickwit_proto::ingest::ingester::IngesterStatus; +use quickwit_proto::types::AvailabilityZone; use tonic::transport::Channel; use crate::member::{ClusterMember, build_cluster_member}; @@ -106,8 +107,8 @@ impl ClusterNode { self.inner.is_self_node } - pub fn availability_zone(&self) -> Option<&str> { - self.inner.member.availability_zone.as_deref() + pub fn availability_zone(&self) -> Option { + self.inner.member.availability_zone.clone() } pub fn enable_standalone_compactors(&self) -> bool { diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index 0f6c5229176..18e082ecc32 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -31,7 +31,7 @@ use quickwit_common::shared_consts::{ use quickwit_common::uri::Uri; use quickwit_proto::indexing::CpuCapacity; use quickwit_proto::tonic::codec::CompressionEncoding; -use quickwit_proto::types::NodeId; +use quickwit_proto::types::{AvailabilityZone, NodeId}; use serde::{Deserialize, Deserializer, Serialize}; use tracing::{info, warn}; @@ -916,7 +916,7 @@ impl Default for JaegerConfig { pub struct NodeConfig { pub cluster_id: String, pub node_id: NodeId, - pub availability_zone: Option, + pub availability_zone: Option, pub enabled_services: HashSet, pub gossip_listen_addr: SocketAddr, pub grpc_listen_addr: SocketAddr, diff --git a/quickwit/quickwit-config/src/node_config/serialize.rs b/quickwit/quickwit-config/src/node_config/serialize.rs index 4e7edefb433..675d7c01031 100644 --- a/quickwit/quickwit-config/src/node_config/serialize.rs +++ b/quickwit/quickwit-config/src/node_config/serialize.rs @@ -24,7 +24,7 @@ use quickwit_common::fs::get_disk_size; use quickwit_common::net::{Host, find_private_ip, get_short_hostname}; use quickwit_common::new_coolid; use quickwit_common::uri::Uri; -use quickwit_proto::types::NodeId; +use quickwit_proto::types::{AvailabilityZone, NodeId}; use serde::{Deserialize, Serialize}; use tracing::{info, warn}; @@ -255,11 +255,17 @@ impl NodeConfigBuilder { .node_id .resolve(env_vars) .map(|node_id_str| NodeId::from_str(&node_id_str))?; - let availability_zone = self - .availability_zone - .resolve_optional(env_vars)? - .map(|availability_zone| availability_zone.trim().to_string()) - .filter(|availability_zone| !availability_zone.is_empty()); + let availability_zone = + self.availability_zone + .resolve_optional(env_vars)? + .and_then(|availability_zone| { + let availability_zone = availability_zone.trim(); + if availability_zone.is_empty() { + None + } else { + Some(AvailabilityZone::from(availability_zone)) + } + }); let enable_standalone_compactors = self.enable_standalone_compactors.resolve(env_vars)?; let docs_clustering_config = @@ -633,7 +639,7 @@ pub fn node_config_for_tests_from_ports( ) -> NodeConfig { let node_id = NodeId::from_str(&default_node_id().unwrap()); let enabled_services = QuickwitService::default_services(); - let availability_zone = Some(String::from("az-1")); + let availability_zone = Some(AvailabilityZone::from("az-1")); let listen_address = Host::default(); let rest_listen_addr = listen_address .with_port(rest_listen_port) @@ -729,7 +735,7 @@ mod tests { assert!(config.is_service_enabled(QuickwitService::Janitor)); assert!(config.is_service_enabled(QuickwitService::Metastore)); - assert_eq!(config.availability_zone.unwrap(), "az-1"); + assert_eq!(config.availability_zone.as_deref(), Some("az-1")); assert_eq!( config.rest_config.listen_addr, SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 1111) diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index 1ef58b6bd44..01c2d2d9fa4 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -35,7 +35,7 @@ use quickwit_proto::indexing::{ use quickwit_proto::ingest::ingester::IngesterStatus; use quickwit_proto::types::NodeId; use scheduling::{ - AvailabilityZone, Eligibility, IndexerInfo, SourceToSchedule, SourceToScheduleType, + Eligibility, IndexerInfo, SourceToSchedule, SourceToScheduleType, compute_max_num_shards_per_pipeline, is_shard_in_same_zone, }; use serde::Serialize; @@ -332,10 +332,7 @@ fn build_indexer_info(indexer: &IndexerPoolEntry, locality_aware: bool) -> Index }; IndexerInfo { cpu_capacity: indexer.indexing_capacity, - availability_zone: indexer - .availability_zone - .as_deref() - .map(AvailabilityZone::from), + availability_zone: indexer.availability_zone.clone(), eligibility, } } @@ -983,7 +980,7 @@ mod tests { use proptest::{prop_compose, proptest}; use quickwit_config::{IndexConfig, KafkaSourceParams, SourceConfig, SourceParams}; use quickwit_metastore::IndexMetadata; - use quickwit_proto::types::{IndexUid, PipelineUid, ShardId, SourceUid}; + use quickwit_proto::types::{AvailabilityZone, IndexUid, PipelineUid, ShardId, SourceUid}; use super::*; use crate::indexing_scheduler::scheduling::{ @@ -1685,7 +1682,7 @@ mod tests { fn test_all_indexers_advertise_availability_zone() { let indexer_pool = IndexerPool::default(); let mut zoned_indexer = mock_indexer_node_info("indexer-zoned", IngesterStatus::Ready); - zoned_indexer.availability_zone = Some("az-a".to_string()); + zoned_indexer.availability_zone = Some(AvailabilityZone::from("az-a")); indexer_pool.insert(zoned_indexer.node_id.clone(), zoned_indexer); assert!(all_indexers_advertise_availability_zone(&indexer_pool)); @@ -1702,7 +1699,7 @@ mod tests { let locality_aware = true; { let mut ready = mock_indexer_node_info("indexer-ready", IngesterStatus::Ready); - ready.availability_zone = Some("az-a".to_string()); + ready.availability_zone = Some(AvailabilityZone::from("az-a")); let retiring = mock_indexer_node_info("indexer-retiring", IngesterStatus::Retiring); let decommissioning = mock_indexer_node_info("indexer-decommissioning", IngesterStatus::Decommissioning); @@ -1743,9 +1740,9 @@ mod tests { } { let mut ready = mock_indexer_node_info("indexer-ready", IngesterStatus::Ready); - ready.availability_zone = Some("az-a".to_string()); + ready.availability_zone = Some(AvailabilityZone::from("az-a")); let mut retiring = mock_indexer_node_info("indexer-retiring", IngesterStatus::Retiring); - retiring.availability_zone = Some("az-b".to_string()); + retiring.availability_zone = Some(AvailabilityZone::from("az-b")); let indexers = vec![ready, retiring]; let locality_unaware = false; 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 badf9122ce0..8b0470875a0 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs @@ -20,12 +20,11 @@ pub mod scheduling_logic_model; use std::collections::HashMap; use std::num::NonZeroU32; -use std::sync::Arc; use fnv::{FnvHashMap, FnvHashSet}; use quickwit_common::rate_limited_debug; use quickwit_proto::indexing::{CpuCapacity, IndexingTask}; -use quickwit_proto::types::{NodeId, PipelineUid, ShardId, SourceUid}; +use quickwit_proto::types::{AvailabilityZone, NodeId, PipelineUid, ShardId, SourceUid}; pub use scheduling_logic_model::Eligibility; use scheduling_logic_model::{IndexerLocality, IndexerOrd, LocalityGroup, SourceOrd}; use tracing::{error, warn}; @@ -38,8 +37,6 @@ use crate::indexing_scheduler::scheduling::scheduling_logic_model::{ }; use crate::model::ShardLocations; -pub type AvailabilityZone = Arc; - /// If we have several pipelines below this threshold we /// reduce the number of pipelines. /// diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs index c8457086b70..adc436df37b 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/optimization.rs @@ -17,13 +17,13 @@ use std::time::Instant; use fnv::FnvHashMap; use quickwit_proto::indexing::IndexingTask; use quickwit_proto::ingest::ingester::IngesterStatus; -use quickwit_proto::types::{NodeId, ShardId, SourceUid}; +use quickwit_proto::types::{AvailabilityZone, NodeId, ShardId, SourceUid}; use rand::Rng; use rand::seq::SliceRandom; use super::{ - AvailabilityZone, Eligibility, IndexerInfo, SourceToSchedule, - compute_max_num_shards_per_pipeline, shard_availability_zone, + Eligibility, IndexerInfo, SourceToSchedule, compute_max_num_shards_per_pipeline, + shard_availability_zone, }; use crate::IndexerPoolEntry; use crate::indexing_plan::PhysicalIndexingPlan; @@ -364,7 +364,9 @@ mod tests { use fnv::FnvHashMap; use quickwit_proto::indexing::{IndexingTask, mcpu}; use quickwit_proto::ingest::ingester::IngesterStatus; - use quickwit_proto::types::{IndexUid, NodeId, PipelineUid, ShardId, SourceUid}; + use quickwit_proto::types::{ + AvailabilityZone, IndexUid, NodeId, PipelineUid, ShardId, SourceUid, + }; use rand::SeedableRng; use rand::rngs::StdRng; @@ -376,8 +378,7 @@ mod tests { }; use crate::indexing_plan::PhysicalIndexingPlan; use crate::indexing_scheduler::scheduling::{ - AvailabilityZone, Eligibility, IndexerInfo, SourceToSchedule, SourceToScheduleType, - shard_ids_for_indexer, + Eligibility, IndexerInfo, SourceToSchedule, SourceToScheduleType, shard_ids_for_indexer, }; use crate::indexing_scheduler::{ IndexingSchedulerState, MIN_DURATION_BETWEEN_SCHEDULING, get_indexing_plan_density, diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs index 50ca0f9f172..28dcedd227f 100644 --- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs +++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs @@ -12,11 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::fmt; use std::future::Future; use std::num::NonZeroUsize; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use fnv::FnvHashSet; @@ -44,11 +45,11 @@ use quickwit_proto::metastore::{ MetastoreResult, MetastoreService, MetastoreServiceClient, OpenShardSubrequest, OpenShardsRequest, OpenShardsResponse, serde_utils, }; -use quickwit_proto::types::{IndexUid, NodeId, Position, ShardId, SourceUid}; +use quickwit_proto::types::{AvailabilityZone, IndexUid, NodeId, Position, ShardId, SourceUid}; use rand::prelude::IndexedRandom; use rand::rngs::ThreadRng; use rand::seq::SliceRandom; -use rand::{Rng, RngExt, rng}; +use rand::{Rng, rng}; use serde::{Deserialize, Serialize}; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tracing::{Level, debug, enabled, error, info, instrument, warn}; @@ -91,52 +92,203 @@ fn fire_and_forget( }); } -/// Pick an ingester from `ingester_ids_by_num_shards`. -/// We prioritize ingesters with the least number of shards and break ties randomly. -/// -/// Once an ingester has been found, we update `ingester_ids_by_num_shards` to reflect the new -/// state. In particular, the ingester node is moved from its previous num_shards level to its new -/// num_shards level. In particular, a num_shards entry that is empty should be removed from the -/// BTreeMap. -fn pick_least_loaded_ingester<'a>( - ingester_ids_by_num_shards: &mut BTreeMap>, +type SourceShardCount = HashMap; + +fn total_shards(source_shard_counts: &SourceShardCount) -> usize { + source_shard_counts.values().sum() +} + +/// In some cases, it could be advantageous to prefer to place shards in a specific zone (such as +/// when rebalancing the shards of a decommissioned ingester). Otherwise, the default is to spread +/// new shards evenly. +enum ShardPlacement { + Balanced(SourceShardCount), + Zoned(HashMap, SourceShardCount>), +} + +struct EligibleIngester { + node_id: NodeId, + zone: Option, + num_open_shards: AtomicUsize, +} + +/// Find the globally minimally loaded ingester. Break ties with the requested zone, if provided. +fn pick_least_loaded<'a>( + eligible_ingesters: &'a [EligibleIngester], + requested_zone: Option<&AvailabilityZone>, rng: &mut ThreadRng, -) -> Option<&'a NodeId> { - let (&num_shards, ingester_ids) = ingester_ids_by_num_shards.iter_mut().next()?; - let position = rng.random_range(0..ingester_ids.len()); +) -> &'a EligibleIngester { + assert!(!eligible_ingesters.is_empty()); + let mut min_load = usize::MAX; + let mut minima = Vec::new(); + let mut same_zone_minima = Vec::new(); + + for ingester in eligible_ingesters { + let load = ingester.num_open_shards.load(Ordering::Relaxed); + if load > min_load { + continue; + } + if load < min_load { + min_load = load; + minima.clear(); + same_zone_minima.clear(); + } + minima.push(ingester); + if requested_zone.is_some() && ingester.zone.as_ref() == requested_zone { + same_zone_minima.push(ingester); + } + } + let candidates = if !same_zone_minima.is_empty() { + same_zone_minima + } else { + minima + }; + candidates + .choose(rng) + .expect("Candidates came from the list of eligible ingesters, which is not empty") +} - let ingester_id = ingester_ids.swap_remove(position); - let new_num_shards = num_shards + 1; - let should_remove_entry = ingester_ids.is_empty(); +fn all_ingesters_advertise_availability_zone(ingesters: &IngesterPool) -> bool { + ingesters + .keys_values() + .iter() + .all(|(_node_id, ingester)| ingester.availability_zone.is_some()) +} - if should_remove_entry { - ingester_ids_by_num_shards.remove(&num_shards); +fn eligible_ingesters( + ingester_pool: &IngesterPool, + unavailable_ingesters: &FnvHashSet, + model: &ControlPlaneModel, + zonal_placement_enabled: bool, +) -> Vec { + let mut num_open_shards_by_ingester_id: HashMap = HashMap::new(); + for shard in model.all_shards() { + if shard.is_open() { + *num_open_shards_by_ingester_id + .entry(shard.ingester_id.clone()) + .or_default() += 1; + } } - ingester_ids_by_num_shards - .entry(new_num_shards) - .or_default() - .push(ingester_id); - Some(ingester_id) + ingester_pool + .keys_values() + .into_iter() + .filter(|(id, ingester)| ingester.status.is_ready() && !unavailable_ingesters.contains(id)) + .map(|(node_id, ingester)| EligibleIngester { + num_open_shards: AtomicUsize::new( + num_open_shards_by_ingester_id + .get(node_id.as_str()) + .copied() + .unwrap_or(0), + ), + node_id, + // If zonal placement is disabled, every ingester's zone is set to None, creating one + // global "zonal" group. + zone: ingester + .availability_zone + .clone() + .filter(|_| zonal_placement_enabled), + }) + .collect() } fn allocate_shards( - num_shards_by_ingester_id: &HashMap, + eligible_ingesters: &[EligibleIngester], + requested_zone: Option, num_shards: usize, -) -> Option> { - let mut ingester_ids_by_num_shards: BTreeMap> = BTreeMap::default(); - for (ingester_id, &num_shards) in num_shards_by_ingester_id { - ingester_ids_by_num_shards - .entry(num_shards) - .or_default() - .push(ingester_id); - } +) -> Vec { let mut rng = rng(); let mut ingester_ids = Vec::with_capacity(num_shards); for _ in 0..num_shards { - let ingester_id = pick_least_loaded_ingester(&mut ingester_ids_by_num_shards, &mut rng)?; - ingester_ids.push(ingester_id); + let picked = pick_least_loaded(eligible_ingesters, requested_zone.as_ref(), &mut rng); + picked.num_open_shards.fetch_add(1, Ordering::Relaxed); + ingester_ids.push(picked.node_id.clone()); } - Some(ingester_ids) + ingester_ids +} + +fn distribute_shards_across_zones( + num_to_open: usize, + zones: &HashSet, +) -> HashMap, usize> { + if num_to_open == 0 { + return HashMap::new(); + } + if zones.is_empty() { + return HashMap::from([(None, num_to_open)]); + } + let mut shuffled: Vec<&AvailabilityZone> = zones.iter().collect(); + shuffled.shuffle(&mut rng()); + shuffled + .iter() + .cycle() + .take(num_to_open) + .map(|zone| Some((*zone).clone())) + .counts() +} + +/// For each source's requested count, group shards to open by zone. +fn balance_shards_to_open_across_zones( + source_shard_counts: SourceShardCount, + eligible_ingesters: &[EligibleIngester], +) -> HashMap, SourceShardCount> { + let zones: HashSet = eligible_ingesters + .iter() + .filter_map(|ingester| ingester.zone.clone()) + .collect(); + let mut num_shards_by_source_by_zone: HashMap, SourceShardCount> = + HashMap::new(); + for (source_uid, num_shards) in source_shard_counts { + // Number of shards to open for this source in each zone. + for (zone, count) in distribute_shards_across_zones(num_shards, &zones) { + num_shards_by_source_by_zone + .entry(zone) + .or_default() + .insert(source_uid.clone(), count); + } + } + num_shards_by_source_by_zone +} + +/// Rebalancing shards consists of opening new replacements, and then upon success, closing the +/// originals. For the newly opened shards, we find a relevant shard that needed rebalanceing, so +/// that we can close it. +/// +/// The ingester pool can change while replacements are being opened. A shard whose ingester has +/// disappeared from the pool can no longer be closed directly, so it is deliberately left for the +/// control plane's self-healing mechanisms instead of turning normal cluster churn into a panic. +fn match_shards_to_close( + ingester_pool: &IngesterPool, + opened_by_original_zone: &HashMap, SourceShardCount>, + shards_to_rebalance: &mut Vec, +) -> Vec { + let mut shards_to_close: Vec = Vec::new(); + for (original_zone, num_opened_by_source) in opened_by_original_zone { + for (source_uid, &num_opened) in num_opened_by_source { + for num_matched in 0..num_opened { + let Some(position) = shards_to_rebalance.iter().position(|shard| { + shard.source_uid() == *source_uid + && ingester_pool + .get(shard.ingester_id.as_str()) + .and_then(|ingester| ingester.availability_zone.clone()) + == *original_zone + }) else { + // This would only happen if the ingester pool changed underneath after shards + // were opened, and is unlikely, but it is possible. + warn!( + index_uid = %source_uid.index_uid, + source_id = %source_uid.source_id, + ?original_zone, + num_opened, + num_matched, + "could not match every replacement shard to a live predecessor" + ); + break; + }; + shards_to_close.push(shards_to_rebalance.swap_remove(position)); + } + } + } + shards_to_close } #[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)] @@ -405,7 +557,7 @@ impl IngestController { let mut get_or_create_open_shards_successes = Vec::with_capacity(num_subrequests); let mut get_or_create_open_shards_failures = Vec::new(); - let mut num_shards_to_open_by_source = HashMap::new(); + let mut num_shards_to_open_by_source = SourceShardCount::new(); let unavailable_ingesters: FnvHashSet = get_open_shards_request .unavailable_ingesters @@ -446,7 +598,7 @@ impl IngestController { if let Err(metastore_error) = self .try_open_shards( - num_shards_to_open_by_source, + ShardPlacement::Balanced(num_shards_to_open_by_source), model, &unavailable_ingesters, progress, @@ -493,53 +645,6 @@ impl IngestController { Ok(response) } - /// Allocates and assigns new shards to ingesters. - fn allocate_shards( - &self, - num_shards_to_allocate: usize, - unavailable_ingesters: &FnvHashSet, - model: &ControlPlaneModel, - ) -> Option> { - // Count of open shards per available ingester node (including the ingester with 0 open - // shards). - let mut num_open_shards_by_ingester_id: HashMap = self - .ingester_pool - .keys_values() - .into_iter() - .filter(|(ingester_id, ingester)| { - ingester.status.is_ready() && !unavailable_ingesters.contains(ingester_id) - }) - .map(|(ingester_id, _)| (ingester_id, 0)) - .collect(); - - let num_ingesters = num_open_shards_by_ingester_id.len(); - - if num_ingesters == 0 { - warn!("failed to allocate {num_shards_to_allocate} shards: no ingesters available"); - return None; - } - - for shard in model.all_shards() { - if shard.is_open() && !unavailable_ingesters.contains(shard.ingester_id.as_str()) { - if let Some(num_shards) = - num_open_shards_by_ingester_id.get_mut(shard.ingester_id.as_str()) - { - *num_shards += 1; - } else { - // The shard is not present in `num_open_shards_by_ingester_id`. - // This is normal. It just means an ingester is temporarily unavailable, - // either from the control plane view (not present in the indexer pool, - // because as a result of information from chitchat), or because it is in the - // unavailable ingesters map. - } - } - } - - let ingester_ids = - allocate_shards(&num_open_shards_by_ingester_id, num_shards_to_allocate)?; - Some(ingester_ids.into_iter().cloned().collect()) - } - /// Calls init shards on the ingesters hosting newly opened shards. async fn init_shards( &self, @@ -633,22 +738,20 @@ impl IngestController { return Ok(()); } let new_num_open_shards = shard_stats.num_open_shards + num_shards_to_open; - let num_shards_to_open_by_source: HashMap = + let num_shards_to_open_by_source: SourceShardCount = HashMap::from_iter([(source_uid.clone(), num_shards_to_open)]); - let num_opened_shards_by_source_result = self + let try_open_shards_result = self .try_open_shards( - num_shards_to_open_by_source, + ShardPlacement::Balanced(num_shards_to_open_by_source), model, &Default::default(), progress, ) .await; - match num_opened_shards_by_source_result { - Ok(num_opened_shards_by_source) => { - assert!(num_opened_shards_by_source.len() <= 1); - - if num_opened_shards_by_source.is_empty() { + match try_open_shards_result { + Ok(opened_shards) => { + if opened_shards.is_empty() { // We did not manage to create the shard. // We can release our permit. model.release_scaling_permits(&source_uid, ScalingMode::Up(num_shards_to_open)); @@ -681,6 +784,73 @@ impl IngestController { } } + /// Try to open shards given the counts by source and optional requested zones. + async fn try_open_shards( + &mut self, + placement: ShardPlacement, + model: &mut ControlPlaneModel, + unavailable_ingesters: &FnvHashSet, + progress: &Progress, + ) -> MetastoreResult, SourceShardCount>> { + // Zonal aware placement is enabled only after every ingester has advertised its + // zone. If not, every ingester's zone is set to None and the global balancing logic + // applies. + let zonal_placement_enabled = + all_ingesters_advertise_availability_zone(&self.ingester_pool); + let eligible_ingesters = eligible_ingesters( + &self.ingester_pool, + unavailable_ingesters, + model, + zonal_placement_enabled, + ); + if eligible_ingesters.is_empty() { + warn!("failed to open shards: no ingesters available"); + return Ok(HashMap::new()); + } + let num_shards_to_open_by_source_by_zone = match placement { + ShardPlacement::Balanced(num_shards_by_source) => { + balance_shards_to_open_across_zones(num_shards_by_source, &eligible_ingesters) + } + ShardPlacement::Zoned(num_shards_by_source_by_zone) => num_shards_by_source_by_zone, + }; + self.open_shards( + num_shards_to_open_by_source_by_zone, + &eligible_ingesters, + model, + progress, + ) + .await + } + + /// Iterates the per-zone groups, calling [`Self::try_open_shards_by_zone`] for each. The + /// per-AZ structure is preserved in the result so callers can match opens back to + /// their originating (source, AZ) bucket. + async fn open_shards( + &mut self, + num_shards_by_source_by_zone: HashMap, SourceShardCount>, + eligible_ingesters: &[EligibleIngester], + model: &mut ControlPlaneModel, + progress: &Progress, + ) -> MetastoreResult, SourceShardCount>> { + let mut opened_by_zone: HashMap, SourceShardCount> = + HashMap::new(); + for (requested_zone, num_shards_by_source) in num_shards_by_source_by_zone { + let opened = self + .try_open_shards_by_zone( + num_shards_by_source, + requested_zone.clone(), + eligible_ingesters, + model, + progress, + ) + .await?; + if !opened.is_empty() { + opened_by_zone.insert(requested_zone, opened); + } + } + Ok(opened_by_zone) + } + /// Attempts to open shards for different sources /// The values in `num_shards_to_open_by_source` specify how many shards to open for each /// source. @@ -698,23 +868,24 @@ impl IngestController { /// plane model. /// /// The number of successfully open shards is returned. - async fn try_open_shards( + async fn try_open_shards_by_zone( &mut self, - num_shards_to_open_by_source: HashMap, + num_shards_to_open_by_source: SourceShardCount, + requested_zone: Option, + eligible_ingesters: &[EligibleIngester], model: &mut ControlPlaneModel, - unavailable_ingesters: &FnvHashSet, progress: &Progress, - ) -> MetastoreResult> { - let total_num_shards_to_open: usize = num_shards_to_open_by_source.values().sum(); + ) -> MetastoreResult { + let num_shards_to_open = total_shards(&num_shards_to_open_by_source); - if total_num_shards_to_open == 0 { + if num_shards_to_open == 0 { return Ok(HashMap::new()); } - let Some(ingester_ids) = - self.allocate_shards(total_num_shards_to_open, unavailable_ingesters, model) - else { - return Ok(HashMap::new()); - }; + // By now, we've asserted that there's at least one eligible ingester. We have to have + // something to open a shard on. + assert!(!eligible_ingesters.is_empty()); + let ingester_ids = allocate_shards(eligible_ingesters, requested_zone, num_shards_to_open); + let source_uids_with_multiplicity = num_shards_to_open_by_source .iter() .flat_map(|(source_uid, &num_shards)| std::iter::repeat_n(source_uid, num_shards)); @@ -766,12 +937,11 @@ impl IngestController { let open_shard_subrequests = init_shards_response .successes .into_iter() - .enumerate() - .map(|(subrequest_id, init_shard_success)| { + .map(|init_shard_success| { let shard = init_shard_success.shard(); OpenShardSubrequest { - subrequest_id: subrequest_id as u32, + subrequest_id: init_shard_success.subrequest_id, index_uid: shard.index_uid.clone(), source_id: shard.source_id.clone(), shard_id: shard.shard_id.clone(), @@ -791,7 +961,7 @@ impl IngestController { )) .await?; - let mut num_opened_shards_by_source: HashMap = HashMap::new(); + let mut num_opened_shards_by_source: SourceShardCount = HashMap::new(); for open_shard_subresponse in open_shards_response.subresponses { let source_uid = open_shard_subresponse.open_shard().source_uid(); @@ -939,8 +1109,9 @@ impl IngestController { } /// Rebalances shards from ingesters with too many shards to ingesters with too few shards. - /// Moving a shard consists of closing the shard on the source ingester and opening a new - /// one on the target ingester. + /// Moving a shard consists of opening a new one on the target ingester and closing the shard + /// on the source ingester. We attempt to open new shards in the zone in which the original was + /// closed, but fall back to cross-zonal if it creates better global balance. /// /// This method uses a single semaphore permit to ensure that only one rebalance operation is /// performed at a time. @@ -957,7 +1128,7 @@ impl IngestController { }; self.stats.num_rebalance_shards_ops += 1; - let shards_to_rebalance: Vec = self.compute_shards_to_rebalance(model); + let mut shards_to_rebalance: Vec = self.compute_shards_to_rebalance(model); REBALANCE_SHARDS.set(shards_to_rebalance.len() as f64); @@ -965,16 +1136,23 @@ impl IngestController { debug!("skipping rebalance: no shards to rebalance"); return Ok(0); } - let mut num_shards_to_open_by_source: HashMap = HashMap::new(); - + let mut replacement_counts_by_zone: HashMap, SourceShardCount> = + HashMap::new(); for shard in &shards_to_rebalance { - *num_shards_to_open_by_source + let zone = self + .ingester_pool + .get(shard.ingester_id.as_str()) + .and_then(|ingester| ingester.availability_zone.clone()); + *replacement_counts_by_zone + .entry(zone) + .or_default() .entry(shard.source_uid()) .or_default() += 1; } - let mut num_opened_shards_by_source: HashMap = self + + let opened_by_zone = self .try_open_shards( - num_shards_to_open_by_source, + ShardPlacement::Zoned(replacement_counts_by_zone), model, &Default::default(), progress, @@ -985,31 +1163,24 @@ impl IngestController { REBALANCE_SHARDS.set(0.0); })?; - let num_opened_shards: usize = num_opened_shards_by_source.values().sum(); + let num_opened_shards: usize = opened_by_zone.values().map(total_shards).sum(); + let shards_to_close = match_shards_to_close( + &self.ingester_pool, + &opened_by_zone, + &mut shards_to_rebalance, + ); REBALANCE_SHARDS.set(num_opened_shards as f64); - for source_uid in num_opened_shards_by_source.keys() { + for source_uid in opened_by_zone + .values() + .flat_map(|opened_by_source| opened_by_source.keys()) + .unique() + { // We temporarily disable the ability the scale down the number of shards for // the source to avoid closing the shards we just opened. model.drain_scaling_permits(source_uid, ScalingMode::Down); } - - // Close as many shards as we opened. Because `try_open_shards` might fail partially, we - // must only close the shards that we successfully opened. - let mut shards_to_close = Vec::with_capacity(shards_to_rebalance.len()); - - for shard in shards_to_rebalance { - let source_uid = shard.source_uid(); - let Some(num_open_shards) = num_opened_shards_by_source.get_mut(&source_uid) else { - continue; - }; - if *num_open_shards == 0 { - continue; - }; - *num_open_shards -= 1; - shards_to_close.push(shard); - } let close_shards_fut = self.close_shards(shards_to_close); let mailbox_clone = mailbox.clone(); @@ -1271,6 +1442,30 @@ mod tests { const TEST_SHARD_THROUGHPUT_LIMIT_MIB: f32 = DEFAULT_SHARD_THROUGHPUT_LIMIT.as_u64() as f32 / quickwit_common::shared_consts::MIB as f32; + fn ingester_pool_entry( + status: IngesterStatus, + availability_zone: Option<&str>, + ) -> IngesterPoolEntry { + IngesterPoolEntry { + client: IngesterServiceClient::mocked(), + status, + availability_zone: availability_zone.map(AvailabilityZone::from), + generation_id: GenerationId::from(1u64), + } + } + + fn eligible_ingester( + node_id: &str, + zone: Option<&str>, + num_open_shards: usize, + ) -> EligibleIngester { + EligibleIngester { + node_id: NodeId::from_str(node_id), + zone: zone.map(AvailabilityZone::from), + num_open_shards: AtomicUsize::new(num_open_shards), + } + } + #[tokio::test] async fn test_ingest_controller_get_or_create_open_shards() { let source_id: &'static str = "test-source"; @@ -1701,56 +1896,138 @@ mod tests { } #[test] - fn test_ingest_controller_allocate_shards() { - let metastore = MetastoreServiceClient::mocked(); + fn test_eligible_ingesters_preserve_zones_when_all_ingesters_are_zoned() { let ingester_pool = IngesterPool::default(); - - let controller = IngestController::new( - metastore, - ingester_pool.clone(), - TEST_SHARD_THROUGHPUT_LIMIT_MIB, - 1.001, + let ingester_id_0 = NodeId::from_str("test-ingester-0"); + let ingester_id_1 = NodeId::from_str("test-ingester-1"); + let ingester_id_2 = NodeId::from_str("test-ingester-2"); + ingester_pool.insert( + ingester_id_0.clone(), + ingester_pool_entry(IngesterStatus::Ready, Some("az-a")), + ); + ingester_pool.insert( + ingester_id_1.clone(), + ingester_pool_entry(IngesterStatus::Ready, Some("az-b")), + ); + ingester_pool.insert( + ingester_id_2.clone(), + ingester_pool_entry(IngesterStatus::Ready, Some("az-c")), ); - let mut model = ControlPlaneModel::default(); - - let ingester_ids_opt = controller.allocate_shards(0, &FnvHashSet::default(), &model); + let zonal_placement_enabled = all_ingesters_advertise_availability_zone(&ingester_pool); + let eligible_ingesters = eligible_ingesters( + &ingester_pool, + &FnvHashSet::default(), + &ControlPlaneModel::default(), + zonal_placement_enabled, + ); + let zones_by_ingester: HashMap> = eligible_ingesters + .into_iter() + .map(|ingester| (ingester.node_id, ingester.zone)) + .collect(); - // We have no ingesters available, so we can't find any solution. - assert!(ingester_ids_opt.is_none()); + assert_eq!(zones_by_ingester[&ingester_id_0].as_deref(), Some("az-a")); + assert_eq!(zones_by_ingester[&ingester_id_1].as_deref(), Some("az-b")); + assert_eq!(zones_by_ingester[&ingester_id_2].as_deref(), Some("az-c")); + } + #[test] + fn test_eligible_ingesters_mask_zones_when_one_ready_ingester_is_unzoned() { + let ingester_pool = IngesterPool::default(); + ingester_pool.insert( + NodeId::from_str("test-ingester-0"), + ingester_pool_entry(IngesterStatus::Ready, Some("az-a")), + ); ingester_pool.insert( NodeId::from_str("test-ingester-1"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::mocked()), + ingester_pool_entry(IngesterStatus::Ready, Some("az-b")), + ); + ingester_pool.insert( + NodeId::from_str("test-ingester-2"), + ingester_pool_entry(IngesterStatus::Ready, Some("az-c")), + ); + ingester_pool.insert( + NodeId::from_str("test-ingester-unzoned"), + ingester_pool_entry(IngesterStatus::Ready, None), ); - let ingester_ids = controller - .allocate_shards(0, &FnvHashSet::default(), &model) - .unwrap(); + let zonal_placement_enabled = all_ingesters_advertise_availability_zone(&ingester_pool); + let eligible_ingesters = eligible_ingesters( + &ingester_pool, + &FnvHashSet::default(), + &ControlPlaneModel::default(), + zonal_placement_enabled, + ); - // We tried to allocate 0 shards, so an empty vec makes sense. - assert!(ingester_ids.is_empty()); + assert_eq!(eligible_ingesters.len(), 4); + assert!( + eligible_ingesters + .iter() + .all(|ingester| ingester.zone.is_none()) + ); + } - let ingester_ids = controller - .allocate_shards(1, &FnvHashSet::default(), &model) - .unwrap(); + #[test] + fn test_eligible_ingesters_mask_zones_for_unzoned_non_ready_ingester() { + let ingester_pool = IngesterPool::default(); + let ready_ingester_id = NodeId::from_str("test-ingester-0"); + ingester_pool.insert( + ready_ingester_id.clone(), + ingester_pool_entry(IngesterStatus::Ready, Some("az-a")), + ); + ingester_pool.insert( + NodeId::from_str("test-ingester-1"), + ingester_pool_entry(IngesterStatus::Ready, Some("az-b")), + ); + ingester_pool.insert( + NodeId::from_str("test-ingester-2"), + ingester_pool_entry(IngesterStatus::Ready, Some("az-c")), + ); + ingester_pool.insert( + NodeId::from_str("test-ingester-unzoned"), + ingester_pool_entry(IngesterStatus::Retiring, None), + ); - assert_eq!(ingester_ids, ["test-ingester-1"]); + let zonal_placement_enabled = all_ingesters_advertise_availability_zone(&ingester_pool); + let eligible_ingesters = eligible_ingesters( + &ingester_pool, + &FnvHashSet::default(), + &ControlPlaneModel::default(), + zonal_placement_enabled, + ); + assert_eq!(eligible_ingesters.len(), 3); + assert!( + eligible_ingesters + .iter() + .any(|ingester| ingester.node_id == ready_ingester_id) + ); + assert!( + eligible_ingesters + .iter() + .all(|ingester| ingester.zone.is_none()) + ); + } + + #[test] + fn test_eligible_ingesters_and_allocate_shards() { + let ingester_pool = IngesterPool::default(); + ingester_pool.insert( + NodeId::from_str("test-ingester-1"), + ingester_pool_entry(IngesterStatus::Ready, Some("az-a")), + ); ingester_pool.insert( NodeId::from_str("test-ingester-2"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::mocked()), + ingester_pool_entry(IngesterStatus::Ready, Some("az-b")), + ); + ingester_pool.insert( + NodeId::from_str("test-ingester-3"), + ingester_pool_entry(IngesterStatus::Ready, Some("az-c")), ); - - let mut ingester_ids = controller - .allocate_shards(2, &FnvHashSet::default(), &model) - .unwrap(); - ingester_ids.sort_unstable(); - assert_eq!(ingester_ids, ["test-ingester-1", "test-ingester-2"]); let index_uid = IndexUid::for_test("test-index", 0); let source_id: SourceId = "test-source".to_string(); - let open_shards = vec![ + let shards = vec![ Shard { index_uid: Some(index_uid.clone()), source_id: source_id.clone(), @@ -1767,23 +2044,37 @@ mod tests { ingester_id: "test-ingester-1".to_string(), ..Default::default() }, + Shard { + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(3)), + shard_state: ShardState::Closed as i32, + ingester_id: "test-ingester-3".to_string(), + ..Default::default() + }, ]; - model.insert_shards(&index_uid, &source_id, open_shards); - - let ingester_ids = controller - .allocate_shards(1, &FnvHashSet::default(), &model) - .unwrap(); - // Ingester 1 already has two shards, so ingester 2 is picked as ingester. - assert_eq!(ingester_ids, ["test-ingester-2"]); + let mut model = ControlPlaneModel::default(); + model.insert_shards(&index_uid, &source_id, shards); - ingester_pool.insert( - NodeId::from_str("test-ingester-3"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::mocked()), - ); let unavailable_ingesters = FnvHashSet::from_iter([NodeId::from_str("test-ingester-2")]); - let ingester_ids = controller - .allocate_shards(4, &unavailable_ingesters, &model) - .unwrap(); + let eligible_ingesters = + eligible_ingesters(&ingester_pool, &unavailable_ingesters, &model, true); + assert_eq!(eligible_ingesters.len(), 2); + let initial_loads: HashMap<&str, usize> = eligible_ingesters + .iter() + .map(|ingester| { + ( + ingester.node_id.as_str(), + ingester.num_open_shards.load(Ordering::Relaxed), + ) + }) + .collect(); + assert_eq!(initial_loads.get("test-ingester-1"), Some(&2)); + assert_eq!(initial_loads.get("test-ingester-3"), Some(&0)); + + assert!(allocate_shards(&eligible_ingesters, None, 0).is_empty()); + let ingester_ids = allocate_shards(&eligible_ingesters, None, 4); + // Ingester 2 is unavailable. Ingester 1 already has 2 open shards, ingester 3 has none, so // shards are allocated to balance the load between ingester 1 and ingester 3: ingester 3 // ends up with 3 more shards and ingester 1 with 1 more. @@ -2103,14 +2394,14 @@ mod tests { NodeId::from_str("test-ingester-1"), IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)), ); - let num_shards_to_open_by_source: HashMap = + let num_shards_to_open_by_source: SourceShardCount = HashMap::from_iter([(source_uid.clone(), 1)]); let unavailable_ingesters = FnvHashSet::default(); let progress = Progress::default(); - let num_opened_shards_by_source = controller + let opened_by_zone = controller .try_open_shards( - num_shards_to_open_by_source, + ShardPlacement::Balanced(num_shards_to_open_by_source), &mut model, &unavailable_ingesters, &progress, @@ -2118,8 +2409,111 @@ mod tests { .await .unwrap(); - assert_eq!(num_opened_shards_by_source.len(), 1); - assert_eq!(*num_opened_shards_by_source.get(&source_uid).unwrap(), 1); + assert_eq!(opened_by_zone.len(), 1); + assert_eq!(opened_by_zone[&None][&source_uid], 1); + } + + #[tokio::test] + async fn test_try_open_shards_spreads_across_three_zones() { + let index_uid = IndexUid::for_test("test-index", 0); + let source_id = INGEST_V2_SOURCE_ID.to_string(); + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: source_id.clone(), + }; + let mut model = ControlPlaneModel::default(); + model.add_index(IndexMetadata::for_test("test-index", "ram://test-index")); + model + .add_source(&index_uid, SourceConfig::ingest_v2()) + .unwrap(); + + let mut mock_ingester = MockIngesterService::new(); + mock_ingester + .expect_init_shards() + .times(3) + .returning(|request| { + assert_eq!(request.subrequests.len(), 1); + let subrequest = &request.subrequests[0]; + Ok(InitShardsResponse { + successes: vec![InitShardSuccess { + subrequest_id: subrequest.subrequest_id, + shard: subrequest.shard.clone(), + }], + failures: Vec::new(), + }) + }); + let ingester_client = IngesterServiceClient::from_mock(mock_ingester); + let ingester_pool = IngesterPool::default(); + for (ingester_id, zone) in [ + ("ingester-a", "az-a"), + ("ingester-b", "az-b"), + ("ingester-c", "az-c"), + ] { + ingester_pool.insert( + NodeId::from_str(ingester_id), + IngesterPoolEntry { + client: ingester_client.clone(), + status: IngesterStatus::Ready, + availability_zone: Some(AvailabilityZone::from(zone)), + generation_id: GenerationId::from(1u64), + }, + ); + } + + let mut mock_metastore = MockMetastoreService::new(); + mock_metastore + .expect_open_shards() + .times(3) + .returning(|request| { + assert_eq!(request.subrequests.len(), 1); + let subrequest = &request.subrequests[0]; + Ok(OpenShardsResponse { + subresponses: vec![OpenShardSubresponse { + subrequest_id: subrequest.subrequest_id, + open_shard: Some(Shard { + index_uid: subrequest.index_uid.clone(), + source_id: subrequest.source_id.clone(), + shard_id: subrequest.shard_id.clone(), + ingester_id: subrequest.ingester_id.clone(), + shard_state: ShardState::Open as i32, + doc_mapping_uid: subrequest.doc_mapping_uid, + ..Default::default() + }), + }], + }) + }); + let mut controller = IngestController::new( + MetastoreServiceClient::from_mock(mock_metastore), + ingester_pool, + TEST_SHARD_THROUGHPUT_LIMIT_MIB, + 1.001, + ); + + let opened_by_zone = controller + .try_open_shards( + ShardPlacement::Balanced(HashMap::from([(source_uid.clone(), 3)])), + &mut model, + &FnvHashSet::default(), + &Progress::default(), + ) + .await + .unwrap(); + + assert_eq!(opened_by_zone.len(), 3); + for zone in ["az-a", "az-b", "az-c"] { + assert_eq!( + opened_by_zone[&Some(AvailabilityZone::from(zone))][&source_uid], + 1 + ); + } + let ingester_ids: HashSet<&str> = model + .all_shards() + .map(|shard| shard.ingester_id.as_str()) + .collect(); + assert_eq!( + ingester_ids, + HashSet::from(["ingester-a", "ingester-b", "ingester-c"]) + ); } #[tokio::test] @@ -3347,47 +3741,114 @@ mod tests { assert_eq!(controller.rebalance_semaphore.available_permits(), 1); } - // #[track_caller] - fn assert_allocate_shards_balances_load( - num_shards_by_ingester_id: &HashMap, - num_shards: usize, - ) { - let ingester_ids_opt = super::allocate_shards(num_shards_by_ingester_id, num_shards); - if num_shards == 0 { - assert_eq!(ingester_ids_opt, Some(Vec::new())); - return; + #[test] + fn test_match_shards_to_close_is_zonal_and_tolerates_pool_churn() { + let source_uid = SourceUid { + index_uid: IndexUid::for_test("test-index", 0), + source_id: "test-source".to_string(), + }; + let ingester_pool = IngesterPool::default(); + for (ingester_id, zone) in [ + ("origin-a", "az-a"), + ("origin-b", "az-b"), + ("origin-c", "az-c"), + ] { + ingester_pool.insert( + NodeId::from_str(ingester_id), + ingester_pool_entry(IngesterStatus::Retiring, Some(zone)), + ); } - if num_shards_by_ingester_id.is_empty() { - assert!(ingester_ids_opt.is_none()); + let make_shard = |shard_id, ingester_id: &str| Shard { + index_uid: Some(source_uid.index_uid.clone()), + source_id: source_uid.source_id.clone(), + shard_id: Some(ShardId::from(shard_id)), + ingester_id: ingester_id.to_string(), + shard_state: ShardState::Open as i32, + ..Default::default() + }; + // This shard's ingester disappeared after the replacement request was constructed. + let mut shards_to_rebalance = vec![ + make_shard(1, "departed-origin"), + make_shard(2, "origin-a"), + make_shard(3, "origin-b"), + make_shard(4, "origin-c"), + ]; + let opened_by_zone = HashMap::from([ + ( + Some(AvailabilityZone::from("az-a")), + HashMap::from([(source_uid.clone(), 2)]), + ), + ( + Some(AvailabilityZone::from("az-c")), + HashMap::from([(source_uid, 1)]), + ), + ]); + + let mut shards_to_close = + match_shards_to_close(&ingester_pool, &opened_by_zone, &mut shards_to_rebalance); + shards_to_close.sort_by_key(|shard| shard.shard_id().clone()); + + assert_eq!( + shards_to_close + .iter() + .map(|shard| shard.shard_id()) + .collect_vec(), + [ShardId::from(2), ShardId::from(4)] + ); + assert_eq!(shards_to_rebalance.len(), 2); + } + #[track_caller] + fn assert_allocate_shards_balances_load(initial_num_shards: &[usize], num_shards: usize) { + let eligible_ingesters: Vec = initial_num_shards + .iter() + .enumerate() + .map(|(index, &num_open_shards)| { + eligible_ingester(&format!("ingester-{index}"), None, num_open_shards) + }) + .collect(); + if initial_num_shards.is_empty() { return; } - let ingester_ids = ingester_ids_opt.unwrap(); + let ingester_ids = allocate_shards(&eligible_ingesters, None, num_shards); assert_eq!(ingester_ids.len(), num_shards); - let mut current_num_shards_by_ingester_id = num_shards_by_ingester_id.clone(); + let mut current_num_shards_by_ingester_id: HashMap = initial_num_shards + .iter() + .enumerate() + .map(|(index, &num_open_shards)| { + ( + NodeId::from_str(&format!("ingester-{index}")), + num_open_shards, + ) + }) + .collect(); for ingester in ingester_ids { let min_num_shards = current_num_shards_by_ingester_id .values() .copied() .min() .unwrap(); - let num_shards = current_num_shards_by_ingester_id.get_mut(ingester).unwrap(); + let num_shards = current_num_shards_by_ingester_id + .get_mut(&ingester) + .unwrap(); assert_eq!(*num_shards, min_num_shards); *num_shards += 1; } + for ingester in &eligible_ingesters { + assert_eq!( + ingester.num_open_shards.load(Ordering::Relaxed), + current_num_shards_by_ingester_id[&ingester.node_id] + ); + } } fn assert_allocate_shards_for_initial_counts(initial_num_shards: &[usize]) { - let num_shards_by_ingester_id: HashMap = initial_num_shards - .iter() - .enumerate() - .map(|(index, &num_shards)| (NodeId::from_str(&format!("shard-{index}")), num_shards)) - .collect(); for num_shards in 0..10 { - assert_allocate_shards_balances_load(&num_shards_by_ingester_id, num_shards); + assert_allocate_shards_balances_load(initial_num_shards, num_shards); } } use proptest::prelude::*; + use quickwit_cluster::GenerationId; proptest! { #[test] @@ -3406,47 +3867,74 @@ mod tests { assert_allocate_shards_for_initial_counts(&[2, 3, 2]); assert_allocate_shards_for_initial_counts(&[2, 4, 6]); assert_allocate_shards_for_initial_counts(&[2, 3, 10]); + assert_allocate_shards_for_initial_counts(&[7, 7, 7]); } #[test] - fn test_allocate_shards_prop_test_bug() { - assert_allocate_shards_for_initial_counts(&[7, 7, 7]); + fn test_pick_least_loaded_uses_zone_only_to_break_global_ties() { + let mut rng = rand::rng(); + let tied_ingesters = vec![ + eligible_ingester("ingester-a", Some("az-a"), 1), + eligible_ingester("ingester-b", Some("az-b"), 1), + eligible_ingester("ingester-c", Some("az-c"), 1), + ]; + let az_b = AvailabilityZone::from("az-b"); + let picked = pick_least_loaded(&tied_ingesters, Some(&az_b), &mut rng); + assert_eq!(picked.node_id, "ingester-b"); + + let uneven_ingesters = vec![ + eligible_ingester("ingester-a", Some("az-a"), 2), + eligible_ingester("ingester-b", Some("az-b"), 1), + eligible_ingester("ingester-c", Some("az-c"), 2), + ]; + let az_a = AvailabilityZone::from("az-a"); + let picked = pick_least_loaded(&uneven_ingesters, Some(&az_a), &mut rng); + assert_eq!(picked.node_id, "ingester-b"); + let picked = pick_least_loaded(&uneven_ingesters, None, &mut rng); + assert_eq!(picked.node_id, "ingester-b"); } #[test] - fn test_pick_least_loaded_ingester() { - let ingester_id_1 = NodeId::from_str("ingester-1"); - let ingester_id_2 = NodeId::from_str("ingester-2"); - let mut ingester_ids_by_num_shards = BTreeMap::default(); - ingester_ids_by_num_shards.insert(1, vec![&ingester_id_1, &ingester_id_2]); - let mut rng = rand::rng(); - let ingester_id = - pick_least_loaded_ingester(&mut ingester_ids_by_num_shards, &mut rng).unwrap(); - assert!(ingester_id == &ingester_id_1 || ingester_id == &ingester_id_2); - assert_eq!(ingester_ids_by_num_shards.len(), 2); - let remaining_ingester_id = if ingester_id == &ingester_id_1 { - &ingester_id_2 - } else { - &ingester_id_1 - }; + fn test_balance_shards_across_three_distinct_zones() { + let no_zones = HashSet::new(); + assert!(distribute_shards_across_zones(0, &no_zones).is_empty()); assert_eq!( - &ingester_ids_by_num_shards.get(&1).unwrap()[..], - &[remaining_ingester_id] + distribute_shards_across_zones(5, &no_zones), + HashMap::from([(None, 5)]) ); - assert_eq!( - &ingester_ids_by_num_shards.get(&2).unwrap()[..], - &[ingester_id] - ); - let second_ingester_id = - pick_least_loaded_ingester(&mut ingester_ids_by_num_shards, &mut rng).unwrap(); - assert_eq!(second_ingester_id, remaining_ingester_id); - assert_eq!(ingester_ids_by_num_shards.len(), 1); - assert_eq!( - &ingester_ids_by_num_shards.get(&2).unwrap()[..], - &[ingester_id, second_ingester_id] + + let zones = HashSet::from_iter([ + AvailabilityZone::from("az-a"), + AvailabilityZone::from("az-b"), + AvailabilityZone::from("az-c"), + ]); + for (num_shards, expected_num_zones) in [(2, 2), (3, 3), (8, 3)] { + let distribution = distribute_shards_across_zones(num_shards, &zones); + assert_eq!(distribution.len(), expected_num_zones); + assert_eq!(distribution.values().sum::(), num_shards); + let min_count = distribution.values().min().unwrap(); + let max_count = distribution.values().max().unwrap(); + assert!(max_count - min_count <= 1); + } + + let source_uid = SourceUid { + index_uid: IndexUid::for_test("index", 0), + source_id: "source".to_string(), + }; + // Multiple ingesters in az-a must not give that zone more weight than az-b or az-c. + let eligible_ingesters = vec![ + eligible_ingester("ingester-a-0", Some("az-a"), 0), + eligible_ingester("ingester-a-1", Some("az-a"), 0), + eligible_ingester("ingester-b", Some("az-b"), 0), + eligible_ingester("ingester-c", Some("az-c"), 0), + ]; + let balanced = balance_shards_to_open_across_zones( + HashMap::from([(source_uid.clone(), 6)]), + &eligible_ingesters, ); + assert_eq!(balanced.len(), 3); + assert!(balanced.values().all(|counts| counts[&source_uid] == 2)); } - /// Test helper for compute_shards_to_rebalance. /// The reason for testing both available and unavailable ingesters with open shards is to /// ensure the algorithm holds up when there are open shards diff --git a/quickwit/quickwit-control-plane/src/lib.rs b/quickwit/quickwit-control-plane/src/lib.rs index b4f776095ce..45ee0275d21 100644 --- a/quickwit/quickwit-control-plane/src/lib.rs +++ b/quickwit/quickwit-control-plane/src/lib.rs @@ -22,7 +22,7 @@ pub(crate) mod model; use quickwit_common::tower::Pool; use quickwit_proto::indexing::{CpuCapacity, IndexingServiceClient, IndexingTask}; use quickwit_proto::ingest::ingester::IngesterStatus; -use quickwit_proto::types::NodeId; +use quickwit_proto::types::{AvailabilityZone, NodeId}; /// Indexer-node specific information stored in the pool of available indexer nodes #[derive(Debug, Clone)] @@ -33,7 +33,7 @@ pub struct IndexerPoolEntry { pub indexing_tasks: Vec, pub indexing_capacity: CpuCapacity, pub ingester_status: IngesterStatus, - pub availability_zone: Option, + pub availability_zone: Option, } pub type IndexerPool = Pool; diff --git a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs index 1c9343ebfdb..aa81f5c5747 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs @@ -50,7 +50,7 @@ use quickwit_proto::ingest::ingester::{IngesterServiceClient, IngesterStatus}; use quickwit_proto::ingest::router::{IngestRequestV2, IngestSubrequest}; use quickwit_proto::ingest::{CommitTypeV2, DocBatchV2, DocFormat}; use quickwit_proto::types::{ - DocUid, DocUidGenerator, IndexId, IndexUid, NodeId, SourceId, SubrequestId, + AvailabilityZone, DocUid, DocUidGenerator, IndexId, IndexUid, NodeId, SourceId, SubrequestId, }; use serde::Serialize; use tracing::{error, info}; @@ -70,7 +70,7 @@ pub use self::router::IngestRouter; pub struct IngesterPoolEntry { pub client: IngesterServiceClient, pub status: IngesterStatus, - pub availability_zone: Option, + pub availability_zone: Option, pub generation_id: GenerationId, } diff --git a/quickwit/quickwit-ingest/src/ingest_v2/router.rs b/quickwit/quickwit-ingest/src/ingest_v2/router.rs index 5d58a5e4718..d39641133fa 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/router.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/router.rs @@ -36,7 +36,7 @@ use quickwit_proto::ingest::router::{ IngestFailureReason, IngestRequestV2, IngestResponseV2, IngestRouterService, }; use quickwit_proto::ingest::{CommitTypeV2, IngestV2Error, IngestV2Result, RateLimitingCause}; -use quickwit_proto::types::{NodeId, SubrequestId}; +use quickwit_proto::types::{AvailabilityZone, NodeId, SubrequestId}; use serde_json::{Value as JsonValue, json}; use tokio::sync::{Mutex, Semaphore}; use tokio::time::error::Elapsed; @@ -126,7 +126,7 @@ impl IngestRouter { control_plane: ControlPlaneServiceClient, ingester_pool: IngesterPool, event_broker: EventBroker, - self_availability_zone: Option, + self_availability_zone: Option, ) -> Self { let state = Arc::new(Mutex::new(RouterState { debouncer: GetOrCreateOpenShardsRequestDebouncer::default(), @@ -654,7 +654,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let mut workbench = IngestWorkbench::default(); let (get_or_create_open_shard_request_opt, rendezvous) = router @@ -894,7 +894,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![ IngestSubrequest { @@ -991,7 +991,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![IngestSubrequest { subrequest_id: 0, @@ -1049,7 +1049,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![IngestSubrequest { subrequest_id: 0, @@ -1078,7 +1078,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![IngestSubrequest { subrequest_id: 0, @@ -1135,7 +1135,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![IngestSubrequest { subrequest_id: 0, @@ -1201,7 +1201,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![ IngestSubrequest { @@ -1286,7 +1286,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let index_uid_0: IndexUid = IndexUid::for_test("test-index-0", 0); @@ -1448,7 +1448,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let index_uid: IndexUid = IndexUid::for_test("test-index-0", 0); ingester_pool.insert( @@ -1560,7 +1560,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let index_uid_0: IndexUid = IndexUid::for_test("test-index-0", 0); let index_uid_1: IndexUid = IndexUid::for_test("test-index-1", 0); @@ -1625,7 +1625,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let index_uid: IndexUid = IndexUid::for_test("test-index-0", 0); ingester_pool.insert( @@ -1722,7 +1722,7 @@ mod tests { ControlPlaneServiceClient::from_mock(MockControlPlaneService::new()), ingester_pool.clone(), event_broker.clone(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); router.subscribe(); @@ -1776,7 +1776,7 @@ mod tests { ControlPlaneServiceClient::from_mock(MockControlPlaneService::new()), IngesterPool::default(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![ IngestSubrequest { @@ -1871,7 +1871,7 @@ mod tests { ControlPlaneServiceClient::from_mock(MockControlPlaneService::new()), ingester_pool.clone(), EventBroker::default(), - Some("test-az".to_string()), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![IngestSubrequest { subrequest_id: 0, diff --git a/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs b/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs index d3ce2b8bfbe..f774318fd71 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs @@ -18,7 +18,7 @@ use std::collections::{HashMap, HashSet}; use itertools::Itertools; use quickwit_cluster::GenerationId; use quickwit_proto::ingest::Shard; -use quickwit_proto::types::{IndexId, IndexUid, NodeId, SourceId}; +use quickwit_proto::types::{AvailabilityZone, IndexId, IndexUid, NodeId, SourceId}; use rand::rng; use rand::seq::IndexedRandom; @@ -127,7 +127,7 @@ impl RoutingEntry { &self, ingester_pool: &IngesterPool, unavailable_ingesters: &HashSet, - self_availability_zone: &Option, + self_availability_zone: &Option, ) -> Option<&IngesterNode> { let (local_ingesters, remote_ingesters): (Vec<&IngesterNode>, Vec<&IngesterNode>) = self .nodes @@ -136,7 +136,7 @@ impl RoutingEntry { .partition(|node| { let node_az = ingester_pool .get(&node.node_id) - .and_then(|h| h.availability_zone); + .and_then(|h| h.availability_zone.clone()); node_az == *self_availability_zone }); @@ -147,11 +147,11 @@ impl RoutingEntry { #[derive(Debug, Default)] pub(super) struct RoutingTable { table: HashMap<(IndexId, SourceId), RoutingEntry>, - self_availability_zone: Option, + self_availability_zone: Option, } impl RoutingTable { - pub fn new(self_availability_zone: Option) -> Self { + pub fn new(self_availability_zone: Option) -> Self { Self { self_availability_zone, ..Default::default() @@ -184,7 +184,7 @@ impl RoutingTable { }; let target_az = ingester_pool .get(target_node_id) - .and_then(|entry| entry.availability_zone); + .and_then(|entry| entry.availability_zone.clone()); match target_az { Some(ref az) if az == self_az => "same_az", Some(_) => "cross_az", @@ -199,7 +199,9 @@ impl RoutingTable { let mut per_index: HashMap> = HashMap::new(); for ((index_id, source_id), entry) in &self.table { for (node_id, node) in &entry.nodes { - let az = ingester_pool.get(node_id).and_then(|h| h.availability_zone); + let az = ingester_pool + .get(node_id) + .and_then(|h| h.availability_zone.clone()); per_index .entry(index_id.clone()) .or_default() @@ -360,7 +362,7 @@ mod tests { IngesterPoolEntry { client: IngesterServiceClient::mocked(), status: IngesterStatus::Ready, - availability_zone: availability_zone.map(|s| s.to_string()), + availability_zone: availability_zone.map(AvailabilityZone::from), generation_id: GenerationId::from(1u64), } } @@ -704,7 +706,7 @@ mod tests { #[test] fn test_pick_node_prefers_same_az() { - let mut table = RoutingTable::new(Some("az-1".to_string())); + let mut table = RoutingTable::new(Some(AvailabilityZone::from("az-1"))); let pool = IngesterPool::default(); table.apply_capacity_update( @@ -734,7 +736,7 @@ mod tests { #[test] fn test_pick_node_falls_back_to_cross_az() { - let mut table = RoutingTable::new(Some("az-1".to_string())); + let mut table = RoutingTable::new(Some(AvailabilityZone::from("az-1"))); let pool = IngesterPool::default(); table.apply_capacity_update( @@ -776,7 +778,7 @@ mod tests { #[test] fn test_pick_node_missing_entry() { - let table = RoutingTable::new(Some("az-1".to_string())); + let table = RoutingTable::new(Some(AvailabilityZone::from("az-1"))); let pool = IngesterPool::default(); assert!( @@ -971,7 +973,7 @@ mod tests { #[test] fn test_classify_az_locality() { - let table = RoutingTable::new(Some("az-1".to_string())); + let table = RoutingTable::new(Some(AvailabilityZone::from("az-1"))); let pool = IngesterPool::default(); pool.insert( NodeId::from_str("node-local"), diff --git a/quickwit/quickwit-proto/src/types/mod.rs b/quickwit/quickwit-proto/src/types/mod.rs index 4582dac33af..7f83d40191b 100644 --- a/quickwit/quickwit-proto/src/types/mod.rs +++ b/quickwit/quickwit-proto/src/types/mod.rs @@ -49,6 +49,8 @@ pub type SubrequestId = u32; pub type IndexingPlanId = String; +pub type AvailabilityZone = Arc; + /// Uniquely identifies a shard and its underlying mrecordlog queue. pub type QueueId = String; // // diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index 0b47daaaaba..815ae9180c0 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -1280,7 +1280,7 @@ fn build_ingester_insert_change( let pool_entry = IngesterPoolEntry { client: ingester_service, status: node.ingester_status, - availability_zone: node.availability_zone().map(|az| az.to_string()), + availability_zone: node.availability_zone(), generation_id: node.generation_id, }; Change::Insert(node_id, pool_entry) @@ -1514,7 +1514,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()), + availability_zone: node.availability_zone(), }, ) }