grafos_leasekit/

manager.rs

//! The [`RenewalManager`] — poll-driven lease renewal engine.

extern crate alloc;
use alloc::boxed::Box;
use alloc::vec::Vec;

use crate::backoff::Backoff;
use crate::policy::RenewalPolicy;

/// Tracks the renewal state for a single registered lease.
struct Entry {
    lease_id: u128,
    created_at: u64,
    expires_at: u64,
    policy: RenewalPolicy,
    backoff: Backoff,
    /// Absolute time before which we should not retry after a failure.
    next_retry_at: u64,
    /// Whether the lease has been successfully renewed at least once by
    /// this manager (used to compute renewal deadline from latest expiry).
    last_renewed_at: Option<u64>,
}

impl Entry {
    fn renewal_deadline(&self) -> u64 {
        let created = self.last_renewed_at.unwrap_or(self.created_at);
        let jitter_seed = self.lease_id as u64;
        self.policy
            .renewal_deadline(created, self.expires_at, jitter_seed)
    }

    fn is_near_expiry(&self, now: u64) -> bool {
        if now >= self.expires_at {
            return true;
        }
        let remaining = self.expires_at - now;
        let total_ttl = self
            .expires_at
            .saturating_sub(self.last_renewed_at.unwrap_or(self.created_at));
        if total_ttl == 0 {
            return true;
        }
        // Near expiry when less than 10% TTL remains
        remaining < total_ttl / 10
    }
}

/// Summary returned by [`RenewalManager::tick`].
#[derive(Debug, Clone, Default)]
pub struct RenewalSummary {
    /// Number of leases successfully renewed this tick.
    pub renewed: u32,
    /// Number of leases checked but not yet due for renewal.
    pub skipped: u32,
    /// Number of leases where renewal was attempted but failed.
    pub failed: u32,
    /// Lease IDs that are near expiry (< 10% TTL remaining).
    pub near_expiry: Vec<u128>,
}

/// Callback invoked by the manager to actually perform a renewal.
///
/// The manager itself does not hold lease objects; callers provide a
/// `renew_fn` closure to `tick_with` that performs the renewal and
/// returns the new `expires_at` timestamp on success.
///
/// Alternatively, use [`tick`](RenewalManager::tick) for a simpler API
/// that auto-renews by extending `expires_at` by `policy.min_renew_secs`.
pub struct RenewalManager {
    entries: Vec<Entry>,
    revoke_callbacks: Vec<Box<dyn Fn(u128, u8)>>,
}

impl RenewalManager {
    /// Create an empty manager with no registered leases.
    pub fn new() -> Self {
        RenewalManager {
            entries: Vec::new(),
            revoke_callbacks: Vec::new(),
        }
    }

    /// Register a lease for managed renewal.
    ///
    /// `created_at` is inferred as `expires_at - policy.min_renew_secs`
    /// (clamped to 0). For precise control, use `register_with_created_at`.
    pub fn register(&mut self, lease_id: u128, expires_at: u64, policy: RenewalPolicy) {
        let created_at = expires_at.saturating_sub(policy.min_renew_secs);
        self.register_with_created_at(lease_id, created_at, expires_at, policy);
    }

    /// Register a lease with an explicit creation timestamp.
    pub fn register_with_created_at(
        &mut self,
        lease_id: u128,
        created_at: u64,
        expires_at: u64,
        policy: RenewalPolicy,
    ) {
        // Remove any existing entry with the same lease_id
        self.entries.retain(|e| e.lease_id != lease_id);
        let backoff = Backoff::new(1, policy.max_backoff_secs);
        self.entries.push(Entry {
            lease_id,
            created_at,
            expires_at,
            policy,
            backoff,
            next_retry_at: 0,
            last_renewed_at: None,
        });
    }

    /// Unregister a lease.
    pub fn unregister(&mut self, lease_id: u128) {
        self.entries.retain(|e| e.lease_id != lease_id);
    }

    /// Number of leases currently managed.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Returns `true` if no leases are managed.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Drive renewals using a simple model: if a lease is due, extend
    /// `expires_at` by `policy.min_renew_secs`. This is suitable when the
    /// manager owns the lease state (e.g. testing or self-contained use).
    pub fn tick(&mut self, now_unix_secs: u64) -> RenewalSummary {
        self.tick_with(now_unix_secs, |_lease_id, duration| {
            // Simple model: always succeed, return new expiry
            Ok(duration)
        })
    }

    /// Drive renewals with a caller-provided renewal function.
    ///
    /// `renew_fn(lease_id, duration_secs)` should attempt the actual
    /// renewal and return the new `expires_at` on success, or an error.
    pub fn tick_with<F>(&mut self, now_unix_secs: u64, mut renew_fn: F) -> RenewalSummary
    where
        F: FnMut(u128, u64) -> Result<u64, ()>,
    {
        let mut summary = RenewalSummary::default();

        for entry in self.entries.iter_mut() {
            // Check near-expiry
            if entry.is_near_expiry(now_unix_secs) {
                summary.near_expiry.push(entry.lease_id);
            }

            // Already expired — skip
            if now_unix_secs >= entry.expires_at {
                summary.skipped += 1;
                continue;
            }

            // Not yet at renewal deadline — skip
            let deadline = entry.renewal_deadline();
            if now_unix_secs < deadline {
                summary.skipped += 1;
                continue;
            }

            // In backoff — skip
            if now_unix_secs < entry.next_retry_at {
                summary.skipped += 1;
                continue;
            }

            // Attempt renewal
            let duration = entry.policy.min_renew_secs;
            match renew_fn(entry.lease_id, duration) {
                Ok(new_expires_at) => {
                    let new_exp = now_unix_secs.saturating_add(new_expires_at);
                    entry.last_renewed_at = Some(now_unix_secs);
                    entry.expires_at = new_exp;
                    entry.backoff.reset();
                    entry.next_retry_at = 0;
                    summary.renewed += 1;

                    #[cfg(feature = "observe")]
                    emit_renew_success(entry.lease_id);
                }
                Err(()) => {
                    let delay = entry.backoff.next_delay();
                    entry.next_retry_at = now_unix_secs.saturating_add(delay);
                    summary.failed += 1;

                    #[cfg(feature = "observe")]
                    emit_renew_failure(entry.lease_id);
                }
            }
        }

        #[cfg(feature = "observe")]
        emit_tick_metrics(&summary, self.entries.len());

        summary
    }

    /// Returns `true` if the given lease is near expiry (< 10% TTL remaining).
    pub fn is_near_expiry(&self, lease_id: u128, now: u64) -> bool {
        self.entries
            .iter()
            .find(|e| e.lease_id == lease_id)
            .map(|e| e.is_near_expiry(now))
            .unwrap_or(false)
    }

    /// Register a callback invoked when a lease is revoked (e.g. via
    /// WITHDRAW cleanup or REVOKE_BROADCAST).
    ///
    /// The callback receives the lease ID and a reason code (0 = unknown/timeout).
    pub fn on_revoked<F: Fn(u128, u8) + 'static>(&mut self, callback: F) {
        self.revoke_callbacks.push(Box::new(callback));
    }

    /// Notify all registered revocation callbacks for the given lease.
    ///
    /// `reason` is the WITHDRAW reason code (0 = timeout/unknown).
    pub fn notify_revoked(&self, lease_id: u128, reason: u8) {
        for cb in &self.revoke_callbacks {
            cb(lease_id, reason);
        }
    }

    /// Returns the current `expires_at` for a managed lease, if registered.
    pub fn expires_at(&self, lease_id: u128) -> Option<u64> {
        self.entries
            .iter()
            .find(|e| e.lease_id == lease_id)
            .map(|e| e.expires_at)
    }
}

impl Default for RenewalManager {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(feature = "observe")]
fn emit_renew_success(lease_id: u128) {
    let _ = lease_id;
    grafos_observe::FabricMetrics::global().ops_total.inc();
}

#[cfg(feature = "observe")]
fn emit_renew_failure(lease_id: u128) {
    let _ = lease_id;
    grafos_observe::FabricMetrics::global().ops_total.inc();
}

#[cfg(feature = "observe")]
fn emit_tick_metrics(summary: &RenewalSummary, managed_count: usize) {
    let _ = (summary, managed_count);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn register_and_tick_before_deadline() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        // created_at=1000, expires_at=1100, TTL=100
        mgr.register_with_created_at(1, 1000, 1100, policy);

        // At t=1050, only 50% elapsed — should not renew
        let s = mgr.tick(1050);
        assert_eq!(s.renewed, 0);
        assert_eq!(s.skipped, 1);
        assert_eq!(s.failed, 0);
    }

    #[test]
    fn tick_at_deadline_triggers_renewal() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        mgr.register_with_created_at(1, 1000, 1100, policy);

        // At t=1075 (exactly 75% elapsed), should renew
        let s = mgr.tick(1075);
        assert_eq!(s.renewed, 1);
        assert_eq!(s.skipped, 0);
    }

    #[test]
    fn tick_after_expiry_skips() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        mgr.register_with_created_at(1, 1000, 1100, policy);

        let s = mgr.tick(1200);
        assert_eq!(s.renewed, 0);
        assert_eq!(s.skipped, 1);
    }

    #[test]
    fn tick_with_failure_triggers_backoff() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        mgr.register_with_created_at(1, 1000, 1100, policy);

        // First tick at deadline fails
        let s = mgr.tick_with(1075, |_, _| Err(()));
        assert_eq!(s.failed, 1);
        assert_eq!(s.renewed, 0);

        // Immediately retrying should be in backoff (next_retry_at = 1075 + 1 = 1076)
        let s = mgr.tick_with(1075, |_, _| Ok(100));
        assert_eq!(s.skipped, 1);
        assert_eq!(s.renewed, 0);

        // After backoff delay (t=1076), should retry
        let s = mgr.tick_with(1076, |_, _| Ok(100));
        assert_eq!(s.renewed, 1);
    }

    #[test]
    fn multiple_leases() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.50,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };

        mgr.register_with_created_at(1, 1000, 1100, policy);
        mgr.register_with_created_at(2, 1000, 1200, policy);
        mgr.register_with_created_at(3, 1000, 1300, policy);

        // At t=1050: lease 1 (50% at TTL 100) is due, lease 2 (25% at TTL 200) not due
        let s = mgr.tick(1050);
        assert_eq!(s.renewed, 1);
        assert_eq!(s.skipped, 2);

        // At t=1100: lease 2 is now due (50% of 200), lease 3 not yet
        // Lease 1 was renewed, so it has new expiry
        let s = mgr.tick(1100);
        assert!(s.renewed >= 1);
    }

    #[test]
    fn unregister_removes_lease() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy::default();
        mgr.register(1, 1100, policy);
        mgr.register(2, 1200, policy);
        assert_eq!(mgr.len(), 2);

        mgr.unregister(1);
        assert_eq!(mgr.len(), 1);
        assert!(mgr.expires_at(1).is_none());
        assert!(mgr.expires_at(2).is_some());
    }

    #[test]
    fn is_near_expiry_predicate() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        mgr.register_with_created_at(1, 1000, 1100, policy);

        // At t=1050: 50% remaining, not near expiry
        assert!(!mgr.is_near_expiry(1, 1050));

        // At t=1091: 9% remaining, near expiry (< 10%)
        assert!(mgr.is_near_expiry(1, 1091));

        // At t=1100: expired, near expiry
        assert!(mgr.is_near_expiry(1, 1100));
    }

    #[test]
    fn near_expiry_in_summary() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        mgr.register_with_created_at(1, 1000, 1100, policy);

        let s = mgr.tick(1091);
        assert!(s.near_expiry.contains(&1));
    }

    #[test]
    fn re_register_replaces_entry() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy::default();
        mgr.register(1, 1100, policy);
        assert_eq!(mgr.expires_at(1), Some(1100));

        mgr.register(1, 2200, policy);
        assert_eq!(mgr.len(), 1);
        assert_eq!(mgr.expires_at(1), Some(2200));
    }

    #[test]
    fn unknown_lease_is_not_near_expiry() {
        let mgr = RenewalManager::new();
        assert!(!mgr.is_near_expiry(999, 5000));
    }

    #[test]
    fn renewal_extends_expiry() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        mgr.register_with_created_at(1, 1000, 1100, policy);

        let old_exp = mgr.expires_at(1).unwrap();
        assert_eq!(old_exp, 1100);

        // Renew at t=1080 with tick (auto-extends by min_renew_secs=100)
        mgr.tick(1080);
        let new_exp = mgr.expires_at(1).unwrap();
        // new_expires_at = now + duration = 1080 + 100 = 1180
        assert_eq!(new_exp, 1180);
    }

    #[test]
    fn on_revoked_callback_fires() {
        use alloc::sync::Arc;
        use core::sync::atomic::{AtomicBool, Ordering};

        let mut mgr = RenewalManager::new();
        let fired = Arc::new(AtomicBool::new(false));
        let fired_clone = Arc::clone(&fired);
        mgr.on_revoked(move |_lease_id, _reason| {
            fired_clone.store(true, Ordering::SeqCst);
        });

        mgr.notify_revoked(42, 1);
        assert!(fired.load(Ordering::SeqCst));
    }

    #[test]
    fn multiple_revocation_callbacks_fire() {
        use alloc::sync::Arc;
        use core::sync::atomic::{AtomicBool, Ordering};

        let mut mgr = RenewalManager::new();

        let fired_a = Arc::new(AtomicBool::new(false));
        let fired_b = Arc::new(AtomicBool::new(false));

        let a = Arc::clone(&fired_a);
        mgr.on_revoked(move |_, _| {
            a.store(true, Ordering::SeqCst);
        });

        let b = Arc::clone(&fired_b);
        mgr.on_revoked(move |_, _| {
            b.store(true, Ordering::SeqCst);
        });

        mgr.notify_revoked(99, 2);
        assert!(fired_a.load(Ordering::SeqCst));
        assert!(fired_b.load(Ordering::SeqCst));
    }

    #[test]
    fn notify_revoked_with_no_callbacks_is_noop() {
        let mgr = RenewalManager::new();
        // Should not panic
        mgr.notify_revoked(7, 0);
    }

    #[test]
    fn revocation_callback_receives_correct_args() {
        use alloc::sync::Arc;
        use std::sync::Mutex;

        let mut mgr = RenewalManager::new();
        let log: Arc<Mutex<Vec<(u128, u8)>>> = Arc::new(Mutex::new(Vec::new()));
        let log_clone = Arc::clone(&log);
        mgr.on_revoked(move |lease_id, reason| {
            log_clone.lock().unwrap().push((lease_id, reason));
        });

        mgr.notify_revoked(42, 1);
        mgr.notify_revoked(1000, 255);

        let entries = log.lock().unwrap();
        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0], (42, 1));
        assert_eq!(entries[1], (1000, 255));
    }

    #[test]
    fn backoff_escalation_on_repeated_failure() {
        let mut mgr = RenewalManager::new();
        let policy = RenewalPolicy {
            renew_at_fraction: 0.75,
            jitter_fraction: 0.0,
            min_renew_secs: 100,
            max_backoff_secs: 5,
        };
        mgr.register_with_created_at(1, 1000, 1100, policy);

        // Fail at t=1075 => backoff 1s, next_retry_at=1076
        let s = mgr.tick_with(1075, |_, _| Err(()));
        assert_eq!(s.failed, 1);

        // Fail at t=1076 => backoff 2s, next_retry_at=1078
        let s = mgr.tick_with(1076, |_, _| Err(()));
        assert_eq!(s.failed, 1);

        // Skip at t=1077 (in backoff)
        let s = mgr.tick_with(1077, |_, _| Ok(100));
        assert_eq!(s.skipped, 1);

        // Succeed at t=1078 (past backoff)
        let s = mgr.tick_with(1078, |_, _| Ok(100));
        assert_eq!(s.renewed, 1);
    }
}