diff --git a/quickwit/quickwit-control-plane/src/control_plane.rs b/quickwit/quickwit-control-plane/src/control_plane.rs index d17c8faf005..793f5d3b066 100644 --- a/quickwit/quickwit-control-plane/src/control_plane.rs +++ b/quickwit/quickwit-control-plane/src/control_plane.rs @@ -59,8 +59,8 @@ use crate::IndexerPool; use crate::cooldown_map::{CooldownMap, CooldownStatus}; use crate::debouncer::Debouncer; use crate::indexing_scheduler::{IndexingScheduler, IndexingSchedulerState}; -use crate::ingest::IngestController; use crate::ingest::ingest_controller::{IngestControllerStats, RebalanceShardsCallback}; +use crate::ingest::{IngestController, LegacyScalingController}; use crate::metrics::{METASTORE_ERROR_ABORTED, METASTORE_ERROR_MAYBE_EXECUTED, RESTART_TOTAL}; use crate::model::ControlPlaneModel; @@ -98,6 +98,7 @@ pub struct ControlPlane { // the different ingesters. indexing_scheduler: IndexingScheduler, ingest_controller: IngestController, + legacy_scaling_controller: LegacyScalingController, metastore: MetastoreServiceClient, model: ControlPlaneModel, prune_shard_cooldown: CooldownMap<(IndexId, SourceId)>, @@ -131,9 +132,9 @@ impl ControlPlane { let cluster_id = cluster_config.cluster_id.clone(); let indexing_scheduler = IndexingScheduler::new(cluster_id, self_node_id.clone(), indexer_pool.clone()); - let ingest_controller = IngestController::new( - metastore.clone(), - ingester_pool.clone(), + let ingest_controller = + IngestController::new(metastore.clone(), ingester_pool.clone()); + let legacy_scaling_controller = LegacyScalingController::new( cluster_config.shard_throughput_limit, cluster_config.shard_scale_up_factor, ); @@ -145,6 +146,7 @@ impl ControlPlane { cluster_config: cluster_config.clone(), indexing_scheduler, ingest_controller, + legacy_scaling_controller, metastore: metastore.clone(), model: Default::default(), prune_shard_cooldown: CooldownMap::new(NonZeroUsize::new(1024).unwrap()), @@ -950,8 +952,9 @@ impl Handler for ControlPlane { ctx: &ActorContext, ) -> Result { if let Err(metastore_error) = self - .ingest_controller - .update_source_shards( + .legacy_scaling_controller + .update_local_shards( + &mut self.ingest_controller, local_shards_update.source_uid, &local_shards_update.shard_infos, &mut self.model, @@ -993,8 +996,9 @@ impl DeferableReplyHandler for ControlPlane { } if let Some(shards_update) = request.shards_update && let Err(metastore_error) = self - .ingest_controller + .legacy_scaling_controller .handle_shards_update( + &mut self.ingest_controller, &request.node_id, request.generation_id, shards_update, @@ -1178,9 +1182,8 @@ mod tests { NodeId::from_str("cp"), indexer_pool, ), - ingest_controller: IngestController::new( - metastore.clone(), - ingester_pool, + ingest_controller: IngestController::new(metastore.clone(), ingester_pool), + legacy_scaling_controller: LegacyScalingController::new( cluster_config.shard_throughput_limit, cluster_config.shard_scale_up_factor, ), diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs index 743d6ed9828..e8975068e03 100644 --- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs +++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs @@ -15,12 +15,10 @@ 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 bytesize::ByteSize; use fnv::FnvHashSet; use futures::StreamExt; use futures::stream::FuturesUnordered; @@ -28,11 +26,11 @@ use itertools::{Itertools as _, MinMaxResult}; use quickwit_actors::Mailbox; use quickwit_common::Progress; use quickwit_common::pretty::PrettySample; -use quickwit_ingest::{IngesterPool, ShardInfos, SourceShardReport}; +use quickwit_ingest::IngesterPool; use quickwit_proto::control_plane::{ AdviseResetShardsRequest, AdviseResetShardsResponse, GetOrCreateOpenShardsFailureReason, GetOrCreateOpenShardsRequest, GetOrCreateOpenShardsResponse, GetOrCreateOpenShardsSubrequest, - GetOrCreateOpenShardsSuccess, ShardsUpdate, + GetOrCreateOpenShardsSuccess, }; use quickwit_proto::ingest::ingester::{ CloseShardsRequest, CloseShardsResponse, IngesterService, IngesterStatus, InitShardFailure, @@ -48,19 +46,18 @@ use quickwit_proto::metastore::{ }; use quickwit_proto::types::{AvailabilityZone, IndexUid, NodeId, Position, ShardId, SourceUid}; use rand::prelude::IndexedRandom; +use rand::rng; use rand::rngs::ThreadRng; use rand::seq::SliceRandom; -use rand::{Rng, rng}; use serde::{Deserialize, Serialize}; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tracing::{Level, debug, enabled, error, info, instrument, warn}; use ulid::Ulid; -use super::scaling_arbiter::ScalingArbiter; use crate::control_plane::ControlPlane; use crate::ingest::wait_handle::WaitHandle; use crate::metrics::REBALANCE_SHARDS; -use crate::model::{ControlPlaneModel, ScalingMode, ShardEntry, ShardStats}; +use crate::model::{ControlPlaneModel, ScalingMode}; const CLOSE_SHARDS_REQUEST_TIMEOUT: Duration = if cfg!(test) { Duration::from_millis(50) @@ -303,7 +300,6 @@ pub struct IngestController { metastore: MetastoreServiceClient, // This semaphore ensures that only one rebalance operation is performed at a time. rebalance_semaphore: Arc, - scaling_arbiter: ScalingArbiter, } impl fmt::Debug for IngestController { @@ -390,21 +386,12 @@ fn get_open_shard_from_model( } impl IngestController { - pub fn new( - metastore: MetastoreServiceClient, - ingester_pool: IngesterPool, - max_shard_ingestion_throughput: ByteSize, - shard_scale_up_factor: f32, - ) -> Self { + pub fn new(metastore: MetastoreServiceClient, ingester_pool: IngesterPool) -> Self { IngestController { metastore, ingester_pool, rebalance_semaphore: Arc::new(Semaphore::new(1)), stats: IngestControllerStats::default(), - scaling_arbiter: ScalingArbiter::with_max_shard_ingestion_throughput( - max_shard_ingestion_throughput, - shard_scale_up_factor, - ), } } @@ -491,74 +478,6 @@ impl IngestController { } } - pub(crate) async fn handle_shards_update( - &mut self, - node_id: &str, - generation_id: u64, - shards_update: ShardsUpdate, - model: &mut ControlPlaneModel, - progress: &Progress, - ) -> MetastoreResult<()> { - if let Some(ingester) = self.ingester_pool.get(node_id) - && generation_id != ingester.generation_id.as_u64() - { - return Ok(()); - } - for source_shard_infos in &shards_update.shard_infos_by_source { - let SourceShardReport { - source_uid, - shard_infos, - } = source_shard_infos.into(); - - if let Err(metastore_error) = self - .update_source_shards(source_uid, &shard_infos, model, progress) - .await - { - if !metastore_error.is_transaction_certainly_aborted() { - return Err(metastore_error); - } - error!(error=?metastore_error, "failed to update source shards"); - } - } - Ok(()) - } - - pub(crate) async fn update_source_shards( - &mut self, - source_uid: SourceUid, - shard_infos: &ShardInfos, - model: &mut ControlPlaneModel, - progress: &Progress, - ) -> MetastoreResult<()> { - let shard_stats = model.update_shards(&source_uid, shard_infos); - let Some(min_shards) = model - .index_metadata(&source_uid.index_uid) - .map(|index_metadata| index_metadata.index_config.ingest_settings.min_shards) - else { - warn!( - index_uid=%source_uid.index_uid, - "ignoring local shards update for a deleted index" - ); - return Ok(()); - }; - - let Some(scaling_mode) = self.scaling_arbiter.should_scale(shard_stats, min_shards) else { - return Ok(()); - }; - match scaling_mode { - ScalingMode::Up(num_shards) => { - self.try_scale_up_shards(source_uid, shard_stats, model, progress, num_shards) - .await?; - } - ScalingMode::Down => { - self.try_scale_down_shards(source_uid, shard_stats, min_shards, model, progress) - .await?; - } - } - - Ok(()) - } - /// Finds the open shards that satisfies the [`GetOrCreateOpenShardsRequest`] request sent by an /// ingest router. First, the control plane checks its internal shard table to find /// candidates. If it does not contain any, the control plane will ask @@ -743,63 +662,6 @@ impl IngestController { } } - /// Attempts to increase the number of shards. This operation is rate limited to avoid creating - /// to many shards in a short period of time. As a result, this method may not create any - /// shard. - async fn try_scale_up_shards( - &mut self, - source_uid: SourceUid, - shard_stats: ShardStats, - model: &mut ControlPlaneModel, - progress: &Progress, - num_shards_to_open: usize, - ) -> MetastoreResult<()> { - if !model - .acquire_scaling_permits(&source_uid, ScalingMode::Up(num_shards_to_open)) - .unwrap_or(false) - { - return Ok(()); - } - let new_num_open_shards = shard_stats.num_open_shards + num_shards_to_open; - let open_shards_result = self - .open_shards_for_source(&source_uid, num_shards_to_open, model, progress) - .await; - - match open_shards_result { - Ok(num_opened_shards) => { - if num_opened_shards == 0 { - // 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)); - warn!( - index_uid=%source_uid.index_uid, - source_id=%source_uid.source_id, - "scaling up number of shards to {new_num_open_shards} failed: shard initialization failure" - ); - } else { - info!( - index_id=%source_uid.index_uid.index_id, - source_id=%source_uid.source_id, - "successfully scaled up number of shards to {new_num_open_shards}" - ); - } - Ok(()) - } - Err(metastore_error) => { - // We did not manage to create the shard. - // We can release our permit, but we also need to return the error to the caller, in - // order to restart the control plane actor if necessary. - warn!( - index_id=%source_uid.index_uid.index_id, - source_id=%source_uid.source_id, - "scaling up number of shards to {new_num_open_shards} failed: {metastore_error:?}" - ); - model.release_scaling_permits(&source_uid, ScalingMode::Up(num_shards_to_open)); - Err(metastore_error) - } - } - } - /// Try to open shards given the counts by source and optional requested zones. async fn try_open_shards( &mut self, @@ -987,54 +849,6 @@ impl IngestController { Ok(num_opened_shards_by_source) } - /// Attempts to decrease the number of shards. This operation is rate limited to avoid closing - /// shards too aggressively. As a result, this method may not close any shard. - async fn try_scale_down_shards( - &self, - source_uid: SourceUid, - shard_stats: ShardStats, - min_shards: NonZeroUsize, - model: &mut ControlPlaneModel, - progress: &Progress, - ) -> MetastoreResult<()> { - // The scaling arbiter should not suggest scaling down if the number of shards is already - // below the minimum, but we're just being defensive here. - if shard_stats.num_open_shards <= min_shards.get() { - return Ok(()); - } - if !model - .acquire_scaling_permits(&source_uid, ScalingMode::Down) - .unwrap_or(false) - { - return Ok(()); - } - let new_num_open_shards = shard_stats.num_open_shards - 1; - - info!( - index_id=%source_uid.index_uid.index_id, - source_id=%source_uid.source_id, - "scaling down number of shards to {new_num_open_shards}" - ); - let Some(shard) = find_scale_down_candidate(&source_uid, model) else { - model.release_scaling_permits(&source_uid, ScalingMode::Down); - return Ok(()); - }; - info!( - "scaling down shard {} from {}", - shard.shard_id(), - shard.ingester_id - ); - let closed_shard_ids = self - .close_source_shards(&source_uid, vec![shard], model, progress) - .await; - - if closed_shard_ids.is_empty() { - warn!("failed to scale down number of shards"); - model.release_scaling_permits(&source_uid, ScalingMode::Down); - } - Ok(()) - } - pub(crate) fn advise_reset_shards( &self, request: AdviseResetShardsRequest, @@ -1423,31 +1237,6 @@ pub(crate) struct RebalanceShardsCallback { pub rebalance_permit: OwnedSemaphorePermit, } -/// Finds a shard on the ingester with the highest number of open -/// shards for this source. -/// -/// If multiple shards are hosted on that ingester, the shard with the lowest (oldest) -/// shard ID is chosen. -fn find_scale_down_candidate(source_uid: &SourceUid, model: &ControlPlaneModel) -> Option { - let mut shard_entries_by_ingester_id: HashMap> = HashMap::new(); - let mut rng = rng(); - - for shard in model.get_shards_for_source(source_uid)?.values() { - if shard.is_open() { - shard_entries_by_ingester_id - .entry(NodeId::from_str(&shard.ingester_id)) - .or_default() - .push(shard); - } - } - shard_entries_by_ingester_id - .into_iter() - // We use a random number to break ties... The HashMap is randomly seeded so this is - // should not make much difference, but we might want to be as explicit as possible. - .max_by_key(|(_ingester_id, shard_entries)| (shard_entries.len(), rng.next_u32())) - .map(|(_ingester_id, shard_entries)| shard_entries.choose(&mut rng).unwrap().shard.clone()) -} - #[cfg(test)] mod tests { use std::collections::BTreeSet; @@ -1458,10 +1247,9 @@ mod tests { use itertools::Itertools; use quickwit_actors::Universe; use quickwit_common::setup_logging_for_tests; - use quickwit_common::shared_consts::DEFAULT_SHARD_THROUGHPUT_LIMIT; use quickwit_common::tower::DelayLayer; use quickwit_config::{DocMapping, INGEST_V2_SOURCE_ID, SourceConfig}; - use quickwit_ingest::{IngesterPoolEntry, ShardInfo}; + use quickwit_ingest::IngesterPoolEntry; use quickwit_metastore::IndexMetadata; use quickwit_proto::control_plane::GetOrCreateOpenShardsSubrequest; use quickwit_proto::ingest::ingester::{ @@ -1500,185 +1288,6 @@ mod tests { } } - fn shard_reports_for_test(min_shards: usize) -> (ControlPlaneModel, ShardsUpdate) { - let mut model = ControlPlaneModel::default(); - let mut metadata = IndexMetadata::for_test("index", "ram:///index"); - metadata.index_config.ingest_settings.min_shards = - std::num::NonZeroUsize::new(min_shards).unwrap(); - let index_uid = metadata.index_uid.clone(); - model.add_index(metadata); - let mut update = ShardsUpdate::default(); - for source_id in ["source-a", "source-b"] { - model - .add_source( - &index_uid, - SourceConfig::for_test(source_id, quickwit_config::SourceParams::void()), - ) - .unwrap(); - model.insert_shards( - &index_uid, - &source_id.to_string(), - vec![Shard { - index_uid: Some(index_uid.clone()), - source_id: source_id.to_string(), - shard_id: Some(ShardId::from(1)), - ingester_id: "ingester".to_string(), - shard_state: ShardState::Open as i32, - ..Default::default() - }], - ); - update.shard_infos_by_source.push( - SourceShardReport { - source_uid: SourceUid { - index_uid: index_uid.clone(), - source_id: source_id.to_string(), - }, - shard_infos: BTreeSet::from([ShardInfo { - shard_id: ShardId::from(1), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::b(123), - long_term_ingestion_rate: ByteSize::b(456), - }]), - } - .into(), - ); - } - (model, update) - } - - #[tokio::test] - async fn test_shards_update_sources_and_generation() { - let pool = IngesterPool::default(); - let ingester = IngesterPoolEntry::ready_with_client(IngesterServiceClient::mocked()); - let generation = ingester.generation_id.as_u64(); - pool.insert(NodeId::from_str("ingester"), ingester); - let mut controller = IngestController::new( - MetastoreServiceClient::mocked(), - pool, - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.5, - ); - let progress = Progress::default(); - let (mut model, update) = shard_reports_for_test(1); - for wrong_generation in [generation - 1, generation + 1] { - controller - .handle_shards_update( - "ingester", - wrong_generation, - update.clone(), - &mut model, - &progress, - ) - .await - .unwrap(); - assert!( - model - .all_shards() - .all(|shard| shard.short_term_ingestion_rate == ByteSize::default()) - ); - } - controller - .handle_shards_update( - "ingester", - generation, - update.clone(), - &mut model, - &progress, - ) - .await - .unwrap(); - assert_eq!(model.all_shards().count(), 2); - assert!( - model - .all_shards() - .all(|shard| shard.short_term_ingestion_rate == ByteSize::b(123) - && shard.long_term_ingestion_rate == ByteSize::b(456)) - ); - - let (mut model, _) = shard_reports_for_test(1); - controller - .handle_shards_update( - "joining-ingester", - generation, - update, - &mut model, - &progress, - ) - .await - .unwrap(); - assert!( - model - .all_shards() - .all(|shard| shard.short_term_ingestion_rate == ByteSize::b(123)) - ); - } - - #[tokio::test] - async fn test_shards_update_continues_only_after_certainly_aborted_errors() { - for certainly_aborted in [true, false] { - let (mut model, update) = shard_reports_for_test(2); - let mut mock_metastore = MockMetastoreService::new(); - mock_metastore - .expect_open_shards() - .times(if certainly_aborted { 2 } else { 1 }) - .returning(move |_| { - if certainly_aborted { - Err(MetastoreError::InvalidArgument { - message: "aborted".to_string(), - }) - } else { - Err(MetastoreError::Connection { - message: "uncertain".to_string(), - }) - } - }); - let mut mock_ingester = MockIngesterService::new(); - mock_ingester - .expect_init_shards() - .times(if certainly_aborted { 2 } else { 1 }) - .returning(|request| { - Ok(InitShardsResponse { - successes: request - .subrequests - .into_iter() - .map(|request| InitShardSuccess { - subrequest_id: request.subrequest_id, - shard: request.shard, - }) - .collect(), - failures: Vec::new(), - }) - }); - let pool = IngesterPool::default(); - pool.insert( - NodeId::from_str("ingester"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock( - mock_ingester, - )), - ); - let mut controller = IngestController::new( - MetastoreServiceClient::from_mock(mock_metastore), - pool, - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.5, - ); - let result = controller - .handle_shards_update("ingester", 1, update, &mut model, &Progress::default()) - .await; - assert_eq!(result.is_ok(), certainly_aborted); - let second_source = model - .all_shards() - .find(|shard| shard.source_id == "source-b") - .unwrap(); - let expected = if certainly_aborted { - ByteSize::b(123) - } else { - ByteSize::default() - }; - assert_eq!(second_source.short_term_ingestion_rate, expected); - } - } - #[tokio::test] async fn test_ingest_controller_get_or_create_open_shards() { let source_id: &'static str = "test-source"; @@ -1769,12 +1378,7 @@ mod tests { IngesterPoolEntry::ready_with_client(ingester.clone()), ); - let mut controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let mut controller = IngestController::new(metastore, ingester_pool.clone()); let mut model = ControlPlaneModel::default(); model.add_index(index_metadata_0.clone()); @@ -1956,12 +1560,7 @@ mod tests { IngesterPoolEntry::ready_with_client(client.clone()), ); - let mut controller = IngestController::new( - metastore, - ingester_pool, - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let mut controller = IngestController::new(metastore, ingester_pool); let mut model = ControlPlaneModel::default(); model.add_index(index_metadata_0.clone()); @@ -1997,12 +1596,7 @@ mod tests { let metastore = MetastoreServiceClient::mocked(); let ingester_pool = IngesterPool::default(); - let mut controller = IngestController::new( - metastore, - ingester_pool, - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let mut controller = IngestController::new(metastore, ingester_pool); let mut model = ControlPlaneModel::default(); let index_uid = IndexUid::for_test("test-index-0", 0); @@ -2307,12 +1901,7 @@ mod tests { let metastore = MetastoreServiceClient::mocked(); let ingester_pool = IngesterPool::default(); - let controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let controller = IngestController::new(metastore, ingester_pool.clone()); let ingester_id_0 = NodeId::from_str("test-ingester-0"); let mut mock_ingester_0 = MockIngesterService::new(); @@ -2536,12 +2125,7 @@ mod tests { let metastore = MetastoreServiceClient::from_mock(mock_metastore); let ingester_pool = IngesterPool::default(); - let mut controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let mut controller = IngestController::new(metastore, ingester_pool.clone()); let index_uid = IndexUid::for_test("test-index", 0); let source_id = "test-source".to_string(); @@ -2698,8 +2282,6 @@ mod tests { let mut controller = IngestController::new( MetastoreServiceClient::from_mock(mock_metastore), ingester_pool, - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, ); let opened_by_zone = controller @@ -2730,849 +2312,91 @@ mod tests { } #[tokio::test] - async fn test_ingest_controller_update_source_shards() { - let mut mock_metastore = MockMetastoreService::new(); - mock_metastore - .expect_open_shards() - .once() - .returning(|request| { - assert_eq!(request.subrequests.len(), 1); - let subrequest = &request.subrequests[0]; + async fn test_sync_with_ingesters() { + let metastore = MetastoreServiceClient::mocked(); + let ingester_pool = IngesterPool::default(); - assert_eq!(subrequest.index_uid(), &IndexUid::for_test("test-index", 0)); - assert_eq!(subrequest.source_id, "test-source"); - assert_eq!(subrequest.ingester_id, "test-ingester"); + let controller = IngestController::new(metastore, ingester_pool.clone()); - Err(MetastoreError::InvalidArgument { - message: "failed to open shards".to_string(), - }) - }); - mock_metastore - .expect_open_shards() - .once() - .returning(|request| { - assert_eq!(request.subrequests.len(), 1); - let subrequest: &OpenShardSubrequest = &request.subrequests[0]; + let index_uid = IndexUid::for_test("test-index", 0); + let source_id: SourceId = "test-source".to_string(); + let mut model = ControlPlaneModel::default(); + let shards = vec![ + Shard { + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(1)), + shard_state: ShardState::Open as i32, + ingester_id: "node-1".to_string(), + ..Default::default() + }, + Shard { + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(2)), + shard_state: ShardState::Open as i32, + ingester_id: "node-2".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::Open as i32, + ingester_id: "node-2".to_string(), + ..Default::default() + }, + ]; + model.insert_shards(&index_uid, &source_id, shards); - assert_eq!(subrequest.index_uid(), &IndexUid::for_test("test-index", 0)); - assert_eq!(subrequest.source_id, "test-source"); - assert_eq!(subrequest.ingester_id, "test-ingester"); + let mut mock_ingester_1 = MockIngesterService::new(); + let mock_ingester_2 = MockIngesterService::new(); + let mock_ingester_3 = MockIngesterService::new(); - let shard = Shard { - index_uid: subrequest.index_uid.clone(), - source_id: subrequest.source_id.clone(), - shard_id: subrequest.shard_id.clone(), - shard_state: ShardState::Open as i32, - ingester_id: subrequest.ingester_id.clone(), - doc_mapping_uid: subrequest.doc_mapping_uid, - publish_position_inclusive: Some(Position::Beginning), - publish_token: None, - update_timestamp: 1724158996, - }; - let response = OpenShardsResponse { - subresponses: vec![OpenShardSubresponse { - subrequest_id: subrequest.subrequest_id, - open_shard: Some(shard), - }], - }; - Ok(response) + let count_calls = Arc::new(AtomicUsize::new(0)); + let count_calls_clone = count_calls.clone(); + mock_ingester_1 + .expect_retain_shards() + .once() + .returning(move |request| { + assert_eq!(request.retain_shards_for_sources.len(), 1); + assert_eq!( + request.retain_shards_for_sources[0].shard_ids, + [ShardId::from(1)] + ); + count_calls_clone.fetch_add(1, Ordering::Release); + Ok(RetainShardsResponse {}) }); - let metastore = MetastoreServiceClient::from_mock(mock_metastore); + ingester_pool.insert( + NodeId::from_str("node-1"), + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester_1)), + ); + ingester_pool.insert( + NodeId::from_str("node-2"), + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester_2)), + ); + ingester_pool.insert( + NodeId::from_str("node-3"), + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester_3)), + ); + let ingester_id = NodeId::from_str("node-1"); + let wait_handle = controller.sync_with_ingester(&ingester_id, &model); + wait_handle.wait().await; + assert_eq!(count_calls.load(Ordering::Acquire), 1); + } + + #[tokio::test] + async fn test_ingest_controller_advise_reset_shards() { + let metastore = MetastoreServiceClient::mocked(); let ingester_pool = IngesterPool::default(); - let mut controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let controller = IngestController::new(metastore, ingester_pool); + + let mut model = ControlPlaneModel::default(); let index_uid = IndexUid::for_test("test-index", 0); - let mut index_metadata = IndexMetadata::for_test("test-index", "ram://indexes/test-index"); - let source_id: SourceId = "test-source".to_string(); - index_metadata.sources.insert( - source_id.clone(), - SourceConfig::for_test(&source_id, quickwit_config::SourceParams::void()), - ); - - let source_uid = SourceUid { - index_uid: index_uid.clone(), - source_id: source_id.clone(), - }; - let mut model = ControlPlaneModel::default(); - model.add_index(index_metadata); - let progress = Progress::default(); - - let shards = vec![Shard { - index_uid: Some(index_uid.clone()), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(1)), - ingester_id: "test-ingester".to_string(), - shard_state: ShardState::Open as i32, - ..Default::default() - }]; - model.insert_shards(&index_uid, &source_id, shards); - let shard_entries: Vec = model.all_shards().cloned().collect(); - - assert_eq!(shard_entries.len(), 1); - assert_eq!(shard_entries[0].short_term_ingestion_rate, ByteSize::mib(0)); - - // Test update shard ingestion rate but no scale down because num open shards is 1. - let shard_infos = BTreeSet::from_iter([ShardInfo { - shard_id: ShardId::from(1), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(1), - long_term_ingestion_rate: ByteSize::mib(1), - }]); - - controller - .update_source_shards(source_uid.clone(), &shard_infos, &mut model, &progress) - .await - .unwrap(); - - let shard_entries: Vec = model.all_shards().cloned().collect(); - assert_eq!(shard_entries.len(), 1); - assert_eq!(shard_entries[0].short_term_ingestion_rate, ByteSize::mib(1)); - - // Test update shard ingestion rate with failing scale down. - let shards = vec![Shard { - index_uid: Some(index_uid.clone()), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(2)), - shard_state: ShardState::Open as i32, - ingester_id: "test-ingester".to_string(), - ..Default::default() - }]; - model.insert_shards(&index_uid, &source_id, shards); - - let shard_entries: Vec = model.all_shards().cloned().collect(); - assert_eq!(shard_entries.len(), 2); - - let mut mock_ingester = MockIngesterService::new(); - - let index_uid_clone = index_uid.clone(); - mock_ingester.expect_init_shards().returning( - move |init_shard_request: InitShardsRequest| { - assert_eq!(init_shard_request.subrequests.len(), 1); - let init_shard_subrequest: &InitShardSubrequest = - &init_shard_request.subrequests[0]; - assert!(init_shard_subrequest.validate_docs); - Ok(InitShardsResponse { - successes: vec![InitShardSuccess { - subrequest_id: init_shard_subrequest.subrequest_id, - shard: init_shard_subrequest.shard.clone(), - }], - failures: Vec::new(), - }) - }, - ); - mock_ingester - .expect_close_shards() - .returning(move |request| { - assert_eq!(request.shard_pkeys.len(), 1); - assert_eq!(request.shard_pkeys[0].index_uid(), &index_uid_clone); - assert_eq!(request.shard_pkeys[0].source_id, "test-source"); - Err(IngestV2Error::Internal( - "failed to close shards".to_string(), - )) - }); - let ingester = - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); - ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); - - let shard_infos = BTreeSet::from_iter([ - ShardInfo { - shard_id: ShardId::from(1), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(1), - long_term_ingestion_rate: ByteSize::mib(1), - }, - ShardInfo { - shard_id: ShardId::from(2), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(1), - long_term_ingestion_rate: ByteSize::mib(1), - }, - ]); - controller - .update_source_shards(source_uid.clone(), &shard_infos, &mut model, &progress) - .await - .unwrap(); - - // Test update shard ingestion rate with failing scale up. - let shard_infos = BTreeSet::from_iter([ - ShardInfo { - shard_id: ShardId::from(1), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(4), - long_term_ingestion_rate: ByteSize::mib(4), - }, - ShardInfo { - shard_id: ShardId::from(2), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(4), - long_term_ingestion_rate: ByteSize::mib(4), - }, - ]); - - // The first request fails due to an error on the metastore. - let MetastoreError::InvalidArgument { .. } = controller - .update_source_shards(source_uid.clone(), &shard_infos, &mut model, &progress) - .await - .unwrap_err() - else { - panic!(); - }; - - // The second request works! - controller - .update_source_shards(source_uid.clone(), &shard_infos, &mut model, &progress) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_ingest_controller_handle_shards_update_for_deleted_index() { - let metastore = MetastoreServiceClient::from_mock(MockMetastoreService::new()); - let ingester_pool = IngesterPool::default(); - - let mut controller = IngestController::new( - metastore, - ingester_pool, - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); - - let mut model = ControlPlaneModel::default(); - let source_shard_report = SourceShardReport { - source_uid: SourceUid { - index_uid: IndexUid::for_test("test-index", 0), - source_id: "test-source".to_string(), - }, - shard_infos: BTreeSet::new(), - }; - controller - .handle_shards_update( - "test-ingester", - 1, - ShardsUpdate { - shard_infos_by_source: vec![source_shard_report.into()], - }, - &mut model, - &Progress::default(), - ) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_ingest_controller_disable_validation_when_vrl() { - let mut mock_metastore = MockMetastoreService::new(); - mock_metastore - .expect_open_shards() - .once() - .returning(|request| { - let subrequest: &OpenShardSubrequest = &request.subrequests[0]; - let shard = Shard { - index_uid: subrequest.index_uid.clone(), - source_id: subrequest.source_id.clone(), - shard_id: subrequest.shard_id.clone(), - shard_state: ShardState::Open as i32, - ingester_id: subrequest.ingester_id.clone(), - doc_mapping_uid: subrequest.doc_mapping_uid, - publish_position_inclusive: Some(Position::Beginning), - publish_token: None, - update_timestamp: 1724158996, - }; - let response = OpenShardsResponse { - subresponses: vec![OpenShardSubresponse { - subrequest_id: subrequest.subrequest_id, - open_shard: Some(shard), - }], - }; - Ok(response) - }); - let metastore = MetastoreServiceClient::from_mock(mock_metastore); - let ingester_pool = IngesterPool::default(); - - let mut controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); - - let index_uid = IndexUid::for_test("test-index", 0); - let mut index_metadata = IndexMetadata::for_test("test-index", "ram://indexes/test-index"); - let source_id: SourceId = "test-source".to_string(); - let mut source_config = - SourceConfig::for_test(&source_id, quickwit_config::SourceParams::void()); - // set a vrl script - source_config.transform_config = - Some(quickwit_config::TransformConfig::new("".to_string(), None)); - index_metadata - .sources - .insert(source_id.clone(), source_config); - - let source_uid = SourceUid { - index_uid: index_uid.clone(), - source_id: source_id.clone(), - }; - let mut model = ControlPlaneModel::default(); - model.add_index(index_metadata); - let progress = Progress::default(); - - let shards = vec![Shard { - index_uid: Some(index_uid.clone()), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(1)), - ingester_id: "test-ingester".to_string(), - shard_state: ShardState::Open as i32, - ..Default::default() - }]; - model.insert_shards(&index_uid, &source_id, shards); - - let mut mock_ingester = MockIngesterService::new(); - - mock_ingester.expect_init_shards().returning( - move |init_shard_request: InitShardsRequest| { - assert_eq!(init_shard_request.subrequests.len(), 1); - let init_shard_subrequest: &InitShardSubrequest = - &init_shard_request.subrequests[0]; - // we have vrl, so no validation - assert!(!init_shard_subrequest.validate_docs); - Ok(InitShardsResponse { - successes: vec![InitShardSuccess { - subrequest_id: init_shard_subrequest.subrequest_id, - shard: init_shard_subrequest.shard.clone(), - }], - failures: Vec::new(), - }) - }, - ); - - let ingester = - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); - ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); - - let shard_infos = BTreeSet::from_iter([ShardInfo { - shard_id: ShardId::from(1), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(4), - long_term_ingestion_rate: ByteSize::mib(4), - }]); - - controller - .update_source_shards(source_uid.clone(), &shard_infos, &mut model, &progress) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_ingest_controller_try_scale_up_shards() { - let mut mock_metastore = MockMetastoreService::new(); - - let index_uid = IndexUid::from_str("test-index:00000000000000000000000000").unwrap(); - let index_uid_clone = index_uid.clone(); - mock_metastore - .expect_open_shards() - .once() - .returning(move |request| { - assert_eq!(request.subrequests.len(), 1); - assert_eq!(request.subrequests[0].index_uid(), &index_uid_clone); - assert_eq!(request.subrequests[0].source_id, INGEST_V2_SOURCE_ID); - assert_eq!(request.subrequests[0].ingester_id, "test-ingester"); - - Err(MetastoreError::InvalidArgument { - message: "failed to open shards".to_string(), - }) - }); - let index_uid_clone = index_uid.clone(); - mock_metastore - .expect_open_shards() - .returning(move |request| { - assert_eq!(request.subrequests.len(), 1); - assert_eq!(request.subrequests[0].index_uid(), &index_uid_clone); - assert_eq!(request.subrequests[0].source_id, INGEST_V2_SOURCE_ID); - assert_eq!(request.subrequests[0].ingester_id, "test-ingester"); - - let subresponses = vec![metastore::OpenShardSubresponse { - subrequest_id: 0, - open_shard: Some(Shard { - index_uid: Some(index_uid.clone()), - source_id: INGEST_V2_SOURCE_ID.to_string(), - shard_id: Some(ShardId::from(1)), - ingester_id: "test-ingester".to_string(), - shard_state: ShardState::Open as i32, - ..Default::default() - }), - }]; - let response = metastore::OpenShardsResponse { subresponses }; - Ok(response) - }); - let metastore = MetastoreServiceClient::from_mock(mock_metastore); - - let ingester_pool = IngesterPool::default(); - - let mut controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); - - let index_uid = IndexUid::for_test("test-index", 0); - let source_id: SourceId = INGEST_V2_SOURCE_ID.to_string(); - - let source_uid = SourceUid { - index_uid: index_uid.clone(), - source_id: source_id.clone(), - }; - let shard_stats = ShardStats { - num_open_shards: 2, - ..Default::default() - }; - let mut model = ControlPlaneModel::default(); - let index_metadata = - IndexMetadata::for_test(&index_uid.index_id, "ram://indexes/test-index:0"); - model.add_index(index_metadata); - - let source_config = SourceConfig::ingest_v2(); - model.add_source(&index_uid, source_config).unwrap(); - - let progress = Progress::default(); - - // Test could not find ingester because no ingester in pool - controller - .try_scale_up_shards(source_uid.clone(), shard_stats, &mut model, &progress, 1) - .await - .unwrap(); - - let mut mock_ingester = MockIngesterService::new(); - - let index_uid_clone = index_uid.clone(); - mock_ingester - .expect_init_shards() - .once() - .returning(move |request| { - assert_eq!(request.subrequests.len(), 1); - - let subrequest = &request.subrequests[0]; - assert_eq!(subrequest.subrequest_id, 0); - - let shard = request.subrequests[0].shard(); - assert_eq!(shard.index_uid(), &index_uid_clone); - assert_eq!(shard.source_id, INGEST_V2_SOURCE_ID); - assert_eq!(shard.ingester_id, "test-ingester"); - - Err(IngestV2Error::Internal("failed to init shards".to_string())) - }); - let index_uid_clone = index_uid.clone(); - mock_ingester - .expect_init_shards() - .returning(move |request| { - assert_eq!(request.subrequests.len(), 1); - - let subrequest = &request.subrequests[0]; - assert_eq!(subrequest.subrequest_id, 0); - - let shard = subrequest.shard(); - assert_eq!(shard.index_uid(), &index_uid_clone); - assert_eq!(shard.source_id, INGEST_V2_SOURCE_ID); - assert_eq!(shard.ingester_id, "test-ingester"); - - let successes = vec![InitShardSuccess { - subrequest_id: request.subrequests[0].subrequest_id, - shard: Some(shard.clone()), - }]; - let response = InitShardsResponse { - successes, - failures: Vec::new(), - }; - Ok(response) - }); - let ingester = - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); - ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); - - // Test failed to open shards. - controller - .try_scale_up_shards(source_uid.clone(), shard_stats, &mut model, &progress, 1) - .await - .unwrap(); - assert_eq!(model.all_shards().count(), 0); - - // Test failed to init shards. - controller - .try_scale_up_shards(source_uid.clone(), shard_stats, &mut model, &progress, 1) - .await - .unwrap_err(); - assert_eq!(model.all_shards().count(), 0); - - // Test successfully opened shard. - controller - .try_scale_up_shards(source_uid.clone(), shard_stats, &mut model, &progress, 1) - .await - .unwrap(); - assert_eq!( - model.all_shards().filter(|shard| shard.is_open()).count(), - 1 - ); - } - - #[tokio::test] - async fn test_ingest_controller_try_scale_down_shards() { - let metastore = MetastoreServiceClient::mocked(); - let ingester_pool = IngesterPool::default(); - - let controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); - - let index_uid = IndexUid::for_test("test-index", 0); - let source_id: SourceId = "test-source".to_string(); - - let source_uid = SourceUid { - index_uid: index_uid.clone(), - source_id: source_id.clone(), - }; - let shard_stats = ShardStats { - num_open_shards: 2, - ..Default::default() - }; - let min_shards = NonZeroUsize::MIN; - let mut model = ControlPlaneModel::default(); - let progress = Progress::default(); - - // Test could not find a scale down candidate. - controller - .try_scale_down_shards( - source_uid.clone(), - shard_stats, - min_shards, - &mut model, - &progress, - ) - .await - .unwrap(); - - let shards = vec![Shard { - shard_id: Some(ShardId::from(1)), - index_uid: Some(index_uid.clone()), - source_id: source_id.clone(), - ingester_id: "test-ingester".to_string(), - shard_state: ShardState::Open as i32, - ..Default::default() - }]; - model.insert_shards(&index_uid, &source_id, shards); - - // Test ingester is unavailable. - controller - .try_scale_down_shards( - source_uid.clone(), - shard_stats, - min_shards, - &mut model, - &progress, - ) - .await - .unwrap(); - - let mut mock_ingester = MockIngesterService::new(); - - let index_uid_clone = index_uid.clone(); - mock_ingester - .expect_close_shards() - .once() - .returning(move |request| { - assert_eq!(request.shard_pkeys.len(), 1); - assert_eq!(request.shard_pkeys[0].index_uid(), &index_uid_clone); - assert_eq!(request.shard_pkeys[0].source_id, "test-source"); - assert_eq!(request.shard_pkeys[0].shard_id(), ShardId::from(1)); - - Err(IngestV2Error::Internal( - "failed to close shards".to_string(), - )) - }); - let index_uid_clone = index_uid.clone(); - mock_ingester - .expect_close_shards() - .once() - .returning(move |request| { - assert_eq!(request.shard_pkeys.len(), 1); - assert_eq!(request.shard_pkeys[0].index_uid(), &index_uid_clone); - assert_eq!(request.shard_pkeys[0].source_id, "test-source"); - assert_eq!(request.shard_pkeys[0].shard_id(), ShardId::from(1)); - - let response = CloseShardsResponse { - successes: request.shard_pkeys, - }; - Ok(response) - }); - let ingester = - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); - ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); - - // Test failed to close shard. - controller - .try_scale_down_shards( - source_uid.clone(), - shard_stats, - min_shards, - &mut model, - &progress, - ) - .await - .unwrap(); - assert!(model.all_shards().all(|shard| shard.is_open())); - - // Test successfully closed shard. - controller - .try_scale_down_shards( - source_uid.clone(), - shard_stats, - min_shards, - &mut model, - &progress, - ) - .await - .unwrap(); - assert!(model.all_shards().all(|shard| shard.is_closed())); - - let shards = vec![Shard { - shard_id: Some(ShardId::from(2)), - index_uid: Some(index_uid.clone()), - source_id: source_id.clone(), - ingester_id: "test-ingester".to_string(), - shard_state: ShardState::Open as i32, - ..Default::default() - }]; - model.insert_shards(&index_uid, &source_id, shards); - - // Test rate limited. - controller - .try_scale_down_shards( - source_uid.clone(), - shard_stats, - min_shards, - &mut model, - &progress, - ) - .await - .unwrap(); - assert!(model.all_shards().any(|shard| shard.is_open())); - } - - #[test] - fn test_find_scale_down_candidate() { - let index_uid = IndexUid::for_test("test-index", 0); - let source_id: SourceId = "test-source".to_string(); - - let source_uid = SourceUid { - index_uid: index_uid.clone(), - source_id: source_id.clone(), - }; - let mut model = ControlPlaneModel::default(); - - assert!(find_scale_down_candidate(&source_uid, &model).is_none()); - - let shards = vec![ - Shard { - index_uid: index_uid.clone().into(), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(1)), - shard_state: ShardState::Open as i32, - ingester_id: "test-ingester-0".to_string(), - ..Default::default() - }, - Shard { - index_uid: index_uid.clone().into(), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(2)), - shard_state: ShardState::Open as i32, - ingester_id: "test-ingester-0".to_string(), - ..Default::default() - }, - Shard { - index_uid: index_uid.clone().into(), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(3)), - shard_state: ShardState::Closed as i32, //< this one is closed - ingester_id: "test-ingester-0".to_string(), - ..Default::default() - }, - Shard { - index_uid: index_uid.clone().into(), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(4)), - shard_state: ShardState::Open as i32, - ingester_id: "test-ingester-1".to_string(), - ..Default::default() - }, - Shard { - index_uid: index_uid.clone().into(), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(5)), - shard_state: ShardState::Open as i32, - ingester_id: "test-ingester-1".to_string(), - ..Default::default() - }, - Shard { - index_uid: index_uid.clone().into(), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(6)), - shard_state: ShardState::Open as i32, - ingester_id: "test-ingester-1".to_string(), - ..Default::default() - }, - ]; - // That's 3 open shards on indexer-1, 2 open shard and one closed shard on indexer-0.. - model.insert_shards(&index_uid, &source_id, shards); - - let shard_infos = BTreeSet::from_iter([ - ShardInfo { - shard_id: ShardId::from(1), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(1), - long_term_ingestion_rate: ByteSize::mib(1), - }, - ShardInfo { - shard_id: ShardId::from(2), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(2), - long_term_ingestion_rate: ByteSize::mib(2), - }, - ShardInfo { - shard_id: ShardId::from(3), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(3), - long_term_ingestion_rate: ByteSize::mib(3), - }, - ShardInfo { - shard_id: ShardId::from(4), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(4), - long_term_ingestion_rate: ByteSize::mib(4), - }, - ShardInfo { - shard_id: ShardId::from(5), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(5), - long_term_ingestion_rate: ByteSize::mib(5), - }, - ShardInfo { - shard_id: ShardId::from(6), - shard_state: ShardState::Open, - short_term_ingestion_rate: ByteSize::mib(6), - long_term_ingestion_rate: ByteSize::mib(6), - }, - ]); - model.update_shards(&source_uid, &shard_infos); - - let shard = find_scale_down_candidate(&source_uid, &model).unwrap(); - // We pick ingester 1 has it has more open shard - assert_eq!(shard.ingester_id, "test-ingester-1"); - } - - #[tokio::test] - async fn test_sync_with_ingesters() { - let metastore = MetastoreServiceClient::mocked(); - let ingester_pool = IngesterPool::default(); - - let controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); - - let index_uid = IndexUid::for_test("test-index", 0); - let source_id: SourceId = "test-source".to_string(); - let mut model = ControlPlaneModel::default(); - let shards = vec![ - Shard { - index_uid: Some(index_uid.clone()), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(1)), - shard_state: ShardState::Open as i32, - ingester_id: "node-1".to_string(), - ..Default::default() - }, - Shard { - index_uid: Some(index_uid.clone()), - source_id: source_id.clone(), - shard_id: Some(ShardId::from(2)), - shard_state: ShardState::Open as i32, - ingester_id: "node-2".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::Open as i32, - ingester_id: "node-2".to_string(), - ..Default::default() - }, - ]; - model.insert_shards(&index_uid, &source_id, shards); - - let mut mock_ingester_1 = MockIngesterService::new(); - let mock_ingester_2 = MockIngesterService::new(); - let mock_ingester_3 = MockIngesterService::new(); - - let count_calls = Arc::new(AtomicUsize::new(0)); - let count_calls_clone = count_calls.clone(); - mock_ingester_1 - .expect_retain_shards() - .once() - .returning(move |request| { - assert_eq!(request.retain_shards_for_sources.len(), 1); - assert_eq!( - request.retain_shards_for_sources[0].shard_ids, - [ShardId::from(1)] - ); - count_calls_clone.fetch_add(1, Ordering::Release); - Ok(RetainShardsResponse {}) - }); - ingester_pool.insert( - NodeId::from_str("node-1"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester_1)), - ); - ingester_pool.insert( - NodeId::from_str("node-2"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester_2)), - ); - ingester_pool.insert( - NodeId::from_str("node-3"), - IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester_3)), - ); - let ingester_id = NodeId::from_str("node-1"); - let wait_handle = controller.sync_with_ingester(&ingester_id, &model); - wait_handle.wait().await; - assert_eq!(count_calls.load(Ordering::Acquire), 1); - } - - #[tokio::test] - async fn test_ingest_controller_advise_reset_shards() { - let metastore = MetastoreServiceClient::mocked(); - let ingester_pool = IngesterPool::default(); - - let controller = IngestController::new( - metastore, - ingester_pool, - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); - - let mut model = ControlPlaneModel::default(); - - let index_uid = IndexUid::for_test("test-index", 0); - let source_id_00: SourceId = "test-source-0".into(); - let source_id_01: SourceId = "test-source-1".into(); + let source_id_00: SourceId = "test-source-0".into(); + let source_id_01: SourceId = "test-source-1".into(); let shards = vec![Shard { index_uid: Some(index_uid.clone()), @@ -3636,12 +2460,7 @@ mod tests { async fn test_ingest_controller_close_shards() { let metastore = MetastoreServiceClient::mocked(); let ingester_pool = IngesterPool::default(); - let controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let controller = IngestController::new(metastore, ingester_pool.clone()); let closed_shards = controller.close_shards(Vec::new()).await; assert_eq!(closed_shards.len(), 0); @@ -3794,12 +2613,7 @@ mod tests { }); let metastore = MetastoreServiceClient::from_mock(mock_metastore); let ingester_pool = IngesterPool::default(); - let mut controller = IngestController::new( - metastore, - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let mut controller = IngestController::new(metastore, ingester_pool.clone()); let mut model = ControlPlaneModel::default(); @@ -4272,12 +3086,8 @@ mod tests { model.insert_shards(&index_uid, &source_id, shards.clone()); - let controller = IngestController::new( - MetastoreServiceClient::mocked(), - ingester_pool.clone(), - DEFAULT_SHARD_THROUGHPUT_LIMIT, - 1.001, - ); + let controller = + IngestController::new(MetastoreServiceClient::mocked(), ingester_pool.clone()); let shards_to_rebalance = controller.compute_shards_to_rebalance(&model); // All shards on retiring ingesters must be rebalanced. diff --git a/quickwit/quickwit-control-plane/src/ingest/scaling_arbiter.rs b/quickwit/quickwit-control-plane/src/ingest/legacy_scaling_arbiter.rs similarity index 96% rename from quickwit/quickwit-control-plane/src/ingest/scaling_arbiter.rs rename to quickwit/quickwit-control-plane/src/ingest/legacy_scaling_arbiter.rs index 9eaac27f201..b8774390f34 100644 --- a/quickwit/quickwit-control-plane/src/ingest/scaling_arbiter.rs +++ b/quickwit/quickwit-control-plane/src/ingest/legacy_scaling_arbiter.rs @@ -18,7 +18,7 @@ use bytesize::ByteSize; use crate::model::{ScalingMode, ShardStats}; -pub(crate) struct ScalingArbiter { +pub(crate) struct LegacyScalingArbiter { // Threshold in B/s below which we decrease the number of shards. scale_down_shards_threshold: ByteSize, @@ -38,13 +38,13 @@ pub(crate) struct ScalingArbiter { shard_scale_up_factor: f32, } -impl ScalingArbiter { +impl LegacyScalingArbiter { pub fn with_max_shard_ingestion_throughput( max_shard_throughput: ByteSize, shard_scale_up_factor: f32, - ) -> ScalingArbiter { + ) -> LegacyScalingArbiter { let max_shard_throughput_bytes = max_shard_throughput.as_u64(); - ScalingArbiter { + LegacyScalingArbiter { scale_up_shards_short_term_threshold: ByteSize::b(max_shard_throughput_bytes * 8 / 10), scale_up_shards_long_term_threshold: ByteSize::b(max_shard_throughput_bytes * 3 / 10), scale_down_shards_threshold: ByteSize::b(max_shard_throughput_bytes * 2 / 10), @@ -120,7 +120,7 @@ mod tests { use bytesize::ByteSize; - use super::ScalingArbiter; + use super::LegacyScalingArbiter; use crate::model::{ScalingMode, ShardStats}; #[test] @@ -128,7 +128,7 @@ mod tests { // use shard throughput 10MiB to simplify calculations // with a factor close to 1 shards are effectively added 1 by 1 let scaling_arbiter = - ScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.01); + LegacyScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.01); assert_eq!( scaling_arbiter.should_scale( ShardStats { @@ -219,7 +219,7 @@ mod tests { fn test_scaling_arbiter_2x() { // use shard throughput 10MiB to simplify calculations let scaling_arbiter = - ScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 2.); + LegacyScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 2.); assert_eq!( scaling_arbiter.should_scale( ShardStats { @@ -323,7 +323,7 @@ mod tests { fn test_scale_up_computations() { // use shard throughput 10MiB to simplify calculations let scaling_arbiter = - ScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.5); + LegacyScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.5); let shard_stats = ShardStats { num_open_shards: 0, @@ -404,7 +404,7 @@ mod tests { #[test] fn test_scaling_arbiter_idle() { let scaling_arbiter = - ScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.5); + LegacyScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.5); let shard_stats = ShardStats { num_open_shards: 0, @@ -430,7 +430,7 @@ mod tests { #[test] fn test_scaling_arbiter_min_shards() { let scaling_arbiter = - ScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.5); + LegacyScalingArbiter::with_max_shard_ingestion_throughput(ByteSize::mib(10), 1.5); let shard_stats = ShardStats { num_open_shards: 1, diff --git a/quickwit/quickwit-control-plane/src/ingest/legacy_scaling_controller.rs b/quickwit/quickwit-control-plane/src/ingest/legacy_scaling_controller.rs new file mode 100644 index 00000000000..5dc9c704f7d --- /dev/null +++ b/quickwit/quickwit-control-plane/src/ingest/legacy_scaling_controller.rs @@ -0,0 +1,1287 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::num::NonZeroUsize; + +use bytesize::ByteSize; +use quickwit_common::Progress; +use quickwit_ingest::{ShardInfos, SourceShardReport}; +use quickwit_proto::control_plane::ShardsUpdate; +use quickwit_proto::ingest::Shard; +use quickwit_proto::metastore::MetastoreResult; +use quickwit_proto::types::{NodeId, SourceUid}; +use rand::prelude::IndexedRandom; +use rand::{Rng, rng}; +use tracing::{error, info, warn}; + +use super::legacy_scaling_arbiter::LegacyScalingArbiter; +use crate::ingest::IngestController; +use crate::model::{ControlPlaneModel, ScalingMode, ShardEntry, ShardStats}; + +pub(crate) struct LegacyScalingController { + legacy_scaling_arbiter: LegacyScalingArbiter, +} + +impl LegacyScalingController { + pub fn new(max_shard_ingestion_throughput: ByteSize, shard_scale_up_factor: f32) -> Self { + LegacyScalingController { + legacy_scaling_arbiter: LegacyScalingArbiter::with_max_shard_ingestion_throughput( + max_shard_ingestion_throughput, + shard_scale_up_factor, + ), + } + } + + pub(crate) async fn handle_shards_update( + &self, + ingest_controller: &mut IngestController, + node_id: &str, + generation_id: u64, + shards_update: ShardsUpdate, + model: &mut ControlPlaneModel, + progress: &Progress, + ) -> MetastoreResult<()> { + if let Some(ingester) = ingest_controller.ingester_pool.get(node_id) + && generation_id != ingester.generation_id.as_u64() + { + return Ok(()); + } + for source_shard_infos in &shards_update.shard_infos_by_source { + let SourceShardReport { + source_uid, + shard_infos, + } = source_shard_infos.into(); + + if let Err(metastore_error) = self + .update_local_shards(ingest_controller, source_uid, &shard_infos, model, progress) + .await + { + if !metastore_error.is_transaction_certainly_aborted() { + return Err(metastore_error); + } + error!(error=?metastore_error, "failed to update source shards"); + } + } + Ok(()) + } + + pub(crate) async fn update_local_shards( + &self, + // TODO: hold the ingest controller on the struct instead of passing it in. + ingest_controller: &mut IngestController, + source_uid: SourceUid, + shard_infos: &ShardInfos, + model: &mut ControlPlaneModel, + progress: &Progress, + ) -> MetastoreResult<()> { + model.update_shards(&source_uid, shard_infos); + self.scale_source_shards(ingest_controller, source_uid, model, progress) + .await + } + + async fn scale_source_shards( + &self, + // TODO: hold the ingest controller on the struct instead of passing it in. + ingest_controller: &mut IngestController, + source_uid: SourceUid, + model: &mut ControlPlaneModel, + progress: &Progress, + ) -> MetastoreResult<()> { + let Some(shard_stats) = model.legacy_shard_stats(&source_uid) else { + return Ok(()); + }; + let Some(min_shards) = model + .index_metadata(&source_uid.index_uid) + .map(|index_metadata| index_metadata.index_config.ingest_settings.min_shards) + else { + warn!( + index_uid=%source_uid.index_uid, + "ignoring local shards update for a deleted index" + ); + return Ok(()); + }; + + let Some(scaling_mode) = self + .legacy_scaling_arbiter + .should_scale(shard_stats, min_shards) + else { + return Ok(()); + }; + match scaling_mode { + ScalingMode::Up(num_shards) => { + self.try_scale_up_shards( + ingest_controller, + source_uid, + shard_stats, + model, + progress, + num_shards, + ) + .await?; + } + ScalingMode::Down => { + self.try_scale_down_shards( + ingest_controller, + source_uid, + shard_stats, + min_shards, + model, + progress, + ) + .await?; + } + } + + Ok(()) + } + + /// Attempts to increase the number of shards. This operation is rate limited to avoid creating + /// to many shards in a short period of time. As a result, this method may not create any + /// shard. + async fn try_scale_up_shards( + &self, + // TODO: hold the ingest controller on the struct instead of passing it in. + ingest_controller: &mut IngestController, + source_uid: SourceUid, + shard_stats: ShardStats, + model: &mut ControlPlaneModel, + progress: &Progress, + num_shards_to_open: usize, + ) -> MetastoreResult<()> { + if !model + .acquire_scaling_permits(&source_uid, ScalingMode::Up(num_shards_to_open)) + .unwrap_or(false) + { + return Ok(()); + } + let new_num_open_shards = shard_stats.num_open_shards + num_shards_to_open; + let open_shards_result = ingest_controller + .open_shards_for_source(&source_uid, num_shards_to_open, model, progress) + .await; + + match open_shards_result { + Ok(num_opened_shards) => { + if num_opened_shards == 0 { + // 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)); + warn!( + index_uid=%source_uid.index_uid, + source_id=%source_uid.source_id, + "scaling up number of shards to {new_num_open_shards} failed: shard initialization failure" + ); + } else { + info!( + index_id=%source_uid.index_uid.index_id, + source_id=%source_uid.source_id, + "successfully scaled up number of shards to {new_num_open_shards}" + ); + } + Ok(()) + } + Err(metastore_error) => { + // We did not manage to create the shard. + // We can release our permit, but we also need to return the error to the caller, in + // order to restart the control plane actor if necessary. + warn!( + index_id=%source_uid.index_uid.index_id, + source_id=%source_uid.source_id, + "scaling up number of shards to {new_num_open_shards} failed: {metastore_error:?}" + ); + model.release_scaling_permits(&source_uid, ScalingMode::Up(num_shards_to_open)); + Err(metastore_error) + } + } + } + + /// Attempts to decrease the number of shards. This operation is rate limited to avoid closing + /// shards too aggressively. As a result, this method may not close any shard. + async fn try_scale_down_shards( + &self, + // TODO: hold the ingest controller on the struct instead of passing it in. + ingest_controller: &IngestController, + source_uid: SourceUid, + shard_stats: ShardStats, + min_shards: NonZeroUsize, + model: &mut ControlPlaneModel, + progress: &Progress, + ) -> MetastoreResult<()> { + // The scaling arbiter should not suggest scaling down if the number of shards is already + // below the minimum, but we're just being defensive here. + if shard_stats.num_open_shards <= min_shards.get() { + return Ok(()); + } + if !model + .acquire_scaling_permits(&source_uid, ScalingMode::Down) + .unwrap_or(false) + { + return Ok(()); + } + let new_num_open_shards = shard_stats.num_open_shards - 1; + + info!( + index_id=%source_uid.index_uid.index_id, + source_id=%source_uid.source_id, + "scaling down number of shards to {new_num_open_shards}" + ); + let Some(shard) = find_scale_down_candidate(&source_uid, model) else { + model.release_scaling_permits(&source_uid, ScalingMode::Down); + return Ok(()); + }; + info!( + "scaling down shard {} from {}", + shard.shard_id(), + shard.ingester_id + ); + let closed_shard_ids = ingest_controller + .close_source_shards(&source_uid, vec![shard], model, progress) + .await; + + if closed_shard_ids.is_empty() { + warn!("failed to scale down number of shards"); + model.release_scaling_permits(&source_uid, ScalingMode::Down); + } + Ok(()) + } +} + +/// Finds a shard on the ingester with the highest number of open +/// shards for this source. +/// +/// If multiple shards are hosted on that ingester, the shard with the lowest (oldest) +/// shard ID is chosen. +fn find_scale_down_candidate(source_uid: &SourceUid, model: &ControlPlaneModel) -> Option { + let mut shard_entries_by_ingester_id: HashMap> = HashMap::new(); + let mut rng = rng(); + + for shard in model.get_shards_for_source(source_uid)?.values() { + if shard.is_open() { + shard_entries_by_ingester_id + .entry(NodeId::from_str(&shard.ingester_id)) + .or_default() + .push(shard); + } + } + shard_entries_by_ingester_id + .into_iter() + // We use a random number to break ties... The HashMap is randomly seeded so this is + // should not make much difference, but we might want to be as explicit as possible. + .max_by_key(|(_ingester_id, shard_entries)| (shard_entries.len(), rng.next_u32())) + .map(|(_ingester_id, shard_entries)| shard_entries.choose(&mut rng).unwrap().shard.clone()) +} +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + use std::str::FromStr; + + use bytesize::ByteSize; + use quickwit_common::Progress; + use quickwit_common::shared_consts::DEFAULT_SHARD_THROUGHPUT_LIMIT; + use quickwit_config::{INGEST_V2_SOURCE_ID, SourceConfig}; + use quickwit_ingest::{IngesterPool, IngesterPoolEntry, ShardInfo, SourceShardReport}; + use quickwit_metastore::IndexMetadata; + use quickwit_proto::control_plane::ShardsUpdate; + use quickwit_proto::ingest::ingester::{ + CloseShardsResponse, IngesterServiceClient, InitShardSubrequest, InitShardSuccess, + InitShardsRequest, InitShardsResponse, MockIngesterService, + }; + use quickwit_proto::ingest::{IngestV2Error, Shard, ShardState}; + use quickwit_proto::metastore; + use quickwit_proto::metastore::{ + MetastoreError, MetastoreServiceClient, MockMetastoreService, OpenShardSubrequest, + OpenShardSubresponse, OpenShardsResponse, + }; + use quickwit_proto::types::{IndexUid, NodeId, Position, ShardId, SourceId, SourceUid}; + + use super::*; + + #[tokio::test] + async fn test_handle_shards_update_for_deleted_index() { + let metastore = MetastoreServiceClient::from_mock(MockMetastoreService::new()); + let ingester_pool = IngesterPool::default(); + let mut controller = IngestController::new(metastore, ingester_pool); + let scaling_controller = + LegacyScalingController::new(DEFAULT_SHARD_THROUGHPUT_LIMIT, 1.001); + let mut model = ControlPlaneModel::default(); + let source_shard_report = SourceShardReport { + source_uid: SourceUid { + index_uid: IndexUid::for_test("test-index", 0), + source_id: "test-source".to_string(), + }, + shard_infos: BTreeSet::new(), + }; + scaling_controller + .handle_shards_update( + &mut controller, + "test-ingester", + 1, + ShardsUpdate { + shard_infos_by_source: vec![source_shard_report.into()], + }, + &mut model, + &Progress::default(), + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_ingest_controller_update_source_shards() { + let mut mock_metastore = MockMetastoreService::new(); + mock_metastore + .expect_open_shards() + .once() + .returning(|request| { + assert_eq!(request.subrequests.len(), 1); + let subrequest = &request.subrequests[0]; + + assert_eq!(subrequest.index_uid(), &IndexUid::for_test("test-index", 0)); + assert_eq!(subrequest.source_id, "test-source"); + assert_eq!(subrequest.ingester_id, "test-ingester"); + + Err(MetastoreError::InvalidArgument { + message: "failed to open shards".to_string(), + }) + }); + mock_metastore + .expect_open_shards() + .once() + .returning(|request| { + assert_eq!(request.subrequests.len(), 1); + let subrequest: &OpenShardSubrequest = &request.subrequests[0]; + + assert_eq!(subrequest.index_uid(), &IndexUid::for_test("test-index", 0)); + assert_eq!(subrequest.source_id, "test-source"); + assert_eq!(subrequest.ingester_id, "test-ingester"); + + let shard = Shard { + index_uid: subrequest.index_uid.clone(), + source_id: subrequest.source_id.clone(), + shard_id: subrequest.shard_id.clone(), + shard_state: ShardState::Open as i32, + ingester_id: subrequest.ingester_id.clone(), + doc_mapping_uid: subrequest.doc_mapping_uid, + publish_position_inclusive: Some(Position::Beginning), + publish_token: None, + update_timestamp: 1724158996, + }; + let response = OpenShardsResponse { + subresponses: vec![OpenShardSubresponse { + subrequest_id: subrequest.subrequest_id, + open_shard: Some(shard), + }], + }; + Ok(response) + }); + let metastore = MetastoreServiceClient::from_mock(mock_metastore); + let ingester_pool = IngesterPool::default(); + + let mut controller = IngestController::new(metastore, ingester_pool.clone()); + let scaling_controller = + LegacyScalingController::new(DEFAULT_SHARD_THROUGHPUT_LIMIT, 1.001); + + let index_uid = IndexUid::for_test("test-index", 0); + let mut index_metadata = IndexMetadata::for_test("test-index", "ram://indexes/test-index"); + let source_id: SourceId = "test-source".to_string(); + index_metadata.sources.insert( + source_id.clone(), + SourceConfig::for_test(&source_id, quickwit_config::SourceParams::void()), + ); + + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: source_id.clone(), + }; + let mut model = ControlPlaneModel::default(); + model.add_index(index_metadata); + let progress = Progress::default(); + + let shards = vec![Shard { + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(1)), + ingester_id: "test-ingester".to_string(), + shard_state: ShardState::Open as i32, + ..Default::default() + }]; + model.insert_shards(&index_uid, &source_id, shards); + let shard_entries: Vec = model.all_shards().cloned().collect(); + + assert_eq!(shard_entries.len(), 1); + assert_eq!(shard_entries[0].short_term_ingestion_rate, ByteSize::mib(0)); + + // Test update shard ingestion rate but no scale down because num open shards is 1. + let shard_infos = BTreeSet::from_iter([ShardInfo { + shard_id: ShardId::from(1), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(1), + long_term_ingestion_rate: ByteSize::mib(1), + }]); + + scaling_controller + .update_local_shards( + &mut controller, + source_uid.clone(), + &shard_infos, + &mut model, + &progress, + ) + .await + .unwrap(); + + let shard_entries: Vec = model.all_shards().cloned().collect(); + assert_eq!(shard_entries.len(), 1); + assert_eq!(shard_entries[0].short_term_ingestion_rate, ByteSize::mib(1)); + + // Test update shard ingestion rate with failing scale down. + let shards = vec![Shard { + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(2)), + shard_state: ShardState::Open as i32, + ingester_id: "test-ingester".to_string(), + ..Default::default() + }]; + model.insert_shards(&index_uid, &source_id, shards); + + let shard_entries: Vec = model.all_shards().cloned().collect(); + assert_eq!(shard_entries.len(), 2); + + let mut mock_ingester = MockIngesterService::new(); + + let index_uid_clone = index_uid.clone(); + mock_ingester.expect_init_shards().returning( + move |init_shard_request: InitShardsRequest| { + assert_eq!(init_shard_request.subrequests.len(), 1); + let init_shard_subrequest: &InitShardSubrequest = + &init_shard_request.subrequests[0]; + assert!(init_shard_subrequest.validate_docs); + Ok(InitShardsResponse { + successes: vec![InitShardSuccess { + subrequest_id: init_shard_subrequest.subrequest_id, + shard: init_shard_subrequest.shard.clone(), + }], + failures: Vec::new(), + }) + }, + ); + mock_ingester + .expect_close_shards() + .returning(move |request| { + assert_eq!(request.shard_pkeys.len(), 1); + assert_eq!(request.shard_pkeys[0].index_uid(), &index_uid_clone); + assert_eq!(request.shard_pkeys[0].source_id, "test-source"); + Err(IngestV2Error::Internal( + "failed to close shards".to_string(), + )) + }); + let ingester = + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); + ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); + + let shard_infos = BTreeSet::from_iter([ + ShardInfo { + shard_id: ShardId::from(1), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(1), + long_term_ingestion_rate: ByteSize::mib(1), + }, + ShardInfo { + shard_id: ShardId::from(2), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(1), + long_term_ingestion_rate: ByteSize::mib(1), + }, + ]); + scaling_controller + .update_local_shards( + &mut controller, + source_uid.clone(), + &shard_infos, + &mut model, + &progress, + ) + .await + .unwrap(); + + // Test update shard ingestion rate with failing scale up. + let shard_infos = BTreeSet::from_iter([ + ShardInfo { + shard_id: ShardId::from(1), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(4), + long_term_ingestion_rate: ByteSize::mib(4), + }, + ShardInfo { + shard_id: ShardId::from(2), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(4), + long_term_ingestion_rate: ByteSize::mib(4), + }, + ]); + + // The first request fails due to an error on the metastore. + let MetastoreError::InvalidArgument { .. } = scaling_controller + .update_local_shards( + &mut controller, + source_uid.clone(), + &shard_infos, + &mut model, + &progress, + ) + .await + .unwrap_err() + else { + panic!(); + }; + + // The second request works! + scaling_controller + .update_local_shards( + &mut controller, + source_uid.clone(), + &shard_infos, + &mut model, + &progress, + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_ingest_controller_disable_validation_when_vrl() { + let mut mock_metastore = MockMetastoreService::new(); + mock_metastore + .expect_open_shards() + .once() + .returning(|request| { + let subrequest: &OpenShardSubrequest = &request.subrequests[0]; + let shard = Shard { + index_uid: subrequest.index_uid.clone(), + source_id: subrequest.source_id.clone(), + shard_id: subrequest.shard_id.clone(), + shard_state: ShardState::Open as i32, + ingester_id: subrequest.ingester_id.clone(), + doc_mapping_uid: subrequest.doc_mapping_uid, + publish_position_inclusive: Some(Position::Beginning), + publish_token: None, + update_timestamp: 1724158996, + }; + let response = OpenShardsResponse { + subresponses: vec![OpenShardSubresponse { + subrequest_id: subrequest.subrequest_id, + open_shard: Some(shard), + }], + }; + Ok(response) + }); + let metastore = MetastoreServiceClient::from_mock(mock_metastore); + let ingester_pool = IngesterPool::default(); + + let mut controller = IngestController::new(metastore, ingester_pool.clone()); + let scaling_controller = + LegacyScalingController::new(DEFAULT_SHARD_THROUGHPUT_LIMIT, 1.001); + + let index_uid = IndexUid::for_test("test-index", 0); + let mut index_metadata = IndexMetadata::for_test("test-index", "ram://indexes/test-index"); + let source_id: SourceId = "test-source".to_string(); + let mut source_config = + SourceConfig::for_test(&source_id, quickwit_config::SourceParams::void()); + // set a vrl script + source_config.transform_config = + Some(quickwit_config::TransformConfig::new("".to_string(), None)); + index_metadata + .sources + .insert(source_id.clone(), source_config); + + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: source_id.clone(), + }; + let mut model = ControlPlaneModel::default(); + model.add_index(index_metadata); + let progress = Progress::default(); + + let shards = vec![Shard { + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(1)), + ingester_id: "test-ingester".to_string(), + shard_state: ShardState::Open as i32, + ..Default::default() + }]; + model.insert_shards(&index_uid, &source_id, shards); + + let mut mock_ingester = MockIngesterService::new(); + + mock_ingester.expect_init_shards().returning( + move |init_shard_request: InitShardsRequest| { + assert_eq!(init_shard_request.subrequests.len(), 1); + let init_shard_subrequest: &InitShardSubrequest = + &init_shard_request.subrequests[0]; + // we have vrl, so no validation + assert!(!init_shard_subrequest.validate_docs); + Ok(InitShardsResponse { + successes: vec![InitShardSuccess { + subrequest_id: init_shard_subrequest.subrequest_id, + shard: init_shard_subrequest.shard.clone(), + }], + failures: Vec::new(), + }) + }, + ); + + let ingester = + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); + ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); + + let shard_infos = BTreeSet::from_iter([ShardInfo { + shard_id: ShardId::from(1), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(4), + long_term_ingestion_rate: ByteSize::mib(4), + }]); + + scaling_controller + .update_local_shards( + &mut controller, + source_uid.clone(), + &shard_infos, + &mut model, + &progress, + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_ingest_controller_try_scale_up_shards() { + let mut mock_metastore = MockMetastoreService::new(); + + let index_uid = IndexUid::from_str("test-index:00000000000000000000000000").unwrap(); + let index_uid_clone = index_uid.clone(); + mock_metastore + .expect_open_shards() + .once() + .returning(move |request| { + assert_eq!(request.subrequests.len(), 1); + assert_eq!(request.subrequests[0].index_uid(), &index_uid_clone); + assert_eq!(request.subrequests[0].source_id, INGEST_V2_SOURCE_ID); + assert_eq!(request.subrequests[0].ingester_id, "test-ingester"); + + Err(MetastoreError::InvalidArgument { + message: "failed to open shards".to_string(), + }) + }); + let index_uid_clone = index_uid.clone(); + mock_metastore + .expect_open_shards() + .returning(move |request| { + assert_eq!(request.subrequests.len(), 1); + assert_eq!(request.subrequests[0].index_uid(), &index_uid_clone); + assert_eq!(request.subrequests[0].source_id, INGEST_V2_SOURCE_ID); + assert_eq!(request.subrequests[0].ingester_id, "test-ingester"); + + let subresponses = vec![metastore::OpenShardSubresponse { + subrequest_id: 0, + open_shard: Some(Shard { + index_uid: Some(index_uid.clone()), + source_id: INGEST_V2_SOURCE_ID.to_string(), + shard_id: Some(ShardId::from(1)), + ingester_id: "test-ingester".to_string(), + shard_state: ShardState::Open as i32, + ..Default::default() + }), + }]; + let response = metastore::OpenShardsResponse { subresponses }; + Ok(response) + }); + let metastore = MetastoreServiceClient::from_mock(mock_metastore); + + let ingester_pool = IngesterPool::default(); + + let mut controller = IngestController::new(metastore, ingester_pool.clone()); + let scaling_controller = + LegacyScalingController::new(DEFAULT_SHARD_THROUGHPUT_LIMIT, 1.001); + + let index_uid = IndexUid::for_test("test-index", 0); + let source_id: SourceId = INGEST_V2_SOURCE_ID.to_string(); + + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: source_id.clone(), + }; + let shard_stats = ShardStats { + num_open_shards: 2, + ..Default::default() + }; + let mut model = ControlPlaneModel::default(); + let index_metadata = + IndexMetadata::for_test(&index_uid.index_id, "ram://indexes/test-index:0"); + model.add_index(index_metadata); + + let source_config = SourceConfig::ingest_v2(); + model.add_source(&index_uid, source_config).unwrap(); + + let progress = Progress::default(); + + // Test could not find ingester because no ingester in pool + scaling_controller + .try_scale_up_shards( + &mut controller, + source_uid.clone(), + shard_stats, + &mut model, + &progress, + 1, + ) + .await + .unwrap(); + + let mut mock_ingester = MockIngesterService::new(); + + let index_uid_clone = index_uid.clone(); + mock_ingester + .expect_init_shards() + .once() + .returning(move |request| { + assert_eq!(request.subrequests.len(), 1); + + let subrequest = &request.subrequests[0]; + assert_eq!(subrequest.subrequest_id, 0); + + let shard = request.subrequests[0].shard(); + assert_eq!(shard.index_uid(), &index_uid_clone); + assert_eq!(shard.source_id, INGEST_V2_SOURCE_ID); + assert_eq!(shard.ingester_id, "test-ingester"); + + Err(IngestV2Error::Internal("failed to init shards".to_string())) + }); + let index_uid_clone = index_uid.clone(); + mock_ingester + .expect_init_shards() + .returning(move |request| { + assert_eq!(request.subrequests.len(), 1); + + let subrequest = &request.subrequests[0]; + assert_eq!(subrequest.subrequest_id, 0); + + let shard = subrequest.shard(); + assert_eq!(shard.index_uid(), &index_uid_clone); + assert_eq!(shard.source_id, INGEST_V2_SOURCE_ID); + assert_eq!(shard.ingester_id, "test-ingester"); + + let successes = vec![InitShardSuccess { + subrequest_id: request.subrequests[0].subrequest_id, + shard: Some(shard.clone()), + }]; + let response = InitShardsResponse { + successes, + failures: Vec::new(), + }; + Ok(response) + }); + let ingester = + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); + ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); + + // Test failed to open shards. + scaling_controller + .try_scale_up_shards( + &mut controller, + source_uid.clone(), + shard_stats, + &mut model, + &progress, + 1, + ) + .await + .unwrap(); + assert_eq!(model.all_shards().count(), 0); + + // Test failed to init shards. + scaling_controller + .try_scale_up_shards( + &mut controller, + source_uid.clone(), + shard_stats, + &mut model, + &progress, + 1, + ) + .await + .unwrap_err(); + assert_eq!(model.all_shards().count(), 0); + + // Test successfully opened shard. + scaling_controller + .try_scale_up_shards( + &mut controller, + source_uid.clone(), + shard_stats, + &mut model, + &progress, + 1, + ) + .await + .unwrap(); + assert_eq!( + model.all_shards().filter(|shard| shard.is_open()).count(), + 1 + ); + } + + #[tokio::test] + async fn test_ingest_controller_try_scale_down_shards() { + let metastore = MetastoreServiceClient::mocked(); + let ingester_pool = IngesterPool::default(); + + let mut controller = IngestController::new(metastore, ingester_pool.clone()); + let scaling_controller = + LegacyScalingController::new(DEFAULT_SHARD_THROUGHPUT_LIMIT, 1.001); + + let index_uid = IndexUid::for_test("test-index", 0); + let source_id: SourceId = "test-source".to_string(); + + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: source_id.clone(), + }; + let shard_stats = ShardStats { + num_open_shards: 2, + ..Default::default() + }; + let min_shards = NonZeroUsize::MIN; + let mut model = ControlPlaneModel::default(); + let progress = Progress::default(); + + // Test could not find a scale down candidate. + scaling_controller + .try_scale_down_shards( + &mut controller, + source_uid.clone(), + shard_stats, + min_shards, + &mut model, + &progress, + ) + .await + .unwrap(); + + let shards = vec![Shard { + shard_id: Some(ShardId::from(1)), + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + ingester_id: "test-ingester".to_string(), + shard_state: ShardState::Open as i32, + ..Default::default() + }]; + model.insert_shards(&index_uid, &source_id, shards); + + // Test ingester is unavailable. + scaling_controller + .try_scale_down_shards( + &mut controller, + source_uid.clone(), + shard_stats, + min_shards, + &mut model, + &progress, + ) + .await + .unwrap(); + + let mut mock_ingester = MockIngesterService::new(); + + let index_uid_clone = index_uid.clone(); + mock_ingester + .expect_close_shards() + .once() + .returning(move |request| { + assert_eq!(request.shard_pkeys.len(), 1); + assert_eq!(request.shard_pkeys[0].index_uid(), &index_uid_clone); + assert_eq!(request.shard_pkeys[0].source_id, "test-source"); + assert_eq!(request.shard_pkeys[0].shard_id(), ShardId::from(1)); + + Err(IngestV2Error::Internal( + "failed to close shards".to_string(), + )) + }); + let index_uid_clone = index_uid.clone(); + mock_ingester + .expect_close_shards() + .once() + .returning(move |request| { + assert_eq!(request.shard_pkeys.len(), 1); + assert_eq!(request.shard_pkeys[0].index_uid(), &index_uid_clone); + assert_eq!(request.shard_pkeys[0].source_id, "test-source"); + assert_eq!(request.shard_pkeys[0].shard_id(), ShardId::from(1)); + + let response = CloseShardsResponse { + successes: request.shard_pkeys, + }; + Ok(response) + }); + let ingester = + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock(mock_ingester)); + ingester_pool.insert(NodeId::from_str("test-ingester"), ingester); + + // Test failed to close shard. + scaling_controller + .try_scale_down_shards( + &mut controller, + source_uid.clone(), + shard_stats, + min_shards, + &mut model, + &progress, + ) + .await + .unwrap(); + assert!(model.all_shards().all(|shard| shard.is_open())); + + // Test successfully closed shard. + scaling_controller + .try_scale_down_shards( + &mut controller, + source_uid.clone(), + shard_stats, + min_shards, + &mut model, + &progress, + ) + .await + .unwrap(); + assert!(model.all_shards().all(|shard| shard.is_closed())); + + let shards = vec![Shard { + shard_id: Some(ShardId::from(2)), + index_uid: Some(index_uid.clone()), + source_id: source_id.clone(), + ingester_id: "test-ingester".to_string(), + shard_state: ShardState::Open as i32, + ..Default::default() + }]; + model.insert_shards(&index_uid, &source_id, shards); + + // Test rate limited. + scaling_controller + .try_scale_down_shards( + &mut controller, + source_uid.clone(), + shard_stats, + min_shards, + &mut model, + &progress, + ) + .await + .unwrap(); + assert!(model.all_shards().any(|shard| shard.is_open())); + } + + #[test] + fn test_find_scale_down_candidate() { + let index_uid = IndexUid::for_test("test-index", 0); + let source_id: SourceId = "test-source".to_string(); + + let source_uid = SourceUid { + index_uid: index_uid.clone(), + source_id: source_id.clone(), + }; + let mut model = ControlPlaneModel::default(); + + assert!(find_scale_down_candidate(&source_uid, &model).is_none()); + + let shards = vec![ + Shard { + index_uid: index_uid.clone().into(), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(1)), + shard_state: ShardState::Open as i32, + ingester_id: "test-ingester-0".to_string(), + ..Default::default() + }, + Shard { + index_uid: index_uid.clone().into(), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(2)), + shard_state: ShardState::Open as i32, + ingester_id: "test-ingester-0".to_string(), + ..Default::default() + }, + Shard { + index_uid: index_uid.clone().into(), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(3)), + shard_state: ShardState::Closed as i32, //< this one is closed + ingester_id: "test-ingester-0".to_string(), + ..Default::default() + }, + Shard { + index_uid: index_uid.clone().into(), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(4)), + shard_state: ShardState::Open as i32, + ingester_id: "test-ingester-1".to_string(), + ..Default::default() + }, + Shard { + index_uid: index_uid.clone().into(), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(5)), + shard_state: ShardState::Open as i32, + ingester_id: "test-ingester-1".to_string(), + ..Default::default() + }, + Shard { + index_uid: index_uid.clone().into(), + source_id: source_id.clone(), + shard_id: Some(ShardId::from(6)), + shard_state: ShardState::Open as i32, + ingester_id: "test-ingester-1".to_string(), + ..Default::default() + }, + ]; + // That's 3 open shards on indexer-1, 2 open shard and one closed shard on indexer-0.. + model.insert_shards(&index_uid, &source_id, shards); + + let shard_infos = BTreeSet::from_iter([ + ShardInfo { + shard_id: ShardId::from(1), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(1), + long_term_ingestion_rate: ByteSize::mib(1), + }, + ShardInfo { + shard_id: ShardId::from(2), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(2), + long_term_ingestion_rate: ByteSize::mib(2), + }, + ShardInfo { + shard_id: ShardId::from(3), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(3), + long_term_ingestion_rate: ByteSize::mib(3), + }, + ShardInfo { + shard_id: ShardId::from(4), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(4), + long_term_ingestion_rate: ByteSize::mib(4), + }, + ShardInfo { + shard_id: ShardId::from(5), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(5), + long_term_ingestion_rate: ByteSize::mib(5), + }, + ShardInfo { + shard_id: ShardId::from(6), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::mib(6), + long_term_ingestion_rate: ByteSize::mib(6), + }, + ]); + model.update_shards(&source_uid, &shard_infos); + + let shard = find_scale_down_candidate(&source_uid, &model).unwrap(); + // We pick ingester 1 has it has more open shard + assert_eq!(shard.ingester_id, "test-ingester-1"); + } + + fn shard_reports_for_test(min_shards: usize) -> (ControlPlaneModel, ShardsUpdate) { + let mut model = ControlPlaneModel::default(); + let mut metadata = IndexMetadata::for_test("index", "ram:///index"); + metadata.index_config.ingest_settings.min_shards = + std::num::NonZeroUsize::new(min_shards).unwrap(); + let index_uid = metadata.index_uid.clone(); + model.add_index(metadata); + let mut update = ShardsUpdate::default(); + for source_id in ["source-a", "source-b"] { + model + .add_source( + &index_uid, + SourceConfig::for_test(source_id, quickwit_config::SourceParams::void()), + ) + .unwrap(); + model.insert_shards( + &index_uid, + &source_id.to_string(), + vec![Shard { + index_uid: Some(index_uid.clone()), + source_id: source_id.to_string(), + shard_id: Some(ShardId::from(1)), + ingester_id: "ingester".to_string(), + shard_state: ShardState::Open as i32, + ..Default::default() + }], + ); + update.shard_infos_by_source.push( + SourceShardReport { + source_uid: SourceUid { + index_uid: index_uid.clone(), + source_id: source_id.to_string(), + }, + shard_infos: BTreeSet::from([ShardInfo { + shard_id: ShardId::from(1), + shard_state: ShardState::Open, + short_term_ingestion_rate: ByteSize::b(123), + long_term_ingestion_rate: ByteSize::b(456), + }]), + } + .into(), + ); + } + (model, update) + } + + #[tokio::test] + async fn test_shards_update_sources_and_generation() { + let pool = IngesterPool::default(); + let ingester = IngesterPoolEntry::ready_with_client(IngesterServiceClient::mocked()); + let generation = ingester.generation_id.as_u64(); + pool.insert(NodeId::from_str("ingester"), ingester); + let mut controller = IngestController::new(MetastoreServiceClient::mocked(), pool); + let scaling_controller = LegacyScalingController::new(DEFAULT_SHARD_THROUGHPUT_LIMIT, 1.5); + let progress = Progress::default(); + let (mut model, update) = shard_reports_for_test(1); + for wrong_generation in [generation - 1, generation + 1] { + scaling_controller + .handle_shards_update( + &mut controller, + "ingester", + wrong_generation, + update.clone(), + &mut model, + &progress, + ) + .await + .unwrap(); + assert!( + model + .all_shards() + .all(|shard| shard.short_term_ingestion_rate == ByteSize::default()) + ); + } + scaling_controller + .handle_shards_update( + &mut controller, + "ingester", + generation, + update.clone(), + &mut model, + &progress, + ) + .await + .unwrap(); + assert_eq!(model.all_shards().count(), 2); + assert!( + model + .all_shards() + .all(|shard| shard.short_term_ingestion_rate == ByteSize::b(123) + && shard.long_term_ingestion_rate == ByteSize::b(456)) + ); + + let (mut model, _) = shard_reports_for_test(1); + scaling_controller + .handle_shards_update( + &mut controller, + "joining-ingester", + generation, + update, + &mut model, + &progress, + ) + .await + .unwrap(); + assert!( + model + .all_shards() + .all(|shard| shard.short_term_ingestion_rate == ByteSize::b(123)) + ); + } + + #[tokio::test] + async fn test_shards_update_continues_only_after_certainly_aborted_errors() { + for certainly_aborted in [true, false] { + let (mut model, update) = shard_reports_for_test(2); + let mut mock_metastore = MockMetastoreService::new(); + mock_metastore + .expect_open_shards() + .times(if certainly_aborted { 2 } else { 1 }) + .returning(move |_| { + if certainly_aborted { + Err(MetastoreError::InvalidArgument { + message: "aborted".to_string(), + }) + } else { + Err(MetastoreError::Connection { + message: "uncertain".to_string(), + }) + } + }); + let mut mock_ingester = MockIngesterService::new(); + mock_ingester + .expect_init_shards() + .times(if certainly_aborted { 2 } else { 1 }) + .returning(|request| { + Ok(InitShardsResponse { + successes: request + .subrequests + .into_iter() + .map(|request| InitShardSuccess { + subrequest_id: request.subrequest_id, + shard: request.shard, + }) + .collect(), + failures: Vec::new(), + }) + }); + let pool = IngesterPool::default(); + pool.insert( + NodeId::from_str("ingester"), + IngesterPoolEntry::ready_with_client(IngesterServiceClient::from_mock( + mock_ingester, + )), + ); + let mut controller = + IngestController::new(MetastoreServiceClient::from_mock(mock_metastore), pool); + let scaling_controller = + LegacyScalingController::new(DEFAULT_SHARD_THROUGHPUT_LIMIT, 1.5); + let result = scaling_controller + .handle_shards_update( + &mut controller, + "ingester", + 1, + update, + &mut model, + &Progress::default(), + ) + .await; + assert_eq!(result.is_ok(), certainly_aborted); + let second_source = model + .all_shards() + .find(|shard| shard.source_id == "source-b") + .unwrap(); + let expected = if certainly_aborted { + ByteSize::b(123) + } else { + ByteSize::default() + }; + assert_eq!(second_source.short_term_ingestion_rate, expected); + } + } +} diff --git a/quickwit/quickwit-control-plane/src/ingest/mod.rs b/quickwit/quickwit-control-plane/src/ingest/mod.rs index 25da92f2916..3aca5497826 100644 --- a/quickwit/quickwit-control-plane/src/ingest/mod.rs +++ b/quickwit/quickwit-control-plane/src/ingest/mod.rs @@ -13,8 +13,10 @@ // limitations under the License. pub(crate) mod ingest_controller; -mod scaling_arbiter; +mod legacy_scaling_arbiter; +mod legacy_scaling_controller; mod wait_handle; pub use ingest_controller::IngestController; +pub(crate) use legacy_scaling_controller::LegacyScalingController; pub use wait_handle::WaitHandle; diff --git a/quickwit/quickwit-control-plane/src/model/mod.rs b/quickwit/quickwit-control-plane/src/model/mod.rs index 3ecfc473651..6972fa5572b 100644 --- a/quickwit/quickwit-control-plane/src/model/mod.rs +++ b/quickwit/quickwit-control-plane/src/model/mod.rs @@ -370,18 +370,18 @@ impl ControlPlaneModel { .find_open_shards(index_uid, source_id, unavailable_ingesters) } + pub fn legacy_shard_stats(&self, source_uid: &SourceUid) -> Option { + self.shard_table.legacy_shard_stats(source_uid) + } + /// Updates the state and ingestion rate of the shards according to the given shard infos. - pub fn update_shards( - &mut self, - source_uid: &SourceUid, - shard_infos: &ShardInfos, - ) -> ShardStats { + pub fn update_shards(&mut self, source_uid: &SourceUid, shard_infos: &ShardInfos) { debug!( index_uid=%source_uid.index_uid, source_id=%source_uid.source_id, "updating shards" ); - self.shard_table.update_shards(source_uid, shard_infos) + self.shard_table.update_shards(source_uid, shard_infos); } /// Sets the state of the shards identified by their index UID, source ID, and shard IDs to diff --git a/quickwit/quickwit-control-plane/src/model/shard_table.rs b/quickwit/quickwit-control-plane/src/model/shard_table.rs index ec9ae06cbe6..59c9849869b 100644 --- a/quickwit/quickwit-control-plane/src/model/shard_table.rs +++ b/quickwit/quickwit-control-plane/src/model/shard_table.rs @@ -451,6 +451,12 @@ impl ShardTable { Some(open_shards) } + pub fn legacy_shard_stats(&self, source_uid: &SourceUid) -> Option { + let table_entry = self.table_entries.get(source_uid)?; + let shard_stats = table_entry.shards_stats(); + Some(shard_stats) + } + pub fn update_shard_metrics_for_source_uid(&self, source_uid: &SourceUid) { let Some(table_entry) = self.table_entries.get(source_uid) else { return; @@ -484,13 +490,9 @@ impl ShardTable { gauge!(parent: CLOSED_SHARDS, labels: [labels]).set(num_closed_shards as f64); } - pub fn update_shards( - &mut self, - source_uid: &SourceUid, - shard_infos: &ShardInfos, - ) -> ShardStats { + pub fn update_shards(&mut self, source_uid: &SourceUid, shard_infos: &ShardInfos) { let Some(table_entry) = self.table_entries.get_mut(source_uid) else { - return ShardStats::default(); + return; }; for shard_info in shard_infos { let ShardInfo { @@ -510,7 +512,6 @@ impl ShardTable { } } } - table_entry.shards_stats() } /// Sets the state of the shards identified by their index UID, source ID, and shard IDs to @@ -924,7 +925,8 @@ mod tests { long_term_ingestion_rate: ByteSize::mib(5), }, ]); - let shard_stats = shard_table.update_shards(&source_uid, &shard_infos); + shard_table.update_shards(&source_uid, &shard_infos); + let shard_stats = shard_table.legacy_shard_stats(&source_uid).unwrap(); assert_eq!(shard_stats.num_open_shards, 2); assert_eq!( shard_stats.avg_short_term_ingestion_rate,