use std::{ collections::{ HashMap, HashSet, hash_map::Entry::{Occupied as HOccupied, Vacant as HVacant}, }, sync::Arc, time::{Duration, Instant}, }; use anyhow::{Result, anyhow}; use log::{debug, info, warn}; use tici_common::{ codec::{codec_v2::get_txn_keyspace_prefix, table_end_key, table_start_key}, index::IndexKey, }; use tici_grpc::{ErrorCode, KeyRange, ShardLocalCacheInfo, ShardManifestHeader}; use tokio::sync::RwLock; use crate::{ ShardInfo, ShardInfoAppendUpdate, config::MetaServiceConfig, meta::{ metrics::{ READER_SHARD_COUNT_GAUGE, SHARD_CACHE_ALLOCATE_COUNTER, SHARD_CACHE_EVENT_COUNTER, SHARD_CACHE_QUERY_DURATION_HISTOGRAM, }, reader::{ AllocateReaderReason, CacheAllocator, IndexMeta, UnifiedShardInfo, cache_scheduler::{CacheSchedulerFactory, ReaderCacheScheduler}, reader_pool::ReaderPool, }, writer::ShardEventIndexInfo, }, }; /// Manages shard cache allocation for each index pub struct ShardCacheManager { // Unified metadata manager that contains both shard metadata and cache allocations pub index_metadata: RwLock>, pub cache_scheduler: Arc, pub config: MetaServiceConfig, auto_merge_range: bool, } impl ShardCacheManager { pub fn new(config: &MetaServiceConfig, auto_merge_range: bool) -> Arc { let cache_scheduling_config = config.cache_scheduling.clone(); let cache_scheduler = CacheSchedulerFactory::create_cache_scheduler(&cache_scheduling_config.strategy); Arc::new(Self { index_metadata: RwLock::new(HashMap::new()), cache_scheduler, config: config.clone(), auto_merge_range, }) } pub fn make_cache_allocator(&self) -> CacheAllocator { CacheAllocator::new( self.cache_scheduler.clone(), self.config.s3.bucket.clone(), self.config.cache_scheduling.allocate_parallelism, self.config.cache_scheduling.max_retry_attempts, ) } pub fn normalize_keyrange( &self, index_key: &IndexKey, tidb_x_already_encoded: bool, key_ranges: &[tici_grpc::KeyRange], ) -> Vec { if key_ranges.is_empty() { return vec![tici_grpc::KeyRange { start_key: table_start_key(index_key), end_key: table_end_key(index_key), }]; } if index_key.is_tidb_x() && !tidb_x_already_encoded { let keyranges: Vec = key_ranges .iter() .map(|range| { // TODO(P2) validate range let start_key = if range.start_key.is_empty() { table_start_key(index_key) } else { let mut start_key = get_txn_keyspace_prefix(index_key.keyspace_id); start_key.extend_from_slice(&range.start_key); start_key }; let end_key = if range.end_key.is_empty() { table_end_key(index_key) } else { let mut end_key = get_txn_keyspace_prefix(index_key.keyspace_id); end_key.extend_from_slice(&range.end_key); end_key }; KeyRange { start_key, end_key } }) .collect(); log::debug!("encoded keyrange={:?}", keyranges); keyranges } else { key_ranges.to_vec() } } // Reader management methods using unified metadata manager pub async fn on_reader_node_down(&self, expired_reader_addr: &str) -> usize { let mut metadata_guard = self.index_metadata.write().await; let mut removed_count = 0; for index_meta in metadata_guard.values_mut() { removed_count += index_meta.on_reader_node_down(expired_reader_addr).await; } if removed_count > 0 { info!( "removed {} cache allocations for expired reader addr={}", removed_count, expired_reader_addr ); } // Update counter: decrease by number of removed allocations if removed_count > 0 { SHARD_CACHE_ALLOCATE_COUNTER.sub(removed_count as i64); } removed_count } } impl ShardCacheManager { pub async fn get_index_meta_cloned(&self, index_key: &IndexKey) -> Option { self.index_metadata.read().await.get(index_key).cloned() } pub async fn find_overlapping_shards( &self, index_key: &IndexKey, key_ranges: &[tici_grpc::KeyRange], reason: AllocateReaderReason, perform_warmup: Option>, ) -> Result<(Vec, usize, bool)> { if let Some(index_meta) = self.index_metadata.read().await.get(index_key) { Ok(index_meta .find_overlapping_shards(key_ranges, reason, perform_warmup, self.auto_merge_range) .await?) } else { Err(anyhow!("index not found, {:?}", index_key)) } } // TODO(P2) Improve the performance pub async fn get_shard_by_id(&self, shard_id: &u64) -> Option { for (_, index) in self.index_metadata.read().await.iter() { if let Some(s) = index.get_shard_by_id(shard_id) { return Some(s.clone()); } } None } pub async fn get_shards_count(&self, index_key: &IndexKey) -> usize { if let Some(x) = self.index_metadata.read().await.get(index_key) { x.shard_map.len() } else { 0 } } pub async fn remedy_shard_caches( &self, not_allocated_shard_ids: HashSet, ) -> (Vec, Vec) { // Some shard_ids are not recorded in read manager, but reported from heartbeat, // we must check: // - If the shard does not exist, then it is merged or splitted, tell reader to // remove them. // - If the shard exists, but not allocated, this must because meta service has // started, and lost all its allocation map. // TODO(P2) After we support failover, we can remove this logic, because all // such shards can be considered under the first circumstance. let mut need_to_remedy = vec![]; let mut nonexist = vec![]; for shard_id in not_allocated_shard_ids { if let Some(us) = self.get_shard_by_id(&shard_id).await { need_to_remedy.push(us); } else { nonexist.push(shard_id); } } (need_to_remedy, nonexist) } pub async fn on_warmup_finished( &self, unified_shard: &UnifiedShardInfo, reason: AllocateReaderReason, ) -> Result> { if let Some(index_meta) = self .index_metadata .write() .await .get_mut(&unified_shard.index_key) { index_meta.on_warmup_finished(unified_shard, reason).await } else { info!( "on_warmup_finished: can't find index for warmup, unified_shard={}", unified_shard.to_long_ident() ); Ok(HashSet::new()) } } pub async fn get_or_allocate_cache( &self, index_key: &IndexKey, normal_ranges: Vec, reader_pool: Arc, limit: usize, reason: AllocateReaderReason, wait_timeout: Option, ) -> Result, tici_grpc::ErrorCode>> { // If the range is malformed, an error may be raised. let start = Instant::now(); let mut retry_count = 0; loop { let (overlapping_shards, unready_count, is_no_reader) = if let Some(index_meta) = self.index_metadata.read().await.get(index_key) { match index_meta .find_overlapping_shards( normal_ranges.as_slice(), reason.clone(), Some(reader_pool.clone()), self.auto_merge_range, ) .await { Ok(x) => x, Err(e) => { info!( "get_or_allocate_cache: index_key={:?}, range={:?}, error={:?}", index_key, normal_ranges, e ); return Ok(either::Right(ErrorCode::UnknownError)); } } } else { // No index found. SHARD_CACHE_EVENT_COUNTER .with_label_values(&["read_failed_no_index"]) .inc(); info!( "get_or_allocate_cache: can't find index, index_key={:?}", index_key ); return Ok(either::Right(ErrorCode::IndexNotFound)); }; if unready_count > 0 { if is_no_reader { info!( "get_or_allocate_cache: failed for no readers, index={:?}, range={:?}", index_key, normal_ranges ); return Ok(either::Right(tici_grpc::ErrorCode::WorkerNotFound)); } if retry_count % 10 == 0 { if retry_count % (10 * 60) == 0 { // Will print for every 30s. info!( "get_or_allocate_cache: retry find_overlapping_shards, index={:?}, range={:?}", index_key, normal_ranges ); } else { debug!( "get_or_allocate_cache: retry find_overlapping_shards, index={:?}, range={:?}", index_key, normal_ranges ); } } } else { // Succeed. let mut vr = vec![]; for shard in overlapping_shards .into_iter() .take(if limit > 0 { limit } else { usize::MAX }) { // Note the range is trimmed by `find_overlapping_shards` vr.push(ShardLocalCacheInfo { shard: Some(ShardManifestHeader { shard_id: shard.shard_id, start_key: shard.start_key().clone().into(), end_key: shard.end_key().clone().into(), epoch: shard.epoch, seq: shard.seq.load(std::sync::atomic::Ordering::SeqCst), }), local_cache_addrs: shard.cloned_tiflash_server_addrs().await, }); } if retry_count == 0 { SHARD_CACHE_QUERY_DURATION_HISTOGRAM .with_label_values(&["hit"]) .observe(start.elapsed().as_secs_f64()); } else { SHARD_CACHE_QUERY_DURATION_HISTOGRAM .with_label_values(&["need_warm"]) .observe(start.elapsed().as_secs_f64()); } return Ok(either::Left(vr)); } // TODO(P1) Make this depend on state change? if let Some(wait_timeout) = wait_timeout { if start.elapsed() > wait_timeout { info!( "get_or_allocate_cache: can not find all ready shards, normal_ranges={:?}, index={:?}", normal_ranges, index_key ); return Err(anyhow!( "get_or_allocate_cache: can not find all ready shards, normal_ranges={:?}, index={:?}", normal_ranges, index_key )); } } retry_count += 1; tokio::time::sleep(std::time::Duration::from_millis(100)).await; } } } impl ShardCacheManager { /// Handle full sync event: Initialize all shard metadata pub async fn handle_full_sync( &self, index_infos: Vec, shards: Vec, ) -> Result<()> { // Create index metadata for each index for index_info in index_infos.iter() { let index_key = IndexKey::from(&index_info.index_info); let mut metadata_guard = self.index_metadata.write().await; let index_meta = IndexMeta::new( index_info.table_info.clone(), index_info.index_info.clone(), self.make_cache_allocator(), self.config.cache_scheduling.optimize_for_fast_scan, ); metadata_guard.insert(index_key, index_meta); } // Add all shards to their respective indexes { let shards = shards.as_slice(); info!("handle_full_sync, len={}", shards.len()); let mut metadata_guard = self.index_metadata.write().await; // New codes let mut m: HashMap> = HashMap::new(); for shard in shards { match m.entry(shard.index_key.clone()) { // For each shard, insert it into the temporary map `m` grouped by index_key, // accumulating shards into a Vec per key. HOccupied(mut entry) => { entry .get_mut() .push(UnifiedShardInfo::from_shard_info(shard)); } HVacant(entry) => { entry.insert(vec![UnifiedShardInfo::from_shard_info(shard)]); } } } for (k, v) in m.into_iter() { if let Some(index_meta) = metadata_guard.get_mut(&k) { index_meta.init(v); } } } info!( "full sync completed: indexes={}, shards={}", index_infos.len(), shards.len() ); // Update reader shard count gauge to total shards after full sync READER_SHARD_COUNT_GAUGE.set(shards.len() as i64); Ok(()) } /// Handle shard add event: Add new shard to metadata pub async fn handle_shard_add( &self, index_info: ShardEventIndexInfo, shard: &ShardInfo, ) -> Result<()> { let index_key = IndexKey::from(&index_info.index_info); // Prepare shard metadata outside the lock let unified_shard = UnifiedShardInfo::from_shard_info(shard); { let mut metadata_guard = self.index_metadata.write().await; let index_meta = metadata_guard.entry(index_key).or_insert_with(|| { // Create new IndexMeta if it doesn't exist IndexMeta::new( index_info.table_info, index_info.index_info, self.make_cache_allocator(), self.config.cache_scheduling.optimize_for_fast_scan, ) }); index_meta.handle_add_shard(unified_shard)?; } info!("added shard to index metadata, {}", shard.to_long_ident()); READER_SHARD_COUNT_GAUGE.inc(); Ok(()) } pub async fn handle_shard_remove( &self, index_info: ShardEventIndexInfo, shard_ids: Vec, reader_node: Arc, ) -> Result<()> { let index_key = IndexKey::from(&index_info.index_info); debug!( "handle_shard_remove, index_key={:?}, shard_ids={:?}", index_key, shard_ids ); let mut metadata_guard = self.index_metadata.write().await; if let Some(index_meta) = metadata_guard.get_mut(&index_key) { if let Err(e) = index_meta.handle_remove_shard(shard_ids.as_slice()) { warn!( "failed to remove shard because shard not found, shard_id={}, index_key={:?}, err={:?}", shard_ids .iter() .map(|x| format!("shard_id={}", x)) .collect::>() .join(";"), index_key, e ); return Err(e); } } else { return Err(anyhow::anyhow!( "index metadata not found, index_key={:?}", &index_key )); } drop(metadata_guard); reader_node.remove_shards(shard_ids.as_slice()).await; info!( "removed shard from metadata and cache, {}", shard_ids .iter() .map(|x| format!("shard_id={}", x)) .collect::>() .join(";") ); // Metrics: -1 shard on successful removal READER_SHARD_COUNT_GAUGE.dec(); Ok(()) } /// Handle shard update event: Update shard metadata (especially seq number) pub async fn handle_shard_update( &self, index_info: ShardEventIndexInfo, shard: &ShardInfo, ) -> Result<()> { let index_key = IndexKey::from(&index_info.index_info); // Use unified metadata manager to update shard debug!( "update_shard, index_key={:?}, shard_id={}", &index_key, shard.shard_id ); let mut metadata_guard = self.index_metadata.write().await; if let Some(index_meta) = metadata_guard.get_mut(&index_key) { index_meta.handle_update_shard(shard).await?; Ok(()) } else { Err(anyhow::anyhow!( "index metadata not found, index_key={:?}", index_key )) } } pub async fn handle_shard_append( &self, index_info: ShardEventIndexInfo, append: &ShardInfoAppendUpdate, ) -> Result<()> { let index_key = IndexKey::from(&index_info.index_info); let metadata_guard = self.index_metadata.read().await; if let Some(index_meta) = metadata_guard.get(&index_key) { index_meta.handle_append_shard(append).await } else { Err(anyhow::anyhow!( "index metadata not found, cache_key={:?}", index_key )) } } pub async fn handle_shard_split( &self, index: ShardEventIndexInfo, original: ShardInfo, new_shards: Vec, reader_pool: Arc, ) { let index_key = IndexKey::from(&index.index_info); if let Some(index_meta) = self.index_metadata.write().await.get_mut(&index_key) { index_meta .handle_split_shard(original, new_shards.as_slice(), reader_pool) .await; } else { info!( "handle_shard_split: can't find index for split, index_key={:?}", index_key ); } // Metrics: reader shard count += (new_shards - 1) if !new_shards.is_empty() { let delta = new_shards.len() as i64 - 1; match delta.cmp(&0) { std::cmp::Ordering::Greater => { READER_SHARD_COUNT_GAUGE.add(delta); } _ => { // delta <= 0 means no actual split occurred, no change // needed } } } } pub async fn handle_shard_merge( &self, index_info: ShardEventIndexInfo, old_shards: Vec, new_shard: ShardInfo, reader_pool: Arc, ) { let index_key = IndexKey::from(&index_info.index_info); let to_merge_len = old_shards.len(); if let Some(index_meta) = self.index_metadata.write().await.get_mut(&index_key) { index_meta .handle_merge_shard(old_shards, new_shard, reader_pool) .await; } else { info!( "handle_shard_merge: can't find index for merge, index_key={:?}", index_key ); } // Metrics: -(old_shards.len() - 1) let delta = to_merge_len as i64 - 1; if delta > 0 { READER_SHARD_COUNT_GAUGE.sub(delta); } // delta <= 0 means no actual merge occurred, no change needed } pub async fn on_msg_get_shard_manifest( &self, request: &tici_grpc::GetShardManifestRequest, ) -> Result { match self.index_metadata.read().await.get(&IndexKey::new( request.keyspace_id, request.table_id, request.index_id, )) { Some(index_meta) => match index_meta.get_shard_by_id(&request.shard_id) { None => Ok(tici_grpc::GetShardManifestResponse { status: ErrorCode::ShardNotFound.into(), ..Default::default() }), Some(x) => { let frag_metas = x.cloned_fragments().await; let seq = x.seq.load(std::sync::atomic::Ordering::SeqCst); info!( "on_msg_get_shard_manifest: shard_id={}, seq={}, frag={}", request.shard_id, seq, frag_metas .iter() .map(|x| x.frag_path.clone()) .collect::>() .join("|") ); Ok(tici_grpc::GetShardManifestResponse { status: ErrorCode::Success.into(), frag_metas, shard_id: x.shard_id, seq, }) } }, None => Ok(tici_grpc::GetShardManifestResponse { status: ErrorCode::IndexNotFound.into(), ..Default::default() }), } } } #[cfg(test)] mod tests { use std::time::{Duration, Instant}; use scopeguard::defer; use tici_grpc::{ ReaderCachedShard, test_utils::{ create_test_heartbeat_request, create_test_heartbeat_request_with_cached_shards, }, }; use tici_system::readable::ReadableDuration; use tici_test_utils::test_utils::{init_logger_for_ut, must_finally, must_finally_ex}; use super::*; use crate::{ config::make_default_test_meta_config, meta::{ metrics::{ SHARD_CACHE_COUNT_GAUGE, SHARD_CACHE_WARM_UP_QUEUE_HISTOGRAM, clear_prom_metrics, }, reader::{ ShardCacheManager, reader_pool::warmup::{WarmupMgr, WarmupTestAction, WarmupTestController}, test_utils::{ create_test_index_info, create_test_shard_info2, create_test_table_info, }, }, writer::ShardEventIndexInfo, }, }; #[derive(Clone)] struct Test { pub index_key: IndexKey, pub reader_pool: Arc, pub manager: Arc, pub index_info: ShardEventIndexInfo, pub shards: Vec, pub warmup_controller: Arc, } async fn setup_tests(config: MetaServiceConfig, auto_merge_range: bool) -> Test { let index_key = IndexKey::new(1, 1, 1); let warmup_controller = Arc::new(WarmupTestController::new()); warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; info!("setup_tests: auto_merge_range={}", auto_merge_range); let warmup_mgr = WarmupMgr::with_test_controller(warmup_controller.clone()); let manager = ShardCacheManager::new(&config, auto_merge_range); let mut reader_pool = ReaderPool::new_with_warmup_mgr(&config, warmup_mgr); reader_pool.init(manager.clone()); reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![], )) .await .unwrap(); let reader_pool = Arc::new(reader_pool); let index_info = ShardEventIndexInfo { table_info: create_test_table_info(&index_key), index_info: create_test_index_info(&index_key), }; let shards_builder: Vec<(u64, &[u8], &[u8], u64)> = vec![(1, b"", b"f", 1), (2, b"f", b"k", 1), (3, b"k", b"", 1)]; let mut shards = vec![]; for (shard_id, start_key, end_key, seq) in shards_builder { let unified_shard = create_test_shard_info2(&index_key, shard_id, start_key, end_key, seq); shards.push(unified_shard); } manager .handle_full_sync(vec![index_info.clone()], shards.clone()) .await .unwrap(); assert_eq!( manager .index_metadata .read() .await .get(&index_key) .unwrap() .shard_map .shards_by_id .len(), 3 ); Test { index_key, reader_pool, manager, index_info, shards, warmup_controller, } } async fn check_find_over_shards_ex( manager: Arc, index_key: &IndexKey, key_ranges: &[KeyRange], ends: Vec>, shard_ids: Vec, warmup: Option>, perform_warmup: Option>, ) -> Vec { let (res, ..) = manager .find_overlapping_shards( index_key, key_ranges, AllocateReaderReason::ReadRequest, perform_warmup.clone(), ) .await .unwrap(); let res = if perform_warmup.is_some() { for shard_id in shard_ids.iter() { must_finally(async || { manager .get_shard_by_id(shard_id) .await .unwrap() .has_ready_readers() .await }) .await; } let (res, ..) = manager .find_overlapping_shards( index_key, key_ranges, AllocateReaderReason::ReadRequest, perform_warmup.clone(), ) .await .unwrap(); res } else { res }; assert_eq!(res.len(), ends.len(), "{:?} <===> {:?}", res, ends); for (x, y) in res.iter().zip(ends.iter()) { assert_eq!( &index_key.trim_prefix(x.range.encoded_end().as_slice()), y, "{:?}", res ); } assert_eq!( res.len(), shard_ids.len(), "actual: {:?} vs expected: {:?}", res, shard_ids ); for (x, y) in res.iter().zip(shard_ids.iter()) { assert_eq!(&x.shard_id, y, "{:?}", res); } if let Some(v) = warmup { assert_eq!(res.len(), v.len()); for (x, y) in res.iter().zip(v.iter()) { assert_eq!(&x.has_ready_readers().await, y, "{:?}", res); } } res } async fn check_find_over_shards( manager: Arc, index_key: &IndexKey, key_ranges: &[KeyRange], ends: Vec>, shard_ids: Vec, ) -> Vec { check_find_over_shards_ex(manager, index_key, key_ranges, ends, shard_ids, None, None).await } async fn must_success_warmed_up( shard: &UnifiedShardInfo, index_key: &IndexKey, _manager: Arc, read_pool: Arc, ) { let handle = read_pool .background_warmup( shard, read_pool.clone(), create_test_table_info(index_key), create_test_index_info(index_key), AllocateReaderReason::ReadRequest, ) .unwrap(); handle.await.unwrap(); must_warmed_up(shard).await; } async fn must_not_warmed_up(shard: &UnifiedShardInfo) { let check = must_finally_ex(async |_| { if !shard.has_ready_readers().await && shard.in_warming.load(std::sync::atomic::Ordering::SeqCst) == 0 { either::Either::Left(()) } else { either::Either::Right(format!( "{:?},{}", shard.get_reader_addrs().await, shard.in_warming.load(std::sync::atomic::Ordering::SeqCst) )) } }) .await; if let either::Either::Right(info) = check { panic!("must_not_warmed_up failed: {info}"); } } async fn must_warmed_up(shard: &UnifiedShardInfo) { let check = must_finally_ex(async |_| { if shard.has_ready_readers().await && shard.in_warming.load(std::sync::atomic::Ordering::SeqCst) == 0 { either::Either::Left(()) } else { either::Either::Right(format!( "{:?},{}", shard.get_reader_addrs().await, shard.in_warming.load(std::sync::atomic::Ordering::SeqCst) )) } }) .await; if let either::Either::Right(info) = check { panic!("must_warmed_up failed: {info}"); } } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_basic_warmup() { clear_prom_metrics(); // We don't test the `do_replace` in this test. fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let test = setup_tests(config, false).await; let Test { index_key, reader_pool, manager, index_info, shards, warmup_controller, } = test.clone(); drop(index_info); drop(shards); let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key), b"".to_vec()].concat(), end_key: [table_end_key(&index_key), b"".to_vec()].concat(), }]; let start = Instant::now(); warmup_controller .set_fallback_action(Some(WarmupTestAction::Fail( "mock warmup failed".to_string(), ))) .await; let r = manager .get_or_allocate_cache( &index_key, key_ranges.clone(), reader_pool.clone(), 0, AllocateReaderReason::ReadRequest, Some(Duration::from_millis(400)), ) .await; assert!(r.is_err(), "{:?}", r); assert!(start.elapsed() > Duration::from_millis(300)); // Wait all warm up to end. tokio::time::sleep(Duration::from_millis(1000)).await; warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; let (res, ..) = manager .find_overlapping_shards( &index_key, key_ranges.as_slice(), AllocateReaderReason::ReadRequest, None, ) .await .unwrap(); assert_eq!(res.len(), 3, "{:?}", res); assert!( manager .find_overlapping_shards(&index_key, &[], AllocateReaderReason::ReadRequest, None,) .await .is_err() ); let shard_1 = res.first().cloned().unwrap(); let shard_2 = res.get(1).cloned().unwrap(); let shard_3 = res.get(2).cloned().unwrap(); let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key), b"c".to_vec()].concat(), end_key: [table_start_key(&index_key), b"h".to_vec()].concat(), }]; let _ = check_find_over_shards( manager.clone(), &index_key, key_ranges.as_slice(), vec![b"f".to_vec(), b"k".to_vec()], vec![1, 2], ) .await; // ===> Mock a warmup on read request must_success_warmed_up(&shard_1, &index_key, manager.clone(), reader_pool.clone()).await; must_not_warmed_up(&shard_2).await; must_not_warmed_up(&shard_3).await; assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 3 ); let handle = reader_pool .background_warmup( &shard_2, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap(); handle.await.unwrap(); let handle = reader_pool .background_warmup( &shard_3, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap(); handle.await.unwrap(); // ===> Mock a heartbeat with a non exist shard let res = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![shard_1.shard_id, 999], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(res.shards_to_keep.len(), 1); assert_eq!(res.shards_to_remove.len(), 1); // 999 // ===> Mock a heartbeat with shard_2, it can remedy assert!(!shard_2.has_ready_readers().await); let res = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr0".to_string(), "addr0_0".to_string(), vec![ ReaderCachedShard { shard_id: shard_1.shard_id, seq: shard_1.seq.load(std::sync::atomic::Ordering::SeqCst), }, ReaderCachedShard { shard_id: shard_2.shard_id, seq: shard_2.seq.load(std::sync::atomic::Ordering::SeqCst), }, ], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(res.shards_to_keep.len(), 2); assert!(shard_2.has_ready_readers().await); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_add_replica() { clear_prom_metrics(); // Avoid coupling this test with shard_map replacement details. fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, .. } = setup_tests(config, false).await; // Register a second reader so the shard can have 2 replicas. reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![], )) .await .unwrap(); let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); assert!(!shard_1.has_ready_readers().await); // Warm up the first replica. reader_pool .background_warmup( &shard_1, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap() .await .unwrap(); must_warmed_up(&shard_1).await; assert_eq!(shard_1.get_reader_addrs().await.len(), 1); let first_addr = shard_1 .cloned_tiflash_server_addrs() .await .first() .unwrap() .clone(); shard_1.set_desired_replica_count(2).await; // Add one more replica on another reader. reader_pool .background_add_replica( &shard_1, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::Manual, ) .unwrap() .await .unwrap(); must_warmed_up(&shard_1).await; let mut addrs = shard_1.get_reader_addrs().await; addrs.sort(); assert_eq!(addrs, vec!["addr0".to_string(), "addr1".to_string()]); let key_ranges = vec![KeyRange { start_key: [table_start_key(&index_key), b"a".to_vec()].concat(), end_key: [table_start_key(&index_key), b"b".to_vec()].concat(), }]; let res = manager .get_or_allocate_cache( &index_key, key_ranges, reader_pool.clone(), 0, AllocateReaderReason::ReadRequest, Some(Duration::from_millis(200)), ) .await .unwrap(); let shards = match res { either::Either::Left(vec) => vec, either::Either::Right(ec) => panic!("expected Ok, got error_code={:?}", ec), }; assert_eq!(shards.len(), 1); let shard = shards.first().expect("checked len == 1"); assert_eq!(shard.shard.as_ref().unwrap().shard_id, shard_1.shard_id); let local_cache_addrs = shard.local_cache_addrs.clone(); assert_eq!(local_cache_addrs.len(), 2); assert_eq!(local_cache_addrs.first().unwrap(), &first_addr); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_heartbeat_reported_replica_after_warmup() { clear_prom_metrics(); // Avoid coupling this test with shard_map replacement details. fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); assert!(!shard_1.has_ready_readers().await); // Warm up shard_1 on the only existing reader(addr0). reader_pool .background_warmup( &shard_1, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap() .await .unwrap(); must_warmed_up(&shard_1).await; assert_eq!(shard_1.get_reader_addrs().await, vec!["addr0".to_string()]); let first_addr = shard_1 .cloned_tiflash_server_addrs() .await .first() .unwrap() .clone(); // Another reader reports it already has shard_1 cache. Since the shard still // only wants one replica, the reported replica is recovered and then // immediately marked as removing by redundant cleanup. let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: shard_1.seq.load(std::sync::atomic::Ordering::SeqCst), }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_remove.len(), 1); assert_eq!(resp.shards_to_keep.len(), 0); assert_eq!(shard_1.get_reader_addrs().await, vec!["addr0".to_string()]); assert_eq!( shard_1.get_removing_reader_addrs().await, vec!["addr1".to_string()] ); let key_ranges = vec![KeyRange { start_key: [table_start_key(&index_key), b"a".to_vec()].concat(), end_key: [table_start_key(&index_key), b"b".to_vec()].concat(), }]; let res = manager .get_or_allocate_cache( &index_key, key_ranges, reader_pool.clone(), 0, AllocateReaderReason::ReadRequest, Some(Duration::from_millis(200)), ) .await .unwrap(); let shards = match res { either::Either::Left(vec) => vec, either::Either::Right(ec) => panic!("expected Ok, got error_code={:?}", ec), }; assert_eq!(shards.len(), 1); let shard = shards.first().expect("checked len == 1"); assert_eq!(shard.shard.as_ref().unwrap().shard_id, shard_1.shard_id); let local_cache_addrs = shard.local_cache_addrs.clone(); assert_eq!(local_cache_addrs, vec![first_addr]); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_legacy_heartbeat_promotes_when_no_ready_reader_exists() { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); assert!(!shard_1.has_ready_readers().await); let legacy_resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![shard_1.shard_id], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(legacy_resp.shards_to_remove.len(), 0); assert_eq!(legacy_resp.shards_to_keep.len(), 1); assert_eq!(shard_1.get_reader_addrs().await, vec!["addr1".to_string()]); assert_eq!(shard_1.get_reader_observed_seq("addr1").await, None); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_legacy_heartbeat_recovered_replica_is_ready() { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); assert!( shard_1 .add_ready_reader("addr0".to_string(), "addr0_0".to_string()) .await ); shard_1.set_desired_replica_count(2).await; let legacy_resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![shard_1.shard_id], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(legacy_resp.shards_to_remove.len(), 0); assert_eq!(legacy_resp.shards_to_keep.len(), 1); let mut addrs = shard_1.get_reader_addrs().await; addrs.sort(); assert_eq!(addrs, vec!["addr0".to_string(), "addr1".to_string()]); assert_eq!(shard_1.get_reader_observed_seq("addr1").await, None); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_heartbeat_updates_reader_observed_seq() { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: 3, }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_remove.len(), 0); assert_eq!(resp.shards_to_keep.len(), 1); assert_eq!(shard_1.get_reader_observed_seq("addr1").await, Some(3)); let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: 7, }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_remove.len(), 0); assert_eq!(resp.shards_to_keep.len(), 1); assert_eq!(shard_1.get_reader_observed_seq("addr1").await, Some(7)); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_heartbeat_recovered_replica_with_seq_is_ready() { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); assert!( shard_1 .add_ready_reader("addr0".to_string(), "addr0_0".to_string()) .await ); shard_1.set_desired_replica_count(2).await; let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: 0, }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_remove.len(), 0); assert_eq!(resp.shards_to_keep.len(), 1); let mut addrs = shard_1.get_reader_addrs().await; addrs.sort(); assert_eq!(addrs, vec!["addr0".to_string(), "addr1".to_string()]); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_heartbeat_marks_redundant_reader_removing() { clear_prom_metrics(); fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); reader_pool .background_warmup( &shard_1, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap() .await .unwrap(); must_warmed_up(&shard_1).await; let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: shard_1.seq.load(std::sync::atomic::Ordering::SeqCst), }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_keep.len(), 0); assert_eq!(resp.shards_to_remove.len(), 1); assert_eq!(resp.shards_to_remove[0].shard_id, shard_1.shard_id); assert_eq!(shard_1.get_reader_addrs().await, vec!["addr0".to_string()]); assert_eq!( shard_1.get_removing_reader_addrs().await, vec!["addr1".to_string()] ); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_heartbeat_keeps_highest_seq_anchor_when_balanced_target_lags() { clear_prom_metrics(); fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![], )) .await .unwrap(); let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); let shard_2 = manager.get_shard_by_id(&2).await.unwrap().clone(); let addr0_reader = reader_pool .get_all_readers() .await .into_iter() .find(|reader| reader.addr == "addr0") .unwrap(); assert!( shard_1 .add_ready_reader_with_observed_seq( "addr0".to_string(), "addr0_0".to_string(), Some(shard_1.seq.load(std::sync::atomic::Ordering::SeqCst)), ) .await ); addr0_reader.add_shard_cache(shard_1.clone()).await; addr0_reader.add_shard_cache(shard_2).await; let lagging_resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: shard_1.seq.load(std::sync::atomic::Ordering::SeqCst) - 1, }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(lagging_resp.shards_to_remove.len(), 0); assert_eq!(lagging_resp.shards_to_keep.len(), 1); let mut addrs = shard_1.get_reader_addrs().await; addrs.sort(); assert_eq!(addrs, vec!["addr0".to_string(), "addr1".to_string()]); assert!(shard_1.get_removing_reader_addrs().await.is_empty()); let caught_up_resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: shard_1.seq.load(std::sync::atomic::Ordering::SeqCst), }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(caught_up_resp.shards_to_remove.len(), 0); assert_eq!(caught_up_resp.shards_to_keep.len(), 1); assert_eq!(shard_1.get_reader_addrs().await, vec!["addr1".to_string()]); assert_eq!( shard_1.get_removing_reader_addrs().await, vec!["addr0".to_string()] ); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_batch_reconcile_uses_planned_load_to_avoid_over_cleanup() { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![], )) .await .unwrap(); let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); let shard_2 = manager.get_shard_by_id(&2).await.unwrap().clone(); let shard_3 = manager.get_shard_by_id(&3).await.unwrap().clone(); shard_1 .add_ready_reader_with_observed_seq("addr0".to_string(), "addr0_0".to_string(), Some(5)) .await; shard_2 .add_ready_reader_with_observed_seq( "addr0".to_string(), "addr0_0".to_string(), Some(10), ) .await; let addr0_reader = reader_pool .get_all_readers() .await .into_iter() .find(|reader| reader.addr == "addr0") .unwrap(); addr0_reader.add_shard_cache(shard_1.clone()).await; addr0_reader.add_shard_cache(shard_2.clone()).await; addr0_reader.add_shard_cache(shard_3).await; let addr1_reader = reader_pool .get_all_readers() .await .into_iter() .find(|reader| reader.addr == "addr1") .unwrap(); addr1_reader.add_shard_cache(shard_1.clone()).await; addr1_reader.add_shard_cache(shard_2.clone()).await; let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ ReaderCachedShard { shard_id: shard_1.shard_id, seq: 5, }, ReaderCachedShard { shard_id: shard_2.shard_id, seq: 9, }, ], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_keep.len(), 1); assert_eq!(resp.shards_to_remove.len(), 1); assert_eq!(shard_1.get_reader_addrs().await, vec!["addr1".to_string()]); assert_eq!( shard_1.get_removing_reader_addrs().await, vec!["addr0".to_string()] ); assert_eq!(shard_2.get_reader_addrs().await, vec!["addr0".to_string()]); assert_eq!( shard_2.get_removing_reader_addrs().await, vec!["addr1".to_string()] ); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_heartbeat_skips_redundant_cleanup_when_ready_reader_has_unknown_seq() { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![], )) .await .unwrap(); let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); assert!( shard_1 .add_ready_reader("addr0".to_string(), "addr0_0".to_string()) .await ); let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: 7, }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_remove.len(), 0); assert_eq!(resp.shards_to_keep.len(), 1); let mut addrs = shard_1.get_reader_addrs().await; addrs.sort(); assert_eq!(addrs, vec!["addr0".to_string(), "addr1".to_string()]); assert!(shard_1.get_removing_reader_addrs().await.is_empty()); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_heartbeat_cleanup_ignores_blocked_and_expired_readers() { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::millis(500); let Test { reader_pool, manager, .. } = setup_tests(config, false).await; tokio::time::sleep(Duration::from_millis(1200)).await; reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![], )) .await .unwrap(); reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr2".to_string(), "addr2_0".to_string(), vec![], )) .await .unwrap(); reader_pool.block_readers(&["addr2".to_string()]).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); shard_1.set_desired_replica_count(1).await; assert!( shard_1 .add_ready_reader_with_observed_seq( "addr0".to_string(), "addr0_0".to_string(), Some(20), ) .await ); assert!( shard_1 .add_ready_reader_with_observed_seq( "addr2".to_string(), "addr2_0".to_string(), Some(10), ) .await ); let addr0_reader = reader_pool .get_all_readers() .await .into_iter() .find(|reader| reader.addr == "addr0") .unwrap(); addr0_reader.add_shard_cache(shard_1.clone()).await; let addr2_reader = reader_pool .get_all_readers() .await .into_iter() .find(|reader| reader.addr == "addr2") .unwrap(); addr2_reader.add_shard_cache(shard_1.clone()).await; let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr1".to_string(), "addr1_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: 5, }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_keep.len(), 1); assert_eq!(resp.shards_to_remove.len(), 0); assert_eq!(shard_1.get_reader_addrs().await, vec!["addr1".to_string()]); assert_eq!( shard_1.get_removing_reader_addrs().await, vec!["addr0".to_string(), "addr2".to_string()] ); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_manual_remove_replica_marks_removing_until_heartbeat_clears_it() { clear_prom_metrics(); fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.manual_remove_wait_timeout = ReadableDuration::secs(1); config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); reader_pool .background_warmup( &shard_1, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap() .await .unwrap(); must_warmed_up(&shard_1).await; let shard_id = shard_1.shard_id; let manager_clone = manager.clone(); let reader_pool_clone = reader_pool.clone(); let remove_handle = tokio::spawn(async move { ShardCacheManager::on_msg_manual_remove_replica_by_shard( &manager_clone, reader_pool_clone, shard_id, "addr0".to_string(), ) .await }); tokio::time::timeout(Duration::from_millis(200), async { loop { if shard_1.is_reader_removing("addr0").await { break; } tokio::task::yield_now().await; } }) .await .unwrap(); assert_eq!(shard_1.desired_replica_count().await, 0); assert_eq!(shard_1.get_reader_addrs().await.len(), 0); assert!(shard_1.is_reader_removing("addr0").await); let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request_with_cached_shards( "addr0".to_string(), "addr0_0".to_string(), vec![ReaderCachedShard { shard_id: shard_1.shard_id, seq: shard_1.seq.load(std::sync::atomic::Ordering::SeqCst), }], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_keep.len(), 0); assert_eq!(resp.shards_to_remove.len(), 1); let resp = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request("addr0".to_string(), "addr0_0".to_string(), vec![]), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(resp.shards_to_keep.len(), 0); assert_eq!(resp.shards_to_remove.len(), 0); let (status, _) = remove_handle.await.unwrap().unwrap(); assert_eq!(status, ErrorCode::Success); assert_eq!(shard_1.desired_replica_count().await, 0); assert!(!shard_1.has_reader_addr(&"addr0".to_string()).await); assert_eq!(shard_1.get_removing_reader_addrs().await.len(), 0); } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_manual_remove_replica_times_out_without_cleanup_ack() { clear_prom_metrics(); fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.manual_remove_wait_timeout = ReadableDuration::millis(50); config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, .. } = setup_tests(config, false).await; let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); reader_pool .background_warmup( &shard_1, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap() .await .unwrap(); must_warmed_up(&shard_1).await; let (status, msg) = ShardCacheManager::on_msg_manual_remove_replica_by_shard( &manager, reader_pool.clone(), shard_1.shard_id, "addr0".to_string(), ) .await .unwrap(); assert_eq!(status, ErrorCode::TryAgain); assert!(msg.contains("remove_failed=timeout")); assert_eq!(shard_1.desired_replica_count().await, 0); assert!(shard_1.is_reader_removing("addr0").await); } async fn test_split_warmup_impl(optimize_for_fast_scan: bool, auto_merge_range: bool) { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.optimize_for_fast_scan = optimize_for_fast_scan; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, index_info, shards, warmup_controller, } = setup_tests(config, auto_merge_range).await; let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key), b"c".to_vec()].concat(), end_key: [table_start_key(&index_key), b"h".to_vec()].concat(), }]; let (res, ..) = manager .find_overlapping_shards( &index_key, key_ranges.as_slice(), AllocateReaderReason::ReadRequest, None, ) .await .unwrap(); let to_split = res.first().cloned().unwrap(); // ===> Mock a warmup on read request must_success_warmed_up(&to_split, &index_key, manager.clone(), reader_pool.clone()).await; assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 3 ); // ===> Mock a split on warmed shard // During this time gap, the newly splitted shards are not warmed up, so we can // access them. // TODO(P2) Make this controllable, in warmup manager. warmup_controller .set_fallback_action(Some(WarmupTestAction::SucceedAfter(Duration::from_millis( 1500, )))) .await; let (tx_12, rx_12) = tokio::sync::oneshot::channel(); warmup_controller .set_shard_actions( 12, Some(WarmupTestAction::SucceedCond(Arc::new(RwLock::new(Some( rx_12, ))))), ) .await; ShardCacheManager::handle_shard_split( &manager, index_info.clone(), shards.first().cloned().unwrap(), vec![ create_test_shard_info2(&index_key, 11, b"", b"d", 1), create_test_shard_info2(&index_key, 12, b"d", b"f", 1), ], reader_pool.clone(), ) .await; assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 5 ); // Now is ^..d..f..k..$ let check = async |to_split_shard_count: usize, expected: Vec, auto_merge_range: bool| -> Vec { // Splitted at `d`. let s0 = expected[0]; let s1 = expected[1]; let expected_ends = if auto_merge_range { if s0 == s1 { vec![b"f".to_vec(), b"k".to_vec()] } else { vec![b"d".to_vec(), b"f".to_vec(), b"k".to_vec()] } } else { vec![b"d".to_vec(), b"f".to_vec(), b"k".to_vec()] }; let merged_expected = if auto_merge_range && s0 == s1 { vec![s0, expected[2]] } else { expected }; let res = check_find_over_shards( manager.clone(), &index_key, key_ranges.as_slice(), expected_ends, merged_expected, ) .await; let split_0 = res.first().cloned().unwrap(); if auto_merge_range { assert_eq!(split_0.shard_id, s0); } else { let split_1 = res.get(1).cloned().unwrap(); assert_eq!(split_0.shard_id, s0); assert_eq!(split_1.shard_id, s1); } // ===> Mock a heartbeat w/o split_0, but very big timeout let p = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![], ), Some(ReadableDuration::days(100)), ) .await .unwrap(); // There could be some asynchronous problems, so we consider the shard is still // ready. Note, the `to_split` will still exist, unless it is removed by // `shard_map`. assert_eq!( p.shards_to_keep.len(), 2 + to_split_shard_count, "{:?}", p.shards_to_keep ); assert_eq!( p.shards_to_keep .iter() .find(|x| x.shard_id == s0) .iter() .len(), 1 ); assert!(split_0.has_ready_readers().await); res }; let new_split_0 = manager.get_shard_by_id(&11).await.unwrap().clone(); assert_eq!(new_split_0.shard_id, 11); assert!(!new_split_0.has_ready_readers().await); let new_split_1 = manager.get_shard_by_id(&12).await.unwrap().clone(); assert_eq!(new_split_1.shard_id, 12); assert!(!new_split_1.has_ready_readers().await); // Shards are allocated to read nodes, but not warmed up(ready) yet. must_finally(async || { !reader_pool.get_shard_assigned(&11).await.is_empty() && !reader_pool.get_shard_assigned(&12).await.is_empty() }) .await; // Because both of the new shards are not warmed up, so they are not replaced. let _ = check(1, vec![1, 1, 2], auto_merge_range).await; // ===> Mock a heartbeat w/o split_0, but very small timeout let p = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![to_split.shard_id], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); // If theses failed, increase the warmup wait time. assert!(!new_split_0.has_ready_readers().await); assert!(!new_split_1.has_ready_readers().await); assert_eq!(p.shards_to_keep.len(), 3, "{:?}", p.shards_to_keep); // After the mocked wait, the new shards are ready. must_warmed_up(&new_split_0).await; if optimize_for_fast_scan { let _ = check(1, vec![1, 1, 2], auto_merge_range).await; } else { let _ = check(1, vec![11, 1, 2], auto_merge_range).await; } tx_12.send(()).unwrap(); must_warmed_up(&new_split_1).await; assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 4 ); warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; // ===> After both the shards have been warmed, find_overlapping_shards will // return new shards. must_finally(async || manager.get_shards_count(&index_key).await == 4).await; // Currently, all the shards have readers. check(0, vec![11, 12, 2], auto_merge_range).await; // Now, all shard has no readers. let p = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![to_split.shard_id], ), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!(p.shards_to_keep.len(), 0, "{:?}", p.shards_to_keep); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_split_warmup_replace() { test_split_warmup_impl(false, false).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_split_warmup_whole_replace() { test_split_warmup_impl(true, false).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_split_warmup_wholemerge_replace() { test_split_warmup_impl(true, true).await; } async fn test_split_warmup_while_warming_impl(optimize_for_fast_scan: bool, test_remedy: bool) { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.optimize_for_fast_scan = optimize_for_fast_scan; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, index_info, shards, warmup_controller, } = setup_tests(config, false).await; warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; if test_remedy { warmup_controller.set_warm_derived_while_parent_warming(false); } let to_split = manager.get_shard_by_id(&1).await.unwrap().clone(); let (tx_1, rx_1) = tokio::sync::oneshot::channel(); warmup_controller .set_shard_actions( 1, Some(WarmupTestAction::SucceedCond(Arc::new(RwLock::new(Some( rx_1, ))))), ) .await; let shard_1_warmup = reader_pool .background_warmup( &to_split, reader_pool.clone(), create_test_table_info(&index_key), create_test_index_info(&index_key), AllocateReaderReason::ReadRequest, ) .unwrap(); // Must in warming state, and is not ready. must_finally(async || { !to_split.has_ready_readers().await && to_split .in_warming .load(std::sync::atomic::Ordering::SeqCst) == 1 }) .await; let (tx_11, rx_11) = tokio::sync::oneshot::channel(); warmup_controller .set_shard_actions( 11, Some(WarmupTestAction::SucceedCond(Arc::new(RwLock::new(Some( rx_11, ))))), ) .await; let (tx_12, rx_12) = tokio::sync::oneshot::channel(); warmup_controller .set_shard_actions( 12, Some(WarmupTestAction::SucceedCond(Arc::new(RwLock::new(Some( rx_12, ))))), ) .await; ShardCacheManager::handle_shard_split( &manager, index_info.clone(), shards.first().cloned().unwrap(), vec![ create_test_shard_info2(&index_key, 11, b"", b"d", 1), create_test_shard_info2(&index_key, 12, b"d", b"f", 1), ], reader_pool.clone(), ) .await; // Must not warming, and not ready. let split_11 = manager.get_shard_by_id(&11).await.unwrap().clone(); let split_12 = manager.get_shard_by_id(&12).await.unwrap().clone(); if test_remedy { assert!(!split_11.has_ready_readers().await); assert!(!split_12.has_ready_readers().await); tokio::time::sleep(Duration::from_millis(1000)).await; assert_eq!( split_11 .in_warming .load(std::sync::atomic::Ordering::SeqCst), 0 ); assert_eq!( split_12 .in_warming .load(std::sync::atomic::Ordering::SeqCst), 0 ); } else { tokio::time::sleep(Duration::from_millis(1000)).await; assert_eq!( split_11 .in_warming .load(std::sync::atomic::Ordering::SeqCst), 1 ); assert_eq!( split_12 .in_warming .load(std::sync::atomic::Ordering::SeqCst), 1 ); } tx_1.send(()).unwrap(); shard_1_warmup.await.unwrap(); must_warmed_up(&to_split).await; let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key), b"c".to_vec()].concat(), end_key: [table_start_key(&index_key), b"h".to_vec()].concat(), }]; let _ = manager .find_overlapping_shards( &index_key, key_ranges.as_slice(), AllocateReaderReason::ReadRequest, Some(reader_pool.clone()), ) .await .unwrap(); must_finally(async || { split_11 .in_warming .load(std::sync::atomic::Ordering::SeqCst) == 1 }) .await; must_finally(async || { split_12 .in_warming .load(std::sync::atomic::Ordering::SeqCst) == 1 }) .await; tx_11.send(()).unwrap(); tx_12.send(()).unwrap(); must_warmed_up(&split_11).await; must_warmed_up(&split_12).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_split_warmup_not_ready_slow_noremedy() { test_split_warmup_while_warming_impl(false, false).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_split_warmup_not_ready_fast_noremedy() { test_split_warmup_while_warming_impl(true, false).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_split_warmup_not_ready_slow_remedy() { test_split_warmup_while_warming_impl(false, true).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_split_warmup_not_ready_fast_remedy() { test_split_warmup_while_warming_impl(true, true).await; } async fn test_merge_warmup_impl(optimize_for_fast_scan: bool, auto_merge_range: bool) { clear_prom_metrics(); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.optimize_for_fast_scan = optimize_for_fast_scan; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let Test { index_key, reader_pool, manager, index_info, shards, warmup_controller, } = setup_tests(config, auto_merge_range).await; drop(shards); warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 3 ); // 1 + 2 ==> 11 ShardCacheManager::handle_shard_merge( &manager, index_info.clone(), vec![ create_test_shard_info2(&index_key, 1, b"", b"f", 1), create_test_shard_info2(&index_key, 2, b"f", b"k", 1), ], create_test_shard_info2(&index_key, 11, b"", b"k", 1), reader_pool.clone(), ) .await; assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 4 ); let merged = manager.get_shard_by_id(&11).await.unwrap().clone(); must_not_warmed_up(&merged).await; let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key), b"c".to_vec()].concat(), end_key: [table_start_key(&index_key), b"h".to_vec()].concat(), }]; // ^..f..k..$ // 11 -> ^..f..k // 3 -> k..$ if auto_merge_range { let _ = check_find_over_shards_ex( manager.clone(), &index_key, key_ranges.as_slice(), vec![b"k".to_vec()], vec![11], Some(vec![false]), None, ) .await; } else { let _ = check_find_over_shards_ex( manager.clone(), &index_key, key_ranges.as_slice(), vec![b"f".to_vec(), b"k".to_vec()], vec![11, 11], Some(vec![false, false]), None, ) .await; } // Forcely set below layers to be warmed up. // Note, this is not consistent, because we doesn't handle read_pool for // convenience. let shard_1 = manager.get_shard_by_id(&1).await.unwrap().clone(); // Manually warmup shard 1 shard_1 .add_ready_reader("addr0".to_string(), "addr0_0".to_string()) .await; let _ = check_find_over_shards_ex( manager.clone(), &index_key, key_ranges.as_slice(), vec![b"f".to_vec(), b"k".to_vec()], vec![1, 11], Some(vec![true, false]), None, ) .await; // When the shards are merged, and the merged shard is ready, // a coalesce will happen, thus reduce the returned shard count. manager .find_overlapping_shards( &index_key, key_ranges.as_slice(), // Have to use `Merge` reason. AllocateReaderReason::Merge, Some(reader_pool.clone()), ) .await .unwrap(); must_warmed_up(&merged).await; let _ = check_find_over_shards_ex( manager.clone(), &index_key, key_ranges.as_slice(), vec![b"k".to_vec()], vec![11], Some(vec![true]), Some(reader_pool.clone()), ) .await; assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 2 ); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_merge_warmup_replace() { test_merge_warmup_impl(false, false).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_merge_warmup_whole_replace() { test_merge_warmup_impl(true, false).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_merge_warmup_wholemerge_replace() { test_merge_warmup_impl(true, true).await; } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_add_remove_shard() { clear_prom_metrics(); // We don't test the `do_replace` in this test. fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let index_key = IndexKey::new(1, 1, 1); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let manager = ShardCacheManager::new(&config, false); let warmup_controller = Arc::new(WarmupTestController::new()); warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; let warmup_mgr = WarmupMgr::with_test_controller(warmup_controller.clone()); let mut reader_pool = ReaderPool::new_with_warmup_mgr(&config, warmup_mgr); reader_pool.init(manager.clone()); reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![], )) .await .unwrap(); let reader_pool = Arc::new(reader_pool); let index_info = ShardEventIndexInfo { table_info: create_test_table_info(&index_key), index_info: create_test_index_info(&index_key), }; let init_shard = create_test_shard_info2(&index_key, 1, b"b", b"y", 1); manager .handle_shard_add(index_info.clone(), &init_shard) .await .unwrap(); assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 1 ); let init_shard2 = create_test_shard_info2(&index_key, 2, b"a", b"z", 1); assert!( manager .handle_shard_add(index_info.clone(), &init_shard2) .await .is_err() ); let init_shard3 = create_test_shard_info2(&index_key, 3, b"g", b"h", 1); assert!( manager .handle_shard_add(index_info.clone(), &init_shard3) .await .is_err() ); let uni1 = manager.get_shard_by_id(&1).await.unwrap(); warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; must_success_warmed_up(&uni1, &index_key, manager.clone(), reader_pool.clone()).await; manager .handle_shard_remove(index_info, vec![1], reader_pool.clone()) .await .unwrap(); assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 0 ); assert!(manager.get_shard_by_id(&1).await.is_none()); assert_eq!( manager .get_index_meta_cloned(&index_key) .await .unwrap() .shard_len(), 0 ); let p = reader_pool .on_msg_reader_heartbeat( create_test_heartbeat_request("addr0".to_string(), "addr0_0".to_string(), vec![1]), Some(ReadableDuration::ZERO), ) .await .unwrap(); assert_eq!( p.shards_to_remove, vec![ShardManifestHeader { shard_id: 1, ..Default::default() }] ); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_warmup_rate_limiter() { clear_prom_metrics(); // We don't test the `do_replace` in this test. fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.strategy = "balanced-count".to_string(); config.cache_scheduling.allocate_parallelism = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let test = setup_tests(config, false).await; let Test { index_key, reader_pool, manager, index_info, shards, warmup_controller, } = test.clone(); drop(index_info); drop(shards); warmup_controller .set_fallback_action(Some(WarmupTestAction::SucceedAfter(Duration::from_millis( 500, )))) .await; let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key), b"".to_vec()].concat(), end_key: [table_end_key(&index_key), b"".to_vec()].concat(), }]; manager .get_or_allocate_cache( &index_key, key_ranges.clone(), reader_pool.clone(), 0, AllocateReaderReason::ReadRequest, Some(Duration::from_millis(2000)), ) .await .unwrap(); tokio::time::sleep(std::time::Duration::from_millis(500)).await; assert!(SHARD_CACHE_WARM_UP_QUEUE_HISTOGRAM.get_sample_count() > 0); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_get_readers_for_schedule_filters_blocked_and_expired() { init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.strategy = "balanced-count".to_string(); config.reader_pool.ttl_duration = ReadableDuration::millis(500); let manager = ShardCacheManager::new(&config, false); let mut reader_pool = ReaderPool::new(&config); reader_pool.init(manager); reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr_expired".to_string(), "addr_expired_0".to_string(), vec![], )) .await .unwrap(); // Ensure `addr_expired` becomes expired when `ttl_duration == 0`. tokio::time::sleep(Duration::from_millis(1200)).await; reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr_active".to_string(), "addr_active_0".to_string(), vec![], )) .await .unwrap(); reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr_blocked".to_string(), "addr_blocked_0".to_string(), vec![], )) .await .unwrap(); reader_pool .block_readers(&["addr_blocked".to_string()]) .await; let mut readers = reader_pool .get_readers_for_schedule() .await .into_iter() .map(|r| r.addr.clone()) .collect::>(); readers.sort(); assert_eq!(readers, vec!["addr_active".to_string()]); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_balanced_count_scheduler_avoids_loaded_reader() { // We don't test the `do_replace` in this test. fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let index_key = IndexKey::new(1, 1, 1); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.cache_scheduling.strategy = "balanced-count".to_string(); config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let warmup_controller = Arc::new(WarmupTestController::new()); warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; let warmup_mgr = WarmupMgr::with_test_controller(warmup_controller); let manager = ShardCacheManager::new(&config, false); let mut reader_pool = ReaderPool::new_with_warmup_mgr(&config, warmup_mgr); reader_pool.init(manager.clone()); let r0 = reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![], )) .await .unwrap(); let _r1 = reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr1".to_string(), "addr1_0".to_string(), vec![], )) .await .unwrap(); let _r2 = reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr2".to_string(), "addr2_0".to_string(), vec![], )) .await .unwrap(); let reader_pool = Arc::new(reader_pool); // Pre-assign some dummy shards to addr0, so the balanced scheduler should avoid // it. for shard_id in 1000..1005 { let dummy = create_test_shard_info2(&index_key, shard_id, b"a", b"b", 1); r0.add_shard_cache(UnifiedShardInfo::from_shard_info(&dummy)) .await; } let index_info = ShardEventIndexInfo { table_info: create_test_table_info(&index_key), index_info: create_test_index_info(&index_key), }; let shards = vec![ create_test_shard_info2(&index_key, 1, b"", b"f", 1), create_test_shard_info2(&index_key, 2, b"f", b"k", 1), create_test_shard_info2(&index_key, 3, b"k", b"", 1), ]; manager .handle_full_sync(vec![index_info], shards) .await .unwrap(); for shard_id in 1..=3 { let shard = manager.get_shard_by_id(&shard_id).await.unwrap(); must_success_warmed_up(&shard, &index_key, manager.clone(), reader_pool.clone()).await; } // Each shard should be assigned to exactly one reader, and should not be // assigned to addr0. The subsequent shards should be spread across // addr1 and addr2. let mut addr1_assigned = 0usize; let mut addr2_assigned = 0usize; for shard_id in 1..=3 { let assigned = reader_pool.get_shard_assigned(&shard_id).await; assert_eq!(assigned.len(), 1, "shard_id={}", shard_id); let addr = assigned[0].addr.as_str(); assert_ne!(addr, "addr0", "shard_id={}", shard_id); match addr { "addr1" => addr1_assigned += 1, "addr2" => addr2_assigned += 1, _ => panic!("unexpected reader addr={}, shard_id={}", addr, shard_id), } } assert!( addr1_assigned > 0 && addr2_assigned > 0, "addr1_assigned={}, addr2_assigned={}", addr1_assigned, addr2_assigned ); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_multiple_index_query() { clear_prom_metrics(); // We don't test the `do_replace` in this test. fail::cfg("on_warmup_finished_start", "return").unwrap(); defer!({ fail::remove("on_warmup_finished_start"); }); init_logger_for_ut(); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let index_key1 = IndexKey::new(1, 1, 1); let index_key2 = IndexKey::new(2, 2, 2); let warmup_controller = Arc::new(WarmupTestController::new()); warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; let warmup_mgr = WarmupMgr::with_test_controller(warmup_controller.clone()); let manager = ShardCacheManager::new(&config, false); let mut reader_pool = ReaderPool::new_with_warmup_mgr(&config, warmup_mgr); reader_pool.init(manager.clone()); reader_pool .register_or_update_reader(&create_test_heartbeat_request( "addr0".to_string(), "addr0_0".to_string(), vec![], )) .await .unwrap(); let reader_pool = Arc::new(reader_pool); let index_info1 = ShardEventIndexInfo { table_info: create_test_table_info(&index_key1), index_info: create_test_index_info(&index_key1), }; let index_info2 = ShardEventIndexInfo { table_info: create_test_table_info(&index_key2), index_info: create_test_index_info(&index_key2), }; let shards_builder: Vec<(u64, &[u8], &[u8], u64)> = vec![(1, b"", b"f", 1), (2, b"f", b"k", 1), (3, b"k", b"", 1)]; let mut shards = vec![]; for (shard_id, start_key, end_key, seq) in shards_builder.iter() { let shard = create_test_shard_info2(&index_key1, *shard_id, start_key, end_key, *seq); shards.push(shard); } for (shard_id, start_key, end_key, seq) in shards_builder.iter() { let shard = create_test_shard_info2(&index_key2, shard_id + 10, start_key, end_key, *seq); shards.push(shard); } manager .handle_full_sync( vec![index_info1.clone(), index_info2.clone()], shards.clone(), ) .await .unwrap(); // Warmup all shards warmup_controller .set_fallback_action(Some(WarmupTestAction::Succeed)) .await; let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key1), b"".to_vec()].concat(), end_key: [table_end_key(&index_key1), b"".to_vec()].concat(), }]; let r = manager .get_or_allocate_cache( &index_key1, key_ranges.clone(), reader_pool.clone(), 0, AllocateReaderReason::ReadRequest, None, ) .await .unwrap() .unwrap_left(); assert_eq!(r.len(), 3); let r = manager .get_or_allocate_cache( &index_key2, key_ranges.clone(), reader_pool.clone(), 0, AllocateReaderReason::ReadRequest, None, ) .await .unwrap() .unwrap_left(); assert_eq!(r.len(), 0); } #[allow(clippy::identity_op)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_warmup_no_reader() { clear_prom_metrics(); init_logger_for_ut(); let index_key = IndexKey::new(1, 1, 1); let mut config = make_default_test_meta_config(false); config.cache_scheduling.max_retry_attempts = 1; config.reader_pool.ttl_duration = ReadableDuration::secs(10000); let manager = ShardCacheManager::new(&config, false); let mut reader_pool = ReaderPool::new(&config); reader_pool.init(manager.clone()); let reader_pool = Arc::new(reader_pool); let index_info = ShardEventIndexInfo { table_info: create_test_table_info(&index_key), index_info: create_test_index_info(&index_key), }; let init_shard = create_test_shard_info2(&index_key, 1, b"b", b"y", 1); manager .handle_shard_add(index_info.clone(), &init_shard) .await .unwrap(); assert_eq!( SHARD_CACHE_COUNT_GAUGE .get_metric_with_label_values(&["reg"]) .unwrap() .get(), 1 ); let key_ranges = vec![tici_grpc::KeyRange { start_key: [table_start_key(&index_key), b"c".to_vec()].concat(), end_key: [table_start_key(&index_key), b"h".to_vec()].concat(), }]; let (_, b, is_no_reader) = manager .find_overlapping_shards( &index_key, key_ranges.as_slice(), AllocateReaderReason::ReadRequest, Some(reader_pool.clone()), ) .await .unwrap(); assert_eq!(b, 1); assert!(is_no_reader); assert!( SHARD_CACHE_EVENT_COUNTER .with_label_values(&["warm_failed_no_reader"]) .get() > 0 ); } }