mesofact_core/proxy/
worker_client.rs1use 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 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 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 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 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 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 pub async fn kill(&self) -> std::io::Result<()> {
187 self.child.lock().await.kill().await
188 }
189
190 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}