mito2/cache/
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
15pub mod bloom_filter_index;
16pub mod inverted_index;
17pub mod result_cache;
18
19use std::future::Future;
20use std::hash::Hash;
21use std::io;
22use std::ops::Range;
23use std::sync::Arc;
24
25use async_trait::async_trait;
26use bytes::Bytes;
27use common_base::range_read::{Metadata, RangeReader};
28use common_cache::paged::PagedCache;
29
30use crate::metrics::CACHE_BYTES;
31
32/// Metrics for index metadata.
33const INDEX_METADATA_TYPE: &str = "index_metadata";
34/// Metrics for index content.
35const INDEX_CONTENT_TYPE: &str = "index_content";
36
37/// Metrics collected from IndexCache operations.
38#[derive(Debug, Default, Clone)]
39pub struct IndexCacheMetrics {
40    /// Number of cache hits.
41    pub cache_hit: usize,
42    /// Number of cache misses.
43    pub cache_miss: usize,
44    /// Number of pages accessed.
45    pub num_pages: usize,
46    /// Total bytes from pages.
47    pub page_bytes: u64,
48}
49
50/// Cache for index metadata and content.
51pub struct IndexCache<K, M> {
52    /// Cache for index metadata
53    index_metadata: common_cache::Cache<K, Arc<M>>,
54    /// Cache for index content.
55    index: PagedCache<K>,
56    index_type: &'static str,
57}
58
59impl<K, M> IndexCache<K, M>
60where
61    K: Hash + Eq + Clone + Send + Sync + 'static,
62    M: Send + Sync + 'static,
63{
64    pub fn new_with_weighter(
65        index_metadata_cap: u64,
66        index_content_cap: u64,
67        page_size: u64,
68        index_type: &'static str,
69        metadata_type: &'static str,
70        weight_of_metadata: fn(&K, &Arc<M>) -> u32,
71        weight_of_content: fn(&(K, u64), &Bytes) -> u32,
72    ) -> Self {
73        common_telemetry::debug!(
74            "Building IndexCache with metadata size: {index_metadata_cap}, content size: {index_content_cap}, page size: {page_size}, index type: {index_type}"
75        );
76        let index_metadata = common_cache::Cache::new(
77            &format!("index_metadata_{index_type}"),
78            index_metadata_cap,
79            weight_of_metadata,
80            |_, _, _| {},
81        )
82        .with_size_metric(CACHE_BYTES.bind(&[INDEX_METADATA_TYPE]))
83        .with_metrics(crate::metrics::cache_loading_metrics(metadata_type));
84        let index_cache = PagedCache::new(
85            &format!("index_content_{index_type}"),
86            index_content_cap,
87            page_size,
88            weight_of_content,
89            crate::metrics::cache_page_observer(index_type, INDEX_CONTENT_TYPE, |hits, misses| {
90                (hits, misses)
91            }),
92        );
93        Self {
94            index_metadata,
95            index: index_cache,
96            index_type,
97        }
98    }
99}
100
101impl<K, M> IndexCache<K, M>
102where
103    K: Hash + Eq + Clone + Send + Sync + 'static,
104    M: Send + Sync + 'static,
105{
106    /// Looks up completed metadata without starting or waiting for a load.
107    pub async fn get_metadata(&self, key: K) -> Option<Arc<M>> {
108        self.index_metadata.get(&key).await
109    }
110
111    /// Inserts completed metadata for cache eviction and accounting tests.
112    #[cfg(test)]
113    pub async fn put_metadata(&self, key: K, metadata: Arc<M>) {
114        self.index_metadata.insert(key, metadata).await;
115    }
116
117    /// Shares loads after the caller's initial lookup missed. Native request counts
118    /// describe these submissions, including waiters, rather than completed hits.
119    async fn get_or_load_metadata<F, E>(&self, key: K, load: F) -> Result<Arc<M>, Arc<E>>
120    where
121        F: Future<Output = Result<Arc<M>, E>>,
122        E: Send + Sync + 'static,
123    {
124        if self.index_metadata.capacity() == 0 {
125            return load.await.map_err(Arc::new);
126        }
127        self.index_metadata.try_get_with(&key, load).await
128    }
129
130    /// Fetches one byte range through fixed-size pages.
131    async fn get_or_load<F, Fut, E>(
132        &self,
133        key: K,
134        file_size: u64,
135        offset: u64,
136        size: u32,
137        load: F,
138    ) -> Result<(Bytes, IndexCacheMetrics), Arc<io::Error>>
139    where
140        F: Fn(Vec<Range<u64>>) -> Fut + Send + Sync,
141        Fut: Future<Output = Result<Vec<Bytes>, E>> + Send,
142        E: std::error::Error + Send + Sync + 'static,
143    {
144        let end = offset.checked_add(u64::from(size)).ok_or_else(|| {
145            Arc::new(io::Error::new(
146                io::ErrorKind::InvalidInput,
147                "Index range overflow",
148            ))
149        })?;
150        let (pages, metrics) = self
151            .get_or_load_vec(key, file_size, std::slice::from_ref(&(offset..end)), load)
152            .await?;
153        Ok((pages.into_iter().next().unwrap_or_default(), metrics))
154    }
155
156    /// Shares missing page loads while retaining legacy unique-page metrics.
157    async fn get_or_load_vec<F, Fut, E>(
158        &self,
159        key: K,
160        file_size: u64,
161        ranges: &[Range<u64>],
162        load: F,
163    ) -> Result<(Vec<Bytes>, IndexCacheMetrics), Arc<io::Error>>
164    where
165        F: Fn(Vec<Range<u64>>) -> Fut + Send + Sync,
166        Fut: Future<Output = Result<Vec<Bytes>, E>> + Send,
167        E: std::error::Error + Send + Sync + 'static,
168    {
169        let reader = IndexPageReader {
170            load: &load,
171            file_size,
172        };
173        let (pages, stats) = self.index.fetch(key, file_size, ranges, &reader).await?;
174        crate::metrics::cache_page_fetch_metrics(self.index_type, &stats);
175        Ok((
176            pages,
177            IndexCacheMetrics {
178                cache_hit: stats.cache_hit,
179                cache_miss: stats.cache_miss,
180                num_pages: stats.cache_hit + stats.cache_miss,
181                page_bytes: stats.aligned_bytes,
182            },
183        ))
184    }
185
186    /// Invalidates all cache entries whose keys satisfy `predicate`.
187    pub fn invalidate_if<F>(&self, predicate: F)
188    where
189        F: Fn(&K) -> bool + Send + Sync + 'static,
190    {
191        let predicate = Arc::new(predicate);
192        let metadata_predicate = Arc::clone(&predicate);
193
194        self.index_metadata
195            .invalidate_if(move |key, _| metadata_predicate(key))
196            .expect("cache should support invalidation closures");
197
198        self.index
199            .invalidate_if(move |key| predicate(key))
200            .expect("cache should support invalidation closures");
201    }
202}
203
204/// Adapts borrowed index readers to the common range-reader interface.
205struct IndexPageReader<'a, F> {
206    load: &'a F,
207    file_size: u64,
208}
209
210#[async_trait]
211impl<F, Fut, E> RangeReader for IndexPageReader<'_, F>
212where
213    F: Fn(Vec<Range<u64>>) -> Fut + Send + Sync,
214    Fut: Future<Output = Result<Vec<Bytes>, E>> + Send,
215    E: std::error::Error + Send + Sync + 'static,
216{
217    async fn metadata(&self) -> io::Result<Metadata> {
218        Ok(Metadata {
219            content_length: self.file_size,
220        })
221    }
222
223    async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
224        let mut reads = self.read_vec(&[range]).await?;
225        if reads.len() != 1 {
226            return Err(io::Error::new(
227                io::ErrorKind::InvalidData,
228                "Index reader returned an invalid result count",
229            ));
230        }
231        Ok(reads.remove(0))
232    }
233
234    async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
235        (self.load)(ranges.to_vec())
236            .await
237            .map_err(|error| io::Error::other(Arc::new(error)))
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use tokio::sync::Notify;
244
245    use super::*;
246
247    fn metadata_cache() -> IndexCache<u64, u64> {
248        IndexCache::new_with_weighter(
249            1024,
250            1024,
251            4,
252            "test",
253            "test_index_metadata",
254            |_, _| 8,
255            |_, bytes| bytes.len() as u32,
256        )
257    }
258
259    #[tokio::test]
260    async fn metadata_shares_same_key_without_blocking_other_keys() {
261        let cache = metadata_cache();
262        let release = Notify::new();
263        let mut owner = Box::pin(cache.get_or_load_metadata(1, async {
264            release.notified().await;
265            Ok::<_, std::io::Error>(Arc::new(11))
266        }));
267        assert!(futures::poll!(&mut owner).is_pending());
268        let mut waiter = Box::pin(cache.get_or_load_metadata(1, async {
269            Err::<Arc<u64>, _>(std::io::Error::other("unexpected duplicate load"))
270        }));
271        assert!(futures::poll!(&mut waiter).is_pending());
272        let other = cache
273            .get_or_load_metadata(2, async { Ok::<_, std::io::Error>(Arc::new(33)) })
274            .await
275            .unwrap();
276        assert_eq!(*other, 33);
277        release.notify_one();
278        let owner = owner.await.unwrap();
279        let waiter = waiter.await.unwrap();
280        assert!(Arc::ptr_eq(&owner, &waiter));
281        assert_eq!(*cache.get_metadata(1).await.unwrap(), 11);
282    }
283
284    #[tokio::test]
285    async fn metadata_failure_is_shared_but_not_cached() {
286        let cache = metadata_cache();
287        let release = Notify::new();
288        let mut owner = Box::pin(cache.get_or_load_metadata(1, async {
289            release.notified().await;
290            Err::<Arc<u64>, _>(std::io::Error::other("metadata read failed"))
291        }));
292        assert!(futures::poll!(&mut owner).is_pending());
293        let mut waiter = Box::pin(
294            cache.get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(22)) }),
295        );
296        assert!(futures::poll!(&mut waiter).is_pending());
297        release.notify_one();
298        let error = owner.await.unwrap_err();
299        let shared = waiter.await.unwrap_err();
300        assert!(Arc::ptr_eq(&error, &shared));
301        assert!(cache.get_metadata(1).await.is_none());
302        let retried = cache
303            .get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(33)) })
304            .await
305            .unwrap();
306        assert_eq!(*retried, 33);
307    }
308
309    #[tokio::test]
310    async fn cancelled_metadata_loader_allows_waiter_to_take_over() {
311        let cache = metadata_cache();
312        let mut owner = Box::pin(cache.get_or_load_metadata(
313            1,
314            std::future::pending::<Result<Arc<u64>, std::io::Error>>(),
315        ));
316        assert!(futures::poll!(&mut owner).is_pending());
317        let mut waiter = Box::pin(
318            cache.get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(22)) }),
319        );
320        assert!(futures::poll!(&mut waiter).is_pending());
321        drop(owner);
322        let value = tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
323            .await
324            .unwrap()
325            .unwrap();
326        assert_eq!(*value, 22);
327    }
328
329    #[tokio::test]
330    async fn metadata_invalidation_releases_charge_with_live_reader() {
331        let cache = metadata_cache();
332        let value = Arc::new(11);
333        cache.index_metadata.insert(1, value.clone()).await;
334        cache.put_metadata(2, Arc::new(22)).await;
335        cache.index_metadata.run_pending_tasks().await;
336        assert_eq!(cache.index_metadata.entry_count(), 2);
337        cache.invalidate_if(|key| *key == 1);
338        assert!(cache.get_metadata(1).await.is_none());
339        assert_eq!(*cache.get_metadata(2).await.unwrap(), 22);
340        cache.index_metadata.run_pending_tasks().await;
341        assert_eq!(cache.index_metadata.entry_count(), 1);
342        assert_eq!(*value, 11);
343        drop(value);
344        assert_eq!(cache.index_metadata.entry_count(), 1);
345    }
346
347    #[tokio::test]
348    async fn disabled_metadata_cache_does_not_share_loads() {
349        let cache = IndexCache::new_with_weighter(
350            0,
351            1024,
352            4,
353            "test_disabled",
354            "test_disabled_index_metadata",
355            |_: &u64, _: &Arc<u64>| 8,
356            |_, bytes| bytes.len() as u32,
357        );
358        let mut owner = Box::pin(cache.get_or_load_metadata(
359            1,
360            std::future::pending::<Result<Arc<u64>, std::io::Error>>(),
361        ));
362        assert!(futures::poll!(&mut owner).is_pending());
363        let mut independent = Box::pin(
364            cache.get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(22)) }),
365        );
366        assert!(futures::poll!(&mut independent).is_ready());
367    }
368
369    #[tokio::test]
370    async fn metadata_metrics_distinguish_requests_from_loads() {
371        let metrics = crate::metrics::cache_loading_metrics("test_index_metadata_metrics");
372        let mut cache = metadata_cache();
373        cache.index_metadata = cache.index_metadata.with_metrics(metrics.clone());
374        let release = Notify::new();
375        let mut owner = Box::pin(cache.get_or_load_metadata(1, async {
376            release.notified().await;
377            Ok::<_, std::io::Error>(Arc::new(11))
378        }));
379        let mut waiter = Box::pin(cache.get_or_load_metadata(1, async {
380            Err::<Arc<u64>, _>(std::io::Error::other("duplicate metadata load"))
381        }));
382        assert!(futures::poll!(&mut owner).is_pending());
383        assert!(futures::poll!(&mut waiter).is_pending());
384        assert_eq!(metrics.requests.active.get(), 2);
385        assert_eq!(metrics.loads.active.get(), 1);
386        release.notify_one();
387        let (owner, waiter) = tokio::join!(owner, waiter);
388        assert!(Arc::ptr_eq(&owner.unwrap(), &waiter.unwrap()));
389        assert_eq!(metrics.requests.succeeded.get(), 2);
390        assert_eq!(metrics.loads.succeeded.get(), 1);
391        assert_eq!(metrics.requests.active.get(), 0);
392        assert_eq!(metrics.loads.active.get(), 0);
393    }
394    #[tokio::test]
395    async fn overlapping_page_requests_share_bytes_and_preserve_legacy_metrics() {
396        let cache = metadata_cache();
397        let release = Notify::new();
398        let reads = std::sync::Mutex::new(Vec::new());
399        let load = |ranges: Vec<Range<u64>>| {
400            reads.lock().unwrap().push(ranges.clone());
401            let release = &release;
402            async move {
403                release.notified().await;
404                Ok::<_, io::Error>(
405                    ranges
406                        .iter()
407                        .map(|range| {
408                            Bytes::from(range.clone().map(|v| v as u8).collect::<Vec<_>>())
409                        })
410                        .collect(),
411                )
412            }
413        };
414        let first_ranges = std::slice::from_ref(&(0..8));
415        let second_ranges = std::slice::from_ref(&(4..12));
416        let mut first = Box::pin(cache.get_or_load_vec(1, 12, first_ranges, &load));
417        assert!(futures::poll!(&mut first).is_pending());
418        let mut second = Box::pin(cache.get_or_load_vec(1, 12, second_ranges, &load));
419        assert!(futures::poll!(&mut second).is_pending());
420        release.notify_waiters();
421        let (first, second) = tokio::join!(first, second);
422        let (first, first_metrics) = first.unwrap();
423        let (second, second_metrics) = second.unwrap();
424        assert_eq!(first, vec![Bytes::from_static(&[0, 1, 2, 3, 4, 5, 6, 7])]);
425        assert_eq!(
426            second,
427            vec![Bytes::from_static(&[4, 5, 6, 7, 8, 9, 10, 11])]
428        );
429        for metrics in [first_metrics, second_metrics] {
430            assert_eq!(
431                (
432                    metrics.cache_hit,
433                    metrics.cache_miss,
434                    metrics.num_pages,
435                    metrics.page_bytes
436                ),
437                (0, 2, 2, 8)
438            );
439        }
440        let batches = reads.lock().unwrap().clone();
441        assert_eq!(batches, vec![vec![0..4, 4..8], vec![8..12]]);
442        let mut physical = batches.into_iter().flatten().collect::<Vec<_>>();
443        physical.sort_by_key(|range| range.start);
444        assert!(physical.windows(2).all(|pair| pair[0].end <= pair[1].start));
445        assert_eq!(
446            physical
447                .iter()
448                .map(|range| range.end - range.start)
449                .sum::<u64>(),
450            12
451        );
452        let (hot, metrics) = cache
453            .get_or_load_vec(1, 12, &[1..6, 4..12], |_| async {
454                Err::<Vec<Bytes>, _>(io::Error::other("unexpected hot read"))
455            })
456            .await
457            .unwrap();
458        assert_eq!(hot[0], Bytes::from_static(&[1, 2, 3, 4, 5]));
459        assert_eq!(
460            (
461                metrics.cache_hit,
462                metrics.cache_miss,
463                metrics.num_pages,
464                metrics.page_bytes
465            ),
466            (3, 0, 3, 12)
467        );
468        cache.invalidate_if(|key| *key == 1);
469        assert!(!cache.index.contains(&1, 0).await);
470    }
471
472    #[tokio::test]
473    async fn index_pages_preserve_shared_typed_read_errors_and_retry() {
474        use common_error::ext::ErrorExt;
475        use index::inverted_index::error::{Error, IndexNotFoundSnafu};
476        let cache = metadata_cache();
477        let error = cache
478            .get_or_load_vec(1, 8, std::slice::from_ref(&(0..8)), |_| async {
479                Err::<Vec<Bytes>, _>(IndexNotFoundSnafu { name: "missing" }.build())
480            })
481            .await
482            .unwrap_err();
483        let retained = crate::cache::find_shared_page_source::<Error>(&error).unwrap();
484        assert_eq!(
485            retained.status_code(),
486            common_error::status_code::StatusCode::InvalidArguments
487        );
488        assert!(!cache.index.contains(&1, 0).await);
489        let (pages, _) = cache
490            .get_or_load_vec(1, 8, std::slice::from_ref(&(0..8)), |ranges| async move {
491                Ok::<_, io::Error>(
492                    ranges
493                        .iter()
494                        .map(|r| Bytes::from(vec![7; (r.end - r.start) as usize]))
495                        .collect(),
496                )
497            })
498            .await
499            .unwrap();
500        assert_eq!(pages, vec![Bytes::from(vec![7; 8])]);
501    }
502    #[tokio::test]
503    async fn single_page_range_retains_shared_bytes() {
504        let cache = metadata_cache();
505        let (first, _) = cache
506            .get_or_load(1, 8, 1, 2, |ranges| async move {
507                Ok::<_, io::Error>(
508                    ranges
509                        .into_iter()
510                        .map(|range| Bytes::from(range.map(|n| n as u8).collect::<Vec<_>>()))
511                        .collect(),
512                )
513            })
514            .await
515            .unwrap();
516        let (second, metrics) = cache
517            .get_or_load(1, 8, 1, 2, |_| async {
518                Err::<Vec<Bytes>, _>(io::Error::other("unexpected read on cache hit"))
519            })
520            .await
521            .unwrap();
522        assert_eq!(&first[..], &[1, 2]);
523        assert_eq!(first, second);
524        assert_eq!(first.as_ptr(), second.as_ptr());
525        assert_eq!(metrics.cache_hit, 1);
526        assert_eq!(metrics.cache_miss, 0);
527    }
528}
mito2/cache/index.rs

mito2/cache/
index.rs
