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        if self.released {
217            return;
218        }
219        #[cfg(not(windows))]
220        release_client_blocking_best_effort(&self.socket, &self.lease_id);
221    }
222}
223
224pub struct PoolClientLease {
225    pub socket: String,
226    pub image: Option<String>,
227    pub volumes: Vec<String>,
228    pub vcpus: u32,
229    pub memory_mb: u32,
230}
231
232pub struct PoolLeaseExec {
233    pub cmd: Vec<String>,
234    pub timeout_ns: Option<u64>,
235    pub env: Vec<String>,
236    pub working_dir: Option<String>,
237    pub rootfs: Option<String>,
238    pub stdin: Option<Vec<u8>>,
239    pub user: Option<String>,
240}
241
242#[cfg(not(windows))]
243pub async fn run_client(req: PoolClientRun) -> Result<PoolClientOutput> {
244    use tokio::net::UnixStream;
245
246    let mut stream = UnixStream::connect(&req.socket).await.map_err(|e| {
247        BoxError::PoolError(format!(
248            "Failed to connect to pool daemon at {} ({}). Is `a3s-box pool start` running?",
249            req.socket, e
250        ))
251    })?;
252
253    write_frame(
254        &mut stream,
255        &serde_json::to_vec(&PoolRequest::Run(PoolRunRequest {
256            image: req.image,
257            user: req.user,
258            workdir: req.workdir,
259            rootfs: req.rootfs,
260            env: req.env,
261            volumes: req.volumes,
262            vcpus: Some(req.vcpus),
263            memory_mb: Some(req.memory_mb),
264            exec: req.exec,
265            timeout_ns: req.timeout_ns,
266            cmd: req.cmd,
267        }))?,
268    )
269    .await?;
270    let resp: PoolRunResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
271
272    if let Some(err) = resp.error {
273        return Err(BoxError::PoolError(err));
274    }
275
276    Ok(PoolClientOutput {
277        stdout: resp.stdout,
278        stderr: resp.stderr,
279        exit_code: resp.exit_code,
280    })
281}
282
283#[cfg(windows)]
284pub async fn run_client(_req: PoolClientRun) -> Result<PoolClientOutput> {
285    Err(BoxError::PoolError(
286        "`pool run` is not supported on Windows".to_string(),
287    ))
288}
289
290#[cfg(not(windows))]
291pub async fn status_client(socket: &str) -> Result<PoolStatusResponse> {
292    use tokio::net::UnixStream;
293
294    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
295        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
296    })?;
297    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Status)?).await?;
298    Ok(serde_json::from_slice(&read_frame(&mut stream).await?)?)
299}
300
301#[cfg(not(windows))]
302pub async fn stop_client(socket: &str) -> Result<()> {
303    use tokio::net::UnixStream;
304
305    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
306        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
307    })?;
308    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Stop)?).await?;
309    let resp: PoolStopResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
310    if let Some(error) = resp.error {
311        return Err(BoxError::PoolError(error));
312    }
313    Ok(())
314}
315
316#[cfg(windows)]
317pub async fn stop_client(_socket: &str) -> Result<()> {
318    Err(BoxError::PoolError(
319        "`pool stop` is not supported on Windows".to_string(),
320    ))
321}
322
323#[cfg(not(windows))]
324async fn lease_client(req: &PoolClientLease) -> Result<PoolLeaseResponse> {
325    use tokio::net::UnixStream;
326
327    let mut stream = UnixStream::connect(&req.socket).await.map_err(|e| {
328        BoxError::PoolError(format!(
329            "Failed to connect to pool daemon at {} ({}). Is `a3s-box pool start` running?",
330            req.socket, e
331        ))
332    })?;
333    write_frame(
334        &mut stream,
335        &serde_json::to_vec(&PoolRequest::Lease(PoolLeaseRequest {
336            image: req.image.clone(),
337            volumes: req.volumes.clone(),
338            vcpus: Some(req.vcpus),
339            memory_mb: Some(req.memory_mb),
340        }))?,
341    )
342    .await?;
343    let resp: PoolLeaseResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
344    if let Some(error) = resp.error.as_ref() {
345        return Err(BoxError::PoolError(error.clone()));
346    }
347    Ok(resp)
348}
349
350#[cfg(windows)]
351async fn lease_client(_req: &PoolClientLease) -> Result<PoolLeaseResponse> {
352    Err(BoxError::PoolError(
353        "warm-pool leases are not supported on Windows".to_string(),
354    ))
355}
356
357#[cfg(not(windows))]
358async fn lease_exec_client(socket: &str, req: PoolLeaseExecRequest) -> Result<PoolClientOutput> {
359    use tokio::net::UnixStream;
360
361    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
362        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
363    })?;
364    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Exec(req))?).await?;
365    let resp: PoolRunResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
366    if let Some(error) = resp.error {
367        return Err(BoxError::PoolError(error));
368    }
369    Ok(PoolClientOutput {
370        stdout: resp.stdout,
371        stderr: resp.stderr,
372        exit_code: resp.exit_code,
373    })
374}
375
376#[cfg(windows)]
377async fn lease_exec_client(_socket: &str, _req: PoolLeaseExecRequest) -> Result<PoolClientOutput> {
378    Err(BoxError::PoolError(
379        "warm-pool leases are not supported on Windows".to_string(),
380    ))
381}
382
383#[cfg(not(windows))]
384async fn release_client(socket: &str, lease_id: &str) -> Result<()> {
385    use tokio::net::UnixStream;
386
387    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
388        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
389    })?;
390    write_frame(
391        &mut stream,
392        &serde_json::to_vec(&PoolRequest::Release(PoolLeaseReleaseRequest {
393            lease_id: lease_id.to_string(),
394        }))?,
395    )
396    .await?;
397    let resp: PoolLeaseReleaseResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
398    if let Some(error) = resp.error {
399        return Err(BoxError::PoolError(error));
400    }
401    Ok(())
402}
403
404#[cfg(windows)]
405async fn release_client(_socket: &str, _lease_id: &str) -> Result<()> {
406    Err(BoxError::PoolError(
407        "warm-pool leases are not supported on Windows".to_string(),
408    ))
409}
410
411#[cfg(not(windows))]
412fn release_client_blocking_best_effort(socket: &str, lease_id: &str) {
413    use std::io::Write;
414    use std::os::unix::net::UnixStream;
415    use std::time::Duration;
416
417    let Ok(mut stream) = UnixStream::connect(socket) else {
418        return;
419    };
420    let timeout = Some(Duration::from_millis(500));
421    let _ = stream.set_read_timeout(timeout);
422    let _ = stream.set_write_timeout(timeout);
423
424    let Ok(payload) = serde_json::to_vec(&PoolRequest::Release(PoolLeaseReleaseRequest {
425        lease_id: lease_id.to_string(),
426    })) else {
427        return;
428    };
429    let _ = stream
430        .write_all(&(payload.len() as u32).to_le_bytes())
431        .and_then(|_| stream.write_all(&payload))
432        .and_then(|_| stream.flush());
433}
434
435/// Length-prefixed (u32 LE) framing for the pool Unix-socket protocol.
436#[cfg(not(windows))]
437pub async fn write_frame<W>(w: &mut W, data: &[u8]) -> std::io::Result<()>
438where
439    W: tokio::io::AsyncWrite + Unpin,
440{
441    use tokio::io::AsyncWriteExt;
442
443    w.write_all(&(data.len() as u32).to_le_bytes()).await?;
444    w.write_all(data).await?;
445    w.flush().await
446}
447
448#[cfg(not(windows))]
449pub async fn read_frame<R>(r: &mut R) -> std::io::Result<Vec<u8>>
450where
451    R: tokio::io::AsyncRead + Unpin,
452{
453    use tokio::io::AsyncReadExt;
454
455    let mut len = [0u8; 4];
456    r.read_exact(&mut len).await?;
457    let mut buf = vec![0u8; u32::from_le_bytes(len) as usize];
458    r.read_exact(&mut buf).await?;
459    Ok(buf)
460}
461
462#[cfg(test)]
463mod tests {
464    #[cfg(not(windows))]
465    #[test]
466    fn lease_drop_releases_synchronously() {
467        use super::*;
468        use std::io::Read;
469        use std::os::unix::net::UnixListener;
470
471        let tmp = tempfile::TempDir::new().unwrap();
472        let socket = tmp.path().join("pool.sock");
473        let listener = UnixListener::bind(&socket).unwrap();
474        let socket_arg = socket.to_string_lossy().to_string();
475
476        let server = std::thread::spawn(move || {
477            let (mut stream, _) = listener.accept().unwrap();
478            let mut len = [0_u8; 4];
479            stream.read_exact(&mut len).unwrap();
480            let mut request = vec![0_u8; u32::from_le_bytes(len) as usize];
481            stream.read_exact(&mut request).unwrap();
482            let request: PoolRequest = serde_json::from_slice(&request).unwrap();
483            match request {
484                PoolRequest::Release(req) => assert_eq!(req.lease_id, "lease-drop"),
485                _ => panic!("drop should send release request"),
486            }
487        });
488
489        let lease = PoolLeaseClient {
490            socket: socket_arg,
491            lease_id: "lease-drop".to_string(),
492            released: false,
493        };
494        drop(lease);
495
496        server.join().unwrap();
497    }
498}
a3s_box_runtime/pool/client.rs

a3s_box_runtime/pool/
client.rs
