remora/crates/remora-stream/src/timing.rs

use std::time::Duration;

/// Adapter-selected timing bounds below the JobProtocol attempt budget.
///
/// Connect, response-start, and no-progress timeouts map to structured runtime
/// failures. There is no elapsed wall-clock job deadline.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RetryTiming {
    pub connect_timeout: Duration,
    pub response_start_timeout: Duration,
    pub no_progress_timeout: Duration,
    pub base_backoff: Duration,
    pub max_backoff: Duration,
    pub max_server_hint: Duration,
}

impl RetryTiming {
    /// Conservative defaults suitable for synthetic and small-object transfers.
    #[must_use]
    pub const fn default_bounds() -> Self {
        Self {
            connect_timeout: Duration::from_secs(10),
            response_start_timeout: Duration::from_secs(30),
            no_progress_timeout: Duration::from_secs(60),
            base_backoff: Duration::from_millis(100),
            max_backoff: Duration::from_secs(30),
            max_server_hint: Duration::from_secs(60),
        }
    }

    /// Validate and clamp configuration; zero or inverted bounds become defaults.
    #[must_use]
    pub fn sanitized(self) -> Self {
        let defaults = Self::default_bounds();
        let base = if self.base_backoff.is_zero() {
            defaults.base_backoff
        } else {
            self.base_backoff
        };
        let max_backoff = if self.max_backoff < base {
            base
        } else {
            self.max_backoff
        };
        let max_server_hint = if self.max_server_hint.is_zero() {
            defaults.max_server_hint
        } else {
            self.max_server_hint
        };
        Self {
            connect_timeout: nonzero_or(self.connect_timeout, defaults.connect_timeout),
            response_start_timeout: nonzero_or(
                self.response_start_timeout,
                defaults.response_start_timeout,
            ),
            no_progress_timeout: nonzero_or(self.no_progress_timeout, defaults.no_progress_timeout),
            base_backoff: base,
            max_backoff,
            max_server_hint,
        }
    }
}

fn nonzero_or(value: Duration, fallback: Duration) -> Duration {
    if value.is_zero() { fallback } else { value }
}

/// Exponential backoff ceiling before jitter: `min(max, base * 2^attempt)`.
#[must_use]
pub fn saturating_backoff(attempt: u32, timing: &RetryTiming) -> Duration {
    let timing = timing.sanitized();
    let shift = attempt.min(16);
    let factor = 1u64.checked_shl(shift).unwrap_or(u64::MAX);
    let base_ms = u64::try_from(timing.base_backoff.as_millis()).unwrap_or(u64::MAX);
    let scaled = base_ms.saturating_mul(factor);
    let max_ms = u64::try_from(timing.max_backoff.as_millis()).unwrap_or(u64::MAX);
    Duration::from_millis(scaled.min(max_ms))
}

/// Full-jitter delay in `[0, saturating_backoff]`.
///
/// If the OS RNG is unavailable, returns the full backoff ceiling so retries
/// do not collapse to a synchronized zero delay.
#[must_use]
pub fn full_jitter_delay(attempt: u32, timing: &RetryTiming) -> Duration {
    let ceiling = saturating_backoff(attempt, timing);
    match random_u64() {
        Some(random) => full_jitter_delay_with(ceiling, random),
        None => ceiling,
    }
}

/// Full-jitter with an injected random value for deterministic tests.
#[must_use]
pub fn full_jitter_delay_with(ceiling: Duration, random: u64) -> Duration {
    if ceiling.is_zero() {
        return Duration::ZERO;
    }
    let ceiling_ms = u64::try_from(ceiling.as_millis()).unwrap_or(u64::MAX);
    if ceiling_ms == 0 {
        return Duration::ZERO;
    }
    let chosen = random % (ceiling_ms.saturating_add(1));
    Duration::from_millis(chosen)
}

/// Cap or reject a server-supplied retry hint.
///
/// Zero, missing, or excessive hints yield `None` so adapters ignore them.
#[must_use]
pub fn capped_server_hint(hint: Option<Duration>, timing: &RetryTiming) -> Option<Duration> {
    let timing = timing.sanitized();
    let hint = hint?;
    if hint.is_zero() || hint > timing.max_server_hint {
        return None;
    }
    Some(hint)
}

fn random_u64() -> Option<u64> {
    let mut bytes = [0u8; 8];
    getrandom::fill(&mut bytes).ok()?;
    Some(u64::from_le_bytes(bytes))
}

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

    #[test]
    fn backoff_grows_then_caps() {
        let timing = RetryTiming {
            base_backoff: Duration::from_millis(100),
            max_backoff: Duration::from_secs(1),
            ..RetryTiming::default_bounds()
        };
        assert_eq!(saturating_backoff(0, &timing), Duration::from_millis(100));
        assert_eq!(saturating_backoff(1, &timing), Duration::from_millis(200));
        assert_eq!(saturating_backoff(2, &timing), Duration::from_millis(400));
        assert_eq!(saturating_backoff(10, &timing), Duration::from_secs(1));
    }

    #[test]
    fn full_jitter_stays_within_ceiling() {
        let ceiling = Duration::from_millis(1000);
        for random in [0, 1, 500, 999, 1000, u64::MAX] {
            let delay = full_jitter_delay_with(ceiling, random);
            assert!(delay <= ceiling);
        }
    }

    #[test]
    fn server_hints_are_capped_or_ignored() {
        let timing = RetryTiming {
            max_server_hint: Duration::from_secs(5),
            ..RetryTiming::default_bounds()
        };
        assert_eq!(
            capped_server_hint(Some(Duration::from_secs(3)), &timing),
            Some(Duration::from_secs(3))
        );
        assert_eq!(
            capped_server_hint(Some(Duration::from_secs(6)), &timing),
            None
        );
        assert_eq!(capped_server_hint(Some(Duration::ZERO), &timing), None);
        assert_eq!(capped_server_hint(None, &timing), None);
    }
}