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    /// A computation in flight across a `clear()` was started with pre-clear inputs; its
125    /// result must not repopulate the cleared cache. `StateManager::repair_tipset_lookup`
126    /// relies on this `quick_cache` placeholder behavior when it evicts potentially
127    /// tainted results: fills already computing with poisoned inputs are discarded, not
128    /// inserted after the clear.
129    #[tokio::test]
130    async fn test_clear_discards_in_flight_computation() {
131        let cache: Arc<ForestCache<TipsetKey, String>> = Arc::new(ForestCache::new("test"));
132        let key = create_test_tipset_key(1);
133
134        // `entered` fires from inside the fill future, i.e. strictly after
135        // `get_or_insert_async` has installed its placeholder; `gate` then holds the
136        // fill in flight until the cache has been cleared.
137        let entered = Arc::new(tokio::sync::Notify::new());
138        let gate = Arc::new(tokio::sync::Notify::new());
139        let handle = tokio::spawn({
140            let cache = Arc::clone(&cache);
141            let key = key.clone();
142            let entered = Arc::clone(&entered);
143            let gate = Arc::clone(&gate);
144            async move {
145                cache
146                    .get_or_insert_async(&key, async {
147                        entered.notify_one();
148                        gate.notified().await;
149                        anyhow::Ok("stale_value".to_string())
150                    })
151                    .await
152            }
153        });
154        entered.notified().await;
155
156        cache.clear();
157        gate.notify_one();
158        let stale = handle.await.unwrap().unwrap();
159        assert_eq!(stale, "stale_value");
160
161        assert!(
162            cache.get(&key).is_none(),
163            "in-flight fill must not repopulate a cleared cache"
164        );
165    }
166
167    #[tokio::test]
168    async fn test_concurrent_different_keys() {
169        let cache: Arc<ForestCache<TipsetKey, String>> = Arc::new(ForestCache::new("test"));
170        let computation_count = Arc::new(AtomicU8::new(0));
171
172        let mut handles = vec![];
173        for i in 0..10 {
174            let cache_clone = Arc::clone(&cache);
175            let key = create_test_tipset_key(i);
176            let count_clone = Arc::clone(&computation_count);
177
178            let handle = tokio::spawn(async move {
179                cache_clone
180                    .get_or_insert_async(&key, {
181                        let count = Arc::clone(&count_clone);
182                        async move {
183                            count.fetch_add(1, Ordering::SeqCst);
184                            tokio::time::sleep(Duration::from_millis(5)).await;
185                            anyhow::Ok(format!("value_{i}"))
186                        }
187                    })
188                    .await
189            });
190            handles.push(handle);
191        }
192
193        let results: Vec<_> = futures::future::join_all(handles)
194            .await
195            .into_iter()
196            .collect::<Result<Vec<_>, _>>()
197            .unwrap();
198
199        assert_eq!(computation_count.load(Ordering::SeqCst), 10);
200
201        for (i, result) in results.iter().enumerate() {
202            assert_eq!(result.as_ref().unwrap(), &format!("value_{i}"));
203        }
204    }
205}
forest/state_manager/cache.rs

forest/state_manager/
cache.rs
