Skip to main content

risingwave_storage/hummock/event_handler/
refiller.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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    /// Cache refill timeout.
50    pub timeout: Duration,
51
52    /// Data file cache refill levels.
53    pub data_refill_levels: HashSet<u32>,
54
55    /// Meta file cache refill concurrency.
56    pub meta_refill_concurrency: usize,
57
58    /// Data file cache refill concurrency.
59    pub concurrency: usize,
60
61    /// Data file cache refill unit (blocks).
62    pub unit: usize,
63
64    /// Data file cache reill unit threshold.
65    ///
66    /// Only units whose admit rate > threshold will be refilled.
67    pub threshold: f64,
68
69    /// Skip recent filter.
70    pub skip_recent_filter: bool,
71
72    /// Skip inheritance filter.
73    pub skip_inheritance_filter: bool,
74
75    /// Default table cache refill policy.
76    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/// Per-table metadata captured for a refill task to decide whether an sstable block should be
114/// refilled. Mutable runtime state used to build this snapshot is owned by `CacheRefiller`.
115#[derive(Clone)]
116pub struct TableCacheRefillContext {
117    /// Vnodes covered by local streaming read versions on this compute node.
118    pub streaming_vnode_bitmap: Option<Bitmap>,
119    /// Vnodes served by this compute node according to the serving vnode mapping.
120    pub serving_vnode_bitmap: Option<Bitmap>,
121    /// Effective refill policy after applying the default policy and per-table overrides.
122    pub policy: CacheRefillPolicy,
123}
124
125/// Read-only data cloned from `CacheRefiller` for monitor/debugging APIs.
126///
127/// Streaming vnode mapping is the table-level union maintained by the refiller.
128#[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        // Insert-only deltas have no delete-side evidence for recent/inheritance filters.
164        // Only serving-owned blocks need refill, because streaming writers already populated
165        // their local cache.
166        (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        // A table switch always starts a new block. The next table's smallest key has an
190        // unrelated vnode, so use the current table's terminal range instead.
191        Some(next_block_meta) if next_block_meta.table_id() != block_meta.table_id() => {
192            Bound::Unbounded
193        }
194        // Full-key versions of the same table key may span adjacent blocks. After projecting
195        // away the epoch, the boundary vnode therefore remains part of the current block.
196        Some(next_block_meta) => Bound::Included(
197            FullKey::decode(&next_block_meta.smallest_key)
198                .user_key
199                .table_key,
200        ),
201        // `SstableMeta::largest_key` is the actual last key, unlike the next block's smallest
202        // key above. Keep it inclusive, especially for singleton tables whose key contains only
203        // the vnode prefix.
204        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    // Block-meta separators may shorten the table key below the vnode prefix. They are valid
216    // full-key search boundaries but cannot identify a vnode, so fail open instead of panicking
217    // or dropping a block that may belong to this worker.
218    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
230/// A cache refiller for hummock data.
231pub(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            // Writer-appended L0 SSTs are already warm on the streaming side. Their data refill
286            // may therefore only be needed by serving workers.
287            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            // This is deliberately a whole-SST admission check before Meta load. The serving
296            // mapping key set identifies result tables, but bitmap contents remain for the exact
297            // block/vnode decision in DataCacheRefillTaskGenerator after Meta load. An SST stays
298            // when any contained table matches either the serving or streaming refill lane.
299            delta.insert_sst_infos.retain(|sst| {
300                sst.table_ids.iter().any(|table_id| {
301                    // A missing entry means there is no table override, not that the table is
302                    // absent. Preserve the configured legacy/default policy in that case.
303                    let policy = self
304                        .table_cache_refill_policies
305                        .get(table_id)
306                        .copied()
307                        .unwrap_or(self.default_policy);
308
309                    // Enabled preserves legacy full refill. For scoped policies, mapping keys are
310                    // only a whole-SST coarse gate; bitmap bits still filter blocks post-Meta.
311                    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    /// Replaces the complete policy snapshot applicable to this worker.
352    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    /// Replaces the complete serving vnode mapping snapshot.
360    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                // `Enabled` normally does not use bitmap filtering. The only exception is L0
404                // insert-only refill, where serving workers still need serving-locality evidence.
405                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        // Skip data cache refill if data disk cache is not enabled.
456        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        // CN-written SSTs are appended to L0 without replacing parent SSTs. Other inserted SSTs
466        // need delete-side evidence for recent and inheritance filtering.
467        debug_assert!(has_parent_ssts || self.delta.insert_sst_level == 0);
468
469        // Return if the target level is not in the refill levels
470        if !self
471            .context
472            .config
473            .data_refill_levels
474            .contains(&self.delta.insert_sst_level)
475        {
476            return tasks;
477        }
478
479        // Cache refill units must not cross a table boundary. A logical SST projection still
480        // decides whether to admit each single-table unit.
481        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                // SstableBuilder ends a block before the table ID changes, so block metadata
490                // defines the exact physical boundary. `table_ids` below only admits logical
491                // projections and must not make a unit span another table.
492                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        // Policy/vnode ownership defines refill responsibility first, but it does not bypass
527        // recent admission for normal insert+delete refill.
528        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        // Skipping the recent filter selects full refill. Inheritance filtering only applies to
546        // non-L0 normal refill after real recent-filter admission.
547        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    // Return if recent filter is required and no deleted sst ids are in the recent filter.
560    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        // Get parent sst metas from cache.
576        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        // assert units in asc order
602        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                    // `largest_key` can be included or excluded, both are treated as included here
618                    &psst.meta.largest_key
619                } else {
620                    &psst.meta.block_metas[pblk + 1].smallest_key
621                };
622
623                // partition point: unit.right < pblk.left
624                let uleft = originals.partition_point(|task| {
625                    KeyComparator::compare_encoded_full_key(task.largest_key(), pleft)
626                        == std::cmp::Ordering::Less
627                });
628                // partition point: unit.left <= pblk.right
629                let uright = originals.partition_point(|task| {
630                    KeyComparator::compare_encoded_full_key(task.smallest_key(), pright)
631                        != std::cmp::Ordering::Greater
632                });
633
634                // overlapping: uleft..uright
635                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                    // Main counts after recent admission but before inheritance.
718                    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            // update filter for sst id only
772            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                            // The entry should always be `Some(..)`, use if here for compatible.
831                            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            // Multiple SSTs in a delta may refer to the same table.
1216            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            // A synthetic delete marks this as a normal delta; recent and inheritance filters
1914            // are disabled below, so the test does not rely on a matching parent SST.
1915            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}