From 953170b311f998bf3a4c17aee4f9a7a98b204938 Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Sun, 13 Sep 2026 21:08:10 -0400 Subject: [PATCH 1/7] Zonally aware ingest controller --- .../src/ingest/ingest_controller.rs | 933 +++++++++++++----- 1 file changed, 707 insertions(+), 226 deletions(-) diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs index 50ca0f9f172..9fa642464b8 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; @@ -48,7 +49,7 @@ use quickwit_proto::types::{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,54 +92,200 @@ 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 Zone = String; + +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<&Zone>, 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()); +) -> Option<&'a EligibleIngester> { + let min_load = eligible_ingesters + .iter() + .map(|ingester| ingester.num_open_shards.load(Ordering::Relaxed)) + .min()?; + let minima: Vec<&EligibleIngester> = eligible_ingesters + .iter() + .filter(|ingester| ingester.num_open_shards.load(Ordering::Relaxed) == min_load) + .collect(); + let same_zone_minima: Vec<&EligibleIngester> = minima + .iter() + .copied() + .filter(|ingester| requested_zone.is_some() && ingester.zone.as_ref() == requested_zone) + .collect(); + let candidates = if !same_zone_minima.is_empty() { + same_zone_minima + } else { + minima + }; + candidates.choose(rng).copied() +} - 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 + .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); +) -> Option> { + if eligible_ingesters.is_empty() { + return None; } 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) + .expect("eligible ingesters non-empty"); + picked.num_open_shards.fetch_add(1, Ordering::Relaxed); + ingester_ids.push(picked.node_id.clone()); } Some(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<&Zone> = 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 +} + +/// Matches successful replacement opens to shards that are still hosted by live ingesters. +/// +/// 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_zone: &HashMap, SourceShardCount>, + shards_to_rebalance: &mut Vec, +) -> Vec { + let mut shards_to_close: Vec = Vec::new(); + for (requested_zone, opened_by_source) in opened_by_zone { + for (source_uid, &num_opened) in opened_by_source { + let mut num_matched = 0; + for _ 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) + == *requested_zone + }) else { + break; + }; + shards_to_close.push(shards_to_rebalance.swap_remove(position)); + num_matched += 1; + } + if num_matched < num_opened { + warn!( + index_uid = %source_uid.index_uid, + source_id = %source_uid.source_id, + ?requested_zone, + num_opened, + num_matched, + "could not match every replacement shard to a live predecessor" + ); + } + } + } + shards_to_close +} + #[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)] pub struct IngestControllerStats { pub num_rebalance_shards_ops: usize, @@ -405,7 +552,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 +593,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 +640,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 +733,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 +779,72 @@ 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,20 +862,21 @@ 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) + allocate_shards(eligible_ingesters, requested_zone, num_shards_to_open) else { return Ok(HashMap::new()); }; @@ -766,12 +931,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 +955,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 +1103,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 +1122,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 +1130,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); + *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 +1157,27 @@ impl IngestController { REBALANCE_SHARDS.set(0.0); })?; - let num_opened_shards: usize = num_opened_shards_by_source.values().sum(); + // For every shard we successfully opened, close an equivalent from the zone we requested + // it in. The preferred zone is not necessarily the zone in which the replacement landed. + // Pool membership may have changed while opening replacements, so matching is best effort. + 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 +1439,29 @@ 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(str::to_string), + } + } + + 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(str::to_string), + 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 +1892,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 +2040,41 @@ 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 mut model = ControlPlaneModel::default(); + model.insert_shards(&index_uid, &source_id, 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 unavailable_ingesters = FnvHashSet::from_iter([NodeId::from_str("test-ingester-2")]); + 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)); - ingester_pool.insert( - NodeId::from_str("test-ingester-3"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::mocked()), + assert!(allocate_shards(&[], None, 0).is_none()); + assert_eq!( + allocate_shards(&eligible_ingesters, None, 0), + Some(Vec::new()) ); - let unavailable_ingesters = FnvHashSet::from_iter([NodeId::from_str("test-ingester-2")]); - let ingester_ids = controller - .allocate_shards(4, &unavailable_ingesters, &model) - .unwrap(); + let ingester_ids = allocate_shards(&eligible_ingesters, None, 4).unwrap(); + // 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,107 @@ 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(zone.to_string()), + }, + ); + } + + 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(zone.to_string())][&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,43 +3737,108 @@ 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() { + 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("az-a".to_string()), + HashMap::from([(source_uid.clone(), 2)]), + ), + (Some("az-c".to_string()), 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(); + let ingester_ids_opt = allocate_shards(&eligible_ingesters, None, num_shards); + if initial_num_shards.is_empty() { assert!(ingester_ids_opt.is_none()); return; } let ingester_ids = ingester_ids_opt.unwrap(); 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); } } @@ -3406,47 +3861,73 @@ 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(); + assert!(pick_least_loaded(&[], None, &mut rng).is_none()); + + 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 = "az-b".to_string(); + let picked = pick_least_loaded(&tied_ingesters, Some(&az_b), &mut rng).unwrap(); + 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 = "az-a".to_string(); + let picked = pick_least_loaded(&uneven_ingesters, Some(&az_a), &mut rng).unwrap(); + assert_eq!(picked.node_id, "ingester-b"); + let picked = pick_least_loaded(&uneven_ingesters, None, &mut rng).unwrap(); + 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(["az-a".to_string(), "az-b".to_string(), "az-c".to_string()]); + 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 From 384d00c77f3397ed20cf51d12742b8f0c9009795 Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Mon, 14 Sep 2026 14:34:37 -0400 Subject: [PATCH 2/7] PR comments --- .../src/ingest/ingest_controller.rs | 103 ++++++++---------- quickwit/quickwit-ingest/src/ingest_v2/mod.rs | 3 +- .../quickwit-ingest/src/ingest_v2/router.rs | 30 ++--- .../src/ingest_v2/routing_table.rs | 17 +-- quickwit/quickwit-serve/src/lib.rs | 4 +- 5 files changed, 75 insertions(+), 82 deletions(-) diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs index 9fa642464b8..174d7d95bef 100644 --- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs +++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs @@ -92,7 +92,7 @@ fn fire_and_forget( }); } -type Zone = String; +type Zone = Arc; type SourceShardCount = HashMap; @@ -119,11 +119,13 @@ fn pick_least_loaded<'a>( eligible_ingesters: &'a [EligibleIngester], requested_zone: Option<&Zone>, rng: &mut ThreadRng, -) -> Option<&'a EligibleIngester> { +) -> &'a EligibleIngester { + assert!(!eligible_ingesters.is_empty()); let min_load = eligible_ingesters .iter() .map(|ingester| ingester.num_open_shards.load(Ordering::Relaxed)) - .min()?; + .min() + .expect("There should be at least one eligible ingester"); let minima: Vec<&EligibleIngester> = eligible_ingesters .iter() .filter(|ingester| ingester.num_open_shards.load(Ordering::Relaxed) == min_load) @@ -133,12 +135,12 @@ fn pick_least_loaded<'a>( .copied() .filter(|ingester| requested_zone.is_some() && ingester.zone.as_ref() == requested_zone) .collect(); - let candidates = if !same_zone_minima.is_empty() { + let candidates = if same_zone_minima.len() > 0 { same_zone_minima } else { minima }; - candidates.choose(rng).copied() + candidates.choose(rng).expect("Candidates came from the list of eligible ingesters, which is not empty") } fn all_ingesters_advertise_availability_zone(ingesters: &IngesterPool) -> bool { @@ -187,19 +189,15 @@ fn allocate_shards( eligible_ingesters: &[EligibleIngester], requested_zone: Option, num_shards: usize, -) -> Option> { - if eligible_ingesters.is_empty() { - return None; - } +) -> Vec { let mut rng = rng(); let mut ingester_ids = Vec::with_capacity(num_shards); for _ in 0..num_shards { - let picked = pick_least_loaded(eligible_ingesters, requested_zone.as_ref(), &mut rng) - .expect("eligible ingesters non-empty"); + 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( @@ -257,29 +255,27 @@ fn match_shards_to_close( let mut shards_to_close: Vec = Vec::new(); for (requested_zone, opened_by_source) in opened_by_zone { for (source_uid, &num_opened) in opened_by_source { - let mut num_matched = 0; - for _ in 0..num_opened { - let Some(position) = shards_to_rebalance.iter().position(|shard| { + 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) == *requested_zone - }) else { + ) 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, + ?requested_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)); - num_matched += 1; - } - if num_matched < num_opened { - warn!( - index_uid = %source_uid.index_uid, - source_id = %source_uid.source_id, - ?requested_zone, - num_opened, - num_matched, - "could not match every replacement shard to a live predecessor" - ); } } } @@ -875,11 +871,11 @@ impl IngestController { if num_shards_to_open == 0 { return Ok(HashMap::new()); } - let Some(ingester_ids) = - allocate_shards(eligible_ingesters, requested_zone, num_shards_to_open) - 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)); @@ -1446,7 +1442,7 @@ mod tests { IngesterPoolEntry { client: IngesterServiceClient::mocked(), status, - availability_zone: availability_zone.map(str::to_string), + availability_zone: availability_zone.map(Arc::from), } } @@ -1457,7 +1453,7 @@ mod tests { ) -> EligibleIngester { EligibleIngester { node_id: NodeId::from_str(node_id), - zone: zone.map(str::to_string), + zone: zone.map(Arc::from), num_open_shards: AtomicUsize::new(num_open_shards), } } @@ -2068,12 +2064,8 @@ mod tests { assert_eq!(initial_loads.get("test-ingester-1"), Some(&2)); assert_eq!(initial_loads.get("test-ingester-3"), Some(&0)); - assert!(allocate_shards(&[], None, 0).is_none()); - assert_eq!( - allocate_shards(&eligible_ingesters, None, 0), - Some(Vec::new()) - ); - let ingester_ids = allocate_shards(&eligible_ingesters, None, 4).unwrap(); + 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 @@ -2454,7 +2446,7 @@ mod tests { IngesterPoolEntry { client: ingester_client.clone(), status: IngesterStatus::Ready, - availability_zone: Some(zone.to_string()), + availability_zone: Some(Arc::from(zone)), }, ); } @@ -2500,7 +2492,7 @@ mod tests { assert_eq!(opened_by_zone.len(), 3); for zone in ["az-a", "az-b", "az-c"] { - assert_eq!(opened_by_zone[&Some(zone.to_string())][&source_uid], 1); + assert_eq!(opened_by_zone[&Some(Arc::from(zone))][&source_uid], 1); } let ingester_ids: HashSet<&str> = model .all_shards() @@ -3771,10 +3763,10 @@ mod tests { ]; let opened_by_zone = HashMap::from([ ( - Some("az-a".to_string()), + Some(Arc::from("az-a")), HashMap::from([(source_uid.clone(), 2)]), ), - (Some("az-c".to_string()), HashMap::from([(source_uid, 1)])), + (Some(Arc::from("az-c")), HashMap::from([(source_uid, 1)])), ]); let mut shards_to_close = @@ -3799,12 +3791,10 @@ mod tests { eligible_ingester(&format!("ingester-{index}"), None, num_open_shards) }) .collect(); - let ingester_ids_opt = allocate_shards(&eligible_ingesters, None, num_shards); if initial_num_shards.is_empty() { - assert!(ingester_ids_opt.is_none()); 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: HashMap = initial_num_shards .iter() @@ -3867,15 +3857,13 @@ mod tests { #[test] fn test_pick_least_loaded_uses_zone_only_to_break_global_ties() { let mut rng = rand::rng(); - assert!(pick_least_loaded(&[], None, &mut rng).is_none()); - 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 = "az-b".to_string(); - let picked = pick_least_loaded(&tied_ingesters, Some(&az_b), &mut rng).unwrap(); + let az_b = Arc::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![ @@ -3883,10 +3871,10 @@ mod tests { eligible_ingester("ingester-b", Some("az-b"), 1), eligible_ingester("ingester-c", Some("az-c"), 2), ]; - let az_a = "az-a".to_string(); - let picked = pick_least_loaded(&uneven_ingesters, Some(&az_a), &mut rng).unwrap(); + let az_a = Arc::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).unwrap(); + let picked = pick_least_loaded(&uneven_ingesters, None, &mut rng); assert_eq!(picked.node_id, "ingester-b"); } @@ -3899,8 +3887,11 @@ mod tests { HashMap::from([(None, 5)]) ); - let zones = - HashSet::from_iter(["az-a".to_string(), "az-b".to_string(), "az-c".to_string()]); + let zones = HashSet::from_iter([ + Arc::from("az-a"), + Arc::from("az-b"), + Arc::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); diff --git a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs index 1c9343ebfdb..cf78a63dfa5 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs @@ -34,6 +34,7 @@ mod workbench; use std::collections::HashMap; use std::collections::hash_map::Entry; use std::ops::{Add, AddAssign}; +use std::sync::Arc; use std::time::Duration; use std::{env, fmt}; @@ -70,7 +71,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..48740e7acbc 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/router.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/router.rs @@ -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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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(Arc::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..599e683d5ce 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs @@ -14,6 +14,7 @@ use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use itertools::Itertools; use quickwit_cluster::GenerationId; @@ -127,7 +128,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 @@ -147,11 +148,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() @@ -360,7 +361,7 @@ mod tests { IngesterPoolEntry { client: IngesterServiceClient::mocked(), status: IngesterStatus::Ready, - availability_zone: availability_zone.map(|s| s.to_string()), + availability_zone: availability_zone.map(Arc::from), generation_id: GenerationId::from(1u64), } } @@ -704,7 +705,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(Arc::from("az-1"))); let pool = IngesterPool::default(); table.apply_capacity_update( @@ -734,7 +735,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(Arc::from("az-1"))); let pool = IngesterPool::default(); table.apply_capacity_update( @@ -776,7 +777,7 @@ mod tests { #[test] fn test_pick_node_missing_entry() { - let table = RoutingTable::new(Some("az-1".to_string())); + let table = RoutingTable::new(Some(Arc::from("az-1"))); let pool = IngesterPool::default(); assert!( @@ -971,7 +972,7 @@ mod tests { #[test] fn test_classify_az_locality() { - let table = RoutingTable::new(Some("az-1".to_string())); + let table = RoutingTable::new(Some(Arc::from("az-1"))); let pool = IngesterPool::default(); pool.insert( NodeId::from_str("node-local"), diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index 0b47daaaaba..e10ef930829 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -1155,7 +1155,7 @@ async fn setup_ingest_v2( control_plane.clone(), ingester_pool.clone(), event_broker.clone(), - node_config.availability_zone.clone(), + node_config.availability_zone.as_deref().map(Arc::from), ); ingest_router.subscribe(); setup_ingester_capacity_update_listener(cluster.clone(), event_broker.clone()) @@ -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().map(Arc::from), generation_id: node.generation_id, }; Change::Insert(node_id, pool_entry) From bf044953923f84f7e049c7397ba4804893186db8 Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Mon, 14 Sep 2026 14:45:02 -0400 Subject: [PATCH 3/7] Comments, naming, lints --- .../src/ingest/ingest_controller.rs | 35 +++++++++---------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs index 174d7d95bef..8a350dc422c 100644 --- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs +++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs @@ -135,12 +135,14 @@ fn pick_least_loaded<'a>( .copied() .filter(|ingester| requested_zone.is_some() && ingester.zone.as_ref() == requested_zone) .collect(); - let candidates = if same_zone_minima.len() > 0 { + 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") + candidates + .choose(rng) + .expect("Candidates came from the list of eligible ingesters, which is not empty") } fn all_ingesters_advertise_availability_zone(ingesters: &IngesterPool) -> bool { @@ -242,33 +244,35 @@ fn balance_shards_to_open_across_zones( num_shards_by_source_by_zone } -/// Matches successful replacement opens to shards that are still hosted by live ingesters. +/// 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_zone: &HashMap, SourceShardCount>, + opened_by_original_zone: &HashMap, SourceShardCount>, shards_to_rebalance: &mut Vec, ) -> Vec { let mut shards_to_close: Vec = Vec::new(); - for (requested_zone, opened_by_source) in opened_by_zone { - for (source_uid, &num_opened) in opened_by_source { + 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| + 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) - == *requested_zone - ) else { + == *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, - ?requested_zone, + ?original_zone, num_opened, num_matched, "could not match every replacement shard to a live predecessor" @@ -874,7 +878,7 @@ impl IngestController { // 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 ingester_ids = allocate_shards(eligible_ingesters, requested_zone, num_shards_to_open); let source_uids_with_multiplicity = num_shards_to_open_by_source .iter() @@ -1153,9 +1157,6 @@ impl IngestController { REBALANCE_SHARDS.set(0.0); })?; - // For every shard we successfully opened, close an equivalent from the zone we requested - // it in. The preferred zone is not necessarily the zone in which the replacement landed. - // Pool membership may have changed while opening replacements, so matching is best effort. let num_opened_shards: usize = opened_by_zone.values().map(total_shards).sum(); let shards_to_close = match_shards_to_close( &self.ingester_pool, @@ -3887,11 +3888,7 @@ mod tests { HashMap::from([(None, 5)]) ); - let zones = HashSet::from_iter([ - Arc::from("az-a"), - Arc::from("az-b"), - Arc::from("az-c"), - ]); + let zones = HashSet::from_iter([Arc::from("az-a"), Arc::from("az-b"), Arc::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); From 112cb69088cb724b3122e2d5c941645d4c223a5a Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Wed, 16 Sep 2026 13:56:35 -0400 Subject: [PATCH 4/7] One pass for picking least loaded --- .../src/ingest/ingest_controller.rs | 33 +++++++++++-------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs index 8a350dc422c..100f2957558 100644 --- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs +++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs @@ -121,20 +121,25 @@ fn pick_least_loaded<'a>( rng: &mut ThreadRng, ) -> &'a EligibleIngester { assert!(!eligible_ingesters.is_empty()); - let min_load = eligible_ingesters - .iter() - .map(|ingester| ingester.num_open_shards.load(Ordering::Relaxed)) - .min() - .expect("There should be at least one eligible ingester"); - let minima: Vec<&EligibleIngester> = eligible_ingesters - .iter() - .filter(|ingester| ingester.num_open_shards.load(Ordering::Relaxed) == min_load) - .collect(); - let same_zone_minima: Vec<&EligibleIngester> = minima - .iter() - .copied() - .filter(|ingester| requested_zone.is_some() && ingester.zone.as_ref() == requested_zone) - .collect(); + 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 { From f5c81e12107bd64e06c9a85f3cb74dda30f2a9cb Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Wed, 16 Sep 2026 14:19:20 -0400 Subject: [PATCH 5/7] PR comments, make az globally arc --- quickwit/quickwit-cluster/src/cluster.rs | 2 +- quickwit/quickwit-cluster/src/member.rs | 10 +-- quickwit/quickwit-cluster/src/node.rs | 5 +- .../quickwit-config/src/node_config/mod.rs | 4 +- .../src/node_config/serialize.rs | 16 ++-- .../src/indexing_scheduler/mod.rs | 17 ++-- .../src/indexing_scheduler/scheduling/mod.rs | 5 +- .../scheduling/optimization.rs | 13 ++-- .../src/ingest/ingest_controller.rs | 78 +++++++++++-------- quickwit/quickwit-control-plane/src/lib.rs | 4 +- quickwit/quickwit-ingest/src/ingest_v2/mod.rs | 5 +- .../quickwit-ingest/src/ingest_v2/router.rs | 32 ++++---- .../src/ingest_v2/routing_table.rs | 27 +++---- quickwit/quickwit-proto/src/types/mod.rs | 2 + quickwit/quickwit-serve/src/lib.rs | 6 +- 15 files changed, 122 insertions(+), 104 deletions(-) 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..5ad8d073c7e 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}; @@ -258,8 +258,14 @@ impl NodeConfigBuilder { 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()); + .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 100f2957558..28dcedd227f 100644 --- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs +++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs @@ -45,7 +45,7 @@ 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; @@ -92,8 +92,6 @@ fn fire_and_forget( }); } -type Zone = Arc; - type SourceShardCount = HashMap; fn total_shards(source_shard_counts: &SourceShardCount) -> usize { @@ -105,19 +103,19 @@ fn total_shards(source_shard_counts: &SourceShardCount) -> usize { /// new shards evenly. enum ShardPlacement { Balanced(SourceShardCount), - Zoned(HashMap, SourceShardCount>), + Zoned(HashMap, SourceShardCount>), } struct EligibleIngester { node_id: NodeId, - zone: Option, + 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<&Zone>, + requested_zone: Option<&AvailabilityZone>, rng: &mut ThreadRng, ) -> &'a EligibleIngester { assert!(!eligible_ingesters.is_empty()); @@ -187,6 +185,7 @@ fn eligible_ingesters( // global "zonal" group. zone: ingester .availability_zone + .clone() .filter(|_| zonal_placement_enabled), }) .collect() @@ -194,7 +193,7 @@ fn eligible_ingesters( fn allocate_shards( eligible_ingesters: &[EligibleIngester], - requested_zone: Option, + requested_zone: Option, num_shards: usize, ) -> Vec { let mut rng = rng(); @@ -209,15 +208,15 @@ fn allocate_shards( fn distribute_shards_across_zones( num_to_open: usize, - zones: &HashSet, -) -> HashMap, 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<&Zone> = zones.iter().collect(); + let mut shuffled: Vec<&AvailabilityZone> = zones.iter().collect(); shuffled.shuffle(&mut rng()); shuffled .iter() @@ -231,12 +230,13 @@ fn distribute_shards_across_zones( fn balance_shards_to_open_across_zones( source_shard_counts: SourceShardCount, eligible_ingesters: &[EligibleIngester], -) -> HashMap, SourceShardCount> { - let zones: HashSet = eligible_ingesters +) -> 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(); + 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) { @@ -258,7 +258,7 @@ fn balance_shards_to_open_across_zones( /// 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>, + opened_by_original_zone: &HashMap, SourceShardCount>, shards_to_rebalance: &mut Vec, ) -> Vec { let mut shards_to_close: Vec = Vec::new(); @@ -269,7 +269,7 @@ fn match_shards_to_close( shard.source_uid() == *source_uid && ingester_pool .get(shard.ingester_id.as_str()) - .and_then(|ingester| ingester.availability_zone) + .and_then(|ingester| ingester.availability_zone.clone()) == *original_zone }) else { // This would only happen if the ingester pool changed underneath after shards @@ -791,7 +791,7 @@ impl IngestController { model: &mut ControlPlaneModel, unavailable_ingesters: &FnvHashSet, progress: &Progress, - ) -> MetastoreResult, SourceShardCount>> { + ) -> 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. @@ -827,12 +827,13 @@ impl IngestController { /// their originating (source, AZ) bucket. async fn open_shards( &mut self, - num_shards_by_source_by_zone: HashMap, SourceShardCount>, + 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(); + ) -> 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( @@ -870,7 +871,7 @@ impl IngestController { async fn try_open_shards_by_zone( &mut self, num_shards_to_open_by_source: SourceShardCount, - requested_zone: Option, + requested_zone: Option, eligible_ingesters: &[EligibleIngester], model: &mut ControlPlaneModel, progress: &Progress, @@ -1135,13 +1136,13 @@ impl IngestController { debug!("skipping rebalance: no shards to rebalance"); return Ok(0); } - let mut replacement_counts_by_zone: HashMap, SourceShardCount> = + let mut replacement_counts_by_zone: HashMap, SourceShardCount> = HashMap::new(); for shard in &shards_to_rebalance { let zone = self .ingester_pool .get(shard.ingester_id.as_str()) - .and_then(|ingester| ingester.availability_zone); + .and_then(|ingester| ingester.availability_zone.clone()); *replacement_counts_by_zone .entry(zone) .or_default() @@ -1448,7 +1449,8 @@ mod tests { IngesterPoolEntry { client: IngesterServiceClient::mocked(), status, - availability_zone: availability_zone.map(Arc::from), + availability_zone: availability_zone.map(AvailabilityZone::from), + generation_id: GenerationId::from(1u64), } } @@ -1459,7 +1461,7 @@ mod tests { ) -> EligibleIngester { EligibleIngester { node_id: NodeId::from_str(node_id), - zone: zone.map(Arc::from), + zone: zone.map(AvailabilityZone::from), num_open_shards: AtomicUsize::new(num_open_shards), } } @@ -1919,7 +1921,7 @@ mod tests { &ControlPlaneModel::default(), zonal_placement_enabled, ); - let zones_by_ingester: HashMap> = eligible_ingesters + let zones_by_ingester: HashMap> = eligible_ingesters .into_iter() .map(|ingester| (ingester.node_id, ingester.zone)) .collect(); @@ -2452,7 +2454,8 @@ mod tests { IngesterPoolEntry { client: ingester_client.clone(), status: IngesterStatus::Ready, - availability_zone: Some(Arc::from(zone)), + availability_zone: Some(AvailabilityZone::from(zone)), + generation_id: GenerationId::from(1u64), }, ); } @@ -2498,7 +2501,10 @@ mod tests { assert_eq!(opened_by_zone.len(), 3); for zone in ["az-a", "az-b", "az-c"] { - assert_eq!(opened_by_zone[&Some(Arc::from(zone))][&source_uid], 1); + assert_eq!( + opened_by_zone[&Some(AvailabilityZone::from(zone))][&source_uid], + 1 + ); } let ingester_ids: HashSet<&str> = model .all_shards() @@ -3769,10 +3775,13 @@ mod tests { ]; let opened_by_zone = HashMap::from([ ( - Some(Arc::from("az-a")), + Some(AvailabilityZone::from("az-a")), HashMap::from([(source_uid.clone(), 2)]), ), - (Some(Arc::from("az-c")), HashMap::from([(source_uid, 1)])), + ( + Some(AvailabilityZone::from("az-c")), + HashMap::from([(source_uid, 1)]), + ), ]); let mut shards_to_close = @@ -3839,6 +3848,7 @@ mod tests { } use proptest::prelude::*; + use quickwit_cluster::GenerationId; proptest! { #[test] @@ -3868,7 +3878,7 @@ mod tests { eligible_ingester("ingester-b", Some("az-b"), 1), eligible_ingester("ingester-c", Some("az-c"), 1), ]; - let az_b = Arc::from("az-b"); + 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"); @@ -3877,7 +3887,7 @@ mod tests { eligible_ingester("ingester-b", Some("az-b"), 1), eligible_ingester("ingester-c", Some("az-c"), 2), ]; - let az_a = Arc::from("az-a"); + 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); @@ -3893,7 +3903,11 @@ mod tests { HashMap::from([(None, 5)]) ); - let zones = HashSet::from_iter([Arc::from("az-a"), Arc::from("az-b"), Arc::from("az-c")]); + 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); 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 cf78a63dfa5..aa81f5c5747 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs @@ -34,7 +34,6 @@ mod workbench; use std::collections::HashMap; use std::collections::hash_map::Entry; use std::ops::{Add, AddAssign}; -use std::sync::Arc; use std::time::Duration; use std::{env, fmt}; @@ -51,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}; @@ -71,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 48740e7acbc..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(Arc::from("test-az")), + 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(Arc::from("test-az")), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![ IngestSubrequest { @@ -991,7 +991,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some(Arc::from("test-az")), + 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(Arc::from("test-az")), + 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(Arc::from("test-az")), + 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(Arc::from("test-az")), + 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(Arc::from("test-az")), + Some(AvailabilityZone::from("test-az")), ); let ingest_subrequests = vec![ IngestSubrequest { @@ -1286,7 +1286,7 @@ mod tests { control_plane, ingester_pool.clone(), EventBroker::default(), - Some(Arc::from("test-az")), + 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(Arc::from("test-az")), + 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(Arc::from("test-az")), + 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(Arc::from("test-az")), + 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(Arc::from("test-az")), + Some(AvailabilityZone::from("test-az")), ); router.subscribe(); @@ -1776,7 +1776,7 @@ mod tests { ControlPlaneServiceClient::from_mock(MockControlPlaneService::new()), IngesterPool::default(), EventBroker::default(), - Some(Arc::from("test-az")), + 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(Arc::from("test-az")), + 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 599e683d5ce..f774318fd71 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs @@ -14,12 +14,11 @@ use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; -use std::sync::Arc; 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; @@ -128,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 @@ -137,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 }); @@ -148,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() @@ -185,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", @@ -200,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() @@ -361,7 +362,7 @@ mod tests { IngesterPoolEntry { client: IngesterServiceClient::mocked(), status: IngesterStatus::Ready, - availability_zone: availability_zone.map(Arc::from), + availability_zone: availability_zone.map(AvailabilityZone::from), generation_id: GenerationId::from(1u64), } } @@ -705,7 +706,7 @@ mod tests { #[test] fn test_pick_node_prefers_same_az() { - let mut table = RoutingTable::new(Some(Arc::from("az-1"))); + let mut table = RoutingTable::new(Some(AvailabilityZone::from("az-1"))); let pool = IngesterPool::default(); table.apply_capacity_update( @@ -735,7 +736,7 @@ mod tests { #[test] fn test_pick_node_falls_back_to_cross_az() { - let mut table = RoutingTable::new(Some(Arc::from("az-1"))); + let mut table = RoutingTable::new(Some(AvailabilityZone::from("az-1"))); let pool = IngesterPool::default(); table.apply_capacity_update( @@ -777,7 +778,7 @@ mod tests { #[test] fn test_pick_node_missing_entry() { - let table = RoutingTable::new(Some(Arc::from("az-1"))); + let table = RoutingTable::new(Some(AvailabilityZone::from("az-1"))); let pool = IngesterPool::default(); assert!( @@ -972,7 +973,7 @@ mod tests { #[test] fn test_classify_az_locality() { - let table = RoutingTable::new(Some(Arc::from("az-1"))); + 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 e10ef930829..b33a14bc959 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -1155,7 +1155,7 @@ async fn setup_ingest_v2( control_plane.clone(), ingester_pool.clone(), event_broker.clone(), - node_config.availability_zone.as_deref().map(Arc::from), + node_config.availability_zone.clone(), ); ingest_router.subscribe(); setup_ingester_capacity_update_listener(cluster.clone(), event_broker.clone()) @@ -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(Arc::from), + availability_zone: node.availability_zone().cloned(), 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().cloned(), }, ) } From ec57bacb361825ffcadc83520d123f6d7140ac73 Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Wed, 16 Sep 2026 14:37:14 -0400 Subject: [PATCH 6/7] compile --- quickwit/quickwit-serve/src/lib.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index b33a14bc959..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().cloned(), + 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().cloned(), + availability_zone: node.availability_zone(), }, ) } From 44fd47f3dcf61435be89bb9b8346c3d88d1a5de2 Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Wed, 16 Sep 2026 14:47:33 -0400 Subject: [PATCH 7/7] lints ffs --- .../src/node_config/serialize.rs | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/quickwit/quickwit-config/src/node_config/serialize.rs b/quickwit/quickwit-config/src/node_config/serialize.rs index 5ad8d073c7e..675d7c01031 100644 --- a/quickwit/quickwit-config/src/node_config/serialize.rs +++ b/quickwit/quickwit-config/src/node_config/serialize.rs @@ -255,17 +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)? - .and_then(|availability_zone| { - let availability_zone = availability_zone.trim(); - if availability_zone.is_empty() { - None - } else { - Some(AvailabilityZone::from(availability_zone)) - } - }); + 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 =