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 root_key(&self) -> Result<Vec<u8>, Self::Error> {
157        self.store.root_key()
158    }
159
160    async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
161        let Some(cache) = &self.cache else {
162            return self.store.read_value_bytes(key).await;
163        };
164        // First inquiring in the read_value_bytes LRU
165        {
166            let mut cache = cache.lock().unwrap();
167            if let Some(value) = cache.query_read_value(key) {
168                #[cfg(with_metrics)]
169                metrics::READ_VALUE_CACHE_HIT_COUNT
170                    .with_label_values(&[])
171                    .inc();
172                return Ok(value);
173            }
174        }
175        #[cfg(with_metrics)]
176        metrics::READ_VALUE_CACHE_MISS_COUNT
177            .with_label_values(&[])
178            .inc();
179        let value = self.store.read_value_bytes(key).await?;
180        let mut cache = cache.lock().unwrap();
181        cache.insert_read_value(key, &value);
182        Ok(value)
183    }
184
185    async fn contains_key(&self, key: &[u8]) -> Result<bool, Self::Error> {
186        let Some(cache) = &self.cache else {
187            return self.store.contains_key(key).await;
188        };
189        {
190            let mut cache = cache.lock().unwrap();
191            if let Some(value) = cache.query_contains_key(key) {
192                #[cfg(with_metrics)]
193                metrics::CONTAINS_KEY_CACHE_HIT_COUNT
194                    .with_label_values(&[])
195                    .inc();
196                return Ok(value);
197            }
198        }
199        #[cfg(with_metrics)]
200        metrics::CONTAINS_KEY_CACHE_MISS_COUNT
201            .with_label_values(&[])
202            .inc();
203        let result = self.store.contains_key(key).await?;
204        let mut cache = cache.lock().unwrap();
205        cache.insert_contains_key(key, result);
206        Ok(result)
207    }
208
209    async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, Self::Error> {
210        let Some(cache) = &self.cache else {
211            return self.store.contains_keys(keys).await;
212        };
213        let size = keys.len();
214        let mut results = vec![false; size];
215        let mut indices = Vec::new();
216        let mut key_requests = Vec::new();
217        {
218            let mut cache = cache.lock().unwrap();
219            for i in 0..size {
220                if let Some(value) = cache.query_contains_key(&keys[i]) {
221                    #[cfg(with_metrics)]
222                    metrics::CONTAINS_KEY_CACHE_HIT_COUNT
223                        .with_label_values(&[])
224                        .inc();
225                    results[i] = value;
226                } else {
227                    #[cfg(with_metrics)]
228                    metrics::CONTAINS_KEY_CACHE_MISS_COUNT
229                        .with_label_values(&[])
230                        .inc();
231                    indices.push(i);
232                    key_requests.push(keys[i].clone());
233                }
234            }
235        }
236        if !key_requests.is_empty() {
237            let key_results = self.store.contains_keys(&key_requests).await?;
238            let mut cache = cache.lock().unwrap();
239            for ((index, result), key) in indices.into_iter().zip(key_results).zip(key_requests) {
240                results[index] = result;
241                cache.insert_contains_key(&key, result);
242            }
243        }
244        Ok(results)
245    }
246
247    async fn read_multi_values_bytes(
248        &self,
249        keys: &[Vec<u8>],
250    ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
251        let Some(cache) = &self.cache else {
252            return self.store.read_multi_values_bytes(keys).await;
253        };
254
255        let mut result = Vec::with_capacity(keys.len());
256        let mut cache_miss_indices = Vec::new();
257        let mut miss_keys = Vec::new();
258        {
259            let mut cache = cache.lock().unwrap();
260            for (i, key) in keys.iter().enumerate() {
261                if let Some(value) = cache.query_read_value(key) {
262                    #[cfg(with_metrics)]
263                    metrics::READ_VALUE_CACHE_HIT_COUNT
264                        .with_label_values(&[])
265                        .inc();
266                    result.push(value);
267                } else {
268                    #[cfg(with_metrics)]
269                    metrics::READ_VALUE_CACHE_MISS_COUNT
270                        .with_label_values(&[])
271                        .inc();
272                    result.push(None);
273                    cache_miss_indices.push(i);
274                    miss_keys.push(key.clone());
275                }
276            }
277        }
278        if !miss_keys.is_empty() {
279            let values = self.store.read_multi_values_bytes(&miss_keys).await?;
280            let mut cache = cache.lock().unwrap();
281            for (i, (key, value)) in cache_miss_indices
282                .into_iter()
283                .zip(miss_keys.into_iter().zip(values))
284            {
285                cache.insert_read_value(&key, &value);
286                result[i] = value;
287            }
288        }
289        Ok(result)
290    }
291
292    async fn find_keys_by_prefix(&self, key_prefix: &[u8]) -> Result<Vec<Vec<u8>>, Self::Error> {
293        let Some(cache) = self.get_exclusive_cache() else {
294            return self.store.find_keys_by_prefix(key_prefix).await;
295        };
296        {
297            let mut cache = cache.lock().unwrap();
298            if let Some(value) = cache.query_find_keys(key_prefix) {
299                #[cfg(with_metrics)]
300                metrics::FIND_KEYS_BY_PREFIX_CACHE_HIT_COUNT
301                    .with_label_values(&[])
302                    .inc();
303                return Ok(value);
304            }
305        }
306        #[cfg(with_metrics)]
307        metrics::FIND_KEYS_BY_PREFIX_CACHE_MISS_COUNT
308            .with_label_values(&[])
309            .inc();
310        let keys = self.store.find_keys_by_prefix(key_prefix).await?;
311        let mut cache = cache.lock().unwrap();
312        cache.insert_find_keys(key_prefix.to_vec(), &keys);
313        Ok(keys)
314    }
315
316    async fn find_key_values_by_prefix(
317        &self,
318        key_prefix: &[u8],
319    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Self::Error> {
320        let Some(cache) = self.get_exclusive_cache() else {
321            return self.store.find_key_values_by_prefix(key_prefix).await;
322        };
323        {
324            let mut cache = cache.lock().unwrap();
325            if let Some(value) = cache.query_find_key_values(key_prefix) {
326                #[cfg(with_metrics)]
327                metrics::FIND_KEY_VALUES_BY_PREFIX_CACHE_HIT_COUNT
328                    .with_label_values(&[])
329                    .inc();
330                return Ok(value);
331            }
332        }
333        #[cfg(with_metrics)]
334        metrics::FIND_KEY_VALUES_BY_PREFIX_CACHE_MISS_COUNT
335            .with_label_values(&[])
336            .inc();
337        let key_values = self.store.find_key_values_by_prefix(key_prefix).await?;
338        let mut cache = cache.lock().unwrap();
339        cache.insert_find_key_values(key_prefix.to_vec(), &key_values);
340        Ok(key_values)
341    }
342}
343
344impl<K> WritableKeyValueStore for LruCachingStore<K>
345where
346    K: WritableKeyValueStore,
347{
348    // The LRU cache does not change the underlying store's size limits.
349    const MAX_VALUE_SIZE: usize = K::MAX_VALUE_SIZE;
350
351    async fn write_batch(&self, batch: Batch) -> Result<(), Self::Error> {
352        self.store.write_batch(batch.clone()).await?;
353        if let Some(cache) = &self.cache {
354            let mut cache = cache.lock().unwrap();
355            for operation in &batch.operations {
356                match operation {
357                    WriteOperation::Put { key, value } => {
358                        cache.put_key_value(key, value);
359                    }
360                    WriteOperation::Delete { key } => {
361                        cache.delete_key(key);
362                    }
363                    WriteOperation::DeletePrefix { key_prefix } => {
364                        cache.delete_prefix(key_prefix);
365                    }
366                }
367            }
368        }
369        Ok(())
370    }
371
372    async fn clear_journal(&self) -> Result<(), Self::Error> {
373        self.store.clear_journal().await
374    }
375}
376
377/// The configuration type for the `LruCachingStore`.
378#[derive(Debug, Clone, Serialize, Deserialize)]
379pub struct LruCachingConfig<C> {
380    /// The inner configuration of the `LruCachingStore`.
381    pub inner_config: C,
382    /// The cache size being used.
383    pub storage_cache_config: StorageCacheConfig,
384}
385
386impl<D> KeyValueDatabase for LruCachingDatabase<D>
387where
388    D: KeyValueDatabase,
389{
390    type Config = LruCachingConfig<D::Config>;
391
392    type Store = LruCachingStore<D::Store>;
393
394    fn get_name() -> String {
395        format!("lru caching {}", D::get_name())
396    }
397
398    async fn connect(config: &Self::Config, namespace: &str) -> Result<Self, Self::Error> {
399        let database = D::connect(&config.inner_config, namespace).await?;
400        Ok(LruCachingDatabase {
401            database,
402            config: config.storage_cache_config.clone(),
403        })
404    }
405
406    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
407        let store = self.database.open_shared(root_key)?;
408        // Caching for immutable data is handled in DbStorage.
409        Ok(LruCachingStore { store, cache: None })
410    }
411
412    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
413        let store = self.database.open_exclusive(root_key)?;
414        let store = LruCachingStore::new(
415            store,
416            self.config.clone(),
417            /* has_exclusive_access */ true,
418        );
419        Ok(store)
420    }
421
422    async fn list_all(config: &Self::Config) -> Result<Vec<String>, Self::Error> {
423        D::list_all(&config.inner_config).await
424    }
425
426    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
427        self.database.list_root_keys().await
428    }
429
430    async fn delete_all(config: &Self::Config) -> Result<(), Self::Error> {
431        D::delete_all(&config.inner_config).await
432    }
433
434    async fn exists(config: &Self::Config, namespace: &str) -> Result<bool, Self::Error> {
435        D::exists(&config.inner_config, namespace).await
436    }
437
438    async fn create(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
439        D::create(&config.inner_config, namespace).await
440    }
441
442    async fn delete(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
443        D::delete(&config.inner_config, namespace).await
444    }
445}
446
447impl<S> LruCachingStore<S> {
448    /// Creates a new key-value store that provides LRU caching at top of the given store.
449    fn new(store: S, config: StorageCacheConfig, has_exclusive_access: bool) -> Self {
450        let cache = {
451            if config.max_cache_entries == 0 {
452                None
453            } else {
454                Some(Arc::new(Mutex::new(LruPrefixCache::new(
455                    config,
456                    has_exclusive_access,
457                ))))
458            }
459        };
460        Self { store, cache }
461    }
462
463    /// Returns a cache with exclusive access if one exists.
464    fn get_exclusive_cache(&self) -> Option<&Arc<Mutex<LruPrefixCache>>> {
465        let Some(cache) = &self.cache else {
466            return None;
467        };
468        let has_exclusive_access = {
469            let cache = cache.lock().unwrap();
470            cache.has_exclusive_access()
471        };
472        if has_exclusive_access {
473            Some(cache)
474        } else {
475            None
476        }
477    }
478}
479
480/// A memory database with caching.
481#[cfg(with_testing)]
482pub type LruCachingMemoryDatabase = LruCachingDatabase<MemoryDatabase>;
483
484#[cfg(with_testing)]
485impl<D> TestKeyValueDatabase for LruCachingDatabase<D>
486where
487    D: TestKeyValueDatabase,
488{
489    async fn new_test_config() -> Result<LruCachingConfig<D::Config>, D::Error> {
490        let inner_config = D::new_test_config().await?;
491        let storage_cache_config = DEFAULT_STORAGE_CACHE_CONFIG;
492        Ok(LruCachingConfig {
493            inner_config,
494            storage_cache_config,
495        })
496    }
497}
498
499#[cfg(with_testing)]
500impl<D: crate::backends::DatabaseBackup> crate::backends::DatabaseBackup for LruCachingDatabase<D> {
501    fn backup_to(&self, dir: &std::path::Path) -> anyhow::Result<()> {
502        self.database.backup_to(dir)
503    }
504}
linera_views/backends/lru_caching.rs

linera_views/backends/
lru_caching.rs
