Skip to main content

ipfrs_network/
provider_renewal.rs

1//! DHT provider record auto-renewal scheduler.
2//!
3//! This module implements [`ProviderRenewalScheduler`], which tracks CIDs that
4//! have been announced to the Kademlia DHT as provider records and schedules
5//! automatic re-announcement before the standard 24-hour TTL expires.
6//!
7//! ## Background
8//!
9//! The Kademlia specification (and the libp2p/IPFS implementation) assigns a
10//! 24-hour TTL to provider records.  If a node does not re-publish a record
11//! before the TTL elapses, remote peers will stop discovering that node as a
12//! provider for the corresponding CID.  The scheduler triggers renewal at 80%
13//! of the TTL (configurable), giving a comfortable safety margin.
14//!
15//! ## Usage
16//!
17//! ```rust
18//! use ipfrs_network::provider_renewal::{ProviderRenewalScheduler, RenewalConfig};
19//! use std::time::{SystemTime, UNIX_EPOCH};
20//!
21//! let scheduler = ProviderRenewalScheduler::new(RenewalConfig::default());
22//!
23//! let now = SystemTime::now()
24//!     .duration_since(UNIX_EPOCH)
25//!     .unwrap_or_default()
26//!     .as_secs();
27//!
28//! scheduler.track("bafybeihdwdcefgh4dqkjv67uzcmw7ojee6xedzdetojuzjevtenxquvyku", now);
29//!
30//! // Later, check what needs renewal:
31//! let due = scheduler.due_for_renewal(now + 70_000);
32//! scheduler.mark_renewed(&due, now + 70_000);
33//! ```
34
35use std::collections::HashMap;
36use std::sync::atomic::{AtomicU64, Ordering};
37use std::sync::Arc;
38
39use parking_lot::RwLock;
40
41// ─────────────────────────────────────────────────────────────────────────────
42// Constants
43// ─────────────────────────────────────────────────────────────────────────────
44
45/// Default DHT provider record TTL (24 hours in seconds, per Kademlia spec).
46pub const DEFAULT_PROVIDER_TTL_SECS: u64 = 86_400;
47
48/// Renew when this fraction of TTL has elapsed (80 %).
49pub const DEFAULT_RENEWAL_THRESHOLD: f64 = 0.80;
50
51/// Default maximum number of CIDs to track simultaneously.
52const DEFAULT_MAX_TRACKED_CIDS: usize = 100_000;
53
54// ─────────────────────────────────────────────────────────────────────────────
55// RenewalConfig
56// ─────────────────────────────────────────────────────────────────────────────
57
58/// Configuration for the [`ProviderRenewalScheduler`].
59#[derive(Debug, Clone)]
60pub struct RenewalConfig {
61    /// Provider record TTL in seconds.  Default: 86 400 (24 h).
62    pub ttl_secs: u64,
63
64    /// Fraction of TTL at which to trigger renewal.  Default: 0.80.
65    pub renewal_threshold: f64,
66
67    /// Maximum CIDs to track simultaneously.  Default: 100 000.
68    pub max_tracked_cids: usize,
69}
70
71impl Default for RenewalConfig {
72    fn default() -> Self {
73        Self {
74            ttl_secs: DEFAULT_PROVIDER_TTL_SECS,
75            renewal_threshold: DEFAULT_RENEWAL_THRESHOLD,
76            max_tracked_cids: DEFAULT_MAX_TRACKED_CIDS,
77        }
78    }
79}
80
81impl RenewalConfig {
82    /// Number of seconds after which a record should be renewed.
83    ///
84    /// Computed as `ttl_secs * renewal_threshold`, truncated to an integer.
85    pub fn renewal_interval_secs(&self) -> u64 {
86        (self.ttl_secs as f64 * self.renewal_threshold) as u64
87    }
88}
89
90// ─────────────────────────────────────────────────────────────────────────────
91// ProviderRecord
92// ─────────────────────────────────────────────────────────────────────────────
93
94/// State for a single provided CID tracked by the scheduler.
95#[derive(Debug, Clone)]
96pub struct ProviderRecord {
97    /// String representation of the CID.
98    pub cid_str: String,
99
100    /// Unix timestamp (seconds) at which `track()` was first called.
101    pub provided_at_secs: u64,
102
103    /// Unix timestamp (seconds) of the most recent successful renewal.
104    /// `0` means the record has never been explicitly renewed since tracking
105    /// began.
106    pub last_renewal_secs: u64,
107
108    /// How many times this record has been successfully renewed.
109    pub renewal_count: u64,
110}
111
112impl ProviderRecord {
113    /// Create a new record anchored at `now_secs`.
114    pub fn new(cid_str: impl Into<String>, now_secs: u64) -> Self {
115        Self {
116            cid_str: cid_str.into(),
117            provided_at_secs: now_secs,
118            last_renewal_secs: 0,
119            renewal_count: 0,
120        }
121    }
122
123    /// Returns `true` when the record is old enough to require renewal.
124    ///
125    /// Age is measured from the *later* of `provided_at_secs` and
126    /// `last_renewal_secs` (treating `0` as "never renewed").
127    pub fn needs_renewal(&self, now_secs: u64, config: &RenewalConfig) -> bool {
128        let baseline = if self.last_renewal_secs == 0 {
129            self.provided_at_secs
130        } else {
131            self.last_renewal_secs.max(self.provided_at_secs)
132        };
133        let age = now_secs.saturating_sub(baseline);
134        age >= config.renewal_interval_secs()
135    }
136
137    /// Update the record after a successful renewal.
138    pub fn mark_renewed(&mut self, now_secs: u64) {
139        self.last_renewal_secs = now_secs;
140        self.renewal_count += 1;
141    }
142}
143
144// ─────────────────────────────────────────────────────────────────────────────
145// ProviderRenewalScheduler
146// ─────────────────────────────────────────────────────────────────────────────
147
148/// Manages scheduled renewal of DHT provider records.
149///
150/// The scheduler is cheap to clone because it is `Arc`-wrapped internally.
151/// All methods take `&self` to allow concurrent access from multiple tasks.
152pub struct ProviderRenewalScheduler {
153    config: RenewalConfig,
154    records: RwLock<HashMap<String, ProviderRecord>>,
155    total_renewals: AtomicU64,
156}
157
158impl ProviderRenewalScheduler {
159    /// Create a new scheduler wrapped in an [`Arc`].
160    pub fn new(config: RenewalConfig) -> Arc<Self> {
161        Arc::new(Self {
162            config,
163            records: RwLock::new(HashMap::new()),
164            total_renewals: AtomicU64::new(0),
165        })
166    }
167
168    /// Register a CID for renewal tracking.
169    ///
170    /// Returns `true` if the CID was newly inserted or was already tracked.
171    /// Returns `false` when the capacity limit is reached and the CID is new.
172    pub fn track(&self, cid_str: impl Into<String>, now_secs: u64) -> bool {
173        let key: String = cid_str.into();
174
175        // Fast-path: already tracked — just update the timestamp.
176        {
177            let mut records = self.records.write();
178            if records.contains_key(&key) {
179                // Refresh the anchor timestamp so existing records remain valid.
180                if let Some(r) = records.get_mut(&key) {
181                    r.provided_at_secs = now_secs;
182                }
183                return true;
184            }
185
186            // Capacity check.
187            if records.len() >= self.config.max_tracked_cids {
188                return false;
189            }
190
191            records.insert(key.clone(), ProviderRecord::new(key, now_secs));
192        }
193
194        true
195    }
196
197    /// Stop tracking a CID (e.g. when the corresponding block is deleted).
198    ///
199    /// Returns `true` if the record existed and was removed.
200    pub fn untrack(&self, cid_str: &str) -> bool {
201        self.records.write().remove(cid_str).is_some()
202    }
203
204    /// Returns all CIDs that are due for renewal at `now_secs`.
205    pub fn due_for_renewal(&self, now_secs: u64) -> Vec<String> {
206        self.records
207            .read()
208            .values()
209            .filter(|r| r.needs_renewal(now_secs, &self.config))
210            .map(|r| r.cid_str.clone())
211            .collect()
212    }
213
214    /// Mark a slice of CIDs as successfully renewed at `now_secs`.
215    ///
216    /// Returns the number of records that were actually updated (i.e. were
217    /// present in the tracker).
218    pub fn mark_renewed(&self, cid_strs: &[String], now_secs: u64) -> usize {
219        let mut count = 0usize;
220        let mut records = self.records.write();
221
222        for key in cid_strs {
223            if let Some(r) = records.get_mut(key.as_str()) {
224                r.mark_renewed(now_secs);
225                count += 1;
226            }
227        }
228
229        if count > 0 {
230            self.total_renewals
231                .fetch_add(count as u64, Ordering::Relaxed);
232        }
233
234        count
235    }
236
237    /// Total number of CIDs currently tracked.
238    pub fn tracked_count(&self) -> usize {
239        self.records.read().len()
240    }
241
242    /// Total successful renewals across all CIDs since the scheduler was
243    /// created.
244    pub fn total_renewals(&self) -> u64 {
245        self.total_renewals.load(Ordering::Relaxed)
246    }
247
248    /// Return a point-in-time snapshot of all tracked records.
249    pub fn snapshot(&self) -> Vec<ProviderRecord> {
250        self.records.read().values().cloned().collect()
251    }
252
253    /// Remove records that are definitively stale: the record has *never* been
254    /// renewed (`last_renewal_secs == 0`) and its age exceeds `2 × TTL`.
255    ///
256    /// Returns the number of records pruned.
257    pub fn prune_stale(&self, now_secs: u64) -> usize {
258        let cutoff = self.config.ttl_secs.saturating_mul(2);
259        let mut records = self.records.write();
260        let before = records.len();
261
262        records.retain(|_, r| {
263            if r.last_renewal_secs == 0 {
264                let age = now_secs.saturating_sub(r.provided_at_secs);
265                age <= cutoff
266            } else {
267                true
268            }
269        });
270
271        before - records.len()
272    }
273}
274
275// ─────────────────────────────────────────────────────────────────────────────
276// Tests
277// ─────────────────────────────────────────────────────────────────────────────
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282
283    // Convenience: epoch-like base timestamp that is large enough to avoid
284    // saturating-sub oddities.
285    const T0: u64 = 1_700_000_000;
286
287    fn default_scheduler() -> Arc<ProviderRenewalScheduler> {
288        ProviderRenewalScheduler::new(RenewalConfig::default())
289    }
290
291    // ── Config ────────────────────────────────────────────────────────────────
292
293    #[test]
294    fn test_default_config_renewal_interval() {
295        let cfg = RenewalConfig::default();
296        // 86_400 * 0.80 = 69_120
297        assert_eq!(cfg.renewal_interval_secs(), 69_120);
298    }
299
300    #[test]
301    fn test_renewal_config_custom() {
302        let cfg = RenewalConfig {
303            ttl_secs: 3_600,
304            renewal_threshold: 0.75,
305            max_tracked_cids: 50,
306        };
307        // 3600 * 0.75 = 2700
308        assert_eq!(cfg.renewal_interval_secs(), 2_700);
309        assert_eq!(cfg.max_tracked_cids, 50);
310    }
311
312    // ── Track / Untrack ───────────────────────────────────────────────────────
313
314    #[test]
315    fn test_track_and_count() {
316        let sched = default_scheduler();
317        assert_eq!(sched.tracked_count(), 0);
318
319        assert!(sched.track("cid1", T0));
320        assert!(sched.track("cid2", T0));
321        assert_eq!(sched.tracked_count(), 2);
322    }
323
324    #[test]
325    fn test_untrack() {
326        let sched = default_scheduler();
327        sched.track("cid1", T0);
328        assert_eq!(sched.tracked_count(), 1);
329
330        assert!(sched.untrack("cid1"));
331        assert_eq!(sched.tracked_count(), 0);
332
333        // Removing a non-existent key returns false.
334        assert!(!sched.untrack("cid1"));
335    }
336
337    // ── Renewal detection ─────────────────────────────────────────────────────
338
339    #[test]
340    fn test_due_for_renewal_empty_initially() {
341        let sched = default_scheduler();
342        sched.track("cid1", T0);
343
344        // Immediately after tracking, nothing should be due.
345        let due = sched.due_for_renewal(T0);
346        assert!(
347            due.is_empty(),
348            "expected no renewals immediately: got {due:?}"
349        );
350    }
351
352    #[test]
353    fn test_due_for_renewal_after_threshold() {
354        let sched = default_scheduler();
355        sched.track("cid1", T0);
356
357        // Advance time by exactly the renewal interval.
358        let interval = sched.config.renewal_interval_secs();
359        let due = sched.due_for_renewal(T0 + interval);
360        assert_eq!(due.len(), 1);
361        assert_eq!(due[0], "cid1");
362    }
363
364    #[test]
365    fn test_mark_renewed_resets_timer() {
366        let sched = default_scheduler();
367        sched.track("cid1", T0);
368
369        let interval = sched.config.renewal_interval_secs();
370        let t1 = T0 + interval; // due now
371
372        // Confirm it is due.
373        assert_eq!(sched.due_for_renewal(t1).len(), 1);
374
375        // Renew it.
376        sched.mark_renewed(&["cid1".to_string()], t1);
377
378        // Should no longer be due right after renewal.
379        assert!(sched.due_for_renewal(t1).is_empty());
380
381        // Should become due again after another full interval.
382        let t2 = t1 + interval;
383        assert_eq!(sched.due_for_renewal(t2).len(), 1);
384    }
385
386    #[test]
387    fn test_mark_renewed_increments_count() {
388        let sched = default_scheduler();
389        sched.track("cid1", T0);
390
391        let interval = sched.config.renewal_interval_secs();
392
393        // Renew twice.
394        sched.mark_renewed(&["cid1".to_string()], T0 + interval);
395        sched.mark_renewed(&["cid1".to_string()], T0 + interval * 2);
396
397        let snap = sched.snapshot();
398        let rec = snap
399            .iter()
400            .find(|r| r.cid_str == "cid1")
401            .expect("record missing");
402        assert_eq!(rec.renewal_count, 2);
403    }
404
405    // ── Total-renewals counter ─────────────────────────────────────────────────
406
407    #[test]
408    fn test_total_renewals_counter() {
409        let sched = default_scheduler();
410        sched.track("cid1", T0);
411        sched.track("cid2", T0);
412        sched.track("cid3", T0);
413
414        assert_eq!(sched.total_renewals(), 0);
415
416        let interval = sched.config.renewal_interval_secs();
417        let renewed = sched.mark_renewed(
418            &["cid1".to_string(), "cid2".to_string(), "cid3".to_string()],
419            T0 + interval,
420        );
421        assert_eq!(renewed, 3);
422        assert_eq!(sched.total_renewals(), 3);
423
424        // Second batch.
425        let renewed2 = sched.mark_renewed(&["cid1".to_string()], T0 + interval * 2);
426        assert_eq!(renewed2, 1);
427        assert_eq!(sched.total_renewals(), 4);
428    }
429
430    // ── Capacity limit ────────────────────────────────────────────────────────
431
432    #[test]
433    fn test_capacity_limit() {
434        let cfg = RenewalConfig {
435            max_tracked_cids: 3,
436            ..Default::default()
437        };
438        let sched = ProviderRenewalScheduler::new(cfg);
439
440        assert!(sched.track("cid1", T0));
441        assert!(sched.track("cid2", T0));
442        assert!(sched.track("cid3", T0));
443
444        // 4th new CID must be rejected.
445        assert!(!sched.track("cid4", T0));
446        assert_eq!(sched.tracked_count(), 3);
447
448        // Existing CIDs must still be re-trackable (idempotent).
449        assert!(sched.track("cid1", T0 + 1));
450        assert_eq!(sched.tracked_count(), 3);
451    }
452
453    // ── Prune stale ───────────────────────────────────────────────────────────
454
455    #[test]
456    fn test_prune_stale_removes_old_records() {
457        let sched = default_scheduler();
458        let ttl = sched.config.ttl_secs;
459
460        sched.track("old_cid", T0);
461        sched.track("fresh_cid", T0);
462
463        // Advance time so "old_cid" is older than 2×TTL but "fresh_cid" is not.
464        let old_enough = T0 + ttl * 2 + 1;
465
466        // Manually update fresh_cid so it was tracked recently.
467        sched.track("fresh_cid", old_enough - 100); // re-track with recent time
468
469        let pruned = sched.prune_stale(old_enough);
470        assert_eq!(pruned, 1, "expected exactly 1 stale record pruned");
471        assert_eq!(sched.tracked_count(), 1);
472
473        let snap = sched.snapshot();
474        assert_eq!(snap[0].cid_str, "fresh_cid");
475    }
476
477    // ── Snapshot ─────────────────────────────────────────────────────────────
478
479    #[test]
480    fn test_snapshot() {
481        let sched = default_scheduler();
482        sched.track("a", T0);
483        sched.track("b", T0);
484        sched.track("c", T0);
485
486        let snap = sched.snapshot();
487        assert_eq!(snap.len(), 3);
488
489        // All CIDs must be present.
490        let mut cids: Vec<&str> = snap.iter().map(|r| r.cid_str.as_str()).collect();
491        cids.sort_unstable();
492        assert_eq!(cids, vec!["a", "b", "c"]);
493    }
494}