mz_persist_client/internal/
cache.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! In-process caches of [Blob].
11
12use std::sync::{Arc, Mutex};
13
14use async_trait::async_trait;
15use bytes::Bytes;
16use mz_dyncfg::{Config, ConfigSet, ParameterScope};
17use mz_ore::bytes::SegmentedBytes;
18use mz_ore::cast::CastFrom;
19use mz_persist::location::{Blob, BlobMetadata, ExternalError};
20
21use crate::cfg::PersistConfig;
22use crate::internal::metrics::Metrics;
23
24// In-memory cache for [Blob].
25#[derive(Debug)]
26pub struct BlobMemCache {
27    /// [`ConfigSet`] of dynamic configs.
28    cfg: Arc<ConfigSet>,
29    /// Number of 'workers' or threads in the current process, used for dynamic sizing.
30    num_workers: usize,
31    metrics: Arc<Metrics>,
32    cache: Mutex<lru::Lru<String, SegmentedBytes>>,
33    blob: Arc<dyn Blob>,
34}
35
36pub(crate) const BLOB_CACHE_MEM_LIMIT_BYTES: Config<usize> = Config::new(
37    "persist_blob_cache_mem_limit_bytes",
38    // 128MiB
39    128 * 1024 * 1024,
40    "Capacity of in-mem blob cache in bytes (Materialize).",
41    ParameterScope::Environment,
42);
43
44pub(crate) const BLOB_CACHE_SCALE_WITH_THREADS: Config<bool> = Config::new(
45    "persist_blob_cache_scale_with_threads",
46    false,
47    "Whether or not the size of the in-mem blob cache scales with the number of threads in the current process (Materialize).",
48    ParameterScope::Environment,
49);
50
51pub(crate) const BLOB_CACHE_SCALE_FACTOR_BYTES: Config<usize> = Config::new(
52    "persist_blob_cache_scale_factor_bytes",
53    // 32MiB
54    32 * 1024 * 1024,
55    "Scale factor for the in-mem blob cache, in bytes, if scaling with threads (Materialize).",
56    ParameterScope::Environment,
57);
58
59impl BlobMemCache {
60    pub fn new(cfg: &PersistConfig, metrics: Arc<Metrics>, blob: Arc<dyn Blob>) -> Arc<dyn Blob> {
61        let eviction_metrics = Arc::clone(&metrics);
62        let capacity_bytes =
63            BlobMemCache::get_capacity_bytes(&cfg.configs, cfg.isolated_runtime_worker_threads);
64        let cache = lru::Lru::new(capacity_bytes, move |_, _, _| {
65            eviction_metrics.blob_cache_mem.evictions.inc()
66        });
67        let blob = BlobMemCache {
68            cfg: Arc::clone(&cfg.configs),
69            num_workers: cfg.isolated_runtime_worker_threads,
70            metrics,
71            cache: Mutex::new(cache),
72            blob,
73        };
74        Arc::new(blob)
75    }
76
77    fn resize_and_update_size_metrics(&self, cache: &mut lru::Lru<String, SegmentedBytes>) {
78        let capacity_bytes = BlobMemCache::get_capacity_bytes(&self.cfg, self.num_workers);
79        cache.update_capacity(capacity_bytes);
80        self.metrics
81            .blob_cache_mem
82            .size_blobs
83            .set(u64::cast_from(cache.entry_count()));
84        self.metrics
85            .blob_cache_mem
86            .size_bytes
87            .set(u64::cast_from(cache.entry_weight()));
88    }
89
90    fn get_capacity_bytes(cfg: &Arc<ConfigSet>, num_workers: usize) -> usize {
91        // Note(parkmycar): To prevent regressing the size of the cache in
92        // small processes we use the static size as a minimum.
93        let static_size = BLOB_CACHE_MEM_LIMIT_BYTES.get(cfg);
94
95        if BLOB_CACHE_SCALE_WITH_THREADS.get(cfg) {
96            let per_thread_const = BLOB_CACHE_SCALE_FACTOR_BYTES.get(cfg);
97            let dynamic_size = num_workers.saturating_mul(per_thread_const);
98            std::cmp::max(dynamic_size, static_size)
99        } else {
100            static_size
101        }
102    }
103}
104
105#[async_trait]
106impl Blob for BlobMemCache {
107    async fn get(&self, key: &str) -> Result<Option<SegmentedBytes>, ExternalError> {
108        // First check if the blob is in the cache. If it is, return it. If not,
109        // fetch it and put it in the cache.
110        //
111        // Blobs are write-once modify-never, so we don't have to worry about
112        // any races or cache invalidations here. If the value is in the cache,
113        // any value in S3 is guaranteed to match (if not, then there's a
114        // horrible bug somewhere else).
115        if let Some((_, cached_value)) = self.cache.lock().expect("lock poisoned").get(key) {
116            self.metrics.blob_cache_mem.hits_blobs.inc();
117            self.metrics
118                .blob_cache_mem
119                .hits_bytes
120                .inc_by(u64::cast_from(cached_value.len()));
121            return Ok(Some(cached_value.clone()));
122        }
123
124        let res = self.blob.get(key).await?;
125        if let Some(blob) = res.as_ref() {
126            // TODO: It would likely be useful to allow a caller to opt out of
127            // adding the data to the cache (e.g. compaction inputs, perhaps
128            // some read handles).
129            let mut cache = self.cache.lock().expect("lock poisoned");
130            // If the weight of this single blob is greater than the capacity of
131            // the cache, it will push out everything in the cache and then
132            // immediately get evicted itself. So, skip adding it in that case.
133            if blob.len() <= cache.capacity() {
134                cache.insert(key.to_owned(), blob.clone(), blob.len());
135                self.resize_and_update_size_metrics(&mut cache);
136            }
137        }
138        Ok(res)
139    }
140
141    async fn list_keys_and_metadata(
142        &self,
143        key_prefix: &str,
144        f: &mut (dyn FnMut(BlobMetadata) + Send + Sync),
145    ) -> Result<(), ExternalError> {
146        self.blob.list_keys_and_metadata(key_prefix, f).await
147    }
148
149    async fn set(&self, key: &str, value: Bytes) -> Result<(), ExternalError> {
150        let () = self.blob.set(key, value.clone()).await?;
151        let weight = value.len();
152        let mut cache = self.cache.lock().expect("lock poisoned");
153        // If the weight of this single blob is greater than the capacity of
154        // the cache, it will push out everything in the cache and then
155        // immediately get evicted itself. So, skip adding it in that case.
156        if weight <= cache.capacity() {
157            cache.insert(key.to_owned(), SegmentedBytes::from(value), weight);
158            self.resize_and_update_size_metrics(&mut cache);
159        }
160        Ok(())
161    }
162
163    async fn delete(&self, key: &str) -> Result<Option<usize>, ExternalError> {
164        let res = self.blob.delete(key).await;
165        let mut cache = self.cache.lock().expect("lock poisoned");
166        cache.remove(key);
167        self.resize_and_update_size_metrics(&mut cache);
168        res
169    }
170
171    async fn restore(&self, key: &str) -> Result<(), ExternalError> {
172        self.blob.restore(key).await
173    }
174}
175
176mod lru {
177    use std::borrow::Borrow;
178    use std::collections::BTreeMap;
179    use std::hash::Hash;
180
181    use mz_ore::collections::HashMap;
182
183    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
184    pub struct Weight(usize);
185
186    #[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
187    pub struct Time(usize);
188
189    /// A weighted cache, evicting the least recently used entries.
190    ///
191    /// This is reimplemented here, instead of using an existing crate, because
192    /// existing options seem to either not support weights or they use unsafe.
193    pub struct Lru<K, V> {
194        evict_fn: Box<dyn Fn(K, V, usize) + Send>,
195        capacity: Weight,
196
197        next_time: Time,
198        entries: HashMap<K, (V, Weight, Time)>,
199        by_time: BTreeMap<Time, K>,
200        total_weight: Weight,
201    }
202
203    impl<K: Hash + Eq + Clone, V> Lru<K, V> {
204        /// Returns a new [Lru] with the requested configuration.
205        ///
206        /// `evict_fn` is called for every entry evicted by the least recently
207        /// used policy. It is not called for entries replaced by the same key
208        /// in `insert` nor entries explictly removed by `remove`.
209        pub fn new<F>(capacity: usize, evict_fn: F) -> Self
210        where
211            F: Fn(K, V, usize) + Send + 'static,
212        {
213            Lru {
214                evict_fn: Box::new(evict_fn),
215                capacity: Weight(capacity),
216                next_time: Time::default(),
217                entries: HashMap::new(),
218                by_time: BTreeMap::new(),
219                total_weight: Weight(0),
220            }
221        }
222
223        /// Returns the capacity of the cache.
224        pub fn capacity(&self) -> usize {
225            self.capacity.0
226        }
227
228        /// Returns the total number of entries in the cache.
229        pub fn entry_count(&self) -> usize {
230            mz_ore::soft_assert_eq_no_log!(self.entries.len(), self.by_time.len());
231            self.entries.len()
232        }
233
234        /// Returns the sum of weights of entries in the cache.
235        pub fn entry_weight(&self) -> usize {
236            self.total_weight.0
237        }
238
239        /// Changes the weighted capacity of the cache, evicting as necessary if
240        /// the new value is smaller.
241        pub fn update_capacity(&mut self, capacity: usize) {
242            self.capacity = Weight(capacity);
243            self.resize();
244            assert!(self.total_weight <= self.capacity);
245
246            // Intentionally not run with debug_assert, because validate is
247            // `O(n)` in the size of the cache.
248            #[cfg(test)]
249            self.validate();
250        }
251
252        /// Returns a reference to entry with the given key, if present, marking
253        /// it as most recently used.
254        pub fn get<Q>(&mut self, key: &Q) -> Option<(&K, &V)>
255        where
256            K: Borrow<Q>,
257            Q: Hash + Eq + ?Sized,
258        {
259            {
260                let (key, val, weight) = self.remove(key)?;
261                self.insert_not_exists(key, val, Weight(weight));
262            }
263            let (key, (val, _, _)) = self
264                .entries
265                .get_key_value(key)
266                .expect("internal lru invariant violated");
267
268            // Intentionally not run with debug_assert, because validate is
269            // `O(n)` in the size of the cache.
270            #[cfg(test)]
271            self.validate();
272
273            Some((key, val))
274        }
275
276        /// Inserts the given key and value into the cache, marking it as most
277        /// recently used.
278        ///
279        /// If the key already exists in the cache, the existing value and
280        /// weight are first removed.
281        pub fn insert(&mut self, key: K, val: V, weight: usize) {
282            let _ = self.remove(&key);
283            self.insert_not_exists(key, val, Weight(weight));
284
285            // Intentionally not run with debug_assert, because validate is
286            // `O(n)` in the size of the cache.
287            #[cfg(test)]
288            self.validate();
289        }
290
291        /// Removes the entry with the given key from the cache, if present.
292        ///
293        /// Returns None if the entry was not in the cache.
294        pub fn remove<Q>(&mut self, k: &Q) -> Option<(K, V, usize)>
295        where
296            K: Borrow<Q>,
297            Q: Hash + Eq + ?Sized,
298        {
299            let (_, _, time) = self.entries.get(k)?;
300            let (key, val, weight) = self.remove_exists(time.clone());
301
302            // Intentionally not run with debug_assert, because validate is
303            // `O(n)` in the size of the cache.
304            #[cfg(test)]
305            self.validate();
306
307            Some((key, val, weight.0))
308        }
309
310        /// Returns an iterator over the entries in the cache in order from most
311        /// recently used to least.
312        #[allow(dead_code)]
313        pub(crate) fn iter(&self) -> impl Iterator<Item = (&K, &V, usize)> {
314            self.by_time.iter().rev().map(|(_, key)| {
315                let (val, _, weight) = self
316                    .entries
317                    .get(key)
318                    .expect("internal lru invariant violated");
319                (key, val, weight.0)
320            })
321        }
322
323        fn insert_not_exists(&mut self, key: K, val: V, weight: Weight) {
324            let time = self.next_time.clone();
325            self.next_time.0 += 1;
326
327            self.total_weight.0 = self
328                .total_weight
329                .0
330                .checked_add(weight.0)
331                .expect("weight overflow");
332            assert!(
333                self.entries
334                    .insert(key.clone(), (val, weight, time.clone()))
335                    .is_none()
336            );
337            assert!(self.by_time.insert(time, key).is_none());
338            self.resize();
339        }
340
341        fn remove_exists(&mut self, time: Time) -> (K, V, Weight) {
342            let key = self
343                .by_time
344                .remove(&time)
345                .expect("internal list invariant violated");
346            let (val, weight, _time) = self
347                .entries
348                .remove(&key)
349                .expect("internal list invariant violated");
350            self.total_weight.0 = self
351                .total_weight
352                .0
353                .checked_sub(weight.0)
354                .expect("internal lru invariant violated");
355
356            (key, val, weight)
357        }
358
359        fn resize(&mut self) {
360            while self.total_weight > self.capacity {
361                let (time, _) = self
362                    .by_time
363                    .first_key_value()
364                    .expect("internal lru invariant violated");
365                let (key, val, weight) = self.remove_exists(time.clone());
366                (self.evict_fn)(key, val, weight.0);
367            }
368        }
369
370        /// Checks internal invariants.
371        ///
372        /// TODO: Give this persist's usual `-> Result<(), String>` signature
373        /// instead of panic-ing.
374        #[cfg(test)]
375        pub(crate) fn validate(&self) {
376            assert!(self.total_weight <= self.capacity);
377
378            let mut count = 0;
379            let mut weight = 0;
380            for (time, key) in self.by_time.iter() {
381                let (_val, w, t) = self
382                    .entries
383                    .get(key)
384                    .expect("internal lru invariant violated");
385                count += 1;
386                weight += w.0;
387                assert_eq!(time, t);
388            }
389            assert_eq!(count, self.by_time.len());
390            assert_eq!(weight, self.total_weight.0);
391        }
392    }
393
394    impl<K: std::fmt::Debug, V> std::fmt::Debug for Lru<K, V> {
395        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
396            let Lru {
397                evict_fn: _,
398                capacity,
399                next_time,
400                entries: _,
401                by_time,
402                total_weight,
403            } = self;
404            f.debug_struct("Lru")
405                .field("capacity", &capacity)
406                .field("total_weight", &total_weight)
407                .field("next_time", &next_time)
408                .field("by_time", &by_time)
409                .finish_non_exhaustive()
410        }
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use mz_ore::assert_none;
417    use proptest::arbitrary::any;
418    use proptest::proptest;
419    use proptest_derive::Arbitrary;
420
421    use super::lru::*;
422
423    #[derive(Debug, Arbitrary)]
424    enum LruOp {
425        Get { key: u8 },
426        Insert { key: u8, weight: u8 },
427        Remove { key: u8 },
428    }
429
430    fn prop_testcase(ops: Vec<LruOp>) {
431        // In the long run, we'd expect maybe 1/2s of the `u8::MAX` possible
432        // keys to be present and for the average weight to be `u8::MAX / 2`.
433        // Select a capacity that is somewhat less than that.
434        let capacity = usize::from(u8::MAX / 2) * usize::from(u8::MAX / 2) / 2;
435        let mut cache = Lru::new(capacity, |_, _, _| {});
436        for op in ops {
437            match op {
438                LruOp::Get { key } => {
439                    let _ = cache.get(&key);
440                }
441                LruOp::Insert { key, weight } => {
442                    cache.insert(key, (), usize::from(weight));
443                }
444                LruOp::Remove { key } => {
445                    let _ = cache.remove(&key);
446                }
447            }
448            cache.validate();
449        }
450    }
451
452    #[mz_ore::test]
453    #[cfg_attr(miri, ignore)] // too slow
454    fn lru_cache_prop() {
455        proptest!(|(state in proptest::collection::vec(any::<LruOp>(), 0..100))| prop_testcase(state));
456    }
457
458    impl Lru<&'static str, ()> {
459        fn keys(&self) -> Vec<&'static str> {
460            self.iter().map(|(k, _, _)| *k).collect()
461        }
462    }
463
464    #[mz_ore::test]
465    #[cfg_attr(miri, ignore)]
466    fn lru_cache_usage() {
467        let mut cache = Lru::<&'static str, ()>::new(3, |_, _, _| {});
468
469        // Empty
470        assert_eq!(cache.entry_count(), 0);
471        assert_eq!(cache.entry_weight(), 0);
472
473        // Insert into empty.
474        cache.insert("a", (), 2);
475        assert_eq!(cache.entry_count(), 1);
476        assert_eq!(cache.entry_weight(), 2);
477        assert_eq!(cache.keys(), &["a"]);
478
479        // Insert and push out previous.
480        cache.insert("b", (), 2);
481        assert_eq!(cache.entry_count(), 1);
482        assert_eq!(cache.entry_weight(), 2);
483        assert_eq!(cache.keys(), &["b"]);
484
485        // Insert and don't push out previous.
486        cache.insert("c", (), 1);
487        assert_eq!(cache.entry_count(), 2);
488        assert_eq!(cache.entry_weight(), 3);
489        assert_eq!(cache.keys(), &["c", "b"]);
490
491        // More than two elements.
492        cache.insert("d", (), 1);
493        cache.insert("e", (), 1);
494        assert_eq!(cache.entry_count(), 3);
495        assert_eq!(cache.entry_weight(), 3);
496        assert_eq!(cache.keys(), &["e", "d", "c"]);
497
498        // Get the head.
499        cache.get("e");
500        assert_eq!(cache.entry_count(), 3);
501        assert_eq!(cache.entry_weight(), 3);
502        assert_eq!(cache.keys(), &["e", "d", "c"]);
503
504        // Get the tail.
505        cache.get("c");
506        assert_eq!(cache.entry_count(), 3);
507        assert_eq!(cache.entry_weight(), 3);
508        assert_eq!(cache.keys(), &["c", "e", "d"]);
509
510        // Get the mid.
511        cache.get("e");
512        assert_eq!(cache.entry_count(), 3);
513        assert_eq!(cache.entry_weight(), 3);
514        assert_eq!(cache.keys(), &["e", "c", "d"]);
515
516        // Get a non-existent element.
517        cache.get("f");
518        assert_eq!(cache.entry_count(), 3);
519        assert_eq!(cache.entry_weight(), 3);
520        assert_eq!(cache.keys(), &["e", "c", "d"]);
521
522        // Remove an element.
523        assert!(cache.remove("c").is_some());
524        assert_eq!(cache.entry_count(), 2);
525        assert_eq!(cache.entry_weight(), 2);
526        assert_eq!(cache.keys(), &["e", "d"]);
527
528        // Remove a non-existent element.
529        assert_none!(cache.remove("f"));
530        assert_eq!(cache.entry_count(), 2);
531        assert_eq!(cache.entry_weight(), 2);
532        assert_eq!(cache.keys(), &["e", "d"]);
533
534        // Push out everything with a big weight
535        cache.insert("f", (), 3);
536        assert_eq!(cache.entry_count(), 1);
537        assert_eq!(cache.entry_weight(), 3);
538        assert_eq!(cache.keys(), &["f"]);
539
540        // Push out everything with a weight so big it doesn't even fit. (Is
541        // this even the behavior we want?)
542        cache.insert("g", (), 4);
543        assert_eq!(cache.entry_count(), 0);
544        assert_eq!(cache.entry_weight(), 0);
545
546        // Resize up
547        cache.insert("h", (), 2);
548        cache.insert("i", (), 1);
549        cache.update_capacity(4);
550        cache.insert("j", (), 1);
551        assert_eq!(cache.entry_count(), 3);
552        assert_eq!(cache.entry_weight(), 4);
553        assert_eq!(cache.keys(), &["j", "i", "h"]);
554
555        // Resize down
556        cache.update_capacity(2);
557        assert_eq!(cache.entry_count(), 2);
558        assert_eq!(cache.entry_weight(), 2);
559        assert_eq!(cache.keys(), &["j", "i"]);
560    }
561}
mz_persist_client/internal/cache.rs

mz_persist_client/internal/
cache.rs
