Skip to main content

mesofact_core/proxy/
worker_client.rs

1//! WorkerClient โ€” UDS connection to one Bun render-pool worker.
2//!
3//! Wire protocol: NDJSON-framed, one JSON object per line. `id = 0` is
4//! reserved for lifecycle messages (ready / ping / pong / drain).
5//! See `.yah/docs/architecture/mesofact.md` ยง"IPC protocol".
6//!
7//! Concurrency note: all I/O is serialized through `io: Mutex<WorkerIo>`.
8//! For P7 this is fine โ€” Mode 2 renders are stubbed (501), so the only
9//! callers are the watchdog (ping) and shutdown (drain). Refactor in P9
10//! when concurrent renders share the socket.
11
12use serde_json::{json, Value};
13use std::collections::HashMap;
14use std::path::{Path, PathBuf};
15use std::process::Stdio;
16use std::time::Duration;
17use thiserror::Error;
18use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
19use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
20use tokio::net::UnixStream;
21use tokio::process::{Child, Command};
22use tokio::sync::Mutex;
23use tokio::time::{sleep, timeout};
24
25const PONG_DEADLINE: Duration = Duration::from_secs(5);
26const CONNECT_DEADLINE: Duration = Duration::from_secs(10);
27const READY_DEADLINE: Duration = Duration::from_secs(15);
28
29#[derive(Debug, Error)]
30pub enum WorkerError {
31    #[error("io: {0}")]
32    Io(#[from] std::io::Error),
33    #[error("json: {0}")]
34    Json(#[from] serde_json::Error),
35    #[error("pong not received within 5 s")]
36    PongTimeout,
37    #[error("worker closed the connection")]
38    Closed,
39    #[error("render error '{code}' (retryable={retryable})")]
40    Render { code: String, retryable: bool },
41    #[error("worker queue is full (503)")]
42    QueueOverflow,
43    #[error("worker is draining")]
44    Draining,
45}
46
47#[derive(Debug, Clone)]
48pub struct RenderResult {
49    pub html: String,
50    pub headers: HashMap<String, String>,
51    pub cache_ttl: Option<u64>,
52    pub cache_tags: Vec<String>,
53}
54
55struct WorkerIo {
56    reader: BufReader<OwnedReadHalf>,
57    writer: OwnedWriteHalf,
58}
59
60pub struct WorkerClient {
61    io: Mutex<WorkerIo>,
62    pub socket_path: PathBuf,
63    child: Mutex<Child>,
64}
65
66impl WorkerClient {
67    /// Spawn a Bun worker and connect over UDS. Returns after the worker sends
68    /// its `ready` lifecycle message.
69    pub async fn spawn(
70        socket_path: PathBuf,
71        manifest_path: &Path,
72        worker_entry: &Path,
73        config_path: Option<&Path>,
74    ) -> Result<Self, WorkerError> {
75        let mut cmd = Command::new("bun");
76        cmd.arg(worker_entry)
77            .arg("--socket")
78            .arg(&socket_path)
79            .arg("--manifest")
80            .arg(manifest_path)
81            .stdout(Stdio::null())
82            .stderr(Stdio::inherit())
83            .kill_on_drop(true);
84        if let Some(cfg) = config_path {
85            // Worker registers adapters from `mesofact.config.toml` at boot so a
86            // render's `sqlite('db')` / `r2('assets')` resolves at request time.
87            cmd.arg("--config").arg(cfg);
88        }
89
90        let child = cmd.spawn()?;
91        let stream = connect_with_retry(&socket_path, CONNECT_DEADLINE).await?;
92        let (r, w) = stream.into_split();
93
94        let client = Self {
95            io: Mutex::new(WorkerIo {
96                reader: BufReader::new(r),
97                writer: w,
98            }),
99            socket_path,
100            child: Mutex::new(child),
101        };
102
103        let msg = client.recv_with_timeout(READY_DEADLINE).await?;
104        if msg.get("kind").and_then(|v| v.as_str()) != Some("ready") {
105            return Err(WorkerError::Io(std::io::Error::other(format!(
106                "expected ready, got {msg}"
107            ))));
108        }
109        Ok(client)
110    }
111
112    /// Send a ping and wait up to 5 s for a pong.
113    pub async fn ping(&self) -> Result<(), WorkerError> {
114        let mut io = self.io.lock().await;
115        send(&mut io.writer, &json!({ "id": 0, "kind": "ping" })).await?;
116
117        timeout(PONG_DEADLINE, async {
118            loop {
119                let v = recv(&mut io.reader).await?;
120                if v.get("kind").and_then(|k| k.as_str()) == Some("pong") {
121                    return Ok(());
122                }
123            }
124        })
125        .await
126        .map_err(|_| WorkerError::PongTimeout)?
127    }
128
129    /// Ask the worker to finish in-flight renders and then exit gracefully.
130    pub async fn drain(&self) -> Result<(), WorkerError> {
131        let mut io = self.io.lock().await;
132        send(&mut io.writer, &json!({ "id": 0, "kind": "drain" })).await
133    }
134
135    /// Invoke render on the worker (used by Mode 2 in P9+; stubbed at 501 in P7).
136    pub async fn render(
137        &self,
138        id: u32,
139        route: &str,
140        req: Value,
141        deadline_ms: u64,
142    ) -> Result<RenderResult, WorkerError> {
143        let msg = json!({
144            "id": id,
145            "kind": "render",
146            "route": route,
147            "req": req,
148            "deadline_ms": deadline_ms,
149        });
150        let mut io = self.io.lock().await;
151        send(&mut io.writer, &msg).await?;
152
153        loop {
154            let v = recv(&mut io.reader).await?;
155            if v.get("id").and_then(|i| i.as_u64()) != Some(u64::from(id)) {
156                continue;
157            }
158            return match v.get("kind").and_then(|k| k.as_str()) {
159                Some("ok") => Ok(RenderResult {
160                    html: v["html"].as_str().unwrap_or("").to_string(),
161                    headers: serde_json::from_value(v["headers"].clone()).unwrap_or_default(),
162                    cache_ttl: v["cache"]["ttl"].as_u64(),
163                    cache_tags: serde_json::from_value(v["cache"]["tags"].clone())
164                        .unwrap_or_default(),
165                }),
166                Some("err") => {
167                    let code = v["error"]["code"]
168                        .as_str()
169                        .unwrap_or("render_failed")
170                        .to_string();
171                    let retryable = v["error"]["retryable"].as_bool().unwrap_or(false);
172                    Err(match code.as_str() {
173                        "queue_overflow" => WorkerError::QueueOverflow,
174                        "draining" => WorkerError::Draining,
175                        _ => WorkerError::Render { code, retryable },
176                    })
177                }
178                _ => Err(WorkerError::Io(std::io::Error::other(
179                    "unexpected message kind",
180                ))),
181            };
182        }
183    }
184
185    /// Kill the worker process immediately (no graceful drain).
186    pub async fn kill(&self) -> std::io::Result<()> {
187        self.child.lock().await.kill().await
188    }
189
190    /// Wait for the worker process to exit.
191    pub async fn wait(&self) -> std::io::Result<std::process::ExitStatus> {
192        self.child.lock().await.wait().await
193    }
194
195    async fn recv_with_timeout(&self, dur: Duration) -> Result<Value, WorkerError> {
196        let mut io = self.io.lock().await;
197        timeout(dur, recv(&mut io.reader))
198            .await
199            .map_err(|_| WorkerError::PongTimeout)?
200    }
201}
202
203async fn send(writer: &mut OwnedWriteHalf, msg: &Value) -> Result<(), WorkerError> {
204    let mut b = serde_json::to_vec(msg)?;
205    b.push(b'\n');
206    writer.write_all(&b).await?;
207    writer.flush().await?;
208    Ok(())
209}
210
211async fn recv(reader: &mut BufReader<OwnedReadHalf>) -> Result<Value, WorkerError> {
212    let mut line = String::new();
213    let n = reader.read_line(&mut line).await?;
214    if n == 0 {
215        return Err(WorkerError::Closed);
216    }
217    Ok(serde_json::from_str(&line)?)
218}
219
220async fn connect_with_retry(path: &Path, deadline: Duration) -> Result<UnixStream, WorkerError> {
221    let start = std::time::Instant::now();
222    loop {
223        match UnixStream::connect(path).await {
224            Ok(s) => return Ok(s),
225            Err(e) if start.elapsed() >= deadline => return Err(e.into()),
226            Err(_) => sleep(Duration::from_millis(50)).await,
227        }
228    }
229}