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