grafos_leasekit/
manager.rs

1//! The [`RenewalManager`] — poll-driven lease renewal engine.
2
3extern crate alloc;
4use alloc::boxed::Box;
5use alloc::rc::Rc;
6use alloc::vec::Vec;
7use core::cell::Cell;
8use core::fmt;
9
10use crate::backoff::Backoff;
11use crate::policy::RenewalPolicy;
12pub use grafos_core::RevokeState;
13
14/// Workload-facing handle for observing a registered lease.
15///
16/// A handle is intentionally narrow: it exposes the lease id and the
17/// canonical revoke state for that lease, but it does not own renewal,
18/// admission, or scheduler policy.
19#[derive(Clone, Debug)]
20pub struct LeaseHandle {
21    lease_id: u128,
22    revoke_state: Rc<Cell<RevokeState>>,
23}
24
25impl LeaseHandle {
26    fn new(lease_id: u128) -> Self {
27        Self {
28            lease_id,
29            revoke_state: Rc::new(Cell::new(RevokeState::Active)),
30        }
31    }
32
33    /// Returns the lease identifier associated with this handle.
34    pub fn lease_id(&self) -> u128 {
35        self.lease_id
36    }
37
38    /// Returns the latest revoke state observed by the renewal manager.
39    pub fn revoke_state(&self) -> RevokeState {
40        self.revoke_state.get()
41    }
42
43    /// Returns `true` when the lease has reached a terminal revoke state.
44    pub fn is_revoke_terminal(&self) -> bool {
45        self.revoke_state().is_terminal()
46    }
47}
48
49/// Error returned when a lease revoke-state transition cannot be applied.
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub enum LeaseRevokeTransitionError {
52    /// The lease is not registered with this manager.
53    LeaseNotRegistered { lease_id: u128 },
54    /// The requested direct state transition is not legal for `RevokeState`.
55    IllegalTransition {
56        lease_id: u128,
57        from: RevokeState,
58        to: RevokeState,
59    },
60}
61
62impl fmt::Display for LeaseRevokeTransitionError {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        match self {
65            Self::LeaseNotRegistered { lease_id } => {
66                write!(f, "lease {lease_id} is not registered")
67            }
68            Self::IllegalTransition { lease_id, from, to } => {
69                write!(
70                    f,
71                    "illegal revoke-state transition for lease {lease_id}: {from} -> {to}"
72                )
73            }
74        }
75    }
76}
77
78#[cfg(feature = "std")]
79impl std::error::Error for LeaseRevokeTransitionError {}
80
81/// Tracks the renewal state for a single registered lease.
82struct Entry {
83    lease_id: u128,
84    created_at: u64,
85    expires_at: u64,
86    policy: RenewalPolicy,
87    backoff: Backoff,
88    /// Absolute time before which we should not retry after a failure.
89    next_retry_at: u64,
90    /// Whether the lease has been successfully renewed at least once by
91    /// this manager (used to compute renewal deadline from latest expiry).
92    last_renewed_at: Option<u64>,
93    /// Shared workload-facing revoke state for this registered lease.
94    revoke_state: Rc<Cell<RevokeState>>,
95}
96
97impl Entry {
98    fn renewal_deadline(&self) -> u64 {
99        let created = self.last_renewed_at.unwrap_or(self.created_at);
100        let jitter_seed = self.lease_id as u64;
101        self.policy
102            .renewal_deadline(created, self.expires_at, jitter_seed)
103    }
104
105    fn is_near_expiry(&self, now: u64) -> bool {
106        if now >= self.expires_at {
107            return true;
108        }
109        let remaining = self.expires_at - now;
110        let total_ttl = self
111            .expires_at
112            .saturating_sub(self.last_renewed_at.unwrap_or(self.created_at));
113        if total_ttl == 0 {
114            return true;
115        }
116        // Near expiry when less than 10% TTL remains
117        remaining < total_ttl / 10
118    }
119
120    fn handle(&self) -> LeaseHandle {
121        LeaseHandle {
122            lease_id: self.lease_id,
123            revoke_state: Rc::clone(&self.revoke_state),
124        }
125    }
126
127    fn transition_revoke_state(
128        &self,
129        next: RevokeState,
130    ) -> Result<RevokeState, LeaseRevokeTransitionError> {
131        let current = self.revoke_state.get();
132        if current == next {
133            return Ok(current);
134        }
135        if !current.legal_transition_to(next) {
136            return Err(LeaseRevokeTransitionError::IllegalTransition {
137                lease_id: self.lease_id,
138                from: current,
139                to: next,
140            });
141        }
142        self.revoke_state.set(next);
143        Ok(next)
144    }
145}
146
147/// Summary returned by [`RenewalManager::tick`].
148#[derive(Debug, Clone, Default)]
149pub struct RenewalSummary {
150    /// Number of leases successfully renewed this tick.
151    pub renewed: u32,
152    /// Number of leases checked but not yet due for renewal.
153    pub skipped: u32,
154    /// Number of leases where renewal was attempted but failed.
155    pub failed: u32,
156    /// Lease IDs that are near expiry (< 10% TTL remaining).
157    pub near_expiry: Vec<u128>,
158}
159
160/// Callback invoked by the manager to actually perform a renewal.
161///
162/// The manager itself does not hold lease objects; callers provide a
163/// `renew_fn` closure to `tick_with` that performs the renewal and
164/// returns the new `expires_at` timestamp on success.
165///
166/// Alternatively, use [`tick`](RenewalManager::tick) for a simpler API
167/// that auto-renews by extending `expires_at` by `policy.min_renew_secs`.
168pub struct RenewalManager {
169    entries: Vec<Entry>,
170    revoke_callbacks: Vec<Box<dyn Fn(u128, u8)>>,
171}
172
173impl RenewalManager {
174    /// Create an empty manager with no registered leases.
175    pub fn new() -> Self {
176        RenewalManager {
177            entries: Vec::new(),
178            revoke_callbacks: Vec::new(),
179        }
180    }
181
182    /// Register a lease for managed renewal.
183    ///
184    /// `created_at` is inferred as `expires_at - policy.min_renew_secs`
185    /// (clamped to 0). For precise control, use `register_with_created_at`.
186    pub fn register(&mut self, lease_id: u128, expires_at: u64, policy: RenewalPolicy) {
187        let created_at = expires_at.saturating_sub(policy.min_renew_secs);
188        let _ = self.register_entry(lease_id, created_at, expires_at, policy);
189    }
190
191    /// Register a lease and return a workload-facing handle for observing it.
192    pub fn register_handle(
193        &mut self,
194        lease_id: u128,
195        expires_at: u64,
196        policy: RenewalPolicy,
197    ) -> LeaseHandle {
198        let created_at = expires_at.saturating_sub(policy.min_renew_secs);
199        self.register_entry(lease_id, created_at, expires_at, policy)
200    }
201
202    /// Register a lease with an explicit creation timestamp.
203    pub fn register_with_created_at(
204        &mut self,
205        lease_id: u128,
206        created_at: u64,
207        expires_at: u64,
208        policy: RenewalPolicy,
209    ) {
210        let _ = self.register_entry(lease_id, created_at, expires_at, policy);
211    }
212
213    /// Register a lease with an explicit creation timestamp and return a
214    /// workload-facing handle for observing it.
215    pub fn register_with_created_at_handle(
216        &mut self,
217        lease_id: u128,
218        created_at: u64,
219        expires_at: u64,
220        policy: RenewalPolicy,
221    ) -> LeaseHandle {
222        self.register_entry(lease_id, created_at, expires_at, policy)
223    }
224
225    fn register_entry(
226        &mut self,
227        lease_id: u128,
228        created_at: u64,
229        expires_at: u64,
230        policy: RenewalPolicy,
231    ) -> LeaseHandle {
232        // Remove any existing entry with the same lease_id
233        self.entries.retain(|e| e.lease_id != lease_id);
234        let backoff = Backoff::new(1, policy.max_backoff_secs);
235        let handle = LeaseHandle::new(lease_id);
236        self.entries.push(Entry {
237            lease_id,
238            created_at,
239            expires_at,
240            policy,
241            backoff,
242            next_retry_at: 0,
243            last_renewed_at: None,
244            revoke_state: Rc::clone(&handle.revoke_state),
245        });
246        handle
247    }
248
249    /// Unregister a lease.
250    pub fn unregister(&mut self, lease_id: u128) {
251        self.entries.retain(|e| e.lease_id != lease_id);
252    }
253
254    /// Return a workload-facing handle for a managed lease, if registered.
255    pub fn handle(&self, lease_id: u128) -> Option<LeaseHandle> {
256        self.entries
257            .iter()
258            .find(|e| e.lease_id == lease_id)
259            .map(Entry::handle)
260    }
261
262    /// Return the current revoke state for a managed lease, if registered.
263    pub fn revoke_state(&self, lease_id: u128) -> Option<RevokeState> {
264        self.handle(lease_id).map(|handle| handle.revoke_state())
265    }
266
267    /// Number of leases currently managed.
268    pub fn len(&self) -> usize {
269        self.entries.len()
270    }
271
272    /// Returns `true` if no leases are managed.
273    pub fn is_empty(&self) -> bool {
274        self.entries.is_empty()
275    }
276
277    /// Drive renewals using a simple model: if a lease is due, extend
278    /// `expires_at` by `policy.min_renew_secs`. This is suitable when the
279    /// manager owns the lease state (e.g. testing or self-contained use).
280    pub fn tick(&mut self, now_unix_secs: u64) -> RenewalSummary {
281        self.tick_with(now_unix_secs, |_lease_id, duration| {
282            // Simple model: always succeed, return new expiry
283            Ok(duration)
284        })
285    }
286
287    /// Drive renewals with a caller-provided renewal function.
288    ///
289    /// `renew_fn(lease_id, duration_secs)` should attempt the actual
290    /// renewal and return the new `expires_at` on success, or an error.
291    pub fn tick_with<F>(&mut self, now_unix_secs: u64, mut renew_fn: F) -> RenewalSummary
292    where
293        F: FnMut(u128, u64) -> Result<u64, ()>,
294    {
295        let mut summary = RenewalSummary::default();
296
297        for entry in self.entries.iter_mut() {
298            // Check near-expiry
299            if entry.is_near_expiry(now_unix_secs) {
300                summary.near_expiry.push(entry.lease_id);
301            }
302
303            // Already expired — skip
304            if now_unix_secs >= entry.expires_at {
305                if entry.revoke_state.get() == RevokeState::Active {
306                    let _ = entry.transition_revoke_state(RevokeState::Expired);
307                }
308                summary.skipped += 1;
309                continue;
310            }
311
312            // Not yet at renewal deadline — skip
313            let deadline = entry.renewal_deadline();
314            if now_unix_secs < deadline {
315                summary.skipped += 1;
316                continue;
317            }
318
319            // In backoff — skip
320            if now_unix_secs < entry.next_retry_at {
321                summary.skipped += 1;
322                continue;
323            }
324
325            // Attempt renewal
326            let duration = entry.policy.min_renew_secs;
327            match renew_fn(entry.lease_id, duration) {
328                Ok(new_expires_at) => {
329                    let new_exp = now_unix_secs.saturating_add(new_expires_at);
330                    entry.last_renewed_at = Some(now_unix_secs);
331                    entry.expires_at = new_exp;
332                    entry.backoff.reset();
333                    entry.next_retry_at = 0;
334                    summary.renewed += 1;
335
336                    #[cfg(feature = "observe")]
337                    emit_renew_success(entry.lease_id);
338                }
339                Err(()) => {
340                    let delay = entry.backoff.next_delay();
341                    entry.next_retry_at = now_unix_secs.saturating_add(delay);
342                    summary.failed += 1;
343
344                    #[cfg(feature = "observe")]
345                    emit_renew_failure(entry.lease_id);
346                }
347            }
348        }
349
350        #[cfg(feature = "observe")]
351        emit_tick_metrics(&summary, self.entries.len());
352
353        summary
354    }
355
356    /// Returns `true` if the given lease is near expiry (< 10% TTL remaining).
357    pub fn is_near_expiry(&self, lease_id: u128, now: u64) -> bool {
358        self.entries
359            .iter()
360            .find(|e| e.lease_id == lease_id)
361            .map(|e| e.is_near_expiry(now))
362            .unwrap_or(false)
363    }
364
365    /// Register a callback invoked when a lease is revoked (e.g. via
366    /// WITHDRAW cleanup or REVOKE_BROADCAST).
367    ///
368    /// The callback receives the lease ID and a reason code (0 = unknown/timeout).
369    pub fn on_revoked<F: Fn(u128, u8) + 'static>(&mut self, callback: F) {
370        self.revoke_callbacks.push(Box::new(callback));
371    }
372
373    /// Apply a typed revoke-state transition to a managed lease.
374    ///
375    /// Duplicate transitions are treated as idempotent observations. Illegal
376    /// direct transitions are rejected and leave the previous state intact.
377    pub fn transition_revoke_state(
378        &self,
379        lease_id: u128,
380        next: RevokeState,
381    ) -> Result<RevokeState, LeaseRevokeTransitionError> {
382        let entry = self
383            .entries
384            .iter()
385            .find(|e| e.lease_id == lease_id)
386            .ok_or(LeaseRevokeTransitionError::LeaseNotRegistered { lease_id })?;
387        entry.transition_revoke_state(next)
388    }
389
390    /// Notify all registered revocation callbacks for the given lease.
391    ///
392    /// `reason` is the WITHDRAW reason code (0 = timeout/unknown).
393    pub fn notify_revoked(&self, lease_id: u128, reason: u8) {
394        if let Some(entry) = self.entries.iter().find(|e| e.lease_id == lease_id) {
395            if entry.revoke_state.get() == RevokeState::Active {
396                let _ = entry.transition_revoke_state(RevokeState::RevokeWarning);
397            }
398        }
399        for cb in &self.revoke_callbacks {
400            cb(lease_id, reason);
401        }
402    }
403
404    /// Returns the current `expires_at` for a managed lease, if registered.
405    pub fn expires_at(&self, lease_id: u128) -> Option<u64> {
406        self.entries
407            .iter()
408            .find(|e| e.lease_id == lease_id)
409            .map(|e| e.expires_at)
410    }
411}
412
413impl Default for RenewalManager {
414    fn default() -> Self {
415        Self::new()
416    }
417}
418
419#[cfg(feature = "observe")]
420fn emit_renew_success(lease_id: u128) {
421    let _ = lease_id;
422    grafos_observe::FabricMetrics::global().ops_total.inc();
423}
424
425#[cfg(feature = "observe")]
426fn emit_renew_failure(lease_id: u128) {
427    let _ = lease_id;
428    grafos_observe::FabricMetrics::global().ops_total.inc();
429}
430
431#[cfg(feature = "observe")]
432fn emit_tick_metrics(summary: &RenewalSummary, managed_count: usize) {
433    let _ = (summary, managed_count);
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439
440    #[test]
441    fn register_and_tick_before_deadline() {
442        let mut mgr = RenewalManager::new();
443        let policy = RenewalPolicy {
444            renew_at_fraction: 0.75,
445            jitter_fraction: 0.0,
446            min_renew_secs: 100,
447            max_backoff_secs: 5,
448        };
449        // created_at=1000, expires_at=1100, TTL=100
450        mgr.register_with_created_at(1, 1000, 1100, policy);
451
452        // At t=1050, only 50% elapsed — should not renew
453        let s = mgr.tick(1050);
454        assert_eq!(s.renewed, 0);
455        assert_eq!(s.skipped, 1);
456        assert_eq!(s.failed, 0);
457    }
458
459    #[test]
460    fn tick_at_deadline_triggers_renewal() {
461        let mut mgr = RenewalManager::new();
462        let policy = RenewalPolicy {
463            renew_at_fraction: 0.75,
464            jitter_fraction: 0.0,
465            min_renew_secs: 100,
466            max_backoff_secs: 5,
467        };
468        mgr.register_with_created_at(1, 1000, 1100, policy);
469
470        // At t=1075 (exactly 75% elapsed), should renew
471        let s = mgr.tick(1075);
472        assert_eq!(s.renewed, 1);
473        assert_eq!(s.skipped, 0);
474    }
475
476    #[test]
477    fn tick_after_expiry_skips() {
478        let mut mgr = RenewalManager::new();
479        let policy = RenewalPolicy {
480            renew_at_fraction: 0.75,
481            jitter_fraction: 0.0,
482            min_renew_secs: 100,
483            max_backoff_secs: 5,
484        };
485        mgr.register_with_created_at(1, 1000, 1100, policy);
486
487        let s = mgr.tick(1200);
488        assert_eq!(s.renewed, 0);
489        assert_eq!(s.skipped, 1);
490    }
491
492    #[test]
493    fn register_handle_exposes_active_revoke_state() {
494        let mut mgr = RenewalManager::new();
495        let handle = mgr.register_handle(1, 1100, RenewalPolicy::default());
496
497        assert_eq!(handle.lease_id(), 1);
498        assert_eq!(handle.revoke_state(), RevokeState::Active);
499        assert_eq!(mgr.revoke_state(1), Some(RevokeState::Active));
500        assert!(!handle.is_revoke_terminal());
501    }
502
503    #[test]
504    fn handle_observes_cooperative_revoke_path() {
505        let mut mgr = RenewalManager::new();
506        let handle = mgr.register_with_created_at_handle(1, 1000, 1100, RenewalPolicy::default());
507
508        assert_eq!(
509            mgr.transition_revoke_state(1, RevokeState::RevokeWarning),
510            Ok(RevokeState::RevokeWarning)
511        );
512        assert_eq!(handle.revoke_state(), RevokeState::RevokeWarning);
513
514        assert_eq!(
515            mgr.transition_revoke_state(1, RevokeState::GraceRunning),
516            Ok(RevokeState::GraceRunning)
517        );
518        assert_eq!(handle.revoke_state(), RevokeState::GraceRunning);
519
520        assert_eq!(
521            mgr.transition_revoke_state(1, RevokeState::CheckpointReported),
522            Ok(RevokeState::CheckpointReported)
523        );
524        assert_eq!(handle.revoke_state(), RevokeState::CheckpointReported);
525
526        assert_eq!(
527            mgr.transition_revoke_state(1, RevokeState::Torndown),
528            Ok(RevokeState::Torndown)
529        );
530        assert_eq!(handle.revoke_state(), RevokeState::Torndown);
531        assert!(handle.is_revoke_terminal());
532    }
533
534    #[test]
535    fn handle_observes_forced_teardown_path() {
536        let mut mgr = RenewalManager::new();
537        let handle = mgr.register_with_created_at_handle(1, 1000, 1100, RenewalPolicy::default());
538
539        assert_eq!(
540            mgr.transition_revoke_state(1, RevokeState::ForcedTeardown),
541            Ok(RevokeState::ForcedTeardown)
542        );
543        assert_eq!(handle.revoke_state(), RevokeState::ForcedTeardown);
544
545        assert_eq!(
546            mgr.transition_revoke_state(1, RevokeState::Torndown),
547            Ok(RevokeState::Torndown)
548        );
549        assert_eq!(handle.revoke_state(), RevokeState::Torndown);
550    }
551
552    #[test]
553    fn handle_observes_expiry_when_active_lease_ticks_past_ttl() {
554        let mut mgr = RenewalManager::new();
555        let handle = mgr.register_with_created_at_handle(1, 1000, 1100, RenewalPolicy::default());
556
557        let s = mgr.tick(1100);
558
559        assert_eq!(s.skipped, 1);
560        assert_eq!(handle.revoke_state(), RevokeState::Expired);
561        assert!(handle.is_revoke_terminal());
562    }
563
564    #[test]
565    fn transition_rejects_illegal_direct_path_without_mutating() {
566        let mut mgr = RenewalManager::new();
567        let handle = mgr.register_with_created_at_handle(1, 1000, 1100, RenewalPolicy::default());
568
569        let err = mgr
570            .transition_revoke_state(1, RevokeState::GraceRunning)
571            .unwrap_err();
572
573        assert_eq!(
574            err,
575            LeaseRevokeTransitionError::IllegalTransition {
576                lease_id: 1,
577                from: RevokeState::Active,
578                to: RevokeState::GraceRunning,
579            }
580        );
581        assert_eq!(handle.revoke_state(), RevokeState::Active);
582    }
583
584    #[test]
585    fn transition_unknown_lease_returns_typed_error() {
586        let mgr = RenewalManager::new();
587
588        assert_eq!(
589            mgr.transition_revoke_state(404, RevokeState::RevokeWarning),
590            Err(LeaseRevokeTransitionError::LeaseNotRegistered { lease_id: 404 })
591        );
592    }
593
594    #[test]
595    fn tick_with_failure_triggers_backoff() {
596        let mut mgr = RenewalManager::new();
597        let policy = RenewalPolicy {
598            renew_at_fraction: 0.75,
599            jitter_fraction: 0.0,
600            min_renew_secs: 100,
601            max_backoff_secs: 5,
602        };
603        mgr.register_with_created_at(1, 1000, 1100, policy);
604
605        // First tick at deadline fails
606        let s = mgr.tick_with(1075, |_, _| Err(()));
607        assert_eq!(s.failed, 1);
608        assert_eq!(s.renewed, 0);
609
610        // Immediately retrying should be in backoff (next_retry_at = 1075 + 1 = 1076)
611        let s = mgr.tick_with(1075, |_, _| Ok(100));
612        assert_eq!(s.skipped, 1);
613        assert_eq!(s.renewed, 0);
614
615        // After backoff delay (t=1076), should retry
616        let s = mgr.tick_with(1076, |_, _| Ok(100));
617        assert_eq!(s.renewed, 1);
618    }
619
620    #[test]
621    fn multiple_leases() {
622        let mut mgr = RenewalManager::new();
623        let policy = RenewalPolicy {
624            renew_at_fraction: 0.50,
625            jitter_fraction: 0.0,
626            min_renew_secs: 100,
627            max_backoff_secs: 5,
628        };
629
630        mgr.register_with_created_at(1, 1000, 1100, policy);
631        mgr.register_with_created_at(2, 1000, 1200, policy);
632        mgr.register_with_created_at(3, 1000, 1300, policy);
633
634        // At t=1050: lease 1 (50% at TTL 100) is due, lease 2 (25% at TTL 200) not due
635        let s = mgr.tick(1050);
636        assert_eq!(s.renewed, 1);
637        assert_eq!(s.skipped, 2);
638
639        // At t=1100: lease 2 is now due (50% of 200), lease 3 not yet
640        // Lease 1 was renewed, so it has new expiry
641        let s = mgr.tick(1100);
642        assert!(s.renewed >= 1);
643    }
644
645    #[test]
646    fn unregister_removes_lease() {
647        let mut mgr = RenewalManager::new();
648        let policy = RenewalPolicy::default();
649        mgr.register(1, 1100, policy);
650        mgr.register(2, 1200, policy);
651        assert_eq!(mgr.len(), 2);
652
653        mgr.unregister(1);
654        assert_eq!(mgr.len(), 1);
655        assert!(mgr.expires_at(1).is_none());
656        assert!(mgr.expires_at(2).is_some());
657    }
658
659    #[test]
660    fn is_near_expiry_predicate() {
661        let mut mgr = RenewalManager::new();
662        let policy = RenewalPolicy {
663            renew_at_fraction: 0.75,
664            jitter_fraction: 0.0,
665            min_renew_secs: 100,
666            max_backoff_secs: 5,
667        };
668        mgr.register_with_created_at(1, 1000, 1100, policy);
669
670        // At t=1050: 50% remaining, not near expiry
671        assert!(!mgr.is_near_expiry(1, 1050));
672
673        // At t=1091: 9% remaining, near expiry (< 10%)
674        assert!(mgr.is_near_expiry(1, 1091));
675
676        // At t=1100: expired, near expiry
677        assert!(mgr.is_near_expiry(1, 1100));
678    }
679
680    #[test]
681    fn near_expiry_in_summary() {
682        let mut mgr = RenewalManager::new();
683        let policy = RenewalPolicy {
684            renew_at_fraction: 0.75,
685            jitter_fraction: 0.0,
686            min_renew_secs: 100,
687            max_backoff_secs: 5,
688        };
689        mgr.register_with_created_at(1, 1000, 1100, policy);
690
691        let s = mgr.tick(1091);
692        assert!(s.near_expiry.contains(&1));
693    }
694
695    #[test]
696    fn re_register_replaces_entry() {
697        let mut mgr = RenewalManager::new();
698        let policy = RenewalPolicy::default();
699        mgr.register(1, 1100, policy);
700        assert_eq!(mgr.expires_at(1), Some(1100));
701
702        mgr.register(1, 2200, policy);
703        assert_eq!(mgr.len(), 1);
704        assert_eq!(mgr.expires_at(1), Some(2200));
705    }
706
707    #[test]
708    fn unknown_lease_is_not_near_expiry() {
709        let mgr = RenewalManager::new();
710        assert!(!mgr.is_near_expiry(999, 5000));
711    }
712
713    #[test]
714    fn renewal_extends_expiry() {
715        let mut mgr = RenewalManager::new();
716        let policy = RenewalPolicy {
717            renew_at_fraction: 0.75,
718            jitter_fraction: 0.0,
719            min_renew_secs: 100,
720            max_backoff_secs: 5,
721        };
722        mgr.register_with_created_at(1, 1000, 1100, policy);
723
724        let old_exp = mgr.expires_at(1).unwrap();
725        assert_eq!(old_exp, 1100);
726
727        // Renew at t=1080 with tick (auto-extends by min_renew_secs=100)
728        mgr.tick(1080);
729        let new_exp = mgr.expires_at(1).unwrap();
730        // new_expires_at = now + duration = 1080 + 100 = 1180
731        assert_eq!(new_exp, 1180);
732    }
733
734    #[test]
735    fn on_revoked_callback_fires() {
736        use alloc::sync::Arc;
737        use core::sync::atomic::{AtomicBool, Ordering};
738
739        let mut mgr = RenewalManager::new();
740        let handle = mgr.register_handle(42, 1100, RenewalPolicy::default());
741        let fired = Arc::new(AtomicBool::new(false));
742        let fired_clone = Arc::clone(&fired);
743        mgr.on_revoked(move |_lease_id, _reason| {
744            fired_clone.store(true, Ordering::SeqCst);
745        });
746
747        mgr.notify_revoked(42, 1);
748        assert!(fired.load(Ordering::SeqCst));
749        assert_eq!(handle.revoke_state(), RevokeState::RevokeWarning);
750    }
751
752    #[test]
753    fn multiple_revocation_callbacks_fire() {
754        use alloc::sync::Arc;
755        use core::sync::atomic::{AtomicBool, Ordering};
756
757        let mut mgr = RenewalManager::new();
758
759        let fired_a = Arc::new(AtomicBool::new(false));
760        let fired_b = Arc::new(AtomicBool::new(false));
761
762        let a = Arc::clone(&fired_a);
763        mgr.on_revoked(move |_, _| {
764            a.store(true, Ordering::SeqCst);
765        });
766
767        let b = Arc::clone(&fired_b);
768        mgr.on_revoked(move |_, _| {
769            b.store(true, Ordering::SeqCst);
770        });
771
772        mgr.notify_revoked(99, 2);
773        assert!(fired_a.load(Ordering::SeqCst));
774        assert!(fired_b.load(Ordering::SeqCst));
775    }
776
777    #[test]
778    fn notify_revoked_with_no_callbacks_is_noop() {
779        let mgr = RenewalManager::new();
780        // Should not panic
781        mgr.notify_revoked(7, 0);
782    }
783
784    #[test]
785    fn revocation_callback_receives_correct_args() {
786        use alloc::sync::Arc;
787        use core::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
788
789        let mut mgr = RenewalManager::new();
790        let count = Arc::new(AtomicUsize::new(0));
791        let first = Arc::new(AtomicU64::new(0));
792        let second = Arc::new(AtomicU64::new(0));
793        let count_clone = Arc::clone(&count);
794        let first_clone = Arc::clone(&first);
795        let second_clone = Arc::clone(&second);
796        mgr.on_revoked(move |lease_id, reason| {
797            let packed = ((lease_id as u64) << 8) | u64::from(reason);
798            match count_clone.fetch_add(1, Ordering::SeqCst) {
799                0 => first_clone.store(packed, Ordering::SeqCst),
800                1 => second_clone.store(packed, Ordering::SeqCst),
801                _ => {}
802            }
803        });
804
805        mgr.notify_revoked(42, 1);
806        mgr.notify_revoked(1000, 255);
807
808        assert_eq!(count.load(Ordering::SeqCst), 2);
809        assert_eq!(first.load(Ordering::SeqCst), (42 << 8) | 1);
810        assert_eq!(second.load(Ordering::SeqCst), (1000 << 8) | 255);
811    }
812
813    #[test]
814    fn backoff_escalation_on_repeated_failure() {
815        let mut mgr = RenewalManager::new();
816        let policy = RenewalPolicy {
817            renew_at_fraction: 0.75,
818            jitter_fraction: 0.0,
819            min_renew_secs: 100,
820            max_backoff_secs: 5,
821        };
822        mgr.register_with_created_at(1, 1000, 1100, policy);
823
824        // Fail at t=1075 => backoff 1s, next_retry_at=1076
825        let s = mgr.tick_with(1075, |_, _| Err(()));
826        assert_eq!(s.failed, 1);
827
828        // Fail at t=1076 => backoff 2s, next_retry_at=1078
829        let s = mgr.tick_with(1076, |_, _| Err(()));
830        assert_eq!(s.failed, 1);
831
832        // Skip at t=1077 (in backoff)
833        let s = mgr.tick_with(1077, |_, _| Ok(100));
834        assert_eq!(s.skipped, 1);
835
836        // Succeed at t=1078 (past backoff)
837        let s = mgr.tick_with(1078, |_, _| Ok(100));
838        assert_eq!(s.renewed, 1);
839    }
840}
grafos_leasekit/manager.rs

grafos_leasekit/
manager.rs
