risingwave_storage/hummock/
sstable_store.rs

1// Copyright 2022 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::clone::Clone;
16use std::collections::VecDeque;
17use std::ops::Deref;
18use std::pin::Pin;
19use std::sync::Arc;
20use std::sync::atomic::{AtomicUsize, Ordering};
21
22use await_tree::{InstrumentAwait, SpanExt};
23use bytes::Bytes;
24use fail::fail_point;
25use foyer::{
26    Cache, CacheBuilder, CacheEntry, EventListener, Hint, HybridCache, HybridCacheBuilder,
27    HybridCacheEntry, HybridCacheProperties,
28};
29use futures::{FutureExt, StreamExt, future};
30use prost::Message;
31use risingwave_hummock_sdk::sstable_info::SstableInfo;
32use risingwave_hummock_sdk::vector_index::{HnswGraphFileInfo, VectorFileInfo};
33use risingwave_hummock_sdk::{
34    HummockHnswGraphFileId, HummockObjectId, HummockRawObjectId, HummockSstableObjectId,
35    HummockVectorFileId, SST_OBJECT_SUFFIX,
36};
37use risingwave_hummock_trace::TracedCachePolicy;
38use risingwave_object_store::object::{
39    ObjectError, ObjectMetadataIter, ObjectResult, ObjectStoreRef, ObjectStreamingUploader,
40};
41use risingwave_pb::hummock::PbHnswGraph;
42use serde::{Deserialize, Serialize};
43use thiserror_ext::AsReport;
44use tokio::time::Instant;
45
46use super::{
47    BatchUploadWriter, Block, BlockMeta, BlockResponse, RecentFilter, Sstable, SstableMeta,
48    SstableWriterOptions,
49};
50use crate::hummock::block_stream::{BlockDataStream, MemoryUsageTracker, PrefetchBlockStream};
51use crate::hummock::none::NoneRecentFilter;
52use crate::hummock::vector::file::{VectorBlock, VectorBlockMeta, VectorFileMeta};
53use crate::hummock::vector::monitor::VectorStoreCacheStats;
54use crate::hummock::{BlockEntry, BlockHolder, HummockError, HummockResult, RecentFilterTrait};
55use crate::monitor::{HummockStateStoreMetrics, StoreLocalStatistic};
56
57macro_rules! impl_vector_index_meta_file {
58    ($($type_name:ident),+) => {
59        pub enum HummockVectorIndexMetaFile {
60            $(
61                $type_name(Pin<Box<$type_name>>),
62            )+
63        }
64
65        $(
66            impl From<$type_name> for HummockVectorIndexMetaFile {
67                fn from(v: $type_name) -> Self {
68                    Self::$type_name(Box::pin(v))
69                }
70            }
71
72            unsafe impl Send for VectorMetaFileHolder<$type_name> {}
73
74            impl VectorMetaFileHolder<$type_name> {
75                fn try_from_entry(
76                    entry: CacheEntry<HummockRawObjectId, HummockVectorIndexMetaFile>,
77                    object_id: HummockRawObjectId
78                ) -> HummockResult<Self> {
79                    let HummockVectorIndexMetaFile::$type_name(file_meta) = &*entry else {
80                        return Err(HummockError::decode_error(format!(
81                            "expect {} for object {}",
82                            stringify!($type_name),
83                            object_id
84                        )));
85                    };
86                    let ptr = file_meta.as_ref().get_ref() as *const _;
87                    Ok(VectorMetaFileHolder {
88                        _cache_entry: entry,
89                        ptr,
90                    })
91                }
92            }
93        )+
94    };
95}
96
97impl_vector_index_meta_file!(VectorFileMeta, PbHnswGraph);
98
99pub struct VectorMetaFileHolder<T> {
100    _cache_entry: CacheEntry<HummockRawObjectId, HummockVectorIndexMetaFile>,
101    ptr: *const T,
102}
103
104impl<T> Deref for VectorMetaFileHolder<T> {
105    type Target = T;
106
107    fn deref(&self) -> &Self::Target {
108        // SAFETY: VectorFileHolder is exposed only as immutable, and `VectorFileMeta` is pinned via box
109        unsafe { &*self.ptr }
110    }
111}
112
113/// Meta-cache value with the existing bincode format and allocation-free filter sizing.
114///
115/// A separate type lets this cache override `Code::estimated_size` without changing
116/// `Sstable`'s serde implementation or conflicting with Foyer's serde blanket impl.
117#[derive(Debug)]
118pub struct SstableMetaCacheValue(Box<Sstable>);
119
120impl From<Sstable> for SstableMetaCacheValue {
121    fn from(sstable: Sstable) -> Self {
122        Self(Box::new(sstable))
123    }
124}
125
126impl Deref for SstableMetaCacheValue {
127    type Target = Sstable;
128
129    fn deref(&self) -> &Self::Target {
130        &self.0
131    }
132}
133
134impl AsRef<Sstable> for SstableMetaCacheValue {
135    fn as_ref(&self) -> &Sstable {
136        &self.0
137    }
138}
139
140impl foyer::Code for SstableMetaCacheValue {
141    fn encode(&self, writer: &mut impl std::io::Write) -> foyer::Result<()> {
142        bincode::serialize_into(writer, &self.0).map_err(foyer::Error::bincode_error)
143    }
144
145    fn decode(reader: &mut impl std::io::Read) -> foyer::Result<Self> {
146        bincode::deserialize_from(reader)
147            .map(Self)
148            .map_err(foyer::Error::bincode_error)
149    }
150
151    fn estimated_size(&self) -> usize {
152        self.0.serde_encoded_size()
153    }
154}
155
156pub type TableHolder = HybridCacheEntry<HummockSstableObjectId, SstableMetaCacheValue>;
157
158pub type VectorBlockHolder = CacheEntry<(HummockVectorFileId, usize), Box<VectorBlock>>;
159
160pub type VectorFileHolder = VectorMetaFileHolder<VectorFileMeta>;
161pub type HnswGraphFileHolder = VectorMetaFileHolder<PbHnswGraph>;
162
163#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash, Serialize, Deserialize)]
164pub struct SstableBlockIndex {
165    pub sst_id: HummockSstableObjectId,
166    pub block_idx: u64,
167}
168
169pub struct BlockCacheEventListener {
170    metrics: Arc<HummockStateStoreMetrics>,
171}
172
173impl BlockCacheEventListener {
174    pub fn new(metrics: Arc<HummockStateStoreMetrics>) -> Self {
175        Self { metrics }
176    }
177}
178
179impl EventListener for BlockCacheEventListener {
180    type Key = SstableBlockIndex;
181    type Value = Box<Block>;
182
183    fn on_leave(&self, _reason: foyer::Event, _key: &Self::Key, value: &Self::Value)
184    where
185        Self::Key: foyer::Key,
186        Self::Value: foyer::Value,
187    {
188        self.metrics
189            .block_efficiency_histogram
190            .observe(value.efficiency());
191    }
192}
193
194// TODO: Define policy based on use cases (read / compaction / ...).
195#[derive(Clone, Copy, Eq, PartialEq)]
196pub enum CachePolicy {
197    /// Disable read cache and not fill the cache afterwards.
198    Disable,
199    /// Try reading the cache and fill the cache afterwards.
200    Fill(Hint),
201    /// Read the cache but not fill the cache afterwards.
202    NotFill,
203}
204
205impl Default for CachePolicy {
206    fn default() -> Self {
207        CachePolicy::Fill(Hint::Normal)
208    }
209}
210
211impl From<TracedCachePolicy> for CachePolicy {
212    fn from(policy: TracedCachePolicy) -> Self {
213        match policy {
214            TracedCachePolicy::Disable => Self::Disable,
215            TracedCachePolicy::Fill(priority) => Self::Fill(priority.into()),
216            TracedCachePolicy::NotFill => Self::NotFill,
217        }
218    }
219}
220
221impl From<CachePolicy> for TracedCachePolicy {
222    fn from(policy: CachePolicy) -> Self {
223        match policy {
224            CachePolicy::Disable => Self::Disable,
225            CachePolicy::Fill(priority) => Self::Fill(priority.into()),
226            CachePolicy::NotFill => Self::NotFill,
227        }
228    }
229}
230
231pub struct SstableStoreConfig {
232    pub store: ObjectStoreRef,
233    pub path: String,
234
235    pub prefetch_buffer_capacity: usize,
236    pub max_prefetch_block_number: usize,
237    pub recent_filter: Arc<RecentFilter<(HummockSstableObjectId, usize)>>,
238    pub state_store_metrics: Arc<HummockStateStoreMetrics>,
239    pub use_new_object_prefix_strategy: bool,
240    pub skip_bloom_filter_in_serde: bool,
241
242    pub meta_cache: HybridCache<HummockSstableObjectId, SstableMetaCacheValue>,
243    pub block_cache: HybridCache<SstableBlockIndex, Box<Block>>,
244
245    pub vector_meta_cache: Cache<HummockRawObjectId, HummockVectorIndexMetaFile>,
246    pub vector_block_cache: Cache<(HummockVectorFileId, usize), Box<VectorBlock>>,
247}
248
249pub struct SstableStore {
250    path: String,
251    store: ObjectStoreRef,
252
253    meta_cache: HybridCache<HummockSstableObjectId, SstableMetaCacheValue>,
254    block_cache: HybridCache<SstableBlockIndex, Box<Block>>,
255    pub vector_meta_cache: Cache<HummockRawObjectId, HummockVectorIndexMetaFile>,
256    pub vector_block_cache: Cache<(HummockVectorFileId, usize), Box<VectorBlock>>,
257
258    /// Recent filter for `(sst_obj_id, blk_idx)`.
259    ///
260    /// `blk_idx == USIZE::MAX` stands for `sst_obj_id` only entry.
261    recent_filter: Arc<RecentFilter<(HummockSstableObjectId, usize)>>,
262    prefetch_buffer_usage: Arc<AtomicUsize>,
263    prefetch_buffer_capacity: usize,
264    max_prefetch_block_number: usize,
265    /// Whether the object store is divided into prefixes depends on two factors:
266    ///   1. The specific object store type.
267    ///   2. Whether the existing cluster is a new cluster.
268    ///
269    /// The value of `use_new_object_prefix_strategy` is determined by the `use_new_object_prefix_strategy` field in the system parameters.
270    /// For a new cluster, `use_new_object_prefix_strategy` is set to True.
271    /// For an old cluster, `use_new_object_prefix_strategy` is set to False.
272    /// The final decision of whether to divide prefixes is based on this field and the specific object store type, this approach is implemented to ensure backward compatibility.
273    use_new_object_prefix_strategy: bool,
274
275    /// sst serde happens when a sst meta is written to meta disk cache.
276    /// Excluding the SST filter from serde can reduce the meta disk cache entry size
277    /// and reduce disk IO throughput at the cost of making the SST filter useless.
278    skip_bloom_filter_in_serde: bool,
279}
280
281impl SstableStore {
282    pub fn new(config: SstableStoreConfig) -> Self {
283        // TODO: We should validate path early. Otherwise object store won't report invalid path
284        // error until first write attempt.
285
286        Self {
287            path: config.path,
288            store: config.store,
289
290            meta_cache: config.meta_cache,
291            block_cache: config.block_cache,
292            vector_meta_cache: config.vector_meta_cache,
293            vector_block_cache: config.vector_block_cache,
294
295            recent_filter: config.recent_filter,
296            prefetch_buffer_usage: Arc::new(AtomicUsize::new(0)),
297            prefetch_buffer_capacity: config.prefetch_buffer_capacity,
298            max_prefetch_block_number: config.max_prefetch_block_number,
299            use_new_object_prefix_strategy: config.use_new_object_prefix_strategy,
300            skip_bloom_filter_in_serde: config.skip_bloom_filter_in_serde,
301        }
302    }
303
304    /// For compactor, we do not need a high concurrency load for cache. Instead, we need the cache
305    ///  can be evict more effective.
306    #[expect(clippy::borrowed_box)]
307    pub async fn for_compactor(
308        store: ObjectStoreRef,
309        path: String,
310        block_cache_capacity: usize,
311        meta_cache_capacity: usize,
312        use_new_object_prefix_strategy: bool,
313    ) -> HummockResult<Self> {
314        let meta_cache = HybridCacheBuilder::new()
315            .memory(meta_cache_capacity)
316            .with_shards(1)
317            .with_weighter(
318                |_: &HummockSstableObjectId, value: &SstableMetaCacheValue| {
319                    std::mem::size_of::<HummockSstableObjectId>()
320                        + value.estimated_meta_cache_memory_weight()
321                },
322            )
323            .storage()
324            .build()
325            .await
326            .map_err(HummockError::foyer_error)?;
327
328        let block_cache = HybridCacheBuilder::new()
329            .memory(block_cache_capacity)
330            .with_shards(1)
331            .with_weighter(|_: &SstableBlockIndex, value: &Box<Block>| {
332                std::mem::size_of::<SstableBlockIndex>() + value.estimated_memory_weight()
333            })
334            .storage()
335            .build()
336            .await
337            .map_err(HummockError::foyer_error)?;
338
339        Ok(Self {
340            path,
341            store,
342
343            prefetch_buffer_usage: Arc::new(AtomicUsize::new(0)),
344            prefetch_buffer_capacity: block_cache_capacity,
345            max_prefetch_block_number: 16, /* compactor won't use this parameter, so just assign a default value. */
346            recent_filter: Arc::new(NoneRecentFilter::default().into()),
347            use_new_object_prefix_strategy,
348            skip_bloom_filter_in_serde: false,
349
350            meta_cache,
351            block_cache,
352            vector_meta_cache: CacheBuilder::new(1 << 10).build(),
353            vector_block_cache: CacheBuilder::new(1 << 10).build(),
354        })
355    }
356
357    pub async fn delete(&self, object_id: HummockSstableObjectId) -> HummockResult<()> {
358        self.store
359            .delete(self.get_sst_data_path(object_id).as_str())
360            .await?;
361        self.meta_cache.remove(&object_id);
362        // TODO(MrCroxx): support group remove in foyer.
363        Ok(())
364    }
365
366    pub(crate) async fn put_sst_data(
367        &self,
368        object_id: HummockSstableObjectId,
369        data: Bytes,
370    ) -> HummockResult<()> {
371        let data_path = self.get_sst_data_path(object_id);
372        self.store
373            .upload(&data_path, data)
374            .await
375            .map_err(Into::into)
376    }
377
378    /// Buffers consecutive decoded blocks starting at `block_index`, up to `end_index` (exclusive).
379    /// Cache hits, the memory budget, and the prefetch limit may shorten the returned sequence.
380    /// All I/O and decoding errors are returned here, before the buffered blocks are consumed.
381    pub async fn prefetch_blocks(
382        &self,
383        sst: &Sstable,
384        block_index: usize,
385        end_index: usize,
386        policy: CachePolicy,
387        stats: &mut StoreLocalStatistic,
388    ) -> HummockResult<Box<PrefetchBlockStream>> {
389        let object_id = sst.id;
390        if self.prefetch_buffer_usage.load(Ordering::Acquire) > self.prefetch_buffer_capacity {
391            let block = self.get(sst, block_index, policy, stats).await?;
392            return Ok(Box::new(PrefetchBlockStream::new(
393                VecDeque::from([block]),
394                block_index,
395                None,
396            )));
397        }
398        if policy != CachePolicy::Disable
399            && let Some(entry) = self
400                .block_cache
401                .get(&SstableBlockIndex {
402                    sst_id: object_id,
403                    block_idx: block_index as _,
404                })
405                .await
406                .map_err(HummockError::foyer_error)?
407        {
408            stats.cache_data_block_total += 1;
409            if entry.source() == foyer::Source::Outer {
410                stats.cache_data_block_miss += 1;
411            }
412            let block = BlockHolder::from_hybrid_cache_entry(entry);
413            return Ok(Box::new(PrefetchBlockStream::new(
414                VecDeque::from([block]),
415                block_index,
416                None,
417            )));
418        }
419        let end_index = std::cmp::min(
420            end_index,
421            block_index.saturating_add(self.max_prefetch_block_number),
422        );
423        let mut end_index = std::cmp::min(end_index, sst.meta.block_metas.len());
424        let start_offset = sst.meta.block_metas[block_index].offset as usize;
425        if policy != CachePolicy::Disable {
426            let mut min_hit_index = end_index;
427            let mut hit_count = 0;
428            for idx in block_index..end_index {
429                if self.block_cache.contains(&SstableBlockIndex {
430                    sst_id: object_id,
431                    block_idx: idx as _,
432                }) {
433                    if min_hit_index > idx && idx > block_index {
434                        min_hit_index = idx;
435                    }
436                    hit_count += 1;
437                }
438            }
439
440            if hit_count * 3 >= (end_index - block_index)
441                || min_hit_index * 2 > block_index + end_index
442            {
443                end_index = min_hit_index;
444            }
445        }
446        stats.cache_data_prefetch_count += 1;
447        stats.cache_data_prefetch_block_count += (end_index - block_index) as u64;
448        let end_offset = start_offset
449            + sst.meta.block_metas[block_index..end_index]
450                .iter()
451                .map(|meta| meta.len as usize)
452                .sum::<usize>();
453        let data_path = self.get_sst_data_path(object_id);
454        let memory_usage = end_offset - start_offset;
455        let tracker = MemoryUsageTracker::new(self.prefetch_buffer_usage.clone(), memory_usage);
456        let span = await_tree::span!("Prefetch SST-{}", object_id).verbose();
457        let store = self.store.clone();
458        let join_handle = tokio::spawn(async move {
459            store
460                .read(&data_path, start_offset..end_offset)
461                .instrument_await(span)
462                .await
463        });
464        let buf = match join_handle.await {
465            Ok(Ok(data)) => data,
466            Ok(Err(e)) => {
467                tracing::error!(
468                    "prefetch meet error when read {}..{} from sst-{} ({})",
469                    start_offset,
470                    end_offset,
471                    object_id,
472                    sst.meta.estimated_size,
473                );
474                return Err(e.into());
475            }
476            Err(_) => {
477                return Err(HummockError::other("cancel by other thread"));
478            }
479        };
480        let mut offset = 0;
481        let mut blocks = VecDeque::default();
482        for idx in block_index..end_index {
483            let end = offset + sst.meta.block_metas[idx].len as usize;
484            if end > buf.len() {
485                return Err(ObjectError::internal("read unexpected EOF").into());
486            }
487            // copy again to avoid holding a large data in memory.
488            let block = Block::decode_with_copy(
489                buf.slice(offset..end),
490                sst.meta.block_metas[idx].uncompressed_size as usize,
491                true,
492            )?;
493            let holder = if let CachePolicy::Fill(hint) = policy {
494                let hint = if idx == block_index { hint } else { Hint::Low };
495                let entry = self.block_cache.insert_with_properties(
496                    SstableBlockIndex {
497                        sst_id: object_id,
498                        block_idx: idx as _,
499                    },
500                    Box::new(block),
501                    HybridCacheProperties::default().with_hint(hint),
502                );
503                BlockHolder::from_hybrid_cache_entry(entry)
504            } else {
505                BlockHolder::from_owned_block(Box::new(block))
506            };
507
508            blocks.push_back(holder);
509            offset = end;
510        }
511        Ok(Box::new(PrefetchBlockStream::new(
512            blocks,
513            block_index,
514            Some(tracker),
515        )))
516    }
517
518    pub async fn get_block_response(
519        &self,
520        sst: &Sstable,
521        block_index: usize,
522        policy: CachePolicy,
523    ) -> HummockResult<BlockResponse> {
524        let object_id = sst.id;
525        let (range, uncompressed_capacity) = sst.calculate_block_info(block_index);
526        let store = self.store.clone();
527
528        let file_size = sst.meta.estimated_size;
529        let data_path = Arc::new(self.get_sst_data_path(object_id));
530
531        let disable_cache: fn() -> bool = || {
532            fail_point!("disable_block_cache", |_| true);
533            false
534        };
535
536        let policy = if disable_cache() {
537            CachePolicy::Disable
538        } else {
539            policy
540        };
541
542        let idx = SstableBlockIndex {
543            sst_id: object_id,
544            block_idx: block_index as _,
545        };
546
547        self.recent_filter
548            .extend([(object_id, usize::MAX), (object_id, block_index)]);
549
550        // future: fetch block if hybrid cache miss
551        let fetch_block = async move {
552            let block_data = match store
553                .read(&data_path, range.clone())
554                .instrument_await("get_block_response".verbose())
555                .await
556            {
557                Ok(data) => data,
558                Err(e) => {
559                    tracing::error!(
560                        "get_block_response meet error when read {:?} from sst-{}, total length: {}",
561                        range,
562                        object_id,
563                        file_size
564                    );
565                    return Err(HummockError::from(e));
566                }
567            };
568            let block = Box::new(Block::decode(block_data, uncompressed_capacity)?);
569            Ok(block)
570        };
571
572        match policy {
573            CachePolicy::Fill(hint) => {
574                let properties = HybridCacheProperties::default().with_hint(hint);
575                let fetch = self.block_cache.get_or_fetch(&idx, || {
576                    fetch_block.map(|res| res.map(|block| (block, properties)))
577                });
578                Ok(BlockResponse::Fetch(fetch))
579            }
580            CachePolicy::NotFill => {
581                match self
582                    .block_cache
583                    .get(&idx)
584                    .await
585                    .map_err(HummockError::foyer_error)?
586                {
587                    Some(entry) => Ok(BlockResponse::Block(BlockHolder::from_hybrid_cache_entry(
588                        entry,
589                    ))),
590                    _ => {
591                        let block = fetch_block.await?;
592                        Ok(BlockResponse::Block(BlockHolder::from_owned_block(block)))
593                    }
594                }
595            }
596            CachePolicy::Disable => {
597                let block = fetch_block.await?;
598                Ok(BlockResponse::Block(BlockHolder::from_owned_block(block)))
599            }
600        }
601    }
602
603    pub async fn get(
604        &self,
605        sst: &Sstable,
606        block_index: usize,
607        policy: CachePolicy,
608        stats: &mut StoreLocalStatistic,
609    ) -> HummockResult<BlockHolder> {
610        let block_response = self.get_block_response(sst, block_index, policy).await?;
611        let block_holder = block_response.wait().await?;
612        stats.cache_data_block_total += 1;
613        if let BlockEntry::HybridCache(entry) = block_holder.entry()
614            && entry.source() == foyer::Source::Outer
615        {
616            stats.cache_data_block_miss += 1;
617        }
618        Ok(block_holder)
619    }
620
621    pub async fn get_vector_file_meta(
622        &self,
623        vector_file: &VectorFileInfo,
624        stats: &mut VectorStoreCacheStats,
625    ) -> HummockResult<VectorFileHolder> {
626        let store = self.store.clone();
627        let path = self.get_object_data_path(HummockObjectId::VectorFile(vector_file.object_id));
628        let meta_offset = vector_file.meta_offset;
629        let entry = self
630            .vector_meta_cache
631            .get_or_fetch(&vector_file.object_id.as_raw(), || async move {
632                let encoded_footer = store.read(&path, meta_offset..).await?;
633                let meta = VectorFileMeta::decode_footer(&encoded_footer)?;
634                Ok::<_, anyhow::Error>(HummockVectorIndexMetaFile::from(meta))
635            })
636            .await?;
637        stats.file_meta_total += 1;
638        if entry.source() == foyer::Source::Outer {
639            stats.file_meta_miss += 1;
640        }
641        VectorFileHolder::try_from_entry(entry, vector_file.object_id.as_raw())
642    }
643
644    pub async fn get_vector_block(
645        &self,
646        vector_file: &VectorFileInfo,
647        block_idx: usize,
648        block_meta: &VectorBlockMeta,
649        stats: &mut VectorStoreCacheStats,
650    ) -> HummockResult<VectorBlockHolder> {
651        let store = self.store.clone();
652        let path = self.get_object_data_path(HummockObjectId::VectorFile(vector_file.object_id));
653        let start_offset = block_meta.offset;
654        let end_offset = start_offset + block_meta.block_size;
655        let entry = self
656            .vector_block_cache
657            .get_or_fetch(&(vector_file.object_id, block_idx), || async move {
658                let encoded_block = store.read(&path, start_offset..end_offset).await?;
659                let block = VectorBlock::decode(&encoded_block)?;
660                Ok::<_, anyhow::Error>(Box::new(block))
661            })
662            .await
663            .map_err(HummockError::foyer_error)?;
664
665        stats.file_block_total += 1;
666        if entry.source() == foyer::Source::Outer {
667            stats.file_block_miss += 1;
668        }
669        Ok(entry)
670    }
671
672    pub fn insert_vector_cache(
673        &self,
674        object_id: HummockVectorFileId,
675        meta: VectorFileMeta,
676        blocks: Vec<VectorBlock>,
677    ) {
678        self.vector_meta_cache
679            .insert(object_id.as_raw(), meta.into());
680        for (idx, block) in blocks.into_iter().enumerate() {
681            self.vector_block_cache
682                .insert((object_id, idx), Box::new(block));
683        }
684    }
685
686    pub fn insert_hnsw_graph_cache(&self, object_id: HummockHnswGraphFileId, graph: PbHnswGraph) {
687        self.vector_meta_cache
688            .insert(object_id.as_raw(), graph.into());
689    }
690
691    pub async fn get_hnsw_graph(
692        &self,
693        graph_file: &HnswGraphFileInfo,
694        stats: &mut VectorStoreCacheStats,
695    ) -> HummockResult<HnswGraphFileHolder> {
696        let store = self.store.clone();
697        let graph_file_path =
698            self.get_object_data_path(HummockObjectId::HnswGraphFile(graph_file.object_id));
699        let entry = self
700            .vector_meta_cache
701            .get_or_fetch(&graph_file.object_id.as_raw(), || async move {
702                let encoded_graph = store.read(&graph_file_path, ..).await?;
703                let graph = PbHnswGraph::decode(encoded_graph.as_ref())?;
704                Ok::<_, anyhow::Error>(HummockVectorIndexMetaFile::from(graph))
705            })
706            .await
707            .map_err(HummockError::foyer_error)?;
708        stats.hnsw_graph_total += 1;
709        if entry.source() == foyer::Source::Outer {
710            stats.hnsw_graph_miss += 1;
711        }
712        HnswGraphFileHolder::try_from_entry(entry, graph_file.object_id.as_raw())
713    }
714
715    pub fn get_sst_data_path(&self, object_id: impl Into<HummockSstableObjectId>) -> String {
716        self.get_object_data_path(HummockObjectId::Sstable(object_id.into()))
717    }
718
719    pub fn get_object_data_path(&self, object_id: HummockObjectId) -> String {
720        let obj_prefix = self.store.get_object_prefix(
721            object_id.as_raw().as_raw_id(),
722            self.use_new_object_prefix_strategy,
723        );
724        risingwave_hummock_sdk::get_object_data_path(&obj_prefix, &self.path, object_id)
725    }
726
727    pub fn get_object_id_from_path(path: &str) -> HummockObjectId {
728        risingwave_hummock_sdk::get_object_id_from_path(path)
729    }
730
731    pub fn store(&self) -> ObjectStoreRef {
732        self.store.clone()
733    }
734
735    #[cfg(any(test, feature = "test"))]
736    pub async fn clear_block_cache(&self) -> HummockResult<()> {
737        self.block_cache
738            .clear()
739            .await
740            .map_err(HummockError::foyer_error)
741    }
742
743    #[cfg(any(test, feature = "test"))]
744    pub async fn clear_meta_cache(&self) -> HummockResult<()> {
745        self.meta_cache
746            .clear()
747            .await
748            .map_err(HummockError::foyer_error)
749    }
750
751    pub async fn sstable_cached(
752        &self,
753        sst_obj_id: HummockSstableObjectId,
754    ) -> HummockResult<Option<TableHolder>> {
755        self.meta_cache
756            .get(&sst_obj_id)
757            .await
758            .map_err(HummockError::foyer_error)
759    }
760
761    /// Returns `table_holder`
762    pub async fn sstable(
763        &self,
764        sstable_info_ref: &SstableInfo,
765        stats: &mut StoreLocalStatistic,
766    ) -> HummockResult<TableHolder> {
767        let object_id = sstable_info_ref.object_id;
768        let store = self.store.clone();
769        let meta_path = self.get_sst_data_path(object_id);
770        let stats_ptr = stats.remote_io_time.clone();
771        let range = sstable_info_ref.meta_offset as usize..;
772        let skip_bloom_filter_in_serde = self.skip_bloom_filter_in_serde;
773
774        let fetch = self.meta_cache.get_or_fetch(&object_id, || async move {
775            let now = Instant::now();
776            let buf = store
777                .read(&meta_path, range)
778                .instrument_await("get_meta_response".verbose())
779                .await?;
780            let sst = Sstable::from_meta_bytes(object_id, &buf, skip_bloom_filter_in_serde)?;
781            let add = (now.elapsed().as_secs_f64() * 1000.0).ceil();
782            stats_ptr.fetch_add(add as u64, Ordering::Relaxed);
783            Ok::<_, anyhow::Error>(SstableMetaCacheValue::from(sst))
784        });
785
786        stats.cache_meta_block_total += 1;
787        let entry = fetch
788            .instrument_await("fetch_meta".verbose())
789            .await
790            .map_err(HummockError::foyer_error);
791        if let Ok(ref entry) = entry
792            && entry.source() == foyer::Source::Outer
793        {
794            stats.cache_meta_block_miss += 1;
795        }
796        entry
797    }
798
799    pub async fn list_sst_object_metadata_from_object_store(
800        &self,
801        prefix: Option<String>,
802        start_after: Option<String>,
803        limit: Option<usize>,
804    ) -> HummockResult<ObjectMetadataIter> {
805        let list_path = format!("{}/{}", self.path, prefix.unwrap_or("".into()));
806        let raw_iter = self.store.list(&list_path, start_after, limit).await?;
807        let iter = raw_iter.filter(|r| match r {
808            Ok(i) => future::ready(i.key.ends_with(&format!(".{}", SST_OBJECT_SUFFIX))),
809            Err(_) => future::ready(true),
810        });
811        Ok(Box::pin(iter))
812    }
813
814    pub fn create_sst_writer(
815        self: Arc<Self>,
816        object_id: impl Into<HummockSstableObjectId>,
817        options: SstableWriterOptions,
818    ) -> BatchUploadWriter {
819        BatchUploadWriter::new(object_id, self, options)
820    }
821
822    pub fn insert_meta_cache(&self, object_id: HummockSstableObjectId, meta: SstableMeta) {
823        let sst = Sstable::new(object_id, meta, self.skip_bloom_filter_in_serde);
824        self.meta_cache
825            .insert(object_id, SstableMetaCacheValue::from(sst));
826    }
827
828    pub fn get_prefetch_memory_usage(&self) -> usize {
829        self.prefetch_buffer_usage.load(Ordering::Acquire)
830    }
831
832    pub async fn get_stream_for_blocks(
833        &self,
834        object_id: HummockSstableObjectId,
835        metas: &[BlockMeta],
836    ) -> HummockResult<BlockDataStream> {
837        fail_point!("get_stream_err");
838        let data_path = self.get_sst_data_path(object_id);
839        let store = self.store();
840        let block_meta = &metas[0];
841        let start_pos = block_meta.offset as usize;
842        let end_pos = metas.iter().map(|meta| meta.len as usize).sum::<usize>() + start_pos;
843        let range = start_pos..end_pos;
844        // spawn to tokio pool because the object-storage sdk may not be safe to cancel.
845        let ret = tokio::spawn(async move { store.streaming_read(&data_path, range).await }).await;
846
847        let reader = match ret {
848            Ok(Ok(reader)) => reader,
849            Ok(Err(e)) => return Err(HummockError::from(e)),
850            Err(e) => {
851                return Err(HummockError::other(format!(
852                    "failed to get result, this read request may be canceled: {}",
853                    e.as_report()
854                )));
855            }
856        };
857        Ok(BlockDataStream::new(reader, metas))
858    }
859
860    pub fn meta_cache(&self) -> &HybridCache<HummockSstableObjectId, SstableMetaCacheValue> {
861        &self.meta_cache
862    }
863
864    pub fn block_cache(&self) -> &HybridCache<SstableBlockIndex, Box<Block>> {
865        &self.block_cache
866    }
867
868    pub fn recent_filter(&self) -> &Arc<RecentFilter<(HummockSstableObjectId, usize)>> {
869        &self.recent_filter
870    }
871
872    pub async fn create_streaming_uploader(
873        &self,
874        path: &str,
875    ) -> ObjectResult<ObjectStreamingUploader> {
876        self.store.streaming_upload(path).await
877    }
878}
879
880pub type SstableStoreRef = Arc<SstableStore>;
881#[cfg(test)]
882mod tests {
883    use std::ops::Range;
884    use std::sync::Arc;
885
886    use risingwave_hummock_sdk::HummockObjectId;
887    use risingwave_hummock_sdk::sstable_info::SstableInfo;
888
889    use super::{SstableStoreRef, SstableWriterOptions};
890    use crate::hummock::iterator::HummockIterator;
891    use crate::hummock::iterator::test_utils::{iterator_test_key_of, mock_sstable_store};
892    use crate::hummock::sstable::SstableIteratorReadOptions;
893    use crate::hummock::test_utils::{
894        default_builder_opt_for_test, gen_default_test_sstable, gen_test_sstable_data, put_sst,
895        test_key_of,
896    };
897    use crate::hummock::value::HummockValue;
898    use crate::hummock::{CachePolicy, SstableIterator, SstableMeta, SstableStore};
899    use crate::monitor::StoreLocalStatistic;
900
901    const SST_ID: u64 = 1;
902
903    fn get_hummock_value(x: usize) -> HummockValue<Vec<u8>> {
904        HummockValue::put(format!("overlapped_new_{}", x).as_bytes().to_vec())
905    }
906
907    async fn validate_sst(
908        sstable_store: SstableStoreRef,
909        info: &SstableInfo,
910        mut meta: SstableMeta,
911        x_range: Range<usize>,
912    ) {
913        let mut stats = StoreLocalStatistic::default();
914        let holder = sstable_store.sstable(info, &mut stats).await.unwrap();
915        std::mem::take(&mut meta.bloom_filter);
916        assert_eq!(holder.meta, meta);
917        let holder = sstable_store.sstable(info, &mut stats).await.unwrap();
918        assert_eq!(holder.meta, meta);
919        let mut iter = SstableIterator::new(
920            holder,
921            sstable_store,
922            Arc::new(SstableIteratorReadOptions::default()),
923            info,
924        );
925        iter.rewind().await.unwrap();
926        for i in x_range {
927            let key = iter.key();
928            let value = iter.value();
929            assert_eq!(key, iterator_test_key_of(i).to_ref());
930            assert_eq!(value, get_hummock_value(i).as_slice());
931            iter.next().await.unwrap();
932        }
933    }
934
935    #[tokio::test]
936    async fn test_batch_upload() {
937        let sstable_store = mock_sstable_store().await;
938        let x_range = 0..100;
939        let (data, meta) = gen_test_sstable_data(
940            default_builder_opt_for_test(),
941            x_range
942                .clone()
943                .map(|x| (iterator_test_key_of(x), get_hummock_value(x))),
944        )
945        .await;
946        let writer_opts = SstableWriterOptions {
947            capacity_hint: None,
948            tracker: None,
949            policy: CachePolicy::Disable,
950        };
951        let info = put_sst(
952            SST_ID,
953            data.clone(),
954            meta.clone(),
955            sstable_store.clone(),
956            writer_opts,
957            vec![0],
958        )
959        .await
960        .unwrap();
961
962        validate_sst(sstable_store, &info, meta, x_range).await;
963    }
964
965    #[tokio::test]
966    async fn test_streaming_upload() {
967        // Generate test data.
968        let sstable_store = mock_sstable_store().await;
969        let x_range = 0..100;
970        let (data, meta) = gen_test_sstable_data(
971            default_builder_opt_for_test(),
972            x_range
973                .clone()
974                .map(|x| (iterator_test_key_of(x), get_hummock_value(x))),
975        )
976        .await;
977        let writer_opts = SstableWriterOptions {
978            capacity_hint: None,
979            tracker: None,
980            policy: CachePolicy::Disable,
981        };
982        let info = put_sst(
983            SST_ID,
984            data.clone(),
985            meta.clone(),
986            sstable_store.clone(),
987            writer_opts,
988            vec![0],
989        )
990        .await
991        .unwrap();
992
993        validate_sst(sstable_store, &info, meta, x_range).await;
994    }
995
996    #[tokio::test]
997    async fn test_empty_prefetch_falls_back_to_block_get() {
998        let mut sstable_store = mock_sstable_store().await;
999        // Direct SstableStoreConfig construction can bypass the serde nonzero check.
1000        Arc::get_mut(&mut sstable_store)
1001            .unwrap()
1002            .max_prefetch_block_number = 0;
1003        let (sstable, info) =
1004            gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
1005                .await;
1006        sstable_store.clear_block_cache().await.unwrap();
1007        let mut iter = SstableIterator::new(
1008            sstable,
1009            sstable_store.clone(),
1010            Arc::new(SstableIteratorReadOptions {
1011                prefetch: true,
1012                cache_policy: CachePolicy::Disable,
1013                ..Default::default()
1014            }),
1015            &info,
1016        );
1017        tokio::time::timeout(std::time::Duration::from_secs(5), iter.rewind())
1018            .await
1019            .expect("empty prefetch must not cause an unbounded refill loop")
1020            .unwrap();
1021        assert_eq!(iter.key(), test_key_of(0).to_ref());
1022        let mut stats = StoreLocalStatistic::default();
1023        iter.collect_local_statistic(&mut stats);
1024        assert_eq!(stats.cache_data_prefetch_count, 1);
1025        assert_eq!(stats.cache_data_block_total, 1);
1026        assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
1027    }
1028
1029    #[tokio::test]
1030    async fn test_prefetch_failure_falls_back_and_releases_tracker() {
1031        for truncate in [true, false] {
1032            let sstable_store = mock_sstable_store().await;
1033            let (sstable, info) =
1034                gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
1035                    .await;
1036            assert!(sstable.meta.block_metas.len() > 1);
1037            sstable_store.clear_block_cache().await.unwrap();
1038            let path = sstable_store.get_sst_data_path(sstable.id);
1039            let mut data = sstable_store.store.read(&path, ..).await.unwrap().to_vec();
1040            let second_offset = sstable.meta.block_metas[1].offset as usize;
1041            if truncate {
1042                // A multi-block read fails, but reading block 0 still succeeds.
1043                data.truncate(second_offset);
1044            } else {
1045                // Decoding block 1 fails its checksum after block 0 was already decoded.
1046                data[second_offset] ^= 1;
1047            }
1048            sstable_store
1049                .store
1050                .upload(&path, data.into())
1051                .await
1052                .unwrap();
1053            let mut iter = SstableIterator::new(
1054                sstable.clone(),
1055                sstable_store.clone(),
1056                Arc::new(SstableIteratorReadOptions {
1057                    prefetch: true,
1058                    cache_policy: CachePolicy::Disable,
1059                    ..Default::default()
1060                }),
1061                &info,
1062            );
1063            iter.rewind().await.unwrap();
1064            assert_eq!(iter.key(), test_key_of(0).to_ref());
1065            assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
1066            let mut stats = StoreLocalStatistic::default();
1067            iter.collect_local_statistic(&mut stats);
1068            assert_eq!(stats.cache_data_prefetch_count, 1);
1069            assert_eq!(stats.cache_data_block_total, 1);
1070            // If the single-block read also fails, the error must still reach the caller.
1071            let second_key = risingwave_hummock_sdk::key::FullKey::decode(
1072                &sstable.meta.block_metas[1].smallest_key,
1073            );
1074            assert!(iter.seek(second_key).await.is_err());
1075            assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
1076        }
1077    }
1078
1079    #[tokio::test]
1080    async fn test_prefetch_releases_old_budget_before_refill() {
1081        let mut sstable_store = mock_sstable_store().await;
1082        // The first prefetch may start at zero usage. A retained old tracker would force
1083        // the next prefetch into the single-block path instead of reading another batch.
1084        Arc::get_mut(&mut sstable_store)
1085            .unwrap()
1086            .prefetch_buffer_capacity = 0;
1087        let (sstable, info) =
1088            gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
1089                .await;
1090        sstable_store.clear_block_cache().await.unwrap();
1091        let next_batch = sstable_store.max_prefetch_block_number;
1092        assert!(sstable.meta.block_metas.len() > next_batch + 1);
1093        let mut iter = SstableIterator::new(
1094            sstable.clone(),
1095            sstable_store.clone(),
1096            Arc::new(SstableIteratorReadOptions {
1097                prefetch: true,
1098                cache_policy: CachePolicy::Disable,
1099                ..Default::default()
1100            }),
1101            &info,
1102        );
1103        for idx in [0, next_batch] {
1104            let key = risingwave_hummock_sdk::key::FullKey::decode(
1105                &sstable.meta.block_metas[idx].smallest_key,
1106            );
1107            iter.seek(key).await.unwrap();
1108            assert_eq!(iter.key(), key);
1109            assert!(sstable_store.get_prefetch_memory_usage() > 0);
1110        }
1111        let mut stats = StoreLocalStatistic::default();
1112        iter.collect_local_statistic(&mut stats);
1113        assert_eq!(stats.cache_data_prefetch_count, 2);
1114        assert_eq!(stats.cache_data_block_total, 0);
1115        drop(iter);
1116        assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
1117    }
1118
1119    #[tokio::test]
1120    async fn test_prefetch_producer_large_limit() {
1121        let mut sstable_store = mock_sstable_store().await;
1122        Arc::get_mut(&mut sstable_store)
1123            .unwrap()
1124            .max_prefetch_block_number = usize::MAX;
1125        let (sstable, _) =
1126            gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
1127                .await;
1128        sstable_store.clear_block_cache().await.unwrap();
1129        let mut stats = StoreLocalStatistic::default();
1130        let mut stream = sstable_store
1131            .prefetch_blocks(&sstable, 1, 3, CachePolicy::Disable, &mut stats)
1132            .await
1133            .unwrap();
1134        assert_eq!(stats.cache_data_prefetch_block_count, 2);
1135        for idx in [1, 2] {
1136            assert!(matches!(
1137                stream.take_block(idx),
1138                crate::hummock::block_stream::PrefetchLookup::Hit(_)
1139            ));
1140        }
1141        assert!(matches!(
1142            stream.take_block(3),
1143            crate::hummock::block_stream::PrefetchLookup::Exhausted
1144        ));
1145        drop(stream);
1146        assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
1147    }
1148
1149    #[tokio::test]
1150    async fn test_prefetch_producer_respects_disabled_cache() {
1151        let sstable_store = mock_sstable_store().await;
1152        let (sstable, _) =
1153            gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
1154                .await;
1155        sstable_store.clear_block_cache().await.unwrap();
1156        // A cached block in the requested range must not shorten a cache-disabled prefetch.
1157        sstable_store
1158            .get(
1159                &sstable,
1160                1,
1161                CachePolicy::Fill(foyer::Hint::Normal),
1162                &mut StoreLocalStatistic::default(),
1163            )
1164            .await
1165            .unwrap();
1166        let mut stats = StoreLocalStatistic::default();
1167        let stream = sstable_store
1168            .prefetch_blocks(&sstable, 0, 3, CachePolicy::Disable, &mut stats)
1169            .await
1170            .unwrap();
1171        assert_eq!(stats.cache_data_prefetch_block_count, 3);
1172        drop(stream);
1173        assert!(
1174            !sstable_store
1175                .block_cache
1176                .contains(&super::SstableBlockIndex {
1177                    sst_id: sstable.id,
1178                    block_idx: 0,
1179                })
1180        );
1181
1182        // With block 0 cached and the object removed, NotFill may succeed but Disable must fail.
1183        sstable_store
1184            .get(
1185                &sstable,
1186                0,
1187                CachePolicy::Fill(foyer::Hint::Normal),
1188                &mut StoreLocalStatistic::default(),
1189            )
1190            .await
1191            .unwrap();
1192        sstable_store.delete(sstable.id).await.unwrap();
1193        assert!(
1194            sstable_store
1195                .prefetch_blocks(
1196                    &sstable,
1197                    0,
1198                    3,
1199                    CachePolicy::NotFill,
1200                    &mut StoreLocalStatistic::default()
1201                )
1202                .await
1203                .is_ok()
1204        );
1205        assert!(
1206            sstable_store
1207                .prefetch_blocks(
1208                    &sstable,
1209                    0,
1210                    3,
1211                    CachePolicy::Disable,
1212                    &mut StoreLocalStatistic::default()
1213                )
1214                .await
1215                .is_err()
1216        );
1217        assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
1218    }
1219
1220    #[tokio::test]
1221    async fn test_basic() {
1222        let sstable_store = mock_sstable_store().await;
1223        let object_id = 123;
1224        let data_path = sstable_store.get_sst_data_path(object_id);
1225        assert_eq!(data_path, "test/123.data");
1226        assert_eq!(
1227            SstableStore::get_object_id_from_path(&data_path),
1228            HummockObjectId::Sstable(object_id.into())
1229        );
1230    }
1231
1232    #[tokio::test]
1233    async fn test_clear_file_cache_and_refetch() {
1234        let sstable_store = mock_sstable_store().await;
1235        let x_range = 0..10;
1236        let (data, meta) = gen_test_sstable_data(
1237            default_builder_opt_for_test(),
1238            x_range.map(|x| (iterator_test_key_of(x), get_hummock_value(x))),
1239        )
1240        .await;
1241        let info = put_sst(
1242            SST_ID,
1243            data,
1244            meta,
1245            sstable_store.clone(),
1246            SstableWriterOptions {
1247                capacity_hint: None,
1248                tracker: None,
1249                policy: CachePolicy::Fill(foyer::Hint::Normal),
1250            },
1251            vec![0],
1252        )
1253        .await
1254        .unwrap();
1255
1256        let mut stats = StoreLocalStatistic::default();
1257        assert!(
1258            sstable_store
1259                .sstable_cached(info.object_id)
1260                .await
1261                .unwrap()
1262                .is_some()
1263        );
1264        let sstable = sstable_store.sstable(&info, &mut stats).await.unwrap();
1265        sstable_store
1266            .get(
1267                &sstable,
1268                0,
1269                CachePolicy::Fill(foyer::Hint::Normal),
1270                &mut stats,
1271            )
1272            .await
1273            .unwrap();
1274        assert!(
1275            sstable_store
1276                .block_cache()
1277                .get(&super::SstableBlockIndex {
1278                    sst_id: info.object_id,
1279                    block_idx: 0,
1280                })
1281                .await
1282                .unwrap()
1283                .is_some()
1284        );
1285
1286        sstable_store.clear_meta_cache().await.unwrap();
1287        assert!(
1288            sstable_store
1289                .sstable_cached(info.object_id)
1290                .await
1291                .unwrap()
1292                .is_none()
1293        );
1294        sstable_store.sstable(&info, &mut stats).await.unwrap();
1295        assert!(
1296            sstable_store
1297                .sstable_cached(info.object_id)
1298                .await
1299                .unwrap()
1300                .is_some()
1301        );
1302
1303        sstable_store.clear_block_cache().await.unwrap();
1304        assert!(
1305            sstable_store
1306                .block_cache()
1307                .get(&super::SstableBlockIndex {
1308                    sst_id: info.object_id,
1309                    block_idx: 0,
1310                })
1311                .await
1312                .unwrap()
1313                .is_none()
1314        );
1315        sstable_store
1316            .get(
1317                &sstable,
1318                0,
1319                CachePolicy::Fill(foyer::Hint::Normal),
1320                &mut stats,
1321            )
1322            .await
1323            .unwrap();
1324        assert!(
1325            sstable_store
1326                .block_cache()
1327                .get(&super::SstableBlockIndex {
1328                    sst_id: info.object_id,
1329                    block_idx: 0,
1330                })
1331                .await
1332                .unwrap()
1333                .is_some()
1334        );
1335    }
1336}
risingwave_storage/hummock/sstable_store.rs

risingwave_storage/hummock/
sstable_store.rs
