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

linera_views/backends/
lru_caching.rs
