Skip to main content

arcbox_computer_runtime/sandbox/
execution.rs

1//! Addressable, resumable workload executions (CORE-55).
2//!
3//! An execution is a workload started inside a sandbox that outlives any
4//! particular client connection. The manager retains offset-addressed
5//! stdout/stderr per execution so a client can disconnect, reconnect, and
6//! resume both streams without loss, and gives the process an identity that
7//! can be written to, signalled, resized, and awaited independently of any
8//! stream. Stdin writes are offset-idempotent: a retried write of already
9//! accepted bytes is deduplicated instead of double-fed to the process.
10
11use std::collections::VecDeque;
12
13use tokio::sync::{mpsc, watch};
14
15use super::workload::{WorkloadClaim, WorkloadSlot};
16use super::*;
17use crate::lifecycle::actor::WorkloadOutcome;
18
19/// Retained bytes per stdio channel. Older bytes are dropped and the buffer's
20/// base offset advances; an attach below the base resumes from the earliest
21/// retained byte (chunk offsets expose the gap to the client).
22const CHANNEL_RETENTION: usize = 8 * 1024 * 1024;
23
24/// How long an exited execution stays attachable/awaitable before it is
25/// dropped from the registry.
26const EXITED_RETENTION: Duration = Duration::from_secs(5 * 60);
27
28/// Largest single chunk forwarded on an attach stream.
29const ATTACH_CHUNK: usize = 64 * 1024;
30
31/// Buffered chunks between an execution's buffers and one attach consumer.
32const ATTACH_QUEUE: usize = 64;
33
34/// How many times a start re-checks the registry after awaiting an identical
35/// in-flight start. One round is the normal case.
36const MAX_START_RESERVE_ROUNDS: usize = 8;
37
38/// Parameters for starting an execution.
39#[derive(Debug, Clone, Default)]
40pub struct ExecutionSpec {
41    /// Caller-supplied id for idempotent retry (auto-generated when `None`).
42    pub id: Option<String>,
43    /// Command and arguments.
44    pub cmd: Vec<String>,
45    /// Environment variable overrides.
46    pub env: HashMap<String, String>,
47    /// Working directory (empty = rootfs default).
48    pub working_dir: String,
49    /// User to run as (empty = root).
50    pub user: String,
51    /// Allocate a pseudo-TTY (stdout carries the merged PTY stream).
52    pub tty: bool,
53    /// Initial terminal size when `tty` is set.
54    pub tty_size: Option<(u16, u16)>,
55    /// Kill the process after this many seconds (0 = no timeout).
56    pub timeout_seconds: u32,
57    /// Keep stdin open for [`ComputerManager::write_stdin`]. When false the
58    /// process starts with stdin already at EOF (run semantics).
59    pub stdin: bool,
60}
61
62/// Point-in-time view of an execution.
63#[derive(Debug, Clone)]
64pub struct ExecutionSnapshot {
65    /// Execution id, unique within its sandbox.
66    pub id: String,
67    /// Owning sandbox.
68    pub sandbox_id: ComputerId,
69    /// Whether the workload runs on a pseudo-TTY.
70    pub tty: bool,
71    /// When the workload was dispatched.
72    pub started_at: DateTime<Utc>,
73    /// When the workload terminated (`None` while running).
74    pub exited_at: Option<DateTime<Utc>>,
75    /// How the process terminated. `None` while running, and also when the
76    /// session broke before an exit was observed (see `error`).
77    pub exit_status: Option<ExitStatus>,
78    /// Set when the session broke before the process reported an exit.
79    pub error: Option<String>,
80    /// Total bytes ever produced on stdout (monotonic, beyond retention).
81    pub stdout_len: u64,
82    /// Total bytes ever produced on stderr.
83    pub stderr_len: u64,
84    /// Stdin acceptance state.
85    pub stdin: StdinState,
86}
87
88impl ExecutionSnapshot {
89    /// True until the workload exits (or its session breaks).
90    #[must_use]
91    pub const fn is_running(&self) -> bool {
92        self.exited_at.is_none()
93    }
94}
95
96/// Which stdio channel an attach chunk belongs to. For TTY executions the
97/// merged PTY stream arrives on `Stdout`.
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99pub enum ExecutionChannel {
100    /// Standard output (or the PTY stream for `tty` executions).
101    Stdout,
102    /// Standard error (never produced by `tty` executions).
103    Stderr,
104}
105
106/// One offset-addressed chunk delivered on an attach stream.
107#[derive(Debug, Clone)]
108pub struct ExecutionOutput {
109    /// Source channel.
110    pub channel: ExecutionChannel,
111    /// Absolute offset of `data[0]` within the channel's byte stream. May be
112    /// greater than the requested resume offset when retention dropped bytes.
113    pub offset: u64,
114    /// Raw output bytes.
115    pub data: Vec<u8>,
116}
117
118/// Stdin acceptance state, the anchor for offset-resumable writes.
119#[derive(Debug, Clone, Copy, Default)]
120pub struct StdinState {
121    /// Bytes accepted and forwarded so far — the offset at which the next
122    /// write must start.
123    pub bytes_written: u64,
124    /// Whether stdin has been closed.
125    pub closed: bool,
126}
127
128/// Retained suffix of one stdio channel, addressed by absolute offsets.
129#[derive(Default)]
130struct ChannelBuffer {
131    /// Absolute offset of `data[0]`.
132    base: u64,
133    data: VecDeque<u8>,
134}
135
136impl ChannelBuffer {
137    /// One past the last retained byte — the channel's total length so far.
138    fn end(&self) -> u64 {
139        self.base + self.data.len() as u64
140    }
141
142    fn append(&mut self, bytes: &[u8]) {
143        self.data.extend(bytes);
144        if self.data.len() > CHANNEL_RETENTION {
145            let drop = self.data.len() - CHANNEL_RETENTION;
146            self.data.drain(..drop);
147            self.base += drop as u64;
148        }
149    }
150
151    /// Read up to `max` bytes starting at `offset`, clamped to the retained
152    /// range. Returns the actual start offset alongside the bytes.
153    fn read_from(&self, offset: u64, max: usize) -> (u64, Vec<u8>) {
154        let start = offset.max(self.base);
155        if start >= self.end() {
156            return (start, Vec::new());
157        }
158        // The buffer is capped at CHANNEL_RETENTION, so the skip fits usize.
159        let skip = usize::try_from(start - self.base).expect("offset within retained buffer");
160        let take = (self.data.len() - skip).min(max);
161        let (front, back) = self.data.as_slices();
162        let mut out = Vec::with_capacity(take);
163        if skip < front.len() {
164            let n = (front.len() - skip).min(take);
165            out.extend_from_slice(&front[skip..skip + n]);
166            out.extend_from_slice(&back[..take - n]);
167        } else {
168            let off = skip - front.len();
169            out.extend_from_slice(&back[off..off + take]);
170        }
171        (start, out)
172    }
173}
174
175/// Mutable state of one execution. Guarded by a std mutex; never held across
176/// an await point.
177#[derive(Default)]
178struct ExecState {
179    stdout: ChannelBuffer,
180    stderr: ChannelBuffer,
181    stdin_written: u64,
182    stdin_closed: bool,
183    exited_at: Option<DateTime<Utc>>,
184    exit_status: Option<ExitStatus>,
185    error: Option<String>,
186}
187
188/// A live or recently-exited execution.
189pub(super) struct Execution {
190    id: String,
191    sandbox_id: ComputerId,
192    tty: bool,
193    started_at: DateTime<Utc>,
194    /// The started command, kept to distinguish an idempotent start retry
195    /// (same id, same command) from an id collision.
196    cmd: Vec<String>,
197    state: Mutex<ExecState>,
198    /// Bumped after every state mutation; attach and wait subscribers park on
199    /// it instead of polling.
200    version: watch::Sender<u64>,
201    /// Session input queue (stdin bytes / resize / signal / EOF).
202    input_tx: mpsc::Sender<ExecInputMsg>,
203    /// Serialises `write_stdin` offset accounting across concurrent callers.
204    stdin_gate: tokio::sync::Mutex<()>,
205}
206
207impl Execution {
208    fn new(
209        id: String,
210        sandbox_id: ComputerId,
211        spec: &ExecutionSpec,
212        input_tx: mpsc::Sender<ExecInputMsg>,
213    ) -> Self {
214        Self {
215            id,
216            sandbox_id,
217            tty: spec.tty,
218            started_at: Utc::now(),
219            cmd: spec.cmd.clone(),
220            state: Mutex::new(ExecState {
221                stdin_closed: !spec.stdin,
222                ..ExecState::default()
223            }),
224            version: watch::channel(0).0,
225            input_tx,
226            stdin_gate: tokio::sync::Mutex::new(()),
227        }
228    }
229
230    fn snapshot(&self) -> ExecutionSnapshot {
231        let st = self.state.lock().unwrap();
232        ExecutionSnapshot {
233            id: self.id.clone(),
234            sandbox_id: self.sandbox_id.clone(),
235            tty: self.tty,
236            started_at: self.started_at,
237            exited_at: st.exited_at,
238            exit_status: st.exit_status,
239            error: st.error.clone(),
240            stdout_len: st.stdout.end(),
241            stderr_len: st.stderr.end(),
242            stdin: StdinState {
243                bytes_written: st.stdin_written,
244                closed: st.stdin_closed,
245            },
246        }
247    }
248
249    fn bump(&self) {
250        self.version.send_modify(|v| *v += 1);
251    }
252
253    /// Record the session's end and wake all subscribers. Idempotent: the
254    /// first outcome wins (a teardown purge and the session's own exit path
255    /// can race).
256    fn mark_exited(&self, outcome: &std::result::Result<ExitStatus, String>) {
257        let mut st = self.state.lock().unwrap();
258        if st.exited_at.is_some() {
259            return;
260        }
261        st.exited_at = Some(Utc::now());
262        match outcome {
263            Ok(status) => st.exit_status = Some(*status),
264            Err(e) => st.error = Some(e.clone()),
265        }
266        drop(st);
267        self.bump();
268    }
269
270    fn has_exited(&self) -> bool {
271        self.state.lock().unwrap().exited_at.is_some()
272    }
273
274    /// Error for operations that need the workload alive.
275    fn exited_error(&self) -> ComputerError {
276        ComputerError::WrongState {
277            id: format!("execution '{}'", self.id),
278            expected: "running".into(),
279            actual: "exited".into(),
280        }
281    }
282
283    /// Spawn a forwarder that replays buffered output from the given offsets
284    /// and follows live output until exit.
285    fn attach(
286        self: &Arc<Self>,
287        stdout_offset: u64,
288        stderr_offset: u64,
289    ) -> (ExecutionSnapshot, mpsc::Receiver<ExecutionOutput>) {
290        let snapshot = self.snapshot();
291        let (tx, rx) = mpsc::channel(ATTACH_QUEUE);
292        tokio::spawn(pump_attach(
293            Arc::clone(self),
294            stdout_offset,
295            stderr_offset,
296            tx,
297        ));
298        (snapshot, rx)
299    }
300
301    /// Offset-idempotent stdin write; see [`ComputerManager::write_stdin`].
302    async fn write_stdin(&self, offset: u64, data: &[u8], eof: bool) -> Result<StdinState> {
303        if eof && self.tty {
304            return Err(ComputerError::Config(
305                "cannot close stdin of a TTY execution; send Ctrl-D (0x04) instead".into(),
306            ));
307        }
308        // One writer at a time: the accept-forward-account sequence below
309        // must not interleave between concurrent callers.
310        let _gate = self.stdin_gate.lock().await;
311
312        let fresh = {
313            let st = self.state.lock().unwrap();
314            if st.exited_at.is_some() {
315                return Err(self.exited_error());
316            }
317            if offset > st.stdin_written {
318                return Err(ComputerError::StdinGap {
319                    accepted: st.stdin_written,
320                    offset,
321                });
322            }
323            let overlap = st.stdin_written - offset;
324            if overlap >= data.len() as u64 {
325                // Pure retry of already-accepted bytes.
326                Vec::new()
327            } else {
328                if st.stdin_closed {
329                    return Err(ComputerError::Config("stdin is already closed".into()));
330                }
331                #[allow(
332                    clippy::cast_possible_truncation,
333                    reason = "overlap < data.len() which is usize"
334                )]
335                let skip = overlap as usize;
336                data[skip..].to_vec()
337            }
338        };
339
340        if !fresh.is_empty() {
341            let accepted = fresh.len() as u64;
342            self.input_tx
343                .send(ExecInputMsg::Stdin(fresh))
344                .await
345                .map_err(|_| self.exited_error())?;
346            self.state.lock().unwrap().stdin_written += accepted;
347            self.bump();
348        }
349        if eof && !self.state.lock().unwrap().stdin_closed {
350            // Best-effort: the process may exit between the check and the
351            // send; the closed flag still records what the client asked for.
352            let _ = self.input_tx.send(ExecInputMsg::Eof).await;
353            self.state.lock().unwrap().stdin_closed = true;
354            self.bump();
355        }
356
357        Ok(self.snapshot().stdin)
358    }
359
360    /// Deliver a POSIX signal to the running workload's process group.
361    async fn signal(&self, signal: i32) -> Result<()> {
362        if !(1..=64).contains(&signal) {
363            return Err(ComputerError::Config(format!("invalid signal {signal}")));
364        }
365        if self.has_exited() {
366            return Err(self.exited_error());
367        }
368        self.input_tx
369            .send(ExecInputMsg::Signal(signal))
370            .await
371            .map_err(|_| self.exited_error())
372    }
373
374    /// Resize the workload's pseudo-TTY.
375    async fn resize(&self, width: u16, height: u16) -> Result<()> {
376        if !self.tty {
377            return Err(ComputerError::Config(
378                "execution has no TTY to resize".into(),
379            ));
380        }
381        if self.has_exited() {
382            return Err(self.exited_error());
383        }
384        self.input_tx
385            .send(ExecInputMsg::Resize { width, height })
386            .await
387            .map_err(|_| self.exited_error())
388    }
389
390    /// Wait until the workload exits or `timeout` elapses; zero polls.
391    async fn wait(&self, timeout: Duration) -> ExecutionSnapshot {
392        let mut version = self.version.subscribe();
393        let deadline = tokio::time::Instant::now() + timeout;
394        loop {
395            // Mark the version seen before inspecting state, so a mutation in
396            // between leaves a pending change notification.
397            version.borrow_and_update();
398            if self.has_exited() || timeout.is_zero() {
399                return self.snapshot();
400            }
401            if tokio::time::timeout_at(deadline, version.changed())
402                .await
403                .is_err()
404            {
405                return self.snapshot();
406            }
407        }
408    }
409}
410
411type ExecKey = (ComputerId, String);
412
413/// Registry of executions, keyed by `(sandbox_id, execution_id)`.
414#[derive(Default)]
415pub(super) struct ExecutionRegistry {
416    inner: Mutex<RegistryInner>,
417}
418
419#[derive(Default)]
420struct RegistryInner {
421    live: HashMap<ExecKey, Arc<Execution>>,
422    /// Keys reserved by an in-flight start, so a duplicate id can never
423    /// dispatch a second process (mirrors `reserve_id` for sandboxes).
424    pending: HashMap<ExecKey, PendingStart>,
425}
426
427/// A start that has reserved its key but not finished dispatching.
428struct PendingStart {
429    /// The command being started, so a retry can tell an idempotent replay
430    /// from an id collision before the execution exists.
431    cmd: Vec<String>,
432    /// Fires once the start commits or unwinds, waking matching retries.
433    done: watch::Sender<bool>,
434}
435
436/// Outcome of reserving an execution slot.
437enum Reserve {
438    /// The id already ran this exact command — an idempotent start retry.
439    Existing(ExecutionSnapshot),
440    /// The slot is reserved; commit or drop to release.
441    Slot(SlotGuard),
442    /// An identical start is still in flight; await it, then re-reserve.
443    AwaitPending(watch::Receiver<bool>),
444}
445
446/// RAII reservation for an execution key; removed on drop unless committed.
447struct SlotGuard {
448    registry: Arc<ExecutionRegistry>,
449    key: ExecKey,
450    committed: bool,
451}
452
453impl SlotGuard {
454    fn commit(mut self, exec: &Arc<Execution>) {
455        let mut inner = self.registry.inner.lock().unwrap();
456        let pending = inner.pending.remove(&self.key);
457        inner.live.insert(self.key.clone(), Arc::clone(exec));
458        drop(inner);
459        // Wake matching retries only after the execution is visible, so the
460        // one that wakes finds it live rather than racing back to pending.
461        if let Some(pending) = pending {
462            let _ = pending.done.send(true);
463        }
464        self.committed = true;
465    }
466}
467
468impl Drop for SlotGuard {
469    fn drop(&mut self) {
470        if !self.committed {
471            let pending = self
472                .registry
473                .inner
474                .lock()
475                .unwrap()
476                .pending
477                .remove(&self.key);
478            if let Some(pending) = pending {
479                let _ = pending.done.send(true);
480            }
481        }
482    }
483}
484
485impl ExecutionRegistry {
486    fn reserve(self: &Arc<Self>, key: ExecKey, cmd: &[String]) -> Result<Reserve> {
487        let mut inner = self.inner.lock().unwrap();
488        if let Some(existing) = inner.live.get(&key) {
489            if existing.cmd == cmd {
490                return Ok(Reserve::Existing(existing.snapshot()));
491            }
492            return Err(ComputerError::AlreadyExists(format!(
493                "execution '{}' (id reused for a different command)",
494                key.1
495            )));
496        }
497        if let Some(pending) = inner.pending.get(&key) {
498            if pending.cmd == cmd {
499                // A retry inside the start window is still an idempotent
500                // retry: wait for the original rather than failing, or the
501                // contract only holds once the process is already running.
502                return Ok(Reserve::AwaitPending(pending.done.subscribe()));
503            }
504            return Err(ComputerError::AlreadyExists(format!(
505                "execution '{}' (id reused for a different command)",
506                key.1
507            )));
508        }
509        inner.pending.insert(
510            key.clone(),
511            PendingStart {
512                cmd: cmd.to_vec(),
513                done: watch::channel(false).0,
514            },
515        );
516        Ok(Reserve::Slot(SlotGuard {
517            registry: Arc::clone(self),
518            key,
519            committed: false,
520        }))
521    }
522
523    fn get(&self, sandbox_id: &str, execution_id: &str) -> Result<Arc<Execution>> {
524        self.inner
525            .lock()
526            .unwrap()
527            .live
528            .get(&(sandbox_id.to_owned(), execution_id.to_owned()))
529            .cloned()
530            .ok_or_else(|| {
531                ComputerError::NotFound(format!(
532                    "execution '{execution_id}' in sandbox '{sandbox_id}'"
533                ))
534            })
535    }
536
537    /// Snapshots of every retained execution of one sandbox — running and
538    /// exited, until the retention GC drops them.
539    fn list(&self, sandbox_id: &str) -> Vec<ExecutionSnapshot> {
540        self.inner
541            .lock()
542            .unwrap()
543            .live
544            .iter()
545            .filter(|((sid, _), _)| sid == sandbox_id)
546            .map(|(_, exec)| exec.snapshot())
547            .collect()
548    }
549
550    /// Remove `exec`'s entry if it is still the registered generation.
551    fn remove_generation(&self, exec: &Arc<Execution>) {
552        let key = (exec.sandbox_id.clone(), exec.id.clone());
553        let mut inner = self.inner.lock().unwrap();
554        if inner.live.get(&key).is_some_and(|e| Arc::ptr_eq(e, exec)) {
555            inner.live.remove(&key);
556        }
557    }
558
559    /// Mark every still-running execution of a sandbox as torn down so
560    /// parked attach/wait subscribers resolve. Entries stay registered —
561    /// their buffered output remains readable until the per-execution
562    /// retention GC drops them.
563    fn interrupt_computer(&self, sandbox_id: &str) {
564        let executions: Vec<Arc<Execution>> = {
565            let inner = self.inner.lock().unwrap();
566            inner
567                .live
568                .iter()
569                .filter(|((sid, _), _)| sid == sandbox_id)
570                .map(|(_, exec)| Arc::clone(exec))
571                .collect()
572        };
573        for exec in executions {
574            exec.mark_exited(&Err("sandbox stopped".to_owned()));
575        }
576    }
577}
578
579/// Interrupt executions when their sandbox reaches a terminal state. Covers
580/// every teardown path (stop, remove, TTL expiry, boot failure) without
581/// threading the registry through each of them.
582pub(super) fn spawn_teardown_purge(
583    registry: Arc<ExecutionRegistry>,
584    mut events: broadcast::Receiver<ComputerEvent>,
585) {
586    tokio::spawn(async move {
587        loop {
588            match events.recv().await {
589                Ok(ev) if ev.is_terminal() => {
590                    registry.interrupt_computer(&ev.sandbox_id);
591                }
592                Ok(_) | Err(broadcast::error::RecvError::Lagged(_)) => {}
593                Err(broadcast::error::RecvError::Closed) => break,
594            }
595        }
596    });
597}
598
599/// Drain a session's output into the execution's buffers until it exits,
600/// then keep the exited record around for late attach/wait before GC.
601async fn run_session(
602    exec: Arc<Execution>,
603    mut output_rx: mpsc::Receiver<Result<OutputChunk>>,
604    slot: Arc<dyn WorkloadSlot>,
605    registry: Arc<ExecutionRegistry>,
606) {
607    let outcome: std::result::Result<ExitStatus, String> = loop {
608        match output_rx.recv().await {
609            Some(Ok(OutputChunk::Stdout(data))) => {
610                exec.state.lock().unwrap().stdout.append(&data);
611                exec.bump();
612            }
613            Some(Ok(OutputChunk::Stderr(data))) => {
614                exec.state.lock().unwrap().stderr.append(&data);
615                exec.bump();
616            }
617            Some(Ok(OutputChunk::Exit(status))) => break Ok(status),
618            Some(Err(e)) => break Err(e.to_string()),
619            None => break Err("execution session closed before exit".to_owned()),
620        }
621    };
622
623    match &outcome {
624        Ok(status) => slot.exited(WorkloadOutcome::Exited(*status)),
625        // The guest vm-agent kills the workload when its host connection
626        // drops, so the computer is idle either way — there is just no status
627        // to report, only why.
628        Err(e) => {
629            warn!(sandbox_id = %exec.sandbox_id, execution_id = %exec.id, error = %e,
630                  "execution session broke before exit");
631            slot.exited(WorkloadOutcome::Broke(e.clone()));
632        }
633    }
634    exec.mark_exited(&outcome);
635
636    tokio::time::sleep(EXITED_RETENTION).await;
637    registry.remove_generation(&exec);
638}
639
640impl ComputerManager {
641    /// Start an execution inside a `Ready` sandbox.
642    ///
643    /// The execution id (caller-supplied or generated) addresses the process
644    /// afterwards: attach, stdin writes, signals, resize, and wait all go
645    /// through it. Retrying a start with the same id and command returns the
646    /// existing execution instead of dispatching a second process.
647    pub async fn start_execution(
648        &self,
649        sandbox_id: &ComputerId,
650        spec: ExecutionSpec,
651    ) -> Result<ExecutionSnapshot> {
652        let id = match &spec.id {
653            Some(id) if !id.is_empty() => {
654                super::validate_id("execution id", id)?;
655                id.clone()
656            }
657            _ => Uuid::new_v4().to_string(),
658        };
659        if spec.cmd.is_empty() {
660            return Err(ComputerError::Config(
661                "execution command must not be empty".into(),
662            ));
663        }
664
665        // Resolve the reservation, awaiting an identical in-flight start
666        // rather than failing it. One await is the normal case (the original
667        // commits or unwinds); the bound only stops a pathological interleave
668        // of repeated failing starts from spinning here forever.
669        let mut reservation = None;
670        for _ in 0..MAX_START_RESERVE_ROUNDS {
671            match self
672                .executions
673                .reserve((sandbox_id.clone(), id.clone()), &spec.cmd)?
674            {
675                Reserve::Existing(snapshot) => return Ok(snapshot),
676                Reserve::Slot(guard) => {
677                    reservation = Some(guard);
678                    break;
679                }
680                Reserve::AwaitPending(mut done) => {
681                    let _ = done.changed().await;
682                }
683            }
684        }
685        let Some(reservation) = reservation else {
686            return Err(ComputerError::AlreadyExists(format!(
687                "execution '{id}' (concurrent starts did not settle)"
688            )));
689        };
690
691        let agent = self.require_ready_agent(sandbox_id)?;
692        let start = StartCommand {
693            cmd: spec.cmd.clone(),
694            env: spec.env.clone(),
695            working_dir: spec.working_dir.clone(),
696            user: spec.user.clone(),
697            tty: spec.tty,
698            tty_width: spec.tty_size.map_or(80, |(w, _)| w),
699            tty_height: spec.tty_size.map_or(24, |(_, h)| h),
700            timeout_seconds: spec.timeout_seconds,
701        };
702
703        // Claim the computer (Ready → Running) before dispatching, so a
704        // losing racer never launches a process (same discipline as
705        // workload.rs).
706        let slot = self.workload_slot(sandbox_id)?;
707        slot.claim(WorkloadClaim::Api).await?;
708
709        let (input_tx, output_rx) = match agent.exec(start).await {
710            Ok(pair) => pair,
711            Err(e) => {
712                slot.release();
713                return Err(e);
714            }
715        };
716        if !spec.stdin && !spec.tty {
717            // Run semantics: the process starts with stdin at EOF. TTY
718            // executions are exempt — a PTY has no out-of-band EOF, and the
719            // vm-agent's TTY session stops reading input frames (signals,
720            // resizes) once it sees one.
721            let _ = input_tx.send(ExecInputMsg::Eof).await;
722        }
723
724        let exec = Arc::new(Execution::new(id, sandbox_id.clone(), &spec, input_tx));
725        reservation.commit(&exec);
726
727        tokio::spawn(run_session(
728            Arc::clone(&exec),
729            output_rx,
730            slot,
731            Arc::clone(&self.executions),
732        ));
733
734        Ok(exec.snapshot())
735    }
736
737    /// Attach to an execution's output from the given per-channel offsets.
738    ///
739    /// The receiver yields buffered output first (starting at the earliest
740    /// retained byte if retention already dropped the requested offset), then
741    /// live output, and closes once the execution has exited and both
742    /// channels are drained. Read the final state with
743    /// [`ComputerManager::wait_execution`] afterwards.
744    pub fn attach_execution(
745        &self,
746        sandbox_id: &str,
747        execution_id: &str,
748        stdout_offset: u64,
749        stderr_offset: u64,
750    ) -> Result<(ExecutionSnapshot, mpsc::Receiver<ExecutionOutput>)> {
751        let exec = self.executions.get(sandbox_id, execution_id)?;
752        Ok(exec.attach(stdout_offset, stderr_offset))
753    }
754
755    /// Write stdin bytes at an absolute offset (offset-idempotent).
756    ///
757    /// Bytes below the accepted count are deduplicated so a retried write
758    /// never double-feeds the process; an offset beyond the accepted count is
759    /// a gap and is rejected with the resume point. `eof` closes stdin after
760    /// the write — rejected for TTY executions, where EOF is a byte (Ctrl-D)
761    /// the client sends through the stream itself.
762    pub async fn write_stdin(
763        &self,
764        sandbox_id: &str,
765        execution_id: &str,
766        offset: u64,
767        data: &[u8],
768        eof: bool,
769    ) -> Result<StdinState> {
770        self.executions
771            .get(sandbox_id, execution_id)?
772            .write_stdin(offset, data, eof)
773            .await
774    }
775
776    /// Current stdin acceptance state — the resume point after a lost write.
777    pub fn stdin_status(&self, sandbox_id: &str, execution_id: &str) -> Result<StdinState> {
778        Ok(self
779            .executions
780            .get(sandbox_id, execution_id)?
781            .snapshot()
782            .stdin)
783    }
784
785    /// Deliver a POSIX signal to a running execution's process group.
786    pub async fn signal_execution(
787        &self,
788        sandbox_id: &str,
789        execution_id: &str,
790        signal: i32,
791    ) -> Result<()> {
792        self.executions
793            .get(sandbox_id, execution_id)?
794            .signal(signal)
795            .await
796    }
797
798    /// Resize a running TTY execution's terminal.
799    pub async fn resize_execution(
800        &self,
801        sandbox_id: &str,
802        execution_id: &str,
803        width: u16,
804        height: u16,
805    ) -> Result<()> {
806        self.executions
807            .get(sandbox_id, execution_id)?
808            .resize(width, height)
809            .await
810    }
811
812    /// Wait until an execution exits or `timeout` elapses, then return its
813    /// state. A zero timeout polls the current state.
814    pub async fn wait_execution(
815        &self,
816        sandbox_id: &str,
817        execution_id: &str,
818        timeout: Duration,
819    ) -> Result<ExecutionSnapshot> {
820        Ok(self
821            .executions
822            .get(sandbox_id, execution_id)?
823            .wait(timeout)
824            .await)
825    }
826
827    /// List a sandbox's retained executions, running and exited, ordered by
828    /// start time (then id, for a stable order under equal timestamps).
829    ///
830    /// The sandbox must exist; an unknown id is `NotFound` rather than an
831    /// empty list, so a caller can tell "no executions" from a typo.
832    pub fn list_executions(&self, sandbox_id: &ComputerId) -> Result<Vec<ExecutionSnapshot>> {
833        self.computer(sandbox_id)?;
834        let mut snapshots = self.executions.list(sandbox_id);
835        snapshots.sort_by(|a, b| {
836            a.started_at
837                .cmp(&b.started_at)
838                .then_with(|| a.id.cmp(&b.id))
839        });
840        Ok(snapshots)
841    }
842
843    /// Wait until something inside an alive sandbox listens on TCP `port`,
844    /// or fail with [`ComputerError::DeadlineExceeded`] once `timeout` elapses.
845    ///
846    /// The vm-agent watches the guest's own listen table in-process — no
847    /// connect probes that would perturb the workload with spurious
848    /// accepted connections.
849    pub async fn wait_computer_port(
850        &self,
851        id: &ComputerId,
852        port: u16,
853        timeout: Duration,
854    ) -> Result<()> {
855        let agent = self.require_alive_agent(id)?;
856        match agent.wait_for_port(port, timeout).await? {
857            crate::agent::PortWait::Listening => Ok(()),
858            crate::agent::PortWait::Deadline => Err(ComputerError::DeadlineExceeded(format!(
859                "no listener on port {port} in sandbox '{id}' within {}s",
860                timeout.as_secs()
861            ))),
862        }
863    }
864}
865
866/// Forward buffered + live output to one attach consumer.
867async fn pump_attach(
868    exec: Arc<Execution>,
869    mut stdout_cursor: u64,
870    mut stderr_cursor: u64,
871    tx: mpsc::Sender<ExecutionOutput>,
872) {
873    let mut version = exec.version.subscribe();
874    loop {
875        // Mark the version seen before reading, so an append between the read
876        // and the wait below leaves a pending change notification.
877        version.borrow_and_update();
878        loop {
879            let (chunk, exited) = {
880                let st = exec.state.lock().unwrap();
881                let (start, data) = st.stdout.read_from(stdout_cursor, ATTACH_CHUNK);
882                if data.is_empty() {
883                    let (start, data) = st.stderr.read_from(stderr_cursor, ATTACH_CHUNK);
884                    if data.is_empty() {
885                        (None, st.exited_at.is_some())
886                    } else {
887                        stderr_cursor = start + data.len() as u64;
888                        (
889                            Some(ExecutionOutput {
890                                channel: ExecutionChannel::Stderr,
891                                offset: start,
892                                data,
893                            }),
894                            false,
895                        )
896                    }
897                } else {
898                    stdout_cursor = start + data.len() as u64;
899                    (
900                        Some(ExecutionOutput {
901                            channel: ExecutionChannel::Stdout,
902                            offset: start,
903                            data,
904                        }),
905                        false,
906                    )
907                }
908            };
909            match chunk {
910                Some(chunk) => {
911                    if tx.send(chunk).await.is_err() {
912                        return; // consumer gone
913                    }
914                }
915                None if exited => return, // drained past the exit
916                None => break,            // caught up; wait for more
917            }
918        }
919        // The sender lives inside `exec`, which we hold — never closed.
920        let _ = version.changed().await;
921    }
922}
923
924#[cfg(test)]
925mod tests {
926    use super::*;
927
928    fn stdin_spec() -> ExecutionSpec {
929        ExecutionSpec {
930            cmd: vec!["cat".into()],
931            stdin: true,
932            ..ExecutionSpec::default()
933        }
934    }
935
936    fn test_execution(spec: &ExecutionSpec) -> (Arc<Execution>, mpsc::Receiver<ExecInputMsg>) {
937        let (tx, rx) = mpsc::channel(8);
938        let exec = Arc::new(Execution::new(
939            "exec-1".into(),
940            "sandbox-1".into(),
941            spec,
942            tx,
943        ));
944        (exec, rx)
945    }
946
947    #[test]
948    fn channel_buffer_addresses_bytes_absolutely() {
949        let mut buf = ChannelBuffer::default();
950        buf.append(b"hello ");
951        buf.append(b"world");
952        assert_eq!(buf.end(), 11);
953
954        let (start, data) = buf.read_from(0, 1024);
955        assert_eq!((start, data.as_slice()), (0, b"hello world".as_slice()));
956
957        let (start, data) = buf.read_from(6, 3);
958        assert_eq!((start, data.as_slice()), (6, b"wor".as_slice()));
959
960        // Reads at or past the end are empty at the requested offset.
961        let (start, data) = buf.read_from(11, 1024);
962        assert_eq!((start, data.len()), (11, 0));
963    }
964
965    #[test]
966    fn channel_buffer_trims_to_retention_and_advances_base() {
967        let mut buf = ChannelBuffer::default();
968        let chunk = vec![7u8; CHANNEL_RETENTION];
969        buf.append(&chunk);
970        buf.append(b"tail");
971
972        assert_eq!(buf.base, 4);
973        assert_eq!(buf.end(), CHANNEL_RETENTION as u64 + 4);
974
975        // A read below the base is clamped to the earliest retained byte —
976        // the returned offset exposes the gap.
977        let (start, data) = buf.read_from(0, 8);
978        assert_eq!(start, 4);
979        assert_eq!(data.len(), 8);
980    }
981
982    #[tokio::test]
983    async fn write_stdin_dedupes_retries_and_rejects_gaps() {
984        let (exec, mut rx) = test_execution(&stdin_spec());
985
986        let st = exec.write_stdin(0, b"hello", false).await.unwrap();
987        assert_eq!(st.bytes_written, 5);
988
989        // Full retry of the same write: accepted, nothing re-forwarded.
990        let st = exec.write_stdin(0, b"hello", false).await.unwrap();
991        assert_eq!(st.bytes_written, 5);
992
993        // Overlapping retry: only the unseen suffix is forwarded.
994        let st = exec.write_stdin(3, b"lo world", false).await.unwrap();
995        assert_eq!(st.bytes_written, 11);
996
997        // A gap is rejected with the resume point.
998        let err = exec.write_stdin(99, b"x", false).await.unwrap_err();
999        assert!(matches!(err, ComputerError::StdinGap { accepted: 11, .. }));
1000
1001        // Exactly the deduplicated byte stream reached the process.
1002        let mut forwarded = Vec::new();
1003        while let Ok(msg) = rx.try_recv() {
1004            match msg {
1005                ExecInputMsg::Stdin(data) => forwarded.extend_from_slice(&data),
1006                other => panic!("unexpected input message: {other:?}"),
1007            }
1008        }
1009        assert_eq!(forwarded, b"hello world");
1010    }
1011
1012    #[tokio::test]
1013    async fn write_stdin_eof_closes_and_tty_rejects_eof() {
1014        let (exec, mut rx) = test_execution(&stdin_spec());
1015
1016        let st = exec.write_stdin(0, b"in", true).await.unwrap();
1017        assert!(st.closed);
1018        assert!(matches!(rx.recv().await, Some(ExecInputMsg::Stdin(_))));
1019        assert!(matches!(rx.recv().await, Some(ExecInputMsg::Eof)));
1020
1021        // Further data after EOF is rejected; a pure duplicate is not.
1022        let err = exec.write_stdin(2, b"more", false).await.unwrap_err();
1023        assert!(matches!(err, ComputerError::Config(_)));
1024        let st = exec.write_stdin(0, b"in", false).await.unwrap();
1025        assert_eq!(st.bytes_written, 2);
1026
1027        // TTY executions cannot close stdin via eof.
1028        let tty_spec = ExecutionSpec {
1029            tty: true,
1030            ..stdin_spec()
1031        };
1032        let (tty_exec, _tty_rx) = test_execution(&tty_spec);
1033        let err = tty_exec.write_stdin(0, b"", true).await.unwrap_err();
1034        assert!(matches!(err, ComputerError::Config(_)));
1035    }
1036
1037    #[tokio::test]
1038    async fn attach_replays_from_offsets_and_closes_after_exit() {
1039        let (exec, _rx) = test_execution(&stdin_spec());
1040
1041        {
1042            let mut st = exec.state.lock().unwrap();
1043            st.stdout.append(b"out-data");
1044            st.stderr.append(b"err");
1045        }
1046        exec.bump();
1047
1048        let (snapshot, mut out) = exec.attach(2, 0);
1049        assert!(snapshot.is_running());
1050
1051        let first = out.recv().await.unwrap();
1052        assert_eq!(first.channel, ExecutionChannel::Stdout);
1053        assert_eq!(first.offset, 2);
1054        assert_eq!(first.data, b"t-data");
1055        let second = out.recv().await.unwrap();
1056        assert_eq!(second.channel, ExecutionChannel::Stderr);
1057        assert_eq!(second.offset, 0);
1058        assert_eq!(second.data, b"err");
1059
1060        // Live append after the replay is delivered too.
1061        exec.state.lock().unwrap().stdout.append(b"+live");
1062        exec.bump();
1063        let live = out.recv().await.unwrap();
1064        assert_eq!(live.offset, 8);
1065        assert_eq!(live.data, b"+live");
1066
1067        // Exit closes the stream once drained.
1068        exec.mark_exited(&Ok(ExitStatus::Code(0)));
1069        assert!(out.recv().await.is_none());
1070
1071        let done = exec.wait(Duration::ZERO).await;
1072        assert_eq!(done.exit_status, Some(ExitStatus::Code(0)));
1073    }
1074
1075    #[tokio::test]
1076    async fn wait_times_out_then_resolves_on_exit() {
1077        let (exec, _rx) = test_execution(&stdin_spec());
1078
1079        // Zero timeout polls the running state.
1080        assert!(exec.wait(Duration::ZERO).await.is_running());
1081        // A short timeout returns the still-running state.
1082        assert!(exec.wait(Duration::from_millis(20)).await.is_running());
1083
1084        // A concurrent exit resolves a parked wait.
1085        let exiter = {
1086            let exec = Arc::clone(&exec);
1087            tokio::spawn(async move {
1088                tokio::time::sleep(Duration::from_millis(20)).await;
1089                exec.mark_exited(&Ok(ExitStatus::Signaled(9)));
1090            })
1091        };
1092        let snap = exec.wait(Duration::from_secs(5)).await;
1093        assert_eq!(snap.exit_status, Some(ExitStatus::Signaled(9)));
1094        exiter.await.unwrap();
1095    }
1096
1097    #[tokio::test]
1098    async fn signal_and_resize_validate_state() {
1099        let (exec, mut rx) = test_execution(&stdin_spec());
1100
1101        exec.signal(15).await.unwrap();
1102        assert!(matches!(rx.recv().await, Some(ExecInputMsg::Signal(15))));
1103        assert!(matches!(
1104            exec.signal(0).await.unwrap_err(),
1105            ComputerError::Config(_)
1106        ));
1107        // Non-TTY executions have nothing to resize.
1108        assert!(matches!(
1109            exec.resize(80, 24).await.unwrap_err(),
1110            ComputerError::Config(_)
1111        ));
1112
1113        exec.mark_exited(&Ok(ExitStatus::Code(0)));
1114        assert!(matches!(
1115            exec.signal(9).await.unwrap_err(),
1116            ComputerError::WrongState { .. }
1117        ));
1118    }
1119
1120    #[tokio::test]
1121    async fn interrupt_sandbox_marks_running_executions_but_keeps_them_readable() {
1122        let registry = Arc::new(ExecutionRegistry::default());
1123        let (exec, _rx) = test_execution(&stdin_spec());
1124        registry
1125            .inner
1126            .lock()
1127            .unwrap()
1128            .live
1129            .insert(("sandbox-1".into(), "exec-1".into()), Arc::clone(&exec));
1130
1131        registry.interrupt_computer("sandbox-1");
1132        // Still registered: buffered output stays readable until the
1133        // retention GC, but the execution is resolved as torn down.
1134        assert!(registry.get("sandbox-1", "exec-1").is_ok());
1135        let snap = exec.snapshot();
1136        assert!(!snap.is_running());
1137        assert_eq!(snap.error.as_deref(), Some("sandbox stopped"));
1138
1139        // An already-exited execution keeps its original outcome.
1140        let (done, _rx2) = test_execution(&stdin_spec());
1141        done.mark_exited(&Ok(ExitStatus::Code(3)));
1142        registry
1143            .inner
1144            .lock()
1145            .unwrap()
1146            .live
1147            .insert(("sandbox-1".into(), "exec-2".into()), Arc::clone(&done));
1148        registry.interrupt_computer("sandbox-1");
1149        assert_eq!(done.snapshot().exit_status, Some(ExitStatus::Code(3)));
1150    }
1151
1152    #[tokio::test]
1153    async fn start_reservation_is_idempotent_per_command() {
1154        let registry = Arc::new(ExecutionRegistry::default());
1155        let key = ("s".to_owned(), "e".to_owned());
1156        let cmd = vec!["echo".to_owned()];
1157
1158        // Reserve → commit.
1159        let Reserve::Slot(slot) = registry.reserve(key.clone(), &cmd).unwrap() else {
1160            panic!("expected a fresh slot");
1161        };
1162        let (tx, _rx) = mpsc::channel(1);
1163        let exec = Arc::new(Execution::new(
1164            "e".into(),
1165            "s".into(),
1166            &ExecutionSpec {
1167                cmd: cmd.clone(),
1168                ..ExecutionSpec::default()
1169            },
1170            tx,
1171        ));
1172        slot.commit(&exec);
1173
1174        // Same command → idempotent retry hands back the existing execution.
1175        assert!(matches!(
1176            registry.reserve(key.clone(), &cmd).unwrap(),
1177            Reserve::Existing(_)
1178        ));
1179        // Different command under the same id → collision.
1180        assert!(matches!(
1181            registry.reserve(key, &["other".to_owned()]),
1182            Err(ComputerError::AlreadyExists(_))
1183        ));
1184
1185        // While a start is in flight, a matching retry waits for it instead
1186        // of failing — the idempotency contract holds inside the start
1187        // window, not only once the process is already running. A differing
1188        // command is still a collision.
1189        let key2 = ("s".to_owned(), "e2".to_owned());
1190        {
1191            let Reserve::Slot(_slot) = registry.reserve(key2.clone(), &cmd).unwrap() else {
1192                panic!("expected a fresh slot");
1193            };
1194            assert!(matches!(
1195                registry.reserve(key2.clone(), &cmd).unwrap(),
1196                Reserve::AwaitPending(_)
1197            ));
1198            assert!(matches!(
1199                registry.reserve(key2.clone(), &["other".to_owned()]),
1200                Err(ComputerError::AlreadyExists(_))
1201            ));
1202        }
1203        // A dropped (uncommitted) reservation frees the slot.
1204        assert!(matches!(
1205            registry.reserve(key2, &cmd).unwrap(),
1206            Reserve::Slot(_)
1207        ));
1208    }
1209
1210    #[tokio::test]
1211    async fn a_pending_start_wakes_matching_retries_with_the_committed_execution() {
1212        let registry = Arc::new(ExecutionRegistry::default());
1213        let key = ("s".to_owned(), "e".to_owned());
1214        let cmd = vec!["sleep".to_owned()];
1215
1216        let Reserve::Slot(slot) = registry.reserve(key.clone(), &cmd).unwrap() else {
1217            panic!("expected a fresh slot");
1218        };
1219        let Reserve::AwaitPending(mut waiter) = registry.reserve(key.clone(), &cmd).unwrap() else {
1220            panic!("a matching retry must await the in-flight start");
1221        };
1222
1223        // Commit the original; the waiter must then see the live execution.
1224        let (tx, _rx) = mpsc::channel(1);
1225        let exec = Arc::new(Execution::new(
1226            "e".into(),
1227            "s".into(),
1228            &ExecutionSpec {
1229                cmd: cmd.clone(),
1230                ..ExecutionSpec::default()
1231            },
1232            tx,
1233        ));
1234        slot.commit(&exec);
1235
1236        waiter.changed().await.expect("pending start signals");
1237        assert!(matches!(
1238            registry.reserve(key.clone(), &cmd).unwrap(),
1239            Reserve::Existing(_)
1240        ));
1241
1242        // A start that unwinds instead wakes waiters too, and the key is
1243        // free for the retry to claim.
1244        let key2 = ("s".to_owned(), "e2".to_owned());
1245        let Reserve::Slot(failed) = registry.reserve(key2.clone(), &cmd).unwrap() else {
1246            panic!("expected a fresh slot");
1247        };
1248        let Reserve::AwaitPending(mut waiter) = registry.reserve(key2.clone(), &cmd).unwrap()
1249        else {
1250            panic!("a matching retry must await the in-flight start");
1251        };
1252        drop(failed);
1253        waiter.changed().await.expect("unwound start signals");
1254        assert!(matches!(
1255            registry.reserve(key2, &cmd).unwrap(),
1256            Reserve::Slot(_)
1257        ));
1258    }
1259}