hotshot_libp2p_networking/network/behaviours/dht/store/
persistent.rs

1//! This file contains the `PersistentStore` struct, which is a wrapper around a `RecordStore`
2//! that occasionally saves the DHT to a persistent storage.
3
4use std::{
5    sync::{
6        Arc,
7        atomic::{AtomicU64, Ordering},
8    },
9    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
10};
11
12use anyhow::Context;
13use async_trait::async_trait;
14use delegate::delegate;
15use libp2p::kad::store::{RecordStore, Result};
16use serde::{Deserialize, Serialize};
17use tokio::{sync::Semaphore, time::timeout};
18use tracing::{debug, warn};
19
20/// A trait that we use to save and load the DHT to a file on disk
21/// or other storage medium
22#[async_trait]
23pub trait DhtPersistentStorage: Send + Sync + 'static + Clone {
24    /// Save the DHT (as a list of serializable records) to the persistent storage
25    ///
26    /// # Errors
27    /// - If we fail to save the DHT to the persistent storage provider
28    async fn save(&self, _records: Vec<SerializableRecord>) -> anyhow::Result<()>;
29
30    /// Load the DHT (as a list of serializable records) from the persistent storage
31    ///
32    /// # Errors
33    /// - If we fail to load the DHT from the persistent storage provider
34    async fn load(&self) -> anyhow::Result<Vec<SerializableRecord>>;
35}
36
37/// A no-op `PersistentStorage` that does not persist the DHT
38#[derive(Clone)]
39pub struct DhtNoPersistence;
40
41#[async_trait]
42impl DhtPersistentStorage for DhtNoPersistence {
43    async fn save(&self, _records: Vec<SerializableRecord>) -> anyhow::Result<()> {
44        Ok(())
45    }
46
47    async fn load(&self) -> anyhow::Result<Vec<SerializableRecord>> {
48        Ok(vec![])
49    }
50}
51
52#[async_trait]
53impl<D: DhtPersistentStorage> DhtPersistentStorage for Arc<D> {
54    async fn save(&self, records: Vec<SerializableRecord>) -> anyhow::Result<()> {
55        self.as_ref().save(records).await
56    }
57
58    async fn load(&self) -> anyhow::Result<Vec<SerializableRecord>> {
59        self.as_ref().load().await
60    }
61}
62
63/// A `PersistentStorage` that persists the DHT to a file on disk. Used mostly for
64/// testing.
65#[derive(Clone)]
66pub struct DhtFilePersistence {
67    /// The path to the file on disk
68    path: String,
69}
70
71impl DhtFilePersistence {
72    /// Create a new `DhtFilePersistence` with the given path
73    #[must_use]
74    pub fn new(path: String) -> Self {
75        Self { path }
76    }
77}
78
79#[async_trait]
80impl DhtPersistentStorage for DhtFilePersistence {
81    /// Save the DHT to the file on disk
82    ///
83    /// # Errors
84    /// - If we fail to serialize the records
85    /// - If we fail to write the serialized records to the file
86    async fn save(&self, records: Vec<SerializableRecord>) -> anyhow::Result<()> {
87        // Bincode-serialize the records
88        let to_save =
89            bincode::serialize(&records).with_context(|| "Failed to serialize records")?;
90
91        // Write the serialized records to the file
92        std::fs::write(&self.path, to_save).with_context(|| "Failed to write records to file")?;
93
94        Ok(())
95    }
96
97    /// Load the DHT from the file on disk
98    ///
99    /// # Errors
100    /// - If we fail to read the file
101    /// - If we fail to deserialize the records
102    async fn load(&self) -> anyhow::Result<Vec<SerializableRecord>> {
103        // Read the contents of the file
104        let contents =
105            std::fs::read(&self.path).with_context(|| "Failed to read records from file")?;
106
107        // Deserialize the contents
108        let records: Vec<SerializableRecord> =
109            bincode::deserialize(&contents).with_context(|| "Failed to deserialize records")?;
110
111        Ok(records)
112    }
113}
114
115/// A `RecordStore` wrapper that occasionally saves the DHT to a persistent storage.
116pub struct PersistentStore<R: RecordStore, D: DhtPersistentStorage> {
117    /// The underlying record store
118    underlying_record_store: R,
119
120    /// The persistent storage
121    persistent_storage: D,
122
123    /// The semaphore for limiting the number of concurrent operations (to one)
124    semaphore: Arc<Semaphore>,
125
126    /// The maximum number of records that can be added to the store before the store is saved to the persistent storage
127    max_record_delta: u64,
128
129    /// The running delta between the records in the persistent storage and the records in the underlying store
130    record_delta: Arc<AtomicU64>,
131}
132
133/// A serializable version of a Libp2p `Record`
134#[derive(Serialize, Deserialize)]
135pub struct SerializableRecord {
136    /// The key of the record
137    pub key: libp2p::kad::RecordKey,
138    /// The value of the record
139    #[serde(with = "serde_bytes")]
140    pub value: Vec<u8>,
141    /// The (original) publisher of the record.
142    pub publisher: Option<libp2p::PeerId>,
143    /// The record expiration time in seconds since the Unix epoch
144    ///
145    /// This is an approximation of the expiration time because we can't
146    /// serialize an `Instant` directly.
147    pub expires_unix_secs: Option<u64>,
148}
149
150/// Approximate an `Instant` to the number of seconds since the Unix epoch
151fn instant_to_unix_seconds(instant: Instant) -> anyhow::Result<u64> {
152    // Get the current instant and system time
153    let now_instant = Instant::now();
154    let now_system = SystemTime::now();
155
156    // Get the duration of time between the instant and now
157    if instant > now_instant {
158        Ok(now_system
159            .checked_add(instant - now_instant)
160            .with_context(|| "Overflow when approximating expiration time")?
161            .duration_since(UNIX_EPOCH)
162            .with_context(|| "Failed to get duration since Unix epoch")?
163            .as_secs())
164    } else {
165        Ok(now_system
166            .checked_sub(now_instant - instant)
167            .with_context(|| "Underflow when approximating expiration time")?
168            .duration_since(UNIX_EPOCH)
169            .with_context(|| "Failed to get duration since Unix epoch")?
170            .as_secs())
171    }
172}
173
174/// Convert a unix-second timestamp to an `Instant`
175fn unix_seconds_to_instant(unix_secs: u64) -> anyhow::Result<Instant> {
176    // Get the current instant and unix time
177    let now_instant = Instant::now();
178    let unix_secs_now = SystemTime::now()
179        .duration_since(UNIX_EPOCH)
180        .with_context(|| "Failed to get duration since Unix epoch")?
181        .as_secs();
182
183    if unix_secs > unix_secs_now {
184        // If the instant is in the future, add the duration to the current time
185        now_instant
186            .checked_add(Duration::from_secs(unix_secs - unix_secs_now))
187            .with_context(|| "Overflow when calculating future instant")
188    } else {
189        // If the instant is in the past, subtract the duration from the current time
190        now_instant
191            .checked_sub(Duration::from_secs(unix_secs_now - unix_secs))
192            .with_context(|| "Underflow when calculating past instant")
193    }
194}
195
196/// Allow conversion from a `libp2p::kad::Record` to a `SerializableRecord`
197impl TryFrom<libp2p::kad::Record> for SerializableRecord {
198    type Error = anyhow::Error;
199
200    fn try_from(record: libp2p::kad::Record) -> anyhow::Result<Self> {
201        Ok(SerializableRecord {
202            key: record.key,
203            value: record.value,
204            publisher: record.publisher,
205            expires_unix_secs: record.expires.map(instant_to_unix_seconds).transpose()?,
206        })
207    }
208}
209
210/// Allow conversion from a `SerializableRecord` to a `libp2p::kad::Record`
211impl TryFrom<SerializableRecord> for libp2p::kad::Record {
212    type Error = anyhow::Error;
213
214    fn try_from(record: SerializableRecord) -> anyhow::Result<Self> {
215        Ok(libp2p::kad::Record {
216            key: record.key,
217            value: record.value,
218            publisher: record.publisher,
219            expires: record
220                .expires_unix_secs
221                .map(unix_seconds_to_instant)
222                .transpose()?,
223        })
224    }
225}
226
227impl<R: RecordStore, D: DhtPersistentStorage> PersistentStore<R, D> {
228    /// Create a new `PersistentStore` with the given underlying store and path.
229    /// On creation, the DHT is restored from the persistent storage if possible.
230    ///
231    /// `max_record_delta` is the maximum number of records that can be added to the store before
232    /// the store is saved to the persistent storage.
233    pub async fn new(
234        underlying_record_store: R,
235        persistent_storage: D,
236        max_record_delta: u64,
237    ) -> Self {
238        // Create the new store
239        let mut store = PersistentStore {
240            underlying_record_store,
241            persistent_storage,
242            max_record_delta,
243            record_delta: Arc::new(AtomicU64::new(0)),
244            semaphore: Arc::new(Semaphore::new(1)),
245        };
246
247        // Try to restore the DHT from the persistent store. If it fails, log and start with an empty store
248        if let Err(err) = store.restore_from_persistent_storage().await {
249            debug!(
250                "Failed to restore DHT from persistent storage: {err}. Starting with empty store",
251            );
252        }
253
254        // Return the new store
255        store
256    }
257
258    /// Try saving the DHT to persistent storage if a task is not already in progress.
259    ///
260    /// Returns `true` if the DHT was saved, `false` otherwise.
261    fn try_save_to_persistent_storage(&mut self) -> bool {
262        // Try to acquire the semaphore; if another save operation is already in progress, skip.
263        let Ok(permit) = Arc::clone(&self.semaphore).try_acquire_owned() else {
264            debug!(
265                "Skipping DHT save to persistent storage - another save operation is already in \
266                 progress"
267            );
268            return false;
269        };
270
271        // Get all records and convert them to their serializable counterparts
272        let serializable_records: Vec<_> = self
273            .underlying_record_store
274            .records()
275            .filter_map(|record| {
276                SerializableRecord::try_from(record.into_owned())
277                    .map_err(|err| {
278                        warn!("Failed to convert record to serializable record: {:?}", err);
279                    })
280                    .ok()
281            })
282            .collect();
283
284        // Spawn a task to save the DHT to the persistent storage
285        let persistent_storage = self.persistent_storage.clone();
286        let record_delta = Arc::clone(&self.record_delta);
287        tokio::spawn(async move {
288            debug!("Saving DHT to persistent storage");
289
290            // Save the DHT to the persistent storage
291            match timeout(
292                Duration::from_secs(10),
293                persistent_storage.save(serializable_records),
294            )
295            .await
296            .map_err(|_| anyhow::anyhow!("save operation timed out"))
297            {
298                Ok(Ok(())) => {},
299                Ok(Err(error)) | Err(error) => {
300                    warn!("Failed to save DHT to persistent storage: {error}");
301                },
302            };
303
304            // Reset the record delta
305            record_delta.store(0, Ordering::Release);
306
307            drop(permit);
308
309            debug!("Saved DHT to persistent storage");
310        });
311
312        true
313    }
314
315    /// Attempt to restore the DHT to the underlying store from the persistent storage
316    ///
317    /// # Errors
318    /// - If we fail to load from the persistent storage
319    pub async fn restore_from_persistent_storage(&mut self) -> anyhow::Result<()> {
320        debug!("Restoring DHT from persistent storage");
321
322        // Read the contents of the persistent store
323        let serializable_records = self
324            .persistent_storage
325            .load()
326            .await
327            .with_context(|| "Failed to read DHT from persistent storage")?;
328
329        // Put all records into the new store
330        for serializable_record in serializable_records {
331            // Convert the serializable record back to a `libp2p::kad::Record`
332            match libp2p::kad::Record::try_from(serializable_record) {
333                Ok(record) => {
334                    // If the record doesn't have an expiration time, or has expired, skip it
335                    if record.expires.is_none() || record.expires.unwrap() < Instant::now() {
336                        continue;
337                    }
338
339                    // Put the record into the new store
340                    if let Err(err) = self.underlying_record_store.put(record) {
341                        warn!(
342                            "Failed to restore record from persistent storage: {:?}",
343                            err
344                        );
345                    }
346                },
347                Err(err) => {
348                    warn!("Failed to parse record from persistent storage: {:?}", err);
349                },
350            };
351        }
352
353        debug!("Restored DHT from persistent storage");
354
355        Ok(())
356    }
357}
358
359/// Implement the `RecordStore` trait for `PersistentStore`
360impl<R: RecordStore, D: DhtPersistentStorage> RecordStore for PersistentStore<R, D> {
361    type ProvidedIter<'a>
362        = R::ProvidedIter<'a>
363    where
364        R: 'a,
365        D: 'a;
366    type RecordsIter<'a>
367        = R::RecordsIter<'a>
368    where
369        R: 'a,
370        D: 'a;
371    // Delegate all `RecordStore` methods except `put` to the inner store
372    delegate! {
373        to self.underlying_record_store {
374            fn add_provider(&mut self, record: libp2p::kad::ProviderRecord) -> libp2p::kad::store::Result<()>;
375            fn get(&self, k: &libp2p::kad::RecordKey) -> Option<std::borrow::Cow<'_, libp2p::kad::Record>>;
376            fn provided(&self) -> Self::ProvidedIter<'_>;
377            fn providers(&self, key: &libp2p::kad::RecordKey) -> Vec<libp2p::kad::ProviderRecord>;
378            fn records(&self) -> Self::RecordsIter<'_>;
379            fn remove_provider(&mut self, k: &libp2p::kad::RecordKey, p: &libp2p::PeerId);
380        }
381    }
382
383    /// Override the `put` method to potentially sync the DHT to the persistent store
384    fn put(&mut self, record: libp2p::kad::Record) -> Result<()> {
385        // Try to write to the underlying store
386        let result = self.underlying_record_store.put(record);
387
388        // If the record was successfully written,
389        if result.is_ok() {
390            // Update the record delta
391            self.record_delta.fetch_add(1, Ordering::Relaxed);
392
393            // Check if it's above the maximum record delta
394            if self.record_delta.load(Ordering::Relaxed) > self.max_record_delta {
395                // Try to save the DHT to persistent storage
396                self.try_save_to_persistent_storage();
397            }
398        }
399
400        result
401    }
402
403    /// Overwrite the `remove` method to potentially sync the DHT to the persistent store
404    fn remove(&mut self, k: &libp2p::kad::RecordKey) {
405        // Remove the record from the underlying store
406        self.underlying_record_store.remove(k);
407
408        // Update the record delta
409        self.record_delta.fetch_add(1, Ordering::Relaxed);
410
411        // Check if it's above the maximum record delta
412        if self.record_delta.load(Ordering::Relaxed) > self.max_record_delta {
413            // Try to save the DHT to persistent storage
414            self.try_save_to_persistent_storage();
415        }
416    }
417}
418
419#[cfg(test)]
420mod tests {
421    use libp2p::{
422        PeerId,
423        kad::{RecordKey, store::MemoryStore},
424    };
425    use tracing_subscriber::EnvFilter;
426
427    use super::*;
428
429    #[tokio::test]
430    async fn test_save_and_restore() {
431        // Try initializing tracing
432        let _ = tracing_subscriber::fmt()
433            .with_env_filter(EnvFilter::from_default_env())
434            .try_init();
435
436        // Create a test store
437        let mut store = PersistentStore::new(
438            MemoryStore::new(PeerId::random()),
439            DhtFilePersistence::new("/tmp/test1.dht".to_string()),
440            10,
441        )
442        .await;
443
444        // The key is a random 16-byte array
445        let key = RecordKey::new(&rand::random::<[u8; 16]>().to_vec());
446
447        // The value is a random 16-byte array
448        let random_value = rand::random::<[u8; 16]>();
449
450        // Create a new record
451        let mut record = libp2p::kad::Record::new(key.clone(), random_value.to_vec());
452
453        // Set the expiry time to 10 seconds in the future
454        record.expires = Some(Instant::now() + Duration::from_secs(10));
455
456        // Put a record into the store
457        store.put(record).expect("Failed to put record into store");
458
459        // Try to save the store to a persistent storage
460        assert!(store.try_save_to_persistent_storage());
461
462        // Wait a bit for the save to complete
463        tokio::time::sleep(Duration::from_millis(100)).await;
464
465        // Create a new store from the persistent storage
466        let new_store = PersistentStore::new(
467            MemoryStore::new(PeerId::random()),
468            DhtFilePersistence::new("/tmp/test1.dht".to_string()),
469            10,
470        )
471        .await;
472
473        // Check that the new store has the record
474        let restored_record = new_store
475            .get(&key)
476            .expect("Failed to get record from store");
477
478        // Check that the restored record has the same value as the original record
479        assert_eq!(restored_record.value, random_value.to_vec());
480    }
481
482    #[tokio::test]
483    async fn test_record_delta() {
484        // Try initializing tracing
485        let _ = tracing_subscriber::fmt()
486            .with_env_filter(EnvFilter::from_default_env())
487            .try_init();
488
489        // Create a test store
490        let mut store = PersistentStore::new(
491            MemoryStore::new(PeerId::random()),
492            DhtFilePersistence::new("/tmp/test2.dht".to_string()),
493            10,
494        )
495        .await;
496
497        let mut keys = Vec::new();
498        let mut values = Vec::new();
499
500        // Put 10 records into the store
501        for _ in 0..10 {
502            // Create a random key and value
503            let key = RecordKey::new(&rand::random::<[u8; 16]>().to_vec());
504            let value = rand::random::<[u8; 16]>();
505
506            keys.push(key.clone());
507            values.push(value);
508
509            // Create a new record
510            let mut record = libp2p::kad::Record::new(key, value.to_vec());
511
512            // Set the expiry time to 10 seconds in the future
513            record.expires = Some(Instant::now() + Duration::from_secs(10));
514
515            store.put(record).expect("Failed to put record into store");
516        }
517
518        // Create a new store from the allegedly unpersisted DHT
519        let new_store = PersistentStore::new(
520            MemoryStore::new(PeerId::random()),
521            DhtFilePersistence::new("/tmp/test2.dht".to_string()),
522            10,
523        )
524        .await;
525
526        // Check that the new store has none of the records
527        for key in &keys {
528            assert!(new_store.get(key).is_none());
529        }
530
531        // Create a new record
532        let mut record = libp2p::kad::Record::new(keys[0].clone(), values[0].to_vec());
533
534        // Set the expiry time to 10 seconds in the future
535        record.expires = Some(Instant::now() + Duration::from_secs(10));
536
537        // Store one more record into the new store
538        store.put(record).expect("Failed to put record into store");
539
540        // Wait a bit for the save to complete
541        tokio::time::sleep(Duration::from_millis(100)).await;
542
543        // Create a new store from the allegedly saved DHT
544        let new_store = PersistentStore::new(
545            MemoryStore::new(PeerId::random()),
546            DhtFilePersistence::new("/tmp/test2.dht".to_string()),
547            10,
548        )
549        .await;
550
551        // Check that the new store has all of the records
552        for (i, key) in keys.iter().enumerate() {
553            let restored_record = new_store.get(key).expect("Failed to get record from store");
554            assert_eq!(restored_record.value, values[i]);
555        }
556
557        // Check that the record delta is 0
558        assert_eq!(store.record_delta.load(Ordering::Relaxed), 0);
559    }
560
561    #[test]
562    fn test_approximate_instant() {
563        // Create an expiry time in the future
564        let expiry_future = Instant::now() + Duration::from_secs(10);
565
566        // Approximate the expiry time
567        let approximate_expiry =
568            unix_seconds_to_instant(instant_to_unix_seconds(expiry_future).unwrap())
569                .unwrap()
570                .duration_since(Instant::now());
571
572        // Make sure it's close to 10 seconds in the future
573        assert!(approximate_expiry >= Duration::from_secs(9));
574        assert!(approximate_expiry <= Duration::from_secs(11));
575
576        // Create an expiry time in the past
577        let expiry_past = Instant::now().checked_sub(Duration::from_secs(10)).unwrap();
578
579        // Approximate the expiry time
580        let approximate_expiry =
581            unix_seconds_to_instant(instant_to_unix_seconds(expiry_past).unwrap()).unwrap();
582        let time_difference = approximate_expiry.elapsed();
583
584        // Make sure it's close to 10 seconds in the past
585        assert!(time_difference >= Duration::from_secs(9));
586        assert!(time_difference <= Duration::from_secs(11));
587    }
588}
hotshot_libp2p_networking/network/behaviours/dht/store/persistent.rs

hotshot_libp2p_networking/network/behaviours/dht/store/
persistent.rs
