orchard/engine/

lifecycle.rs

//! Inference engine process lifecycle management.

use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Mutex;
use std::time::{Duration, Instant};

use fs4::fs_std::FileExt;
use nng::options::Options;
use nng::{Protocol, Socket};
use serde_json::json;

use crate::engine::fetch::EngineFetcher;
use crate::engine::multiprocess::{
    pid_is_alive, pid_is_engine, read_pid_file, reap_engine_process, stop_engine_process,
    write_pid_file,
};
use crate::error::{Error, Result};
use crate::ipc::endpoints::{management_url, response_url, EVENT_TOPIC_PREFIX};

const DEFAULT_STARTUP_TIMEOUT_SECS: u64 = 60;
const LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(5);
const LOCK_BACKOFF_INITIAL: Duration = Duration::from_millis(10);
const LOCK_BACKOFF_MAX: Duration = Duration::from_millis(250);

/// Remove a file, logging non-NotFound errors.
fn remove_if_exists(path: &std::path::Path) {
    if let Err(e) = std::fs::remove_file(path) {
        if e.kind() != std::io::ErrorKind::NotFound {
            tracing::warn!("Failed to remove {}: {}", path.display(), e);
        }
    }
}

fn lock_exclusive_with_timeout(lock_file: &std::fs::File) -> Result<()> {
    let deadline = Instant::now() + LOCK_ACQUIRE_TIMEOUT;
    let mut sleep = LOCK_BACKOFF_INITIAL;

    loop {
        match lock_file.try_lock_exclusive() {
            Ok(()) => return Ok(()),
            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                if Instant::now() >= deadline {
                    return Err(Error::LockFailed(format!(
                        "Timed out acquiring engine lock after {:?}",
                        LOCK_ACQUIRE_TIMEOUT
                    )));
                }
                std::thread::sleep(sleep);
                sleep = std::cmp::min(sleep * 2, LOCK_BACKOFF_MAX);
            }
            Err(e) => return Err(Error::LockFailed(e.to_string())),
        }
    }
}

/// File paths used by the engine lifecycle manager.
#[derive(Debug, Clone)]
pub struct EnginePaths {
    pub cache_dir: PathBuf,
    pub pid_file: PathBuf,
    pub lock_file: PathBuf,
    pub ready_file: PathBuf,
    pub engine_log_file: PathBuf,
    pub client_log_file: PathBuf,
}

impl EnginePaths {
    pub fn new() -> Result<Self> {
        let cache_dir = dirs::cache_dir()
            .ok_or_else(|| Error::Internal("Cannot determine cache directory".into()))?
            .join("com.theproxycompany");

        Ok(Self {
            pid_file: cache_dir.join("engine.pid"),
            lock_file: cache_dir.join("engine.lock"),
            ready_file: cache_dir.join("engine.ready"),
            engine_log_file: cache_dir.join("engine.log"),
            client_log_file: cache_dir.join("client.log"),
            cache_dir,
        })
    }
}

/// Global context shared across all InferenceEngine instances in a process.
struct GlobalContext {
    ref_count: AtomicU32,
    initialized: AtomicBool,
    pid_file: Mutex<Option<PathBuf>>,
}

static GLOBAL_CONTEXT: GlobalContext = GlobalContext {
    ref_count: AtomicU32::new(0),
    initialized: AtomicBool::new(false),
    pid_file: Mutex::new(None),
};

pub(crate) fn current_engine_pid_file() -> Option<PathBuf> {
    GLOBAL_CONTEXT
        .pid_file
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .clone()
}

fn set_current_engine_pid_file(pid_file: Option<PathBuf>) {
    *GLOBAL_CONTEXT
        .pid_file
        .lock()
        .unwrap_or_else(|e| e.into_inner()) = pid_file;
}

/// Manages the PIE (Proxy Inference Engine) process lifecycle.
///
/// Handles:
/// - Binary fetching if not installed
/// - Process spawning with proper daemonization
/// - Per-process lease tracking for orchard-rs clients
/// - Registration with PIE's management plane
pub struct InferenceEngine {
    paths: EnginePaths,
    fetcher: EngineFetcher,
    startup_timeout: Duration,
    lease_active: bool,
    closed: bool,
    launch_process: Option<Child>,
}

impl InferenceEngine {
    /// Create a new InferenceEngine and connect to (or spawn) the engine process.
    pub async fn new() -> Result<Self> {
        Self::with_options(EnginePaths::new()?, None).await
    }

    /// Create with custom options.
    pub async fn with_options(
        paths: EnginePaths,
        startup_timeout: Option<Duration>,
    ) -> Result<Self> {
        let fetcher = EngineFetcher::new();
        let startup_timeout =
            startup_timeout.unwrap_or(Duration::from_secs(DEFAULT_STARTUP_TIMEOUT_SECS));

        let mut engine = Self {
            paths,
            fetcher,
            startup_timeout,
            lease_active: false,
            closed: false,
            launch_process: None,
        };

        engine.acquire_lease().await?;

        Ok(engine)
    }

    /// Close this engine instance.
    ///
    /// Releases this process's orchard-rs lease and best-effort deregisters
    /// the client with PIE. Normal disconnect never sends signals to PIE.
    pub fn close(&mut self) -> Result<()> {
        if self.closed {
            return Ok(());
        }

        let should_release = if self.lease_active {
            let prev = GLOBAL_CONTEXT.ref_count.fetch_sub(1, Ordering::SeqCst);
            prev == 1 // Was the last reference
        } else {
            false
        };

        if !self.lease_active || !should_release {
            self.closed = true;
            self.lease_active = false;
            return Ok(());
        }

        // Acquire lock while checking the shared engine PID file so we don't race
        // a concurrent startup or explicit shutdown in another process.
        let lock_file = std::fs::File::create(&self.paths.lock_file)?;
        lock_exclusive_with_timeout(&lock_file)?;

        let engine_pid = read_pid_file(&self.paths.pid_file);
        let engine_running = engine_pid.map(pid_is_alive).unwrap_or(false);

        if engine_running {
            if let Err(e) =
                self.send_client_lifecycle_command("client_deregister", Duration::from_secs(5))
            {
                tracing::warn!(
                    "Failed to deregister orchard-rs client PID {} from PIE: {}",
                    std::process::id(),
                    e
                );
            }
        } else {
            remove_if_exists(&self.paths.pid_file);
            remove_if_exists(&self.paths.ready_file);
            remove_if_exists(&self.paths.cache_dir.join("engine.refs"));
        }

        drop(lock_file);

        self.lease_active = false;
        self.closed = true;
        GLOBAL_CONTEXT.initialized.store(false, Ordering::SeqCst);
        set_current_engine_pid_file(None);

        Ok(())
    }

    /// Force shutdown the engine regardless of reference count.
    pub fn shutdown(timeout: Duration) -> Result<()> {
        let paths = EnginePaths::new()?;

        let lock_file = std::fs::File::create(&paths.lock_file)?;
        lock_exclusive_with_timeout(&lock_file)?;

        // Helper to clean up all state and socket files
        let cleanup_all = |paths: &EnginePaths| {
            remove_if_exists(&paths.pid_file);
            remove_if_exists(&paths.ready_file);
            remove_if_exists(&paths.cache_dir.join("engine.refs"));
            let ipc_dir = paths.cache_dir.join("ipc");
            if ipc_dir.exists() {
                remove_if_exists(&ipc_dir.join("pie_requests.ipc"));
                remove_if_exists(&ipc_dir.join("pie_responses.ipc"));
                remove_if_exists(&ipc_dir.join("pie_management.ipc"));
            }
        };

        let pid = match read_pid_file(&paths.pid_file) {
            Some(p) if pid_is_alive(p) => p,
            _ => {
                tracing::info!("Engine is not running. Cleaning up stale files.");
                cleanup_all(&paths);
                GLOBAL_CONTEXT.initialized.store(false, Ordering::SeqCst);
                set_current_engine_pid_file(None);
                return Ok(());
            }
        };

        if !pid_is_engine(pid) {
            tracing::warn!(
                "PID {} does not belong to proxy_inference_engine; cleaning stale files.",
                pid
            );
            cleanup_all(&paths);
            GLOBAL_CONTEXT.initialized.store(false, Ordering::SeqCst);
            set_current_engine_pid_file(None);
            return Ok(());
        }
        tracing::info!("Sending shutdown signal to engine process {}", pid);

        if stop_engine_process(pid, timeout) {
            cleanup_all(&paths);
            reap_engine_process(pid);
            GLOBAL_CONTEXT.initialized.store(false, Ordering::SeqCst);
            set_current_engine_pid_file(None);
            tracing::info!("Engine process {} terminated gracefully", pid);
            Ok(())
        } else {
            Err(Error::ShutdownFailed(format!(
                "Failed to stop engine process {}",
                pid
            )))
        }
    }

    async fn acquire_lease(&mut self) -> Result<()> {
        if self.closed || self.lease_active {
            return Ok(());
        }

        // Ensure cache directory exists
        std::fs::create_dir_all(&self.paths.cache_dir)?;

        // Acquire file lock
        let lock_file = std::fs::File::create(&self.paths.lock_file)?;
        lock_exclusive_with_timeout(&lock_file)?;

        let engine_pid = read_pid_file(&self.paths.pid_file);
        let engine_running = engine_pid
            .map(|pid| pid_is_alive(pid) && pid_is_engine(pid))
            .unwrap_or(false);
        let mut launched_engine = false;

        // Launch engine if needed
        if !engine_running {
            tracing::debug!("Inference engine not running. Launching new instance.");

            // Clean up stale state files
            remove_if_exists(&self.paths.pid_file);
            remove_if_exists(&self.paths.ready_file);
            remove_if_exists(&self.paths.cache_dir.join("engine.refs"));

            // Clean up stale IPC socket files (left over from crashed engine)
            let ipc_dir = self.paths.cache_dir.join("ipc");
            if ipc_dir.exists() {
                remove_if_exists(&ipc_dir.join("pie_requests.ipc"));
                remove_if_exists(&ipc_dir.join("pie_responses.ipc"));
                remove_if_exists(&ipc_dir.join("pie_management.ipc"));
            }

            self.launch_engine().await?;
            self.wait_for_engine_ready().await?;
            launched_engine = true;
        }

        if let Err(e) =
            self.send_client_lifecycle_command("client_register", Duration::from_secs(5))
        {
            if launched_engine {
                if let Some(pid) = read_pid_file(&self.paths.pid_file) {
                    if let Err(stop_err) = self.stop_engine_locked(pid) {
                        tracing::warn!(
                            "Failed to clean up newly launched engine {} after register failure: {}",
                            pid,
                            stop_err
                        );
                    }
                }
            }
            return Err(e);
        }

        // Update global context
        GLOBAL_CONTEXT.ref_count.fetch_add(1, Ordering::SeqCst);
        GLOBAL_CONTEXT.initialized.store(true, Ordering::SeqCst);
        set_current_engine_pid_file(Some(self.paths.pid_file.clone()));

        drop(lock_file);

        self.lease_active = true;
        Ok(())
    }

    fn send_client_lifecycle_command(&self, command_type: &str, timeout: Duration) -> Result<()> {
        let socket = Socket::new(Protocol::Req0)?;
        socket.set_opt::<nng::options::RecvTimeout>(Some(timeout))?;
        socket.set_opt::<nng::options::SendTimeout>(Some(timeout))?;
        socket.dial(&management_url())?;

        let payload = json!({
            "type": command_type,
            "client_pid": std::process::id(),
        });
        let data = serde_json::to_vec(&payload)?;
        let msg = nng::Message::from(data.as_slice());
        socket.send(msg).map_err(|(_, e)| Error::Nng(e))?;

        let response = socket.recv()?;
        let json: serde_json::Value = serde_json::from_slice(&response)?;
        let status = json.get("status").and_then(|value| value.as_str());
        if matches!(status, Some("ok") | Some("accepted")) {
            return Ok(());
        }

        let message = json
            .get("message")
            .and_then(|value| value.as_str())
            .unwrap_or("unknown error");
        Err(Error::Other(format!(
            "Engine rejected {} for PID {}: {}",
            command_type,
            std::process::id(),
            message
        )))
    }

    async fn launch_engine(&mut self) -> Result<()> {
        let engine_path = self.fetcher.get_engine_path().await?;

        tracing::info!("Launching PIE from {:?}", engine_path);

        // Open log file for engine output
        let log_file = std::fs::File::create(&self.paths.engine_log_file)?;

        let child = Command::new(&engine_path)
            .stdout(Stdio::from(log_file.try_clone()?))
            .stderr(Stdio::from(log_file))
            .spawn()
            .map_err(|e| Error::StartupFailed(format!("Failed to spawn engine: {}", e)))?;

        self.launch_process = Some(child);
        Ok(())
    }

    async fn wait_for_engine_ready(&self) -> Result<()> {
        tracing::info!("Waiting for telemetry heartbeat from engine...");

        // Subscribe to telemetry topic via NNG
        let telemetry_topic = [EVENT_TOPIC_PREFIX, b"telemetry"].concat();
        let response_url = response_url();

        let socket = nng::Socket::new(nng::Protocol::Sub0)
            .map_err(|e| Error::StartupFailed(format!("Failed to create Sub0 socket: {}", e)))?;

        socket
            .set_opt::<nng::options::protocol::pubsub::Subscribe>(telemetry_topic.clone())
            .map_err(|e| Error::StartupFailed(format!("Failed to subscribe: {}", e)))?;

        socket
            .set_opt::<nng::options::RecvTimeout>(Some(Duration::from_millis(250)))
            .map_err(|e| Error::StartupFailed(format!("Failed to set timeout: {}", e)))?;

        // Non-blocking dial - nng will reconnect automatically in background
        socket
            .dial_async(&response_url)
            .map_err(|e| Error::StartupFailed(format!("Failed to dial {}: {}", response_url, e)))?;

        let deadline = Instant::now() + self.startup_timeout;

        while Instant::now() < deadline {
            // Check if launched process died
            if let Some(ref child) = self.launch_process {
                if !pid_is_alive(child.id()) {
                    return Err(Error::StartupFailed(
                        "Engine process exited before signaling readiness; check the engine log"
                            .into(),
                    ));
                }
            }

            // Try to receive a message
            let msg = match socket.recv() {
                Ok(msg) => msg,
                Err(nng::Error::TimedOut) => continue,
                Err(e) => {
                    tracing::debug!("Error receiving telemetry: {}", e);
                    continue;
                }
            };

            // Parse message: topic\x00json_body
            let bytes = msg.as_slice();
            let parts: Vec<&[u8]> = bytes.splitn(2, |&b| b == 0).collect();
            if parts.len() < 2 {
                tracing::warn!("Discarding malformed event message while waiting for telemetry");
                continue;
            }

            let (topic_part, json_body) = (parts[0], parts[1]);
            if topic_part != telemetry_topic.as_slice() {
                tracing::debug!(
                    "Ignoring unexpected startup topic '{}'",
                    String::from_utf8_lossy(topic_part)
                );
                continue;
            }

            // Parse JSON payload
            let payload: serde_json::Value = match serde_json::from_slice(json_body) {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!("Discarding malformed telemetry payload: {}", e);
                    continue;
                }
            };

            // Extract health.pid
            let engine_pid = payload
                .get("health")
                .and_then(|h| h.get("pid"))
                .and_then(|p| p.as_u64())
                .map(|p| p as u32);

            match engine_pid {
                Some(pid) if pid > 0 => {
                    if let Err(e) = write_pid_file(&self.paths.pid_file, pid) {
                        tracing::warn!("Failed to write PID file: {}", e);
                    }
                    tracing::info!("Received telemetry heartbeat. Engine PID {} recorded.", pid);
                    return Ok(());
                }
                _ => {
                    tracing::warn!(
                        "Telemetry payload missing valid PID; waiting for next heartbeat"
                    );
                    continue;
                }
            }
        }

        Err(Error::StartupFailed(format!(
            "Timed out after {:?}s waiting for telemetry heartbeat from engine",
            self.startup_timeout.as_secs()
        )))
    }

    fn stop_engine_locked(&mut self, pid: u32) -> Result<()> {
        let cleanup_files = || {
            remove_if_exists(&self.paths.pid_file);
            remove_if_exists(&self.paths.ready_file);
            let ipc_dir = self.paths.cache_dir.join("ipc");
            if ipc_dir.exists() {
                remove_if_exists(&ipc_dir.join("pie_requests.ipc"));
                remove_if_exists(&ipc_dir.join("pie_responses.ipc"));
                remove_if_exists(&ipc_dir.join("pie_management.ipc"));
            }
        };

        if !pid_is_alive(pid) {
            tracing::debug!("Engine PID {} already exited", pid);
            cleanup_files();
            return Ok(());
        }

        if !pid_is_engine(pid) {
            tracing::warn!(
                "PID {} does not belong to proxy_inference_engine; cleaning stale files.",
                pid
            );
            cleanup_files();
            return Ok(());
        }

        if !stop_engine_process(pid, Duration::from_secs(5)) {
            return Err(Error::ShutdownFailed(format!(
                "Failed to stop engine PID {}",
                pid
            )));
        }

        reap_engine_process(pid);
        cleanup_files();

        tracing::info!("Engine PID {} stopped", pid);
        Ok(())
    }

    /// Generate a unique response channel ID for this client.
    ///
    /// Format: (PID << 32) | random_32_bits
    /// Uses true randomness to avoid collisions between rapid successive calls.
    pub fn generate_response_channel_id() -> u64 {
        use rand::Rng;

        let pid = std::process::id() as u64 & 0xFFFFFFFF;
        let random: u32 = rand::thread_rng().gen();

        let channel_id = (pid << 32) | (random as u64);
        if channel_id == 0 {
            1
        } else {
            channel_id
        }
    }
}

impl Drop for InferenceEngine {
    fn drop(&mut self) {
        if !self.closed {
            if let Err(e) = self.close() {
                tracing::error!("Failed to close InferenceEngine: {}", e);
            }
        }
    }
}

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

    #[test]
    fn test_engine_paths() {
        let paths = EnginePaths::new().expect("cache dir should be available");
        assert!(paths
            .cache_dir
            .to_string_lossy()
            .contains("com.theproxycompany"));
    }

    #[test]
    fn test_generate_channel_id_uniqueness() {
        use std::collections::HashSet;

        // Generate 1000 channel IDs in rapid succession
        let ids: HashSet<u64> = (0..1000)
            .map(|_| InferenceEngine::generate_response_channel_id())
            .collect();

        // All IDs must be unique (HashSet dedupes)
        assert_eq!(
            ids.len(),
            1000,
            "Channel IDs must be unique across rapid calls"
        );

        // All IDs must be non-zero
        assert!(!ids.contains(&0), "Channel ID must never be zero");

        // All IDs should have the current PID in upper 32 bits
        let expected_pid = std::process::id() as u64 & 0xFFFFFFFF;
        for id in &ids {
            let id_pid = id >> 32;
            assert_eq!(id_pid, expected_pid, "Upper 32 bits must be current PID");
        }
    }
}