a3s_box_runtime/pool/
client.rs

1//! Socket protocol and client helpers for the warm-pool daemon.
2
3use a3s_box_core::error::{BoxError, Result};
4use serde::{Deserialize, Serialize};
5
6/// Wire protocol for the `pool` Unix socket.
7///
8/// Client→daemon request: run a one-shot command, query status, stop the
9/// daemon, or manage a short-lived leased VM session.
10#[derive(Serialize, Deserialize)]
11#[serde(tag = "op", rename_all = "snake_case")]
12pub enum PoolRequest {
13    Run(PoolRunRequest),
14    Status,
15    Stop,
16    Lease(PoolLeaseRequest),
17    Exec(PoolLeaseExecRequest),
18    Release(PoolLeaseReleaseRequest),
19}
20
21#[derive(Serialize, Deserialize)]
22pub struct PoolRunRequest {
23    /// Image to run in; `None` means use the daemon's default image.
24    #[serde(default)]
25    pub image: Option<String>,
26    /// User to run as (`uid[:gid]` or name); `None` runs as the image default.
27    #[serde(default)]
28    pub user: Option<String>,
29    /// Working directory inside the sandbox.
30    #[serde(default)]
31    pub workdir: Option<String>,
32    /// Optional guest-visible rootfs to chroot into before executing.
33    #[serde(default)]
34    pub rootfs: Option<String>,
35    /// Extra KEY=VALUE environment entries.
36    #[serde(default)]
37    pub env: Vec<String>,
38    /// Boot-time volume specs for this sandbox pool.
39    #[serde(default)]
40    pub volumes: Vec<String>,
41    /// Boot-time vCPU count for lazily-created pools.
42    #[serde(default)]
43    pub vcpus: Option<u32>,
44    /// Boot-time memory size for lazily-created pools.
45    #[serde(default)]
46    pub memory_mb: Option<u32>,
47    /// Force exec mode for this request.
48    #[serde(default)]
49    pub exec: bool,
50    /// Guest-side execution timeout in nanoseconds.
51    #[serde(default)]
52    pub timeout_ns: Option<u64>,
53    pub cmd: Vec<String>,
54}
55
56#[derive(Serialize, Deserialize)]
57pub struct PoolRunResponse {
58    pub stdout: Vec<u8>,
59    pub stderr: Vec<u8>,
60    pub exit_code: i32,
61    pub error: Option<String>,
62}
63
64#[derive(Serialize, Deserialize)]
65pub struct PoolLeaseRequest {
66    /// Image for the helper VM; `None` means use the daemon's default image.
67    #[serde(default)]
68    pub image: Option<String>,
69    /// Boot-time volume specs for this leased VM.
70    #[serde(default)]
71    pub volumes: Vec<String>,
72    /// Boot-time vCPU count.
73    #[serde(default)]
74    pub vcpus: Option<u32>,
75    /// Boot-time memory size.
76    #[serde(default)]
77    pub memory_mb: Option<u32>,
78}
79
80#[derive(Serialize, Deserialize)]
81pub struct PoolLeaseResponse {
82    pub lease_id: Option<String>,
83    pub error: Option<String>,
84}
85
86#[derive(Serialize, Deserialize)]
87pub struct PoolLeaseExecRequest {
88    pub lease_id: String,
89    pub cmd: Vec<String>,
90    #[serde(default)]
91    pub timeout_ns: Option<u64>,
92    #[serde(default)]
93    pub env: Vec<String>,
94    #[serde(default)]
95    pub working_dir: Option<String>,
96    #[serde(default)]
97    pub rootfs: Option<String>,
98    #[serde(default)]
99    pub stdin: Option<Vec<u8>>,
100    #[serde(default)]
101    pub user: Option<String>,
102}
103
104#[derive(Serialize, Deserialize)]
105pub struct PoolLeaseReleaseRequest {
106    pub lease_id: String,
107}
108
109#[derive(Serialize, Deserialize)]
110pub struct PoolLeaseReleaseResponse {
111    pub error: Option<String>,
112}
113
114/// Live stats for one image's warm pool.
115#[derive(Serialize, Deserialize)]
116pub struct PoolImageStat {
117    pub image: String,
118    pub pool: String,
119    /// Maximum concurrent sandboxes for this pool key.
120    #[serde(default)]
121    pub max: usize,
122    pub idle: usize,
123    /// Sandboxes currently checked out by one-shot runs or leases.
124    #[serde(default)]
125    pub active: usize,
126    /// Active sandboxes held by lease clients.
127    #[serde(default)]
128    pub leased: usize,
129    pub total_created: u64,
130    pub total_acquired: u64,
131    pub total_evicted: u64,
132}
133
134#[derive(Serialize, Deserialize)]
135pub struct PoolStatusResponse {
136    pub images: Vec<PoolImageStat>,
137}
138
139#[derive(Serialize, Deserialize)]
140pub struct PoolStopResponse {
141    pub error: Option<String>,
142}
143
144pub struct PoolClientRun {
145    pub socket: String,
146    pub image: Option<String>,
147    pub user: Option<String>,
148    pub workdir: Option<String>,
149    pub rootfs: Option<String>,
150    pub env: Vec<String>,
151    pub volumes: Vec<String>,
152    pub vcpus: u32,
153    pub memory_mb: u32,
154    pub exec: bool,
155    pub timeout_ns: Option<u64>,
156    pub cmd: Vec<String>,
157}
158
159pub struct PoolClientOutput {
160    pub stdout: Vec<u8>,
161    pub stderr: Vec<u8>,
162    pub exit_code: i32,
163}
164
165pub struct PoolLeaseClient {
166    socket: String,
167    lease_id: String,
168    released: bool,
169}
170
171impl PoolLeaseClient {
172    pub fn lease_id(&self) -> &str {
173        &self.lease_id
174    }
175
176    pub async fn acquire(req: PoolClientLease) -> Result<Self> {
177        let response = lease_client(&req).await?;
178        let lease_id = response.lease_id.ok_or_else(|| {
179            BoxError::PoolError("pool lease response did not include a lease id".to_string())
180        })?;
181        Ok(Self {
182            socket: req.socket,
183            lease_id,
184            released: false,
185        })
186    }
187
188    pub async fn exec(&self, req: PoolLeaseExec) -> Result<PoolClientOutput> {
189        lease_exec_client(
190            &self.socket,
191            PoolLeaseExecRequest {
192                lease_id: self.lease_id.clone(),
193                cmd: req.cmd,
194                timeout_ns: req.timeout_ns,
195                env: req.env,
196                working_dir: req.working_dir,
197                rootfs: req.rootfs,
198                stdin: req.stdin,
199                user: req.user,
200            },
201        )
202        .await
203    }
204
205    pub async fn release(mut self) -> Result<()> {
206        let result = release_client(&self.socket, &self.lease_id).await;
207        if result.is_ok() {
208            self.released = true;
209        }
210        result
211    }
212}
213
214impl Drop for PoolLeaseClient {
215    fn drop(&mut self) {
216        #[cfg(not(windows))]
217        if !self.released {
218            release_client_blocking_best_effort(&self.socket, &self.lease_id);
219        }
220    }
221}
222
223pub struct PoolClientLease {
224    pub socket: String,
225    pub image: Option<String>,
226    pub volumes: Vec<String>,
227    pub vcpus: u32,
228    pub memory_mb: u32,
229}
230
231pub struct PoolLeaseExec {
232    pub cmd: Vec<String>,
233    pub timeout_ns: Option<u64>,
234    pub env: Vec<String>,
235    pub working_dir: Option<String>,
236    pub rootfs: Option<String>,
237    pub stdin: Option<Vec<u8>>,
238    pub user: Option<String>,
239}
240
241#[cfg(not(windows))]
242pub async fn run_client(req: PoolClientRun) -> Result<PoolClientOutput> {
243    use tokio::net::UnixStream;
244
245    let mut stream = UnixStream::connect(&req.socket).await.map_err(|e| {
246        BoxError::PoolError(format!(
247            "Failed to connect to pool daemon at {} ({}). Is `a3s-box pool start` running?",
248            req.socket, e
249        ))
250    })?;
251
252    write_frame(
253        &mut stream,
254        &serde_json::to_vec(&PoolRequest::Run(PoolRunRequest {
255            image: req.image,
256            user: req.user,
257            workdir: req.workdir,
258            rootfs: req.rootfs,
259            env: req.env,
260            volumes: req.volumes,
261            vcpus: Some(req.vcpus),
262            memory_mb: Some(req.memory_mb),
263            exec: req.exec,
264            timeout_ns: req.timeout_ns,
265            cmd: req.cmd,
266        }))?,
267    )
268    .await?;
269    let resp: PoolRunResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
270
271    if let Some(err) = resp.error {
272        return Err(BoxError::PoolError(err));
273    }
274
275    Ok(PoolClientOutput {
276        stdout: resp.stdout,
277        stderr: resp.stderr,
278        exit_code: resp.exit_code,
279    })
280}
281
282#[cfg(windows)]
283pub async fn run_client(_req: PoolClientRun) -> Result<PoolClientOutput> {
284    Err(BoxError::PoolError(
285        "`pool run` is not supported on Windows".to_string(),
286    ))
287}
288
289#[cfg(not(windows))]
290pub async fn status_client(socket: &str) -> Result<PoolStatusResponse> {
291    use tokio::net::UnixStream;
292
293    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
294        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
295    })?;
296    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Status)?).await?;
297    Ok(serde_json::from_slice(&read_frame(&mut stream).await?)?)
298}
299
300#[cfg(not(windows))]
301pub async fn stop_client(socket: &str) -> Result<()> {
302    use tokio::net::UnixStream;
303
304    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
305        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
306    })?;
307    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Stop)?).await?;
308    let resp: PoolStopResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
309    if let Some(error) = resp.error {
310        return Err(BoxError::PoolError(error));
311    }
312    Ok(())
313}
314
315#[cfg(windows)]
316pub async fn stop_client(_socket: &str) -> Result<()> {
317    Err(BoxError::PoolError(
318        "`pool stop` is not supported on Windows".to_string(),
319    ))
320}
321
322#[cfg(not(windows))]
323async fn lease_client(req: &PoolClientLease) -> Result<PoolLeaseResponse> {
324    use tokio::net::UnixStream;
325
326    let mut stream = UnixStream::connect(&req.socket).await.map_err(|e| {
327        BoxError::PoolError(format!(
328            "Failed to connect to pool daemon at {} ({}). Is `a3s-box pool start` running?",
329            req.socket, e
330        ))
331    })?;
332    write_frame(
333        &mut stream,
334        &serde_json::to_vec(&PoolRequest::Lease(PoolLeaseRequest {
335            image: req.image.clone(),
336            volumes: req.volumes.clone(),
337            vcpus: Some(req.vcpus),
338            memory_mb: Some(req.memory_mb),
339        }))?,
340    )
341    .await?;
342    let resp: PoolLeaseResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
343    if let Some(error) = resp.error.as_ref() {
344        return Err(BoxError::PoolError(error.clone()));
345    }
346    Ok(resp)
347}
348
349#[cfg(windows)]
350async fn lease_client(_req: &PoolClientLease) -> Result<PoolLeaseResponse> {
351    Err(BoxError::PoolError(
352        "warm-pool leases are not supported on Windows".to_string(),
353    ))
354}
355
356#[cfg(not(windows))]
357async fn lease_exec_client(socket: &str, req: PoolLeaseExecRequest) -> Result<PoolClientOutput> {
358    use tokio::net::UnixStream;
359
360    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
361        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
362    })?;
363    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Exec(req))?).await?;
364    let resp: PoolRunResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
365    if let Some(error) = resp.error {
366        return Err(BoxError::PoolError(error));
367    }
368    Ok(PoolClientOutput {
369        stdout: resp.stdout,
370        stderr: resp.stderr,
371        exit_code: resp.exit_code,
372    })
373}
374
375#[cfg(windows)]
376async fn lease_exec_client(_socket: &str, _req: PoolLeaseExecRequest) -> Result<PoolClientOutput> {
377    Err(BoxError::PoolError(
378        "warm-pool leases are not supported on Windows".to_string(),
379    ))
380}
381
382#[cfg(not(windows))]
383async fn release_client(socket: &str, lease_id: &str) -> Result<()> {
384    use tokio::net::UnixStream;
385
386    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
387        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
388    })?;
389    write_frame(
390        &mut stream,
391        &serde_json::to_vec(&PoolRequest::Release(PoolLeaseReleaseRequest {
392            lease_id: lease_id.to_string(),
393        }))?,
394    )
395    .await?;
396    let resp: PoolLeaseReleaseResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
397    if let Some(error) = resp.error {
398        return Err(BoxError::PoolError(error));
399    }
400    Ok(())
401}
402
403#[cfg(windows)]
404async fn release_client(_socket: &str, _lease_id: &str) -> Result<()> {
405    Err(BoxError::PoolError(
406        "warm-pool leases are not supported on Windows".to_string(),
407    ))
408}
409
410#[cfg(not(windows))]
411fn release_client_blocking_best_effort(socket: &str, lease_id: &str) {
412    use std::io::Write;
413    use std::os::unix::net::UnixStream;
414    use std::time::Duration;
415
416    let Ok(mut stream) = UnixStream::connect(socket) else {
417        return;
418    };
419    let timeout = Some(Duration::from_millis(500));
420    let _ = stream.set_read_timeout(timeout);
421    let _ = stream.set_write_timeout(timeout);
422
423    let Ok(payload) = serde_json::to_vec(&PoolRequest::Release(PoolLeaseReleaseRequest {
424        lease_id: lease_id.to_string(),
425    })) else {
426        return;
427    };
428    let _ = stream
429        .write_all(&(payload.len() as u32).to_le_bytes())
430        .and_then(|_| stream.write_all(&payload))
431        .and_then(|_| stream.flush());
432}
433
434/// Length-prefixed (u32 LE) framing for the pool Unix-socket protocol.
435#[cfg(not(windows))]
436pub async fn write_frame<W>(w: &mut W, data: &[u8]) -> std::io::Result<()>
437where
438    W: tokio::io::AsyncWrite + Unpin,
439{
440    use tokio::io::AsyncWriteExt;
441
442    w.write_all(&(data.len() as u32).to_le_bytes()).await?;
443    w.write_all(data).await?;
444    w.flush().await
445}
446
447#[cfg(not(windows))]
448pub async fn read_frame<R>(r: &mut R) -> std::io::Result<Vec<u8>>
449where
450    R: tokio::io::AsyncRead + Unpin,
451{
452    use tokio::io::AsyncReadExt;
453
454    let mut len = [0u8; 4];
455    r.read_exact(&mut len).await?;
456    let mut buf = vec![0u8; u32::from_le_bytes(len) as usize];
457    r.read_exact(&mut buf).await?;
458    Ok(buf)
459}
460
461#[cfg(test)]
462mod tests {
463    #[cfg(not(windows))]
464    #[test]
465    fn lease_drop_releases_synchronously() {
466        use super::*;
467        use std::io::Read;
468        use std::os::unix::net::UnixListener;
469
470        let tmp = tempfile::TempDir::new().unwrap();
471        let socket = tmp.path().join("pool.sock");
472        let listener = UnixListener::bind(&socket).unwrap();
473        let socket_arg = socket.to_string_lossy().to_string();
474
475        let server = std::thread::spawn(move || {
476            let (mut stream, _) = listener.accept().unwrap();
477            let mut len = [0_u8; 4];
478            stream.read_exact(&mut len).unwrap();
479            let mut request = vec![0_u8; u32::from_le_bytes(len) as usize];
480            stream.read_exact(&mut request).unwrap();
481            let request: PoolRequest = serde_json::from_slice(&request).unwrap();
482            match request {
483                PoolRequest::Release(req) => assert_eq!(req.lease_id, "lease-drop"),
484                _ => panic!("drop should send release request"),
485            }
486        });
487
488        let lease = PoolLeaseClient {
489            socket: socket_arg,
490            lease_id: "lease-drop".to_string(),
491            released: false,
492        };
493        drop(lease);
494
495        server.join().unwrap();
496    }
497}
a3s_box_runtime/pool/client.rs

a3s_box_runtime/pool/
client.rs
