PiCloud/crates/executor-core/src/sdk/retry.rs

//! `retry::*` Rhai bridge — caller-controlled retry shape (v1.1.9).
//!
//! ```rhai
//! let p = retry::policy(#{
//!     max_attempts: 3,
//!     backoff: "exponential",  // "exponential" | "linear" | "constant"
//!     base_ms: 500,
//!     jitter_pct: 20,
//! });
//!
//! let result = retry::run(p, || {
//!     invoke("/payments/charge", #{ order_id: 1 })
//! });
//!
//! // Optional error-code filter:
//! let p2 = retry::on_codes(p, ["http: 503", "http: 504"]);
//! ```
//!
//! NOTE: the brief called for `retry::with(policy, closure)` but `with`
//! AND `call` are both Rhai reserved keywords (rejected at the parser
//! before any registration would matter) — so we ship `retry::run(...)`
//! instead. Documented in HANDBACK §7.
//!
//! `retry::run(policy, closure)` calls the closure; on throw it sleeps
//! per the policy and re-invokes. Sleeps run via `tokio::time::sleep`
//! through `TokioHandle::block_on` — the bridge is already inside the
//! caller's `spawn_blocking` thread, so blocking is fine.
//!
//! Closures share the caller's engine + cx (Rhai's `FnPtr` is invoked
//! via `call_within_context` on the same engine). If the closure calls
//! `invoke()`, the inner call gets a fresh cx with `trigger_depth + 1`
//! via the invoke bridge — retry doesn't see or interfere with that.

use std::collections::hash_map::DefaultHasher;
use std::hash::Hasher;
use std::sync::Arc;
use std::time::Duration;

use picloud_shared::{SdkCallCx, Services};
use rhai::{Array, Dynamic, Engine as RhaiEngine, EvalAltResult, FnPtr, Module, NativeCallContext};
use tokio::runtime::Handle as TokioHandle;

/// Policy value scripts construct via `retry::policy(opts)`. Custom Rhai
/// type — exposed via `register_type_with_name`.
#[derive(Debug, Clone)]
pub struct Policy {
    pub max_attempts: u32,
    pub backoff: BackoffShape,
    pub base_ms: u32,
    pub jitter_pct: u32,
    /// Empty → retry on any throw. Non-empty → retry only when the
    /// error string starts with one of these prefixes.
    pub on_codes: Vec<String>,
}

impl Default for Policy {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            backoff: BackoffShape::Exponential,
            base_ms: 500,
            jitter_pct: 20,
            on_codes: Vec::new(),
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub enum BackoffShape {
    Constant,
    Linear,
    Exponential,
}

impl BackoffShape {
    fn from_wire(s: &str) -> Option<Self> {
        match s {
            "constant" => Some(Self::Constant),
            "linear" => Some(Self::Linear),
            "exponential" => Some(Self::Exponential),
            _ => None,
        }
    }
}

pub(super) fn register(engine: &mut RhaiEngine, _services: &Services, _cx: Arc<SdkCallCx>) {
    // Custom type so scripts can pass Policy values around.
    engine.register_type_with_name::<Policy>("Policy");

    let mut module = Module::new();

    // `retry::policy(opts)` — opts is a Rhai Map.
    module.set_native_fn(
        "policy",
        move |opts: rhai::Map| -> Result<Policy, Box<EvalAltResult>> { build_policy(opts) },
    );

    // `retry::on_codes(policy, codes)` — returns a new policy with the
    // filter set. Takes ownership of the policy (Rhai semantics) and
    // returns the modified one.
    module.set_native_fn(
        "on_codes",
        move |policy: Policy, codes: Array| -> Result<Policy, Box<EvalAltResult>> {
            let mut p = policy;
            let mut out = Vec::with_capacity(codes.len());
            for c in codes {
                let s = c.into_string().map_err(|_| -> Box<EvalAltResult> {
                    EvalAltResult::ErrorRuntime(
                        "retry::on_codes: codes must be strings".into(),
                        rhai::Position::NONE,
                    )
                    .into()
                })?;
                out.push(s);
            }
            p.on_codes = out;
            Ok(p)
        },
    );

    // `retry::run(policy, fn_ptr)` — registered with NativeCallContext
    // so the bridge can re-invoke the FnPtr in a loop. Named `call` and
    // not `with` because `with` is a Rhai reserved keyword.
    module.set_native_fn(
        "run",
        move |ctx: NativeCallContext,
              policy: Policy,
              fn_ptr: FnPtr|
              -> Result<Dynamic, Box<EvalAltResult>> { retry_run(&ctx, &policy, &fn_ptr) },
    );

    engine.register_static_module("retry", module.into());
}

fn build_policy(opts: rhai::Map) -> Result<Policy, Box<EvalAltResult>> {
    let mut p = Policy::default();
    if let Some(v) = opts.get("max_attempts") {
        let n = v.as_int().map_err(|_| -> Box<EvalAltResult> {
            EvalAltResult::ErrorRuntime(
                "retry::policy: max_attempts must be an integer".into(),
                rhai::Position::NONE,
            )
            .into()
        })?;
        let n = u32::try_from(n).unwrap_or(0);
        p.max_attempts = n.clamp(1, 20);
    }
    if let Some(v) = opts.get("backoff") {
        let s = v.clone().into_string().map_err(|_| -> Box<EvalAltResult> {
            EvalAltResult::ErrorRuntime(
                "retry::policy: backoff must be a string".into(),
                rhai::Position::NONE,
            )
            .into()
        })?;
        p.backoff = BackoffShape::from_wire(&s).ok_or_else(|| -> Box<EvalAltResult> {
            EvalAltResult::ErrorRuntime(
                "retry::policy: backoff must be \"exponential\" | \"linear\" | \"constant\"".into(),
                rhai::Position::NONE,
            )
            .into()
        })?;
    }
    if let Some(v) = opts.get("base_ms") {
        let n = v.as_int().map_err(|_| -> Box<EvalAltResult> {
            EvalAltResult::ErrorRuntime(
                "retry::policy: base_ms must be an integer".into(),
                rhai::Position::NONE,
            )
            .into()
        })?;
        let n = u32::try_from(n).unwrap_or(0);
        p.base_ms = n.clamp(1, 60_000);
    }
    if let Some(v) = opts.get("jitter_pct") {
        let n = v.as_int().map_err(|_| -> Box<EvalAltResult> {
            EvalAltResult::ErrorRuntime(
                "retry::policy: jitter_pct must be an integer".into(),
                rhai::Position::NONE,
            )
            .into()
        })?;
        let n = u32::try_from(n).unwrap_or(0);
        p.jitter_pct = n.clamp(0, 100);
    }
    Ok(p)
}

fn retry_run(
    ctx: &NativeCallContext,
    policy: &Policy,
    fn_ptr: &FnPtr,
) -> Result<Dynamic, Box<EvalAltResult>> {
    let mut attempt: u32 = 0;
    loop {
        attempt += 1;
        // Rhai 1.24's FnPtr only carries the callee's name + args; we
        // invoke it through the engine the NativeCallContext exposes.
        // `()` as args: the closure is nullary — the script-side helper
        // is `|| { ... }`.
        let res: Result<Dynamic, Box<EvalAltResult>> = fn_ptr.call_within_context(ctx, ());
        match res {
            Ok(v) => return Ok(v),
            Err(e) => {
                // on_codes filter: if non-empty, only retry when the
                // error string starts with one of the codes.
                if !policy.on_codes.is_empty() {
                    let msg = e.to_string();
                    let matched = policy.on_codes.iter().any(|c| msg.contains(c));
                    if !matched {
                        return Err(e);
                    }
                }
                if attempt >= policy.max_attempts {
                    return Err(e);
                }
                let delay = compute_delay(policy, attempt);
                let _ = sleep_blocking(delay);
            }
        }
    }
}

fn compute_delay(policy: &Policy, attempt: u32) -> Duration {
    let base_ms = u64::from(policy.base_ms);
    let exp_pow = u64::from(attempt.saturating_sub(1));
    let raw = match policy.backoff {
        BackoffShape::Constant => base_ms,
        BackoffShape::Linear => base_ms.saturating_mul(u64::from(attempt)),
        BackoffShape::Exponential => base_ms.saturating_mul(1u64 << exp_pow.min(20)),
    };
    let raw = raw.min(u64::from(u32::MAX));
    let jittered = apply_jitter(u32::try_from(raw).unwrap_or(u32::MAX), policy.jitter_pct);
    Duration::from_millis(u64::from(jittered))
}

fn apply_jitter(raw: u32, pct: u32) -> u32 {
    if pct == 0 {
        return raw;
    }
    let pct = pct.min(100);
    let span = u64::from(raw) * u64::from(pct) / 100;
    if span == 0 {
        return raw;
    }
    // Deterministic-enough jitter without pulling rand into a hot path:
    // pick a small offset from the high bits of attempt's hash. Random
    // distribution doesn't matter for retry — bounded ± is what we want.
    let mut h = DefaultHasher::new();
    h.write_u64(span);
    h.write_u32(raw);
    let r = h.finish();
    // span is bounded by raw*100/100 = raw ≤ u32::MAX, so 2*span+1 fits
    // in i64 without wrapping. `cast_signed` makes the intent explicit
    // for clippy::cast_possible_wrap.
    let modded = r % (2 * span + 1);
    let offset =
        i64::try_from(modded).unwrap_or(i64::MAX) - i64::try_from(span).unwrap_or(i64::MAX);
    let signed = i64::from(raw).saturating_add(offset).max(0);
    u32::try_from(signed.min(i64::from(u32::MAX))).unwrap_or(u32::MAX)
}

/// Sleep blocking for `delay`. We're inside the caller's spawn_blocking
/// thread, so block_on of a tokio::time::sleep is fine — the runtime
/// thread isn't being held.
fn sleep_blocking(delay: Duration) -> Result<(), Box<EvalAltResult>> {
    let handle = TokioHandle::try_current().map_err(|e| -> Box<EvalAltResult> {
        EvalAltResult::ErrorRuntime(
            format!("retry: no tokio runtime: {e}").into(),
            rhai::Position::NONE,
        )
        .into()
    })?;
    handle.block_on(async move { tokio::time::sleep(delay).await });
    Ok(())
}

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

    #[test]
    fn policy_default_is_reasonable() {
        let p = Policy::default();
        assert_eq!(p.max_attempts, 3);
        assert_eq!(p.base_ms, 500);
        assert_eq!(p.jitter_pct, 20);
        assert!(p.on_codes.is_empty());
    }

    #[test]
    fn policy_clamps_max_attempts() {
        let mut m = rhai::Map::new();
        m.insert("max_attempts".into(), Dynamic::from(0_i64));
        let p = build_policy(m).unwrap();
        assert_eq!(p.max_attempts, 1);

        let mut m = rhai::Map::new();
        m.insert("max_attempts".into(), Dynamic::from(999_i64));
        let p = build_policy(m).unwrap();
        assert_eq!(p.max_attempts, 20);
    }

    #[test]
    fn policy_clamps_base_ms_and_jitter() {
        let mut m = rhai::Map::new();
        m.insert("base_ms".into(), Dynamic::from(0_i64));
        m.insert("jitter_pct".into(), Dynamic::from(200_i64));
        let p = build_policy(m).unwrap();
        assert_eq!(p.base_ms, 1);
        assert_eq!(p.jitter_pct, 100);
    }

    #[test]
    fn policy_rejects_bogus_backoff_shape() {
        let mut m = rhai::Map::new();
        m.insert("backoff".into(), Dynamic::from("bogus"));
        assert!(build_policy(m).is_err());
    }

    #[test]
    fn compute_delay_exponential_doubles() {
        let p = Policy {
            max_attempts: 5,
            backoff: BackoffShape::Exponential,
            base_ms: 1000,
            jitter_pct: 0,
            on_codes: Vec::new(),
        };
        assert_eq!(compute_delay(&p, 1).as_millis(), 1000);
        assert_eq!(compute_delay(&p, 2).as_millis(), 2000);
        assert_eq!(compute_delay(&p, 3).as_millis(), 4000);
    }

    #[test]
    fn compute_delay_linear() {
        let p = Policy {
            max_attempts: 5,
            backoff: BackoffShape::Linear,
            base_ms: 100,
            jitter_pct: 0,
            on_codes: Vec::new(),
        };
        assert_eq!(compute_delay(&p, 1).as_millis(), 100);
        assert_eq!(compute_delay(&p, 2).as_millis(), 200);
        assert_eq!(compute_delay(&p, 5).as_millis(), 500);
    }

    #[test]
    fn compute_delay_constant() {
        let p = Policy {
            max_attempts: 5,
            backoff: BackoffShape::Constant,
            base_ms: 250,
            jitter_pct: 0,
            on_codes: Vec::new(),
        };
        assert_eq!(compute_delay(&p, 1).as_millis(), 250);
        assert_eq!(compute_delay(&p, 4).as_millis(), 250);
    }
}