pe_graph/

retry.rs

//! Retry policy -- per-phase retry with exponential backoff and jitter.
//!
//! Used by the `node!` DSL `#[retry]` annotation and directly by users
//! who want retry semantics on fallible async operations within nodes.

use pe_core::node::NodeResult;
use pe_core::state::StateUpdate;
use serde::{Deserialize, Serialize};
use std::future::Future;
use std::time::Duration;

/// Configuration for retrying a fallible operation with backoff.
///
/// # Example
///
/// ```ignore
/// use pe_graph::retry::{RetryPolicy, with_retry};
///
/// let result = with_retry(&RetryPolicy::default(), || {
///     Box::pin(async { call_llm().await })
/// }).await;
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetryPolicy {
    /// Maximum number of retry attempts (not counting the initial attempt).
    pub max_attempts: u32,
    /// Delay before the first retry.
    pub initial_interval: Duration,
    /// Multiplier applied to the delay after each retry.
    pub backoff_factor: f64,
    /// Maximum delay between retries (caps exponential growth).
    pub max_interval: Duration,
    /// Whether to add random jitter to delays (prevents thundering herd).
    pub jitter: bool,
}

impl Default for RetryPolicy {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            initial_interval: Duration::from_millis(200),
            backoff_factor: 2.0,
            max_interval: Duration::from_secs(10),
            jitter: true,
        }
    }
}

/// Execute an async operation with retry semantics.
///
/// Retries only on [`PeError::is_retryable`](pe_core::PeError::is_retryable) errors. Non-retryable errors
/// and all non-error results (`Update`, `Interrupt`, `Converge`) are
/// returned immediately.
///
/// The closure `f` is called for each attempt. It must return a pinned future.
///
/// # Example
///
/// ```ignore
/// let result = with_retry(&RetryPolicy::default(), || {
///     Box::pin(async { do_work().await })
/// }).await;
/// ```
pub async fn with_retry<F, Fut, U>(policy: &RetryPolicy, f: F) -> NodeResult<U>
where
    F: Fn() -> Fut,
    Fut: Future<Output = NodeResult<U>>,
    U: StateUpdate,
{
    let mut attempts = 0u32;
    let mut delay = policy.initial_interval;

    loop {
        let result = f().await;

        match &result {
            NodeResult::Error(e) if e.is_retryable() && attempts < policy.max_attempts => {
                attempts += 1;
                let sleep_dur = if policy.jitter {
                    apply_jitter(delay)
                } else {
                    delay
                };
                tokio::time::sleep(sleep_dur).await;
                delay = next_delay(delay, policy.backoff_factor, policy.max_interval);
            }
            _ => return result,
        }
    }
}

/// Compute the next backoff delay, capped at max_interval.
pub(crate) fn next_delay(current: Duration, factor: f64, max: Duration) -> Duration {
    let next = current.mul_f64(factor);
    if next > max { max } else { next }
}

/// Apply jitter: random value in [50%, 150%) of the delay.
pub(crate) fn apply_jitter(delay: Duration) -> Duration {
    let nanos = delay.as_nanos() as u64;
    let jitter_nanos = pseudo_random_u64(nanos.max(1));
    let half = nanos / 2;
    Duration::from_nanos(half + jitter_nanos)
}

/// Pseudo-random u64 in [0, range) using `RandomState` entropy.
/// Not cryptographic -- only used for jitter timing.
/// Uses `RandomState` instead of `SystemTime` to avoid identical
/// values when called at the same nanosecond.
fn pseudo_random_u64(range: u64) -> u64 {
    if range == 0 {
        return 0;
    }
    use std::collections::hash_map::RandomState;
    use std::hash::{BuildHasher, Hasher};
    let mut hasher = RandomState::new().build_hasher();
    hasher.write_u64(range);
    hasher.finish() % range
}

#[cfg(test)]
mod tests {
    use super::*;
    use pe_core::error::PeError;
    use pe_core::node::NodeResult;
    use std::sync::atomic::{AtomicU32, Ordering};

    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    struct DummyUpdate;
    impl pe_core::state::StateUpdate for DummyUpdate {}

    #[tokio::test]
    async fn test_succeeds_on_first_try() {
        let policy = RetryPolicy {
            max_attempts: 3,
            jitter: false,
            ..Default::default()
        };
        let result = with_retry(&policy, || async {
            NodeResult::<DummyUpdate>::Update(DummyUpdate)
        })
        .await;
        assert!(matches!(result, NodeResult::Update(_)));
    }

    #[tokio::test]
    async fn test_retries_then_succeeds() {
        let attempts = AtomicU32::new(0);
        let policy = RetryPolicy {
            max_attempts: 3,
            initial_interval: Duration::from_millis(1),
            jitter: false,
            ..Default::default()
        };

        let result = with_retry(&policy, || {
            let count = attempts.fetch_add(1, Ordering::SeqCst);
            async move {
                if count < 2 {
                    NodeResult::<DummyUpdate>::Error(PeError::Timeout { seconds: 1.0 })
                } else {
                    NodeResult::Update(DummyUpdate)
                }
            }
        })
        .await;

        assert!(matches!(result, NodeResult::Update(_)));
        assert_eq!(attempts.load(Ordering::SeqCst), 3); // initial + 2 retries
    }

    #[tokio::test]
    async fn test_exhausts_retries_returns_error() {
        let attempts = AtomicU32::new(0);
        let policy = RetryPolicy {
            max_attempts: 2,
            initial_interval: Duration::from_millis(1),
            jitter: false,
            ..Default::default()
        };

        let result = with_retry(&policy, || {
            attempts.fetch_add(1, Ordering::SeqCst);
            async { NodeResult::<DummyUpdate>::Error(PeError::Timeout { seconds: 1.0 }) }
        })
        .await;

        assert!(matches!(result, NodeResult::Error(PeError::Timeout { .. })));
        // 1 initial + 2 retries = 3 total attempts
        assert_eq!(attempts.load(Ordering::SeqCst), 3);
    }

    #[tokio::test]
    async fn test_non_retryable_error_not_retried() {
        let attempts = AtomicU32::new(0);
        let policy = RetryPolicy {
            max_attempts: 3,
            initial_interval: Duration::from_millis(1),
            jitter: false,
            ..Default::default()
        };

        let result = with_retry(&policy, || {
            attempts.fetch_add(1, Ordering::SeqCst);
            async {
                NodeResult::<DummyUpdate>::Error(PeError::PermissionDenied {
                    action: "test".into(),
                })
            }
        })
        .await;

        assert!(matches!(
            result,
            NodeResult::Error(PeError::PermissionDenied { .. })
        ));
        assert_eq!(attempts.load(Ordering::SeqCst), 1); // no retries
    }

    #[tokio::test]
    async fn test_interrupt_not_retried() {
        let policy = RetryPolicy {
            max_attempts: 3,
            jitter: false,
            ..Default::default()
        };

        let result = with_retry(&policy, || async {
            NodeResult::<DummyUpdate>::Interrupt(pe_core::node::InterruptRequest {
                reason: "test".into(),
                partial_update: None,
                resume_point: "test:0".into(),
            })
        })
        .await;

        assert!(matches!(result, NodeResult::Interrupt(_)));
    }

    #[test]
    fn test_next_delay_exponential() {
        let d = next_delay(Duration::from_millis(100), 2.0, Duration::from_secs(10));
        assert_eq!(d, Duration::from_millis(200));
    }

    #[test]
    fn test_next_delay_capped_at_max() {
        let d = next_delay(Duration::from_secs(8), 2.0, Duration::from_secs(10));
        assert_eq!(d, Duration::from_secs(10));
    }
}