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::libp2p::PeerId;
6use subspace_networking::utils::multihash::ToMultihash;
7use subspace_networking::KeyWithDistance;
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| piece_indices_to_store.remove(key).is_none())
100            .for_each(|(_piece_index, offset)| {
101                // There is no need to adjust the `last_stored_offset` of the `backend` here,
102                // as the free_offset will be preferentially taken from the dangling free offsets
103                self.dangling_free_offsets.push_back(offset);
104            })
105    }
106
107    pub(super) fn push_dangling_free_offset(&mut self, offset: FarmerCacheOffset) {
108        trace!(?offset, "Pushing dangling free offset");
109        self.dangling_free_offsets.push_back(offset);
110    }
111
112    pub(super) fn get_backend(&self, cache_index: CacheIndex) -> Option<&CacheBackend> {
113        self.backends.get(usize::from(cache_index))
114    }
115
116    pub(super) fn backends(&self) -> impl ExactSizeIterator<Item = &CacheBackend> {
117        self.backends.iter()
118    }
119
120    pub(super) fn reuse(
121        self,
122    ) -> (
123        BTreeMap<KeyWithDistance, FarmerCacheOffset>,
124        VecDeque<FarmerCacheOffset>,
125    ) {
126        let Self {
127            mut stored_pieces,
128            mut dangling_free_offsets,
129            backends: _,
130        } = self;
131
132        stored_pieces.clear();
133        dangling_free_offsets.clear();
134        (stored_pieces, dangling_free_offsets)
135    }
136
137    pub(super) fn should_replace(
138        &mut self,
139        key: &KeyWithDistance,
140    ) -> Option<(KeyWithDistance, FarmerCacheOffset)> {
141        if !self.should_include_key_internal(key) {
142            return None;
143        }
144
145        if !self.has_free_capacity() {
146            self.stored_pieces.pop_last()
147        } else {
148            None
149        }
150    }
151
152    pub(super) fn should_include_key(&self, peer_id: PeerId, key: PieceIndex) -> bool {
153        let key = KeyWithDistance::new(peer_id, key.to_multihash());
154        self.should_include_key_internal(&key)
155    }
156
157    fn has_free_capacity(&self) -> bool {
158        if !self.dangling_free_offsets.is_empty() {
159            return true;
160        }
161
162        self.backends.iter().any(|backend| backend.free_size() > 0)
163    }
164
165    fn should_include_key_internal(&self, key: &KeyWithDistance) -> bool {
166        if self.stored_pieces.contains_key(key) {
167            return false;
168        }
169
170        if self.has_free_capacity() {
171            return true;
172        }
173
174        let top_key = self.stored_pieces.last_key_value().map(|(key, _)| key);
175
176        if let Some(top_key) = top_key {
177            top_key > key
178        } else {
179            false
180        }
181    }
182}
subspace_farmer/farmer_cache/piece_cache_state.rs

subspace_farmer/farmer_cache/
piece_cache_state.rs
