openlatch-client 0.6.23

use std::collections::HashMap;
use std::time::{Duration, Instant};

const BACKOFF_BASE: [u64; 5] = [1, 4, 16, 64, 256];
const CIRCUIT_OPEN_THRESHOLD: u32 = 5;
const CIRCUIT_OPEN_WINDOW: Duration = Duration::from_secs(600);
const CIRCUIT_AUTO_CLOSE: Duration = Duration::from_secs(3600);
const HEALTHY_RESET: Duration = Duration::from_secs(600);
const JITTER_PERCENT: f64 = 0.20;

#[derive(Debug, Clone, PartialEq)]
pub enum HealState {
    Healthy,
    Drifted,
    Healing,
    Backoff { until: Instant },
    CircuitOpen { opened_at: Instant },
}

#[derive(Debug)]
pub struct EntryHealTracker {
    pub state: HealState,
    pub attempt: u32,
    pub last_healthy: Option<Instant>,
    recent_heals: Vec<Instant>,
}

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

impl EntryHealTracker {
    pub fn new() -> Self {
        Self {
            state: HealState::Healthy,
            attempt: 0,
            last_healthy: Some(Instant::now()),
            recent_heals: Vec::new(),
        }
    }

    pub fn mark_healthy(&mut self) {
        let now = Instant::now();
        if let Some(last) = self.last_healthy {
            if now.duration_since(last) >= HEALTHY_RESET {
                self.attempt = 0;
                self.recent_heals.clear();
            }
        }
        self.state = HealState::Healthy;
        self.last_healthy = Some(now);
    }

    pub fn mark_drifted(&mut self) {
        self.mark_drifted_at(Instant::now());
    }

    /// An open circuit holds until [`CIRCUIT_AUTO_CLOSE`] has passed, and a
    /// backoff until its `until`: drift seen inside either is not a new reason
    /// to heal. An expired circuit re-arms with a clean window, so healing
    /// resumes instead of reopening on the attempts that opened it.
    fn mark_drifted_at(&mut self, now: Instant) {
        match self.state {
            HealState::CircuitOpen { opened_at }
                if now.duration_since(opened_at) < CIRCUIT_AUTO_CLOSE =>
            {
                return;
            }
            HealState::CircuitOpen { .. } => self.recent_heals.clear(),
            HealState::Backoff { until } if now < until => return,
            _ => {}
        }
        self.state = HealState::Drifted;
        self.last_healthy = None;
    }

    pub fn should_heal(&self) -> bool {
        matches!(self.state, HealState::Drifted)
    }

    pub fn is_circuit_open(&self) -> bool {
        self.is_circuit_open_at(Instant::now())
    }

    fn is_circuit_open_at(&self, now: Instant) -> bool {
        matches!(self.state, HealState::CircuitOpen { opened_at }
            if now.duration_since(opened_at) < CIRCUIT_AUTO_CLOSE)
    }

    pub fn record_heal_attempt(&mut self) {
        self.record_heal_attempt_at(Instant::now());
    }

    fn record_heal_attempt_at(&mut self, now: Instant) {
        self.attempt += 1;
        self.recent_heals
            .retain(|t| now.duration_since(*t) < CIRCUIT_OPEN_WINDOW);
        self.recent_heals.push(now);

        if self.recent_heals.len() as u32 >= CIRCUIT_OPEN_THRESHOLD {
            self.state = HealState::CircuitOpen { opened_at: now };
            tracing::warn!(
                attempts = self.recent_heals.len(),
                "circuit breaker opened — auto-heal paused for this entry"
            );
            return;
        }

        self.state = HealState::Healing;
    }

    pub fn record_heal_success(&mut self) {
        if matches!(self.state, HealState::CircuitOpen { .. }) {
            return;
        }
        self.state = HealState::Healthy;
        self.last_healthy = Some(Instant::now());
    }

    pub fn record_heal_failure(&mut self) {
        self.record_heal_failure_at(Instant::now());
    }

    /// A no-op while the circuit is open, like [`Self::record_heal_success`]:
    /// the attempt that opened it must not trade the circuit for a backoff of
    /// a few seconds.
    fn record_heal_failure_at(&mut self, now: Instant) {
        if matches!(self.state, HealState::CircuitOpen { .. }) {
            return;
        }
        let idx = (self.attempt as usize).min(BACKOFF_BASE.len()) - 1;
        let base_secs = BACKOFF_BASE[idx];
        let jitter = (base_secs as f64 * JITTER_PERCENT * (rand_fraction() * 2.0 - 1.0)) as i64;
        let delay_secs = (base_secs as i64 + jitter).max(1) as u64;

        self.state = HealState::Backoff {
            until: now + Duration::from_secs(delay_secs),
        };
    }

    pub fn force_close_circuit(&mut self) {
        self.state = HealState::Healthy;
        self.attempt = 0;
        self.recent_heals.clear();
        self.last_healthy = Some(Instant::now());
    }

    pub fn backoff_remaining_ms(&self) -> Option<u64> {
        if let HealState::Backoff { until } = self.state {
            let now = Instant::now();
            if until > now {
                return Some(until.duration_since(now).as_millis() as u64);
            }
        }
        None
    }
}

fn rand_fraction() -> f64 {
    let bytes = uuid::Uuid::new_v4().as_bytes()[0];
    bytes as f64 / 255.0
}

#[derive(Debug, Default)]
pub struct HealManager {
    trackers: HashMap<String, EntryHealTracker>,
}

impl HealManager {
    pub fn new() -> Self {
        Self {
            trackers: HashMap::new(),
        }
    }

    pub fn tracker(&mut self, entry_id: &str) -> &mut EntryHealTracker {
        self.trackers.entry(entry_id.to_string()).or_default()
    }

    pub fn open_circuits(&self) -> Vec<&str> {
        self.trackers
            .iter()
            .filter(|(_, t)| t.is_circuit_open())
            .map(|(id, _)| id.as_str())
            .collect()
    }

    /// Drop what is known about `entry_id`, as though it had never been seen:
    /// for a key found not to be ours to heal after an attempt was recorded.
    pub fn forget(&mut self, entry_id: &str) {
        self.trackers.remove(entry_id);
    }

    pub fn force_close(&mut self, entry_id: &str) -> bool {
        if let Some(tracker) = self.trackers.get_mut(entry_id) {
            tracker.force_close_circuit();
            true
        } else {
            false
        }
    }
}

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

    #[test]
    fn new_tracker_is_healthy() {
        let t = EntryHealTracker::new();
        assert_eq!(t.state, HealState::Healthy);
        assert_eq!(t.attempt, 0);
    }

    #[test]
    fn drifted_then_should_heal() {
        let mut t = EntryHealTracker::new();
        t.mark_drifted();
        assert!(t.should_heal());
    }

    #[test]
    fn circuit_opens_after_rapid_heals() {
        let mut t = EntryHealTracker::new();
        for _ in 0..5 {
            t.mark_drifted();
            t.record_heal_attempt();
            t.record_heal_success();
            t.mark_drifted();
        }
        assert!(t.is_circuit_open());
    }

    #[test]
    fn circuit_blocks_drift() {
        let mut t = EntryHealTracker::new();
        for _ in 0..5 {
            t.mark_drifted();
            t.record_heal_attempt();
            t.record_heal_success();
            t.mark_drifted();
        }
        assert!(t.is_circuit_open());
        t.mark_drifted();
        assert!(!t.should_heal());
    }

    #[test]
    fn force_close_resets_circuit() {
        let mut t = EntryHealTracker::new();
        for _ in 0..5 {
            t.mark_drifted();
            t.record_heal_attempt();
            t.record_heal_success();
            t.mark_drifted();
        }
        assert!(t.is_circuit_open());
        t.force_close_circuit();
        assert!(!t.is_circuit_open());
        assert_eq!(t.attempt, 0);
    }

    /// One drift pass whose heal fails, at `now`; whether it attempted.
    fn failed_pass(t: &mut EntryHealTracker, now: Instant) -> bool {
        t.mark_drifted_at(now);
        if !t.should_heal() {
            return false;
        }
        t.record_heal_attempt_at(now);
        t.record_heal_failure_at(now);
        true
    }

    /// The heal that keeps failing: five attempts open the circuit, and it
    /// stays open — every later pass attempts nothing, and the backoff the
    /// failure path used to write cannot take its place.
    #[test]
    fn failed_heals_open_the_circuit_and_it_stays_open() {
        let start = Instant::now();
        let mut t = EntryHealTracker::new();
        let mut now = start;
        let mut attempts = 0;
        // 30 s passes for five minutes: each backoff (≤ 77 s) is over by the
        // time the pass after it is due, or the one after that.
        for _ in 0..10 {
            attempts += u32::from(failed_pass(&mut t, now));
            now += Duration::from_secs(30);
        }
        assert_eq!(attempts, CIRCUIT_OPEN_THRESHOLD);
        assert!(t.is_circuit_open_at(now));

        for _ in 0..100 {
            assert!(
                !failed_pass(&mut t, now),
                "an open circuit attempts nothing"
            );
            now += Duration::from_secs(30);
        }
        assert!(matches!(t.state, HealState::CircuitOpen { .. }));
        assert_eq!(t.attempt, CIRCUIT_OPEN_THRESHOLD);
    }

    /// A failure waits out its backoff, then re-arms.
    #[test]
    fn backoff_delays_the_next_attempt_until_it_ends() {
        let now = Instant::now();
        let mut t = EntryHealTracker::new();
        assert!(failed_pass(&mut t, now));
        let HealState::Backoff { until } = t.state else {
            panic!("a failed heal backs off: {:?}", t.state);
        };
        assert!(until > now);

        t.mark_drifted_at(until - Duration::from_millis(1));
        assert!(!t.should_heal(), "still inside the backoff");
        assert!(matches!(t.state, HealState::Backoff { .. }));

        t.mark_drifted_at(until);
        assert!(t.should_heal(), "the backoff is over");
    }

    /// An open circuit closes by itself after `CIRCUIT_AUTO_CLOSE`: one new
    /// attempt, from a clean window, and not a second circuit-opening run.
    #[test]
    fn an_open_circuit_re_arms_once_after_auto_close() {
        let start = Instant::now();
        let mut t = EntryHealTracker::new();
        for i in 0..CIRCUIT_OPEN_THRESHOLD {
            t.mark_drifted_at(start + Duration::from_secs(u64::from(i)));
            t.record_heal_attempt_at(start + Duration::from_secs(u64::from(i)));
        }
        let HealState::CircuitOpen { opened_at } = t.state else {
            panic!("the circuit opens: {:?}", t.state);
        };

        let closes = opened_at + CIRCUIT_AUTO_CLOSE;
        assert!(!failed_pass(&mut t, closes - Duration::from_secs(1)));
        assert!(failed_pass(&mut t, closes), "re-armed after auto-close");
        assert!(
            matches!(t.state, HealState::Backoff { .. }),
            "one failure backs off; it does not reopen on the old window: {:?}",
            t.state
        );
        assert!(!failed_pass(&mut t, closes), "and the backoff holds");
    }

    /// A heal that works inside the window is still a plain success.
    #[test]
    fn a_success_inside_the_window_is_unchanged() {
        let now = Instant::now();
        let mut t = EntryHealTracker::new();
        assert!(failed_pass(&mut t, now));
        let HealState::Backoff { until } = t.state else {
            panic!("backed off: {:?}", t.state);
        };
        t.mark_drifted_at(until);
        t.record_heal_attempt_at(until);
        t.record_heal_success();
        assert_eq!(t.state, HealState::Healthy);
        assert!(!t.is_circuit_open_at(until));
    }

    #[test]
    fn heal_manager_tracks_multiple_entries() {
        let mut m = HealManager::new();
        m.tracker("entry-1").mark_drifted();
        m.tracker("entry-2").mark_healthy();
        assert!(m.tracker("entry-1").should_heal());
        assert!(!m.tracker("entry-2").should_heal());
    }

    #[test]
    fn heal_manager_open_circuits() {
        let mut m = HealManager::new();
        for _ in 0..5 {
            m.tracker("entry-1").mark_drifted();
            m.tracker("entry-1").record_heal_attempt();
            m.tracker("entry-1").record_heal_success();
            m.tracker("entry-1").mark_drifted();
        }
        let open = m.open_circuits();
        assert!(open.contains(&"entry-1"));
    }

    #[test]
    fn force_close_via_manager() {
        let mut m = HealManager::new();
        for _ in 0..5 {
            m.tracker("entry-1").mark_drifted();
            m.tracker("entry-1").record_heal_attempt();
            m.tracker("entry-1").record_heal_success();
            m.tracker("entry-1").mark_drifted();
        }
        assert!(m.force_close("entry-1"));
        assert!(m.open_circuits().is_empty());
    }
}