mito2/cache/index/
inverted_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 core::ops::Range;
16use std::sync::Arc;
17use std::time::Instant;
18
19use api::v1::index::InvertedIndexMetas;
20use async_trait::async_trait;
21use bytes::Bytes;
22use index::inverted_index::error::{CommonIoSnafu, Error, Result, SharedReadSnafu};
23use index::inverted_index::format::reader::{InvertedIndexReadMetrics, InvertedIndexReader};
24use prost::Message;
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_INVERTED_INDEX: &str = "inverted_index";
32
33/// Cache for inverted index.
34pub type InvertedIndexCache = IndexCache<(FileId, IndexVersion), InvertedIndexMetas>;
35pub type InvertedIndexCacheRef = Arc<InvertedIndexCache>;
36
37impl InvertedIndexCache {
38    /// Creates a new inverted index cache.
39    pub fn new(index_metadata_cap: u64, index_content_cap: u64, page_size: u64) -> Self {
40        Self::new_with_weighter(
41            index_metadata_cap,
42            index_content_cap,
43            page_size,
44            INDEX_TYPE_INVERTED_INDEX,
45            "inverted_index_metadata",
46            inverted_index_metadata_weight,
47            inverted_index_content_weight,
48        )
49    }
50
51    /// Removes all cached entries for the given `file_id`.
52    pub fn invalidate_file(&self, file_id: FileId) {
53        self.invalidate_if(move |key| key.0 == file_id);
54    }
55}
56
57/// Calculates weight for inverted index metadata.
58fn inverted_index_metadata_weight(k: &(FileId, IndexVersion), v: &Arc<InvertedIndexMetas>) -> u32 {
59    (k.0.as_bytes().len() + size_of::<IndexVersion>() + v.encoded_len()) as u32
60}
61
62/// Calculates weight for inverted index content.
63fn inverted_index_content_weight((k, _): &((FileId, IndexVersion), u64), v: &Bytes) -> u32 {
64    (k.0.as_bytes().len() + size_of::<IndexVersion>() + v.len()) as u32
65}
66
67/// Inverted index blob reader with cache.
68pub struct CachedInvertedIndexBlobReader<R> {
69    file_id: FileId,
70    index_version: IndexVersion,
71    blob_size: u64,
72    inner: R,
73    cache: InvertedIndexCacheRef,
74}
75
76impl<R> CachedInvertedIndexBlobReader<R> {
77    /// Creates a new inverted index blob reader with cache.
78    pub fn new(
79        file_id: FileId,
80        index_version: IndexVersion,
81        blob_size: u64,
82        inner: R,
83        cache: InvertedIndexCacheRef,
84    ) -> Self {
85        Self {
86            file_id,
87            index_version,
88            blob_size,
89            inner,
90            cache,
91        }
92    }
93}
94
95#[async_trait]
96impl<R: InvertedIndexReader> InvertedIndexReader for CachedInvertedIndexBlobReader<R> {
97    async fn range_read<'a>(
98        &self,
99        offset: u64,
100        size: u32,
101        metrics: Option<&'a mut InvertedIndexReadMetrics>,
102    ) -> Result<Vec<u8>> {
103        let start = metrics.as_ref().map(|_| Instant::now());
104
105        let inner = &self.inner;
106        let (result, cache_metrics) = self
107            .cache
108            .get_or_load(
109                (self.file_id, self.index_version),
110                self.blob_size,
111                offset,
112                size,
113                move |ranges| async move { inner.read_vec(&ranges, None).await },
114            )
115            .await
116            .map_err(map_page_error)?;
117
118        if let Some(m) = metrics {
119            m.total_bytes += cache_metrics.page_bytes;
120            m.total_ranges += cache_metrics.num_pages;
121            m.cache_hit += cache_metrics.cache_hit;
122            m.cache_miss += cache_metrics.cache_miss;
123            m.fetch_elapsed += start.unwrap().elapsed();
124        }
125
126        Ok(result.to_vec())
127    }
128
129    async fn read_vec<'a>(
130        &self,
131        ranges: &[Range<u64>],
132        metrics: Option<&'a mut InvertedIndexReadMetrics>,
133    ) -> Result<Vec<Bytes>> {
134        let start = metrics.as_ref().map(|_| Instant::now());
135
136        let inner = &self.inner;
137        let (pages, total_cache_metrics) = self
138            .cache
139            .get_or_load_vec(
140                (self.file_id, self.index_version),
141                self.blob_size,
142                ranges,
143                move |ranges| async move { inner.read_vec(&ranges, None).await },
144            )
145            .await
146            .map_err(map_page_error)?;
147        if let Some(m) = metrics {
148            m.total_bytes += total_cache_metrics.page_bytes;
149            m.total_ranges += total_cache_metrics.num_pages;
150            m.cache_hit += total_cache_metrics.cache_hit;
151            m.cache_miss += total_cache_metrics.cache_miss;
152            m.fetch_elapsed += start.unwrap().elapsed();
153        }
154
155        Ok(pages)
156    }
157
158    async fn metadata<'a>(
159        &self,
160        metrics: Option<&'a mut InvertedIndexReadMetrics>,
161    ) -> Result<Arc<InvertedIndexMetas>> {
162        let key = (self.file_id, self.index_version);
163        if let Some(cached) = self.cache.get_metadata(key).await {
164            CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc();
165            if let Some(m) = metrics {
166                m.cache_hit += 1;
167            }
168            return Ok(cached);
169        }
170
171        let start = metrics.as_ref().map(|_| Instant::now());
172        let collect_metrics = metrics.is_some();
173        let mut loaded = InvertedIndexReadMetrics::default();
174        let result = self
175            .cache
176            .get_or_load_metadata(key, async {
177                self.inner
178                    .metadata(collect_metrics.then_some(&mut loaded))
179                    .await
180            })
181            .await;
182        if let Some(m) = metrics {
183            // Only the selected loader performs I/O. Waiters retain zero I/O
184            // counts. fetch_elapsed is per-caller elapsed (including wait and error),
185            // not physical I/O time summed across callers.
186            m.total_bytes += loaded.total_bytes;
187            m.total_ranges += loaded.total_ranges;
188            m.cache_hit += loaded.cache_hit;
189            m.cache_miss += loaded.cache_miss;
190            if let Some(start) = start {
191                m.fetch_elapsed += start.elapsed();
192            }
193        }
194        let meta = result.map_err(|source| SharedReadSnafu.into_error(source))?;
195        CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc();
196        Ok(meta)
197    }
198}
199
200/// Recovers the typed source retained by the range-reader adapter.
201fn map_page_error(error: Arc<std::io::Error>) -> Error {
202    if let Some(source) = crate::cache::find_shared_page_source::<Error>(&error) {
203        SharedReadSnafu.into_error(source)
204    } else {
205        CommonIoSnafu.into_error(std::io::Error::other(error))
206    }
207}
208
209#[cfg(test)]
210mod test {
211    use std::num::NonZeroUsize;
212
213    use futures::stream;
214    use index::Bytes;
215    use index::bitmap::{Bitmap, BitmapType};
216    use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader};
217    use index::inverted_index::format::writer::{InvertedIndexBlobWriter, InvertedIndexWriter};
218    use prometheus::register_int_counter_vec;
219    use rand::{Rng, RngCore};
220
221    use super::*;
222    use crate::sst::index::store::InstrumentedStore;
223    use crate::test_util::TestEnv;
224
225    // Repeat times for following little fuzz tests.
226    const FUZZ_REPEAT_TIMES: usize = 100;
227
228    struct MetadataReader {
229        release: tokio::sync::Notify,
230        fail: std::sync::atomic::AtomicBool,
231    }
232
233    #[async_trait]
234    impl InvertedIndexReader for MetadataReader {
235        async fn range_read<'a>(
236            &self,
237            _: u64,
238            _: u32,
239            _: Option<&'a mut InvertedIndexReadMetrics>,
240        ) -> Result<Vec<u8>> {
241            unreachable!("metadata test does not read content pages")
242        }
243
244        async fn read_vec<'a>(
245            &self,
246            _: &[Range<u64>],
247            _: Option<&'a mut InvertedIndexReadMetrics>,
248        ) -> Result<Vec<bytes::Bytes>> {
249            unreachable!("metadata test does not read content pages")
250        }
251
252        async fn metadata<'a>(
253            &self,
254            metrics: Option<&'a mut InvertedIndexReadMetrics>,
255        ) -> Result<Arc<InvertedIndexMetas>> {
256            self.release.notified().await;
257            if let Some(metrics) = metrics {
258                metrics.total_bytes += 8192;
259                metrics.total_ranges += 1;
260                // The wrapper must replace, not add, this physical-fetch duration.
261                metrics.fetch_elapsed += std::time::Duration::from_secs(3600);
262            }
263            if self.fail.load(std::sync::atomic::Ordering::Relaxed) {
264                return Err(CommonIoSnafu.into_error(std::io::Error::other("metadata failed")));
265            }
266            Ok(Arc::new(InvertedIndexMetas::default()))
267        }
268    }
269
270    #[tokio::test]
271    async fn metadata_waiters_do_not_duplicate_io_metrics() {
272        for fail in [false, true] {
273            let cache = Arc::new(InvertedIndexCache::new(4096, 4096, 1024));
274            let reader = CachedInvertedIndexBlobReader::new(
275                FileId::random(),
276                0,
277                8192,
278                MetadataReader {
279                    release: tokio::sync::Notify::new(),
280                    fail: std::sync::atomic::AtomicBool::new(fail),
281                },
282                cache,
283            );
284            let mut owner_metrics = InvertedIndexReadMetrics::default();
285            let mut waiter_metrics = InvertedIndexReadMetrics::default();
286            let mut owner = Box::pin(reader.metadata(Some(&mut owner_metrics)));
287            assert!(futures::poll!(&mut owner).is_pending());
288            let mut waiter = Box::pin(reader.metadata(Some(&mut waiter_metrics)));
289            assert!(futures::poll!(&mut waiter).is_pending());
290            reader.inner.release.notify_one();
291            let owner_result = owner.await;
292            let waiter_result = waiter.await;
293            assert_eq!(owner_result.is_err(), fail);
294            assert_eq!(waiter_result.is_err(), fail);
295            assert_eq!(
296                (owner_metrics.total_bytes, owner_metrics.total_ranges),
297                (8192, 1)
298            );
299            assert_eq!(
300                (waiter_metrics.total_bytes, waiter_metrics.total_ranges),
301                (0, 0)
302            );
303            assert_eq!((owner_metrics.cache_hit, waiter_metrics.cache_hit), (0, 0));
304            for metrics in [&owner_metrics, &waiter_metrics] {
305                // Per-caller elapsed includes waiting, including on error; it is
306                // not a sum of physical I/O durations across concurrent callers.
307                assert!(metrics.fetch_elapsed > std::time::Duration::ZERO);
308                assert!(metrics.fetch_elapsed < std::time::Duration::from_secs(3600));
309            }
310            if fail {
311                let Error::SharedRead { source: owner, .. } = owner_result.unwrap_err() else {
312                    panic!("owner must preserve the typed shared read error");
313                };
314                let Error::SharedRead { source: waiter, .. } = waiter_result.unwrap_err() else {
315                    panic!("waiter must preserve the typed shared read error");
316                };
317                assert!(Arc::ptr_eq(&owner, &waiter));
318            } else {
319                let mut hot_metrics = InvertedIndexReadMetrics::default();
320                reader.metadata(Some(&mut hot_metrics)).await.unwrap();
321                assert_eq!(hot_metrics.cache_hit, 1);
322                assert_eq!(hot_metrics.total_bytes, 0);
323            }
324        }
325    }
326
327    // Fuzz test for index data page key
328    #[test]
329    fn fuzz_index_calculation() {
330        use common_cache::page::PageLayout;
331        use std::num::NonZeroU64;
332
333        let mut rng = rand::rng();
334        let mut data = vec![0u8; 1024 * 1024];
335        rng.fill_bytes(&mut data);
336        for _ in 0..FUZZ_REPEAT_TIMES {
337            let offset = rng.random_range(0..data.len() as u64);
338            let size = rng.random_range(0..data.len() as u64 - offset);
339            let page_size = NonZeroU64::new(rng.random_range(1_u64..1024)).unwrap();
340            let layout = PageLayout::new(page_size, data.len() as u64);
341            let requested = offset..offset + size;
342            let mut read = Vec::with_capacity(size as usize);
343            for id in layout.page_ids(requested.clone()).unwrap() {
344                let range = layout.page_range(id).unwrap();
345                let page = &data[range.start as usize..range.end as usize];
346                let slice = layout
347                    .request_slice(id, requested.clone())
348                    .unwrap()
349                    .unwrap();
350                read.extend_from_slice(&page[slice]);
351            }
352            assert_eq!(read, data[offset as usize..(offset + size) as usize]);
353        }
354    }
355
356    fn unpack(fst_value: u64) -> [u32; 2] {
357        bytemuck::cast::<u64, [u32; 2]>(fst_value)
358    }
359
360    async fn create_inverted_index_blob() -> Vec<u8> {
361        let mut blob = Vec::new();
362        let mut writer = InvertedIndexBlobWriter::new(&mut blob);
363        writer
364            .add_index(
365                "tag0".to_string(),
366                Bitmap::from_lsb0_bytes(&[0b0000_0001, 0b0000_0000], BitmapType::Roaring),
367                Box::new(stream::iter(vec![
368                    Ok((
369                        Bytes::from("a"),
370                        Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring),
371                    )),
372                    Ok((
373                        Bytes::from("b"),
374                        Bitmap::from_lsb0_bytes(&[0b0010_0000], BitmapType::Roaring),
375                    )),
376                    Ok((
377                        Bytes::from("c"),
378                        Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring),
379                    )),
380                ])),
381                index::bitmap::BitmapType::Roaring,
382            )
383            .await
384            .unwrap();
385        writer
386            .add_index(
387                "tag1".to_string(),
388                Bitmap::from_lsb0_bytes(&[0b0000_0001, 0b0000_0000], BitmapType::Roaring),
389                Box::new(stream::iter(vec![
390                    Ok((
391                        Bytes::from("x"),
392                        Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring),
393                    )),
394                    Ok((
395                        Bytes::from("y"),
396                        Bitmap::from_lsb0_bytes(&[0b0010_0000], BitmapType::Roaring),
397                    )),
398                    Ok((
399                        Bytes::from("z"),
400                        Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring),
401                    )),
402                ])),
403                index::bitmap::BitmapType::Roaring,
404            )
405            .await
406            .unwrap();
407        writer
408            .finish(8, NonZeroUsize::new(1).unwrap())
409            .await
410            .unwrap();
411
412        blob
413    }
414
415    #[tokio::test]
416    async fn test_inverted_index_cache() {
417        let blob = create_inverted_index_blob().await;
418
419        // Init a test range reader in local fs.
420        let mut env = TestEnv::new().await;
421        let file_size = blob.len() as u64;
422        let index_version = 0;
423        let store = env.init_object_store_manager();
424        let temp_path = "data";
425        store.write(temp_path, blob).await.unwrap();
426        let store = InstrumentedStore::new(store);
427        let metric =
428            register_int_counter_vec!("test_bytes", "a counter for test", &["test"]).unwrap();
429        let counter = metric.with_label_values(&["test"]);
430        let range_reader = store
431            .range_reader("data", &counter, &counter)
432            .await
433            .unwrap();
434
435        let reader = InvertedIndexBlobReader::new(range_reader);
436        let cached_reader = CachedInvertedIndexBlobReader::new(
437            FileId::random(),
438            index_version,
439            file_size,
440            reader,
441            Arc::new(InvertedIndexCache::new(8192, 8192, 50)),
442        );
443        let metadata = cached_reader.metadata(None).await.unwrap();
444        assert_eq!(metadata.total_row_count, 8);
445        assert_eq!(metadata.segment_row_count, 1);
446        assert_eq!(metadata.metas.len(), 2);
447        // tag0
448        let tag0 = metadata.metas.get("tag0").unwrap();
449        let stats0 = tag0.stats.as_ref().unwrap();
450        assert_eq!(stats0.distinct_count, 3);
451        assert_eq!(stats0.null_count, 1);
452        assert_eq!(stats0.min_value, Bytes::from("a"));
453        assert_eq!(stats0.max_value, Bytes::from("c"));
454        let fst0 = cached_reader
455            .fst(
456                tag0.base_offset + tag0.relative_fst_offset as u64,
457                tag0.fst_size,
458                None,
459            )
460            .await
461            .unwrap();
462        assert_eq!(fst0.len(), 3);
463        let [offset, size] = unpack(fst0.get(b"a").unwrap());
464        let bitmap = cached_reader
465            .bitmap(
466                tag0.base_offset + offset as u64,
467                size,
468                BitmapType::Roaring,
469                None,
470            )
471            .await
472            .unwrap();
473        assert_eq!(
474            bitmap,
475            Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring)
476        );
477        let [offset, size] = unpack(fst0.get(b"b").unwrap());
478        let bitmap = cached_reader
479            .bitmap(
480                tag0.base_offset + offset as u64,
481                size,
482                BitmapType::Roaring,
483                None,
484            )
485            .await
486            .unwrap();
487        assert_eq!(
488            bitmap,
489            Bitmap::from_lsb0_bytes(&[0b0010_0000], BitmapType::Roaring)
490        );
491        let [offset, size] = unpack(fst0.get(b"c").unwrap());
492        let bitmap = cached_reader
493            .bitmap(
494                tag0.base_offset + offset as u64,
495                size,
496                BitmapType::Roaring,
497                None,
498            )
499            .await
500            .unwrap();
501        assert_eq!(
502            bitmap,
503            Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring)
504        );
505
506        // tag1
507        let tag1 = metadata.metas.get("tag1").unwrap();
508        let stats1 = tag1.stats.as_ref().unwrap();
509        assert_eq!(stats1.distinct_count, 3);
510        assert_eq!(stats1.null_count, 1);
511        assert_eq!(stats1.min_value, Bytes::from("x"));
512        assert_eq!(stats1.max_value, Bytes::from("z"));
513        let fst1 = cached_reader
514            .fst(
515                tag1.base_offset + tag1.relative_fst_offset as u64,
516                tag1.fst_size,
517                None,
518            )
519            .await
520            .unwrap();
521        assert_eq!(fst1.len(), 3);
522        let [offset, size] = unpack(fst1.get(b"x").unwrap());
523        let bitmap = cached_reader
524            .bitmap(
525                tag1.base_offset + offset as u64,
526                size,
527                BitmapType::Roaring,
528                None,
529            )
530            .await
531            .unwrap();
532        assert_eq!(
533            bitmap,
534            Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring)
535        );
536        let [offset, size] = unpack(fst1.get(b"y").unwrap());
537        let bitmap = cached_reader
538            .bitmap(
539                tag1.base_offset + offset as u64,
540                size,
541                BitmapType::Roaring,
542                None,
543            )
544            .await
545            .unwrap();
546        assert_eq!(
547            bitmap,
548            Bitmap::from_lsb0_bytes(&[0b0010_0000], BitmapType::Roaring)
549        );
550        let [offset, size] = unpack(fst1.get(b"z").unwrap());
551        let bitmap = cached_reader
552            .bitmap(
553                tag1.base_offset + offset as u64,
554                size,
555                BitmapType::Roaring,
556                None,
557            )
558            .await
559            .unwrap();
560        assert_eq!(
561            bitmap,
562            Bitmap::from_lsb0_bytes(&[0b0000_0001], BitmapType::Roaring)
563        );
564
565        // fuzz test
566        let mut rng = rand::rng();
567        for _ in 0..FUZZ_REPEAT_TIMES {
568            let offset = rng.random_range(0..file_size);
569            let size = rng.random_range(0..file_size as u32 - offset as u32);
570            let expected = cached_reader.range_read(offset, size, None).await.unwrap();
571            let inner = &cached_reader.inner;
572            let (read, _cache_metrics) = cached_reader
573                .cache
574                .get_or_load(
575                    (cached_reader.file_id, cached_reader.index_version),
576                    file_size,
577                    offset,
578                    size,
579                    |ranges| async move { inner.read_vec(&ranges, None).await },
580                )
581                .await
582                .unwrap();
583            assert_eq!(read.as_ref(), expected.as_slice());
584        }
585    }
586
587    #[tokio::test]
588    async fn test_get_or_load_vec_loads_missing_pages_once() {
589        use std::sync::atomic::{AtomicUsize, Ordering};
590
591        let mut rng = rand::rng();
592        let mut data = vec![0u8; 64 * 1024];
593        rng.fill_bytes(&mut data);
594        let file_size = data.len() as u64;
595        let cache = InvertedIndexCache::new(1024 * 1024, 1024 * 1024, 1000);
596        let key = (FileId::random(), 0);
597
598        for _ in 0..FUZZ_REPEAT_TIMES {
599            cache.invalidate_file(key.0);
600            let ranges = (0..rng.random_range(1..20))
601                .map(|_| {
602                    let start = rng.random_range(0..file_size);
603                    start..rng.random_range(start..=file_size.min(start + 5000))
604                })
605                .collect::<Vec<_>>();
606            let expected = ranges
607                .iter()
608                .map(|r| bytes::Bytes::copy_from_slice(&data[r.start as usize..r.end as usize]))
609                .collect::<Vec<_>>();
610
611            let loads = AtomicUsize::new(0);
612            let physical_ranges = std::sync::Mutex::new(Vec::new());
613            let load = |pages: Vec<Range<u64>>| {
614                loads.fetch_add(1, Ordering::Relaxed);
615                physical_ranges
616                    .lock()
617                    .unwrap()
618                    .extend(pages.iter().cloned());
619                let data = &data;
620                async move {
621                    Ok::<_, std::io::Error>(
622                        pages
623                            .iter()
624                            .map(|r| {
625                                bytes::Bytes::copy_from_slice(
626                                    &data[r.start as usize..r.end as usize],
627                                )
628                            })
629                            .collect(),
630                    )
631                }
632            };
633            let (cold, cold_metrics) = cache
634                .get_or_load_vec(key, file_size, &ranges, load)
635                .await
636                .unwrap();
637            assert_eq!(cold, expected);
638            let cold_loads = loads.load(Ordering::Relaxed);
639            assert_eq!(cold_loads, usize::from(cold_metrics.cache_miss > 0));
640            let mut observed = physical_ranges.lock().unwrap().clone();
641            observed.sort_by_key(|r| r.start);
642            assert!(observed.windows(2).all(|w| w[0].end <= w[1].start));
643            assert_eq!(
644                observed.iter().map(|r| r.end - r.start).sum::<u64>(),
645                cold_metrics.page_bytes
646            );
647
648            let (warm, metrics) = cache
649                .get_or_load_vec(key, file_size, &ranges, load)
650                .await
651                .unwrap();
652            assert_eq!(warm, expected);
653            assert_eq!(metrics.cache_miss, 0);
654            assert_eq!(loads.load(Ordering::Relaxed), cold_loads);
655        }
656    }
657}
mito2/cache/index/inverted_index.rs

mito2/cache/index/
inverted_index.rs
