PiCloud/crates/orchestrator-core/src/client.rs

use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use lru::LruCache;
use picloud_executor_core::{Engine, ExecError, ExecRequest, ExecResponse};
use picloud_shared::ScriptId;

use crate::gate::{AcquireError, ExecutionGate};

/// Maximum wall-clock time we'll wait for a single invocation, regardless
/// of the per-script `timeout_seconds`. Provides a hard ceiling on
/// resource usage independent of misconfigured scripts.
const HARD_TIMEOUT_CAP: Duration = Duration::from_secs(300);

/// Default capacity for the top-level script AST cache. Override via
/// `PICLOUD_SCRIPT_CACHE_SIZE`. Sized assuming a few hundred distinct
/// endpoint scripts per process.
const DEFAULT_SCRIPT_CACHE_SIZE: usize = 256;

/// Identity used by [`ExecutorClient::execute_with_identity`] to key
/// the AST cache. `updated_at` is the freshness comparator — an edit
/// that bumps `scripts.updated_at` invalidates the cached AST on the
/// next lookup, no explicit pub/sub.
#[derive(Debug, Clone, Copy)]
pub struct ScriptIdentity {
    pub script_id: ScriptId,
    pub updated_at: DateTime<Utc>,
}

/// The seam between the orchestrator and the executor.
///
/// Single-node mode plugs in `LocalExecutorClient`, which calls
/// `executor-core` in-process via `spawn_blocking`. Cluster mode plugs
/// in `RemoteExecutorClient`, which forwards over HTTP to an executor
/// node. Everything else in orchestrator-core depends only on this trait.
#[async_trait]
pub trait ExecutorClient: Send + Sync {
    async fn execute(
        &self,
        source: &str,
        req: ExecRequest,
        timeout: Duration,
    ) -> Result<ExecResponse, ExecError>;

    /// v1.1.3: identity-aware variant for caching. Callers that already
    /// know the script's `(id, updated_at)` should use this so the local
    /// executor can reuse a compiled `rhai::AST` across invocations.
    /// Default impl forwards to `execute` so `RemoteExecutorClient` (and
    /// any future transport) keeps working without bespoke caching.
    async fn execute_with_identity(
        &self,
        _identity: ScriptIdentity,
        source: &str,
        req: ExecRequest,
        timeout: Duration,
    ) -> Result<ExecResponse, ExecError> {
        self.execute(source, req, timeout).await
    }
}

/// In-process executor — wraps `executor-core::Engine` directly.
///
/// `executor-core::Engine::execute` is synchronous; we offload it to a
/// blocking thread so it doesn't park a Tokio worker, and apply the
/// wall-clock timeout here.
///
/// Holds an `ExecutionGate` and acquires a permit before `spawn_blocking`
/// so a script storm can't drain the blocking-thread pool. The permit
/// drops with the future, returning the slot.
///
/// v1.1.3 adds a top-level AST cache keyed by `ScriptId`. On
/// `execute_with_identity`, the client compares the caller's
/// `updated_at` against the cached entry's; a match reuses the
/// `Arc<rhai::AST>` and skips Rhai's parser. A mismatch (or absence)
/// triggers a fresh `Engine::compile` + replace.
pub struct LocalExecutorClient {
    engine: Arc<Engine>,
    gate: Arc<ExecutionGate>,
    /// `(updated_at, Arc<rhai::AST>)` keyed by `ScriptId`. `Mutex`
    /// because the cache is shared across invocations of this client;
    /// LRU eviction caps memory growth.
    script_cache: Arc<Mutex<LruCache<ScriptId, CachedScript>>>,
}

pub struct CachedScript {
    pub updated_at: DateTime<Utc>,
    pub ast: Arc<rhai::AST>,
}

impl LocalExecutorClient {
    #[must_use]
    pub fn new(engine: Arc<Engine>, gate: Arc<ExecutionGate>) -> Self {
        let cap = std::env::var("PICLOUD_SCRIPT_CACHE_SIZE")
            .ok()
            .and_then(|s| s.parse::<usize>().ok())
            .unwrap_or(DEFAULT_SCRIPT_CACHE_SIZE);
        Self::with_script_cache_capacity(engine, gate, cap)
    }

    /// Explicit capacity for tests that exercise LRU eviction.
    #[must_use]
    pub fn with_script_cache_capacity(
        engine: Arc<Engine>,
        gate: Arc<ExecutionGate>,
        cap: usize,
    ) -> Self {
        let cap = NonZeroUsize::new(cap.max(1)).expect("max(1) is non-zero");
        Self {
            engine,
            gate,
            script_cache: Arc::new(Mutex::new(LruCache::new(cap))),
        }
    }

    /// Cache lookup with `updated_at` freshness check. Returns the
    /// cached AST on hit; compiles, inserts, returns the fresh AST on
    /// miss or stale. Public so tests can introspect the cache.
    pub fn get_or_compile(
        &self,
        identity: ScriptIdentity,
        source: &str,
    ) -> Result<Arc<rhai::AST>, ExecError> {
        {
            let mut cache = self
                .script_cache
                .lock()
                .expect("script cache lock poisoned");
            if let Some(cached) = cache.get(&identity.script_id) {
                if cached.updated_at == identity.updated_at {
                    tracing::debug!(
                        target = "picloud::scripts::cache",
                        script_id = %identity.script_id,
                        "cache hit"
                    );
                    return Ok(cached.ast.clone());
                }
                tracing::debug!(
                    target = "picloud::scripts::cache",
                    script_id = %identity.script_id,
                    "cache stale; recompiling"
                );
            } else {
                tracing::debug!(
                    target = "picloud::scripts::cache",
                    script_id = %identity.script_id,
                    "cache miss"
                );
            }
        }
        let ast = self.engine.compile(source)?;
        let mut cache = self
            .script_cache
            .lock()
            .expect("script cache lock poisoned");
        cache.put(
            identity.script_id,
            CachedScript {
                updated_at: identity.updated_at,
                ast: ast.clone(),
            },
        );
        Ok(ast)
    }

    /// Shared script-AST cache. Exposed so tests can introspect cache
    /// state (length / contents) under a Mutex lock.
    #[must_use]
    pub fn script_cache(&self) -> &Arc<Mutex<LruCache<ScriptId, CachedScript>>> {
        &self.script_cache
    }
}

#[async_trait]
impl ExecutorClient for LocalExecutorClient {
    async fn execute(
        &self,
        source: &str,
        req: ExecRequest,
        timeout: Duration,
    ) -> Result<ExecResponse, ExecError> {
        // Acquire before spending any wall-clock budget. The permit is
        // held by this future; on `tokio::time::timeout` firing, the
        // future drops and the permit returns to the pool. Audit
        // 2026-06-11 H-C1: the detached `spawn_blocking` thread is now
        // also bounded by the Rhai `on_progress` deadline hook (installed
        // via `execute_with_deadline`), so a runaway script terminates
        // soon after the deadline instead of holding its OS thread until
        // the per-op budget self-exhausts.
        let _permit =
            self.gate
                .try_acquire()
                .map_err(
                    |AcquireError::Overloaded { retry_after_secs }| ExecError::Overloaded {
                        retry_after_secs,
                    },
                )?;

        let timeout = timeout.min(HARD_TIMEOUT_CAP);
        let timeout_secs = u32::try_from(timeout.as_secs()).unwrap_or(u32::MAX);
        let deadline = Instant::now() + timeout;

        let engine = self.engine.clone();
        let source = source.to_string();
        let join = tokio::task::spawn_blocking(move || {
            engine.execute_with_deadline(&source, req, Some(deadline))
        });

        match tokio::time::timeout(timeout, join).await {
            Err(_) => Err(ExecError::Timeout(timeout_secs)),
            Ok(Err(join_err)) => Err(ExecError::Runtime(format!(
                "execution task panicked: {join_err}"
            ))),
            Ok(Ok(res)) => res,
        }
    }

    async fn execute_with_identity(
        &self,
        identity: ScriptIdentity,
        source: &str,
        req: ExecRequest,
        timeout: Duration,
    ) -> Result<ExecResponse, ExecError> {
        let _permit =
            self.gate
                .try_acquire()
                .map_err(
                    |AcquireError::Overloaded { retry_after_secs }| ExecError::Overloaded {
                        retry_after_secs,
                    },
                )?;

        let ast = self.get_or_compile(identity, source)?;

        let timeout = timeout.min(HARD_TIMEOUT_CAP);
        let timeout_secs = u32::try_from(timeout.as_secs()).unwrap_or(u32::MAX);
        let deadline = Instant::now() + timeout;

        let engine = self.engine.clone();
        let join = tokio::task::spawn_blocking(move || {
            engine.execute_ast_with_deadline(&ast, req, Some(deadline))
        });

        match tokio::time::timeout(timeout, join).await {
            Err(_) => Err(ExecError::Timeout(timeout_secs)),
            Ok(Err(join_err)) => Err(ExecError::Runtime(format!(
                "execution task panicked: {join_err}"
            ))),
            Ok(Ok(res)) => res,
        }
    }
}

/// Remote executor — forwards to a peer executor node over HTTP.
///
/// Skeleton only; fleshed out when cluster mode lands.
pub struct RemoteExecutorClient {
    _client: reqwest::Client,
    _base_url: String,
}

impl RemoteExecutorClient {
    #[must_use]
    pub fn new(base_url: impl Into<String>) -> Self {
        Self {
            _client: reqwest::Client::new(),
            _base_url: base_url.into(),
        }
    }
}

#[async_trait]
impl ExecutorClient for RemoteExecutorClient {
    async fn execute(
        &self,
        _source: &str,
        _req: ExecRequest,
        _timeout: Duration,
    ) -> Result<ExecResponse, ExecError> {
        Err(ExecError::Runtime(
            "RemoteExecutorClient not implemented (cluster mode is v1.3+)".into(),
        ))
    }
}

#[cfg(test)]
mod cache_tests {
    use super::*;
    use picloud_executor_core::Limits;
    use picloud_shared::Services;

    fn engine() -> Arc<Engine> {
        Arc::new(Engine::new(Limits::default(), Services::default()))
    }

    fn client_with_cap(cap: usize) -> LocalExecutorClient {
        LocalExecutorClient::with_script_cache_capacity(
            engine(),
            Arc::new(ExecutionGate::new(32)),
            cap,
        )
    }

    fn identity_at(t: DateTime<Utc>) -> ScriptIdentity {
        ScriptIdentity {
            script_id: ScriptId::new(),
            updated_at: t,
        }
    }

    #[test]
    fn cache_hit_when_identity_matches() {
        let client = client_with_cap(8);
        let identity = identity_at(Utc::now());
        let src = "fn f() { 1 }";

        let ast_a = client.get_or_compile(identity, src).unwrap();
        let ast_b = client.get_or_compile(identity, src).unwrap();

        // Same Arc — cache served the second call without recompiling.
        assert!(
            Arc::ptr_eq(&ast_a, &ast_b),
            "expected identical Arc<AST> from cache hit"
        );
    }

    #[test]
    fn cache_invalidated_when_updated_at_changes() {
        let client = client_with_cap(8);
        let script_id = ScriptId::new();
        let t0 = Utc::now() - chrono::Duration::seconds(10);
        let t1 = Utc::now();

        let ast_a = client
            .get_or_compile(
                ScriptIdentity {
                    script_id,
                    updated_at: t0,
                },
                "fn f() { 1 }",
            )
            .unwrap();
        let ast_b = client
            .get_or_compile(
                ScriptIdentity {
                    script_id,
                    updated_at: t1,
                },
                "fn f() { 2 }",
            )
            .unwrap();

        // Different Arc — cache miss forced recompile.
        assert!(
            !Arc::ptr_eq(&ast_a, &ast_b),
            "expected recompile on updated_at change"
        );
    }

    #[test]
    fn distinct_script_ids_cache_independently() {
        let client = client_with_cap(8);
        let now = Utc::now();
        let a = identity_at(now);
        let b = identity_at(now);
        client.get_or_compile(a, "fn x() { 1 }").unwrap();
        client.get_or_compile(b, "fn x() { 1 }").unwrap();

        let cache = client.script_cache().lock().unwrap();
        assert_eq!(
            cache.len(),
            2,
            "distinct script_ids should yield two entries"
        );
    }

    #[test]
    fn lru_eviction_caps_cache_size() {
        // Capacity 1 — every new script evicts the previous.
        let client = client_with_cap(1);
        client
            .get_or_compile(identity_at(Utc::now()), "fn a() { 1 }")
            .unwrap();
        client
            .get_or_compile(identity_at(Utc::now()), "fn b() { 2 }")
            .unwrap();
        client
            .get_or_compile(identity_at(Utc::now()), "fn c() { 3 }")
            .unwrap();
        assert_eq!(client.script_cache().lock().unwrap().len(), 1);
    }

    #[test]
    fn script_identity_is_copy() {
        // Copy is load-bearing — many call sites pass it by value.
        let id = identity_at(Utc::now());
        let _ = id;
        let _ = id; // should still be usable
    }

    #[test]
    fn compile_error_does_not_poison_cache() {
        let client = client_with_cap(8);
        let identity = identity_at(Utc::now());
        // Bad source — should error and not insert anything.
        let res = client.get_or_compile(identity, "@@@ not valid rhai @@@");
        assert!(res.is_err(), "garbage source should fail to compile");
        // A subsequent good compile under a fresh identity must still work.
        let good = client.get_or_compile(identity_at(Utc::now()), "fn ok() { 1 }");
        assert!(good.is_ok());
    }
}