forest/state_manager/
cache.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use crate::prelude::*;
5use crate::state_manager::DEFAULT_TIPSET_CACHE_SIZE;
6use crate::utils::cache::{CacheKeyConstraints, CacheValueConstraints, SizeTrackingCache};
7use std::borrow::Cow;
8use std::num::NonZeroUsize;
9
10/// A cache that handles concurrent access and computation for tipset-related
11/// data. Coalesces concurrent computations of the same key, so only one caller
12/// actually runs the `compute` future and the rest wait on its result.
13#[derive(derive_more::Deref)]
14pub(crate) struct ForestCache<K: CacheKeyConstraints, V: CacheValueConstraints> {
15    cache: SizeTrackingCache<K, V>,
16}
17
18impl<K: CacheKeyConstraints, V: CacheValueConstraints> ShallowClone for ForestCache<K, V> {
19    fn shallow_clone(&self) -> Self {
20        Self {
21            cache: self.cache.shallow_clone(),
22        }
23    }
24}
25
26impl<K: CacheKeyConstraints, V: CacheValueConstraints> ForestCache<K, V> {
27    pub fn new(cache_identifier: impl Into<Cow<'static, str>>) -> Self {
28        Self::with_size(cache_identifier, DEFAULT_TIPSET_CACHE_SIZE)
29    }
30
31    pub fn with_size(
32        cache_identifier: impl Into<Cow<'static, str>>,
33        cache_size: NonZeroUsize,
34    ) -> Self {
35        Self {
36            cache: SizeTrackingCache::new_with_metrics(cache_identifier, cache_size),
37        }
38    }
39
40    pub fn get_map<T>(&self, key: &K, mapper: impl FnOnce(&V) -> T) -> Option<T> {
41        self.cache.get(key).as_ref().map(mapper)
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48    use crate::blocks::TipsetKey;
49    use cid::Cid;
50    use fvm_ipld_encoding::DAG_CBOR;
51    use multihash_derive::MultihashDigest;
52    use std::sync::Arc;
53    use std::sync::atomic::{AtomicU8, Ordering};
54    use std::time::Duration;
55
56    fn create_test_tipset_key(i: u64) -> TipsetKey {
57        let bytes = i.to_le_bytes().to_vec();
58        let cid = Cid::new_v1(
59            DAG_CBOR,
60            crate::utils::multihash::MultihashCode::Blake2b256.digest(&bytes),
61        );
62        TipsetKey::from(nunny::vec![cid])
63    }
64
65    #[tokio::test]
66    async fn test_tipset_cache_basic_functionality() {
67        let cache: ForestCache<TipsetKey, String> = ForestCache::new("test");
68        let key = create_test_tipset_key(1);
69
70        let result = cache
71            .get_or_insert_async(&key, async { anyhow::Ok("computed_value".to_string()) })
72            .await
73            .unwrap();
74        assert_eq!(result, "computed_value");
75
76        let result = cache
77            .get_or_insert_async(&key, async { anyhow::Ok("should_not_compute".to_string()) })
78            .await
79            .unwrap();
80        assert_eq!(result, "computed_value");
81    }
82
83    #[tokio::test]
84    async fn test_concurrent_same_key_computation() {
85        let cache: Arc<ForestCache<TipsetKey, String>> = Arc::new(ForestCache::new("test"));
86        let key = create_test_tipset_key(1);
87        let computation_count = Arc::new(AtomicU8::new(0));
88
89        let mut handles = vec![];
90        for i in 0..10 {
91            let cache_clone = Arc::clone(&cache);
92            let key_clone = key.clone();
93            let count_clone = Arc::clone(&computation_count);
94
95            let handle = tokio::spawn(async move {
96                cache_clone
97                    .get_or_insert_async(&key_clone, {
98                        let count = Arc::clone(&count_clone);
99                        async move {
100                            count.fetch_add(1, Ordering::SeqCst);
101                            tokio::time::sleep(Duration::from_millis(10)).await;
102                            anyhow::Ok(format!("computed_value_{i}"))
103                        }
104                    })
105                    .await
106            });
107            handles.push(handle);
108        }
109
110        let results: Vec<_> = futures::future::join_all(handles)
111            .await
112            .into_iter()
113            .collect::<Result<Vec<_>, _>>()
114            .unwrap();
115
116        assert_eq!(computation_count.load(Ordering::SeqCst), 1);
117
118        let first_result = results[0].as_ref().unwrap();
119        for result in &results {
120            assert_eq!(result.as_ref().unwrap(), first_result);
121        }
122    }
123
124    #[tokio::test]
125    async fn test_concurrent_different_keys() {
126        let cache: Arc<ForestCache<TipsetKey, String>> = Arc::new(ForestCache::new("test"));
127        let computation_count = Arc::new(AtomicU8::new(0));
128
129        let mut handles = vec![];
130        for i in 0..10 {
131            let cache_clone = Arc::clone(&cache);
132            let key = create_test_tipset_key(i);
133            let count_clone = Arc::clone(&computation_count);
134
135            let handle = tokio::spawn(async move {
136                cache_clone
137                    .get_or_insert_async(&key, {
138                        let count = Arc::clone(&count_clone);
139                        async move {
140                            count.fetch_add(1, Ordering::SeqCst);
141                            tokio::time::sleep(Duration::from_millis(5)).await;
142                            anyhow::Ok(format!("value_{i}"))
143                        }
144                    })
145                    .await
146            });
147            handles.push(handle);
148        }
149
150        let results: Vec<_> = futures::future::join_all(handles)
151            .await
152            .into_iter()
153            .collect::<Result<Vec<_>, _>>()
154            .unwrap();
155
156        assert_eq!(computation_count.load(Ordering::SeqCst), 10);
157
158        for (i, result) in results.iter().enumerate() {
159            assert_eq!(result.as_ref().unwrap(), &format!("value_{i}"));
160        }
161    }
162}
forest/state_manager/cache.rs

forest/state_manager/
cache.rs
