Fix/workload logs and metrics

agent/csfx-guest-init/src/main.rs

@@ -136,7 +136,8 @@ async fn run() -> Result<()> {
         stage = "start_container",
         "Starting container via libcontainer"
     );
-    let (stdout_read, stderr_read, container_pid) = start_container(&mmds_data.env).await?;
+    let (stdout_read, stderr_read, container_pid, _stdout_write, _stderr_write) =
+        start_container(&mmds_data.env).await?;
     debug!(
         stage = "container_started",
         container_pid = container_pid.as_raw(),
@@ -200,7 +201,7 @@ fn apply_extra_env(extra_env: &std::collections::HashMap<String, String>) -> Res
 
 async fn start_container(
     extra_env: &std::collections::HashMap<String, String>,
-) -> Result<(AsyncFd<OwnedFd>, AsyncFd<OwnedFd>, Pid)> {
+) -> Result<(AsyncFd<OwnedFd>, AsyncFd<OwnedFd>, Pid, OwnedFd, OwnedFd)> {
     apply_extra_env(extra_env).context("Failed to apply mmds env to container config")?;
     write_container_etc_hosts();
 
@@ -212,12 +213,19 @@ async fn start_container(
     set_nonblocking(&stdout_read).context("Failed to set stdout pipe non-blocking")?;
     set_nonblocking(&stderr_read).context("Failed to set stderr pipe non-blocking")?;
 
+    let stdout_write_container = stdout_write
+        .try_clone()
+        .context("Failed to clone stdout pipe for container")?;
+    let stderr_write_container = stderr_write
+        .try_clone()
+        .context("Failed to clone stderr pipe for container")?;
+
     let container_pid = tokio::task::spawn_blocking(move || -> Result<Pid> {
         let mut container = ContainerBuilder::new(CONTAINER_ID.to_string(), SyscallType::default())
             .with_root_path(CONTAINER_STATE_ROOT)
             .context("Invalid container state root path")?
-            .with_stdout(stdout_write)
-            .with_stderr(stderr_write)
+            .with_stdout(stdout_write_container)
+            .with_stderr(stderr_write_container)
             .as_init(CONTAINER_BUNDLE_PATH)
             .with_systemd(false)
             .with_detach(true)
@@ -238,6 +246,8 @@ async fn start_container(
         AsyncFd::new(stdout_read).context("Failed to register stdout pipe")?,
         AsyncFd::new(stderr_read).context("Failed to register stderr pipe")?,
         container_pid,
+        stdout_write,
+        stderr_write,
     ))
 }
 

agent/src/firecracker/runtime.rs

@@ -606,16 +606,25 @@ async fn read_cpu_usage_percent(handle: &VmHandle) -> Option<f64> {
 }
 
 async fn read_memory_usage_bytes(cgroup_path: &Path) -> Option<i64> {
-    let content = tokio::fs::read_to_string(cgroup_path.join("memory.current"))
-        .await
-        .ok()?;
-    content.trim().parse::<i64>().ok()
+    let path = cgroup_path.join("memory.current");
+    match tokio::fs::read_to_string(&path).await {
+        Ok(content) => content.trim().parse::<i64>().ok(),
+        Err(e) => {
+            warn!(path = %path.display(), error = %e, "failed to read workload memory cgroup");
+            None
+        }
+    }
 }
 
 async fn read_cgroup_u64(cgroup_path: &Path, file: &str, key: &str) -> Option<u64> {
-    let content = tokio::fs::read_to_string(cgroup_path.join(file))
-        .await
-        .ok()?;
+    let path = cgroup_path.join(file);
+    let content = match tokio::fs::read_to_string(&path).await {
+        Ok(content) => content,
+        Err(e) => {
+            warn!(path = %path.display(), error = %e, "failed to read workload cgroup file");
+            return None;
+        }
+    };
 
     content.lines().find_map(|line| {
         let (line_key, value) = line.split_once(' ')?;
@@ -900,7 +909,10 @@ fn create_metrics_fifo(path: &Path) -> Result<()> {
 }
 
 fn cgroup_path(workload_id: &str) -> PathBuf {
-    Path::new(CGROUP_ROOT).join(CGROUP_PARENT).join(workload_id)
+    let unit_name = jailer_unit_name(&jailer_short_id(workload_id));
+    Path::new(CGROUP_ROOT)
+        .join("system.slice")
+        .join(format!("{}.service", unit_name))
 }
 
 async fn apply_port_dnat(workload_id: &str, guest_ip: &str, spec: &WorkloadSpec) -> Result<()> {

agent/src/server.rs

@@ -3,6 +3,7 @@ use axum::{
     extract::ws::{Message, WebSocket, WebSocketUpgrade},
     extract::{ConnectInfo, Path, State},
     http::{HeaderMap, StatusCode},
+    response::IntoResponse,
     routing::{get, post},
     Json, Router,
 };
@@ -158,7 +159,7 @@ async fn logs_handler(
     Path(workload_id): Path<String>,
     ConnectInfo(addr): ConnectInfo<SocketAddr>,
     headers: HeaderMap,
-) -> Result<axum::body::Body, (StatusCode, String)> {
+) -> Result<impl IntoResponse, (StatusCode, String)> {
     info!(workload_id = %workload_id, source = %addr, "log stream request received");
 
     if !is_internal_source(&addr) {
@@ -186,8 +187,17 @@ async fn logs_handler(
         })?;
 
     info!(workload_id = %workload_id, container_id = %container_id, "opening log stream to guest");
-    let stream = state.firecracker.logs(&container_id);
-    Ok(axum::body::Body::from_stream(stream))
+    let stream =
+        futures_util::stream::once(async { Ok::<_, std::io::Error>(axum::body::Bytes::new()) })
+            .chain(state.firecracker.logs(&container_id));
+
+    Ok((
+        [(
+            axum::http::header::CONTENT_TYPE,
+            "text/plain; charset=utf-8",
+        )],
+        axum::body::Body::from_stream(stream),
+    ))
 }
 
 async fn exec_handler(

control-plane/api-gateway/src/routes/agent_proxy.rs

@@ -199,8 +199,17 @@ pub async fn stream_workload_logs(
         ));
     }
 
-    let stream = resp.bytes_stream();
-    Ok(Body::from_stream(stream))
+    let stream =
+        futures_util::stream::once(async { Ok::<_, reqwest::Error>(axum::body::Bytes::new()) })
+            .chain(resp.bytes_stream());
+
+    Ok((
+        [(
+            axum::http::header::CONTENT_TYPE,
+            "text/plain; charset=utf-8",
+        )],
+        Body::from_stream(stream),
+    ))
 }
 
 pub async fn issue_exec_ticket(