perf(agent): bound the generation-6 copy path

Cargo.toml

@@ -88,8 +88,9 @@ microsandbox-protocol = { version = "=0.6.9", path = "crates/protocol" }
 microsandbox-runtime = { version = "=0.6.9", path = "crates/runtime", default-features = false }
 microsandbox-utils = { version = "=0.6.9", path = "crates/utils" }
 microsandbox-vsock = { version = "=0.6.9", path = "crates/vsock" }
-msb_krun = "=0.1.31"
-msb_krun_utils = "=0.1.31"
+# Keep CI on the exact prerequisite revision until its console API is released to crates.io.
+msb_krun = { git = "https://github.com/superradcompany/libkrun", rev = "ff087a1a0add4006f7ad4d753316cde1015a28b2" }
+msb_krun_utils = { git = "https://github.com/superradcompany/libkrun", rev = "ff087a1a0add4006f7ad4d753316cde1015a28b2" }
 test-macros = { path = "crates/testing/macros" }
 test-utils = { path = "crates/testing/utils" }
 

crates/agentd/Cargo.toml

@@ -16,6 +16,7 @@ path = "lib/lib.rs"
 
 [dependencies]
 base64.workspace = true
+bytes.workspace = true
 chrono.workspace = true
 ciborium.workspace = true
 libc.workspace = true

crates/agentd/lib/agent.rs

@@ -8,9 +8,10 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Instant;
 
+use bytes::BytesMut;
 use chrono::Utc;
 use tokio::io::unix::AsyncFd;
-use tokio::sync::{mpsc, watch};
+use tokio::sync::watch;
 use tokio::time::{self, Duration};
 
 use microsandbox_protocol::HANDOFF_POWEROFF_TIMEOUT;
@@ -34,7 +35,8 @@ use crate::fs::{FsReadSession, FsState, FsStreamSession, FsWriteSession};
 use crate::process::ProcessManager;
 use crate::serial::AGENT_PORT_NAME;
 use crate::session::{
-    ExecSession, RawActivity, RawSessionCompletion, SessionOutput, resolve_default_user,
+    ExecSession, RawActivity, RawSessionCompletion, SessionOutput, SessionOutputSender,
+    resolve_default_user,
 };
 use crate::tcp::TcpSession;
 use crate::{clock, fs, handoff, heartbeat, serial};
@@ -155,13 +157,13 @@ pub async fn run(
 
     // Buffer for serial reads.
     let mut read_buf = vec![0u8; SERIAL_READ_BUF_SIZE];
-    let mut serial_in_buf = Vec::new();
+    let mut serial_in_buf = BytesMut::new();
     let mut serial_out_buf = Vec::new();
 
     let mut state = AgentState::default();
 
     // Channel for session output events.
-    let (session_tx, mut session_rx) = mpsc::unbounded_channel::<(u32, SessionOutput)>();
+    let (session_tx, mut session_rx) = SessionOutputSender::channel();
 
     // Heartbeat/activity state.
     let mut activity = ActivityTracker::new();
@@ -238,19 +240,9 @@ pub async fn run(
                             // message-level failures are reported on the same
                             // correlation ID with `core.error`; unrecoverable
                             // frame-level failures still close the agent loop.
-                            while let Some(frame) = codec::try_decode_raw_from_buf(&mut serial_in_buf)
+                            while let Some(msg) = codec::try_decode_from_bytes(&mut serial_in_buf)
                                 .map_err(|e| AgentdError::ExecSession(format!("decode frame: {e}")))?
                             {
-                                let id = frame.id;
-                                let msg = match codec::raw_frame_to_message(frame) {
-                                    Ok(msg) => msg,
-                                    Err(e) => {
-                                        return Err(AgentdError::ExecSession(format!(
-                                            "decode message for id {id}: {e}"
-                                        )));
-                                    }
-                                };
-
                                 if msg.flags != msg.t.flags() {
                                     let out_before = serial_out_buf.len();
                                     encode_core_error_if_supported(
@@ -316,8 +308,9 @@ pub async fn run(
             }
 
             // Receive output events from session reader tasks.
-            Some((id, output)) = session_rx.recv() => {
-                match output {
+            Some(envelope) = session_rx.recv() => {
+                let id = envelope.id;
+                match envelope.output {
                     SessionOutput::Stdout(data) => {
                         let len = data.len();
                         let msg = Message::with_payload(MessageType::ExecStdout, id, &ExecStdout { data })
@@ -352,8 +345,13 @@ pub async fn run(
                             &mut state.read_sessions,
                             &mut state.tcp_sessions,
                         );
-                        // Pre-encoded frame — write directly to output buffer.
-                        serial_out_buf.extend_from_slice(&output.frame);
+                        // The producer already owns an encoded frame. Write from that allocation
+                        // directly so multi-megabyte FS/TCP frames are not copied into a second
+                        // serial staging buffer.
+                        if !serial_out_buf.is_empty() {
+                            flush_write_buf(&async_port, &mut serial_out_buf).await?;
+                        }
+                        write_all_async_fd(&async_port, &output.frame).await?;
                     }
                 }
                 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
@@ -413,7 +411,7 @@ async fn handle_message(
     msg: Message,
     state: &mut AgentState,
     activity: &mut ActivityTracker,
-    session_tx: &mpsc::UnboundedSender<(u32, SessionOutput)>,
+    session_tx: &SessionOutputSender,
     out_buf: &mut Vec<u8>,
     config: &AgentdConfig,
 ) -> AgentdResult<()> {
@@ -1171,11 +1169,22 @@ fn write_all_to_fd(fd: i32, mut buf: &[u8], deadline: Instant) -> AgentdResult<(
 
 /// Flushes the write buffer to the async fd.
 async fn flush_write_buf(fd: &AsyncFd<std::fs::File>, buf: &mut Vec<u8>) -> AgentdResult<()> {
-    while !buf.is_empty() {
+    write_all_async_fd(fd, buf).await?;
+    buf.clear();
+    Ok(())
+}
+
+/// Write an immutable region to the nonblocking serial descriptor with cursor advancement.
+async fn write_all_async_fd(fd: &AsyncFd<std::fs::File>, buf: &[u8]) -> AgentdResult<()> {
+    let mut written = 0;
+    while written < buf.len() {
         let mut guard = fd.writable().await?;
-        match guard.try_io(|inner| write_to_fd(inner.get_ref().as_raw_fd(), buf)) {
+        match guard.try_io(|inner| write_to_fd(inner.get_ref().as_raw_fd(), &buf[written..])) {
             Ok(Ok(n)) => {
-                buf.drain(..n);
+                if n == 0 {
+                    return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into());
+                }
+                written += n;
             }
             Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => continue,
             Ok(Err(e)) => return Err(e.into()),

crates/agentd/lib/fs.rs

@@ -19,10 +19,12 @@ use microsandbox_protocol::fs::{
 };
 use microsandbox_protocol::message::{Message, MessageType};
 use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
-use tokio::sync::{Mutex, mpsc};
+use tokio::sync::Mutex;
 use tokio::task::JoinHandle;
 
-use crate::session::{RawActivity, RawSessionCompletion, RawSessionOutput, SessionOutput};
+use crate::session::{
+    RawActivity, RawSessionCompletion, RawSessionOutput, SessionOutput, SessionOutputSender,
+};
 
 //--------------------------------------------------------------------------------------------------
 // Constants
@@ -302,7 +304,7 @@ pub async fn handle_fs_request(
     req: FsRequest,
     state: &mut FsState,
     out_buf: &mut Vec<u8>,
-    session_tx: &mpsc::UnboundedSender<(u32, SessionOutput)>,
+    session_tx: &SessionOutputSender,
 ) -> Result<Option<FsStreamSession>, String> {
     match req.op {
         FsOp::RealPath { path } => {
@@ -749,11 +751,11 @@ async fn handle_read_stream(
     file: Arc<Mutex<tokio::fs::File>>,
     offset: u64,
     len: Option<u64>,
-    tx: &mpsc::UnboundedSender<(u32, SessionOutput)>,
+    tx: &SessionOutputSender,
 ) {
     let mut file = file.lock().await;
     if let Err(e) = file.seek(std::io::SeekFrom::Start(offset)).await {
-        send_raw_response(id, false, Some(format!("seek: {e}")), None, tx);
+        send_raw_response(id, false, Some(format!("seek: {e}")), None, tx).await;
         return;
     }
 
@@ -762,6 +764,11 @@ async fn handle_read_stream(
     let mut buf = Vec::new();
 
     loop {
+        // Reserve before the file read and CBOR materialization so concurrent streams cannot
+        // each create an uncharged maximum-sized frame while aggregate output is saturated.
+        let Some(permit) = tx.reserve(codec::MAX_FRAME_SIZE as usize + 4).await else {
+            return;
+        };
         let read_len = match remaining {
             Some(0) => break,
             Some(n) => chunk.len().min(n as usize),
@@ -780,7 +787,8 @@ async fn handle_read_stream(
                 let msg = match Message::with_payload(MessageType::FsData, id, &data) {
                     Ok(msg) => msg,
                     Err(e) => {
-                        send_raw_response(id, false, Some(format!("encode chunk: {e}")), None, tx);
+                        send_raw_response(id, false, Some(format!("encode chunk: {e}")), None, tx)
+                            .await;
                         return;
                     }
                 };
@@ -792,22 +800,27 @@ async fn handle_read_stream(
                         Some(format!("encode chunk frame: {e}")),
                         None,
                         tx,
-                    );
+                    )
+                    .await;
                     return;
                 }
-                let output = RawSessionOutput::new(buf.clone(), RawActivity::fs_bytes(n), None);
-                if tx.send((id, SessionOutput::Raw(output))).is_err() {
+                let output =
+                    RawSessionOutput::new(std::mem::take(&mut buf), RawActivity::fs_bytes(n), None);
+                if !tx
+                    .send_reserved(id, SessionOutput::Raw(output), permit)
+                    .await
+                {
                     return;
                 }
             }
             Err(e) => {
-                send_raw_response(id, false, Some(format!("read: {e}")), None, tx);
+                send_raw_response(id, false, Some(format!("read: {e}")), None, tx).await;
                 return;
             }
         }
     }
 
-    send_raw_response(id, true, None, None, tx);
+    send_raw_response(id, true, None, None, tx).await;
 }
 
 //--------------------------------------------------------------------------------------------------
@@ -978,12 +991,12 @@ fn encode_response(id: u32, resp: FsResponse, out_buf: &mut Vec<u8>) -> Result<(
     Ok(())
 }
 
-fn send_raw_response(
+async fn send_raw_response(
     id: u32,
     ok: bool,
     error: Option<String>,
     data: Option<FsResponseData>,
-    tx: &mpsc::UnboundedSender<(u32, SessionOutput)>,
+    tx: &SessionOutputSender,
 ) {
     let resp = FsResponse { ok, error, data };
     match Message::with_payload(MessageType::FsResponse, id, &resp) {
@@ -996,7 +1009,7 @@ fn send_raw_response(
                         RawActivity::guest_message(),
                         Some(RawSessionCompletion::FsRead),
                     );
-                    let _ = tx.send((id, SessionOutput::Raw(output)));
+                    let _ = tx.send(id, SessionOutput::Raw(output)).await;
                 }
                 Err(e) => {
                     eprintln!("failed to encode fs response frame for {id}: {e}");