1use std::collections::HashSet;
16use std::collections::hash_map::HashMap;
17use std::hash::Hash;
18use std::ops::{Bound, Range};
19use std::sync::Arc;
20use std::time::{Duration, Instant};
21
22use foyer::RangeBoundsExt;
23use futures::future::{join_all, try_join_all};
24use itertools::Itertools;
25use risingwave_common::bitmap::Bitmap;
26use risingwave_common::config::Role;
27use risingwave_common::config::streaming::CacheRefillPolicy;
28use risingwave_common::hash::VirtualNode;
29use risingwave_common::license::Feature;
30use risingwave_common::util::iter_util::ZipEqFast;
31use risingwave_hummock_sdk::compaction_group::hummock_version_ext::SstDeltaInfo;
32use risingwave_hummock_sdk::key::{FullKey, vnode_range};
33use risingwave_hummock_sdk::{HummockSstableObjectId, KeyComparator};
34use risingwave_pb::id::TableId;
35use thiserror_ext::AsReport;
36use tokio::sync::Semaphore;
37use tokio::task::JoinHandle;
38
39use crate::hummock::{
40 Block, HummockError, HummockResult, RecentFilterTrait, Sstable, SstableBlockIndex,
41 SstableStoreRef, TableHolder,
42};
43use crate::monitor::StoreLocalStatistic;
44pub use crate::monitor::{CacheRefillMetrics, GLOBAL_CACHE_REFILL_METRICS};
45use crate::opts::StorageOpts;
46
47#[derive(Debug)]
48pub struct CacheRefillConfig {
49 pub timeout: Duration,
51
52 pub data_refill_levels: HashSet<u32>,
54
55 pub meta_refill_concurrency: usize,
57
58 pub concurrency: usize,
60
61 pub unit: usize,
63
64 pub threshold: f64,
68
69 pub skip_recent_filter: bool,
71
72 pub skip_inheritance_filter: bool,
74
75 pub table_cache_refill_default_policy: CacheRefillPolicy,
77}
78
79impl CacheRefillConfig {
80 pub fn from_storage_opts(options: &StorageOpts) -> Self {
81 let data_refill_levels = match Feature::ElasticDiskCache.check_available() {
82 Ok(_) => options
83 .cache_refill_data_refill_levels
84 .iter()
85 .copied()
86 .collect(),
87 Err(e) => {
88 tracing::warn!(error = %e.as_report(), "ElasticDiskCache is not available.");
89 HashSet::new()
90 }
91 };
92
93 Self {
94 timeout: Duration::from_millis(options.cache_refill_timeout_ms),
95 data_refill_levels,
96 concurrency: options.cache_refill_concurrency,
97 meta_refill_concurrency: options.cache_refill_meta_refill_concurrency,
98 unit: options.cache_refill_unit,
99 threshold: options.cache_refill_threshold,
100 skip_recent_filter: options.cache_refill_skip_recent_filter,
101 skip_inheritance_filter: options.cache_refill_skip_inheritance_filter,
102 table_cache_refill_default_policy: options
103 .cache_refill_table_cache_refill_default_policy,
104 }
105 }
106}
107
108pub(crate) type SpawnRefillTask =
109 Arc<dyn Fn(Vec<SstDeltaInfo>, CacheRefillContext) -> JoinHandle<()> + Send + Sync + 'static>;
110
111pub type TableCacheRefillContextMap = HashMap<TableId, TableCacheRefillContext>;
112
113#[derive(Clone)]
116pub struct TableCacheRefillContext {
117 pub streaming_vnode_bitmap: Option<Bitmap>,
119 pub serving_vnode_bitmap: Option<Bitmap>,
121 pub policy: CacheRefillPolicy,
123}
124
125#[derive(Clone)]
129pub struct TableCacheRefillMonitorSnapshot {
130 pub contexts: TableCacheRefillContextMap,
131 pub policies: HashMap<TableId, CacheRefillPolicy>,
132 pub default_policy: CacheRefillPolicy,
133 pub streaming_table_vnode_mapping: HashMap<TableId, Bitmap>,
134 pub serving_table_vnode_mapping: HashMap<TableId, Bitmap>,
135}
136
137fn vnode_range_overlaps_bitmap(vnode_range: (usize, usize), bitmap: &Bitmap) -> bool {
138 assert!(vnode_range.0 <= vnode_range.1);
139 let start = vnode_range.0.min(bitmap.len());
140 let end = vnode_range.1.min(bitmap.len());
141 if start == end || !bitmap.any() {
142 return false;
143 }
144 if bitmap.all() {
145 return true;
146 }
147 (start..end).any(|vnode| bitmap.is_set(vnode))
148}
149
150impl TableCacheRefillContext {
151 fn allows_normal_data_refill_block(&self, sstable: &Sstable, block_index: usize) -> bool {
152 if self.policy.is_unscoped_enabled() {
153 return true;
154 }
155
156 (self.policy.is_streaming_scoped()
157 && self.check_table_refill_streaming_vnodes(sstable, block_index))
158 || (self.policy.is_serving_scoped()
159 && self.check_table_refill_serving_vnodes(sstable, block_index))
160 }
161
162 fn allows_insert_only_data_refill_block(&self, sstable: &Sstable, block_index: usize) -> bool {
163 (self.policy.is_unscoped_enabled() || self.policy.is_serving_scoped())
167 && self.check_table_refill_serving_vnodes(sstable, block_index)
168 }
169
170 fn check_table_refill_streaming_vnodes(&self, sstable: &Sstable, block_index: usize) -> bool {
171 self.streaming_vnode_bitmap.as_ref().is_some_and(|bitmap| {
172 let vnode_range = block_vnode_range(sstable, block_index);
173 vnode_range_overlaps_bitmap(vnode_range, bitmap)
174 })
175 }
176
177 fn check_table_refill_serving_vnodes(&self, sstable: &Sstable, block_index: usize) -> bool {
178 self.serving_vnode_bitmap.as_ref().is_some_and(|bitmap| {
179 let vnode_range = block_vnode_range(sstable, block_index);
180 vnode_range_overlaps_bitmap(vnode_range, bitmap)
181 })
182 }
183}
184
185fn block_vnode_range(sstable: &Sstable, block_index: usize) -> (usize, usize) {
186 let block_meta = &sstable.meta.block_metas[block_index];
187 let block_smallest_key = FullKey::decode(&block_meta.smallest_key);
188 let table_key_end = match sstable.meta.block_metas.get(block_index + 1) {
189 Some(next_block_meta) if next_block_meta.table_id() != block_meta.table_id() => {
192 Bound::Unbounded
193 }
194 Some(next_block_meta) => Bound::Included(
197 FullKey::decode(&next_block_meta.smallest_key)
198 .user_key
199 .table_key,
200 ),
201 None => Bound::Included(
205 FullKey::decode(&sstable.meta.largest_key)
206 .user_key
207 .table_key,
208 ),
209 };
210
211 let table_key_range = (
212 Bound::Included(block_smallest_key.user_key.table_key),
213 table_key_end,
214 );
215 if match &table_key_range.0 {
219 Bound::Included(key) | Bound::Excluded(key) => key.as_ref().len() < VirtualNode::SIZE,
220 Bound::Unbounded => false,
221 } || match &table_key_range.1 {
222 Bound::Included(key) | Bound::Excluded(key) => key.as_ref().len() < VirtualNode::SIZE,
223 Bound::Unbounded => false,
224 } {
225 return (0, VirtualNode::MAX_REPRESENTABLE.to_index() + 1);
226 }
227 vnode_range(&table_key_range)
228}
229
230pub(crate) struct CacheRefiller {
232 spawn_refill_task: SpawnRefillTask,
233
234 config: Arc<CacheRefillConfig>,
235 meta_refill_concurrency: Option<Arc<Semaphore>>,
236 concurrency: Arc<Semaphore>,
237 sstable_store: SstableStoreRef,
238
239 role: Role,
240 default_policy: CacheRefillPolicy,
241 table_cache_refill_policies: HashMap<TableId, CacheRefillPolicy>,
242 streaming_table_vnode_mapping: HashMap<TableId, Bitmap>,
243 serving_table_vnode_mapping: HashMap<TableId, Bitmap>,
244}
245
246impl CacheRefiller {
247 pub(crate) fn new(
248 role: Role,
249 config: CacheRefillConfig,
250 sstable_store: SstableStoreRef,
251 spawn_refill_task: SpawnRefillTask,
252 ) -> Self {
253 let config = Arc::new(config);
254 let concurrency = Arc::new(Semaphore::new(config.concurrency));
255 let default_policy = config.table_cache_refill_default_policy;
256 let meta_refill_concurrency = if config.meta_refill_concurrency == 0 {
257 None
258 } else {
259 Some(Arc::new(Semaphore::new(config.meta_refill_concurrency)))
260 };
261 Self {
262 spawn_refill_task,
263 config,
264 meta_refill_concurrency,
265 concurrency,
266 sstable_store,
267 role,
268 default_policy,
269 table_cache_refill_policies: HashMap::new(),
270 streaming_table_vnode_mapping: HashMap::new(),
271 serving_table_vnode_mapping: HashMap::new(),
272 }
273 }
274
275 pub(crate) fn default_spawn_refill_task() -> SpawnRefillTask {
276 Arc::new(|deltas, context| {
277 let task = CacheRefillTask { deltas, context };
278 tokio::spawn(task.run())
279 })
280 }
281
282 pub(crate) fn start_cache_refill(&mut self, mut deltas: Vec<SstDeltaInfo>) -> JoinHandle<()> {
283 for delta in &mut deltas {
284 let for_serving = self.role.for_serving();
285 let for_streaming =
288 self.role.for_streaming() && !delta.delete_sst_object_ids.is_empty();
289
290 if !for_serving && !for_streaming {
291 delta.insert_sst_infos.clear();
292 continue;
293 }
294
295 delta.insert_sst_infos.retain(|sst| {
300 sst.table_ids.iter().any(|table_id| {
301 let policy = self
304 .table_cache_refill_policies
305 .get(table_id)
306 .copied()
307 .unwrap_or(self.default_policy);
308
309 match policy {
312 CacheRefillPolicy::Enabled => for_streaming || for_serving,
313 CacheRefillPolicy::Disabled => false,
314 CacheRefillPolicy::Streaming => {
315 for_streaming
316 && self.streaming_table_vnode_mapping.contains_key(table_id)
317 }
318 CacheRefillPolicy::Serving => {
319 for_serving && self.serving_table_vnode_mapping.contains_key(table_id)
320 }
321 CacheRefillPolicy::Both => {
322 (for_streaming
323 && self.streaming_table_vnode_mapping.contains_key(table_id))
324 || (for_serving
325 && self.serving_table_vnode_mapping.contains_key(table_id))
326 }
327 }
328 })
329 });
330 }
331 let context = self.new_cache_refill_context(&deltas);
332 (self.spawn_refill_task)(deltas, context)
333 }
334
335 fn new_cache_refill_context(&self, deltas: &[SstDeltaInfo]) -> CacheRefillContext {
336 let table_ids = deltas.iter().flat_map(|delta| {
337 delta
338 .insert_sst_infos
339 .iter()
340 .flat_map(|sst| sst.table_ids.iter().copied())
341 });
342 CacheRefillContext {
343 config: self.config.clone(),
344 meta_refill_concurrency: self.meta_refill_concurrency.clone(),
345 concurrency: self.concurrency.clone(),
346 sstable_store: self.sstable_store.clone(),
347 table_cache_refill_context_map: Arc::new(self.table_cache_refill_contexts(table_ids)),
348 }
349 }
350
351 pub(crate) fn replace_table_cache_refill_policies(
353 &mut self,
354 policies: HashMap<TableId, CacheRefillPolicy>,
355 ) {
356 self.table_cache_refill_policies = policies;
357 }
358
359 pub(crate) fn replace_serving_table_vnode_mapping(
361 &mut self,
362 mapping: HashMap<TableId, Bitmap>,
363 ) {
364 self.serving_table_vnode_mapping = mapping;
365 }
366
367 pub(crate) fn update_streaming_table_vnodes(
368 &mut self,
369 table_id: TableId,
370 streaming_vnodes: Option<Bitmap>,
371 ) {
372 if let Some(streaming_vnodes) = streaming_vnodes {
373 self.streaming_table_vnode_mapping
374 .insert(table_id, streaming_vnodes);
375 } else {
376 self.streaming_table_vnode_mapping.remove(&table_id);
377 }
378 }
379
380 fn table_cache_refill_contexts(
381 &self,
382 table_ids: impl IntoIterator<Item = TableId>,
383 ) -> TableCacheRefillContextMap {
384 let for_streaming = self.role.for_streaming();
385 let for_serving = self.role.for_serving();
386 let mut contexts = HashMap::new();
387 for table_id in table_ids {
388 if for_serving
389 && !for_streaming
390 && !self.serving_table_vnode_mapping.contains_key(&table_id)
391 {
392 continue;
393 }
394 contexts.entry(table_id).or_insert_with(|| {
395 let policy = self
396 .table_cache_refill_policies
397 .get(&table_id)
398 .copied()
399 .unwrap_or(self.default_policy);
400 let streaming_vnode_bitmap = (for_streaming && policy.is_streaming_scoped())
401 .then(|| self.streaming_table_vnode_mapping.get(&table_id).cloned())
402 .flatten();
403 let serving_vnode_bitmap = (for_serving
406 && (policy.is_serving_scoped() || policy.is_unscoped_enabled()))
407 .then(|| self.serving_table_vnode_mapping.get(&table_id).cloned())
408 .flatten();
409 TableCacheRefillContext {
410 streaming_vnode_bitmap,
411 serving_vnode_bitmap,
412 policy,
413 }
414 });
415 }
416 contexts
417 }
418
419 pub(crate) fn table_cache_refill_monitor_snapshot(&self) -> TableCacheRefillMonitorSnapshot {
420 let table_ids = self
421 .table_cache_refill_policies
422 .keys()
423 .chain(self.streaming_table_vnode_mapping.keys())
424 .chain(self.serving_table_vnode_mapping.keys())
425 .copied();
426 TableCacheRefillMonitorSnapshot {
427 contexts: self.table_cache_refill_contexts(table_ids),
428 policies: self.table_cache_refill_policies.clone(),
429 default_policy: self.default_policy,
430 streaming_table_vnode_mapping: self.streaming_table_vnode_mapping.clone(),
431 serving_table_vnode_mapping: self.serving_table_vnode_mapping.clone(),
432 }
433 }
434}
435
436#[derive(Clone)]
437pub(crate) struct CacheRefillContext {
438 config: Arc<CacheRefillConfig>,
439 meta_refill_concurrency: Option<Arc<Semaphore>>,
440 concurrency: Arc<Semaphore>,
441 sstable_store: SstableStoreRef,
442 table_cache_refill_context_map: Arc<TableCacheRefillContextMap>,
443}
444
445struct DataCacheRefillTaskGenerator<'a> {
446 context: &'a CacheRefillContext,
447 delta: &'a SstDeltaInfo,
448 ssts: &'a [TableHolder],
449}
450
451impl DataCacheRefillTaskGenerator<'_> {
452 fn generate_unfiltered_tasks(&self) -> Vec<DataCacheRefillTask> {
453 let mut tasks = Vec::new();
454
455 if !self.context.sstable_store.block_cache().is_hybrid() {
457 return tasks;
458 }
459
460 if self.delta.insert_sst_infos.is_empty() {
461 return tasks;
462 }
463
464 let has_parent_ssts = !self.delta.delete_sst_object_ids.is_empty();
465 debug_assert!(has_parent_ssts || self.delta.insert_sst_level == 0);
468
469 if !self
471 .context
472 .config
473 .data_refill_levels
474 .contains(&self.delta.insert_sst_level)
475 {
476 return tasks;
477 }
478
479 let unit = self.context.config.unit;
482 assert!(unit > 0, "cache refill unit must be positive");
483 let table_cache_refill_context_map = &self.context.table_cache_refill_context_map;
484 for (sst_info, sst) in self.delta.insert_sst_infos.iter().zip_eq_fast(self.ssts) {
485 debug_assert_eq!(sst_info.object_id, sst.id);
486 debug_assert!(sst_info.table_ids.is_sorted());
487 let mut blk_start = 0;
488 while blk_start < sst.block_count() {
489 let table_id = sst.meta.block_metas[blk_start].table_id();
493 let mut blk_end = std::cmp::min(sst.block_count(), blk_start + unit);
494 if let Some(table_boundary) = (blk_start + 1..blk_end)
495 .find(|&block_index| sst.meta.block_metas[block_index].table_id() != table_id)
496 {
497 blk_end = table_boundary;
498 }
499
500 let should_refill = sst_info.table_ids.binary_search(&table_id).is_ok()
501 && (blk_start..blk_end).any(|block_index| {
502 table_cache_refill_context_map
503 .get(&table_id)
504 .is_some_and(|context| {
505 if has_parent_ssts {
506 context.allows_normal_data_refill_block(sst, block_index)
507 } else {
508 context.allows_insert_only_data_refill_block(sst, block_index)
509 }
510 })
511 });
512 if should_refill {
513 tasks.push(DataCacheRefillTask {
514 sst: sst.clone(),
515 blks: blk_start..blk_end,
516 });
517 }
518 blk_start = blk_end;
519 }
520 }
521
522 if tasks.is_empty() {
523 return tasks;
524 }
525
526 if has_parent_ssts
529 && !self.context.config.skip_recent_filter
530 && !self.filter_by_recent_filter()
531 {
532 GLOBAL_CACHE_REFILL_METRICS
533 .data_refill_filtered_total
534 .inc_by(self.delta.delete_sst_object_ids.len() as u64);
535 return vec![];
536 }
537
538 tasks
539 }
540
541 async fn filter_by_inheritance_if_needed(
542 &self,
543 tasks: Vec<DataCacheRefillTask>,
544 ) -> Vec<DataCacheRefillTask> {
545 let should_filter_by_inheritance = !tasks.is_empty()
548 && !self.delta.delete_sst_object_ids.is_empty()
549 && self.delta.insert_sst_level != 0
550 && !self.context.config.skip_recent_filter
551 && !self.context.config.skip_inheritance_filter;
552 if should_filter_by_inheritance {
553 self.filter_by_inheritance_filter(tasks).await
554 } else {
555 tasks
556 }
557 }
558
559 fn filter_by_recent_filter(&self) -> bool {
561 let recent_filter = self.context.sstable_store.recent_filter();
562 let targets = self
563 .delta
564 .delete_sst_object_ids
565 .iter()
566 .map(|id| (*id, usize::MAX))
567 .collect_vec();
568 recent_filter.contains_any(targets.iter())
569 }
570
571 async fn filter_by_inheritance_filter(
572 &self,
573 originals: Vec<DataCacheRefillTask>,
574 ) -> Vec<DataCacheRefillTask> {
575 let sstable_store = self.context.sstable_store.clone();
577 let futures = self.delta.delete_sst_object_ids.iter().map(|sst_obj_id| {
578 let store = &sstable_store;
579 async move {
580 let res = store.sstable_cached(*sst_obj_id).await;
581 match res {
582 Ok(Some(_)) => GLOBAL_CACHE_REFILL_METRICS
583 .data_refill_parent_meta_lookup_hit_total
584 .inc(),
585 Ok(None) => GLOBAL_CACHE_REFILL_METRICS
586 .data_refill_parent_meta_lookup_miss_total
587 .inc(),
588 _ => {}
589 }
590 res
591 }
592 });
593 let parent_ssts = match try_join_all(futures).await {
594 Ok(parent_ssts) => parent_ssts.into_iter().flatten(),
595 Err(e) => {
596 tracing::error!(error = %e.as_report(), "get old meta from cache error");
597 return vec![];
598 }
599 };
600
601 if cfg!(debug_assertions) {
603 originals.iter().tuple_windows().for_each(|(a, b)| {
604 debug_assert_ne!(
605 KeyComparator::compare_encoded_full_key(a.largest_key(), b.smallest_key()),
606 std::cmp::Ordering::Greater
607 )
608 });
609 }
610
611 let mut filtered: HashSet<SstableUnit> = HashSet::default();
612 let recent_filter = self.context.sstable_store.recent_filter();
613 for psst in parent_ssts {
614 for pblk in 0..psst.block_count() {
615 let pleft = &psst.meta.block_metas[pblk].smallest_key;
616 let pright = if pblk + 1 == psst.block_count() {
617 &psst.meta.largest_key
619 } else {
620 &psst.meta.block_metas[pblk + 1].smallest_key
621 };
622
623 let uleft = originals.partition_point(|task| {
625 KeyComparator::compare_encoded_full_key(task.largest_key(), pleft)
626 == std::cmp::Ordering::Less
627 });
628 let uright = originals.partition_point(|task| {
630 KeyComparator::compare_encoded_full_key(task.smallest_key(), pright)
631 != std::cmp::Ordering::Greater
632 });
633
634 for task in originals.iter().take(uright).skip(uleft) {
636 let unit = task.unit();
637 if filtered.contains(&unit) {
638 continue;
639 }
640 if recent_filter.contains(&(psst.id, pblk)) {
641 filtered.insert(unit);
642 }
643 }
644 }
645 }
646
647 let hit = filtered.len();
648 let miss = originals.len() - hit;
649 GLOBAL_CACHE_REFILL_METRICS
650 .data_refill_unit_inheritance_hit_total
651 .inc_by(hit as u64);
652 GLOBAL_CACHE_REFILL_METRICS
653 .data_refill_unit_inheritance_miss_total
654 .inc_by(miss as u64);
655
656 originals
657 .into_iter()
658 .filter(|task| filtered.contains(&task.unit()))
659 .collect()
660 }
661}
662
663#[derive(Debug)]
664struct DataCacheRefillTask {
665 sst: TableHolder,
666 blks: Range<usize>,
667}
668
669impl DataCacheRefillTask {
670 fn unit(&self) -> SstableUnit {
671 SstableUnit {
672 sst_obj_id: self.sst.id,
673 blks: self.blks.clone(),
674 }
675 }
676
677 fn smallest_key(&self) -> &[u8] {
678 &self.sst.meta.block_metas[self.blks.start].smallest_key
679 }
680
681 fn largest_key(&self) -> &[u8] {
682 if self.blks.end == self.sst.block_count() {
683 &self.sst.meta.largest_key
684 } else {
685 &self.sst.meta.block_metas[self.blks.end].smallest_key
686 }
687 }
688}
689
690struct CacheRefillTask {
691 deltas: Vec<SstDeltaInfo>,
692 context: CacheRefillContext,
693}
694
695impl CacheRefillTask {
696 async fn run(self) {
697 let tasks = self
698 .deltas
699 .iter()
700 .map(|delta| {
701 let context = self.context.clone();
702 async move {
703 let holders = match Self::meta_cache_refill(&context, delta).await {
704 Ok(holders) => holders,
705 Err(e) => {
706 tracing::warn!(error = %e.as_report(), "meta cache refill error");
707 return;
708 }
709 };
710 let generator = DataCacheRefillTaskGenerator {
711 context: &context,
712 delta,
713 ssts: &holders,
714 };
715 let tasks = generator.generate_unfiltered_tasks();
716
717 let unfiltered_block_count =
719 tasks.iter().map(|task| task.blks.len() as u64).sum();
720 GLOBAL_CACHE_REFILL_METRICS
721 .data_refill_block_unfiltered_total
722 .inc_by(unfiltered_block_count);
723
724 let tasks = generator.filter_by_inheritance_if_needed(tasks).await;
725 Self::data_cache_refill(&context, tasks).await;
726 }
727 })
728 .collect_vec();
729 let future = join_all(tasks);
730
731 let _ = tokio::time::timeout(self.context.config.timeout, future).await;
732 }
733
734 async fn meta_cache_refill(
735 context: &CacheRefillContext,
736 delta: &SstDeltaInfo,
737 ) -> HummockResult<Vec<TableHolder>> {
738 let tasks = delta
739 .insert_sst_infos
740 .iter()
741 .map(|info| async {
742 let mut stats = StoreLocalStatistic::default();
743 GLOBAL_CACHE_REFILL_METRICS.meta_refill_attempts_total.inc();
744
745 let permit = if let Some(c) = &context.meta_refill_concurrency {
746 Some(c.acquire().await.unwrap())
747 } else {
748 None
749 };
750
751 let now = Instant::now();
752 let res = context.sstable_store.sstable(info, &mut stats).await;
753 stats.discard();
754 if res.is_ok() {
755 GLOBAL_CACHE_REFILL_METRICS
756 .meta_refill_success_duration
757 .observe(now.elapsed().as_secs_f64());
758 }
759 drop(permit);
760
761 res
762 })
763 .collect_vec();
764 let holders = try_join_all(tasks).await?;
765 Ok(holders)
766 }
767
768 async fn data_cache_refill(context: &CacheRefillContext, tasks: Vec<DataCacheRefillTask>) {
769 let mut futures = Vec::with_capacity(tasks.len());
770 for task in tasks {
771 context
773 .sstable_store
774 .recent_filter()
775 .insert((task.sst.id, usize::MAX));
776
777 let blocks = task.blks.len();
778 let mut contexts = Vec::with_capacity(blocks);
779 let mut admits = 0;
780
781 let (range_first, _) = task.sst.calculate_block_info(task.blks.start);
782 let (range_last, _) = task.sst.calculate_block_info(task.blks.end - 1);
783 let range = range_first.start..range_last.end;
784
785 let size = range.size().unwrap();
786
787 GLOBAL_CACHE_REFILL_METRICS
788 .data_refill_ideal_bytes
789 .inc_by(size as _);
790
791 for blk in task.blks {
792 let (range, uncompressed_capacity) = task.sst.calculate_block_info(blk);
793 let key = SstableBlockIndex {
794 sst_id: task.sst.id,
795 block_idx: blk as u64,
796 };
797
798 let mut writer = context.sstable_store.block_cache().storage_writer(key);
799
800 if writer.filter(size).is_admitted() {
801 admits += 1;
802 }
803
804 contexts.push((writer, range, uncompressed_capacity))
805 }
806
807 if admits as f64 / contexts.len() as f64 >= context.config.threshold {
808 let sstable_store = context.sstable_store.clone();
809 let context = context.clone();
810 let future = async move {
811 GLOBAL_CACHE_REFILL_METRICS.data_refill_attempts_total.inc();
812
813 let permit = context.concurrency.acquire().await.unwrap();
814
815 GLOBAL_CACHE_REFILL_METRICS.data_refill_started_total.inc();
816
817 let now = Instant::now();
818
819 let data = sstable_store
820 .store()
821 .read(&sstable_store.get_sst_data_path(task.sst.id), range.clone())
822 .await?;
823 let mut apply_disk_cache_futures = vec![];
824 for (w, r, uc) in contexts {
825 let offset = r.start - range.start;
826 let len = r.end - r.start;
827 let bytes = data.slice(offset..offset + len);
828 let future = async move {
829 let value = Box::new(Block::decode(bytes, uc)?);
830 if let Some(_entry) = w.force().insert(value) {
832 GLOBAL_CACHE_REFILL_METRICS
833 .data_refill_success_bytes
834 .inc_by(len as u64);
835 GLOBAL_CACHE_REFILL_METRICS
836 .data_refill_block_success_total
837 .inc();
838 }
839 Ok::<_, HummockError>(())
840 };
841 apply_disk_cache_futures.push(future);
842 }
843 try_join_all(apply_disk_cache_futures)
844 .await
845 .map_err(HummockError::file_cache)?;
846
847 GLOBAL_CACHE_REFILL_METRICS
848 .data_refill_success_duration
849 .observe(now.elapsed().as_secs_f64());
850 drop(permit);
851
852 Ok::<_, HummockError>(())
853 };
854 futures.push(future);
855 }
856 }
857
858 let futures = futures.into_iter().map(|future| async move {
859 if let Err(e) = future.await {
860 tracing::error!(error = %e.as_report(), "data cache refill task error");
861 }
862 });
863
864 join_all(futures).await;
865 }
866}
867
868#[derive(Debug)]
869pub struct SstableBlock {
870 pub sst_obj_id: HummockSstableObjectId,
871 pub blk_idx: usize,
872}
873
874#[derive(Debug, Hash, PartialEq, Eq)]
875pub struct SstableUnit {
876 pub sst_obj_id: HummockSstableObjectId,
877 pub blks: Range<usize>,
878}
879
880impl Ord for SstableUnit {
881 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
882 match self.sst_obj_id.cmp(&other.sst_obj_id) {
883 std::cmp::Ordering::Equal => {}
884 ord => return ord,
885 }
886 match self.blks.start.cmp(&other.blks.start) {
887 std::cmp::Ordering::Equal => {}
888 ord => return ord,
889 }
890 self.blks.end.cmp(&other.blks.end)
891 }
892}
893
894impl PartialOrd for SstableUnit {
895 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
896 Some(self.cmp(other))
897 }
898}
899
900#[cfg(test)]
901mod tests {
902 use std::collections::{HashMap, HashSet};
903 use std::sync::Arc;
904 use std::time::Duration;
905
906 use bytes::Bytes;
907 use foyer::{
908 BlockEngineConfig, CacheBuilder, DeviceBuilder, FsDeviceBuilder, HybridCacheBuilder,
909 PsyncIoEngineConfig,
910 };
911 use parking_lot::Mutex;
912 use risingwave_common::bitmap::Bitmap;
913 use risingwave_common::config::streaming::CacheRefillPolicy;
914 use risingwave_common::config::{MetricLevel, Role};
915 use risingwave_common::hash::VirtualNode;
916 use risingwave_common::util::epoch::test_epoch;
917 use risingwave_hummock_sdk::compaction_group::group_split::split_sst_with_table_ids;
918 use risingwave_hummock_sdk::key::{FullKey, UserKey, prefix_slice_with_vnode};
919 use risingwave_hummock_sdk::sstable_info::{SstableInfo, SstableInfoInner};
920 use risingwave_hummock_sdk::{EpochWithGap, HummockSstableObjectId};
921 use risingwave_pb::id::TableId;
922
923 use super::{
924 CacheRefillConfig, CacheRefillContext, CacheRefiller, DataCacheRefillTaskGenerator,
925 SpawnRefillTask, SstDeltaInfo, block_vnode_range, vnode_range_overlaps_bitmap,
926 };
927 use crate::hummock::iterator::test_utils::{iterator_test_table_key_of, mock_sstable_store};
928 use crate::hummock::recent_filter::simple::SimpleRecentFilter;
929 use crate::hummock::sstable_store::{SstableStore, SstableStoreConfig};
930 use crate::hummock::test_utils::{
931 default_builder_opt_for_test, gen_test_sstable_with_table_ids,
932 };
933 use crate::hummock::value::HummockValue;
934 use crate::hummock::{RecentFilter, RecentFilterTrait, SstableStoreRef, TableHolder};
935 use crate::monitor::global_hummock_state_store_metrics;
936
937 async fn mock_sstable_store_with_disk_cache(
938 recent_filter: Option<Arc<RecentFilter<(HummockSstableObjectId, usize)>>>,
939 ) -> (SstableStoreRef, tempfile::TempDir) {
940 let cache_dir = tempfile::tempdir().unwrap();
941 let device = FsDeviceBuilder::new(cache_dir.path())
942 .with_capacity(16 << 20)
943 .build()
944 .unwrap();
945 let block_cache = HybridCacheBuilder::new()
946 .memory(64 << 20)
947 .with_shards(2)
948 .storage()
949 .with_io_engine_config(PsyncIoEngineConfig::new())
950 .with_engine_config(
951 BlockEngineConfig::new(device)
952 .with_block_size(1 << 20)
953 .with_buffer_pool_size(4 << 20),
954 )
955 .build()
956 .await
957 .unwrap();
958 assert!(block_cache.is_hybrid());
959 let memory_store = mock_sstable_store().await;
960 let sstable_store = Arc::new(SstableStore::new(SstableStoreConfig {
961 store: memory_store.store(),
962 path: "test".to_owned(),
963 prefetch_buffer_capacity: 64 << 20,
964 max_prefetch_block_number: 16,
965 recent_filter: recent_filter.unwrap_or_else(|| memory_store.recent_filter().clone()),
966 state_store_metrics: Arc::new(global_hummock_state_store_metrics(
967 MetricLevel::Disabled,
968 )),
969 use_new_object_prefix_strategy: true,
970 skip_bloom_filter_in_serde: false,
971 meta_cache: memory_store.meta_cache().clone(),
972 block_cache,
973 vector_meta_cache: CacheBuilder::new(64 << 20).build(),
974 vector_block_cache: CacheBuilder::new(64 << 20).build(),
975 }));
976 (sstable_store, cache_dir)
977 }
978
979 fn test_refill_config(default_policy: CacheRefillPolicy) -> CacheRefillConfig {
980 CacheRefillConfig {
981 timeout: Duration::from_secs(1),
982 data_refill_levels: HashSet::new(),
983 meta_refill_concurrency: 1,
984 concurrency: 1,
985 unit: 1,
986 threshold: 0.0,
987 skip_recent_filter: true,
988 skip_inheritance_filter: true,
989 table_cache_refill_default_policy: default_policy,
990 }
991 }
992
993 async fn gen_test_sst_with_object_id(
994 table_id: TableId,
995 sstable_store: SstableStoreRef,
996 object_id: u64,
997 ) -> (TableHolder, SstableInfo) {
998 gen_test_sstable_with_table_ids(
999 default_builder_opt_for_test(),
1000 object_id,
1001 (0..2).map(|idx| {
1002 (
1003 FullKey {
1004 user_key: risingwave_hummock_sdk::key::UserKey::for_test(
1005 table_id,
1006 iterator_test_table_key_of(idx),
1007 ),
1008 epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(233)),
1009 },
1010 HummockValue::put(vec![idx as u8]),
1011 )
1012 }),
1013 sstable_store,
1014 vec![table_id.as_raw_id()],
1015 )
1016 .await
1017 }
1018
1019 struct DataRefillGeneratorTestFixture {
1020 table_id: TableId,
1021 sstable_store: SstableStoreRef,
1022 sst: TableHolder,
1023 sst_info: SstableInfo,
1024 deleted_sst_object_id: HummockSstableObjectId,
1025 _cache_dir: tempfile::TempDir,
1026 }
1027
1028 impl DataRefillGeneratorTestFixture {
1029 async fn new(
1030 recent_filter: Option<Arc<RecentFilter<(HummockSstableObjectId, usize)>>>,
1031 ) -> Self {
1032 let table_id = TableId::from(233);
1033 let (sstable_store, cache_dir) =
1034 mock_sstable_store_with_disk_cache(recent_filter).await;
1035 let (sst, sst_info) =
1036 gen_test_sst_with_object_id(table_id, sstable_store.clone(), 1).await;
1037 Self {
1038 table_id,
1039 sstable_store,
1040 sst,
1041 sst_info,
1042 deleted_sst_object_id: 2330.into(),
1043 _cache_dir: cache_dir,
1044 }
1045 }
1046
1047 fn context(
1048 &self,
1049 policy: CacheRefillPolicy,
1050 streaming_vnode_bitmap: Option<Bitmap>,
1051 serving_vnode_bitmap: Option<Bitmap>,
1052 configure: impl FnOnce(&mut CacheRefillConfig),
1053 ) -> CacheRefillContext {
1054 let mut config = test_refill_config(CacheRefillPolicy::Enabled);
1055 config.data_refill_levels.insert(0);
1056 configure(&mut config);
1057 CacheRefillContext {
1058 config: Arc::new(config),
1059 meta_refill_concurrency: None,
1060 concurrency: Arc::new(tokio::sync::Semaphore::new(1)),
1061 sstable_store: self.sstable_store.clone(),
1062 table_cache_refill_context_map: Arc::new(HashMap::from([(
1063 self.table_id,
1064 super::TableCacheRefillContext {
1065 streaming_vnode_bitmap,
1066 serving_vnode_bitmap,
1067 policy,
1068 },
1069 )])),
1070 }
1071 }
1072
1073 fn normal_delta(
1074 &self,
1075 insert_sst_level: u32,
1076 deleted_sst_object_id: HummockSstableObjectId,
1077 ) -> SstDeltaInfo {
1078 SstDeltaInfo {
1079 insert_sst_infos: vec![self.sst_info.clone()],
1080 delete_sst_object_ids: vec![deleted_sst_object_id],
1081 insert_sst_level,
1082 }
1083 }
1084
1085 fn normal_l0_delta(&self) -> SstDeltaInfo {
1086 self.normal_delta(0, self.deleted_sst_object_id)
1087 }
1088
1089 async fn generate(
1090 &self,
1091 context: &CacheRefillContext,
1092 delta: &SstDeltaInfo,
1093 ) -> Vec<super::DataCacheRefillTask> {
1094 let generator = DataCacheRefillTaskGenerator {
1095 context,
1096 delta,
1097 ssts: std::slice::from_ref(&self.sst),
1098 };
1099 let tasks = generator.generate_unfiltered_tasks();
1100 generator.filter_by_inheritance_if_needed(tasks).await
1101 }
1102 }
1103
1104 #[tokio::test]
1105 async fn test_table_cache_refill_contexts_by_role_and_policy() {
1106 struct Case {
1107 name: &'static str,
1108 role: Role,
1109 default_policy: CacheRefillPolicy,
1110 policy: Option<CacheRefillPolicy>,
1111 has_streaming_vnodes: bool,
1112 has_serving_vnodes: bool,
1113 expected: Option<(CacheRefillPolicy, bool, bool)>,
1114 }
1115
1116 let cases = [
1117 Case {
1118 name: "streaming role uses streaming side of Both",
1119 role: Role::Streaming,
1120 default_policy: CacheRefillPolicy::Disabled,
1121 policy: Some(CacheRefillPolicy::Both),
1122 has_streaming_vnodes: true,
1123 has_serving_vnodes: true,
1124 expected: Some((CacheRefillPolicy::Both, true, false)),
1125 },
1126 Case {
1127 name: "serving role uses serving side of Both",
1128 role: Role::Serving,
1129 default_policy: CacheRefillPolicy::Disabled,
1130 policy: Some(CacheRefillPolicy::Both),
1131 has_streaming_vnodes: true,
1132 has_serving_vnodes: true,
1133 expected: Some((CacheRefillPolicy::Both, false, true)),
1134 },
1135 Case {
1136 name: "both role keeps both sides",
1137 role: Role::Both,
1138 default_policy: CacheRefillPolicy::Disabled,
1139 policy: Some(CacheRefillPolicy::Both),
1140 has_streaming_vnodes: true,
1141 has_serving_vnodes: true,
1142 expected: Some((CacheRefillPolicy::Both, true, true)),
1143 },
1144 Case {
1145 name: "both role keeps streaming-only ownership",
1146 role: Role::Both,
1147 default_policy: CacheRefillPolicy::Disabled,
1148 policy: Some(CacheRefillPolicy::Both),
1149 has_streaming_vnodes: true,
1150 has_serving_vnodes: false,
1151 expected: Some((CacheRefillPolicy::Both, true, false)),
1152 },
1153 Case {
1154 name: "streaming scope without ownership has no usable bitmap",
1155 role: Role::Streaming,
1156 default_policy: CacheRefillPolicy::Disabled,
1157 policy: Some(CacheRefillPolicy::Streaming),
1158 has_streaming_vnodes: false,
1159 has_serving_vnodes: false,
1160 expected: Some((CacheRefillPolicy::Streaming, false, false)),
1161 },
1162 Case {
1163 name: "pure serving worker excludes unmapped table",
1164 role: Role::Serving,
1165 default_policy: CacheRefillPolicy::Disabled,
1166 policy: Some(CacheRefillPolicy::Serving),
1167 has_streaming_vnodes: true,
1168 has_serving_vnodes: false,
1169 expected: None,
1170 },
1171 Case {
1172 name: "default Enabled retains serving ownership",
1173 role: Role::Serving,
1174 default_policy: CacheRefillPolicy::Enabled,
1175 policy: None,
1176 has_streaming_vnodes: false,
1177 has_serving_vnodes: true,
1178 expected: Some((CacheRefillPolicy::Enabled, false, true)),
1179 },
1180 Case {
1181 name: "explicit policy overrides default",
1182 role: Role::Serving,
1183 default_policy: CacheRefillPolicy::Enabled,
1184 policy: Some(CacheRefillPolicy::Disabled),
1185 has_streaming_vnodes: false,
1186 has_serving_vnodes: true,
1187 expected: Some((CacheRefillPolicy::Disabled, false, false)),
1188 },
1189 ];
1190
1191 let table_id = TableId::from(233);
1192 let streaming_vnodes = Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [1, 3]);
1193 let serving_vnodes = Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [2, 4]);
1194 let sstable_store = mock_sstable_store().await;
1195 for case in cases {
1196 let mut refiller = CacheRefiller::new(
1197 case.role,
1198 test_refill_config(case.default_policy),
1199 sstable_store.clone(),
1200 CacheRefiller::default_spawn_refill_task(),
1201 );
1202 if let Some(policy) = case.policy {
1203 refiller.replace_table_cache_refill_policies(HashMap::from([(table_id, policy)]));
1204 }
1205 if case.has_streaming_vnodes {
1206 refiller.update_streaming_table_vnodes(table_id, Some(streaming_vnodes.clone()));
1207 }
1208 if case.has_serving_vnodes {
1209 refiller.replace_serving_table_vnode_mapping(HashMap::from([(
1210 table_id,
1211 serving_vnodes.clone(),
1212 )]));
1213 }
1214
1215 let contexts = refiller.table_cache_refill_contexts([table_id, table_id]);
1217 assert_eq!(
1218 contexts.len(),
1219 usize::from(case.expected.is_some()),
1220 "{}",
1221 case.name
1222 );
1223 let actual = contexts.get(&table_id).map(|context| {
1224 (
1225 context.policy,
1226 context.streaming_vnode_bitmap.as_ref(),
1227 context.serving_vnode_bitmap.as_ref(),
1228 )
1229 });
1230 let expected = case.expected.map(|(policy, streaming, serving)| {
1231 (
1232 policy,
1233 streaming.then_some(&streaming_vnodes),
1234 serving.then_some(&serving_vnodes),
1235 )
1236 });
1237 assert_eq!(actual, expected, "{}", case.name);
1238 }
1239 }
1240
1241 #[tokio::test]
1242 async fn test_refill_task_captures_runtime_context_snapshot() {
1243 let table_id = TableId::from(233);
1244 let old_vnodes = Bitmap::ones(VirtualNode::COUNT_FOR_TEST);
1245 let new_vnodes = Bitmap::from_range(VirtualNode::COUNT_FOR_TEST, 0..8);
1246 let captured_context = Arc::new(Mutex::new(None::<CacheRefillContext>));
1247 let captured_context_clone = captured_context.clone();
1248 let spawn_refill_task: SpawnRefillTask = Arc::new(move |_, context| {
1249 *captured_context_clone.lock() = Some(context);
1250 tokio::spawn(async {})
1251 });
1252 let mut refiller = CacheRefiller::new(
1253 Role::Serving,
1254 test_refill_config(CacheRefillPolicy::Enabled),
1255 mock_sstable_store().await,
1256 spawn_refill_task,
1257 );
1258
1259 refiller.replace_table_cache_refill_policies(HashMap::from([(
1260 table_id,
1261 CacheRefillPolicy::Serving,
1262 )]));
1263 refiller
1264 .replace_serving_table_vnode_mapping(HashMap::from([(table_id, old_vnodes.clone())]));
1265
1266 let _handle = refiller.start_cache_refill(vec![SstDeltaInfo {
1267 insert_sst_infos: vec![SstableInfo::from(SstableInfoInner {
1268 table_ids: vec![table_id],
1269 ..Default::default()
1270 })],
1271 ..Default::default()
1272 }]);
1273 refiller.replace_table_cache_refill_policies(HashMap::from([(
1274 table_id,
1275 CacheRefillPolicy::Disabled,
1276 )]));
1277 refiller.replace_serving_table_vnode_mapping(HashMap::from([(table_id, new_vnodes)]));
1278
1279 let captured_context = captured_context.lock();
1280 let context = captured_context
1281 .as_ref()
1282 .unwrap()
1283 .table_cache_refill_context_map
1284 .get(&table_id)
1285 .unwrap();
1286 assert_eq!(context.policy, CacheRefillPolicy::Serving);
1287 assert_eq!(context.serving_vnode_bitmap.as_ref(), Some(&old_vnodes));
1288 }
1289
1290 #[tokio::test]
1291 async fn test_cache_refill_prunes_whole_ssts_before_meta_load() {
1292 let streaming_table = TableId::from(1);
1293 let serving_table = TableId::from(2);
1294 let disabled_table = TableId::from(3);
1295 let internal_table = TableId::from(4);
1296 let fallback_table = TableId::from(5);
1297 let serving_vnodes = Bitmap::ones(VirtualNode::COUNT_FOR_TEST);
1298 let sstable_store = mock_sstable_store().await;
1299 let sst_info = |table_ids: Vec<TableId>| {
1300 SstableInfo::from(SstableInfoInner {
1301 table_ids,
1302 ..Default::default()
1303 })
1304 };
1305 let capture_pruned_table_ids = |role,
1306 default_policy,
1307 policies: HashMap<TableId, CacheRefillPolicy>,
1308 streaming_table_vnodes: HashMap<TableId, Bitmap>,
1309 serving_table_vnodes: HashMap<TableId, Bitmap>,
1310 delta| {
1311 let captured_deltas = Arc::new(Mutex::new(None::<Vec<SstDeltaInfo>>));
1312 let captured_deltas_clone = captured_deltas.clone();
1313 let spawn_refill_task: SpawnRefillTask = Arc::new(move |deltas, _| {
1314 *captured_deltas_clone.lock() = Some(deltas);
1315 tokio::spawn(async {})
1316 });
1317 let mut refiller = CacheRefiller::new(
1318 role,
1319 test_refill_config(default_policy),
1320 sstable_store.clone(),
1321 spawn_refill_task,
1322 );
1323 refiller.replace_table_cache_refill_policies(policies);
1324 for (table_id, vnodes) in streaming_table_vnodes {
1325 refiller.update_streaming_table_vnodes(table_id, Some(vnodes));
1326 }
1327 refiller.replace_serving_table_vnode_mapping(serving_table_vnodes);
1328 let _handle = refiller.start_cache_refill(vec![delta]);
1329 captured_deltas
1330 .lock()
1331 .take()
1332 .unwrap()
1333 .pop()
1334 .unwrap()
1335 .insert_sst_infos
1336 .into_iter()
1337 .map(|sst| sst.table_ids.clone())
1338 .collect::<Vec<_>>()
1339 };
1340
1341 let normal_delta = |insert_sst_infos| SstDeltaInfo {
1342 insert_sst_infos,
1343 delete_sst_object_ids: vec![1.into()],
1344 insert_sst_level: 1,
1345 };
1346 let insert_only_delta = |insert_sst_infos| SstDeltaInfo {
1347 insert_sst_infos,
1348 delete_sst_object_ids: vec![],
1349 insert_sst_level: 0,
1350 };
1351
1352 assert_eq!(
1353 capture_pruned_table_ids(
1354 Role::Both,
1355 CacheRefillPolicy::Disabled,
1356 HashMap::from([
1357 (disabled_table, CacheRefillPolicy::Disabled),
1358 (streaming_table, CacheRefillPolicy::Enabled),
1359 ]),
1360 HashMap::new(),
1361 HashMap::new(),
1362 normal_delta(vec![
1363 sst_info(vec![disabled_table]),
1364 sst_info(vec![streaming_table]),
1365 ]),
1366 ),
1367 vec![vec![streaming_table]],
1368 );
1369
1370 assert_eq!(
1371 capture_pruned_table_ids(
1372 Role::Both,
1373 CacheRefillPolicy::Disabled,
1374 HashMap::from([
1375 (streaming_table, CacheRefillPolicy::Streaming),
1376 (internal_table, CacheRefillPolicy::Serving),
1377 ]),
1378 HashMap::from([(streaming_table, serving_vnodes.clone())]),
1379 HashMap::new(),
1380 normal_delta(vec![
1381 sst_info(vec![streaming_table]),
1382 sst_info(vec![internal_table]),
1383 ]),
1384 ),
1385 vec![vec![streaming_table]],
1386 );
1387
1388 assert_eq!(
1389 capture_pruned_table_ids(
1390 Role::Both,
1391 CacheRefillPolicy::Disabled,
1392 HashMap::from([
1393 (streaming_table, CacheRefillPolicy::Streaming),
1394 (serving_table, CacheRefillPolicy::Serving),
1395 (internal_table, CacheRefillPolicy::Serving),
1396 ]),
1397 HashMap::from([(streaming_table, serving_vnodes.clone())]),
1398 HashMap::from([(serving_table, serving_vnodes.clone())]),
1399 insert_only_delta(vec![
1400 sst_info(vec![streaming_table]),
1401 sst_info(vec![serving_table]),
1402 sst_info(vec![internal_table]),
1403 ]),
1404 ),
1405 vec![vec![serving_table]],
1406 );
1407
1408 assert_eq!(
1409 capture_pruned_table_ids(
1410 Role::Serving,
1411 CacheRefillPolicy::Disabled,
1412 HashMap::from([
1413 (streaming_table, CacheRefillPolicy::Streaming),
1414 (serving_table, CacheRefillPolicy::Serving),
1415 ]),
1416 HashMap::new(),
1417 HashMap::from([(serving_table, serving_vnodes.clone())]),
1418 normal_delta(vec![
1419 sst_info(vec![streaming_table]),
1420 sst_info(vec![serving_table]),
1421 ]),
1422 ),
1423 vec![vec![serving_table]],
1424 );
1425
1426 assert_eq!(
1427 capture_pruned_table_ids(
1428 Role::Serving,
1429 CacheRefillPolicy::Disabled,
1430 HashMap::from([
1431 (streaming_table, CacheRefillPolicy::Enabled),
1432 (serving_table, CacheRefillPolicy::Enabled),
1433 ]),
1434 HashMap::new(),
1435 HashMap::from([(serving_table, serving_vnodes.clone())]),
1436 normal_delta(vec![
1437 sst_info(vec![streaming_table]),
1438 sst_info(vec![serving_table]),
1439 sst_info(vec![streaming_table, serving_table]),
1440 ]),
1441 ),
1442 vec![
1443 vec![streaming_table],
1444 vec![serving_table],
1445 vec![streaming_table, serving_table],
1446 ],
1447 );
1448
1449 assert_eq!(
1450 capture_pruned_table_ids(
1451 Role::Both,
1452 CacheRefillPolicy::Disabled,
1453 HashMap::from([
1454 (streaming_table, CacheRefillPolicy::Both),
1455 (serving_table, CacheRefillPolicy::Both),
1456 ]),
1457 HashMap::new(),
1458 HashMap::new(),
1459 normal_delta(vec![
1460 sst_info(vec![streaming_table]),
1461 sst_info(vec![serving_table]),
1462 ]),
1463 ),
1464 Vec::<Vec<TableId>>::new(),
1465 );
1466
1467 assert_eq!(
1468 capture_pruned_table_ids(
1469 Role::Both,
1470 CacheRefillPolicy::Disabled,
1471 HashMap::from([
1472 (streaming_table, CacheRefillPolicy::Both),
1473 (serving_table, CacheRefillPolicy::Both),
1474 ]),
1475 HashMap::from([(streaming_table, serving_vnodes.clone())]),
1476 HashMap::new(),
1477 normal_delta(vec![
1478 sst_info(vec![streaming_table]),
1479 sst_info(vec![serving_table]),
1480 ]),
1481 ),
1482 vec![vec![streaming_table]],
1483 );
1484
1485 assert_eq!(
1486 capture_pruned_table_ids(
1487 Role::Both,
1488 CacheRefillPolicy::Disabled,
1489 HashMap::from([
1490 (streaming_table, CacheRefillPolicy::Both),
1491 (serving_table, CacheRefillPolicy::Both),
1492 ]),
1493 HashMap::new(),
1494 HashMap::from([(serving_table, serving_vnodes.clone())]),
1495 normal_delta(vec![
1496 sst_info(vec![streaming_table]),
1497 sst_info(vec![serving_table]),
1498 ]),
1499 ),
1500 vec![vec![serving_table]],
1501 );
1502
1503 assert_eq!(
1504 capture_pruned_table_ids(
1505 Role::Streaming,
1506 CacheRefillPolicy::Disabled,
1507 HashMap::from([
1508 (streaming_table, CacheRefillPolicy::Streaming),
1509 (serving_table, CacheRefillPolicy::Serving),
1510 ]),
1511 HashMap::new(),
1512 HashMap::new(),
1513 normal_delta(vec![
1514 sst_info(vec![streaming_table]),
1515 sst_info(vec![serving_table]),
1516 ]),
1517 ),
1518 Vec::<Vec<TableId>>::new(),
1519 );
1520
1521 assert_eq!(
1522 capture_pruned_table_ids(
1523 Role::Streaming,
1524 CacheRefillPolicy::Disabled,
1525 HashMap::from([
1526 (streaming_table, CacheRefillPolicy::Streaming),
1527 (serving_table, CacheRefillPolicy::Serving),
1528 ]),
1529 HashMap::new(),
1530 HashMap::new(),
1531 insert_only_delta(vec![
1532 sst_info(vec![streaming_table]),
1533 sst_info(vec![serving_table]),
1534 ]),
1535 ),
1536 Vec::<Vec<TableId>>::new(),
1537 );
1538
1539 assert_eq!(
1540 capture_pruned_table_ids(
1541 Role::Both,
1542 CacheRefillPolicy::Enabled,
1543 HashMap::from([(disabled_table, CacheRefillPolicy::Disabled)]),
1544 HashMap::new(),
1545 HashMap::new(),
1546 normal_delta(vec![
1547 sst_info(vec![disabled_table]),
1548 sst_info(vec![fallback_table]),
1549 sst_info(vec![disabled_table, fallback_table]),
1550 ]),
1551 ),
1552 vec![vec![fallback_table], vec![disabled_table, fallback_table]],
1553 );
1554 }
1555
1556 #[tokio::test]
1557 async fn test_data_refill_requires_disk_cache() {
1558 let fixture = DataRefillGeneratorTestFixture::new(None).await;
1559 let delta = fixture.normal_l0_delta();
1560 let mut context = fixture.context(CacheRefillPolicy::Enabled, None, None, |_| {});
1561 assert!(!fixture.generate(&context, &delta).await.is_empty());
1562
1563 context.sstable_store = mock_sstable_store().await;
1564 assert!(!context.sstable_store.block_cache().is_hybrid());
1565 assert!(fixture.generate(&context, &delta).await.is_empty());
1566 }
1567
1568 #[tokio::test]
1569 async fn test_normal_refill_applies_policy_and_vnode_ownership() {
1570 let fixture = DataRefillGeneratorTestFixture::new(None).await;
1571 let delta = fixture.normal_l0_delta();
1572 let owned = Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [0]);
1573 let unowned = Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [1]);
1574 let cases = vec![
1575 ("Enabled", CacheRefillPolicy::Enabled, None, None, true),
1576 (
1577 "Disabled",
1578 CacheRefillPolicy::Disabled,
1579 Some(owned.clone()),
1580 Some(owned.clone()),
1581 false,
1582 ),
1583 (
1584 "Streaming match",
1585 CacheRefillPolicy::Streaming,
1586 Some(owned.clone()),
1587 None,
1588 true,
1589 ),
1590 (
1591 "Streaming miss",
1592 CacheRefillPolicy::Streaming,
1593 Some(unowned.clone()),
1594 None,
1595 false,
1596 ),
1597 (
1598 "Streaming ownership missing",
1599 CacheRefillPolicy::Streaming,
1600 None,
1601 None,
1602 false,
1603 ),
1604 (
1605 "Serving match",
1606 CacheRefillPolicy::Serving,
1607 None,
1608 Some(owned.clone()),
1609 true,
1610 ),
1611 (
1612 "Serving miss",
1613 CacheRefillPolicy::Serving,
1614 None,
1615 Some(unowned.clone()),
1616 false,
1617 ),
1618 (
1619 "Serving ownership missing",
1620 CacheRefillPolicy::Serving,
1621 None,
1622 None,
1623 false,
1624 ),
1625 (
1626 "Both streaming match",
1627 CacheRefillPolicy::Both,
1628 Some(owned.clone()),
1629 Some(unowned.clone()),
1630 true,
1631 ),
1632 (
1633 "Both serving match",
1634 CacheRefillPolicy::Both,
1635 Some(unowned.clone()),
1636 Some(owned),
1637 true,
1638 ),
1639 (
1640 "Both misses",
1641 CacheRefillPolicy::Both,
1642 Some(unowned.clone()),
1643 Some(unowned),
1644 false,
1645 ),
1646 ];
1647
1648 for (name, policy, streaming_vnodes, serving_vnodes, should_refill) in cases {
1649 let context = fixture.context(policy, streaming_vnodes, serving_vnodes, |_| {});
1650 assert_eq!(
1651 !fixture.generate(&context, &delta).await.is_empty(),
1652 should_refill,
1653 "{name}"
1654 );
1655 }
1656 }
1657
1658 #[tokio::test]
1659 async fn test_normal_refill_applies_recent_and_inheritance_filters() {
1660 let recent_filter = SimpleRecentFilter::new(3, Duration::from_secs(60));
1661 let fixture =
1662 DataRefillGeneratorTestFixture::new(Some(Arc::new(recent_filter.clone().into()))).await;
1663 let delta = fixture.normal_l0_delta();
1664
1665 let serving_context = fixture.context(
1666 CacheRefillPolicy::Serving,
1667 None,
1668 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1669 |config| {
1670 config.skip_recent_filter = false;
1671 },
1672 );
1673 assert!(
1674 fixture.generate(&serving_context, &delta).await.is_empty(),
1675 "explicit Serving policy is not an implicit skip_recent_filter"
1676 );
1677
1678 recent_filter.insert((fixture.deleted_sst_object_id, usize::MAX));
1679 assert!(
1680 !fixture.generate(&serving_context, &delta).await.is_empty(),
1681 "explicit Serving policy should produce tasks after recent admission hits"
1682 );
1683
1684 let (_, parent_sst_info) =
1685 gen_test_sst_with_object_id(fixture.table_id, fixture.sstable_store.clone(), 2).await;
1686 let non_l0_delta = fixture.normal_delta(1, parent_sst_info.object_id);
1687 let non_l0_context = fixture.context(CacheRefillPolicy::Enabled, None, None, |config| {
1688 config.data_refill_levels.insert(1);
1689 config.skip_recent_filter = false;
1690 config.skip_inheritance_filter = false;
1691 });
1692
1693 recent_filter.insert((parent_sst_info.object_id, usize::MAX));
1694 let generator = DataCacheRefillTaskGenerator {
1695 context: &non_l0_context,
1696 delta: &non_l0_delta,
1697 ssts: std::slice::from_ref(&fixture.sst),
1698 };
1699 let unfiltered_tasks = generator.generate_unfiltered_tasks();
1700 assert_eq!(
1701 unfiltered_tasks
1702 .iter()
1703 .map(|task| task.blks.len())
1704 .sum::<usize>(),
1705 fixture.sst.block_count(),
1706 "recent-admitted blocks should reach the inheritance stage"
1707 );
1708 assert!(
1709 generator
1710 .filter_by_inheritance_if_needed(unfiltered_tasks)
1711 .await
1712 .is_empty(),
1713 "after recent admission, parent block recent miss should filter non-L0 normal refill"
1714 );
1715
1716 recent_filter.insert((parent_sst_info.object_id, 0));
1717 let tasks = fixture.generate(&non_l0_context, &non_l0_delta).await;
1718 assert_eq!(tasks.len(), 1);
1719 assert_eq!(tasks[0].sst.id, fixture.sst.id);
1720 assert_eq!(tasks[0].blks, 0..1);
1721 }
1722
1723 #[tokio::test]
1724 async fn test_l0_refill_after_truncation_uses_serving_ownership() {
1725 use risingwave_hummock_sdk::level::Levels;
1726 use risingwave_hummock_sdk::version::{
1727 GroupDelta, GroupDeltas, HummockVersion, HummockVersionStateTableInfo,
1728 };
1729 use risingwave_pb::hummock::{StateTableInfo, StateTableInfoDelta};
1730
1731 let fixture = DataRefillGeneratorTestFixture::new(None).await;
1732 let group = 1.into();
1733 let staging_table = TableId::new(234);
1734 let mut version = HummockVersion::default();
1735 version.levels.insert(
1736 group,
1737 Levels {
1738 group_id: group,
1739 ..Default::default()
1740 },
1741 );
1742 version.state_table_info = HummockVersionStateTableInfo::from_protobuf_owned(
1743 [fixture.table_id, staging_table]
1744 .map(|table_id| {
1745 (
1746 table_id,
1747 StateTableInfo {
1748 compaction_group_id: group,
1749 ..Default::default()
1750 },
1751 )
1752 })
1753 .into(),
1754 );
1755 let mut raw = version.version_delta_after();
1756 for table_id in [fixture.table_id, staging_table] {
1757 raw.state_table_info_delta.insert(
1758 table_id,
1759 StateTableInfoDelta {
1760 committed_epoch: test_epoch(233),
1761 compaction_group_id: group,
1762 },
1763 );
1764 }
1765 raw.group_deltas.insert(
1766 group,
1767 GroupDeltas {
1768 group_deltas: vec![
1769 GroupDelta::NewL0SubLevel(vec![fixture.sst_info.clone()]),
1770 GroupDelta::PruneTableIdsFromSsts([staging_table].into()),
1771 ],
1772 },
1773 );
1774 let infos = version.build_sst_delta_infos(&raw);
1775 let mut config = test_refill_config(CacheRefillPolicy::Serving);
1776 config.data_refill_levels.insert(0);
1777 let mut refiller = CacheRefiller::new(
1778 Role::Serving,
1779 config,
1780 fixture.sstable_store.clone(),
1781 CacheRefiller::default_spawn_refill_task(),
1782 );
1783 refiller.replace_serving_table_vnode_mapping(HashMap::from([(
1784 fixture.table_id,
1785 Bitmap::ones(VirtualNode::COUNT_FOR_TEST),
1786 )]));
1787 fixture.sstable_store.clear_block_cache().await.unwrap();
1788 let handle = refiller.start_cache_refill(infos.clone());
1789 tokio::time::timeout(Duration::from_secs(5), handle)
1790 .await
1791 .unwrap()
1792 .unwrap();
1793 for block_idx in 0..fixture.sst.block_count() {
1794 let key = super::SstableBlockIndex {
1795 sst_id: fixture.sst.id,
1796 block_idx: block_idx as u64,
1797 };
1798 assert!(
1799 fixture
1800 .sstable_store
1801 .block_cache()
1802 .get(&key)
1803 .await
1804 .unwrap()
1805 .is_some(),
1806 "a staging-table truncation must not hide another table's L0 refill"
1807 );
1808 }
1809 let delta = infos.into_iter().next().unwrap();
1810 let cases = [
1811 (
1812 "Enabled + serving overlap",
1813 CacheRefillPolicy::Enabled,
1814 None,
1815 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1816 true,
1817 ),
1818 (
1819 "Enabled without serving ownership",
1820 CacheRefillPolicy::Enabled,
1821 None,
1822 None,
1823 false,
1824 ),
1825 (
1826 "Serving + serving overlap",
1827 CacheRefillPolicy::Serving,
1828 None,
1829 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1830 true,
1831 ),
1832 (
1833 "Streaming + streaming overlap",
1834 CacheRefillPolicy::Streaming,
1835 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1836 None,
1837 false,
1838 ),
1839 (
1840 "Both + streaming overlap + serving non-overlap",
1841 CacheRefillPolicy::Both,
1842 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1843 Some(Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [1])),
1844 false,
1845 ),
1846 (
1847 "Both + serving overlap",
1848 CacheRefillPolicy::Both,
1849 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1850 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1851 true,
1852 ),
1853 (
1854 "Disabled + serving overlap",
1855 CacheRefillPolicy::Disabled,
1856 None,
1857 Some(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
1858 false,
1859 ),
1860 ];
1861
1862 for (name, policy, streaming_vnodes, serving_vnodes, should_refill) in cases {
1863 let context = fixture.context(policy, streaming_vnodes, serving_vnodes, |config| {
1864 config.skip_recent_filter = false;
1865 config.skip_inheritance_filter = false;
1866 });
1867 assert_eq!(
1868 !fixture.generate(&context, &delta).await.is_empty(),
1869 should_refill,
1870 "{name}"
1871 );
1872 }
1873 }
1874
1875 #[tokio::test]
1876 async fn test_refill_units_do_not_cross_table_projection_boundaries() {
1877 let table_a = TableId::from(233);
1878 let table_b = TableId::from(234);
1879 let (sstable_store, _cache_dir) = mock_sstable_store_with_disk_cache(None).await;
1880 let (sst, sst_info) = gen_test_sstable_with_table_ids(
1881 default_builder_opt_for_test(),
1882 1,
1883 [table_a, table_b].into_iter().map(|table_id| {
1884 (
1885 FullKey {
1886 user_key: UserKey::for_test(table_id, iterator_test_table_key_of(0)),
1887 epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(233)),
1888 },
1889 HummockValue::put(b"value".to_vec()),
1890 )
1891 }),
1892 sstable_store.clone(),
1893 vec![table_a.as_raw_id(), table_b.as_raw_id()],
1894 )
1895 .await;
1896 assert_eq!(sst.block_count(), 2, "table switch must form a new block");
1897
1898 let mut next_sst_id = 100.into();
1899 let (table_a_projection, table_b_projection) =
1900 split_sst_with_table_ids(&sst_info, &mut next_sst_id, 1, 1, vec![table_b]);
1901 assert_eq!(table_a_projection.object_id, sst_info.object_id);
1902 assert_eq!(table_b_projection.object_id, sst_info.object_id);
1903 assert_ne!(table_a_projection.sst_id, table_b_projection.sst_id);
1904 assert_eq!(table_a_projection.table_ids, vec![table_a]);
1905 assert_eq!(table_b_projection.table_ids, vec![table_b]);
1906
1907 let deltas = [table_a_projection, table_b_projection].map(|projection| SstDeltaInfo {
1908 insert_sst_infos: vec![projection],
1909 delete_sst_object_ids: vec![],
1910 insert_sst_level: 0,
1911 });
1912 let normal_deltas = deltas.clone().map(|mut delta| {
1913 delta.delete_sst_object_ids = vec![999.into()];
1916 delta
1917 });
1918 let serving_vnodes = Bitmap::ones(VirtualNode::COUNT_FOR_TEST);
1919 let table_cache_refill_context_map = Arc::new(
1920 [table_a, table_b]
1921 .into_iter()
1922 .map(|table_id| {
1923 (
1924 table_id,
1925 super::TableCacheRefillContext {
1926 streaming_vnode_bitmap: None,
1927 serving_vnode_bitmap: Some(serving_vnodes.clone()),
1928 policy: CacheRefillPolicy::Serving,
1929 },
1930 )
1931 })
1932 .collect::<super::TableCacheRefillContextMap>(),
1933 );
1934 let make_context = |unit| {
1935 let mut config = test_refill_config(CacheRefillPolicy::Disabled);
1936 config.data_refill_levels.insert(0);
1937 config.unit = unit;
1938 CacheRefillContext {
1939 config: Arc::new(config),
1940 meta_refill_concurrency: None,
1941 concurrency: Arc::new(tokio::sync::Semaphore::new(1)),
1942 sstable_store: sstable_store.clone(),
1943 table_cache_refill_context_map: table_cache_refill_context_map.clone(),
1944 }
1945 };
1946 let generated_tasks = |context: &CacheRefillContext| {
1947 deltas
1948 .iter()
1949 .map(|delta| {
1950 DataCacheRefillTaskGenerator {
1951 context,
1952 delta,
1953 ssts: std::slice::from_ref(&sst),
1954 }
1955 .generate_unfiltered_tasks()
1956 })
1957 .collect::<Vec<_>>()
1958 };
1959 let generated_ranges = |context: &CacheRefillContext| {
1960 generated_tasks(context)
1961 .into_iter()
1962 .map(|tasks| tasks.into_iter().map(|task| task.blks).collect::<Vec<_>>())
1963 .collect::<Vec<_>>()
1964 };
1965
1966 assert_eq!(
1967 generated_ranges(&make_context(1)),
1968 vec![vec![0..1], vec![1..2]],
1969 "each logical projection must select only its own block"
1970 );
1971
1972 let wide_unit_context = make_context(2);
1973 assert_eq!(
1974 generated_ranges(&wide_unit_context),
1975 vec![vec![0..1], vec![1..2]],
1976 "units are clipped at table boundaries even when unit is larger than a table run"
1977 );
1978
1979 let normal_ranges = normal_deltas
1980 .iter()
1981 .map(|delta| {
1982 DataCacheRefillTaskGenerator {
1983 context: &wide_unit_context,
1984 delta,
1985 ssts: std::slice::from_ref(&sst),
1986 }
1987 .generate_unfiltered_tasks()
1988 .into_iter()
1989 .map(|task| task.blks)
1990 .collect::<Vec<_>>()
1991 })
1992 .collect::<Vec<_>>();
1993 assert_eq!(
1994 normal_ranges,
1995 vec![vec![0..1], vec![1..2]],
1996 "normal refill uses the same table-boundary geometry"
1997 );
1998
1999 for task in generated_tasks(&wide_unit_context).into_iter().flatten() {
2000 assert!(task.blks.len() <= wide_unit_context.config.unit);
2001 assert_eq!(
2002 task.sst.meta.block_metas[task.blks.start].table_id(),
2003 task.sst.meta.block_metas[task.blks.end - 1].table_id(),
2004 "a refill unit must not cross a table boundary"
2005 );
2006 }
2007 }
2008
2009 #[tokio::test]
2010 async fn test_scoped_refill_handles_multi_table_vnode_boundary() {
2011 let table_a = TableId::from(233);
2012 let table_b = TableId::from(234);
2013 let vnode_a = VirtualNode::COUNT_FOR_TEST - 1;
2014 let (sstable_store, _cache_dir) = mock_sstable_store_with_disk_cache(None).await;
2015 let (sst, sst_info) = gen_test_sstable_with_table_ids(
2016 default_builder_opt_for_test(),
2017 1,
2018 [
2019 (
2020 FullKey {
2021 user_key: UserKey::for_test(
2022 table_a,
2023 prefix_slice_with_vnode(VirtualNode::from_index(vnode_a), b"table_a"),
2024 ),
2025 epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(233)),
2026 },
2027 HummockValue::put(Bytes::from_static(b"a")),
2028 ),
2029 (
2030 FullKey {
2031 user_key: UserKey::for_test(
2032 table_b,
2033 prefix_slice_with_vnode(VirtualNode::ZERO, b"table_b"),
2034 ),
2035 epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(233)),
2036 },
2037 HummockValue::put(Bytes::from_static(b"b")),
2038 ),
2039 ]
2040 .into_iter(),
2041 sstable_store.clone(),
2042 vec![table_a.as_raw_id(), table_b.as_raw_id()],
2043 )
2044 .await;
2045 assert_eq!(sst.block_count(), 2, "table switch must form a new block");
2046
2047 let generate = |streaming_vnodes| {
2048 let sstable_store = sstable_store.clone();
2049 let sst = sst.clone();
2050 let sst_info = sst_info.clone();
2051 let mut config = test_refill_config(CacheRefillPolicy::Streaming);
2052 config.data_refill_levels.insert(0);
2053 let context = CacheRefillContext {
2054 config: Arc::new(config),
2055 meta_refill_concurrency: None,
2056 concurrency: Arc::new(tokio::sync::Semaphore::new(1)),
2057 sstable_store,
2058 table_cache_refill_context_map: Arc::new(HashMap::from([(
2059 table_a,
2060 super::TableCacheRefillContext {
2061 streaming_vnode_bitmap: Some(streaming_vnodes),
2062 serving_vnode_bitmap: None,
2063 policy: CacheRefillPolicy::Streaming,
2064 },
2065 )])),
2066 };
2067 async move {
2068 let generator = DataCacheRefillTaskGenerator {
2069 context: &context,
2070 delta: &SstDeltaInfo {
2071 insert_sst_infos: vec![sst_info.clone()],
2072 delete_sst_object_ids: vec![2330.into()],
2073 insert_sst_level: 0,
2074 },
2075 ssts: std::slice::from_ref(&sst),
2076 };
2077 let tasks = generator.generate_unfiltered_tasks();
2078 generator.filter_by_inheritance_if_needed(tasks).await
2079 }
2080 };
2081
2082 let matching_tasks =
2083 generate(Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [vnode_a])).await;
2084 assert_eq!(matching_tasks.len(), 1);
2085 assert_eq!(matching_tasks[0].blks, 0..1);
2086
2087 let non_matching_tasks = generate(Bitmap::from_indices(
2088 VirtualNode::COUNT_FOR_TEST,
2089 [VirtualNode::ZERO.to_index()],
2090 ))
2091 .await;
2092 assert!(non_matching_tasks.is_empty());
2093 }
2094
2095 #[tokio::test]
2096 async fn test_block_vnode_range_handles_vnode_only_block_boundaries() {
2097 let table_id = TableId::from(233);
2098 let vnode = VirtualNode::ZERO;
2099 let sstable_store = mock_sstable_store().await;
2100 let mut builder_options = default_builder_opt_for_test();
2101 builder_options.block_capacity = 1;
2102 let (sst, _) = gen_test_sstable_with_table_ids(
2103 builder_options,
2104 1,
2105 [234, 233].into_iter().map(|epoch| {
2106 (
2107 FullKey {
2108 user_key: UserKey::for_test(table_id, prefix_slice_with_vnode(vnode, b"")),
2109 epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(epoch)),
2110 },
2111 HummockValue::put(Bytes::from_static(b"value")),
2112 )
2113 }),
2114 sstable_store.clone(),
2115 vec![table_id.as_raw_id()],
2116 )
2117 .await;
2118 assert_eq!(sst.block_count(), 2);
2119 let expected = (vnode.to_index(), vnode.to_index() + 1);
2120 assert_eq!(block_vnode_range(&sst, 0), expected);
2121 assert_eq!(block_vnode_range(&sst, 1), expected);
2122 }
2123
2124 #[tokio::test]
2125 async fn test_block_vnode_range_fails_open_for_shortened_meta_keys() {
2126 let table_id = TableId::from(233);
2127 let sstable_store = mock_sstable_store().await;
2128 let mut builder_options = default_builder_opt_for_test();
2129 builder_options.block_capacity = 1;
2130 builder_options.shorten_block_meta_key_threshold = Some(0);
2131 let (sst, _) = gen_test_sstable_with_table_ids(
2132 builder_options,
2133 1,
2134 [255, 256].into_iter().map(|vnode| {
2135 (
2136 FullKey {
2137 user_key: UserKey::for_test(
2138 table_id,
2139 prefix_slice_with_vnode(VirtualNode::from_index(vnode), b"long-key"),
2140 ),
2141 epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(233)),
2142 },
2143 HummockValue::put(Bytes::from_static(b"value")),
2144 )
2145 }),
2146 sstable_store,
2147 vec![table_id.as_raw_id()],
2148 )
2149 .await;
2150 assert_eq!(sst.block_count(), 2);
2151 assert!(
2152 FullKey::decode(&sst.meta.block_metas[1].smallest_key)
2153 .user_key
2154 .table_key
2155 .as_ref()
2156 .len()
2157 < VirtualNode::SIZE
2158 );
2159 let full_range = (0, VirtualNode::MAX_REPRESENTABLE.to_index() + 1);
2160 assert_eq!(block_vnode_range(&sst, 0), full_range);
2161 assert_eq!(block_vnode_range(&sst, 1), full_range);
2162 }
2163
2164 #[test]
2165 fn test_vnode_range_overlaps_bitmap_uses_right_exclusive_end() {
2166 let right_exclusive = Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [12]);
2167 assert!(!vnode_range_overlaps_bitmap((10, 12), &right_exclusive));
2168
2169 let inside_range = Bitmap::from_indices(VirtualNode::COUNT_FOR_TEST, [11]);
2170 assert!(vnode_range_overlaps_bitmap((10, 12), &inside_range));
2171
2172 let last_vnode = Bitmap::from_indices(
2173 VirtualNode::COUNT_FOR_TEST,
2174 [VirtualNode::COUNT_FOR_TEST - 1],
2175 );
2176 assert!(vnode_range_overlaps_bitmap(
2177 (VirtualNode::COUNT_FOR_TEST - 1, VirtualNode::COUNT_FOR_TEST),
2178 &last_vnode
2179 ));
2180 assert!(!vnode_range_overlaps_bitmap(
2181 (VirtualNode::COUNT_FOR_TEST, VirtualNode::COUNT_FOR_TEST + 1),
2182 &last_vnode
2183 ));
2184 }
2185}