use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use anyhow::{Result, anyhow}; use itertools::Itertools; use log::info; use crate::{ HandleRange, RangeBound, meta::{metrics::SHARD_CACHE_COUNT_GAUGE, reader::UnifiedShardInfo}, }; #[derive(Clone, Debug)] pub struct ShardAllocMap { inner: BTreeMap>, pub shards_by_id: HashMap, pub tag: String, optimize_for_fast_scan: bool, } #[derive(Debug)] pub struct QueryResult { pub shard: either::Either, pub valid_range: HandleRange, pub is_trimed: bool, } impl ShardAllocMap { pub fn scan_optimized(&self) -> bool { self.optimize_for_fast_scan } fn union_ranges(ranges: &[&HandleRange]) -> Option { let first = ranges.first()?; let mut min_start = first.sortable_start_ref(); let mut max_end = first.sortable_end_ref(); for r in ranges.iter().skip(1) { if r.sortable_start_ref() < min_start { min_start = r.sortable_start_ref(); } if r.sortable_end_ref() > max_end { max_end = r.sortable_end_ref(); } } Some(HandleRange::new_from_sorted( min_start.clone(), max_end.clone(), )) } pub fn new_with_bootstrap( shard: UnifiedShardInfo, tag: String, optimize_for_fast_scan: bool, ) -> Self { let mut shards_by_id = HashMap::new(); shards_by_id.insert(shard.shard_id, shard.clone()); let mut inner = BTreeMap::new(); let start_key = shard.start_key(); inner.insert(start_key, vec![shard]); Self { inner, shards_by_id, tag, optimize_for_fast_scan, } } pub fn new_uninitialized(tag: String, optimize_for_fast_scan: bool) -> Self { Self { inner: BTreeMap::new(), shards_by_id: HashMap::new(), tag, optimize_for_fast_scan, } } pub fn init(&mut self, v: Vec) { SHARD_CACHE_COUNT_GAUGE .with_label_values(&["reg"]) .add(v.len() as i64); for s in v.into_iter() { self.shards_by_id.insert(s.shard_id, s.clone()); self.inner.insert(s.start_key(), vec![s]); } } pub fn len(&self) -> usize { self.shards_by_id.len() } pub fn segments_len(&self) -> usize { self.inner.len() } pub fn is_empty(&self) -> bool { self.shards_by_id.is_empty() } fn find_intersecting_shards<'a>( &'a self, range: &HandleRange, ) -> Result)>> { let mut result = Vec::new(); let range_start = range.sortable_start_ref(); let range_end = range.sortable_end_ref(); let mut cur = self .inner .lower_bound(std::collections::Bound::Included(range_start)); if let Some((st, _)) = cur.peek_next() { if !st.effective_eq(range_start) { cur.prev().ok_or(anyhow!( "can't locate range start, range={:?}, found={:?}, len={}, keys={:?}, tag={}", range, st, self.segments_len(), self.inner.keys().map(|x| format!("{:?}", x)).join(":"), self.tag, ))?; } } else { cur.prev().ok_or(anyhow!( "can't find lower_bound, maybe not bootstrapped, range={:?}, tag={}", range, self.tag, ))?; } while let Some((st, shard_vec)) = cur.next() { if st >= range_end { break; } result.push((st.clone(), shard_vec)); } Ok(result) } // If a newer shard is warmed up(ready), `check_shard` can be replaced. async fn can_be_whole_replaced(&self, check_shard: &UnifiedShardInfo) -> Result { let check_ready_ranges = self.find_intersecting_shards(&check_shard.range)?; for (_, v) in check_ready_ranges { let mut found_ready_newer_shard = false; for shard in v.iter().rev() { if shard.shard_id == check_shard.shard_id { if found_ready_newer_shard { break; } else { return Ok(false); } } if shard.has_ready_readers().await { found_ready_newer_shard = true; } } // We ignore if the shard doesn't exist, however, this should not // happen. } Ok(true) } pub async fn do_whole_replace(&mut self, shard_id: &u64) -> Result> { let shard = self.shards_by_id.get(shard_id).ok_or(anyhow!( "need to whole replace shard, but new shard is not found, shard_id={}", shard_id ))?; let intersect = self.find_intersecting_shards(&shard.range)?; // Fast path: if an underlayer shard is fully covered by the new ready shard, // it's definitely replaceable by this shard. let mut fast_path_removed: HashMap = HashMap::new(); let mut underlayer_overlapped_shards = HashMap::new(); for (_, shard_vec) in intersect.iter() { let mut found_me = false; // Iterate from newer to older, in order to find all shards older than // `shard_id`. for under_shard in shard_vec.iter().rev() { if shard_id == &under_shard.shard_id { found_me = true; } else if found_me { if under_shard.range.is_subrange_of(&shard.range) { fast_path_removed .entry(under_shard.shard_id) .or_insert_with(|| under_shard.range.clone()); } else { // Boundary cases: partially overlapped shards still need the expensive // `can_be_whole_replaced` check. underlayer_overlapped_shards.insert(under_shard.shard_id, under_shard); } } } } let mut to_remove_ids = HashSet::with_capacity(fast_path_removed.len() + underlayer_overlapped_shards.len()); let mut to_remove_ranges = Vec::with_capacity(fast_path_removed.len() + underlayer_overlapped_shards.len()); for (shard_id, shard_range) in fast_path_removed { to_remove_ids.insert(shard_id); to_remove_ranges.push(shard_range); } // Check all shards in underlayer_overlapped_shards, check whether to replace. for (_, under_shard) in underlayer_overlapped_shards { if self.can_be_whole_replaced(under_shard).await? { log::trace!("can be replaced, shard={:?}, by={:?}", under_shard, shard); to_remove_ids.insert(under_shard.shard_id); to_remove_ranges.push(under_shard.range.clone()); } } let mut removed = HashSet::new(); if to_remove_ids.is_empty() { return Ok(removed); } let to_remove_range_refs = to_remove_ranges.iter().collect::>(); let iter_range = Self::union_ranges(to_remove_range_refs.as_slice()) .ok_or(anyhow!("union_ranges returns None, tag={}", self.tag))?; // Iterate only once over the maximal range covered by all shards to remove. // In each segment, we can drain to the newest shard that should be removed, // which removes all older shards in the same segment as well. for (_, vec) in self .inner .range_mut(iter_range.sortable_start_ref()..iter_range.sortable_end_ref()) { if let Some(pos) = vec .iter() .rposition(|x| to_remove_ids.contains(&x.shard_id)) { vec.drain(..=pos); } } for shard_id in to_remove_ids.into_iter() { removed.insert(shard_id); if self.shards_by_id.remove(&shard_id).is_some() { SHARD_CACHE_COUNT_GAUGE.with_label_values(&["reg"]).dec(); } } Ok(removed) } pub async fn coalesce(&mut self, range: &HandleRange) -> Vec { // Merge adjancent ranges iif all the vectors have the same element, which is // the only one. let bounds = self .inner .range(range.sortable_start_ref()..range.sortable_end_ref()) .collect::>(); let mut to_remove = vec![]; for window in bounds.windows(2) { let (_, fv) = if let Some(x) = window.first() { x } else { continue; }; let (sb, sv) = if let Some(x) = window.get(1) { x } else { continue; }; if fv.len() == 1 && sv.len() == 1 { if let Some(x) = fv.first() { if let Some(y) = sv.first() { if x.shard_id == y.shard_id && x.has_ready_readers().await { to_remove.push((*sb).clone()); } } } } } for key in to_remove.iter() { let _ = self.inner.remove(key); } to_remove } pub fn do_replace(&mut self, shard_id: &u64) -> Result> { let shard = self.shards_by_id.get(shard_id).ok_or(anyhow!( "need to replace shard, but new shard is not found, shard_id={}", shard_id ))?; let len = self.inner.len(); let mut cur = self .inner .lower_bound_mut(std::collections::Bound::Included(shard.start_key_ref())); let st = { let (st, _) = cur.peek_next().ok_or(anyhow!( "do_replace: can't find lower_bound, maybe not bootstrapped, range={:?}, tag={}", shard.range, self.tag, ))?; st.clone() }; // Because we search by shard's start key, so we do not need to `peek_prev` or // `peer_next`, this is just a consistency check. if !st.effective_eq(shard.start_key_ref()) { cur.prev().ok_or(anyhow!( "do_replace: can't locate range start, range={:?}, found={:?}, len={}, tag={}", shard.range, st, len, self.tag, ))?; } let mut removed = HashSet::new(); while let Some((st, shard_vec)) = cur.next() { if st >= shard.end_key_ref() { break; } if let Some(pos) = shard_vec.iter().position(|x| &x.shard_id == shard_id) { for x in shard_vec.drain(..pos) { removed.insert(x.shard_id); } } } // TODO(P2) Avoid this big iteration. for (_, vec) in self.inner.iter() { for v in vec { let _ = removed.remove(&v.shard_id); } } for r in removed.iter() { if self.shards_by_id.remove(r).is_some() { SHARD_CACHE_COUNT_GAUGE.with_label_values(&["reg"]).dec(); } } Ok(removed) } pub fn split( &mut self, original_range: HandleRange, original_id: u64, new_shards: Vec, ) -> Result<()> { let mut merged_bounds = self .inner .range(original_range.sortable_start_ref()..original_range.sortable_end_ref()) .map(|(x, _)| x.clone()) .collect::>(); let mut found = false; let mut new_shards_map = BTreeMap::new(); for shard in new_shards.iter() { let start = shard.range.sortable_start(); if shard.end_key_ref() == original_range.sortable_end_ref() { found = true; } new_shards_map.insert(shard.range.sortable_start(), shard.clone()); merged_bounds.insert(start); } // merged_bounds = inner + new_shards if !found { return Err(anyhow!( "can't split, end_key not match, original_range={:?}, original_shard_id={:?}, new_shards={:?}, tag={}", original_range, original_id, new_shards, self.tag, )); } if let Some(first) = merged_bounds.first() { if first != original_range.sortable_start_ref() { return Err(anyhow!( "can't split, origin node mismatch, first={:?}, original_shard_id={:?}, tag={}", first, original_id, self.tag, )); } if let Some((x, _)) = new_shards_map.first_key_value() { if first != x { return Err(anyhow!( "can't split, new node range mismatch, first={:?}, original_shard_id={:?}, found={:?}, tag={}", first, original_id, x, self.tag, )); } } else { return Err(anyhow!( "can't split, new node is empty, first={:?}, original_shard_id={:?}, tag={}", first, original_id, self.tag, )); } } else { return Err(anyhow!( "split on empty range, original_shard_id={:?}, tag={}", original_id, self.tag, )); } // All checks are performed, and all errors later are unreachable. // Don't remove the old shard here, because new shards are not warmed up. SHARD_CACHE_COUNT_GAUGE .with_label_values(&["reg"]) .add(new_shards.len() as i64); for shard in new_shards.into_iter() { self.shards_by_id.insert(shard.shard_id, shard); } // Do split for bound in merged_bounds.iter() { if let Some((found, v)) = self.inner.range(..=bound.clone()).last() { if found != bound { // Need to split self.inner.insert(bound.clone(), v.clone()); } // Else, the bound already exists, do not touch. } else { return Err(anyhow!( "can't find new node to split, original_shard_id={:?}, original_range={:?}, bound={:?}, tag={}", original_id, original_range, bound, self.tag, )); } } // Add new shards for (bound, v) in self .inner .range_mut(original_range.sortable_start_ref()..original_range.sortable_end_ref()) { if let Some((_, to_add_shard)) = new_shards_map.range(..=bound.clone()).last() { v.push(to_add_shard.clone()); } else { return Err(anyhow!( "can't find new node to insert, original_shard_id={:?}, bound={:?}, tag={}", original_id, bound, self.tag, )); } } Ok(()) } pub fn merge(&mut self, new_shard: UnifiedShardInfo) -> Result<()> { let mut found = false; for (i, (bound_start, v)) in self .inner .range(new_shard.range.sortable_start_ref()..new_shard.range.sortable_end_ref()) .enumerate() { #[allow(clippy::collapsible_if)] if i == 0 { if new_shard.range.sortable_start_ref() != bound_start { return Err(anyhow!( "boundaies must match when merge, new_shard={}, found={:?}, tag={}", new_shard.to_long_ident(), bound_start, self.tag, )); } } if let Some(x) = v.last() { if x.range.sortable_end_ref() == new_shard.end_key_ref() { found = true; } } } if !found { return Err(anyhow!( "can't merge, end_key not match, new_shard_range={:?}, new_shard_id={:?}, tag={}", new_shard.range, new_shard.shard_id, self.tag, )); } for (_, v) in self .inner .range_mut(new_shard.range.sortable_start_ref()..new_shard.range.sortable_end_ref()) { v.push(new_shard.clone()); } // Don't remove the old shard here, because new shards are not warmed up. SHARD_CACHE_COUNT_GAUGE.with_label_values(&["reg"]).inc(); self.shards_by_id.insert(new_shard.shard_id, new_shard); Ok(()) } // Currently only used in handling `bootstrap_shard`. pub fn add_shard(&mut self, new_shard: UnifiedShardInfo) -> Result<()> { if let Some(v) = self.inner.get(new_shard.start_key_ref()) { if v.iter() .filter(|x| x.shard_id == new_shard.shard_id) .count() > 0 { log::warn!( "add_shard: shard already added, shard={}, tag={}", new_shard.to_long_ident(), self.tag, ); return Ok(()); } else { return Err(anyhow!( "can't add_shard, already inited, new_shard={:?}, self={:?}, tag={}", new_shard, self.inner, self.tag, )); } } if self.inner.is_empty() { self.inner .insert(new_shard.range.sortable_start(), vec![new_shard.clone()]); } else { return Err(anyhow!( "can't add_shard, inconsistent detected, new_shard={:?}, self={:?}, tag={}", new_shard, self.inner, self.tag, )); } SHARD_CACHE_COUNT_GAUGE.with_label_values(&["reg"]).inc(); self.shards_by_id.insert(new_shard.shard_id, new_shard); Ok(()) } pub fn remove_shards(&mut self, shard_ids: &[u64]) -> Result { let to_remove: HashSet = shard_ids.iter().copied().collect(); for (_, v) in self.inner.iter_mut() { v.retain(|shard| !to_remove.contains(&shard.shard_id)); } let mut removed = 0; for s in shard_ids.iter() { if self.shards_by_id.remove(s).is_some() { removed += 1; } } SHARD_CACHE_COUNT_GAUGE .with_label_values(&["reg"]) .sub(removed as i64); Ok(removed) } // Left is for ready, Right is for not ready pub async fn query( &self, range: &HandleRange, ) -> Result<(Vec, Vec)> { let intersect = self.find_intersecting_shards(range)?; let mut result: Vec = vec![]; let mut to_remedy_warmup = vec![]; for (index, (st, shard_vec)) in intersect.iter().enumerate() { let mut has_ready = false; let mut handle_mid = |s: &UnifiedShardInfo, is_ready: bool| { if is_ready { let valid_end = if let Some((next_shard_start, _)) = intersect.get(index + 1) { next_shard_start.clone() } else { s.end_key() }; let is_trimed = valid_end != s.end_key(); result.push(QueryResult { shard: either::Left(s.clone()), valid_range: HandleRange::new_from_sorted(st.clone(), valid_end), is_trimed, }); return true; } false }; if self.scan_optimized() { // From old to new for s in shard_vec.iter() { let is_ready = s.has_ready_readers().await; if handle_mid(s, is_ready) { has_ready = true; break; } } } else { // From new to old for s in shard_vec.iter().rev() { let is_ready = s.has_ready_readers().await; if handle_mid(s, is_ready) { has_ready = true; break; } } } if let Some(s) = shard_vec.last() { if !s.is_warming() && !s.has_ready_readers().await { info!( "need to remedy new shard warmup, {}, use={:?}", s.to_long_ident(), result.last() ); to_remedy_warmup.push(s.clone()); } } if !has_ready { let right = shard_vec.last().ok_or(anyhow!( "shard not exist, for range={:?}, tag={}", st, self.tag ))?; let valid_end = if let Some((next_shard_start, _)) = intersect.get(index + 1) { next_shard_start.clone() } else { right.end_key() }; let is_trimed = valid_end != right.end_key(); result.push(QueryResult { shard: either::Right(right.clone()), valid_range: HandleRange::new_from_sorted(st.clone(), valid_end), is_trimed, }); } } Ok((result, to_remedy_warmup)) } } #[cfg(test)] mod tests { use rand::{Rng, SeedableRng, prelude::IndexedRandom, rngs::StdRng}; use tici_common::index::IndexKey; use super::*; use crate::meta::reader::test_utils::create_test_unified_shard_info; fn get_result_map(q: &[QueryResult]) -> BTreeMap { let mut bm = BTreeMap::new(); for QueryResult { shard, .. } in q.iter() { match shard { either::Either::Left(l) => { bm.insert(l.start_key(), l.clone()); } either::Either::Right(r) => { bm.insert(r.start_key(), r.clone()); } } } bm } fn get_vec_map(m: &ShardAllocMap) -> Vec> { let mut r = vec![]; for (_, vec) in m.inner.iter() { r.push(vec.iter().map(|x| x.shard_id).collect::>()); } r } fn get_vec_range(m: &ShardAllocMap) -> Vec> { let mut res: Vec> = vec![]; for (r, _) in m.inner.iter() { res.push(r.encoded_as_start()); } res } fn check_map_consistency(m: &ShardAllocMap) { let _ = m .inner .iter() .map_windows(|x: &[(&RangeBound, &Vec); 2]| { let (_, fvec) = x.first().unwrap(); let (sb, _) = x.get(1).unwrap(); let mut found = false; for s in fvec.iter() { assert!(s.end_key_ref() >= sb); if s.end_key_ref() == *sb { found = true; } } assert!(found); }); } async fn check_query_range(m: &ShardAllocMap, s: &[u8], e: &[u8], expected: Vec) { let (res, _) = m .query(&HandleRange::new(s.to_vec(), e.to_vec())) .await .unwrap(); assert_eq!(res.len(), expected.len()); for (x, y) in res.iter().zip(expected.iter()) { match &x.shard { either::Either::Left(xx) => { assert_eq!(&xx.shard_id, y); } either::Either::Right(xx) => { assert_eq!(&xx.shard_id, y); } } } } async fn check_query_end_key(m: &ShardAllocMap, s: &[u8], e: &[u8], expected: Vec<&[u8]>) { let (res, _) = m .query(&HandleRange::new(s.to_vec(), e.to_vec())) .await .unwrap(); assert_eq!(res.len(), expected.len()); for (i, (x, y)) in res.iter().zip(expected.iter()).enumerate() { match &x.shard { either::Either::Left(xx) => { assert_eq!( &x.valid_range.encoded_end().as_slice(), y, "{} shard_id={}: {:?}..{:?} vs {:?}", i, xx.shard_id, xx.range.encoded_start(), xx.range.encoded_end(), y, ); } either::Either::Right(xx) => { assert_eq!( &x.valid_range.encoded_end().as_slice(), y, "{} shard_id={}: {:?}..{:?} vs {:?}", i, xx.shard_id, xx.range.encoded_start(), xx.range.encoded_end(), y, ); } } } } fn check_shard_vec(v: &[(RangeBound, &Vec)], shard_ids: Vec>) { assert_eq!(v.len(), shard_ids.len()); for (a, b) in v.iter().zip(shard_ids.iter()) { assert_eq!(a.1.len(), b.len()); for (x, y) in a.1.iter().zip(b.iter()) { assert_eq!(&x.shard_id, y); } } } #[tokio::test] async fn test_find_intersecting() { let mut m = ShardAllocMap::new_uninitialized(String::new(), false); m.init(vec![ create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 2, b"b", b"h", 1), create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 3, b"h", b"y", 1), ]); check_shard_vec( m.find_intersecting_shards(&HandleRange::new(b"c".to_vec(), b"d".to_vec())) .unwrap() .as_slice(), vec![vec![2]], ); check_shard_vec( m.find_intersecting_shards(&HandleRange::new(b"c".to_vec(), b"h".to_vec())) .unwrap() .as_slice(), vec![vec![2]], ); check_shard_vec( m.find_intersecting_shards(&HandleRange::new(b"c".to_vec(), b"h1".to_vec())) .unwrap() .as_slice(), vec![vec![2], vec![3]], ); assert!( m.find_intersecting_shards(&HandleRange::new(b"a".to_vec(), b"a2".to_vec())) .is_err() ); assert!( m.find_intersecting_shards(&HandleRange::new(b"a".to_vec(), b"c".to_vec())) .is_err() ); check_shard_vec( m.find_intersecting_shards(&HandleRange::new(b"x".to_vec(), b"z".to_vec())) .unwrap() .as_slice(), vec![vec![3]], ); check_shard_vec( m.find_intersecting_shards(&HandleRange::new(b"b".to_vec(), b"z".to_vec())) .unwrap() .as_slice(), vec![vec![2], vec![3]], ); check_shard_vec( m.find_intersecting_shards(&HandleRange::new(b"c".to_vec(), b"z".to_vec())) .unwrap() .as_slice(), vec![vec![2], vec![3]], ); assert!( m.find_intersecting_shards(&HandleRange::new(b"a".to_vec(), b"z".to_vec())) .is_err() ); let m: ShardAllocMap = ShardAllocMap::new_uninitialized(String::new(), false); assert!( m.find_intersecting_shards(&HandleRange::new(b"c".to_vec(), b"z".to_vec())) .is_err() ); } #[tokio::test] async fn test_alloc_map_merge_plain() { test_alloc_map_merge(false).await; } #[tokio::test] async fn test_alloc_map_merge_coalesce() { test_alloc_map_merge(true).await; } async fn test_alloc_map_merge(do_coalesce: bool) { let u: UnifiedShardInfo = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"a", b"z", 1); let mut m = ShardAllocMap::new_with_bootstrap(u.clone(), String::new(), false); m.split( u.range.clone(), u.shard_id, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 11, b"a", b"k", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 12, b"k", b"p", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 13, b"p", b"z", 1, ), ], ) .unwrap(); check_map_consistency(&m); let m_origin = m.clone(); assert!( m.merge(create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 21, b"b", b"n", 1, )) .is_err() ); assert_eq!(get_vec_range(&m_origin), get_vec_range(&m)); assert_eq!(get_vec_map(&m_origin), get_vec_map(&m)); m.merge(create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 21, b"k", b"z", 1, )) .unwrap(); if do_coalesce { let _ = m .shards_by_id .get_mut(&11) .unwrap() .add_ready_reader("addr1".to_string(), "addr1_1".to_string()) .await; let v = m .coalesce(&HandleRange::new(b"a".to_vec(), b"z".to_vec())) .await; assert_eq!(v.len(), 0); } assert_eq!( get_vec_map(&m), vec![ vec![1, 11], // a..k vec![1, 12, 21], // k..p vec![1, 13, 21], // p..z ] ); if do_coalesce { let _ = m .shards_by_id .get_mut(&21) .unwrap() .add_ready_reader("addr1".to_string(), "addr1_1".to_string()) .await; m.coalesce(&HandleRange::new(b"a".to_vec(), b"z".to_vec())) .await; assert_eq!( get_vec_map(&m), vec![ vec![1, 11], // a..k vec![1, 12, 21], // k..p vec![1, 13, 21], // p..z ] ); m.do_replace(&11).unwrap(); m.do_replace(&21).unwrap(); let v = m .coalesce(&HandleRange::new(b"a".to_vec(), b"z".to_vec())) .await; assert_eq!( get_vec_map(&m), vec![ vec![11], // a..k vec![21], // p..z ] ); assert_eq!(v.len(), 1); return; } // a..k, k..z let vr = get_vec_range(&m); assert_eq!(vr, vec![b"a".to_vec(), b"k".to_vec(), b"p".to_vec()]); check_map_consistency(&m); assert!( m.split( HandleRange::new(b"k".to_vec(), b"z".to_vec()), 21, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 31, b"a", b"t", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 32, b"k", b"z", 1, ), ], ) .is_err() ); check_map_consistency(&m); m.split( HandleRange::new(b"k".to_vec(), b"z".to_vec()), 21, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 31, b"k", b"t", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 32, b"t", b"z", 1, ), ], ) .unwrap(); check_map_consistency(&m); assert_eq!( get_vec_map(&m), vec![ vec![1, 11], // a..k vec![1, 12, 21, 31], // k..p vec![1, 13, 21, 31], // p..t vec![1, 13, 21, 32], // t..z ] ); check_query_range(&m, b"a", b"z", vec![11, 31, 31, 32]).await; check_query_end_key(&m, b"a", b"z", vec![b"k", b"p", b"t", b"z"]).await; // a..k, k..t, t..z let vr = get_vec_range(&m); assert_eq!( vr, vec![b"a".to_vec(), b"k".to_vec(), b"p".to_vec(), b"t".to_vec()] ); let m_origin = m.clone(); assert!( m.merge(create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 31, b"a", b"p", 1, ),) .is_err() ); assert_eq!(get_vec_range(&m_origin), get_vec_range(&m)); assert_eq!(get_vec_map(&m_origin), get_vec_map(&m)); m.merge(create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 41, b"a", b"t", 1, )) .unwrap(); check_map_consistency(&m); assert_eq!( get_vec_map(&m), vec![ vec![1, 11, 41], // a..k vec![1, 12, 21, 31, 41], // k..p vec![1, 13, 21, 31, 41], // p..t vec![1, 13, 21, 32], // t..z ] ); check_query_range(&m, b"a", b"z", vec![41, 41, 41, 32]).await; check_query_end_key(&m, b"a", b"z", vec![b"k", b"p", b"t", b"z"]).await; // a..t, t..z } #[tokio::test] async fn test_alloc_map_multi_split() { let u1 = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"a", b"p", 1); let u2 = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 2, b"p", b"z", 1); let mut m = ShardAllocMap::new_uninitialized(String::new(), false); m.shards_by_id.insert(u1.shard_id, u1.clone()); m.shards_by_id.insert(u2.shard_id, u2.clone()); m.inner .insert(RangeBound::make_start(b"a".to_vec()), vec![u1.clone()]); m.inner .insert(RangeBound::make_start(b"k".to_vec()), vec![u1.clone()]); m.inner .insert(RangeBound::make_start(b"p".to_vec()), vec![u2.clone()]); m.split( u1.range.clone(), u1.shard_id, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 11, b"a", b"e", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 12, b"e", b"p", 1, ), ], ) .unwrap(); check_map_consistency(&m); let vr = get_vec_range(&m); assert_eq!( vr, vec![b"a".to_vec(), b"e".to_vec(), b"k".to_vec(), b"p".to_vec()] ); let vm = get_vec_map(&m); assert_eq!( vm, vec![ vec![1, 11], // a..e vec![1, 12], // e..k vec![1, 12], // k..p vec![2], // p..z ] ); } #[tokio::test] async fn test_alloc_map_basic_operations() { let u = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"a", b"z", 1); let mut m = ShardAllocMap::new_with_bootstrap(u.clone(), String::new(), false); let (k, _) = m .query(&HandleRange::new(b"h".to_vec(), b"j".to_vec())) .await .unwrap(); assert_eq!(k.len(), 1); m.do_replace(&1).unwrap(); } #[tokio::test] async fn test_alloc_map_basic_split() { let u = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"a", b"z", 1); let mut m = ShardAllocMap::new_uninitialized(String::new(), false); assert!( m.split(HandleRange::new(b"h".to_vec(), b"k".to_vec()), 1, vec![]) .is_err() ); assert!( m.query(&HandleRange::new(b"a".to_vec(), b"z".to_vec())) .await .is_err() ); // test split and do_delete // a..z -> a..h + h..z let mut m = ShardAllocMap::new_with_bootstrap(u.clone(), String::new(), false); m.split( u.range.clone(), 1, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 2, b"a", b"h", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 3, b"h", b"z", 1, ), ], ) .unwrap(); check_map_consistency(&m); assert_eq!(m.inner.len(), 2); assert_eq!(m.shards_by_id.len(), 3); let vm = get_vec_map(&m); assert_eq!(vm, vec![vec![1, 2], vec![1, 3]]); let (q, _) = m .query(&HandleRange::new(b"a".to_vec(), b"z".to_vec())) .await .unwrap(); assert_eq!(q.len(), 2); let bm = get_result_map(q.as_slice()); assert_eq!( bm.get(&RangeBound::make_start(b"a".to_vec())) .unwrap() .shard_id, 2 ); assert_eq!( bm.get(&RangeBound::make_start(b"h".to_vec())) .unwrap() .shard_id, 3 ); m.do_replace(&2).unwrap(); let vm = get_vec_map(&m); assert_eq!(vm, vec![vec![2], vec![1, 3]]); assert_eq!(m.shards_by_id.len(), 3, "{:?}", &m.shards_by_id); m.do_replace(&3).unwrap(); let vm = get_vec_map(&m); assert_eq!(vm, vec![vec![2], vec![3]]); assert_eq!(m.shards_by_id.len(), 2, "{:?}", &m.shards_by_id); // test split with invalid bound // a..z -> b..h + h..z, should throw let mut m = ShardAllocMap::new_with_bootstrap(u.clone(), String::new(), false); let m_origin = m.clone(); assert!( m.split( u.range.clone(), 1, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 2, b"b", b"h", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 3, b"h", b"z", 1, ), ], ) .is_err() ); assert_eq!(get_vec_range(&m_origin), get_vec_range(&m)); assert_eq!(get_vec_map(&m_origin), get_vec_map(&m)); assert!( m.split( u.range.clone(), 1, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 2, b"a", b"h", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 3, b"h", b"w", 1, ), ], ) .is_err() ); assert_eq!(get_vec_range(&m_origin), get_vec_range(&m)); assert_eq!(get_vec_map(&m_origin), get_vec_map(&m)); assert!( m.split( u.range.clone(), 1, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 2, b"A", b"h", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 3, b"h", b"z", 1, ), ], ) .is_err() ); // a..b..h..z assert_eq!(get_vec_range(&m_origin), get_vec_range(&m)); assert_eq!(get_vec_map(&m_origin), get_vec_map(&m)); // test query let mut m = ShardAllocMap::new_uninitialized(String::new(), false); m.init(vec![ create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 2, b"b", b"h", 1), create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 3, b"h", b"z", 1), ]); assert!( m.query(&HandleRange::new(b"a".to_vec(), b"y".to_vec())) .await .is_err() ); check_query_range(&m, b"c", b"y", vec![2, 3]).await; check_query_end_key(&m, b"c", b"y", vec![b"h", b"z"]).await; check_query_range(&m, b"h", b"y", vec![3]).await; check_query_end_key(&m, b"h", b"y", vec![b"z"]).await; // For example, if the maintained range is [a,b), then we can't query // [inf,-inf), because the range to query is bigger. assert!( m.query(&HandleRange::new_from_sorted( RangeBound::LeftMost, RangeBound::RightMost )) .await .is_err() ); } #[tokio::test] async fn test_alloc_map_whole_replace_waits_for_all_ready() { tici_test_utils::test_utils::init_logger_for_ut(); let u = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"a", b"z", 1); let mut m = ShardAllocMap::new_with_bootstrap(u.clone(), String::new(), true); m.split( u.range.clone(), u.shard_id, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 2, b"a", b"h", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 3, b"h", b"z", 1, ), ], ) .unwrap(); let _ = m .shards_by_id .get_mut(&2) .unwrap() .add_ready_reader("addr1".to_string(), "addr1_1".to_string()) .await; let removed = m.do_whole_replace(&2).await.unwrap(); assert!(removed.is_empty()); assert!(m.shards_by_id.contains_key(&1)); assert_eq!(get_vec_map(&m), vec![vec![1, 2], vec![1, 3]]); let _ = m .shards_by_id .get_mut(&3) .unwrap() .add_ready_reader("addr2".to_string(), "addr2_1".to_string()) .await; let removed = m.do_whole_replace(&3).await.unwrap(); assert_eq!(removed, [1].into_iter().collect::>()); assert!(!m.shards_by_id.contains_key(&1)); assert_eq!(get_vec_map(&m), vec![vec![2], vec![3]]); } #[tokio::test] async fn test_alloc_map_whole_replace_removes_multiple_shards() { tici_test_utils::test_utils::init_logger_for_ut(); let u = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"a", b"z", 1); let mut m = ShardAllocMap::new_with_bootstrap(u.clone(), String::new(), true); m.split( u.range.clone(), u.shard_id, vec![ create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 2, b"a", b"h", 1, ), create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 3, b"h", b"z", 1, ), ], ) .unwrap(); m.merge(create_test_unified_shard_info( &IndexKey::new_nullspace(100, 233), 4, b"a", b"z", 2, )) .unwrap(); let _ = m .shards_by_id .get_mut(&4) .unwrap() .add_ready_reader("addr1".to_string(), "addr1_1".to_string()) .await; let removed = m.do_whole_replace(&4).await.unwrap(); assert_eq!(removed, [1, 2, 3].into_iter().collect::>(),); assert!(!m.shards_by_id.contains_key(&1)); assert!(!m.shards_by_id.contains_key(&2)); assert!(!m.shards_by_id.contains_key(&3)); assert!(m.shards_by_id.contains_key(&4)); assert_eq!(get_vec_map(&m), vec![vec![4], vec![4]]); } #[tokio::test] async fn test_alloc_map_add_shard() { let u = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"a", b"z", 1); let mut m = ShardAllocMap::new_uninitialized(String::new(), false); m.add_shard(u.clone()).unwrap(); let u1 = create_test_unified_shard_info(&IndexKey::new_nullspace(100, 233), 1, b"b", b"z", 1); assert!(m.add_shard(u1.clone()).is_err()); m.remove_shards(&[1]).unwrap(); for (_, v) in m.inner.iter() { assert!(!v.iter().any(|x| x.shard_id == 1)); } } fn extract_query_shard_id(q: &QueryResult) -> u64 { match &q.shard { either::Either::Left(s) => s.shard_id, either::Either::Right(s) => s.shard_id, } } async fn assert_full_range_invariants(m: &ShardAllocMap, full_range: &HandleRange) { let (res, _) = m.query(full_range).await.unwrap(); // Each returned shard appears only once. let mut returned: HashSet = HashSet::with_capacity(res.len()); for q in res.iter() { let shard_id = extract_query_shard_id(q); // TODO(P1) Remove duplicate let _ = returned.insert(shard_id); } // If a shard is returned, no underlayer shard of it can be returned. for q in res.iter() { let shard_id = extract_query_shard_id(q); let mut underlayer = HashSet::new(); for (_, v) in m .inner .range(q.valid_range.sortable_start_ref()..q.valid_range.sortable_end_ref()) { let pos = v .iter() .position(|s| s.shard_id == shard_id) .expect("returned shard must exist in inner"); underlayer.extend(v.iter().take(pos).map(|s| s.shard_id)); } for under_id in underlayer.into_iter() { assert!( !returned.contains(&under_id), "underlayer shard is also returned: under_id={}, by={}, res={:?}", under_id, shard_id, res ); } } // `shards_by_id` and `inner` must contain the same shard ids. let mut inner_ids = HashSet::new(); for (_, v) in m.inner.iter() { for shard in v.iter() { inner_ids.insert(shard.shard_id); } } let map_ids: HashSet = m.shards_by_id.keys().copied().collect(); assert_eq!( inner_ids, map_ids, "inconsistent shards_by_id vs inner: shards_by_id={:?}, inner={:?}", m.shards_by_id.keys().collect::>(), inner_ids ); } fn can_split_range(range: &HandleRange) -> bool { let (RangeBound::Key(s), RangeBound::Key(e)) = (range.sortable_start_ref(), range.sortable_end_ref()) else { return false; }; s.len() == 1 && e.len() == 1 && s[0] + 1 < e[0] } fn random_mid_key(rng: &mut StdRng, range: &HandleRange) -> Vec { let (RangeBound::Key(s), RangeBound::Key(e)) = (range.sortable_start_ref(), range.sortable_end_ref()) else { unreachable!("range must be a single-byte key range"); }; let mid = rng.random_range((s[0] + 1)..e[0]); vec![mid] } #[tokio::test] async fn test_alloc_map_randomized_operations_invariants() { let index_key = IndexKey::new_nullspace(100, 233); let full_range = HandleRange::new(b"a".to_vec(), b"z".to_vec()); let bootstrap = create_test_unified_shard_info(&index_key, 1, b"a", b"z", 1 /* seq */); let mut m = ShardAllocMap::new_with_bootstrap(bootstrap, "test".to_string(), true); let mut rng = StdRng::seed_from_u64(0xD16C_4A3C_51E2_A5D1); let mut next_shard_id: u64 = 2; let n = 300usize; for _ in 0..n { match rng.random_range(0..3) { // split 0 => { let (q, _) = m.query(&full_range).await.unwrap(); let candidates = q .iter() .filter(|x| !x.is_trimed && can_split_range(&x.valid_range)) .collect::>(); if let Some(to_split) = candidates.choose(&mut rng) { let mid = random_mid_key(&mut rng, &to_split.valid_range); let start = to_split.valid_range.encoded_start(); let end = to_split.valid_range.encoded_end(); log::info!( "split, {} => {} {}", extract_query_shard_id(to_split), next_shard_id, next_shard_id + 1 ); m.split( to_split.valid_range.clone(), extract_query_shard_id(to_split), vec![ create_test_unified_shard_info( &index_key, next_shard_id, &start, &mid, next_shard_id, ), create_test_unified_shard_info( &index_key, next_shard_id + 1, &mid, &end, next_shard_id + 1, ), ], ) .unwrap(); next_shard_id += 2; } } // merge 1 => { // Derive adjacent top-layer shards from `inner`, without using `query`. let segs = m .inner .range(full_range.sortable_start_ref()..full_range.sortable_end_ref()) .map(|(st, v)| { let top = v.last().expect("allocation map segment must have shards"); (st.clone(), top.shard_id) }) .collect::>(); if segs.len() >= 2 { // Merge consecutive segments with the same top-layer shard id. // This is equivalent to the full range `query` result (deduped). let mut runs: Vec<(u64, RangeBound, RangeBound)> = Vec::new(); for (i, (st, shard_id)) in segs.iter().enumerate() { let seg_end = if let Some((next_st, _)) = segs.get(i + 1) { next_st.clone() } else { full_range.sortable_end() }; if let Some((last_id, _, last_end)) = runs.last_mut() { if last_id == shard_id { *last_end = seg_end; continue; } } runs.push((*shard_id, st.clone(), seg_end)); } let mut candidates = Vec::new(); for i in 0..runs.len().saturating_sub(1) { let (l_id, l_start, l_end) = &runs[i]; let (r_id, r_start, r_end) = &runs[i + 1]; if !l_end.effective_eq(r_start) { continue; } let Some(l_shard) = m.shards_by_id.get(l_id) else { continue; }; let Some(r_shard) = m.shards_by_id.get(r_id) else { continue; }; // Only merge shards that are not partially overlapped by other newer // shards (i.e., their top-layer run matches their real range). let l_untrimmed = l_shard.start_key_ref().effective_eq(l_start) && l_shard.end_key_ref().effective_eq(l_end); let r_untrimmed = r_shard.start_key_ref().effective_eq(r_start) && r_shard.end_key_ref().effective_eq(r_end); if l_untrimmed && r_untrimmed { candidates.push(i); } } if let Some(i) = candidates.choose(&mut rng) { let start = runs[*i].1.encoded_as_start(); let end = runs[*i + 1].2.encoded_as_end(); log::info!( "merge, {} {} => {}", runs[*i].0, runs[*i + 1].0, next_shard_id ); m.merge(create_test_unified_shard_info( &index_key, next_shard_id, &start, &end, next_shard_id, )) .unwrap(); next_shard_id += 1; } } } // warm up + whole replace + coalesce 2 => { let (q, _) = m.query(&full_range).await.unwrap(); let mut not_ready = Vec::new(); for qr in q.iter() { if matches!(&qr.shard, either::Either::Right(_)) { not_ready.push(extract_query_shard_id(qr)); } } if let Some(to_warm) = not_ready.choose(&mut rng) { let shard = m.shards_by_id.get(to_warm).unwrap(); log::info!("warmup, {}", shard.shard_id); let _ = shard .add_ready_reader( format!("addr_{}", to_warm), format!("tiflash_{}", to_warm), ) .await; let _ = m.do_whole_replace(to_warm).await.unwrap(); let _ = m.coalesce(&full_range).await; } } _ => unreachable!(), } assert_full_range_invariants(&m, &full_range).await; } } }