ant_bootstrap/
cache_store.rs

1// Copyright 2024 MaidSafe.net limited.
2//
3// This SAFE Network Software is licensed to you under The General Public License (GPL), version 3.
4// Unless required by applicable law or agreed to in writing, the SAFE Network Software distributed
5// under the GPL Licence is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
6// KIND, either express or implied. Please review the Licences for the specific language governing
7// permissions and limitations relating to use of the SAFE Network Software.
8
9use crate::{
10    craft_valid_multiaddr, multiaddr_get_peer_id, BootstrapAddr, BootstrapAddresses,
11    BootstrapCacheConfig, Error, InitialPeersConfig, Result,
12};
13use atomic_write_file::AtomicWriteFile;
14use libp2p::{multiaddr::Protocol, Multiaddr, PeerId};
15use serde::{Deserialize, Serialize};
16use std::{
17    collections::{hash_map::Entry, HashMap},
18    fs::{self, OpenOptions},
19    io::{Read, Write},
20    path::PathBuf,
21    time::{Duration, SystemTime},
22};
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct CacheData {
26    pub peers: std::collections::HashMap<PeerId, BootstrapAddresses>,
27    pub last_updated: SystemTime,
28    pub network_version: String,
29}
30
31impl CacheData {
32    pub fn insert(&mut self, peer_id: PeerId, bootstrap_addr: BootstrapAddr) {
33        match self.peers.entry(peer_id) {
34            Entry::Occupied(mut occupied_entry) => {
35                occupied_entry.get_mut().insert_addr(&bootstrap_addr);
36            }
37            Entry::Vacant(vacant_entry) => {
38                vacant_entry.insert(BootstrapAddresses(vec![bootstrap_addr]));
39            }
40        }
41    }
42
43    /// Sync the self cache with another cache. This would just add the 'other' state to self.
44    pub fn sync(&mut self, other: &CacheData) {
45        for (peer, other_addresses_state) in other.peers.iter() {
46            let bootstrap_addresses = self
47                .peers
48                .entry(*peer)
49                .or_insert(other_addresses_state.clone());
50
51            trace!("Syncing {peer:?} from other with addrs count: {:?}. Our in memory state count: {:?}", other_addresses_state.0.len(), bootstrap_addresses.0.len());
52
53            bootstrap_addresses.sync(other_addresses_state);
54        }
55
56        self.last_updated = SystemTime::now();
57    }
58
59    /// Remove the oldest peers until we're under the max_peers limit
60    pub fn try_remove_oldest_peers(&mut self, cfg: &BootstrapCacheConfig) {
61        if self.peers.len() > cfg.max_peers {
62            let mut peer_last_seen_map = HashMap::new();
63            for (peer, addrs) in self.peers.iter() {
64                let mut latest_seen = Duration::from_secs(u64::MAX);
65                for addr in addrs.0.iter() {
66                    if let Ok(elapsed) = addr.last_seen.elapsed() {
67                        trace!("Time elapsed for {addr:?} is {elapsed:?}");
68                        if elapsed < latest_seen {
69                            trace!("Updating latest_seen to {elapsed:?}");
70                            latest_seen = elapsed;
71                        }
72                    }
73                }
74                trace!("Last seen for {peer:?} is {latest_seen:?}");
75                peer_last_seen_map.insert(*peer, latest_seen);
76            }
77
78            while self.peers.len() > cfg.max_peers {
79                // find the peer with the largest last_seen
80                if let Some((&oldest_peer, last_seen)) = peer_last_seen_map
81                    .iter()
82                    .max_by_key(|(_, last_seen)| **last_seen)
83                {
84                    debug!("Found the oldest peer to remove: {oldest_peer:?} with last_seen of {last_seen:?}");
85                    self.peers.remove(&oldest_peer);
86                    peer_last_seen_map.remove(&oldest_peer);
87                }
88            }
89        }
90    }
91}
92
93impl Default for CacheData {
94    fn default() -> Self {
95        Self {
96            peers: std::collections::HashMap::new(),
97            last_updated: SystemTime::now(),
98            network_version: crate::get_network_version(),
99        }
100    }
101}
102
103#[derive(Clone, Debug)]
104pub struct BootstrapCacheStore {
105    pub(crate) cache_path: PathBuf,
106    pub(crate) config: BootstrapCacheConfig,
107    pub(crate) data: CacheData,
108}
109
110impl BootstrapCacheStore {
111    pub fn config(&self) -> &BootstrapCacheConfig {
112        &self.config
113    }
114
115    /// Create an empty CacheStore with the given configuration
116    pub fn new(config: BootstrapCacheConfig) -> Result<Self> {
117        info!("Creating new CacheStore with config: {:?}", config);
118        let cache_path = config.cache_file_path.clone();
119
120        // Create cache directory if it doesn't exist
121        if let Some(parent) = cache_path.parent() {
122            if !parent.exists() {
123                info!("Attempting to create cache directory at {parent:?}");
124                fs::create_dir_all(parent).inspect_err(|err| {
125                    warn!("Failed to create cache directory at {parent:?}: {err}");
126                })?;
127            }
128        }
129
130        let store = Self {
131            cache_path,
132            config,
133            data: CacheData::default(),
134        };
135
136        Ok(store)
137    }
138
139    /// Create an empty CacheStore from the given Initial Peers Configuration.
140    /// This also modifies the `BootstrapCacheConfig` if provided based on the `InitialPeersConfig`.
141    /// And also performs some actions based on the `InitialPeersConfig`.
142    ///
143    /// `InitialPeersConfig::bootstrap_cache_dir` will take precedence over the path provided inside `config`.
144    pub fn new_from_initial_peers_config(
145        init_peers_config: &InitialPeersConfig,
146        config: Option<BootstrapCacheConfig>,
147    ) -> Result<Self> {
148        let mut config = if let Some(cfg) = config {
149            cfg
150        } else {
151            BootstrapCacheConfig::default_config(init_peers_config.local)?
152        };
153        if let Some(bootstrap_cache_path) = init_peers_config.get_bootstrap_cache_path()? {
154            config.cache_file_path = bootstrap_cache_path;
155        }
156
157        let store = Self::new(config)?;
158
159        // If it is the first node, clear the cache.
160        if init_peers_config.first {
161            info!("First node in network, writing empty cache to disk");
162            store.write()?;
163        }
164
165        Ok(store)
166    }
167
168    /// Load cache data from disk
169    /// Make sure to have clean addrs inside the cache as we don't call craft_valid_multiaddr
170    pub fn load_cache_data(cfg: &BootstrapCacheConfig) -> Result<CacheData> {
171        // Try to open the file with read permissions
172        let mut file = OpenOptions::new()
173            .read(true)
174            .open(&cfg.cache_file_path)
175            .inspect_err(|err| warn!("Failed to open cache file: {err}",))?;
176
177        // Read the file contents
178        let mut contents = String::new();
179        file.read_to_string(&mut contents).inspect_err(|err| {
180            warn!("Failed to read cache file: {err}");
181        })?;
182
183        // Parse the cache data
184        let mut data = serde_json::from_str::<CacheData>(&contents).map_err(|err| {
185            warn!("Failed to parse cache data: {err}");
186            Error::FailedToParseCacheData
187        })?;
188
189        data.try_remove_oldest_peers(cfg);
190
191        Ok(data)
192    }
193
194    pub fn peer_count(&self) -> usize {
195        self.data.peers.len()
196    }
197
198    pub fn get_all_addrs(&self) -> impl Iterator<Item = &BootstrapAddr> {
199        self.data
200            .peers
201            .values()
202            .flat_map(|bootstrap_addresses| bootstrap_addresses.0.iter())
203    }
204
205    /// Get a list containing single addr per peer. We use the least faulty addr for each peer.
206    /// This list is sorted by the failure rate of the addr.
207    pub fn get_sorted_addrs(&self) -> impl Iterator<Item = &Multiaddr> {
208        let mut addrs = self
209            .data
210            .peers
211            .values()
212            .flat_map(|bootstrap_addresses| bootstrap_addresses.get_least_faulty())
213            .collect::<Vec<_>>();
214
215        addrs.sort_by_key(|addr| addr.failure_rate() as u64);
216
217        addrs.into_iter().map(|addr| &addr.addr)
218    }
219
220    /// Update the status of an addr in the cache. The peer must be added to the cache first.
221    pub fn update_addr_status(&mut self, addr: &Multiaddr, success: bool) {
222        if let Some(peer_id) = multiaddr_get_peer_id(addr) {
223            debug!("Updating addr status: {addr} (success: {success})");
224            if let Some(bootstrap_addresses) = self.data.peers.get_mut(&peer_id) {
225                bootstrap_addresses.update_addr_status(addr, success);
226            } else {
227                debug!("Peer not found in cache to update: {addr}");
228            }
229        }
230    }
231
232    /// Add a set of addresses to the cache.
233    pub fn add_addr(&mut self, addr: Multiaddr) {
234        debug!("Trying to add new addr: {addr}");
235        let Some(addr) = craft_valid_multiaddr(&addr, false) else {
236            return;
237        };
238        let peer_id = match addr.iter().find(|p| matches!(p, Protocol::P2p(_))) {
239            Some(Protocol::P2p(id)) => id,
240            _ => return,
241        };
242
243        // Check if we already have this peer
244        if let Some(bootstrap_addrs) = self.data.peers.get_mut(&peer_id) {
245            if let Some(bootstrap_addr) = bootstrap_addrs.get_addr_mut(&addr) {
246                debug!("Updating existing peer's last_seen {addr}");
247                bootstrap_addr.last_seen = SystemTime::now();
248                return;
249            } else {
250                let mut bootstrap_addr = BootstrapAddr::new(addr.clone());
251                bootstrap_addr.success_count = 1;
252                bootstrap_addrs.insert_addr(&bootstrap_addr);
253            }
254        } else {
255            let mut bootstrap_addr = BootstrapAddr::new(addr.clone());
256            bootstrap_addr.success_count = 1;
257            self.data
258                .peers
259                .insert(peer_id, BootstrapAddresses(vec![bootstrap_addr]));
260        }
261
262        debug!("Added new peer {addr:?}, performing cleanup of old addrs");
263        self.try_remove_oldest_peers();
264    }
265
266    /// Remove a single address for a peer.
267    pub fn remove_addr(&mut self, addr: &Multiaddr) {
268        if let Some(peer_id) = multiaddr_get_peer_id(addr) {
269            if let Some(bootstrap_addresses) = self.data.peers.get_mut(&peer_id) {
270                bootstrap_addresses.remove_addr(addr);
271            } else {
272                debug!("Peer {peer_id:?} not found in the cache. Not removing addr: {addr:?}")
273            }
274        } else {
275            debug!("Could not obtain PeerId for {addr:?}, not removing addr from cache.");
276        }
277    }
278
279    pub fn try_remove_oldest_peers(&mut self) {
280        self.data.try_remove_oldest_peers(&self.config);
281    }
282
283    /// Flush the cache to disk after syncing with the CacheData from the file.
284    pub fn sync_and_flush_to_disk(&mut self) -> Result<()> {
285        if self.config.disable_cache_writing {
286            info!("Cache writing is disabled, skipping sync to disk");
287            return Ok(());
288        }
289
290        info!(
291            "Flushing cache to disk, with data containing: {} peers",
292            self.data.peers.len(),
293        );
294
295        if let Ok(data_from_file) = Self::load_cache_data(&self.config) {
296            self.data.sync(&data_from_file);
297        } else {
298            warn!("Failed to load cache data from file, overwriting with new data");
299        }
300
301        self.data.try_remove_oldest_peers(&self.config);
302
303        self.write().inspect_err(|e| {
304            error!("Failed to save cache to disk: {e}");
305        })?;
306
307        // Flush after writing
308        self.data.peers.clear();
309
310        Ok(())
311    }
312
313    /// Write the cache to disk atomically. This will overwrite the existing cache file, use sync_and_flush_to_disk to
314    /// sync with the file first.
315    pub fn write(&self) -> Result<()> {
316        debug!("Writing cache to disk: {:?}", self.cache_path);
317        // Create parent directory if it doesn't exist
318        if let Some(parent) = self.cache_path.parent() {
319            fs::create_dir_all(parent)?;
320        }
321
322        let mut file = AtomicWriteFile::options()
323            .open(&self.cache_path)
324            .inspect_err(|err| {
325                error!("Failed to open cache file using AtomicWriteFile: {err}");
326            })?;
327
328        let data = serde_json::to_string_pretty(&self.data).inspect_err(|err| {
329            error!("Failed to serialize cache data: {err}");
330        })?;
331        writeln!(file, "{data}")?;
332        file.commit().inspect_err(|err| {
333            error!("Failed to commit atomic write: {err}");
334        })?;
335
336        info!("Cache written to disk: {:?}", self.cache_path);
337
338        Ok(())
339    }
340}
ant_bootstrap/cache_store.rs

ant_bootstrap/
cache_store.rs
