puffin/puffin_manager/
cache.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::future::Future;
16use std::sync::Arc;
17
18use common_cache::Cache;
19use common_cache::metrics::CacheMetrics;
20use common_cache::metrics::CacheSizeMetrics;
21use common_telemetry::tracing::{self, Instrument};
22use snafu::ResultExt;
23
24use crate::error::{CacheGetSnafu, Result};
25use crate::file_metadata::FileMetadata;
26use crate::puffin_manager::file_key::PuffinFileKey;
27
28/// Metrics for index metadata.
29const PUFFIN_METADATA_TYPE: &str = "puffin_metadata";
30
31pub type PuffinMetadataCacheRef = Arc<PuffinMetadataCache>;
32
33/// A cache for storing the metadata of the index files.
34pub struct PuffinMetadataCache {
35    cache: Cache<PuffinFileKey<String>, Arc<FileMetadata>>,
36}
37
38fn puffin_metadata_weight(k: &PuffinFileKey<String>, v: &Arc<FileMetadata>) -> u32 {
39    u32::try_from(k.estimated_size() + k.handle().capacity() + v.memory_usage()).unwrap_or(u32::MAX)
40}
41
42impl PuffinMetadataCache {
43    pub fn new(capacity: u64, cache_bytes: &CacheSizeMetrics) -> Self {
44        common_telemetry::debug!("Building PuffinMetadataCache with capacity: {capacity}");
45        Self {
46            cache: Cache::new(
47                "puffin_metadata",
48                capacity,
49                |k: &PuffinFileKey<String>, v: &Arc<FileMetadata>| puffin_metadata_weight(k, v),
50                |_, _, _| {},
51            )
52            .with_size_metric(cache_bytes.bind(&[PUFFIN_METADATA_TYPE])),
53        }
54    }
55
56    /// Attaches native loader metrics without changing existing cache byte accounting.
57    pub fn with_load_metrics(mut self, metrics: Arc<CacheMetrics>) -> Self {
58        self.cache = self.cache.with_metrics(metrics);
59        self
60    }
61
62    /// Gets the metadata from the cache.
63    pub async fn get_metadata(&self, file_id: &PuffinFileKey<String>) -> Option<Arc<FileMetadata>> {
64        self.cache.get(file_id).await
65    }
66
67    /// Shares one metadata load among concurrent readers of the same file.
68    #[tracing::instrument(name = "puffin_metadata_lookup", skip_all)]
69    pub async fn try_get_with<F>(
70        &self,
71        file_id: &PuffinFileKey<String>,
72        load: F,
73    ) -> Result<Arc<FileMetadata>>
74    where
75        F: Future<Output = Result<Arc<FileMetadata>>>,
76    {
77        self.cache
78            .try_get_with(file_id, async {
79                let metadata = load
80                    .instrument(tracing::info_span!("puffin_metadata_load"))
81                    .await?;
82                Ok(metadata)
83            })
84            .await
85            .context(CacheGetSnafu)
86    }
87
88    /// Puts the metadata into the cache.
89    pub async fn put_metadata(&self, file_id: PuffinFileKey<String>, metadata: Arc<FileMetadata>) {
90        self.cache.insert(file_id, metadata).await;
91    }
92
93    /// Invalidates the versioned handle across storage scopes when the caller has no manager.
94    /// Pending loaders are not fenced; physical removal runs asynchronously.
95    pub fn invalidate_file(&self, file_id: &str) -> std::io::Result<()> {
96        let file_id = file_id.to_string();
97        self.cache
98            .invalidate_if(move |key, _| key.handle() == &file_id)
99    }
100
101    /// Removes the metadata of the given file from the cache, if present.
102    pub async fn remove(&self, file_id: &PuffinFileKey<String>) {
103        self.cache.invalidate(file_id).await;
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use std::sync::atomic::{AtomicUsize, Ordering};
110
111    use prometheus::Opts;
112    use prometheus::core::Collector;
113
114    use super::*;
115    use crate::error::InvalidPuffinFooterSnafu;
116
117    fn key() -> PuffinFileKey<String> {
118        PuffinFileKey::new(
119            Arc::new(crate::puffin_manager::file_key::PuffinStorageScope::new(
120                "file".into(),
121                String::new(),
122                "/root/".into(),
123            )),
124            "file".into(),
125            "file".into(),
126        )
127    }
128
129    fn new_cache() -> (PuffinMetadataCache, CacheSizeMetrics) {
130        let metrics =
131            CacheSizeMetrics::new(Opts::new("test_metadata_bytes", "test"), &["type"]).unwrap();
132        let capacity = u64::from(puffin_metadata_weight(&key(), &metadata())) * 4;
133        (PuffinMetadataCache::new(capacity, &metrics), metrics)
134    }
135
136    fn sample_bytes(metrics: &CacheSizeMetrics) -> i64 {
137        metrics.collect()[0].get_metric()[0]
138            .get_gauge()
139            .as_ref()
140            .unwrap()
141            .value() as i64
142    }
143
144    fn metadata() -> Arc<FileMetadata> {
145        Arc::new(FileMetadata {
146            blobs: Vec::new(),
147            properties: Default::default(),
148        })
149    }
150
151    // Hold both requests at the same cold-load boundary. Only coalescing differs.
152    #[tokio::test]
153    async fn compare_separate_lookup_with_coalesced_load() {
154        let key = key();
155        for _ in 0..3 {
156            for coalesce in [false, true] {
157                let (cache, _) = new_cache();
158                let calls = AtomicUsize::new(0);
159                let gate = tokio::sync::Notify::new();
160                let request = || async {
161                    let load = async {
162                        calls.fetch_add(1, Ordering::Relaxed);
163                        gate.notified().await;
164                        Ok(metadata())
165                    };
166                    if coalesce {
167                        cache.try_get_with(&key, load).await
168                    } else if let Some(value) = cache.get_metadata(&key).await {
169                        Ok(value)
170                    } else {
171                        let value = load.await?;
172                        cache.put_metadata(key.clone(), value.clone()).await;
173                        Ok(value)
174                    }
175                };
176                let mut first = Box::pin(request());
177                let mut second = Box::pin(request());
178                assert!(futures::poll!(&mut first).is_pending());
179                assert!(futures::poll!(&mut second).is_pending());
180                assert_eq!(calls.load(Ordering::Relaxed), if coalesce { 1 } else { 2 });
181                gate.notify_waiters();
182                let (first, second) = tokio::join!(first, second);
183                assert_eq!(first.unwrap().as_ref(), second.unwrap().as_ref());
184            }
185        }
186    }
187
188    #[tokio::test]
189    async fn concurrent_loads_share_result_and_accounting() {
190        let key = key();
191        let (cache, gauge) = new_cache();
192        let calls = AtomicUsize::new(0);
193        let gate = tokio::sync::Notify::new();
194        let load = || async {
195            calls.fetch_add(1, Ordering::Relaxed);
196            gate.notified().await;
197            Ok(metadata())
198        };
199        let mut first = Box::pin(cache.try_get_with(&key, load()));
200        let mut second = Box::pin(cache.try_get_with(&key, load()));
201        assert!(futures::poll!(&mut first).is_pending());
202        assert!(futures::poll!(&mut second).is_pending());
203        assert_eq!(calls.load(Ordering::Relaxed), 1);
204        gate.notify_waiters();
205        let (first, second) = tokio::join!(first, second);
206        let first = first.unwrap();
207        assert!(Arc::ptr_eq(&first, &second.unwrap()));
208        cache.cache.run_pending_tasks().await;
209        assert_eq!(
210            sample_bytes(&gauge),
211            i64::from(puffin_metadata_weight(&key, &first))
212        );
213        let warm = cache
214            .try_get_with(&key, async { panic!("warm read loaded") })
215            .await
216            .unwrap();
217        assert!(Arc::ptr_eq(&first, &warm));
218        cache.remove(&key).await;
219        cache.cache.run_pending_tasks().await;
220        assert!(cache.get_metadata(&key).await.is_none());
221        assert_eq!(sample_bytes(&gauge), 0);
222    }
223
224    #[tokio::test]
225    async fn shared_failure_can_retry() {
226        let key = key();
227        let (cache, gauge) = new_cache();
228        let calls = AtomicUsize::new(0);
229        let gate = tokio::sync::Notify::new();
230        let load = || async {
231            calls.fetch_add(1, Ordering::Relaxed);
232            gate.notified().await;
233            InvalidPuffinFooterSnafu.fail()
234        };
235        let mut first = Box::pin(cache.try_get_with(&key, load()));
236        let mut second = Box::pin(cache.try_get_with(&key, load()));
237        assert!(futures::poll!(&mut first).is_pending());
238        assert!(futures::poll!(&mut second).is_pending());
239        gate.notify_waiters();
240        let (first, second) = tokio::join!(first, second);
241        assert!(first.is_err());
242        assert!(second.is_err());
243        assert_eq!(calls.load(Ordering::Relaxed), 1);
244        assert_eq!(sample_bytes(&gauge), 0);
245        assert!(
246            cache
247                .try_get_with(&key, async { Ok(metadata()) })
248                .await
249                .is_ok()
250        );
251    }
252
253    #[tokio::test]
254    async fn cancelled_loader_allows_waiter_to_load() {
255        let key = key();
256        let (cache, _) = new_cache();
257        let mut first = Box::pin(cache.try_get_with(&key, std::future::pending()));
258        assert!(futures::poll!(&mut first).is_pending());
259        let mut second = Box::pin(cache.try_get_with(&key, async { Ok(metadata()) }));
260        assert!(futures::poll!(&mut second).is_pending());
261        drop(first);
262        tokio::time::timeout(std::time::Duration::from_secs(5), second)
263            .await
264            .unwrap()
265            .unwrap();
266    }
267
268    #[tokio::test]
269    async fn same_handle_in_different_stores_does_not_share_footer() {
270        let (cache, _) = new_cache();
271        let first = key();
272        let second = PuffinFileKey::new(
273            Arc::new(crate::puffin_manager::file_key::PuffinStorageScope::new(
274                "file".into(),
275                String::new(),
276                "/other-root/".into(),
277            )),
278            "file".into(),
279            "file".into(),
280        );
281        let first_metadata = metadata();
282        cache
283            .put_metadata(first.clone(), first_metadata.clone())
284            .await;
285        assert!(cache.get_metadata(&second).await.is_none());
286        let second_metadata = cache
287            .try_get_with(&second, async {
288                Ok(Arc::new(FileMetadata {
289                    blobs: Vec::new(),
290                    properties: [("store".to_string(), "second".to_string())].into(),
291                }))
292            })
293            .await
294            .unwrap();
295        assert_ne!(first_metadata, second_metadata);
296        cache.remove(&first).await;
297        assert!(cache.get_metadata(&first).await.is_none());
298        assert_eq!(Some(second_metadata), cache.get_metadata(&second).await);
299        cache.invalidate_file("file").unwrap();
300        cache.cache.run_pending_tasks().await;
301        assert!(cache.get_metadata(&second).await.is_none());
302    }
303}
puffin/puffin_manager/cache.rs

puffin/puffin_manager/
cache.rs
