zakura-network 7.0.1

//! Branch-owned body-availability retry episodes.

use std::collections::{BTreeSet, HashMap};

use chrono::{DateTime, Duration, Utc};
use zakura_header_chain::{
    BodyUnavailableSummary, BranchId, Clock, Frontier, HeaderGeneration, SourceId,
};

const ALARM_ATTEMPTS: u32 = 10;
const ALARM_AFTER: Duration = Duration::minutes(10);
const ALARM_PROBE_INTERVAL: Duration = Duration::minutes(10);
const MAX_BACKOFF_SECONDS: i64 = 60;
const MAX_JITTER_PER_THOUSAND: i16 = 100;

/// Injected deterministic jitter, bounded by the scheduler to plus or minus ten percent.
pub trait RetryJitter {
    /// Return signed per-thousand jitter for one stable retry identity and attempt.
    fn offset_per_thousand(&self, branch: BranchId, header: Frontier, attempt: u32) -> i16;
}

/// Stable seeded jitter derived from the exact branch, header, and attempt identity.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct SeededRetryJitter {
    seed: [u8; 32],
}

impl SeededRetryJitter {
    /// Construct deterministic node-local jitter from an authenticated or random seed.
    pub const fn new(seed: [u8; 32]) -> Self {
        Self { seed }
    }
}

impl RetryJitter for SeededRetryJitter {
    fn offset_per_thousand(&self, branch: BranchId, header: Frontier, attempt: u32) -> i16 {
        let mut state = blake2b_simd::Params::new()
            .hash_length(2)
            .personal(b"ZkBodyRetryV1__")
            .to_state();
        state.update(&self.seed);
        state.update(&branch.anchor_hash.0);
        state.update(&branch.target_tip_hash.0);
        state.update(&header.height.0.to_le_bytes());
        state.update(&header.hash.0);
        state.update(&attempt.to_le_bytes());
        let digest = state.finalize();
        let sample = u16::from_le_bytes(
            digest.as_bytes()[..2]
                .try_into()
                .expect("the configured digest contains exactly two bytes"),
        );
        i16::try_from(sample % 201).expect("a value modulo 201 fits in i16") - 100
    }
}

/// Result of recording one supplier failure.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum RetryUpdate {
    /// A repeat failure arrived before that supplier was eligible to retry.
    TooEarly,
    /// Retry remains active at the returned time.
    RetryAt(DateTime<Utc>),
    /// The failed delivery made persistent unavailability visible.
    Alarmed {
        /// Earliest persistent-alarm probe time.
        probe_at: DateTime<Utc>,
    },
    /// An already-alarmed probe failed and the next bounded probe was scheduled.
    ProbeAt(DateTime<Utc>),
}

/// One selected-branch body-unavailability episode.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BodyRetryEpisode {
    /// Exact selected branch whose body is unavailable.
    pub branch: BranchId,
    /// Selected-header generation that owns this episode.
    pub generation: HeaderGeneration,
    /// Exact header whose body is unavailable.
    pub header: Frontier,
    /// Authoritative local start time.
    pub started_at: DateTime<Utc>,
    /// Failed delivery attempts in this episode.
    pub attempts: u32,
    /// Suppliers already tried in this episode.
    pub tried_suppliers: BTreeSet<SourceId>,
    /// Earliest time a repeated supplier or alarm probe is eligible.
    pub next_probe_at: DateTime<Utc>,
    /// Whether persistent unavailability is visible.
    pub alarmed: bool,
    eligible_suppliers: BTreeSet<SourceId>,
}

impl BodyRetryEpisode {
    /// Start a fresh episode using an injected authoritative clock.
    pub fn new<C: Clock>(
        branch: BranchId,
        generation: HeaderGeneration,
        header: Frontier,
        eligible_suppliers: BTreeSet<SourceId>,
        clock: &C,
    ) -> Self {
        let now = clock.now();
        Self {
            branch,
            generation,
            header,
            started_at: now,
            attempts: 0,
            tried_suppliers: BTreeSet::new(),
            next_probe_at: now,
            alarmed: false,
            eligible_suppliers,
        }
    }

    /// Restore a durable episode without weakening its persisted alarm cadence.
    pub fn restore(
        branch: BranchId,
        generation: HeaderGeneration,
        header: Frontier,
        eligible_suppliers: BTreeSet<SourceId>,
        summary: BodyUnavailableSummary,
    ) -> Self {
        let tried_suppliers = if summary.alarmed {
            eligible_suppliers.clone()
        } else {
            BTreeSet::new()
        };
        Self {
            branch,
            generation,
            header,
            started_at: summary.started_at,
            attempts: summary.attempts,
            tried_suppliers,
            next_probe_at: summary.next_probe_at,
            alarmed: summary.alarmed,
            eligible_suppliers,
        }
    }

    /// Merge the current eligible suppliers without resetting this episode's history.
    ///
    /// Supplier membership is ephemeral peer information, not evidence that an
    /// unavailable body became available. Only [`Self::restart`] may reset the
    /// authoritative age, attempts, alarm state, or probe cadence.
    ///
    /// Returns whether the eligible membership changed.
    pub fn refresh_suppliers(&mut self, eligible_suppliers: BTreeSet<SourceId>) -> bool {
        let changed = self.eligible_suppliers != eligible_suppliers;
        self.eligible_suppliers = eligible_suppliers;
        self.tried_suppliers
            .retain(|supplier| self.eligible_suppliers.contains(supplier));
        changed
    }

    /// Start a new episode after an explicit operator retry command.
    pub fn restart<C: Clock>(&mut self, clock: &C) {
        *self = Self::new(
            self.branch,
            self.generation,
            self.header,
            self.eligible_suppliers.clone(),
            clock,
        );
    }

    /// Whether a repeated supplier attempt or persistent-alarm probe is due.
    pub fn is_due<C: Clock>(&self, clock: &C) -> bool {
        clock.now() >= self.next_probe_at
    }

    /// Record one failed supplier attempt and return its bounded scheduling consequence.
    pub fn record_failure<C: Clock, J: RetryJitter>(
        &mut self,
        supplier: SourceId,
        clock: &C,
        jitter: &J,
    ) -> RetryUpdate {
        let now = clock.now();
        if (self.alarmed || self.tried_suppliers.contains(&supplier)) && now < self.next_probe_at {
            return RetryUpdate::TooEarly;
        }

        self.attempts = self.attempts.saturating_add(1);
        self.tried_suppliers.insert(supplier);
        let all_suppliers_tried = !self.eligible_suppliers.is_empty()
            && self
                .eligible_suppliers
                .iter()
                .all(|supplier| self.tried_suppliers.contains(supplier));
        let alarm_due = self.attempts >= ALARM_ATTEMPTS
            || now.signed_duration_since(self.started_at) >= ALARM_AFTER;
        if self.alarmed {
            self.next_probe_at = retry_deadline(now, ALARM_PROBE_INTERVAL);
            return RetryUpdate::ProbeAt(self.next_probe_at);
        }
        if all_suppliers_tried && alarm_due {
            self.alarmed = true;
            self.next_probe_at = retry_deadline(now, ALARM_PROBE_INTERVAL);
            return RetryUpdate::Alarmed {
                probe_at: self.next_probe_at,
            };
        }

        self.next_probe_at = retry_deadline(
            now,
            retry_delay(self.branch, self.header, self.attempts, jitter),
        );
        RetryUpdate::RetryAt(self.next_probe_at)
    }

    /// Return the bounded durable alarm summary for state admission.
    pub fn summary(&self) -> BodyUnavailableSummary {
        BodyUnavailableSummary {
            started_at: self.started_at,
            attempts: self.attempts,
            suppliers: u32::try_from(self.eligible_suppliers.len()).unwrap_or(u32::MAX),
            supplier_set_digest: BodyUnavailableSummary::supplier_set_digest(
                &self.eligible_suppliers,
            ),
            alarmed: self.alarmed,
            next_probe_at: self.next_probe_at,
        }
    }
}

fn retry_deadline(now: DateTime<Utc>, delay: Duration) -> DateTime<Utc> {
    now.checked_add_signed(delay)
        .unwrap_or(DateTime::<Utc>::MAX_UTC)
}

fn retry_delay<J: RetryJitter>(
    branch: BranchId,
    header: Frontier,
    attempt: u32,
    jitter: &J,
) -> Duration {
    let base_seconds = if attempt <= 6 {
        1_i64 << attempt.saturating_sub(1)
    } else {
        MAX_BACKOFF_SECONDS
    };
    let offset = jitter
        .offset_per_thousand(branch, header, attempt)
        .clamp(-MAX_JITTER_PER_THOUSAND, MAX_JITTER_PER_THOUSAND);
    let milliseconds = base_seconds
        .saturating_mul(1_000)
        .saturating_mul(i64::from(1_000_i16.saturating_add(offset)))
        / 1_000;
    let milliseconds = milliseconds.min(MAX_BACKOFF_SECONDS.saturating_mul(1_000));
    Duration::milliseconds(milliseconds)
}

#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
struct BodyRetryKey {
    generation: HeaderGeneration,
    branch: BranchId,
    header_hash: zakura_chain::block::Hash,
}

/// Generation- and branch-owned body retry work.
#[derive(Clone, Debug, Default)]
pub struct BodyRetryQueue(HashMap<BodyRetryKey, BodyRetryEpisode>);

impl BodyRetryQueue {
    /// Insert or replace one exact generation/branch/header episode.
    pub fn insert(&mut self, episode: BodyRetryEpisode) -> Option<BodyRetryEpisode> {
        self.0.insert(
            BodyRetryKey {
                generation: episode.generation,
                branch: episode.branch,
                header_hash: episode.header.hash,
            },
            episode,
        )
    }

    /// Return one exact current episode for scheduling or completion handling.
    pub fn get_mut(
        &mut self,
        generation: HeaderGeneration,
        branch: BranchId,
        header_hash: zakura_chain::block::Hash,
    ) -> Option<&mut BodyRetryEpisode> {
        self.0.get_mut(&BodyRetryKey {
            generation,
            branch,
            header_hash,
        })
    }

    /// Remove one exact completed or canceled episode.
    pub fn remove(
        &mut self,
        generation: HeaderGeneration,
        branch: BranchId,
        header_hash: zakura_chain::block::Hash,
    ) -> Option<BodyRetryEpisode> {
        self.0.remove(&BodyRetryKey {
            generation,
            branch,
            header_hash,
        })
    }

    /// Number of exact retry episodes.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Whether no body retry episode remains scheduled.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Retire episodes outside the exact current generation and finalized anchor.
    pub fn retain_current(&mut self, generation: HeaderGeneration, finalized: Frontier) {
        self.retain_scope(generation, finalized.hash);
    }

    /// Retire episodes outside the exact current generation and finalized anchor hash.
    pub fn retain_scope(
        &mut self,
        generation: HeaderGeneration,
        finalized_hash: zakura_chain::block::Hash,
    ) {
        self.0.retain(|key, _| {
            key.generation == generation && key.branch.anchor_hash == finalized_hash
        });
    }

    /// Reauthorize every retained hash-specific episode under compatible authority.
    pub fn refresh_scope(&mut self, generation: HeaderGeneration, branch: BranchId) {
        let episodes: Vec<_> = self
            .0
            .drain()
            .map(|(_, mut episode)| {
                episode.generation = generation;
                episode.branch = branch;
                episode
            })
            .collect();
        for episode in episodes {
            self.insert(episode);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use zakura_chain::block;

    use super::*;

    struct ManualClock(Mutex<DateTime<Utc>>);

    impl ManualClock {
        fn new(now: DateTime<Utc>) -> Self {
            Self(Mutex::new(now))
        }

        fn advance(&self, duration: Duration) {
            let mut now = self.0.lock().expect("the test clock mutex is not poisoned");
            *now += duration;
        }
    }

    impl Clock for ManualClock {
        fn now(&self) -> DateTime<Utc> {
            *self.0.lock().expect("the test clock mutex is not poisoned")
        }
    }

    struct FixedJitter(i16);

    impl RetryJitter for FixedJitter {
        fn offset_per_thousand(&self, _branch: BranchId, _header: Frontier, _attempt: u32) -> i16 {
            self.0
        }
    }

    fn hash(byte: u8) -> block::Hash {
        block::Hash([byte; 32])
    }

    fn source(byte: u8) -> SourceId {
        SourceId::from_digest([byte; 32])
    }

    fn clock() -> ManualClock {
        ManualClock::new(
            DateTime::from_timestamp(1_700_000_000, 0).expect("the timestamp is valid"),
        )
    }

    fn episode(clock: &ManualClock, suppliers: &[SourceId]) -> BodyRetryEpisode {
        BodyRetryEpisode::new(
            BranchId::new(hash(1), hash(2)),
            HeaderGeneration::new(3),
            Frontier::new(block::Height(20), hash(2)),
            suppliers.iter().copied().collect(),
            clock,
        )
    }

    #[test]
    fn exact_backoff_caps_then_alarms_and_probes_every_ten_minutes() {
        let clock = clock();
        let supplier = source(3);
        let mut episode = episode(&clock, &[supplier]);
        let jitter = FixedJitter(0);

        for expected_seconds in [1, 2, 4, 8, 16, 32, 60, 60, 60] {
            let now = clock.now();
            assert_eq!(
                episode.record_failure(supplier, &clock, &jitter),
                RetryUpdate::RetryAt(now + Duration::seconds(expected_seconds))
            );
            clock.advance(Duration::seconds(expected_seconds));
        }
        let now = clock.now();
        assert_eq!(
            episode.record_failure(supplier, &clock, &jitter),
            RetryUpdate::Alarmed {
                probe_at: now + ALARM_PROBE_INTERVAL
            }
        );
        assert_eq!(
            episode.record_failure(supplier, &clock, &jitter),
            RetryUpdate::TooEarly
        );
        clock.advance(ALARM_PROBE_INTERVAL);
        let now = clock.now();
        assert_eq!(
            episode.record_failure(supplier, &clock, &jitter),
            RetryUpdate::ProbeAt(now + ALARM_PROBE_INTERVAL)
        );
        assert_eq!(
            episode.summary(),
            BodyUnavailableSummary {
                started_at: episode.started_at,
                attempts: 11,
                suppliers: 1,
                supplier_set_digest: BodyUnavailableSummary::supplier_set_digest(
                    &episode.eligible_suppliers,
                ),
                alarmed: true,
                next_probe_at: episode.next_probe_at,
            }
        );
    }

    #[test]
    fn jitter_and_final_delay_are_clamped_and_only_restart_resets_an_episode() {
        let clock = clock();
        let first = source(3);
        let second = source(4);
        let mut episode = episode(&clock, &[first]);
        let now = clock.now();
        assert_eq!(
            episode.record_failure(first, &clock, &FixedJitter(-500)),
            RetryUpdate::RetryAt(now + Duration::milliseconds(900))
        );
        clock.advance(Duration::milliseconds(900));
        let now = clock.now();
        assert_eq!(
            episode.record_failure(first, &clock, &FixedJitter(500)),
            RetryUpdate::RetryAt(now + Duration::milliseconds(2_200))
        );
        assert_eq!(
            retry_delay(episode.branch, episode.header, 7, &FixedJitter(-500)),
            Duration::seconds(54)
        );
        assert_eq!(
            retry_delay(episode.branch, episode.header, 7, &FixedJitter(500)),
            Duration::seconds(60),
            "the jittered delay itself must not exceed the normative cap"
        );
        let started_at = episode.started_at;
        let attempts = episode.attempts;
        let next_probe_at = episode.next_probe_at;
        assert!(episode.refresh_suppliers([first, second].into_iter().collect()));
        assert_eq!(episode.started_at, started_at);
        assert_eq!(episode.attempts, attempts);
        assert_eq!(episode.tried_suppliers, [first].into_iter().collect());
        assert_eq!(episode.next_probe_at, next_probe_at);
        assert!(!episode.alarmed);
        episode.record_failure(second, &clock, &FixedJitter(0));
        assert_eq!(episode.attempts, attempts.saturating_add(1));
        episode.restart(&clock);
        assert_eq!(episode.attempts, 0);
        assert!(episode.is_due(&clock));
    }

    #[test]
    fn supplier_churn_cannot_suppress_a_persistent_alarm() {
        let clock = clock();
        let first = source(3);
        let mut episode = episode(&clock, &[first]);
        let jitter = FixedJitter(0);

        for _ in 0..ALARM_ATTEMPTS {
            assert!(matches!(
                episode.record_failure(first, &clock, &jitter),
                RetryUpdate::RetryAt(_) | RetryUpdate::Alarmed { .. }
            ));
            if !episode.alarmed {
                clock.advance(Duration::seconds(60));
            }
        }
        assert!(episode.alarmed);
        let started_at = episode.started_at;
        let attempts = episode.attempts;
        let next_probe_at = episode.next_probe_at;

        let mut suppliers: BTreeSet<_> = [first].into_iter().collect();
        for value in 4..=10 {
            let supplier = source(value);
            assert!(suppliers.insert(supplier));
            assert!(episode.refresh_suppliers(suppliers.clone()));
            assert_eq!(episode.started_at, started_at);
            assert_eq!(episode.attempts, attempts);
            assert!(episode.alarmed);
            assert_eq!(episode.next_probe_at, next_probe_at);
            assert_eq!(episode.tried_suppliers, [first].into_iter().collect());
        }

        clock.advance(ALARM_PROBE_INTERVAL);
        assert_eq!(
            episode.record_failure(first, &clock, &jitter),
            RetryUpdate::ProbeAt(clock.now() + ALARM_PROBE_INTERVAL)
        );
        assert_eq!(episode.attempts, attempts.saturating_add(1));
        assert!(episode.alarmed);
    }

    #[test]
    fn restored_alarm_preserves_episode_age_attempts_and_probe_cadence() {
        let clock = clock();
        let first = source(3);
        let second = source(4);
        let suppliers = [first, second].into_iter().collect();
        let started_at = clock.now() - Duration::minutes(12);
        let next_probe_at = clock.now() + Duration::minutes(4);
        let summary = BodyUnavailableSummary {
            started_at,
            attempts: 14,
            suppliers: 2,
            supplier_set_digest: BodyUnavailableSummary::supplier_set_digest(&suppliers),
            alarmed: true,
            next_probe_at,
        };
        let mut episode = BodyRetryEpisode::restore(
            BranchId::new(hash(1), hash(2)),
            HeaderGeneration::new(3),
            Frontier::new(block::Height(20), hash(2)),
            suppliers,
            summary,
        );

        assert_eq!(episode.summary(), summary);
        assert!(!episode.is_due(&clock));
        assert_eq!(
            episode.record_failure(first, &clock, &FixedJitter(0)),
            RetryUpdate::TooEarly
        );
        clock.advance(Duration::minutes(4));
        assert_eq!(
            episode.record_failure(second, &clock, &FixedJitter(0)),
            RetryUpdate::ProbeAt(clock.now() + ALARM_PROBE_INTERVAL)
        );
        assert_eq!(episode.started_at, started_at);
        assert_eq!(episode.attempts, 15);
        assert!(episode.alarmed);
    }

    #[test]
    fn seeded_jitter_is_reproducible_and_within_the_normative_bound() {
        let episode = episode(&clock(), &[source(3)]);
        let jitter = SeededRetryJitter::new([7; 32]);
        let first = jitter.offset_per_thousand(episode.branch, episode.header, 1);
        assert_eq!(
            first,
            jitter.offset_per_thousand(episode.branch, episode.header, 1)
        );
        assert!((-100..=100).contains(&first));
        for attempt in 2..=64 {
            assert!((-100..=100).contains(&jitter.offset_per_thousand(
                episode.branch,
                episode.header,
                attempt
            )));
        }
    }

    #[test]
    fn elapsed_alarm_waits_until_every_known_supplier_was_tried() {
        let clock = clock();
        let first = source(3);
        let second = source(4);
        let mut episode = episode(&clock, &[first, second]);
        let jitter = FixedJitter(0);
        assert!(matches!(
            episode.record_failure(first, &clock, &jitter),
            RetryUpdate::RetryAt(_)
        ));
        clock.advance(ALARM_AFTER);
        assert!(matches!(
            episode.record_failure(first, &clock, &jitter),
            RetryUpdate::RetryAt(_)
        ));
        assert!(matches!(
            episode.record_failure(second, &clock, &jitter),
            RetryUpdate::Alarmed { .. }
        ));
    }

    #[test]
    fn generation_or_anchor_change_retires_retry_work_before_reuse() {
        let clock = clock();
        let episode = episode(&clock, &[source(3)]);
        let mut queue = BodyRetryQueue::default();
        assert!(queue.insert(episode.clone()).is_none());
        assert!(queue
            .get_mut(episode.generation, episode.branch, episode.header.hash)
            .is_some());
        queue.retain_current(
            HeaderGeneration::new(4),
            Frontier::new(block::Height(10), episode.branch.anchor_hash),
        );
        assert_eq!(queue.len(), 0);
        assert!(queue.is_empty());

        queue.insert(episode.clone());
        assert_eq!(
            queue.remove(episode.generation, episode.branch, episode.header.hash),
            Some(episode.clone())
        );
        queue.insert(episode.clone());
        queue.retain_current(
            episode.generation,
            Frontier::new(block::Height(11), hash(9)),
        );
        assert_eq!(queue.len(), 0);
    }

    #[test]
    fn retry_deadlines_saturate_at_the_clock_boundary() {
        assert_eq!(
            retry_deadline(DateTime::<Utc>::MAX_UTC, Duration::minutes(10)),
            DateTime::<Utc>::MAX_UTC
        );
    }
}