subspace_farmer/farmer_cache/
piece_cache_state.rs

1use crate::farmer_cache::{CacheBackend, CacheIndex, FarmerCacheOffset};
2use std::collections::btree_map::Values;
3use std::collections::{BTreeMap, HashMap, VecDeque};
4use subspace_core_primitives::pieces::PieceIndex;
5use subspace_networking::KeyWithDistance;
6use subspace_networking::libp2p::PeerId;
7use subspace_networking::utils::multihash::ToMultihash;
8use tracing::{debug, trace};
9
10#[derive(Debug, Default, Clone)]
11pub(super) struct PieceCachesState {
12    stored_pieces: BTreeMap<KeyWithDistance, FarmerCacheOffset>,
13    dangling_free_offsets: VecDeque<FarmerCacheOffset>,
14    backends: Vec<CacheBackend>,
15}
16
17impl PieceCachesState {
18    pub(super) fn new(
19        stored_pieces: BTreeMap<KeyWithDistance, FarmerCacheOffset>,
20        dangling_free_offsets: VecDeque<FarmerCacheOffset>,
21        backends: Vec<CacheBackend>,
22    ) -> Self {
23        Self {
24            stored_pieces,
25            dangling_free_offsets,
26            backends,
27        }
28    }
29
30    pub(super) fn total_capacity(&self) -> usize {
31        self.backends()
32            .fold(0usize, |acc, backend| acc + backend.total_capacity as usize)
33    }
34
35    pub(super) fn pop_free_offset(&mut self) -> Option<FarmerCacheOffset> {
36        match self.dangling_free_offsets.pop_front() {
37            Some(free_offset) => {
38                debug!(?free_offset, "Popped dangling free offset");
39                Some(free_offset)
40            }
41            None => {
42                // TODO: Use simple round robin for uniform filling instead of re-sorting all the
43                //  time
44                // Sort piece caches by number of stored pieces to fill those that are less
45                // populated first
46                let mut sorted_backends = self
47                    .backends
48                    .iter_mut()
49                    .enumerate()
50                    .filter_map(|(cache_index, backend)| {
51                        Some((CacheIndex::try_from(cache_index).ok()?, backend))
52                    })
53                    .collect::<Vec<_>>();
54                sorted_backends.sort_unstable_by_key(|(_, backend)| backend.free_size());
55                sorted_backends
56                    .into_iter()
57                    .rev()
58                    .find_map(|(cache_index, backend)| {
59                        backend
60                            .next_free()
61                            .map(|free_offset| FarmerCacheOffset::new(cache_index, free_offset))
62                    })
63            }
64        }
65    }
66
67    pub(super) fn get_stored_piece(&self, key: &KeyWithDistance) -> Option<&FarmerCacheOffset> {
68        self.stored_pieces.get(key)
69    }
70
71    pub(super) fn contains_stored_piece(&self, key: &KeyWithDistance) -> bool {
72        self.stored_pieces.contains_key(key)
73    }
74
75    pub(super) fn push_stored_piece(
76        &mut self,
77        key: KeyWithDistance,
78        cache_offset: FarmerCacheOffset,
79    ) -> Option<FarmerCacheOffset> {
80        self.stored_pieces.insert(key, cache_offset)
81    }
82
83    pub(super) fn stored_pieces_offsets(&self) -> Values<'_, KeyWithDistance, FarmerCacheOffset> {
84        self.stored_pieces.values()
85    }
86
87    pub(super) fn remove_stored_piece(
88        &mut self,
89        key: &KeyWithDistance,
90    ) -> Option<FarmerCacheOffset> {
91        self.stored_pieces.remove(key)
92    }
93
94    pub(super) fn free_unneeded_stored_pieces(
95        &mut self,
96        piece_indices_to_store: &mut HashMap<KeyWithDistance, PieceIndex>,
97    ) {
98        self.stored_pieces
99            .extract_if(.., |key, _offset| {
100                piece_indices_to_store.remove(key).is_none()
101            })
102            .for_each(|(_piece_index, offset)| {
103                // There is no need to adjust the `last_stored_offset` of the `backend` here,
104                // as the free_offset will be preferentially taken from the dangling free offsets
105                self.dangling_free_offsets.push_back(offset);
106            })
107    }
108
109    pub(super) fn push_dangling_free_offset(&mut self, offset: FarmerCacheOffset) {
110        trace!(?offset, "Pushing dangling free offset");
111        self.dangling_free_offsets.push_back(offset);
112    }
113
114    pub(super) fn get_backend(&self, cache_index: CacheIndex) -> Option<&CacheBackend> {
115        self.backends.get(usize::from(cache_index))
116    }
117
118    pub(super) fn backends(&self) -> impl ExactSizeIterator<Item = &CacheBackend> {
119        self.backends.iter()
120    }
121
122    pub(super) fn reuse(
123        self,
124    ) -> (
125        BTreeMap<KeyWithDistance, FarmerCacheOffset>,
126        VecDeque<FarmerCacheOffset>,
127    ) {
128        let Self {
129            mut stored_pieces,
130            mut dangling_free_offsets,
131            backends: _,
132        } = self;
133
134        stored_pieces.clear();
135        dangling_free_offsets.clear();
136        (stored_pieces, dangling_free_offsets)
137    }
138
139    pub(super) fn should_replace(
140        &mut self,
141        key: &KeyWithDistance,
142    ) -> Option<(KeyWithDistance, FarmerCacheOffset)> {
143        if !self.should_include_key_internal(key) {
144            return None;
145        }
146
147        if !self.has_free_capacity() {
148            self.stored_pieces.pop_last()
149        } else {
150            None
151        }
152    }
153
154    pub(super) fn should_include_key(&self, peer_id: PeerId, key: PieceIndex) -> bool {
155        let key = KeyWithDistance::new(peer_id, key.to_multihash());
156        self.should_include_key_internal(&key)
157    }
158
159    fn has_free_capacity(&self) -> bool {
160        if !self.dangling_free_offsets.is_empty() {
161            return true;
162        }
163
164        self.backends.iter().any(|backend| backend.free_size() > 0)
165    }
166
167    fn should_include_key_internal(&self, key: &KeyWithDistance) -> bool {
168        if self.stored_pieces.contains_key(key) {
169            return false;
170        }
171
172        if self.has_free_capacity() {
173            return true;
174        }
175
176        let top_key = self.stored_pieces.last_key_value().map(|(key, _)| key);
177
178        if let Some(top_key) = top_key {
179            top_key > key
180        } else {
181            false
182        }
183    }
184}
subspace_farmer/farmer_cache/piece_cache_state.rs

subspace_farmer/farmer_cache/
piece_cache_state.rs
