From 136e6624cf07c391e467b45dda719fe68a911858 Mon Sep 17 00:00:00 2001 From: Xuanlin_Mi <179552030+Primary33@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:37:34 +0800 Subject: [PATCH] feat: attach image drives to warm starts Accept attachedDrives on POST /v2/sandboxes and resolve them as cold starts do, instead of backing image mounts with sandbox-scoped volumes. Drives added to a snapshot launch take reserved Firecracker slots ahead of volumes. Captures commit them to the manifest while they keep their slots, so relaunch, fork, and pause/resume need no catalog records. Snapshots with volumes or attached drives reject explicit attachedDrives, including empty arrays. aenv start --image-mount and --image-mount-ro now build attachedDrives, and warm starts use v2. Also mount read-only volumes on cold start through envd. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/aenv/src/client/sandboxes.rs | 66 ++- crates/aenv/src/commands/start.rs | 171 +++++- docs/src/concepts/sandboxes/starting.md | 37 ++ .../src/getting-started/aenv-cli/sandboxes.md | 2 + src/api/generated/src/models.rs | 17 + src/api/impls/sandbox.rs | 506 ++++++++++++------ src/api/openapi.yml | 9 + src/sandbox/extra_drive.rs | 6 +- src/sandbox/firecracker/config.rs | 48 +- src/sandbox/firecracker/sandbox.rs | 265 +++++++-- src/sandbox/mod.rs | 2 +- src/snapshot/backends/tests.rs | 4 +- src/snapshot/store/posix_tests.rs | 2 +- src/snapshot/types/drive.rs | 2 +- src/snapshot/types/manifest.rs | 45 +- src/snapshot/types/snapshot.rs | 11 + src/types/extra_drive.rs | 123 ++++- src/types/mod.rs | 2 +- tests/integration/snapshot_attached_drive.rs | 141 ++++- 19 files changed, 1186 insertions(+), 273 deletions(-) diff --git a/crates/aenv/src/client/sandboxes.rs b/crates/aenv/src/client/sandboxes.rs index ec1797ebd..a9f265069 100644 --- a/crates/aenv/src/client/sandboxes.rs +++ b/crates/aenv/src/client/sandboxes.rs @@ -22,15 +22,32 @@ fn volume_mounts_request( }) } +#[derive(Debug, PartialEq, Serialize)] +pub struct AttachedDrive { + #[serde(rename = "driveID")] + pub drive_id: String, + #[serde(rename = "readOnly")] + pub read_only: bool, + #[serde(rename = "mountPath")] + pub mount_path: String, + pub source: AttachedDriveSource, +} + +#[derive(Debug, PartialEq, Serialize)] +pub struct AttachedDriveSource { + pub image: String, +} + #[derive(Debug, Serialize)] -pub struct NewSandbox<'a> { +pub struct NewSandboxV2<'a> { #[serde(rename = "templateID")] pub template_id: &'a str, #[serde(skip_serializing_if = "Option::is_none")] pub timeout: Option, - pub secure: bool, #[serde(skip_serializing_if = "Option::is_none", rename = "volumeMounts")] pub volume_mounts: Option>, + #[serde(skip_serializing_if = "Option::is_none", rename = "attachedDrives")] + pub attached_drives: Option>, } #[derive(Debug, Serialize)] @@ -47,6 +64,8 @@ pub struct NewColdSandbox<'a> { pub secure: bool, #[serde(skip_serializing_if = "Option::is_none", rename = "volumeMounts")] pub volume_mounts: Option>, + #[serde(skip_serializing_if = "Option::is_none", rename = "attachedDrives")] + pub attached_drives: Option>, } #[derive(Deserialize)] @@ -100,14 +119,15 @@ impl Client { template_id: &str, timeout: Option, volume_mounts: Option>, + attached_drives: Option>, ) -> Result { - let body = NewSandbox { + let body = NewSandboxV2 { template_id, timeout, - secure: true, volume_mounts: volume_mounts_request(volume_mounts), + attached_drives, }; - let resp = handle_status(self.post("/sandboxes").send_json(&body))?; + let resp = handle_status(self.post("/v2/sandboxes").send_json(&body))?; let sandbox: Sandbox = resp.into_json()?; Ok(sandbox) } @@ -121,6 +141,7 @@ impl Client { memory_mb: Option, disk_size_mb: Option, volume_mounts: Option>, + attached_drives: Option>, ) -> Result { let body = NewColdSandbox { image, @@ -130,6 +151,7 @@ impl Client { disk_size_mb, secure: true, volume_mounts: volume_mounts_request(volume_mounts), + attached_drives, }; let resp = handle_status(self.post("/sandboxes-cold").send_json(&body))?; let sandbox: Sandbox = resp.into_json()?; @@ -203,23 +225,48 @@ impl Client { #[cfg(test)] mod tests { - use super::{NewColdSandbox, NewSandbox, RefreshSandbox}; + use super::{AttachedDrive, AttachedDriveSource, NewColdSandbox, NewSandboxV2, RefreshSandbox}; #[test] fn new_sandbox_serializes_template_start() { - let body = NewSandbox { + let body = NewSandboxV2 { template_id: "base-template", timeout: Some(300), - secure: true, volume_mounts: None, + attached_drives: None, }; let value = serde_json::to_value(body).unwrap(); assert_eq!(value["templateID"], "base-template"); assert_eq!(value["timeout"], 300); - assert_eq!(value["secure"], true); + assert!(value.get("secure").is_none()); assert!(value.get("cpuCount").is_none()); assert!(value.get("memoryMB").is_none()); + assert!(value.get("attachedDrives").is_none()); + } + + #[test] + fn new_sandbox_serializes_attached_drives() { + let body = NewSandboxV2 { + template_id: "base-template", + timeout: None, + volume_mounts: None, + attached_drives: Some(vec![AttachedDrive { + drive_id: "image_mnt_models".to_owned(), + read_only: false, + mount_path: "/mnt/models".to_owned(), + source: AttachedDriveSource { + image: "aenv/models:v1".to_owned(), + }, + }]), + }; + + let value = serde_json::to_value(body).unwrap(); + let drive = &value["attachedDrives"][0]; + assert_eq!(drive["driveID"], "image_mnt_models"); + assert_eq!(drive["readOnly"], false); + assert_eq!(drive["mountPath"], "/mnt/models"); + assert_eq!(drive["source"]["image"], "aenv/models:v1"); } #[test] @@ -232,6 +279,7 @@ mod tests { disk_size_mb: Some(8192), secure: true, volume_mounts: None, + attached_drives: None, }; let value = serde_json::to_value(body).unwrap(); diff --git a/crates/aenv/src/commands/start.rs b/crates/aenv/src/commands/start.rs index b3c9b00d1..a9e3bc81d 100644 --- a/crates/aenv/src/commands/start.rs +++ b/crates/aenv/src/commands/start.rs @@ -1,7 +1,8 @@ +use crate::client::sandboxes::{AttachedDrive, AttachedDriveSource}; use crate::client::{volumes::validate_volume_reference, Client}; use anyhow::{Context, Result}; use clap::Args as ClapArgs; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::time::{Duration, Instant}; const ENVD_READY_TIMEOUT: Duration = Duration::from_secs(30); @@ -14,6 +15,8 @@ const ENVD_READY_PROBE_INTERVAL: Duration = Duration::from_millis(100); aenv start --cold docker.io/library/ubuntu:24.04 aenv start --cold docker.io/library/ubuntu:24.04 --cpu 2 --mem 2048 --disk-size-mb 8192 aenv start --cold -d docker.io/library/ubuntu:24.04 + aenv start my-template --image-mount /workspace=aenv/repo:main + aenv start my-template --image-mount-ro /mnt/models=aenv/models:v1 Resource overrides are only supported with --cold.")] pub struct Args { @@ -37,6 +40,12 @@ pub struct Args { /// Mount a persistent volume using MOUNT_PATH=VOLUME_ID_OR_NAME (repeatable) #[arg(long = "volume", value_name = "MOUNT_PATH=VOLUME", action = clap::ArgAction::Append)] volumes: Vec, + /// Attach a writable drive from an image using MOUNT_PATH=IMAGE (repeatable) + #[arg(long = "image-mount", value_name = "MOUNT_PATH=IMAGE", action = clap::ArgAction::Append)] + image_mounts: Vec, + /// Attach a read-only drive from an image using MOUNT_PATH=IMAGE (repeatable) + #[arg(long = "image-mount-ro", value_name = "MOUNT_PATH=IMAGE", action = clap::ArgAction::Append)] + read_only_image_mounts: Vec, } fn parse_disk_size_mb(value: &str) -> std::result::Result { @@ -52,6 +61,11 @@ fn parse_disk_size_mb(value: &str) -> std::result::Result { pub fn run(args: Args) -> Result<()> { let client = Client::from_env()?; let volume_mounts = parse_volume_mounts(&args.volumes)?; + let attached_drives = parse_image_mounts( + &args.image_mounts, + &args.read_only_image_mounts, + volume_mounts.as_ref(), + )?; let sandbox = if args.cold { client.create_cold_sandbox( &args.target, @@ -60,12 +74,18 @@ pub fn run(args: Args) -> Result<()> { args.resources.memory_mb, args.disk_size_mb, volume_mounts, + attached_drives, )? } else { if args.resources.is_set() || args.disk_size_mb.is_some() { anyhow::bail!("--cpu-count, --memory-mb, and --disk-size-mb require --cold"); } - client.create_sandbox(&args.target, Some(args.timeout), volume_mounts)? + client.create_sandbox( + &args.target, + Some(args.timeout), + volume_mounts, + attached_drives, + )? }; let sandbox_id = sandbox.sandbox_id; @@ -111,6 +131,69 @@ fn parse_volume_mounts(values: &[String]) -> Result>, +) -> Result>> { + if writable.is_empty() && read_only.is_empty() { + return Ok(None); + } + let mut paths: Vec = volume_mounts + .map(|mounts| mounts.keys().cloned().collect()) + .unwrap_or_default(); + let mut drive_ids = HashSet::new(); + let mut drives = Vec::with_capacity(writable.len() + read_only.len()); + let values = writable + .iter() + .map(|value| (value, false)) + .chain(read_only.iter().map(|value| (value, true))); + for (value, read_only) in values { + let Some((mount_path, image)) = value + .split_once('=') + .filter(|(_, image)| !image.trim().is_empty()) + else { + anyhow::bail!("image mount must use MOUNT_PATH=IMAGE syntax: {value}"); + }; + let mount_path = normalize_mount_path(mount_path) + .with_context(|| format!("invalid image mount: {value}"))?; + if paths + .iter() + .any(|existing| mount_paths_overlap(existing, &mount_path)) + { + anyhow::bail!("image mount path overlaps another mount: {mount_path}"); + } + paths.push(mount_path.clone()); + drives.push(AttachedDrive { + drive_id: image_drive_id(&mount_path, &mut drive_ids), + read_only, + mount_path, + source: AttachedDriveSource { + image: image.trim().to_owned(), + }, + }); + } + Ok(Some(drives)) +} + +/// Derives a unique drive ID such as `image_mnt_models` from a normalized +/// mount path. The prefix keeps it clear of reserved drive IDs. +fn image_drive_id(mount_path: &str, used: &mut HashSet) -> String { + let base: String = format!("image{mount_path}") + .chars() + .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' }) + .collect(); + let mut drive_id = base.clone(); + let mut suffix = 2; + while !used.insert(drive_id.clone()) { + drive_id = format!("{base}_{suffix}"); + suffix += 1; + } + drive_id +} + fn normalize_mount_path(value: &str) -> Result { if !value.starts_with('/') || value == "/" { anyhow::bail!("volume mount path must be an absolute guest path other than /: {value}"); @@ -186,7 +269,7 @@ async fn wait_for_envd( #[cfg(test)] mod tests { - use super::{parse_volume_mounts, Args}; + use super::{parse_image_mounts, parse_volume_mounts, Args}; use clap::Parser; #[derive(Parser)] @@ -221,4 +304,86 @@ mod tests { .is_err() ); } + + #[test] + fn parses_image_mounts_as_attached_drives() { + assert!(parse_image_mounts(&[], &[], None).unwrap().is_none()); + let drives = parse_image_mounts( + &[ + "/workspace=registry.example.com:5000/repo:main".to_owned(), + "/mnt/a-b=aenv/a".to_owned(), + ], + &[ + "/mnt/models/=aenv/models:v1".to_owned(), + "/mnt/a.b=aenv/b".to_owned(), + ], + None, + ) + .unwrap() + .unwrap(); + let summary: Vec<_> = drives + .iter() + .map(|drive| { + ( + drive.drive_id.as_str(), + drive.mount_path.as_str(), + drive.source.image.as_str(), + drive.read_only, + ) + }) + .collect(); + assert_eq!( + summary, + [ + ( + "image_workspace", + "/workspace", + "registry.example.com:5000/repo:main", + false + ), + ("image_mnt_a_b", "/mnt/a-b", "aenv/a", false), + ("image_mnt_models", "/mnt/models", "aenv/models:v1", true), + ("image_mnt_a_b_2", "/mnt/a.b", "aenv/b", true), + ] + ); + } + + #[test] + fn image_mount_flags_parse_for_warm_and_cold_starts() { + for cold in [false, true] { + let mut values = vec!["test", "my-template"]; + if cold { + values.push("--cold"); + } + values.extend([ + "--image-mount", + "/workspace=aenv/repo:main", + "--image-mount-ro", + "/mnt/models=aenv/models:v1", + ]); + let args = TestCli::try_parse_from(values).unwrap().args; + let drives = parse_image_mounts(&args.image_mounts, &args.read_only_image_mounts, None) + .unwrap() + .unwrap(); + assert_eq!(drives.len(), 2); + assert!(!drives[0].read_only); + assert!(drives[1].read_only); + } + } + + #[test] + fn rejects_invalid_or_overlapping_image_mounts() { + assert!(parse_image_mounts(&["/mnt/models".to_owned()], &[], None).is_err()); + assert!(parse_image_mounts(&["/mnt/models=".to_owned()], &[], None).is_err()); + assert!(parse_image_mounts(&["models=aenv/models".to_owned()], &[], None).is_err()); + assert!(parse_image_mounts( + &["/mnt/models=aenv/a".to_owned()], + &["/mnt/models/sub=aenv/b".to_owned()], + None, + ) + .is_err()); + let volumes = parse_volume_mounts(&["/mnt/data=cache".to_owned()]).unwrap(); + let overlapping = ["/mnt/data=aenv/models".to_owned()]; + assert!(parse_image_mounts(&overlapping, &[], volumes.as_ref()).is_err()); + } } diff --git a/docs/src/concepts/sandboxes/starting.md b/docs/src/concepts/sandboxes/starting.md index 3760536d8..653d11759 100644 --- a/docs/src/concepts/sandboxes/starting.md +++ b/docs/src/concepts/sandboxes/starting.md @@ -31,6 +31,8 @@ Warm-start options: | `` | Required | Template or snapshot ID or alias. | | `--timeout ` | `300` | Set the sandbox TTL. The sandbox auto-pauses when it reaches the TTL; see [Auto-Eviction](./auto-eviction.md). | | `--volume =` | None | Mount a persistent volume by ID or name. Repeat the option to mount multiple volumes. Not allowed when the snapshot already has volume mounts. | +| `--image-mount =` | None | Attach a writable drive from an image; see [Attach image drives](#attach-image-drives). Repeatable. Not allowed when the snapshot already has volume mounts or attached drives. | +| `--image-mount-ro =` | None | Attach a read-only drive from an image. Repeatable. Not allowed when the snapshot already has volume mounts or attached drives. | | `-d`, `--detach` | Off | Print the sandbox ID and exit instead of attaching an interactive shell. | Without `--detach`, `aenv start` waits for the sandbox to become ready and then @@ -75,6 +77,8 @@ Cold-start options: | `--memory ` | `[machine].mem_size_mib` from your config file | Set sandbox memory. Aliases: `--memory-mb`, `--mem`. | | `--disk-size-mb ` | Source image virtual size | Set root filesystem size. The value must be greater than zero and divisible by 1024 MiB. Alias: `--disk-mb`. | | `--volume =` | None | Mount a persistent volume by ID or name. Repeat the option to mount multiple volumes. | +| `--image-mount =` | None | Attach a writable drive from an image; see [Attach image drives](#attach-image-drives). Repeatable. | +| `--image-mount-ro =` | None | Attach a read-only drive from an image. Repeatable. | | `-d`, `--detach` | Off | Print the sandbox ID and exit instead of attaching an interactive shell. | Cold-started sandboxes also use secure sandbox authentication by default. @@ -93,6 +97,39 @@ applies only when creating a fresh writable root filesystem, not to read-only images, images with an existing upper layer, or snapshot resume. Sandbox responses report the effective size as `diskSizeMB`. +## Attach image drives + +Warm and cold starts can attach image filesystems as drives at guest paths: + +```bash +aenv start my-template \ + --image-mount /workspace=aenv/repo:main \ + --image-mount-ro /mnt/models=aenv/models:v1 +``` + +The CLI sends these as `attachedDrives` entries, on `POST /v2/sandboxes` for +warm starts and on `POST /sandboxes-cold` for cold starts. Each entry has a +`driveID`, a `source.image`, and optional `mountPath`, `subPath`, `readOnly`, +and `diskSizeMB`. The API defaults `readOnly` to `true`; the CLI sets it to +`false` for `--image-mount` and derives `driveID` from the mount path, such as +`image_mnt_models` for `/mnt/models`. The image accepts the same references as +a cold-start ``. + +An attached drive belongs to its sandbox and is removed with it. Writes are not +kept elsewhere; use a [volume](../volumes/index.md) for data that must outlive +the sandbox. + +| Operation | Attached drives | +|---|---| +| Snapshot | Captures each drive's current state into the snapshot. | +| Launch from a snapshot | Restores the captured drives. Providing `attachedDrives`, including an empty array, is rejected when the snapshot has volume mounts or attached drives. | +| Fork | Gives each child its own copy of the drive state. | +| Pause and resume | Keeps the drives and their writes. | + +On a warm start, attached drives and new volumes share the snapshot's reserved +drive slots, so their combined count cannot exceed the slots it has free. +Snapshots captured before reserved slots existed have none. + ## Build reusable images To build a standalone image, run `aenv build --image --name my-app ./app`. The diff --git a/docs/src/getting-started/aenv-cli/sandboxes.md b/docs/src/getting-started/aenv-cli/sandboxes.md index d4f2a7661..35cc21aa1 100644 --- a/docs/src/getting-started/aenv-cli/sandboxes.md +++ b/docs/src/getting-started/aenv-cli/sandboxes.md @@ -21,6 +21,8 @@ CLI obtains and manages the access token automatically. | `--memory ` | Memory in MiB; only valid with `--cold`. Defaults to `[machine].mem_size_mib` on the server. Aliases: `--memory-mb`, `--mem`. | | `--disk-size-mb ` | Root filesystem size; only valid with `--cold`. Defaults to the source image's virtual size; an explicit value must be at least 1024 and divisible by 1024 MiB. Alias: `--disk-mb`. | | `--volume ` | Mount a persistent volume by ID or name. Available for warm and cold starts; repeatable. | +| `--image-mount ` | Attach a writable drive from an image. Available for warm and cold starts; repeatable. See [Attach image drives](../../concepts/sandboxes/starting.md#attach-image-drives). | +| `--image-mount-ro ` | Attach a read-only drive from an image. Available for warm and cold starts; repeatable. | | `-d, --detach` | Print the sandbox ID and exit without attaching a shell | CPU, memory, and disk overrides are supported only for cold starts. diff --git a/src/api/generated/src/models.rs b/src/api/generated/src/models.rs index fbca57493..3751741f3 100644 --- a/src/api/generated/src/models.rs +++ b/src/api/generated/src/models.rs @@ -4309,6 +4309,12 @@ pub struct NewSandboxV2 { #[validate(nested)] #[serde(skip_serializing_if = "Option::is_none")] pub volume_mounts: Option>, + + /// Drives to attach at launch. If the snapshot already has volume mounts or attached drives, providing this field, including an empty array, returns HTTP 400. Omit this field to restore the snapshot's drives. If the sandbox is later snapshotted, current drive state is captured into the snapshot. + #[serde(rename = "attachedDrives")] + #[validate(nested)] + #[serde(skip_serializing_if = "Option::is_none")] + pub attached_drives: Option>, } impl NewSandboxV2 { @@ -4326,6 +4332,7 @@ impl NewSandboxV2 { custom_extension_params: None, mcp: None, volume_mounts: None, + attached_drives: None, } } } @@ -4367,6 +4374,8 @@ impl std::fmt::Display for NewSandboxV2 { // Skipping mcp in query parameter serialization // Skipping volumeMounts in query parameter serialization + + // Skipping attachedDrives in query parameter serialization ]; write!( @@ -4400,6 +4409,7 @@ impl std::str::FromStr for NewSandboxV2 { Vec>, pub mcp: Vec>, pub volume_mounts: Vec>, + pub attached_drives: Vec>, } let mut intermediate_rep = IntermediateRep::default(); @@ -4477,6 +4487,12 @@ impl std::str::FromStr for NewSandboxV2 { .to_string(), ); } + "attachedDrives" => { + return std::result::Result::Err( + "Parsing a container in this style is not supported in NewSandboxV2" + .to_string(), + ); + } _ => { return std::result::Result::Err( "Unexpected key while parsing NewSandboxV2".to_string(), @@ -4508,6 +4524,7 @@ impl std::str::FromStr for NewSandboxV2 { "Nullable types not supported in NewSandboxV2".to_string(), )?, volume_mounts: intermediate_rep.volume_mounts.into_iter().next(), + attached_drives: intermediate_rep.attached_drives.into_iter().next(), }) } } diff --git a/src/api/impls/sandbox.rs b/src/api/impls/sandbox.rs index 2dc80b351..93d9946cf 100644 --- a/src/api/impls/sandbox.rs +++ b/src/api/impls/sandbox.rs @@ -18,10 +18,11 @@ use crate::orchestrator::{ }; use crate::prometheus::SandboxStageTimer; use crate::sandbox::custom_extension::HookStatusError; -use crate::sandbox::{normalize_mount_path, CustomExtensionParams, ExtraDrive}; +use crate::sandbox::{normalize_mount_path, CustomExtensionParams, ExtraDrive, ExtraDriveKind}; use crate::sandbox::{BaseSandboxNetworkPolicy, SandboxNetworkEgressPolicy, SandboxNetworkPolicy}; use crate::snapshot::{ - SnapshotAlias, SnapshotId, SnapshotPublishMetadata, SnapshotPublishSource, SnapshotVolume, + CommittedSnapshot, SnapshotAlias, SnapshotId, SnapshotPublishMetadata, SnapshotPublishSource, + SnapshotVolume, }; use crate::types::{ImageConfigs, SandboxId, SandboxResources}; use crate::volume::{VolumeManager, VolumeMode}; @@ -169,6 +170,48 @@ fn check_snapshot_volume_replacement( Ok(()) } +/// The restored guest keeps its captured filesystems, so launch-time attached +/// drives are accepted only by snapshots without volumes or attached drives. +fn check_snapshot_attached_drives( + committed: &CommittedSnapshot, + explicit_drives: bool, +) -> Result<(), models::Error> { + if explicit_drives + && (!committed.volume_snapshots.is_empty() || !committed.attached_drives.is_empty()) + { + return Err(ApiImpl::error( + 400, + "this snapshot has volume mounts or attached drives, which cannot be replaced; \ + omit attachedDrives to restore them", + )); + } + Ok(()) +} + +fn check_attached_drive_overlap( + drives: &[ExtraDrive], + volume_mounts: &HashMap, +) -> Result<(), models::Error> { + for path in volume_mounts.keys() { + let path = normalize_mount_path(path.into()) + .map_err(|error| ApiImpl::error(400, error.to_string()))?; + if let Some(drive) = drives.iter().find(|drive| { + path.starts_with(drive.mount_path()) || drive.mount_path().starts_with(&path) + }) { + return Err(ApiImpl::error( + 400, + format!( + "attached drive '{}' mount path {} overlaps volume mount path {}", + drive.drive_id(), + drive.mount_path().display(), + path.display() + ), + )); + } + } + Ok(()) +} + #[derive(Default)] struct PreparedVolumeMounts { owner: Option, @@ -698,6 +741,232 @@ fn validate_domain_allowlist(policy: &SandboxNetworkPolicy) -> anyhow::Result<() Ok(()) } +impl ApiImpl { + /// Creates a sandbox from a template or snapshot. `attached_drives` are + /// added at launch and are only accepted by the v2 endpoint. + async fn create_sandbox_from_snapshot( + &self, + body: &models::NewSandbox, + attached_drives: Option<&[models::AttachedDrive]>, + ) -> Result { + let timer = SandboxStageTimer::new("create_warm"); + let snapshot = match timer + .time( + "load_snapshot", + self.snapshot_manager.load_runnable(&body.template_id), + ) + .await + { + Ok(Some(snapshot)) => snapshot, + Ok(None) => { + return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( + 400, + format!("template {} not found", body.template_id), + ))); + } + Err(err) => { + warn!(error = ?err, template_id = %body.template_id, "failed to load runnable snapshot"); + return Ok(SandboxesPostResponse::Status500_ServerError( + Self::snapshot_manager_error(&err), + )); + } + }; + + let network_policy = + match network_policy_from_create(body.allow_internet_access, body.network.as_ref()) { + Ok(network) => network, + Err(err) => { + return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( + 400, + err.to_string(), + ))); + } + }; + + let custom_params = body + .custom_extension_params + .as_ref() + .map(params_model_to_map); + if let Err(err) = validate_custom_extension_params(custom_params.as_ref()) { + return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( + 400, + err.to_string(), + ))); + } + + if let Err(error) = check_snapshot_volume_replacement( + &snapshot.committed().volume_snapshots, + body.volume_mounts.is_some(), + ) { + return Ok(SandboxesPostResponse::Status400_BadRequest(error)); + } + if let Err(error) = + check_snapshot_attached_drives(snapshot.committed(), attached_drives.is_some()) + { + return Ok(SandboxesPostResponse::Status400_BadRequest(error)); + } + let attached_drives = attached_drives.unwrap_or_default(); + let requested_drive_count = + attached_drives.len() + body.volume_mounts.as_ref().map_or(0, Vec::len); + let free_slots = snapshot.manifest().free_drive_slots(); + if requested_drive_count > free_slots { + return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( + 400, + format!( + "this snapshot has {free_slots} free drive slots, but \ + {requested_drive_count} attached drives and volumes were requested" + ), + ))); + } + let image_resolver = self.image_resolver(); + let attached_drives: Vec<_> = match timer + .time( + "resolve_attached_drives", + resolve_attached_drives(attached_drives, image_resolver.as_ref()), + ) + .await + { + Ok(resolved) => resolved + .into_iter() + .map(|resolved| resolved.drive.with_kind(ExtraDriveKind::LaunchSlot)) + .collect(), + Err(err) => { + warn!(error = %err.message, "failed to resolve attached drives"); + return Ok(Self::client_or_server_response( + err, + SandboxesPostResponse::Status400_BadRequest, + SandboxesPostResponse::Status500_ServerError, + )); + } + }; + + let extra_drives_in_snapshot = + body.volume_mounts.is_none() && !snapshot.committed().volume_snapshots.is_empty(); + let (requested_volume_mounts, restored_volume_ids) = if let Some(mounts) = + body.volume_mounts.as_deref() + { + let mounts = match volume_mounts_from_model(Some(mounts)) { + Ok(mounts) => mounts, + Err(error) => { + return Ok(SandboxesPostResponse::Status400_BadRequest(error)); + } + }; + if let Some(mounts) = &mounts { + if let Err(error) = check_attached_drive_overlap(&attached_drives, mounts) { + return Ok(SandboxesPostResponse::Status400_BadRequest(error)); + } + } + (mounts, Vec::new()) + } else { + match restore_snapshot_volume_mounts(self, &snapshot).await { + Ok((mounts, volume_ids)) => ((!mounts.is_empty()).then_some(mounts), volume_ids), + Err(error) => { + return Ok(match error.code { + 500 => SandboxesPostResponse::Status500_ServerError(error), + _ => SandboxesPostResponse::Status400_BadRequest(error), + }); + } + } + }; + let PreparedVolumeMounts { + owner: pending_volume_owner, + drives: volume_drives, + mounts: volume_mounts, + volume_ids: reserved_volume_ids, + } = match prepare_volume_mounts(self, requested_volume_mounts.as_ref()).await { + Ok(prepared) => prepared, + Err(error) => { + cleanup_volume_ids(&self.volume_manager, &restored_volume_ids).await; + return Ok(match error.code { + 500.. => SandboxesPostResponse::Status500_ServerError(error), + 409 => SandboxesPostResponse::Status409_Conflict(error), + _ => SandboxesPostResponse::Status400_BadRequest(error), + }); + } + }; + // Attached drives take the first free slots, before volumes. + let mut extra_drives = attached_drives; + extra_drives.extend(volume_drives); + + let request = CreateSandboxRequest { + source: SandboxLaunchSource::Snapshot(Box::new(snapshot)), + extra_drives, + extra_drives_in_snapshot, + timeout: duration_from_secs(body.timeout), + timeout_action: match body.auto_pause { + Some(false) => SandboxTimeoutAction::Delete, + _ => SandboxTimeoutAction::Pause, + }, + auto_resume: body.auto_resume.as_ref().is_some_and(|cfg| cfg.enabled), + user_metadata: body.metadata.clone(), + env_vars: body + .env_vars + .clone() + .filter(|env_vars| !env_vars.is_empty()), + network_policy, + secure: body.secure == Some(true), + custom_extension_params: custom_params, + volume_mounts, + }; + + match timer + .time("create_sandbox", self.orchestrator.create_sandbox(request)) + .await + { + Ok(metadata) => { + if let Err(error) = finish_volume_reservation( + &self.volume_manager, + pending_volume_owner.as_deref(), + Some(metadata.id), + &reserved_volume_ids, + ) + .await + { + let _ = self.orchestrator.delete_sandbox(metadata.id).await; + if let Some(owner) = pending_volume_owner.as_deref() { + let _ = self + .volume_manager + .replace_owner_for(owner, None, &reserved_volume_ids) + .await; + } + cleanup_volume_ids(&self.volume_manager, &restored_volume_ids).await; + return Ok(SandboxesPostResponse::Status500_ServerError(Self::error( + 500, + format!("failed to finalize volume reservation: {error}"), + ))); + } + let sandbox_id = metadata.id.to_string(); + Ok( + SandboxesPostResponse::Status201_TheSandboxWasCreatedSuccessfully { + body: self.sandbox_model(metadata), + x_agentenv_sandbox_id: Some(sandbox_id), + }, + ) + } + Err(err) => { + let _ = finish_volume_reservation( + &self.volume_manager, + pending_volume_owner.as_deref(), + None, + &reserved_volume_ids, + ) + .await; + cleanup_volume_ids(&self.volume_manager, &restored_volume_ids).await; + // The custom extension rejected the start hook with a 4xx + // (start-resume for template-based creates). + if let Some(message) = hook_client_error_message(&err) { + return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( + 400, message, + ))); + } + Ok(SandboxesPostResponse::Status500_ServerError( + Self::internal_error(&err), + )) + } + } + } +} + #[async_trait] impl Sandboxes<()> for ApiImpl { type Claims = super::Claims; @@ -1055,184 +1324,18 @@ impl Sandboxes<()> for ApiImpl { _claims: &Self::Claims, body: &models::NewSandbox, ) -> Result { - let timer = SandboxStageTimer::new("create_warm"); - let snapshot = match timer - .time( - "load_snapshot", - self.snapshot_manager.load_runnable(&body.template_id), - ) - .await - { - Ok(Some(snapshot)) => snapshot, - Ok(None) => { - return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( - 400, - format!("template {} not found", body.template_id), - ))); - } - Err(err) => { - warn!(error = ?err, template_id = %body.template_id, "failed to load runnable snapshot"); - return Ok(SandboxesPostResponse::Status500_ServerError( - Self::snapshot_manager_error(&err), - )); - } - }; - - let network_policy = - match network_policy_from_create(body.allow_internet_access, body.network.as_ref()) { - Ok(network) => network, - Err(err) => { - return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( - 400, - err.to_string(), - ))); - } - }; - - let custom_params = body - .custom_extension_params - .as_ref() - .map(params_model_to_map); - if let Err(err) = validate_custom_extension_params(custom_params.as_ref()) { - return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( - 400, - err.to_string(), - ))); - } - - if let Err(error) = check_snapshot_volume_replacement( - &snapshot.committed().volume_snapshots, - body.volume_mounts.is_some(), - ) { - return Ok(SandboxesPostResponse::Status400_BadRequest(error)); - } - - let extra_drives_in_snapshot = - body.volume_mounts.is_none() && !snapshot.committed().volume_snapshots.is_empty(); - let (requested_volume_mounts, restored_volume_ids) = if let Some(mounts) = - body.volume_mounts.as_deref() - { - let mounts = match volume_mounts_from_model(Some(mounts)) { - Ok(mounts) => mounts, - Err(error) => { - return Ok(SandboxesPostResponse::Status400_BadRequest(error)); - } - }; - (mounts, Vec::new()) - } else { - match restore_snapshot_volume_mounts(self, &snapshot).await { - Ok((mounts, volume_ids)) => ((!mounts.is_empty()).then_some(mounts), volume_ids), - Err(error) => { - return Ok(match error.code { - 500 => SandboxesPostResponse::Status500_ServerError(error), - _ => SandboxesPostResponse::Status400_BadRequest(error), - }); - } - } - }; - let PreparedVolumeMounts { - owner: pending_volume_owner, - drives: volume_drives, - mounts: volume_mounts, - volume_ids: reserved_volume_ids, - } = match prepare_volume_mounts(self, requested_volume_mounts.as_ref()).await { - Ok(prepared) => prepared, - Err(error) => { - cleanup_volume_ids(&self.volume_manager, &restored_volume_ids).await; - return Ok(match error.code { - 500.. => SandboxesPostResponse::Status500_ServerError(error), - 409 => SandboxesPostResponse::Status409_Conflict(error), - _ => SandboxesPostResponse::Status400_BadRequest(error), - }); - } - }; - - let request = CreateSandboxRequest { - source: SandboxLaunchSource::Snapshot(Box::new(snapshot)), - extra_drives: volume_drives, - extra_drives_in_snapshot, - timeout: duration_from_secs(body.timeout), - timeout_action: match body.auto_pause { - Some(false) => SandboxTimeoutAction::Delete, - _ => SandboxTimeoutAction::Pause, - }, - auto_resume: body.auto_resume.as_ref().is_some_and(|cfg| cfg.enabled), - user_metadata: body.metadata.clone(), - env_vars: body - .env_vars - .clone() - .filter(|env_vars| !env_vars.is_empty()), - network_policy, - secure: body.secure == Some(true), - custom_extension_params: custom_params, - volume_mounts, - }; - - match timer - .time("create_sandbox", self.orchestrator.create_sandbox(request)) - .await - { - Ok(metadata) => { - if let Err(error) = finish_volume_reservation( - &self.volume_manager, - pending_volume_owner.as_deref(), - Some(metadata.id), - &reserved_volume_ids, - ) - .await - { - let _ = self.orchestrator.delete_sandbox(metadata.id).await; - if let Some(owner) = pending_volume_owner.as_deref() { - let _ = self - .volume_manager - .replace_owner_for(owner, None, &reserved_volume_ids) - .await; - } - cleanup_volume_ids(&self.volume_manager, &restored_volume_ids).await; - return Ok(SandboxesPostResponse::Status500_ServerError(Self::error( - 500, - format!("failed to finalize volume reservation: {error}"), - ))); - } - let sandbox_id = metadata.id.to_string(); - Ok( - SandboxesPostResponse::Status201_TheSandboxWasCreatedSuccessfully { - body: self.sandbox_model(metadata), - x_agentenv_sandbox_id: Some(sandbox_id), - }, - ) - } - Err(err) => { - let _ = finish_volume_reservation( - &self.volume_manager, - pending_volume_owner.as_deref(), - None, - &reserved_volume_ids, - ) - .await; - cleanup_volume_ids(&self.volume_manager, &restored_volume_ids).await; - // The custom extension rejected the start hook with a 4xx - // (start-resume for template-based creates). - if let Some(message) = hook_client_error_message(&err) { - return Ok(SandboxesPostResponse::Status400_BadRequest(Self::error( - 400, message, - ))); - } - Ok(SandboxesPostResponse::Status500_ServerError( - Self::internal_error(&err), - )) - } - } + self.create_sandbox_from_snapshot(body, None).await } async fn v2_sandboxes_post( &self, - method: &Method, - host: &Host, - cookies: &CookieJar, - claims: &Self::Claims, + _method: &Method, + _host: &Host, + _cookies: &CookieJar, + _claims: &Self::Claims, body: &models::NewSandboxV2, ) -> Result { + let attached_drives = body.attached_drives.as_deref(); let body = models::NewSandbox { template_id: body.template_id.clone(), timeout: Some(body.timeout.unwrap_or(300)), @@ -1249,7 +1352,7 @@ impl Sandboxes<()> for ApiImpl { }; Ok( match self - .sandboxes_post(method, host, cookies, claims, &body) + .create_sandbox_from_snapshot(&body, attached_drives) .await? { SandboxesPostResponse::Status201_TheSandboxWasCreatedSuccessfully { @@ -2297,6 +2400,59 @@ mod tests { assert_eq!(error.code, 400); } + #[test] + fn snapshot_attached_drives_require_a_snapshot_without_mounts() { + let empty = CommittedSnapshot::mock(); + let with_volume = CommittedSnapshot { + volume_snapshots: vec![SnapshotVolume { + mount_path: "/volume".to_string(), + mode: VolumeMode::Exclusive, + size_mb: 16, + layers: Vec::new(), + }], + ..CommittedSnapshot::mock() + }; + let with_drive = CommittedSnapshot { + attached_drives: vec![crate::snapshot::CommittedAttachedDrive::Overlaybd { + drive_id: "data".to_string(), + layers: Vec::new(), + read_only: true, + virtual_size: 4096, + mount_path: "/mnt/data".into(), + sub_path: None, + }], + ..CommittedSnapshot::mock() + }; + + assert!(check_snapshot_attached_drives(&empty, true).is_ok()); + for committed in [&with_volume, &with_drive] { + assert!(check_snapshot_attached_drives(committed, false).is_ok()); + let error = check_snapshot_attached_drives(committed, true) + .expect_err("explicit attached drives must not join existing mounts"); + assert_eq!(error.code, 400); + } + } + + #[test] + fn attached_drive_mount_paths_must_not_overlap_volume_mounts() -> anyhow::Result<()> { + let drives = [ExtraDrive::try_new_overlaybd_with_mount_path( + "data", + "/tmp/image.json", + true, + "/workspace/data", + None::, + )?]; + let mounts = |path: &str| HashMap::from([(path.to_string(), "vol_1".to_string())]); + + assert!(check_attached_drive_overlap(&drives, &mounts("/workspace/logs")).is_ok()); + for path in ["/workspace/data", "/workspace//data/cache", "/workspace"] { + let error = check_attached_drive_overlap(&drives, &mounts(path)) + .expect_err("overlapping mount paths should fail"); + assert_eq!(error.code, 400); + } + Ok(()) + } + #[test] fn volume_mounts_from_model_maps_names_by_path() { let mounts = vec![ diff --git a/src/api/openapi.yml b/src/api/openapi.yml index e35b5337d..d26ec5ffd 100644 --- a/src/api/openapi.yml +++ b/src/api/openapi.yml @@ -738,6 +738,15 @@ components: Omit this field to restore the snapshot's volumes. items: $ref: "#/components/schemas/SandboxVolumeMount" + attachedDrives: + allOf: + - $ref: "#/components/schemas/AttachedDrives" + description: > + Drives to attach at launch. If the snapshot already has volume + mounts or attached drives, providing this field, including an empty + array, returns HTTP 400. Omit this field to restore the snapshot's + drives. If the sandbox is later snapshotted, current drive state is + captured into the snapshot. ResumedSandbox: properties: diff --git a/src/sandbox/extra_drive.rs b/src/sandbox/extra_drive.rs index 222617f86..f774e13cb 100644 --- a/src/sandbox/extra_drive.rs +++ b/src/sandbox/extra_drive.rs @@ -85,18 +85,20 @@ impl PreparedDrives { } } +/// Prepares each drive with the mode `mode` returns for its index. pub(crate) async fn prepare_extra_drives( extra_drives: &[ExtraDrive], global_config_path: &Path, sandbox_work_dir: &Path, runtime_upper_mode: UpperMode, - mode: ExtraDrivePrepareMode, + mode: impl Fn(usize) -> ExtraDrivePrepareMode, ) -> Result { let mut mounts = Vec::with_capacity(extra_drives.len()); let mut cleanup_paths = Vec::with_capacity(extra_drives.len()); let mut runtimes = Vec::with_capacity(extra_drives.len()); - for drive in extra_drives { + for (index, drive) in extra_drives.iter().enumerate() { + let mode = mode(index); let result = async { let runtime_dir = drive.runtime_dir(sandbox_work_dir); let (requested_virtual_size, known_source_virtual_size) = mode.device_sizes(drive); diff --git a/src/sandbox/firecracker/config.rs b/src/sandbox/firecracker/config.rs index 0cb3999aa..40a641068 100644 --- a/src/sandbox/firecracker/config.rs +++ b/src/sandbox/firecracker/config.rs @@ -529,13 +529,7 @@ impl FirecrackerSnapshotConfig { // is already serialised inside vm_state.bin, so re-applying a template would // be incorrect and is rejected by Firecracker anyway. let extra_drives = manifest.extra_drives(); - let physical_extra_drive_count = if manifest.volume_drive_slots == 0 { - // Snapshots created before reserved volume slots existed contain - // only physical attached drives. - extra_drives.len() - } else { - manifest.physical_extra_drive_count - }; + let physical_extra_drive_count = manifest.effective_physical_extra_drive_count(); let snapshot_common = FirecrackerCommonConfig { ublk_config: Some(overlaybd_ublk_config), envd_version: snapshot.committed().runtime_versions.envd_version.clone(), @@ -678,6 +672,7 @@ mod tests { use super::*; use crate::cfg::{UblkOverlaybdTomlConfig, UblkTomlConfig}; use crate::snapshot::{CommittedSnapshot, ResolvedAttachedDrive, SnapshotRecord}; + use crate::types::ExtraDriveKind; use std::fs; use std::sync::{Mutex, OnceLock}; use tempfile::tempdir; @@ -812,7 +807,7 @@ mod tests { virtual_size: None, sub_path: None, snapshot_output_dir: None, - volume: false, + kind: ExtraDriveKind::Physical, }, ExtraDrive::Overlaybd { drive_id: "data".to_string(), @@ -822,7 +817,7 @@ mod tests { virtual_size: None, sub_path: None, snapshot_output_dir: None, - volume: false, + kind: ExtraDriveKind::Physical, }, ]; @@ -856,7 +851,7 @@ mod tests { virtual_size: None, sub_path: None, snapshot_output_dir: None, - volume: false, + kind: ExtraDriveKind::Physical, } }) .collect(); @@ -961,6 +956,39 @@ mod tests { Ok(()) } + #[test] + fn runnable_snapshot_keeps_committed_slot_drives_in_slots() -> Result<()> { + let drive = ResolvedAttachedDrive::Overlaybd { + drive_id: "data".to_owned(), + image_config_path: "/tmp/data.json".into(), + read_only: false, + virtual_size: 4096, + mount_path: ExtraDrive::default_mount_path("data"), + sub_path: None, + }; + let snapshot = RunnableSnapshot::from_test_slot_drive_manifest( + SnapshotRecord::mock_ready(CommittedSnapshot::mock()), + vec![drive], + ); + + let config = FirecrackerSnapshotConfig::from_runnable_snapshot(&snapshot)?; + + assert_eq!(config.common.physical_extra_drive_count, 0); + assert_eq!( + config.common.volume_drive_slots, + snapshot.manifest().volume_drive_slots + ); + assert_eq!( + config.common.extra_drives[0].kind(), + ExtraDriveKind::LaunchSlot + ); + assert_eq!( + snapshot.manifest().free_drive_slots(), + snapshot.manifest().volume_drive_slots - 1 + ); + Ok(()) + } + #[test] fn runnable_snapshot_without_tools_drive_version_is_not_launchable() { let mut committed = CommittedSnapshot::mock(); diff --git a/src/sandbox/firecracker/sandbox.rs b/src/sandbox/firecracker/sandbox.rs index 16359f2d0..f8976cd0b 100644 --- a/src/sandbox/firecracker/sandbox.rs +++ b/src/sandbox/firecracker/sandbox.rs @@ -48,8 +48,8 @@ use crate::sandbox::ublk::{ use crate::sandbox::SandboxLaunchConfig; use crate::snapshot::RunnableSnapshot; use crate::types::{ - ExtraDrive, RuntimeArtifactSet, SandboxId, ROOTFS_DRIVE_ID, USER_ROOTFS_DRIVE_ID, - VOLUME_DRIVE_SLOT_PREFIX, + ExtraDrive, ExtraDriveKind, RuntimeArtifactSet, SandboxId, ROOTFS_DRIVE_ID, + USER_ROOTFS_DRIVE_ID, VOLUME_DRIVE_SLOT_PREFIX, }; // ── Constants ──────────────────────────────────────────────────────────────── @@ -189,8 +189,9 @@ pub struct FirecrackerSandbox { envd_instance: Option, frozen_volume_mounts: Vec<(PathBuf, String)>, /// Drives added to a committed snapshot at launch time are not represented - /// in its guest mount namespace. Mount only those drives after envd is - /// ready; paused-state resumes must preserve the captured mount state. + /// in its guest mount namespace, and fresh read-only volumes must not be + /// mounted by init. Mount only those drives after envd is ready; + /// paused-state resumes must preserve the captured mount state. initial_guest_drive_mounts: Vec<(usize, ExtraDrive)>, current_rootfs_virtual_size: Option, live_snapshot_root: Option>, @@ -656,10 +657,12 @@ impl FirecrackerSandbox { } pub(crate) fn new_with_id(mut config: FirecrackerSandboxConfig, id: SandboxId) -> Result { - let (mut physical_drives, volume_drives): (Vec<_>, Vec<_>) = + let (mut physical_drives, mut volume_drives): (Vec<_>, Vec<_>) = std::mem::take(&mut config.common.extra_drives) .into_iter() - .partition(|drive| !drive.is_volume()); + .partition(|drive| !drive.uses_reserved_slot()); + // Committed slot drives must take the first slots, before volumes. + volume_drives.sort_by_key(ExtraDrive::is_volume); config.common.physical_extra_drive_count = physical_drives.len(); config.common.volume_drive_slots = ConfigManager::global_config() .volume @@ -673,13 +676,23 @@ impl FirecrackerSandbox { ); physical_drives.extend(volume_drives); config.common.extra_drives = physical_drives; + let initial_guest_drive_mounts = config + .common + .extra_drives + .iter() + .cloned() + .enumerate() + .filter(|(_, drive)| mounts_after_boot(drive)) + .collect(); debug!( firecracker_binary = %config.common.firecracker_binary.display(), kernel_image = %config.kernel_image.display(), tools_drive_version = %config.common.tools_drive_version, "creating fresh firecracker sandbox" ); - Self::build(id, LaunchMode::Fresh(config)) + let mut sandbox = Self::build(id, LaunchMode::Fresh(config))?; + sandbox.initial_guest_drive_mounts = initial_guest_drive_mounts; + Ok(sandbox) } /// Create a sandbox handle that resumes from the provided snapshot config. @@ -774,7 +787,9 @@ impl FirecrackerSandbox { return Vec::new(); } - let first_slot_index = snapshot_config.common.physical_extra_drive_count; + // Launch-time drives follow the snapshot's own drives. + let first_slot_index = + snapshot_config.common.extra_drives.len() - launch_config.extra_drives.len(); launch_config .extra_drives .iter() @@ -789,12 +804,23 @@ impl FirecrackerSandbox { launch_config: &SandboxLaunchConfig, ) -> Result { let mut snapshot_config = FirecrackerSnapshotConfig::from_runnable_snapshot(snapshot)?; + let free_slots = snapshot.manifest().free_drive_slots(); anyhow::ensure!( - launch_config.extra_drives.len() <= snapshot_config.common.volume_drive_slots, - "snapshot supports at most {} launch-time volume drives, requested {}", - snapshot_config.common.volume_drive_slots, + launch_config.extra_drives.len() <= free_slots, + "snapshot supports at most {} launch-time drives, requested {}", + free_slots, launch_config.extra_drives.len() ); + anyhow::ensure!( + launch_config + .extra_drives + .iter() + .all(ExtraDrive::uses_reserved_slot) + && launch_config + .extra_drives + .is_sorted_by_key(ExtraDrive::is_volume), + "launch-time drives must use reserved slots, with attached drives before volumes" + ); snapshot_config .common .extra_drives @@ -911,8 +937,9 @@ impl FirecrackerSandbox { ) .await?; - // The snapshot already carries mount state for its existing drives. - // Only drives newly supplied for this launch need a guest-side mount. + // The snapshot already carries mount state for its existing drives, + // and init mounts fresh drives except read-only volumes. Only the + // remaining drives need a guest-side mount. if !self.initial_guest_drive_mounts.is_empty() { let envd = self .envd_instance @@ -1125,7 +1152,6 @@ impl FirecrackerSandbox { ) .context("build firecracker snapshot manifest")?; manifest.volume_drive_slots = snapshot_common.volume_drive_slots; - manifest.physical_extra_drive_count = manifest_extra_drives.len(); let snapshot = FirecrackerSnapshotConfig { common: snapshot_common, @@ -1345,16 +1371,20 @@ fn rootfs_and_extra_drive_image_config_paths(common: &FirecrackerCommonConfig) - paths } +/// Read-only drives in writable reserved slots, such as read-only volumes, +/// would have their superblock written by init's default mount. They are +/// mounted `ro,noload` through envd after boot instead, as on snapshot launches. +fn mounts_after_boot(drive: &ExtraDrive) -> bool { + drive.uses_reserved_slot() && drive.read_only() +} + /// Build the `agentenv_drives=vdc:/mnt/data,vdd:/mnt/logs:sub/path` boot arg -/// for extra drives. Returns `None` if there are no extra drives. +/// for extra drives that init mounts. Returns `None` if there are none. /// Drive letter mapping: vda = tools drive, vdb = user image, vdc = first extra drive. /// Each entry is `vd:[:]`, with optional /// Kubernetes-style `subPath` semantics. API validation rejects `:` in both /// `mountPath` and `subPath`, so the `:` separators are unambiguous. fn build_drives_boot_arg(extra_drives: &[ExtraDrive]) -> Option { - if extra_drives.is_empty() { - return None; - } assert!( extra_drives.len() <= MAX_EXTRA_DRIVES, "too many extra drives for guest naming: {} > {}", @@ -1364,6 +1394,7 @@ fn build_drives_boot_arg(extra_drives: &[ExtraDrive]) -> Option { let entries: Vec = extra_drives .iter() .enumerate() + .filter(|(_, drive)| !mounts_after_boot(drive)) .map(|(i, drive)| { let dev_letter = (b'c' + i as u8) as char; let mut entry = format!("vd{}:{}", dev_letter, drive.mount_path().display()); @@ -1374,7 +1405,7 @@ fn build_drives_boot_arg(extra_drives: &[ExtraDrive]) -> Option { entry }) .collect(); - Some(format!("agentenv_drives={}", entries.join(","))) + (!entries.is_empty()).then(|| format!("agentenv_drives={}", entries.join(","))) } /// Append a space-separated argument to an optional kernel command line. @@ -1526,7 +1557,7 @@ impl FirecrackerSandbox { &overlaybd_global, self.work_dir.path(), runtime_upper_mode, - ExtraDrivePrepareMode::Fresh { allow_shrink }, + |_| ExtraDrivePrepareMode::Fresh { allow_shrink }, ) .await .context("prepare extra drives")? @@ -1805,16 +1836,34 @@ impl FirecrackerSandbox { /// Prepare snapshot-backed extra drives and the reserved volume slots. /// /// Snapshot loading reopens every serialized slot path, so all - /// placeholders must exist. Restored volumes replace their placeholders - /// before load. Volumes newly supplied for this launch are absent from the - /// source state; they are returned as `(slot, drive)` pairs to attach with + /// placeholders must exist. Slot drives in the source state, such as + /// restored volumes, replace their placeholders before load. Drives newly + /// supplied for this launch are absent from the source state; they are + /// returned as `(slot, drive)` pairs to attach with /// [`attach_new_volume_drives`](Self::attach_new_volume_drives) after load. async fn prepare_resume_drives( &mut self, common: &FirecrackerCommonConfig, ) -> Result> { + let new_indices = self + .initial_guest_drive_mounts + .iter() + .map(|(index, _)| *index) + .collect::>(); + let is_new_drive = |index: usize| new_indices.contains(&index); + // Attached drives new to this launch carry the requested size, as on a + // fresh boot; every other drive, including volumes, carries its known size. + let allow_shrink = ConfigManager::global_config().ublk.overlaybd.allow_shrink; let extra_drive_attachments = self - .prepare_snapshot_backing_drives(&common.extra_drives) + .prepare_snapshot_backing_drives(&common.extra_drives, |index| { + if is_new_drive(index) + && common.extra_drives[index].kind() == ExtraDriveKind::LaunchSlot + { + ExtraDrivePrepareMode::Fresh { allow_shrink } + } else { + ExtraDrivePrepareMode::Resume + } + }) .await .context("prepare snapshot-backed extra drives for resume")?; let volume_slots = self @@ -1834,22 +1883,25 @@ impl FirecrackerSandbox { volume_drive_attachments.len() <= volume_slots.len(), "snapshot has more launch-time drives than reserved volume slots" ); - let volume_binds = volume_slots + let mut new_binds = Vec::new(); + for (offset, (slot, drive)) in volume_slots .into_iter() .zip(volume_drive_attachments.iter().cloned()) - .collect::>(); - if !self.initial_guest_drive_mounts.is_empty() { - return Ok(volume_binds); - } - for (slot, drive) in &volume_binds { - self.bind_volume_drive_slot(slot, drive).with_context(|| { - format!( - "bind restored volume {} to reserved drive {} before snapshot load", - drive.drive_id, slot.drive_id - ) - })?; + .enumerate() + { + if is_new_drive(physical_count + offset) { + new_binds.push((slot, drive)); + continue; + } + self.bind_volume_drive_slot(&slot, &drive) + .with_context(|| { + format!( + "bind restored drive {} to reserved drive {} before snapshot load", + drive.drive_id, slot.drive_id + ) + })?; } - Ok(Vec::new()) + Ok(new_binds) } /// A volume absent from the source state was loaded as a placeholder. @@ -2305,6 +2357,7 @@ impl FirecrackerSandbox { async fn prepare_snapshot_backing_drives( &mut self, extra_drives: &[ExtraDrive], + mode: impl Fn(usize) -> ExtraDrivePrepareMode, ) -> Result> { if extra_drives.is_empty() { self.resources.extra_drive_runtimes.clear(); @@ -2320,7 +2373,7 @@ impl FirecrackerSandbox { &overlaybd_global, self.work_dir.path(), runtime_upper_mode, - ExtraDrivePrepareMode::Resume, + mode, ) .await?; let (attachments, extra_drive_runtimes) = prepared_extra_drives.into_parts(); @@ -2815,6 +2868,42 @@ mod tests { assert_eq!(common.physical_extra_drive_count, 1); assert!(common.volume_drive_slots >= 1); assert_eq!(common.extra_drives, vec![physical_drive, volume_drive]); + assert!(sandbox.initial_guest_drive_mounts.is_empty()); + Ok(()) + } + + #[test] + fn fresh_read_only_volumes_are_mounted_after_boot() -> Result<()> { + let work_base = TempDir::new()?; + let physical_drive = ExtraDrive::try_new_overlaybd("data", "/tmp/data-image.json", true)?; + let writable_volume = + ExtraDrive::try_new_overlaybd("volume", "/tmp/volume-image.json", false)? + .with_volume_snapshot_output_dir(None); + let read_only_volume = + ExtraDrive::try_new_overlaybd("image", "/tmp/image-image.json", true)? + .with_volume_snapshot_output_dir(None); + let mut config = fresh_config(); + config.common.firecracker_work_base_dir = Some(work_base.path().to_path_buf()); + config.common.extra_drives = vec![ + read_only_volume.clone(), + physical_drive.clone(), + writable_volume.clone(), + ]; + + let sandbox = FirecrackerSandbox::new(config)?; + + assert_eq!( + sandbox.launch.common().extra_drives, + vec![physical_drive, read_only_volume.clone(), writable_volume] + ); + assert_eq!( + sandbox.initial_guest_drive_mounts, + vec![(1, read_only_volume)] + ); + assert_eq!( + build_drives_boot_arg(&sandbox.launch.common().extra_drives), + Some("agentenv_drives=vdc:/mnt/data,vde:/mnt/volume".to_string()) + ); Ok(()) } @@ -3193,6 +3282,106 @@ mod tests { Ok(()) } + #[test] + fn build_drives_boot_arg_omits_drives_mounted_after_boot() -> Result<()> { + let drives = vec![ + ExtraDrive::try_new_overlaybd("image", "/tmp/image.json", true)? + .with_volume_snapshot_output_dir(None), + ]; + + assert_eq!(build_drives_boot_arg(&drives), None); + Ok(()) + } + + #[test] + fn fresh_slot_drives_precede_volumes() -> Result<()> { + let work_base = TempDir::new()?; + let volume = ExtraDrive::try_new_overlaybd("volume", "/tmp/volume-image.json", false)? + .with_volume_snapshot_output_dir(None); + let slot_drive = ExtraDrive::try_new_overlaybd("data", "/tmp/data-image.json", false)? + .with_kind(ExtraDriveKind::LaunchSlot); + let mut config = fresh_config(); + config.common.firecracker_work_base_dir = Some(work_base.path().to_path_buf()); + config.common.extra_drives = vec![volume.clone(), slot_drive.clone()]; + + let sandbox = FirecrackerSandbox::new(config)?; + + let common = sandbox.launch.common(); + assert_eq!(common.physical_extra_drive_count, 0); + assert_eq!(common.extra_drives, vec![slot_drive, volume]); + Ok(()) + } + + #[test] + fn snapshot_launch_adds_drives_after_committed_slot_drives() -> Result<()> { + let committed_drive = crate::snapshot::ResolvedAttachedDrive::Overlaybd { + drive_id: "data".to_owned(), + image_config_path: "/tmp/data-image.json".into(), + read_only: false, + virtual_size: 4096, + mount_path: ExtraDrive::default_mount_path("data"), + sub_path: None, + }; + let snapshot = RunnableSnapshot::from_test_slot_drive_manifest( + SnapshotRecord::mock_ready(CommittedSnapshot::mock()), + vec![committed_drive], + ); + let volume = ExtraDrive::try_new_overlaybd("volume", "/tmp/volume-image.json", false)? + .with_volume_snapshot_output_dir(None); + let launch_config = SandboxLaunchConfig { + sandbox_id: SandboxId::new(), + snapshot_id: "tpl-test".to_string(), + extra_drives: vec![volume.clone()], + ..SandboxLaunchConfig::default() + }; + + let snapshot_config = + FirecrackerSandbox::snapshot_config_for_launch(&snapshot, &launch_config)?; + + assert_eq!(snapshot_config.common.physical_extra_drive_count, 0); + assert_eq!( + snapshot_config + .common + .extra_drives + .iter() + .map(ExtraDrive::kind) + .collect::>(), + vec![ExtraDriveKind::LaunchSlot, ExtraDriveKind::Volume] + ); + assert_eq!( + FirecrackerSandbox::initial_guest_drive_mounts_for_snapshot_launch( + &snapshot_config, + &launch_config + ), + vec![(1, volume.clone())] + ); + + let free_slots = snapshot.manifest().free_drive_slots(); + let too_many = SandboxLaunchConfig { + extra_drives: vec![volume.clone(); free_slots + 1], + ..launch_config.clone() + }; + assert!(FirecrackerSandbox::snapshot_config_for_launch(&snapshot, &too_many).is_err()); + + let slot_drive = ExtraDrive::try_new_overlaybd("logs", "/tmp/logs-image.json", false)? + .with_kind(ExtraDriveKind::LaunchSlot); + for extra_drives in [ + vec![ExtraDrive::try_new_overlaybd( + "logs", + "/tmp/logs-image.json", + false, + )?], + vec![volume, slot_drive], + ] { + let invalid = SandboxLaunchConfig { + extra_drives, + ..launch_config.clone() + }; + assert!(FirecrackerSandbox::snapshot_config_for_launch(&snapshot, &invalid).is_err()); + } + Ok(()) + } + #[test] fn from_snapshot_merges_launch_env_over_snapshot_env() -> Result<()> { let mut snapshot_env_vars = HashMap::new(); diff --git a/src/sandbox/mod.rs b/src/sandbox/mod.rs index c3a5b61e0..32d23885e 100644 --- a/src/sandbox/mod.rs +++ b/src/sandbox/mod.rs @@ -27,7 +27,7 @@ use crate::types::{ImageConfigs, SandboxId}; pub use crate::snapshot::{SandboxSnapshotManifest, FIRECRACKER_BACKEND}; pub use crate::types::{ normalize_mount_path, normalize_mount_path_for_drive, validate_drive_id, validate_mount_path, - validate_sub_path, ExtraDrive, + validate_sub_path, ExtraDrive, ExtraDriveKind, }; pub use ::envd::process::Signal; pub use access::{EnvdAccessToken, SandboxAccessTokenGenerator}; diff --git a/src/snapshot/backends/tests.rs b/src/snapshot/backends/tests.rs index 5a0328921..e197e6318 100644 --- a/src/snapshot/backends/tests.rs +++ b/src/snapshot/backends/tests.rs @@ -454,7 +454,7 @@ async fn publish_rejects_duplicate_attached_drive_ids_before_touching_catalog() virtual_size: Some(32768), sub_path: None, snapshot_output_dir: None, - volume: false, + kind: crate::types::ExtraDriveKind::Physical, }, ExtraDrive::Overlaybd { drive_id: "data".to_string(), @@ -464,7 +464,7 @@ async fn publish_rejects_duplicate_attached_drive_ids_before_touching_catalog() virtual_size: Some(32768), sub_path: None, snapshot_output_dir: None, - volume: false, + kind: crate::types::ExtraDriveKind::Physical, }, ], ); diff --git a/src/snapshot/store/posix_tests.rs b/src/snapshot/store/posix_tests.rs index a309aa7d6..f8ba99ab6 100644 --- a/src/snapshot/store/posix_tests.rs +++ b/src/snapshot/store/posix_tests.rs @@ -98,7 +98,7 @@ async fn attached_drive_rootfs_layers_are_derived_from_resolved_lower_paths() { virtual_size: Some(4096), sub_path: None, snapshot_output_dir: None, - volume: false, + kind: crate::types::ExtraDriveKind::Physical, }]) .expect("attached drive manifest should include virtual size"); diff --git a/src/snapshot/types/drive.rs b/src/snapshot/types/drive.rs index e6d3870fe..492c8bedc 100644 --- a/src/snapshot/types/drive.rs +++ b/src/snapshot/types/drive.rs @@ -56,7 +56,7 @@ impl ResolvedAttachedDrive { virtual_size: Some(*virtual_size), sub_path: sub_path.clone(), snapshot_output_dir: None, - volume: false, + kind: crate::types::ExtraDriveKind::Physical, }, } } diff --git a/src/snapshot/types/manifest.rs b/src/snapshot/types/manifest.rs index 0a4304dd9..7dc0c2c92 100644 --- a/src/snapshot/types/manifest.rs +++ b/src/snapshot/types/manifest.rs @@ -3,7 +3,7 @@ use std::path::{Path, PathBuf}; use anyhow::{bail, Result}; use serde::{Deserialize, Serialize}; -use crate::types::ExtraDrive; +use crate::types::{ExtraDrive, ExtraDriveKind}; pub(crate) const MANIFEST_FORMAT_VERSION: u32 = 1; @@ -43,11 +43,12 @@ pub struct SandboxSnapshotManifest { pub memory: SnapshotMemoryArtifacts, pub rootfs: SnapshotRootfsArtifacts, pub attached_drives: Vec, - /// Number of reserved virtio-block slots available for launch-time volumes. + /// Number of reserved virtio-block slots for launch-time drives. #[serde(default)] pub volume_drive_slots: usize, /// Number of attached drives represented by IDs in the Firecracker snapshot. - /// Launch-time volumes use the reserved slots that follow these drives. + /// Later attached drives occupy the first reserved slots, and launch-time + /// drives use the slots that follow them. #[serde(default)] pub physical_extra_drive_count: usize, /// Runtime-only startup pack reference resolved from the committed @@ -127,16 +128,40 @@ impl SandboxSnapshotManifest { }, attached_drives: Vec::new(), volume_drive_slots: 0, - physical_extra_drive_count: attached_drives.len(), + physical_extra_drive_count: attached_drives + .iter() + .filter(|drive| !drive.uses_reserved_slot()) + .count(), memory_startup_pack: None, } .with_extra_drives(attached_drives) } + /// Attached drives with their own Firecracker drive IDs. Snapshots created + /// before reserved slots existed contain only such drives. + pub fn effective_physical_extra_drive_count(&self) -> usize { + if self.volume_drive_slots == 0 { + self.attached_drives.len() + } else { + self.physical_extra_drive_count + } + } + + /// Reserved slots not occupied by attached drives. + pub fn free_drive_slots(&self) -> usize { + let slot_drives = self + .attached_drives + .len() + .saturating_sub(self.effective_physical_extra_drive_count()); + self.volume_drive_slots.saturating_sub(slot_drives) + } + pub fn extra_drives(&self) -> Vec { + let physical_count = self.effective_physical_extra_drive_count(); self.attached_drives .iter() - .map(|drive| ExtraDrive::Overlaybd { + .enumerate() + .map(|(index, drive)| ExtraDrive::Overlaybd { drive_id: drive.drive_id.clone(), image_config_path: drive.image_config_path.clone(), read_only: drive.read_only, @@ -148,7 +173,11 @@ impl SandboxSnapshotManifest { .unwrap_or_else(|_| ExtraDrive::default_mount_path(&drive.drive_id)), sub_path: drive.sub_path.clone(), snapshot_output_dir: None, - volume: false, + kind: if index < physical_count { + ExtraDriveKind::Physical + } else { + ExtraDriveKind::LaunchSlot + }, }) .collect() } @@ -291,7 +320,7 @@ mod tests { virtual_size: None, sub_path: None, snapshot_output_dir: None, - volume: false, + kind: ExtraDriveKind::Physical, }; let err = SandboxSnapshotManifest::new( @@ -328,7 +357,7 @@ mod tests { virtual_size: Some(0), sub_path: None, snapshot_output_dir: None, - volume: false, + kind: ExtraDriveKind::Physical, }; let err = manifest diff --git a/src/snapshot/types/snapshot.rs b/src/snapshot/types/snapshot.rs index 9d43158ad..346ab1611 100644 --- a/src/snapshot/types/snapshot.rs +++ b/src/snapshot/types/snapshot.rs @@ -622,6 +622,17 @@ impl RunnableSnapshot { snapshot.manifest.volume_drive_slots = 0; snapshot } + + /// Snapshot whose attached drives all occupy reserved slots. + #[cfg(test)] + pub(crate) fn from_test_slot_drive_manifest( + record: SnapshotRecord, + attached_drives: Vec, + ) -> Self { + let mut snapshot = Self::from_test_manifest(record, attached_drives); + snapshot.manifest.physical_extra_drive_count = 0; + snapshot + } } impl fmt::Debug for RunnableSnapshot { diff --git a/src/types/extra_drive.rs b/src/types/extra_drive.rs index c6e6455ae..5441f06bf 100644 --- a/src/types/extra_drive.rs +++ b/src/types/extra_drive.rs @@ -33,13 +33,68 @@ pub enum ExtraDrive { /// Optional persistent destination for a volume snapshot. #[serde(default, skip_serializing_if = "Option::is_none")] snapshot_output_dir: Option, - /// Persistent volumes are omitted from committed attached-drive - /// manifests and rebound through reserved Firecracker slots. - #[serde(default)] - volume: bool, + #[serde(default, rename = "volume", with = "extra_drive_kind_serde")] + kind: ExtraDriveKind, }, } +/// How a drive is attached to Firecracker and whether captures commit it. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ExtraDriveKind { + /// Attached under its own drive ID and committed to snapshot manifests. + #[default] + Physical, + /// Added to a snapshot launch through a reserved slot. Captures commit it + /// to the manifest, and it keeps its slot in later launches. + LaunchSlot, + /// Persistent volume: omitted from committed attached-drive manifests and + /// rebound through reserved Firecracker slots. + Volume, +} + +/// Physical drives and volumes keep the legacy boolean `volume` field, so +/// persisted configs stay readable in both directions. +mod extra_drive_kind_serde { + use serde::{Deserialize, Deserializer, Serialize, Serializer}; + + use super::ExtraDriveKind; + + #[derive(Serialize, Deserialize)] + #[serde(untagged)] + enum Repr { + Legacy(bool), + Named(Named), + } + + #[derive(Serialize, Deserialize)] + #[serde(rename_all = "snake_case")] + enum Named { + LaunchSlot, + } + + pub(super) fn serialize( + kind: &ExtraDriveKind, + serializer: S, + ) -> Result { + match kind { + ExtraDriveKind::Physical => Repr::Legacy(false), + ExtraDriveKind::LaunchSlot => Repr::Named(Named::LaunchSlot), + ExtraDriveKind::Volume => Repr::Legacy(true), + } + .serialize(serializer) + } + + pub(super) fn deserialize<'de, D: Deserializer<'de>>( + deserializer: D, + ) -> Result { + Ok(match Repr::deserialize(deserializer)? { + Repr::Legacy(false) => ExtraDriveKind::Physical, + Repr::Named(Named::LaunchSlot) => ExtraDriveKind::LaunchSlot, + Repr::Legacy(true) => ExtraDriveKind::Volume, + }) + } +} + impl ExtraDrive { pub fn try_new_overlaybd( drive_id: impl Into, @@ -76,7 +131,7 @@ impl ExtraDrive { virtual_size: None, sub_path, snapshot_output_dir: None, - volume: false, + kind: ExtraDriveKind::Physical, }) } @@ -139,10 +194,27 @@ impl ExtraDrive { } } - pub(crate) fn is_volume(&self) -> bool { + pub fn kind(&self) -> ExtraDriveKind { match self { - Self::Overlaybd { volume, .. } => *volume, + Self::Overlaybd { kind, .. } => *kind, + } + } + + pub(crate) fn is_volume(&self) -> bool { + self.kind() == ExtraDriveKind::Volume + } + + /// Whether the drive is bound to a reserved Firecracker slot. + pub(crate) fn uses_reserved_slot(&self) -> bool { + self.kind() != ExtraDriveKind::Physical + } + + pub fn with_kind(&self, kind: ExtraDriveKind) -> Self { + let mut drive = self.clone(); + match &mut drive { + Self::Overlaybd { kind: current, .. } => *current = kind, } + drive } pub(crate) fn with_volume_snapshot_output_dir(&self, output_dir: Option) -> Self { @@ -163,7 +235,7 @@ impl ExtraDrive { virtual_size: *virtual_size, sub_path: sub_path.clone(), snapshot_output_dir: output_dir, - volume: true, + kind: ExtraDriveKind::Volume, }, } } @@ -177,7 +249,7 @@ impl ExtraDrive { virtual_size, sub_path, snapshot_output_dir, - volume, + kind, .. } => Self::Overlaybd { drive_id: drive_id.clone(), @@ -187,7 +259,7 @@ impl ExtraDrive { virtual_size: *virtual_size, sub_path: sub_path.clone(), snapshot_output_dir: snapshot_output_dir.clone(), - volume: *volume, + kind: *kind, }, } } @@ -204,7 +276,7 @@ impl ExtraDrive { read_only, mount_path, snapshot_output_dir, - volume, + kind, sub_path, .. } => Ok(Self::Overlaybd { @@ -215,7 +287,7 @@ impl ExtraDrive { virtual_size: Some(virtual_size), sub_path: sub_path.clone(), snapshot_output_dir: snapshot_output_dir.clone(), - volume: *volume, + kind: *kind, }), } } @@ -386,6 +458,33 @@ mod tests { assert!(drive.snapshot_output_dir().is_none()); } + #[test] + fn drive_kind_keeps_legacy_volume_field() { + let drive = ExtraDrive::try_new_overlaybd("data", "/tmp/image.json", false) + .expect("drive should parse"); + for (kind, encoded) in [ + (ExtraDriveKind::Physical, serde_json::json!(false)), + (ExtraDriveKind::LaunchSlot, serde_json::json!("launch_slot")), + (ExtraDriveKind::Volume, serde_json::json!(true)), + ] { + let drive = drive.with_kind(kind); + let value = serde_json::to_value(&drive).expect("drive should serialize"); + assert_eq!(value["Overlaybd"]["volume"], encoded); + let decoded: ExtraDrive = + serde_json::from_value(value).expect("drive should deserialize"); + assert_eq!(decoded, drive); + } + + let mut legacy = serde_json::to_value(&drive).expect("drive should serialize"); + legacy["Overlaybd"] + .as_object_mut() + .expect("drive fields") + .remove("volume"); + let decoded: ExtraDrive = + serde_json::from_value(legacy).expect("drive without kind should deserialize"); + assert_eq!(decoded.kind(), ExtraDriveKind::Physical); + } + #[test] fn overlaybd_drive_rejects_firecracker_incompatible_drive_id() { for drive_id in ["resume-data", "resume/data", "resume data", "résumé"] { diff --git a/src/types/mod.rs b/src/types/mod.rs index c41c4b72e..8c2a346dd 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -10,7 +10,7 @@ mod volume_mode; pub use extra_drive::{ normalize_mount_path, normalize_mount_path_for_drive, validate_drive_id, validate_mount_path, - validate_sub_path, ExtraDrive, + validate_sub_path, ExtraDrive, ExtraDriveKind, }; pub(crate) use extra_drive::{ROOTFS_DRIVE_ID, USER_ROOTFS_DRIVE_ID, VOLUME_DRIVE_SLOT_PREFIX}; pub use id::SandboxId; diff --git a/tests/integration/snapshot_attached_drive.rs b/tests/integration/snapshot_attached_drive.rs index df65d14a5..ce8c6d850 100644 --- a/tests/integration/snapshot_attached_drive.rs +++ b/tests/integration/snapshot_attached_drive.rs @@ -6,8 +6,8 @@ use crate::common; use agentenv::cfg::ConfigManager; use agentenv::sandbox::{ - ExtraDrive, FirecrackerSandbox, FirecrackerSandboxConfig, OverlaybdConfig, SandboxBackend, - SandboxExecutor, SandboxLaunchConfig, + ExtraDrive, ExtraDriveKind, FirecrackerSandbox, FirecrackerSandboxConfig, OverlaybdConfig, + SandboxBackend, SandboxExecutor, SandboxLaunchConfig, }; use agentenv::snapshot::{ CommittedAttachedDrive, OverlaybdLayerRef, SnapshotAlias, SnapshotId, SnapshotManager, @@ -131,13 +131,7 @@ async fn assert_snapshot_attached_drive_is_visible( Ok(()) } -async fn publish_sandbox_snapshot_with_attached_drive( - snapshot_manager: &SnapshotManager, - alias: &str, - image_config: &Path, - read_only: bool, - setup_cmd: &str, -) -> Result { +fn small_sandbox_config(image_config: &Path) -> Result { let user_image = OverlaybdConfig { image_config_path: image_config.to_path_buf(), read_only: false, @@ -146,6 +140,17 @@ async fn publish_sandbox_snapshot_with_attached_drive( let mut config = FirecrackerSandboxConfig::from_global_config_with_user_image(user_image)?; config.vcpu_count = 1; config.mem_size_mib = 128; + Ok(config) +} + +async fn publish_sandbox_snapshot_with_attached_drive( + snapshot_manager: &SnapshotManager, + alias: &str, + image_config: &Path, + read_only: bool, + setup_cmd: &str, +) -> Result { + let mut config = small_sandbox_config(image_config)?; config.common.extra_drives = vec![ExtraDrive::try_new_overlaybd( "data", image_config.to_path_buf(), @@ -161,7 +166,14 @@ async fn publish_sandbox_snapshot_with_attached_drive( let captured = SandboxBackend::snapshot(&mut sandbox).await?; sandbox.stop().await?; + publish_capture(snapshot_manager, alias, captured).await +} +async fn publish_capture( + snapshot_manager: &SnapshotManager, + alias: &str, + captured: agentenv::sandbox::CapturedSandboxSnapshot, +) -> Result { let metadata = SnapshotPublishMetadata { id: SnapshotId::generate(), alias: Some(SnapshotAlias::parse(alias)?), @@ -218,7 +230,7 @@ async fn frozen_volumes_survive_deletion_without_vm_snapshot() -> Result<()> { virtual_size: Some(1024 * 1024 * 1024), sub_path: None, snapshot_output_dir: Some(image_config.parent().unwrap().to_path_buf()), - volume: true, + kind: ExtraDriveKind::Volume, }); } let mut worker = FirecrackerSandbox::new(config.clone())?; @@ -499,6 +511,115 @@ async fn writable_attached_overlaybd_drive_preserves_metadata_and_mount_mode() - Ok(()) } +async fn assert_guest_files(sandbox: &FirecrackerSandbox, files: &[&str]) -> Result<()> { + for file in files { + let output = sandbox + .run_command("cat", &[&format!("/mnt/data/{file}")]) + .await?; + anyhow::ensure!( + output.exit_code == 0 && output.stdout.trim() == *file, + "unexpected /mnt/data/{file}: {output:?}" + ); + } + Ok(()) +} + +async fn write_guest_file(sandbox: &FirecrackerSandbox, file: &str) -> Result<()> { + let output = sandbox + .run_command( + "sh", + &["-lc", &format!("echo {file} > /mnt/data/{file} && sync")], + ) + .await?; + anyhow::ensure!(output.exit_code == 0, "{}", output.stderr); + Ok(()) +} + +/// A drive added to a snapshot launch keeps its reserved slot through +/// capture, relaunch, fork, and a second capture. +#[tokio::test] +async fn launch_time_attached_drive_keeps_its_slot_across_captures() -> Result<()> { + common::setup().await; + let store = tempdir()?; + let (_, snapshot_manager) = test_parts(store.path()); + let (_, image_config) = common::default_rootfs_template_build_spec_with_image_config(); + + let mut base = FirecrackerSandbox::new(small_sandbox_config(&image_config)?)?; + base.start().await?; + let captured = SandboxBackend::snapshot(&mut base).await?; + base.stop().await?; + let base_alias = unique_alias_for_test("launch_drive_base"); + publish_capture(&snapshot_manager, &base_alias, captured).await?; + + let runnable = snapshot_manager + .load_runnable(&base_alias) + .await? + .ok_or_else(|| anyhow!("base snapshot should exist"))?; + let launch_config = SandboxLaunchConfig { + extra_drives: vec![ + ExtraDrive::try_new_overlaybd("data", image_config.clone(), false)? + .with_kind(ExtraDriveKind::LaunchSlot), + ], + ..SandboxLaunchConfig::default() + }; + let mut first = FirecrackerSandbox::from_snapshot(&runnable, &launch_config)?; + first.start().await?; + assert_snapshot_attached_drive_is_visible(&first, "data", false).await?; + write_guest_file(&first, "first").await?; + let captured = SandboxBackend::snapshot(&mut first).await?; + first.stop().await?; + let first_alias = unique_alias_for_test("launch_drive_first"); + publish_capture(&snapshot_manager, &first_alias, captured).await?; + + let runnable = snapshot_manager + .load_runnable(&first_alias) + .await? + .ok_or_else(|| anyhow!("first snapshot should exist"))?; + let manifest = runnable.manifest(); + assert_eq!(manifest.attached_drives.len(), 1); + assert_eq!(manifest.effective_physical_extra_drive_count(), 0); + assert_eq!(manifest.free_drive_slots() + 1, manifest.volume_drive_slots); + + let mut second = FirecrackerSandbox::from_snapshot(&runnable, &SandboxLaunchConfig::default())?; + second.start().await?; + let second_result: Result<()> = async { + assert_snapshot_attached_drive_is_visible(&second, "data", false).await?; + assert_guest_files(&second, &["first"]).await?; + write_guest_file(&second, "second").await?; + + // A fork child resumes from the paused state while its parent runs on. + let paused = second.pause().await?; + second.resume().await?; + let mut child = FirecrackerSandbox::resume_from_snapshot_config(&paused).await?; + let child_result = assert_guest_files(&child, &["first", "second"]).await; + child.stop().await?; + child_result?; + + let captured = SandboxBackend::snapshot(&mut second).await?; + let second_alias = unique_alias_for_test("launch_drive_second"); + publish_capture(&snapshot_manager, &second_alias, captured).await?; + let runnable = snapshot_manager + .load_runnable(&second_alias) + .await? + .ok_or_else(|| anyhow!("second snapshot should exist"))?; + assert_eq!( + runnable.manifest().effective_physical_extra_drive_count(), + 0 + ); + let mut third = + FirecrackerSandbox::from_snapshot(&runnable, &SandboxLaunchConfig::default())?; + third.start().await?; + let third_result = assert_guest_files(&third, &["first", "second"]).await; + third.stop().await?; + third_result + } + .await; + let second_stop = second.stop().await; + second_result?; + second_stop?; + Ok(()) +} + fn unique_alias_for_test(prefix: &str) -> String { format!( "{prefix}_{:x}{:x}",