linera_views/backends/
lru_caching.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Add LRU (least recently used) caching to a given store.
5
6use std::sync::{Arc, Mutex};
7
8use serde::{Deserialize, Serialize};
9
10#[cfg(with_testing)]
11use crate::memory::MemoryDatabase;
12#[cfg(with_testing)]
13use crate::store::TestKeyValueDatabase;
14use crate::{
15    batch::{Batch, WriteOperation},
16    lru_prefix_cache::{LruPrefixCache, StorageCacheConfig},
17    store::{KeyValueDatabase, ReadableKeyValueStore, WithError, WritableKeyValueStore},
18};
19
20#[cfg(with_metrics)]
21mod metrics {
22    use std::sync::LazyLock;
23
24    use linera_base::prometheus_util::register_int_counter_vec;
25    use prometheus::IntCounterVec;
26
27    /// The total number of cache read value misses.
28    pub static READ_VALUE_CACHE_MISS_COUNT: LazyLock<IntCounterVec> = LazyLock::new(|| {
29        register_int_counter_vec(
30            "num_read_value_cache_miss",
31            "Number of read value cache misses",
32            &[],
33        )
34    });
35
36    /// The total number of read value cache hits.
37    pub static READ_VALUE_CACHE_HIT_COUNT: LazyLock<IntCounterVec> = LazyLock::new(|| {
38        register_int_counter_vec(
39            "num_read_value_cache_hits",
40            "Number of read value cache hits",
41            &[],
42        )
43    });
44
45    /// The total number of contains key cache misses.
46    pub static CONTAINS_KEY_CACHE_MISS_COUNT: LazyLock<IntCounterVec> = LazyLock::new(|| {
47        register_int_counter_vec(
48            "num_contains_key_cache_miss",
49            "Number of contains key cache misses",
50            &[],
51        )
52    });
53
54    /// The total number of contains key cache hits.
55    pub static CONTAINS_KEY_CACHE_HIT_COUNT: LazyLock<IntCounterVec> = LazyLock::new(|| {
56        register_int_counter_vec(
57            "num_contains_key_cache_hit",
58            "Number of contains key cache hits",
59            &[],
60        )
61    });
62
63    /// The total number of find_keys_by_prefix cache misses.
64    pub static FIND_KEYS_BY_PREFIX_CACHE_MISS_COUNT: LazyLock<IntCounterVec> =
65        LazyLock::new(|| {
66            register_int_counter_vec(
67                "num_find_keys_by_prefix_cache_miss",
68                "Number of find keys by prefix cache misses",
69                &[],
70            )
71        });
72
73    /// The total number of find_keys_by_prefix cache hits.
74    pub static FIND_KEYS_BY_PREFIX_CACHE_HIT_COUNT: LazyLock<IntCounterVec> = LazyLock::new(|| {
75        register_int_counter_vec(
76            "num_find_keys_by_prefix_cache_hit",
77            "Number of find keys by prefix cache hits",
78            &[],
79        )
80    });
81
82    /// The total number of find_key_values_by_prefix cache misses.
83    pub static FIND_KEY_VALUES_BY_PREFIX_CACHE_MISS_COUNT: LazyLock<IntCounterVec> =
84        LazyLock::new(|| {
85            register_int_counter_vec(
86                "num_find_key_values_by_prefix_cache_miss",
87                "Number of find key values by prefix cache misses",
88                &[],
89            )
90        });
91
92    /// The total number of find_key_values_by_prefix cache hits.
93    pub static FIND_KEY_VALUES_BY_PREFIX_CACHE_HIT_COUNT: LazyLock<IntCounterVec> =
94        LazyLock::new(|| {
95            register_int_counter_vec(
96                "num_find_key_values_by_prefix_cache_hit",
97                "Number of find key values by prefix cache hits",
98                &[],
99            )
100        });
101}
102
103/// The maximum number of entries in the cache.
104/// If the number of entries in the cache is too large then the underlying maps
105/// become the limiting factor.
106pub const DEFAULT_STORAGE_CACHE_CONFIG: StorageCacheConfig = StorageCacheConfig {
107    max_cache_size: 10000000,
108    max_value_entry_size: 1000000,
109    max_find_keys_entry_size: 1000000,
110    max_find_key_values_entry_size: 1000000,
111    max_cache_entries: 1000,
112    max_cache_value_size: 10000000,
113    max_cache_find_keys_size: 10000000,
114    max_cache_find_key_values_size: 10000000,
115};
116
117/// A key-value database with added LRU caching.
118#[derive(Clone)]
119pub struct LruCachingDatabase<D> {
120    /// The inner store that is called by the LRU cache one.
121    database: D,
122    /// The configuration.
123    config: StorageCacheConfig,
124}
125
126/// A key-value store with added LRU caching.
127#[derive(Clone)]
128pub struct LruCachingStore<S> {
129    /// The inner store that is called by the LRU cache one.
130    store: S,
131    /// The LRU cache of values.
132    cache: Option<Arc<Mutex<LruPrefixCache>>>,
133}
134
135impl<D> WithError for LruCachingDatabase<D>
136where
137    D: WithError,
138{
139    type Error = D::Error;
140}
141
142impl<S> WithError for LruCachingStore<S>
143where
144    S: WithError,
145{
146    type Error = S::Error;
147}
148
149impl<K> ReadableKeyValueStore for LruCachingStore<K>
150where
151    K: ReadableKeyValueStore,
152{
153    // The LRU cache does not change the underlying store's size limits.
154    const MAX_KEY_SIZE: usize = K::MAX_KEY_SIZE;
155
156    fn max_stream_queries(&self) -> usize {
157        self.store.max_stream_queries()
158    }
159
160    fn root_key(&self) -> Result<Vec<u8>, Self::Error> {
161        self.store.root_key()
162    }
163
164    async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
165        let Some(cache) = &self.cache else {
166            return self.store.read_value_bytes(key).await;
167        };
168        // First inquiring in the read_value_bytes LRU
169        {
170            let mut cache = cache.lock().unwrap();
171            if let Some(value) = cache.query_read_value(key) {
172                #[cfg(with_metrics)]
173                metrics::READ_VALUE_CACHE_HIT_COUNT
174                    .with_label_values(&[])
175                    .inc();
176                return Ok(value);
177            }
178        }
179        #[cfg(with_metrics)]
180        metrics::READ_VALUE_CACHE_MISS_COUNT
181            .with_label_values(&[])
182            .inc();
183        let value = self.store.read_value_bytes(key).await?;
184        let mut cache = cache.lock().unwrap();
185        cache.insert_read_value(key, &value);
186        Ok(value)
187    }
188
189    async fn contains_key(&self, key: &[u8]) -> Result<bool, Self::Error> {
190        let Some(cache) = &self.cache else {
191            return self.store.contains_key(key).await;
192        };
193        {
194            let mut cache = cache.lock().unwrap();
195            if let Some(value) = cache.query_contains_key(key) {
196                #[cfg(with_metrics)]
197                metrics::CONTAINS_KEY_CACHE_HIT_COUNT
198                    .with_label_values(&[])
199                    .inc();
200                return Ok(value);
201            }
202        }
203        #[cfg(with_metrics)]
204        metrics::CONTAINS_KEY_CACHE_MISS_COUNT
205            .with_label_values(&[])
206            .inc();
207        let result = self.store.contains_key(key).await?;
208        let mut cache = cache.lock().unwrap();
209        cache.insert_contains_key(key, result);
210        Ok(result)
211    }
212
213    async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, Self::Error> {
214        let Some(cache) = &self.cache else {
215            return self.store.contains_keys(keys).await;
216        };
217        let size = keys.len();
218        let mut results = vec![false; size];
219        let mut indices = Vec::new();
220        let mut key_requests = Vec::new();
221        {
222            let mut cache = cache.lock().unwrap();
223            for i in 0..size {
224                if let Some(value) = cache.query_contains_key(&keys[i]) {
225                    #[cfg(with_metrics)]
226                    metrics::CONTAINS_KEY_CACHE_HIT_COUNT
227                        .with_label_values(&[])
228                        .inc();
229                    results[i] = value;
230                } else {
231                    #[cfg(with_metrics)]
232                    metrics::CONTAINS_KEY_CACHE_MISS_COUNT
233                        .with_label_values(&[])
234                        .inc();
235                    indices.push(i);
236                    key_requests.push(keys[i].clone());
237                }
238            }
239        }
240        if !key_requests.is_empty() {
241            let key_results = self.store.contains_keys(&key_requests).await?;
242            let mut cache = cache.lock().unwrap();
243            for ((index, result), key) in indices.into_iter().zip(key_results).zip(key_requests) {
244                results[index] = result;
245                cache.insert_contains_key(&key, result);
246            }
247        }
248        Ok(results)
249    }
250
251    async fn read_multi_values_bytes(
252        &self,
253        keys: &[Vec<u8>],
254    ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
255        let Some(cache) = &self.cache else {
256            return self.store.read_multi_values_bytes(keys).await;
257        };
258
259        let mut result = Vec::with_capacity(keys.len());
260        let mut cache_miss_indices = Vec::new();
261        let mut miss_keys = Vec::new();
262        {
263            let mut cache = cache.lock().unwrap();
264            for (i, key) in keys.iter().enumerate() {
265                if let Some(value) = cache.query_read_value(key) {
266                    #[cfg(with_metrics)]
267                    metrics::READ_VALUE_CACHE_HIT_COUNT
268                        .with_label_values(&[])
269                        .inc();
270                    result.push(value);
271                } else {
272                    #[cfg(with_metrics)]
273                    metrics::READ_VALUE_CACHE_MISS_COUNT
274                        .with_label_values(&[])
275                        .inc();
276                    result.push(None);
277                    cache_miss_indices.push(i);
278                    miss_keys.push(key.clone());
279                }
280            }
281        }
282        if !miss_keys.is_empty() {
283            let values = self.store.read_multi_values_bytes(&miss_keys).await?;
284            let mut cache = cache.lock().unwrap();
285            for (i, (key, value)) in cache_miss_indices
286                .into_iter()
287                .zip(miss_keys.into_iter().zip(values))
288            {
289                cache.insert_read_value(&key, &value);
290                result[i] = value;
291            }
292        }
293        Ok(result)
294    }
295
296    async fn find_keys_by_prefix(&self, key_prefix: &[u8]) -> Result<Vec<Vec<u8>>, Self::Error> {
297        let Some(cache) = self.get_exclusive_cache() else {
298            return self.store.find_keys_by_prefix(key_prefix).await;
299        };
300        {
301            let mut cache = cache.lock().unwrap();
302            if let Some(value) = cache.query_find_keys(key_prefix) {
303                #[cfg(with_metrics)]
304                metrics::FIND_KEYS_BY_PREFIX_CACHE_HIT_COUNT
305                    .with_label_values(&[])
306                    .inc();
307                return Ok(value);
308            }
309        }
310        #[cfg(with_metrics)]
311        metrics::FIND_KEYS_BY_PREFIX_CACHE_MISS_COUNT
312            .with_label_values(&[])
313            .inc();
314        let keys = self.store.find_keys_by_prefix(key_prefix).await?;
315        let mut cache = cache.lock().unwrap();
316        cache.insert_find_keys(key_prefix.to_vec(), &keys);
317        Ok(keys)
318    }
319
320    async fn find_key_values_by_prefix(
321        &self,
322        key_prefix: &[u8],
323    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Self::Error> {
324        let Some(cache) = self.get_exclusive_cache() else {
325            return self.store.find_key_values_by_prefix(key_prefix).await;
326        };
327        {
328            let mut cache = cache.lock().unwrap();
329            if let Some(value) = cache.query_find_key_values(key_prefix) {
330                #[cfg(with_metrics)]
331                metrics::FIND_KEY_VALUES_BY_PREFIX_CACHE_HIT_COUNT
332                    .with_label_values(&[])
333                    .inc();
334                return Ok(value);
335            }
336        }
337        #[cfg(with_metrics)]
338        metrics::FIND_KEY_VALUES_BY_PREFIX_CACHE_MISS_COUNT
339            .with_label_values(&[])
340            .inc();
341        let key_values = self.store.find_key_values_by_prefix(key_prefix).await?;
342        let mut cache = cache.lock().unwrap();
343        cache.insert_find_key_values(key_prefix.to_vec(), &key_values);
344        Ok(key_values)
345    }
346}
347
348impl<K> WritableKeyValueStore for LruCachingStore<K>
349where
350    K: WritableKeyValueStore,
351{
352    // The LRU cache does not change the underlying store's size limits.
353    const MAX_VALUE_SIZE: usize = K::MAX_VALUE_SIZE;
354
355    async fn write_batch(&self, batch: Batch) -> Result<(), Self::Error> {
356        self.store.write_batch(batch.clone()).await?;
357        if let Some(cache) = &self.cache {
358            let mut cache = cache.lock().unwrap();
359            for operation in &batch.operations {
360                match operation {
361                    WriteOperation::Put { key, value } => {
362                        cache.put_key_value(key, value);
363                    }
364                    WriteOperation::Delete { key } => {
365                        cache.delete_key(key);
366                    }
367                    WriteOperation::DeletePrefix { key_prefix } => {
368                        cache.delete_prefix(key_prefix);
369                    }
370                }
371            }
372        }
373        Ok(())
374    }
375
376    async fn clear_journal(&self) -> Result<(), Self::Error> {
377        self.store.clear_journal().await
378    }
379}
380
381/// The configuration type for the `LruCachingStore`.
382#[derive(Debug, Clone, Serialize, Deserialize)]
383pub struct LruCachingConfig<C> {
384    /// The inner configuration of the `LruCachingStore`.
385    pub inner_config: C,
386    /// The cache size being used.
387    pub storage_cache_config: StorageCacheConfig,
388}
389
390impl<D> KeyValueDatabase for LruCachingDatabase<D>
391where
392    D: KeyValueDatabase,
393{
394    type Config = LruCachingConfig<D::Config>;
395
396    type Store = LruCachingStore<D::Store>;
397
398    fn get_name() -> String {
399        format!("lru caching {}", D::get_name())
400    }
401
402    async fn connect(config: &Self::Config, namespace: &str) -> Result<Self, Self::Error> {
403        let database = D::connect(&config.inner_config, namespace).await?;
404        Ok(LruCachingDatabase {
405            database,
406            config: config.storage_cache_config.clone(),
407        })
408    }
409
410    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
411        let store = self.database.open_shared(root_key)?;
412        let store = LruCachingStore::new(
413            store,
414            self.config.clone(),
415            /* has_exclusive_access */ false,
416        );
417        Ok(store)
418    }
419
420    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
421        let store = self.database.open_exclusive(root_key)?;
422        let store = LruCachingStore::new(
423            store,
424            self.config.clone(),
425            /* has_exclusive_access */ true,
426        );
427        Ok(store)
428    }
429
430    async fn list_all(config: &Self::Config) -> Result<Vec<String>, Self::Error> {
431        D::list_all(&config.inner_config).await
432    }
433
434    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
435        self.database.list_root_keys().await
436    }
437
438    async fn delete_all(config: &Self::Config) -> Result<(), Self::Error> {
439        D::delete_all(&config.inner_config).await
440    }
441
442    async fn exists(config: &Self::Config, namespace: &str) -> Result<bool, Self::Error> {
443        D::exists(&config.inner_config, namespace).await
444    }
445
446    async fn create(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
447        D::create(&config.inner_config, namespace).await
448    }
449
450    async fn delete(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
451        D::delete(&config.inner_config, namespace).await
452    }
453}
454
455impl<S> LruCachingStore<S> {
456    /// Creates a new key-value store that provides LRU caching at top of the given store.
457    fn new(store: S, config: StorageCacheConfig, has_exclusive_access: bool) -> Self {
458        let cache = {
459            if config.max_cache_entries == 0 {
460                None
461            } else {
462                Some(Arc::new(Mutex::new(LruPrefixCache::new(
463                    config,
464                    has_exclusive_access,
465                ))))
466            }
467        };
468        Self { store, cache }
469    }
470
471    /// Returns a cache with exclusive access if one exists.
472    fn get_exclusive_cache(&self) -> Option<&Arc<Mutex<LruPrefixCache>>> {
473        let Some(cache) = &self.cache else {
474            return None;
475        };
476        let has_exclusive_access = {
477            let cache = cache.lock().unwrap();
478            cache.has_exclusive_access()
479        };
480        if has_exclusive_access {
481            Some(cache)
482        } else {
483            None
484        }
485    }
486}
487
488/// A memory database with caching.
489#[cfg(with_testing)]
490pub type LruCachingMemoryDatabase = LruCachingDatabase<MemoryDatabase>;
491
492#[cfg(with_testing)]
493impl<D> TestKeyValueDatabase for LruCachingDatabase<D>
494where
495    D: TestKeyValueDatabase,
496{
497    async fn new_test_config() -> Result<LruCachingConfig<D::Config>, D::Error> {
498        let inner_config = D::new_test_config().await?;
499        let storage_cache_config = DEFAULT_STORAGE_CACHE_CONFIG;
500        Ok(LruCachingConfig {
501            inner_config,
502            storage_cache_config,
503        })
504    }
505}
linera_views/backends/lru_caching.rs

linera_views/backends/
lru_caching.rs
