xy/crates/xy/src/daemon/handlers.rs

use crate::daemon::registry::Registry;
use crate::paths::Paths;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::Mutex;
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use xy_ipc::Connection;
use xy_ipc::envelope::{Incoming, Request, Response, err_response, ok_response};
use xy_protocol::RpcErrorCode;
use xy_protocol::rpc::{
    LogEnd, LogLine, LogsCancelParams, LogsParams, LogsSubscribed, NameOrAll, RestartResult,
    ServerSummary, StartResult, StatusDetail, StopResult, WaitInfo, methods, notifications,
};
use xy_supervisor::supervisor::{StartAck, StopAck, SupervisorCmd};

pub struct ConnState {
    pub subs: Mutex<HashMap<u64, JoinHandle<()>>>,
    pub next: AtomicU64,
}

impl ConnState {
    pub fn new() -> Self {
        Self {
            subs: Mutex::new(HashMap::new()),
            next: AtomicU64::new(1),
        }
    }
}

pub async fn serve(conn: Arc<Connection>, reg: Registry, _paths: Paths) -> std::io::Result<()> {
    let state = Arc::new(ConnState::new());

    loop {
        let Some(incoming) = conn.read_incoming().await? else {
            let mut subs = state.subs.lock().await;

            for (_, h) in subs.drain() {
                h.abort();
            }

            return Ok(());
        };

        if let Incoming::Request(req) = incoming {
            let (resp, log_ready) = handle_request(req, &reg, &conn, &state).await;

            conn.write_response(&resp).await?;

            if let Some(tx) = log_ready {
                let _ = tx.send(());
            }
        }
    }
}

struct ApiError {
    code: i32,
    message: String,
}

impl ApiError {
    fn rpc(code: RpcErrorCode, msg: impl Into<String>) -> Self {
        Self {
            code: code.as_i32(),
            message: msg.into(),
        }
    }
}

type LogReadyTx = tokio::sync::oneshot::Sender<()>;

async fn handle_request(
    req: Request,
    reg: &Registry,
    conn: &Arc<Connection>,
    state: &Arc<ConnState>,
) -> (Response, Option<LogReadyTx>) {
    let id = req.id.clone();
    let method = req.method.as_str();
    let params = req.params.unwrap_or(serde_json::Value::Null);

    let resp = match method {
        methods::LIST => match list(reg).await {
            Ok(v) => ok_response(id, serde_json::to_value(v).unwrap()),
            Err(err) => err_response(id, err.code, err.message),
        },
        methods::STATUS => {
            let p: xy_protocol::rpc::StatusParams = match serde_json::from_value(params) {
                Ok(p) => p,
                Err(err) => {
                    return (
                        err_response(id, -32602, format!("invalid params: {err}")),
                        None,
                    );
                }
            };

            match status(reg, &p.name).await {
                Ok(v) => ok_response(id, serde_json::to_value(v).unwrap()),
                Err(err) => err_response(id, err.code, err.message),
            }
        }
        methods::START => dispatch_lifecycle(id, params, reg, Op::Start).await,
        methods::STOP => dispatch_lifecycle(id, params, reg, Op::Stop).await,
        methods::RESTART => dispatch_lifecycle(id, params, reg, Op::Restart).await,
        methods::RELOAD => match reload(reg).await {
            Ok(v) => ok_response(id, serde_json::to_value(v).unwrap()),
            Err(e) => err_response(id, e.code, e.message),
        },
        methods::LOGS => {
            let p: LogsParams = match serde_json::from_value(params) {
                Ok(p) => p,
                Err(err) => {
                    return (
                        err_response(id, -32602, format!("invalid params: {err}")),
                        None,
                    );
                }
            };

            match start_log_stream(reg, conn.clone(), state.clone(), p).await {
                Ok((sub_id, ready_tx)) => {
                    let resp = ok_response(
                        id,
                        serde_json::to_value(LogsSubscribed {
                            subscription_id: sub_id,
                        })
                        .unwrap(),
                    );
                    return (resp, Some(ready_tx));
                }
                Err(err) => err_response(id, err.code, err.message),
            }
        }
        methods::LOGS_CANCEL => {
            let p: LogsCancelParams = match serde_json::from_value(params) {
                Ok(p) => p,
                Err(err) => {
                    return (
                        err_response(id, -32602, format!("invalid params: {err}")),
                        None,
                    );
                }
            };

            let mut subs = state.subs.lock().await;

            if let Some(h) = subs.remove(&p.subscription_id) {
                h.abort();
            }

            ok_response(id, serde_json::json!({}))
        }
        other => err_response(id, -32601, format!("unknown method `{other}`")),
    };

    (resp, None)
}

async fn list(reg: &Registry) -> Result<Vec<ServerSummary>, ApiError> {
    let mut out = Vec::new();

    for (name, entry) in reg.snapshot().await {
        let s = entry.handle.status.borrow();

        out.push(ServerSummary {
            name,
            state: s.state,
            pid: s.pid,
            port: s.port,
            uptime_secs: s.started_at.map(|t| t.elapsed().as_secs()),
            restart_count: s.restart_count,
            last_exit: s.last_exit,
        });
    }

    Ok(out)
}

async fn status(reg: &Registry, name: &str) -> Result<StatusDetail, ApiError> {
    let Some(entry) = reg.get(name).await else {
        return Err(ApiError::rpc(
            RpcErrorCode::ServerNotFound,
            format!("no such server `{name}`"),
        ));
    };

    let s = entry.handle.status.borrow();

    Ok(StatusDetail {
        summary: ServerSummary {
            name: entry.handle.name.clone(),
            state: s.state,
            pid: s.pid,
            port: s.port,
            uptime_secs: s.started_at.map(|t| t.elapsed().as_secs()),
            restart_count: s.restart_count,
            last_exit: s.last_exit,
        },
        recent_transitions: Vec::new(),
        wait_for: s.waiting.as_ref().map(|w| WaitInfo {
            description: w.description.clone(),
            elapsed_secs: w.started_at.elapsed().as_secs(),
            timeout_secs: w.timeout.as_secs(),
        }),
    })
}

enum Op {
    Start,
    Stop,
    Restart,
}

async fn dispatch_lifecycle(
    id: serde_json::Value,
    params: serde_json::Value,
    reg: &Registry,
    op: Op,
) -> Response {
    let p: NameOrAll = match serde_json::from_value(params) {
        Ok(p) => p,
        Err(err) => return err_response(id, -32602, format!("invalid params: {err}")),
    };

    let targets: Vec<String> = match p {
        NameOrAll::All { all } if all => reg.names().await,
        NameOrAll::Name { name } => vec![name],
        NameOrAll::All { .. } => return err_response(id, -32602, "must set all=true".into()),
    };

    match op {
        Op::Start => {
            let mut started = Vec::new();
            let mut already = Vec::new();

            for name in targets {
                let Some(entry) = reg.get(&name).await else {
                    return err_response(
                        id,
                        RpcErrorCode::ServerNotFound.as_i32(),
                        format!("no such server `{name}`"),
                    );
                };

                let (tx, rx) = oneshot::channel();

                let _ = entry.handle.tx.send(SupervisorCmd::Start { ack: tx }).await;

                match rx.await {
                    Ok(StartAck::Started) => started.push(name),
                    Ok(StartAck::AlreadyRunning) => already.push(name),
                    Ok(StartAck::SpawnFailed(msg)) => {
                        return err_response(
                            id,
                            RpcErrorCode::SpawnFailed.as_i32(),
                            format!("failed to start `{name}`: {msg}"),
                        );
                    }
                    Err(_) => {
                        return err_response(
                            id,
                            RpcErrorCode::SpawnFailed.as_i32(),
                            format!("supervisor for `{name}` dropped"),
                        );
                    }
                }
            }

            ok_response(
                id,
                serde_json::to_value(StartResult {
                    started,
                    already_running: already,
                })
                .unwrap(),
            )
        }
        Op::Stop => {
            let mut stopped = Vec::new();
            let mut not_running = Vec::new();

            for name in targets {
                let Some(entry) = reg.get(&name).await else {
                    return err_response(
                        id,
                        RpcErrorCode::ServerNotFound.as_i32(),
                        format!("no such server `{name}`"),
                    );
                };

                let (tx, rx) = oneshot::channel();

                let _ = entry.handle.tx.send(SupervisorCmd::Stop { ack: tx }).await;

                match rx.await {
                    Ok(StopAck::Stopped) => stopped.push(name),
                    Ok(StopAck::NotRunning) => not_running.push(name),
                    Err(_) => {
                        return err_response(
                            id,
                            RpcErrorCode::SpawnFailed.as_i32(),
                            format!("supervisor for `{name}` dropped"),
                        );
                    }
                }
            }

            ok_response(
                id,
                serde_json::to_value(StopResult {
                    stopped,
                    not_running,
                })
                .unwrap(),
            )
        }
        Op::Restart => {
            let mut restarted = Vec::new();

            for name in targets {
                let Some(entry) = reg.get(&name).await else {
                    return err_response(
                        id,
                        RpcErrorCode::ServerNotFound.as_i32(),
                        format!("no such server `{name}`"),
                    );
                };

                let (tx, rx) = oneshot::channel();

                let _ = entry
                    .handle
                    .tx
                    .send(SupervisorCmd::Restart { ack: tx })
                    .await;

                let _ = rx.await;

                restarted.push(name);
            }

            ok_response(
                id,
                serde_json::to_value(RestartResult { restarted }).unwrap(),
            )
        }
    }
}

use xy_protocol::rpc::ReloadResult;

async fn reload(reg: &Registry) -> Result<ReloadResult, ApiError> {
    let paths = crate::daemon::PATHS.get().ok_or_else(|| {
        ApiError::rpc(RpcErrorCode::ConfigInvalid, "daemon paths not initialized")
    })?;

    let new_configs = xy_protocol::kdl_parse::load_all_configs(&paths.config_dir)
        .map_err(|err| ApiError::rpc(RpcErrorCode::ConfigInvalid, err.to_string()))?;

    use std::collections::HashMap;

    let new_by_name: HashMap<String, xy_protocol::ServerConfig> = new_configs
        .into_iter()
        .map(|c| (c.name.clone(), c))
        .collect();

    let existing_names: Vec<String> = reg.names().await;

    let mut added = Vec::new();
    let mut removed = Vec::new();
    let mut changed = Vec::new();
    let mut unchanged = Vec::new();

    for name in &existing_names {
        if !new_by_name.contains_key(name)
            && let Some(entry) = reg.remove(name).await
        {
            let (tx, rx) = oneshot::channel();

            let _ = entry
                .handle
                .tx
                .send(SupervisorCmd::Shutdown { ack: tx })
                .await;

            let _ = rx.await;

            removed.push(name.clone());
        }
    }

    for (name, cfg) in new_by_name {
        let new_hash = crate::daemon::config_hash(&cfg);

        match reg.get(&name).await {
            None => {
                let handle = crate::daemon::spawn_supervisor(paths, cfg)
                    .map_err(|err| ApiError::rpc(RpcErrorCode::SpawnFailed, err.to_string()))?;

                reg.insert(
                    name.clone(),
                    crate::daemon::registry::Entry {
                        handle: handle.clone(),
                        config_hash: new_hash,
                    },
                )
                .await;

                let (tx, rx) = oneshot::channel();

                let _ = handle.tx.send(SupervisorCmd::Start { ack: tx }).await;

                let _ = rx.await;

                added.push(name);
            }
            Some(entry) if entry.config_hash != new_hash => {
                let (tx, rx) = oneshot::channel();

                let _ = entry
                    .handle
                    .tx
                    .send(SupervisorCmd::Shutdown { ack: tx })
                    .await;

                let _ = rx.await;

                reg.remove(&name).await;

                let handle = crate::daemon::spawn_supervisor(paths, cfg)
                    .map_err(|err| ApiError::rpc(RpcErrorCode::SpawnFailed, err.to_string()))?;

                reg.insert(
                    name.clone(),
                    crate::daemon::registry::Entry {
                        handle: handle.clone(),
                        config_hash: new_hash,
                    },
                )
                .await;

                let (tx, rx) = oneshot::channel();

                let _ = handle.tx.send(SupervisorCmd::Start { ack: tx }).await;

                let _ = rx.await;

                changed.push(name);
            }
            Some(_) => unchanged.push(name),
        }
    }

    Ok(ReloadResult {
        added,
        removed,
        changed,
        unchanged,
    })
}

async fn start_log_stream(
    reg: &Registry,
    conn: Arc<Connection>,
    state: Arc<ConnState>,
    p: LogsParams,
) -> Result<(u64, LogReadyTx), ApiError> {
    let Some(entry) = reg.get(&p.name).await else {
        return Err(ApiError::rpc(
            RpcErrorCode::ServerNotFound,
            format!("no such server `{}`", p.name),
        ));
    };

    let sub_id = state.next.fetch_add(1, Ordering::Relaxed);
    let sink = entry.handle.log_sink.clone();
    let conn2 = conn.clone();
    let state2 = state.clone();
    let follow = p.follow;
    let tail = p.tail;
    let name = p.name.clone();

    let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();

    let task = tokio::spawn(async move {
        // Wait until `serve` has written the LOGS response before sending any
        // LOG notifications.  Without this the spawned task can race ahead of
        // `conn.write_response` and emit notifications that the client
        // discards while still awaiting the response.
        let _ = ready_rx.await;

        for line in sink.ring.snapshot_tail(tail) {
            let n = xy_ipc::envelope::notification(
                notifications::LOG,
                Some(
                    serde_json::to_value(LogLine {
                        subscription_id: sub_id,
                        name: name.clone(),
                        stream: line.stream,
                        line: line.line,
                        ts_unix_ms: line.ts_unix_ms,
                    })
                    .unwrap(),
                ),
            );

            if conn2.write_notification(&n).await.is_err() {
                return;
            }
        }

        if !follow {
            let end = xy_ipc::envelope::notification(
                notifications::LOG_END,
                Some(
                    serde_json::to_value(LogEnd {
                        subscription_id: sub_id,
                    })
                    .unwrap(),
                ),
            );

            let _ = conn2.write_notification(&end).await;

            state2.subs.lock().await.remove(&sub_id);

            return;
        }

        let mut rx = sink.broadcast.subscribe();

        while let Ok(mut line) = rx.recv().await {
            line.subscription_id = sub_id;
            let n = xy_ipc::envelope::notification(
                notifications::LOG,
                Some(serde_json::to_value(&line).unwrap()),
            );

            if conn2.write_notification(&n).await.is_err() {
                break;
            }
        }
    });

    state.subs.lock().await.insert(sub_id, task);

    Ok((sub_id, ready_tx))
}