Skip to main content

mito2/cache/index/
bloom_filter_index.rs

1// Copyright 2023 Greptime Team
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::ops::Range;
16use std::sync::Arc;
17use std::time::Instant;
18
19use api::v1::index::{BloomFilterLoc, BloomFilterMeta};
20use async_trait::async_trait;
21use bytes::Bytes;
22use index::bloom_filter::error::{Error, IoSnafu, Result, SharedReadSnafu};
23use index::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader};
24use index::target::IndexTarget;
25use snafu::IntoError;
26use store_api::storage::{FileId, IndexVersion};
27
28use crate::cache::index::{INDEX_METADATA_TYPE, IndexCache};
29use crate::metrics::{CACHE_HIT, CACHE_MISS};
30
31const INDEX_TYPE_BLOOM_FILTER_INDEX: &str = "bloom_filter_index";
32
33/// Tag for bloom filter index cache.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
35pub enum Tag {
36    Skipping,
37    Fulltext,
38}
39
40pub type BloomFilterIndexKey = (FileId, IndexVersion, Arc<IndexTarget>, Tag);
41
42/// Cache for bloom filter index.
43pub type BloomFilterIndexCache = IndexCache<BloomFilterIndexKey, BloomFilterMeta>;
44pub type BloomFilterIndexCacheRef = Arc<BloomFilterIndexCache>;
45
46impl BloomFilterIndexCache {
47    /// Creates a new bloom filter index cache.
48    pub fn new(index_metadata_cap: u64, index_content_cap: u64, page_size: u64) -> Self {
49        Self::new_with_weighter(
50            index_metadata_cap,
51            index_content_cap,
52            page_size,
53            INDEX_TYPE_BLOOM_FILTER_INDEX,
54            "bloom_filter_index_metadata",
55            bloom_filter_index_metadata_weight,
56            bloom_filter_index_content_weight,
57        )
58    }
59
60    /// Removes all cached entries for the given `file_id`.
61    pub fn invalidate_file(&self, file_id: FileId) {
62        self.invalidate_if(move |key| key.0 == file_id);
63    }
64}
65
66/// Calculates weight for bloom filter index metadata.
67fn bloom_filter_index_metadata_weight(k: &BloomFilterIndexKey, meta: &Arc<BloomFilterMeta>) -> u32 {
68    let base = k.0.as_bytes().len()
69        + std::mem::size_of::<IndexVersion>()
70        + k.2.approx_mem_size()
71        + std::mem::size_of::<Tag>()
72        + std::mem::size_of::<BloomFilterMeta>();
73
74    let vec_estimated = meta.segment_loc_indices.len() * std::mem::size_of::<u64>()
75        + meta.bloom_filter_locs.len() * std::mem::size_of::<BloomFilterLoc>();
76
77    (base + vec_estimated) as u32
78}
79
80/// Calculates weight for bloom filter index content.
81fn bloom_filter_index_content_weight((k, _): &(BloomFilterIndexKey, u64), v: &Bytes) -> u32 {
82    (k.0.as_bytes().len() + k.2.approx_mem_size() + v.len()) as u32
83}
84
85/// Bloom filter index blob reader with cache.
86pub struct CachedBloomFilterIndexBlobReader<R> {
87    file_id: FileId,
88    index_version: IndexVersion,
89    target: Arc<IndexTarget>,
90    tag: Tag,
91    blob_size: u64,
92    inner: R,
93    cache: BloomFilterIndexCacheRef,
94}
95
96impl<R> CachedBloomFilterIndexBlobReader<R> {
97    /// Creates a new bloom filter index blob reader with cache.
98    pub fn new(
99        file_id: FileId,
100        index_version: IndexVersion,
101        target: IndexTarget,
102        tag: Tag,
103        blob_size: u64,
104        inner: R,
105        cache: BloomFilterIndexCacheRef,
106    ) -> Self {
107        Self {
108            file_id,
109            index_version,
110            target: Arc::new(target),
111            tag,
112            blob_size,
113            inner,
114            cache,
115        }
116    }
117}
118
119#[async_trait]
120impl<R: BloomFilterReader + Send> BloomFilterReader for CachedBloomFilterIndexBlobReader<R> {
121    async fn range_read(
122        &self,
123        offset: u64,
124        size: u32,
125        metrics: Option<&mut BloomFilterReadMetrics>,
126    ) -> Result<Bytes> {
127        let start = metrics.as_ref().map(|_| Instant::now());
128        let inner = &self.inner;
129        let (result, cache_metrics) = self
130            .cache
131            .get_or_load(
132                (
133                    self.file_id,
134                    self.index_version,
135                    self.target.clone(),
136                    self.tag,
137                ),
138                self.blob_size,
139                offset,
140                size,
141                move |ranges| async move { inner.read_vec(&ranges, None).await },
142            )
143            .await
144            .map_err(map_page_error)?;
145
146        if let Some(m) = metrics {
147            m.total_ranges += cache_metrics.num_pages;
148            m.total_bytes += cache_metrics.page_bytes;
149            m.cache_hit += cache_metrics.cache_hit;
150            m.cache_miss += cache_metrics.cache_miss;
151            if let Some(start) = start {
152                m.fetch_elapsed += start.elapsed();
153            }
154        }
155
156        Ok(result)
157    }
158
159    async fn read_vec(
160        &self,
161        ranges: &[Range<u64>],
162        metrics: Option<&mut BloomFilterReadMetrics>,
163    ) -> Result<Vec<Bytes>> {
164        let start = metrics.as_ref().map(|_| Instant::now());
165
166        let inner = &self.inner;
167        let (pages, total_cache_metrics) = self
168            .cache
169            .get_or_load_vec(
170                (
171                    self.file_id,
172                    self.index_version,
173                    self.target.clone(),
174                    self.tag,
175                ),
176                self.blob_size,
177                ranges,
178                move |ranges| async move { inner.read_vec(&ranges, None).await },
179            )
180            .await
181            .map_err(map_page_error)?;
182        if let Some(m) = metrics {
183            m.total_ranges += total_cache_metrics.num_pages;
184            m.total_bytes += total_cache_metrics.page_bytes;
185            m.cache_hit += total_cache_metrics.cache_hit;
186            m.cache_miss += total_cache_metrics.cache_miss;
187            if let Some(start) = start {
188                m.fetch_elapsed += start.elapsed();
189            }
190        }
191
192        Ok(pages)
193    }
194
195    /// Reads the meta information of the bloom filter.
196    async fn metadata(
197        &self,
198        metrics: Option<&mut BloomFilterReadMetrics>,
199    ) -> Result<Arc<BloomFilterMeta>> {
200        let key = (
201            self.file_id,
202            self.index_version,
203            self.target.clone(),
204            self.tag,
205        );
206        if let Some(cached) = self.cache.get_metadata(key.clone()).await {
207            CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc();
208            if let Some(m) = metrics {
209                m.cache_hit += 1;
210            }
211            return Ok(cached);
212        }
213
214        let start = metrics.as_ref().map(|_| Instant::now());
215        let collect_metrics = metrics.is_some();
216        let mut loaded = BloomFilterReadMetrics::default();
217        let result = self
218            .cache
219            .get_or_load_metadata(key, async {
220                self.inner
221                    .metadata(collect_metrics.then_some(&mut loaded))
222                    .await
223            })
224            .await;
225        if let Some(m) = metrics {
226            // Only the selected loader performs I/O. Waiters retain zero I/O
227            // counts. fetch_elapsed is per-caller elapsed (including wait and error),
228            // not physical I/O time summed across callers.
229            m.total_bytes += loaded.total_bytes;
230            m.total_ranges += loaded.total_ranges;
231            m.cache_hit += loaded.cache_hit;
232            m.cache_miss += loaded.cache_miss;
233            if let Some(start) = start {
234                m.fetch_elapsed += start.elapsed();
235            }
236        }
237        let meta = result.map_err(|source| SharedReadSnafu.into_error(source))?;
238        CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc();
239        Ok(meta)
240    }
241}
242
243/// Recovers the typed source retained by the range-reader adapter.
244fn map_page_error(error: Arc<std::io::Error>) -> Error {
245    if let Some(source) = crate::cache::find_shared_page_source::<Error>(&error) {
246        SharedReadSnafu.into_error(source)
247    } else {
248        IoSnafu.into_error(std::io::Error::other(error))
249    }
250}
251
252#[cfg(test)]
253mod test {
254    use rand::{Rng, RngCore};
255    use store_api::storage::ColumnId;
256
257    use super::*;
258
259    const FUZZ_REPEAT_TIMES: usize = 100;
260
261    struct MetadataReader {
262        release: tokio::sync::Notify,
263        fail: std::sync::atomic::AtomicBool,
264    }
265
266    #[async_trait]
267    impl BloomFilterReader for MetadataReader {
268        async fn range_read(
269            &self,
270            _: u64,
271            _: u32,
272            _: Option<&mut BloomFilterReadMetrics>,
273        ) -> Result<Bytes> {
274            unreachable!("metadata test does not read content pages")
275        }
276
277        async fn read_vec(
278            &self,
279            _: &[Range<u64>],
280            _: Option<&mut BloomFilterReadMetrics>,
281        ) -> Result<Vec<Bytes>> {
282            unreachable!("metadata test does not read content pages")
283        }
284
285        async fn metadata(
286            &self,
287            metrics: Option<&mut BloomFilterReadMetrics>,
288        ) -> Result<Arc<BloomFilterMeta>> {
289            self.release.notified().await;
290            if let Some(metrics) = metrics {
291                metrics.total_bytes += 8192;
292                metrics.total_ranges += 1;
293                // The wrapper must replace, not add, this physical-fetch duration.
294                metrics.fetch_elapsed += std::time::Duration::from_secs(3600);
295            }
296            if self.fail.load(std::sync::atomic::Ordering::Relaxed) {
297                return Err(IoSnafu.into_error(std::io::Error::other("metadata failed")));
298            }
299            Ok(Arc::new(BloomFilterMeta::default()))
300        }
301    }
302
303    #[tokio::test]
304    async fn metadata_waiters_do_not_duplicate_io_metrics() {
305        for fail in [false, true] {
306            let cache = Arc::new(BloomFilterIndexCache::new(4096, 4096, 1024));
307            let reader = CachedBloomFilterIndexBlobReader::new(
308                FileId::random(),
309                0,
310                IndexTarget::ColumnId(42),
311                Tag::Skipping,
312                8192,
313                MetadataReader {
314                    release: tokio::sync::Notify::new(),
315                    fail: std::sync::atomic::AtomicBool::new(fail),
316                },
317                cache,
318            );
319            let mut owner_metrics = BloomFilterReadMetrics::default();
320            let mut waiter_metrics = BloomFilterReadMetrics::default();
321            let mut owner = Box::pin(reader.metadata(Some(&mut owner_metrics)));
322            assert!(futures::poll!(&mut owner).is_pending());
323            let mut waiter = Box::pin(reader.metadata(Some(&mut waiter_metrics)));
324            assert!(futures::poll!(&mut waiter).is_pending());
325            reader.inner.release.notify_one();
326            let owner_result = owner.await;
327            let waiter_result = waiter.await;
328            assert_eq!(owner_result.is_err(), fail);
329            assert_eq!(waiter_result.is_err(), fail);
330            assert_eq!(
331                (owner_metrics.total_bytes, owner_metrics.total_ranges),
332                (8192, 1)
333            );
334            assert_eq!(
335                (waiter_metrics.total_bytes, waiter_metrics.total_ranges),
336                (0, 0)
337            );
338            assert_eq!((owner_metrics.cache_hit, waiter_metrics.cache_hit), (0, 0));
339            for metrics in [&owner_metrics, &waiter_metrics] {
340                // Per-caller elapsed includes waiting, including on error; it is
341                // not a sum of physical I/O durations across concurrent callers.
342                assert!(metrics.fetch_elapsed > std::time::Duration::ZERO);
343                assert!(metrics.fetch_elapsed < std::time::Duration::from_secs(3600));
344            }
345            if fail {
346                let Error::SharedRead { source: owner, .. } = owner_result.unwrap_err() else {
347                    panic!("owner must preserve the typed shared read error");
348                };
349                let Error::SharedRead { source: waiter, .. } = waiter_result.unwrap_err() else {
350                    panic!("waiter must preserve the typed shared read error");
351                };
352                assert!(Arc::ptr_eq(&owner, &waiter));
353            } else {
354                let mut hot_metrics = BloomFilterReadMetrics::default();
355                reader.metadata(Some(&mut hot_metrics)).await.unwrap();
356                assert_eq!(hot_metrics.cache_hit, 1);
357                assert_eq!(hot_metrics.total_bytes, 0);
358            }
359        }
360    }
361
362    #[tokio::test]
363    async fn json_paths_have_independent_cache_entries() {
364        let cache = BloomFilterIndexCache::new(1024 * 1024, 1024 * 1024, 1024);
365        let file_id = FileId::random();
366        let first = (
367            file_id,
368            0,
369            Arc::new(IndexTarget::new_json_path(1, vec!["a".into(), "b".into()]).unwrap()),
370            Tag::Skipping,
371        );
372        let second = (
373            file_id,
374            0,
375            Arc::new(IndexTarget::new_json_path(1, vec!["a.b".into()]).unwrap()),
376            Tag::Skipping,
377        );
378        cache
379            .put_metadata(
380                first.clone(),
381                Arc::new(BloomFilterMeta {
382                    row_count: 1,
383                    ..Default::default()
384                }),
385            )
386            .await;
387        cache
388            .put_metadata(
389                second.clone(),
390                Arc::new(BloomFilterMeta {
391                    row_count: 2,
392                    ..Default::default()
393                }),
394            )
395            .await;
396        assert_eq!(
397            cache.get_metadata(first.clone()).await.unwrap().row_count,
398            1
399        );
400        assert_eq!(
401            cache.get_metadata(second.clone()).await.unwrap().row_count,
402            2
403        );
404        cache.invalidate_file(file_id);
405        assert!(cache.get_metadata(first).await.is_none());
406        assert!(cache.get_metadata(second).await.is_none());
407    }
408
409    #[test]
410    fn bloom_filter_metadata_weight_counts_vec_contents() {
411        let file_id = FileId::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
412        let version = 0;
413        let column_id: ColumnId = 42;
414        let tag = Tag::Skipping;
415
416        let meta = BloomFilterMeta {
417            rows_per_segment: 128,
418            segment_count: 2,
419            row_count: 256,
420            bloom_filter_size: 1024,
421            segment_loc_indices: vec![0, 64, 128, 192],
422            bloom_filter_locs: vec![
423                BloomFilterLoc {
424                    offset: 0,
425                    size: 512,
426                    element_count: 1000,
427                },
428                BloomFilterLoc {
429                    offset: 512,
430                    size: 512,
431                    element_count: 1000,
432                },
433            ],
434        };
435
436        let weight = bloom_filter_index_metadata_weight(
437            &(
438                file_id,
439                version,
440                Arc::new(IndexTarget::ColumnId(column_id)),
441                tag,
442            ),
443            &Arc::new(meta.clone()),
444        );
445
446        let base = file_id.as_bytes().len()
447            + std::mem::size_of::<IndexVersion>()
448            + std::mem::size_of::<IndexTarget>()
449            + std::mem::size_of::<Tag>()
450            + std::mem::size_of::<BloomFilterMeta>();
451        let expected_dynamic = meta.segment_loc_indices.len() * std::mem::size_of::<u64>()
452            + meta.bloom_filter_locs.len() * std::mem::size_of::<BloomFilterLoc>();
453
454        assert_eq!(weight as usize, base + expected_dynamic);
455    }
456
457    #[test]
458    fn fuzz_index_calculation() {
459        use common_cache::page::PageLayout;
460        use std::num::NonZeroU64;
461
462        let mut rng = rand::rng();
463        let mut data = vec![0u8; 1024 * 1024];
464        rng.fill_bytes(&mut data);
465        for _ in 0..FUZZ_REPEAT_TIMES {
466            let offset = rng.random_range(0..data.len() as u64);
467            let size = rng.random_range(0..data.len() as u64 - offset);
468            let page_size = NonZeroU64::new(rng.random_range(1_u64..1024)).unwrap();
469            let layout = PageLayout::new(page_size, data.len() as u64);
470            let requested = offset..offset + size;
471            let mut read = Vec::with_capacity(size as usize);
472            for id in layout.page_ids(requested.clone()).unwrap() {
473                let range = layout.page_range(id).unwrap();
474                let page = &data[range.start as usize..range.end as usize];
475                let slice = layout
476                    .request_slice(id, requested.clone())
477                    .unwrap()
478                    .unwrap();
479                read.extend_from_slice(&page[slice]);
480            }
481            assert_eq!(read, data[offset as usize..(offset + size) as usize]);
482        }
483    }
484}