cp2k-rs 0.2.3

//! Pure-Rust worker management: spawning, IPC, and lifecycle.
//!
//! All public items in this module are GIL-free. When calling from Python,
//! the caller should release the GIL (e.g., with `py.detach(...)`) around
//! any blocking calls (`start_worker`, `stop_worker`, `ipc_call`).

use std::collections::HashMap;
use std::io::{Read, Write};
use std::os::unix::net::UnixStream;
use std::process::{Child, Command as StdCommand, Stdio};
use std::sync::Mutex;
use std::time::Duration;

use thiserror::Error;

use crate::worker_protocol::{Command, Payload, Request, Response, Status};

// ─── error type ──────────────────────────────────────────────────────────────

/// Errors that can occur during worker operations.
#[derive(Debug, Error)]
pub enum WorkerError {
    #[error("worker mutex poisoned")]
    MutexPoisoned,
    #[error("CP2K worker is not running; call start_worker() first")]
    NotRunning,
    #[error("a CP2K worker is already running; call stop_worker() first")]
    AlreadyRunning,
    #[error(
        "cp2k_rs_worker binary not found; \
             set CP2K_WORKER_BIN or ensure the binary is on PATH"
    )]
    BinaryNotFound,
    #[error("IPC I/O error: {0}")]
    Io(#[from] std::io::Error),
    #[error("serialization error: {0}")]
    Serialize(String),
    #[error("CP2K error: {0}")]
    Cp2kError(String),
    #[error("{0}")]
    Other(String),
}

// ─── global worker state ─────────────────────────────────────────────────────

struct WorkerState {
    child: Child,
    stream: UnixStream,
    next_id: u64,
    socket_path: String,
    socket_file_path: String,
    ready_path: String,
}

static WORKER: Mutex<Option<WorkerState>> = Mutex::new(None);

// ─── IPC helpers (private) ───────────────────────────────────────────────────

fn read_msg(stream: &mut UnixStream) -> std::io::Result<Vec<u8>> {
    let mut len_buf = [0u8; 4];
    stream.read_exact(&mut len_buf)?;
    let len = u32::from_le_bytes(len_buf) as usize;
    let mut buf = vec![0u8; len];
    stream.read_exact(&mut buf)?;
    Ok(buf)
}

fn write_msg(stream: &mut UnixStream, data: &[u8]) -> std::io::Result<()> {
    let len = data.len() as u32;
    stream.write_all(&len.to_le_bytes())?;
    stream.write_all(data)?;
    stream.flush()
}

// ─── public IPC entry point ───────────────────────────────────────────────────

/// Send a command and receive the response.
///
/// Blocks until the response arrives. GIL-free: safe to call inside
/// `py.detach(...)`.
pub fn ipc_call(command: Command) -> Result<Payload, WorkerError> {
    let mut guard = WORKER.lock().map_err(|_| WorkerError::MutexPoisoned)?;
    let state = guard.as_mut().ok_or(WorkerError::NotRunning)?;

    let req = Request {
        request_id: state.next_id,
        command,
    };
    state.next_id += 1;

    let bytes = bincode::serialize(&req).map_err(|e| WorkerError::Serialize(e.to_string()))?;
    write_msg(&mut state.stream, &bytes)?;

    let raw = read_msg(&mut state.stream).map_err(|e| {
        if e.kind() == std::io::ErrorKind::UnexpectedEof {
            WorkerError::Other(
                "CP2K worker process died unexpectedly during a request \
                 (connection closed mid-read). Check the worker's stderr output \
                 for CP2K error messages."
                    .into(),
            )
        } else {
            WorkerError::Io(e)
        }
    })?;
    let resp: Response =
        bincode::deserialize(&raw).map_err(|e| WorkerError::Serialize(e.to_string()))?;

    match resp.status {
        Status::Ok => Ok(resp.payload),
        Status::Error(msg) => Err(WorkerError::Cp2kError(msg)),
    }
}

// ─── binary discovery ─────────────────────────────────────────────────────────

/// Find the `cp2k_rs_worker` binary.
///
/// Search order:
///   1. `CP2K_WORKER_BIN` environment variable
///   2. Same directory as the current executable
///   3. `PATH`
///
/// Python callers may additionally check the installed `cp2k_rs` Python
/// package directory and pass the found path directly to [`start_worker`].
pub fn find_worker_binary() -> Option<std::path::PathBuf> {
    // 1. Explicit override
    if let Ok(p) = std::env::var("CP2K_WORKER_BIN") {
        let path = std::path::PathBuf::from(p);
        if path.exists() {
            return Some(path);
        }
    }

    // 2. Next to the calling executable
    if let Ok(exe) = std::env::current_exe()
        && let Some(dir) = exe.parent()
    {
        let candidate = dir.join("cp2k_rs_worker");
        if candidate.exists() {
            return Some(candidate);
        }
    }

    // 3. PATH lookup
    if let Ok(path_var) = std::env::var("PATH") {
        for dir in std::env::split_paths(&path_var) {
            let candidate = dir.join("cp2k_rs_worker");
            if candidate.exists() {
                return Some(candidate);
            }
        }
    }

    None
}

fn lookup_env(var: &str, extra_env: Option<&HashMap<String, String>>) -> Option<String> {
    extra_env
        .and_then(|env| env.get(var).cloned())
        .or_else(|| std::env::var(var).ok())
}

/// Detect the default MPI launcher based on the current environment.
///
/// When `nproc` is `None` the launcher is returned without a `-n` flag,
/// letting the scheduler or MPI implementation decide the process count.
///
/// For `mpirun` (OpenMPI), critical environment variables are forwarded
/// explicitly via `-x` flags.  OpenMPI does not guarantee that every
/// variable present in the parent shell is propagated to remote ranks; in
/// particular the `OMPI_MCA_*` transport-selection variables must be
/// forwarded explicitly, otherwise UCX may be selected and crash with
/// SIGILL on CI runners that lack the required hardware (InfiniBand, AVX-512
/// instructions used by PSM3, …).
///
/// `extra_env` can be used to forward variables that are set for the worker
/// process but not exported in the parent environment.
pub fn default_launcher(
    nproc: Option<u32>,
    extra_env: Option<&HashMap<String, String>>,
) -> Vec<String> {
    if std::env::var("SLURM_JOB_ID").is_ok() {
        match nproc {
            Some(n) => vec!["srun".to_string(), "-n".to_string(), n.to_string()],
            None => vec!["srun".to_string()],
        }
    } else {
        let mut v = vec!["mpirun".to_string()];
        if let Some(n) = nproc {
            v.push("-n".to_string());
            v.push(n.to_string());
        }

        // Forward environment variables that OpenMPI does not always propagate
        // automatically.  Using `-x VAR=val` hard-codes the value; using `-x VAR`
        // (no `=`) forwards whatever the parent exported.
        //
        // ── Transport / UCX suppression ───────────────────────────────────────
        // On AlmaLinux/RHEL 8 the system OpenMPI is built with UCX support.
        // UCX probes hardware at initialisation time using CPU instructions
        // (e.g. CPUID-gated AVX-512 paths in PSM3/UD verbs) that are absent on
        // plain VMs and CI runners, causing SIGILL inside libucs.so.  This
        // happens even when pml=ob1 is set because UCX is *also* loaded as the
        // one-sided communication (OSC) component, independently of the PML.
        //
        // Safe-default injection policy:
        //   • If the user has already set *any* OMPI_MCA_* variable we assume
        //     they know their MPI environment and we forward only what they set —
        //     no defaults are injected.  This preserves full performance on real
        //     HPC clusters with InfiniBand / RDMA hardware.
        //   • If no OMPI_MCA_* variable is set AND we cannot detect any
        //     InfiniBand hardware (no /sys/class/infiniband entries) we inject
        //     the conservative TCP-only defaults that prevent the UCX SIGILL.
        //   • UCX_ERROR_SIGNALS="" is always forwarded if the parent set it;
        //     it is injected as a default only together with the other UCX vars.
        //   • OMPI_ALLOW_RUN_AS_ROOT is forwarded if the parent set it but is
        //     never injected as a default — it is a security-relevant setting
        //     that the user should opt into explicitly.

        // Always-forward variables (from parent env or extra_env overrides).
        let passthrough_vars: &[&str] = &[
            "OMPI_ALLOW_RUN_AS_ROOT",
            "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM",
            "CP2K_DATA_DIR",
            "OMP_NUM_THREADS",
            "OMP_WAIT_POLICY",
            "RUST_BACKTRACE",
        ];
        for var in passthrough_vars {
            if let Some(val) = lookup_env(var, extra_env) {
                v.push("-x".to_string());
                v.push(format!("{var}={val}"));
            }
        }

        // The MCA transport variables require more careful handling.
        let mca_vars: &[&str] = &[
            "OMPI_MCA_pml",
            "OMPI_MCA_btl",
            "OMPI_MCA_osc",
            "OMPI_MCA_mtl",
            "UCX_ERROR_SIGNALS",
        ];

        // Check whether the parent already has an opinion about MCA transport.
        let user_set_mca = mca_vars.iter().any(|v| lookup_env(v, extra_env).is_some());

        if user_set_mca {
            // User knows their environment: forward exactly what they set,
            // inject nothing extra.
            for var in mca_vars {
                if let Some(val) = lookup_env(var, extra_env) {
                    v.push("-x".to_string());
                    v.push(format!("{var}={val}"));
                }
            }
        } else if !has_infiniband() {
            // No user configuration and no IB hardware detected: inject the
            // conservative defaults that prevent UCX from crashing.
            let safe_defaults: &[(&str, &str)] = &[
                ("OMPI_MCA_pml", "ob1"),
                ("OMPI_MCA_btl", "tcp,self"),
                ("OMPI_MCA_osc", "pt2pt"),
                ("OMPI_MCA_mtl", "^ofi,psm,psm2"),
                ("UCX_ERROR_SIGNALS", ""),
            ];
            for (var, default) in safe_defaults {
                v.push("-x".to_string());
                v.push(format!("{var}={default}"));
            }
        }
        // else: IB hardware present and no explicit MCA config — let OpenMPI
        // auto-detect; don't constrain the transport choice.

        v
    }
}

/// Return `true` if at least one InfiniBand port is visible in sysfs.
///
/// Presence of `/sys/class/infiniband` with at least one entry is a reliable
/// indicator that UCX/RDMA transports are meaningful on this host.  The check
/// is intentionally conservative: if sysfs is unavailable or the read fails
/// for any reason we return `false` (i.e. assume no IB, inject safe defaults).
fn has_infiniband() -> bool {
    std::fs::read_dir("/sys/class/infiniband")
        .map(|mut d| d.next().is_some())
        .unwrap_or(false)
}

// ─── worker lifecycle ─────────────────────────────────────────────────────────

/// Start the MPI worker process and wait until its socket is ready.
///
/// * `worker_bin`       – Path to the `cp2k_rs_worker` binary.
/// * `nproc`            – MPI rank count. `None` omits `-n` and lets the scheduler decide. Ignored when `launcher_cmd` is given.
/// * `launcher_cmd`     – Custom launcher prefix, e.g. `["srun", "-n", "8"]`.
/// * `env`              – Extra environment variables for the worker.
/// * `working_dir`      – Working directory for the worker process.
/// * `connect_timeout`  – Seconds to wait for the socket to become ready.
///
/// GIL-free: safe to call inside `py.detach(...)`.
pub fn start_worker(
    worker_bin: std::path::PathBuf,
    nproc: Option<u32>,
    launcher_cmd: Option<Vec<String>>,
    env: Option<HashMap<String, String>>,
    working_dir: Option<String>,
    connect_timeout: f64,
) -> Result<(), WorkerError> {
    {
        let guard = WORKER.lock().map_err(|_| WorkerError::MutexPoisoned)?;
        if guard.is_some() {
            return Err(WorkerError::AlreadyRunning);
        }
    }

    let mut extra_env = env.unwrap_or_default();

    // Default to passive OpenMP wait policy so idle threads do not busy-spin
    // at 100 % CPU.  Only inject if the user hasn't set it explicitly.
    extra_env
        .entry("OMP_WAIT_POLICY".to_string())
        .or_insert_with(|| "passive".to_string());

    if let Some(parent) = worker_bin.parent() {
        let candidate = parent.join("data");
        if candidate.exists() {
            extra_env.insert(
                "CP2K_DATA_DIR".to_string(),
                candidate.to_string_lossy().into_owned(),
            );
        }
    }

    let mut cmd_parts = launcher_cmd.unwrap_or_else(|| default_launcher(nproc, Some(&extra_env)));
    cmd_parts.push(worker_bin.to_string_lossy().into_owned());

    let socket_path = format!("/tmp/cp2k_worker_{}_{}.sock", std::process::id(), {
        use std::time::{SystemTime, UNIX_EPOCH};
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .subsec_nanos()
    });
    let socket_file_path = format!("{socket_path}.path");
    let ready_path = format!("{socket_path}.ready");

    let _ = std::fs::remove_file(&socket_path);
    let _ = std::fs::remove_file(&socket_file_path);
    let _ = std::fs::remove_file(&ready_path);

    if let Err(e) = std::fs::write(&socket_file_path, &socket_path) {
        return Err(WorkerError::Other(format!(
            "Failed to write CP2K socket path file '{socket_file_path}': {e}"
        )));
    }

    let mut cmd = StdCommand::new(&cmd_parts[0]);
    cmd.args(&cmd_parts[1..]);
    cmd.env("CP2K_WORKER_SOCKET_FILE", &socket_file_path);
    cmd.env("CP2K_WORKER_READY_FILE", &ready_path);
    cmd.stdout(Stdio::inherit());
    cmd.stderr(Stdio::inherit());

    if !extra_env.is_empty() {
        for (k, v) in extra_env {
            cmd.env(k, v);
        }
    }
    if let Some(dir) = working_dir {
        cmd.current_dir(dir);
    }

    let mut child = cmd
        .spawn()
        .map_err(|e| WorkerError::Other(format!("Failed to spawn cp2k_rs_worker: {e}")))?;

    let timeout = Duration::from_secs_f64(connect_timeout);
    let start = std::time::Instant::now();
    loop {
        if std::path::Path::new(&ready_path).exists() {
            break;
        }
        if start.elapsed() > timeout {
            let default_ready = "/tmp/cp2k_rs_worker.ready";
            let default_ready_exists = std::path::Path::new(default_ready).exists();
            let child_status = match child.try_wait() {
                Ok(Some(status)) => format!("worker exited with status {status}"),
                Ok(None) => "worker still running".to_string(),
                Err(e) => format!("worker status check failed: {e}"),
            };
            let extra = if default_ready_exists {
                format!(
                    " Default ready file exists at {default_ready} (possible handshake mismatch)."
                )
            } else {
                String::new()
            };
            return Err(WorkerError::Other(format!(
                "Timed out waiting for cp2k_rs_worker to become ready ({connect_timeout}s). \
                 The worker process may have crashed during CP2K initialization. \
                 Check that CP2K_DATA_DIR is set (it should point to the CP2K data/ directory) \
                 and that the MPI launcher is available on PATH. Socket: {socket_path}. {child_status}.{extra}"
            )));
        }
        std::thread::sleep(Duration::from_millis(50));
    }

    let deadline = std::time::Instant::now() + timeout;
    let stream = loop {
        match UnixStream::connect(&socket_path) {
            Ok(s) => break s,
            Err(_) if std::time::Instant::now() < deadline => {
                std::thread::sleep(Duration::from_millis(50));
            }
            Err(e) => return Err(WorkerError::Other(format!("Connect failed: {e}"))),
        }
    };

    stream.set_read_timeout(None)?;

    let mut guard = WORKER.lock().map_err(|_| WorkerError::MutexPoisoned)?;
    *guard = Some(WorkerState {
        child,
        stream,
        next_id: 0,
        socket_path,
        socket_file_path,
        ready_path,
    });

    Ok(())
}

// ─── shared-memory helpers ──────────────────────────────────────────────────

/// Read a `SharedArray3` payload from shm into an owned `Vec<f64>`.
///
/// Opens the named shm region, copies the data, then unlinks it.
/// Returns `(dims, data)`.
pub fn read_shared_array3(
    shm_name: &str,
    _dims: [usize; 3],
    byte_size: usize,
) -> Result<Vec<f64>, WorkerError> {
    crate::shm::read_shared_array3(shm_name, _dims, byte_size)
        .map_err(|e| WorkerError::Other(format!("shm open failed: {e}")))
}

/// Read a `SharedArray2` payload from shm into an owned `Vec<f64>`.
///
/// Opens the named shm region, copies the data, then unlinks it.
pub fn read_shared_array2(shm_name: &str, byte_size: usize) -> Result<Vec<f64>, WorkerError> {
    crate::shm::read_shared_array2(shm_name, byte_size)
        .map_err(|e| WorkerError::Other(format!("shm open failed: {e}")))
}

/// Shut down the worker process gracefully.
///
/// GIL-free: safe to call inside `py.detach(...)`.
pub fn stop_worker() -> Result<(), WorkerError> {
    // Best-effort shutdown: ignore errors (worker may already be gone).
    let _ = ipc_call(Command::Shutdown);

    let mut guard = WORKER.lock().map_err(|_| WorkerError::MutexPoisoned)?;
    if let Some(mut state) = guard.take() {
        let grace = Duration::from_secs(10);
        let start = std::time::Instant::now();
        loop {
            match state.child.try_wait() {
                Ok(Some(_)) => break,
                Ok(None) if start.elapsed() < grace => {
                    std::thread::sleep(Duration::from_millis(100));
                }
                _ => {
                    let _ = state.child.kill();
                    break;
                }
            }
        }
        let _ = std::fs::remove_file(&state.socket_path);
        let _ = std::fs::remove_file(&state.socket_file_path);
        let _ = std::fs::remove_file(&state.ready_path);
    }

    Ok(())
}