1use std::collections::HashMap;
36use std::sync::atomic::{AtomicU64, Ordering};
37use std::sync::Arc;
38
39use parking_lot::RwLock;
40
41pub const DEFAULT_PROVIDER_TTL_SECS: u64 = 86_400;
47
48pub const DEFAULT_RENEWAL_THRESHOLD: f64 = 0.80;
50
51const DEFAULT_MAX_TRACKED_CIDS: usize = 100_000;
53
54#[derive(Debug, Clone)]
60pub struct RenewalConfig {
61 pub ttl_secs: u64,
63
64 pub renewal_threshold: f64,
66
67 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 pub fn renewal_interval_secs(&self) -> u64 {
86 (self.ttl_secs as f64 * self.renewal_threshold) as u64
87 }
88}
89
90#[derive(Debug, Clone)]
96pub struct ProviderRecord {
97 pub cid_str: String,
99
100 pub provided_at_secs: u64,
102
103 pub last_renewal_secs: u64,
107
108 pub renewal_count: u64,
110}
111
112impl ProviderRecord {
113 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 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 pub fn mark_renewed(&mut self, now_secs: u64) {
139 self.last_renewal_secs = now_secs;
140 self.renewal_count += 1;
141 }
142}
143
144pub struct ProviderRenewalScheduler {
153 config: RenewalConfig,
154 records: RwLock<HashMap<String, ProviderRecord>>,
155 total_renewals: AtomicU64,
156}
157
158impl ProviderRenewalScheduler {
159 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 pub fn track(&self, cid_str: impl Into<String>, now_secs: u64) -> bool {
173 let key: String = cid_str.into();
174
175 {
177 let mut records = self.records.write();
178 if records.contains_key(&key) {
179 if let Some(r) = records.get_mut(&key) {
181 r.provided_at_secs = now_secs;
182 }
183 return true;
184 }
185
186 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 pub fn untrack(&self, cid_str: &str) -> bool {
201 self.records.write().remove(cid_str).is_some()
202 }
203
204 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 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 pub fn tracked_count(&self) -> usize {
239 self.records.read().len()
240 }
241
242 pub fn total_renewals(&self) -> u64 {
245 self.total_renewals.load(Ordering::Relaxed)
246 }
247
248 pub fn snapshot(&self) -> Vec<ProviderRecord> {
250 self.records.read().values().cloned().collect()
251 }
252
253 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#[cfg(test)]
280mod tests {
281 use super::*;
282
283 const T0: u64 = 1_700_000_000;
286
287 fn default_scheduler() -> Arc<ProviderRenewalScheduler> {
288 ProviderRenewalScheduler::new(RenewalConfig::default())
289 }
290
291 #[test]
294 fn test_default_config_renewal_interval() {
295 let cfg = RenewalConfig::default();
296 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 assert_eq!(cfg.renewal_interval_secs(), 2_700);
309 assert_eq!(cfg.max_tracked_cids, 50);
310 }
311
312 #[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 assert!(!sched.untrack("cid1"));
335 }
336
337 #[test]
340 fn test_due_for_renewal_empty_initially() {
341 let sched = default_scheduler();
342 sched.track("cid1", T0);
343
344 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 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; assert_eq!(sched.due_for_renewal(t1).len(), 1);
374
375 sched.mark_renewed(&["cid1".to_string()], t1);
377
378 assert!(sched.due_for_renewal(t1).is_empty());
380
381 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 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 #[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 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 #[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 assert!(!sched.track("cid4", T0));
446 assert_eq!(sched.tracked_count(), 3);
447
448 assert!(sched.track("cid1", T0 + 1));
450 assert_eq!(sched.tracked_count(), 3);
451 }
452
453 #[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 let old_enough = T0 + ttl * 2 + 1;
465
466 sched.track("fresh_cid", old_enough - 100); 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 #[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 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}