mod capture_receipt;
mod cloud_hypervisor;
mod environment;

use super::cleanup_helpers::*;
use super::disk::*;
use super::node_disk_metrics::SandboxDiskUsageLookup;
use super::validation::*;
use super::*;
use crate::snapshot_request::SnapshotStorage;

use pplx_space_common_core::host::linux::identities::gids;
use pplx_space_common_core::types::SandboxId;
use std::os::unix::fs::{MetadataExt, PermissionsExt};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use pplx_space_common_core::host::linux::paths::INTERNAL_OPERATION_SNAPSHOT_PREFIX as INTERNAL_INSTANCE_SNAPSHOT_PREFIX;

#[cfg(test)]
type CaptureIo = std::sync::Arc<
    dyn Fn() -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), Status>> + Send>>
        + Send
        + Sync,
>;

#[cfg(test)]
#[derive(Clone)]
pub(super) struct EncryptedCaptureTestIo {
    pub capture: CaptureIo,
    pub drain: CaptureIo,
    pub cache: PathBuf,
}

enum CaptureDestination {
    Plain(crate::runtime_paths::CreateSnapshotPaths),
    Encrypted(
        pplx_space_common_core::types::NodeSpacePath,
        capture_receipt::CaptureReceipt,
    ),
}

enum RestoreSource {
    Plain(crate::runtime_paths::LoadSnapshotPaths),
    Encrypted(pplx_space_common_core::types::NodeSpacePath),
}

const BTRFS_SNAPSHOTS_DIR_NAME: &str = "snapshots";

fn prepare_btrfs_snapshot_source(instance_dir: &Path) -> Result<(), Status> {
    ensure_snapshot_export_permissions(instance_dir)
}

fn capture_instance_snapshot(
    instance_dir: &Path,
    snapshots_dir: &Path,
    name: &std::ffi::OsStr,
    writable: bool,
    encrypted: bool,
) -> Result<(), Status> {
    if !instance_dir.exists() {
        return Err(Status::not_found(format!(
            "instance dir not found: {}",
            instance_dir.display()
        )));
    }
    if !encrypted {
        prepare_btrfs_snapshot_source(instance_dir)?;
    }
    let snapshots = super::trusted_paths::create_directory(snapshots_dir)?;
    super::trusted_paths::set_permissions(&snapshots, Some(gids::SPACE_CACHE_RW), 0o2770)?;
    let source = super::trusted_paths::open_directory(instance_dir)?;
    if encrypted {
        return capture_encrypted_instance(&source, &snapshots, name, writable);
    }
    super::artifact_fs::btrfs::snapshot(&source, &snapshots, name, !writable)
        .map_err(super::artifact_fs::error)
}

fn capture_encrypted_instance(
    source: &std::fs::File,
    snapshots: &std::fs::File,
    name: &std::ffi::OsStr,
    writable: bool,
) -> Result<(), Status> {
    use super::{artifact_fs, trusted_paths};
    use nix::sys::stat::{Mode, mkdirat};
    use nix::unistd::{UnlinkatFlags, unlinkat};

    // Export only the copy. The live encrypted backing stays private, and the
    // copy stays inaccessible to cache peers until its permissions and RO flag
    // are complete. Abandoned stages are reclaimed with the sandbox cache.
    let temporary = format!(".encrypted-capture-{}", uuid::Uuid::now_v7());
    mkdirat(
        snapshots,
        temporary.as_str(),
        Mode::from_bits_truncate(0o700),
    )
    .map_err(artifact_fs::error)?;
    let stage =
        artifact_fs::directory(snapshots, temporary.as_ref()).map_err(artifact_fs::error)?;
    let metadata = stage.metadata().map_err(artifact_fs::error)?;
    if metadata.uid() != nix::unistd::geteuid().as_raw() || metadata.mode() & 0o077 != 0 {
        return Err(Status::failed_precondition(
            "encrypted capture stage is not private",
        ));
    }
    let result = (|| {
        artifact_fs::btrfs::snapshot(source, &stage, "capture".as_ref(), false)
            .map_err(artifact_fs::error)?;
        let capture =
            artifact_fs::directory(&stage, "capture".as_ref()).map_err(artifact_fs::error)?;
        trusted_paths::export_directory_permissions(&capture, Some(gids::SPACE_CACHE_RW))?;
        // Publish through SNAP_CREATE_V2: permissions and read-only state are
        // visible together, and an existing name is never replaced. A read-only
        // subvolume cannot be renamed on every supported kernel.
        artifact_fs::btrfs::snapshot(&capture, snapshots, name, !writable)
            .map_err(artifact_fs::error)?;
        snapshots.sync_all().map_err(artifact_fs::error)
    })();
    if let Err(error) = artifact_fs::remove_at(&stage, "capture".as_ref()) {
        tracing::warn!(%error, "encrypted capture stage retained for sandbox cleanup");
    }
    // Only an empty directory may be removed by name; never recursively delete
    // a replacement in the shared parent after releasing the private stage.
    if let Err(error) = unlinkat(snapshots, temporary.as_str(), UnlinkatFlags::RemoveDir) {
        tracing::warn!(%error, "encrypted capture staging directory retained for cleanup");
    }
    result
}

#[cfg(test)]
mod capture_tests;

pub(super) fn btrfs_snapshot_tier_and_timestamp(name: &str) -> Option<(&str, &str)> {
    if name.starts_with("d-") || name.starts_with("f-") {
        return name
            .strip_prefix("d-")
            .or_else(|| name.strip_prefix("f-"))
            .map(|rest| rest.split_once('-').unwrap_or(("", rest)));
    }

    name.split_once('-')
}

pub(super) fn list_btrfs_snapshot_dirs_in(
    sandbox_root: &Path,
    budget: &mut DirectoryBudget,
) -> std::io::Result<Vec<PathBuf>> {
    budget.directories(&sandbox_root.join(BTRFS_SNAPSHOTS_DIR_NAME))
}

/// A failed/expired size probe cannot hide an enumerated identity used for
/// retention, deletion or rollback.
fn measure_snapshots(
    paths: Vec<PathBuf>,
    measurement_budget: Duration,
    mut measure: impl FnMut(&str, Duration) -> Result<Option<BtrfsDiskUsage>, String>,
) -> Vec<crate::proto::BtrfsSnapshotInfo> {
    let started = Instant::now();
    let mut snapshots = Vec::new();
    for path in paths {
        let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
            continue;
        };
        let Some((tier, timestamp)) = btrfs_snapshot_tier_and_timestamp(name) else {
            continue;
        };
        let remaining = measurement_budget.saturating_sub(started.elapsed());
        let usage = if remaining.is_zero() {
            None
        } else {
            match measure(&path.to_string_lossy(), remaining) {
                Ok(usage) => usage,
                Err(error) => {
                    tracing::warn!(path = %path.display(), %error, "snapshot size unavailable; retaining inventory entry");
                    None
                }
            }
        };
        snapshots.push(crate::proto::BtrfsSnapshotInfo {
            name: name.to_owned(),
            tier: tier.to_owned(),
            timestamp: timestamp.to_owned(),
            exclusive_bytes: usage.as_ref().map(|usage| usage.exclusive_bytes),
            referenced_bytes: usage.map(|usage| usage.total_bytes),
        });
    }
    snapshots.sort_by(|a, b| a.name.cmp(&b.name));
    snapshots
}

#[cfg(target_os = "linux")]
pub(super) fn bind_mount_file(source: &Path, target: &Path) -> Result<(), Status> {
    nix::mount::mount(
        Some(source),
        target,
        None::<&Path>,
        nix::mount::MsFlags::MS_BIND,
        None::<&Path>,
    )
    .map_err(|error| {
        Status::internal(format!(
            "mount --bind {} {} failed: {error}",
            source.display(),
            target.display()
        ))
    })
}

#[cfg(not(target_os = "linux"))]
pub(super) fn bind_mount_file(source: &Path, target: &Path) -> Result<(), Status> {
    Err(Status::unimplemented(format!(
        "bind mount {} -> {} is only supported on Linux",
        source.display(),
        target.display()
    )))
}

/// Jail subdirectory the encrypted snapshot volume is bind-mounted at, so the
/// jailed Firecracker reads/writes vmstate+memory inside it. Kept short and
/// chroot-relative; the FC API paths below are `/{ENC_SNAPSHOT_JAIL_SUBDIR}/…`.
pub(super) const ENC_SNAPSHOT_JAIL_SUBDIR: &str = "encsnap";

/// Upper bound on how many stacked `encsnap` binds the seal path drains — a
/// restore leaves one, a re-suspend can transiently stack a second; the cap
/// guards against an unexpected unmount that never reports "not mounted".
const MAX_ENC_SNAPSHOT_BIND_DRAIN: usize = 8;

/// Retries of the post-capture runtime stop. Firecracker can still be in kernel
/// exit (KVM teardown) when the supervisor's cgroup removal gives up; the stop
/// is idempotent, so re-issuing it once the exit finishes completes the seal.
const ENC_CAPTURE_STOP_ATTEMPTS: usize = 8;
const ENC_CAPTURE_STOP_RETRY_DELAY: std::time::Duration = std::time::Duration::from_millis(250);

/// File names Firecracker writes into (suspend) and reads from (restore) the
/// encrypted snapshot volume. Shared by the FC API paths and the load-side
/// re-own list so they cannot drift apart.
const ENC_SNAPSHOT_VMSTATE_NAME: &str = "vmstate";

const ENC_SNAPSHOT_MEMORY_NAME: &str = "memory";

/// Files the encrypted-snapshot **load** must re-own to the current sandbox
/// uid. The files inside the reopened volume carry the ownership of the
/// *suspending* run, and a wake allocates a fresh slot-derived uid — so
/// re-owning only the bind root lets the jailed Firecracker traverse the
/// directory but still fail opening `vmstate`/`memory` with EACCES
/// ("Failed to open snapshot file: Permission denied"). The suspend side
/// does not need this: Firecracker creates the files itself under the
/// then-current uid.
fn reown_encrypted_snapshot_inputs(volume: &std::fs::File, uid: u32) -> Result<(), Status> {
    use super::trusted_paths::{open_regular_file, reown};
    let vmstate = open_regular_file(volume, ENC_SNAPSHOT_VMSTATE_NAME.as_ref())?;
    let memory = open_regular_file(volume, ENC_SNAPSHOT_MEMORY_NAME.as_ref())?;
    reown(&vmstate, uid)?;
    reown(&memory, uid)?;
    reown(volume, uid)
}

/// Pin both ends of the encrypted-snapshot bind before re-owning the volume.
/// The jail root belongs to the sandbox uid, so its mountpoint is opened
/// relative to the pinned jail rather than looked up again by pathname.
fn pin_encrypted_snapshot_bind(
    mount_dir: &Path,
    jail_root: &Path,
    reown: impl FnOnce(&std::fs::File, u32) -> Result<(), Status>,
) -> Result<(std::fs::File, std::fs::File), Status> {
    let volume = super::trusted_paths::open_directory(mount_dir)?;
    let jail = super::trusted_paths::open_directory(jail_root)?;
    let uid = jail
        .metadata()
        .map_err(|e| Status::internal(e.to_string()))?
        .uid();
    let target =
        super::trusted_paths::open_bind_directory(&jail, ENC_SNAPSHOT_JAIL_SUBDIR.as_ref())?;
    reown(&volume, uid)?;
    Ok((volume, target))
}

fn stage_plaintext_snapshot_inputs(
    source_dir: &Path,
    jail_root: &Path,
) -> Result<(std::fs::File, std::fs::File, std::fs::File), Status> {
    let source = super::trusted_paths::open_directory(source_dir)?;
    let memory = super::trusted_paths::open_regular_file(&source, "memory".as_ref())?;
    let jail = super::trusted_paths::open_directory(jail_root)?;
    let vmstate_source = super::trusted_paths::open_regular_file(&source, "vmstate".as_ref())?;
    let vmstate =
        crate::contained_copy::stage_open_file(&vmstate_source, &jail, "vmstate".as_ref())
            .map_err(|e| Status::internal(format!("stage vmstate: {e}")))?;
    Ok((vmstate, memory, jail))
}

fn chgrp_path(path: &str, gid: u32) -> Result<(), Status> {
    std::os::unix::fs::chown(path, None, Some(gid))
        .map_err(|error| Status::internal(format!("chgrp {gid} {path} failed: {error}")))
}

fn chmod_path(path: &str, mode: u32) -> Result<(), Status> {
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(mode))
        .map_err(|error| Status::internal(format!("chmod {mode:o} {path} failed: {error}")))
}

impl NodeManagerService {
    pub(super) async fn get_vm_state(
        &self,
        request: Request<GetVmStateRequest>,
    ) -> Result<Response<GetVmStateResponse>, Status> {
        self.authorize_peer(&request, "GetVmState", ALLOW_NODE_STATUS)?;
        let sandbox_id = request
            .get_ref()
            .sandbox_id
            .parse::<SandboxId>()
            .map_err(|error| Status::invalid_argument(error.to_string()))?;
        let state = tokio::time::timeout(FIRECRACKER_VM_STATE_TIMEOUT, async {
            let execution = self.recorded_execution(sandbox_id.to_string()).await?;
            // Only absent historical metadata selects Firecracker. Present but
            // invalid evidence cannot select a different control protocol.
            let runtime = execution.as_ref().map_or(
                pplx_space_common_core::types::RuntimeKind::Firecracker,
                |execution| execution.runtime().kind,
            );
            let control = self
                .runtime_control(runtime, &sandbox_id.to_string())
                .await?;
            control.state().await.map_err(|error| {
                Status::unavailable(format!("runtime state probe failed: {error}"))
            })
        })
        .await
        .map_err(|_| Status::deadline_exceeded("runtime state probe timed out"))??;
        use pplx_space_node_runtime_core::vm_state::VmState;
        let state = match state {
            VmState::NotStarted => crate::proto::VmState::NotStarted,
            VmState::Running => crate::proto::VmState::Running,
            VmState::Paused => crate::proto::VmState::Paused,
            VmState::Shutdown => crate::proto::VmState::Shutdown,
        };
        Ok(Response::new(GetVmStateResponse {
            state: state.into(),
        }))
    }

    pub(super) async fn pause_sandbox(
        &self,
        request: Request<PauseSandboxRequest>,
    ) -> Result<Response<PauseSandboxResponse>, Status> {
        self.authorize_peer(&request, "PauseSandbox", ALLOW_NODE_ADMIN_OR_SBXMGR)?;
        let req = request.into_inner();
        let balloon_hint_timeout = req.balloon_hint_timeout_ms.map(balloon_hint_timeout);
        self.set_vm_paused(
            req.sandbox_id,
            req.execution_fence,
            true,
            balloon_hint_timeout,
        )
        .await?;
        Ok(Response::new(PauseSandboxResponse {}))
    }

    pub(super) async fn resume_sandbox(
        &self,
        request: Request<ResumeSandboxRequest>,
    ) -> Result<Response<ResumeSandboxResponse>, Status> {
        self.authorize_peer(&request, "ResumeSandbox", ALLOW_NODE_ADMIN_OR_SBXMGR)?;
        let req = request.into_inner();
        self.set_vm_paused(req.sandbox_id, req.execution_fence, false, None)
            .await?;
        Ok(Response::new(ResumeSandboxResponse {}))
    }

    async fn set_vm_paused(
        &self,
        sandbox_id: String,
        fence: Option<crate::proto::ExecutionFence>,
        paused: bool,
        balloon_hint_timeout: Option<std::time::Duration>,
    ) -> Result<(), Status> {
        validate_sandbox_id(&sandbox_id)?;
        let owner = sandbox_id.clone();
        let service = self.clone();
        self.run_owned_sandbox_work(&owner, async move {
            let _gate = service.lock_sandbox_vm_cgroup_launch(&sandbox_id).await;
            service.check_execution(&sandbox_id, fence.as_ref(), false)?;
            let execution = service.recorded_execution(sandbox_id.clone()).await?;
            let runtime = execution.as_ref().map_or(
                pplx_space_common_core::types::RuntimeKind::Firecracker,
                |execution| execution.runtime().kind,
            );
            if runtime == pplx_space_common_core::types::RuntimeKind::Firecracker {
                let api_socket =
                    format!("{JAIL_DIR}/firecracker/{sandbox_id}/root/firecracker.sock");
                let client = crate::fc_api::FcApiClient::new(api_socket);
                if let Some(hint_timeout) = balloon_hint_timeout {
                    prepare_balloon_for_snapshot(&client, &sandbox_id, hint_timeout).await;
                }
                return set_firecracker_paused(&client, paused, VM_PAUSE_STATE_TIMEOUT).await;
            }
            let client = service.runtime_control(runtime, &sandbox_id).await?;
            tokio::time::timeout(VM_PAUSE_STATE_TIMEOUT, client.set_paused(paused))
                .await
                .map_err(|_| Status::deadline_exceeded("runtime pause state change timed out"))?
                .map_err(|error| {
                    super::runtime_control::operation_error("set runtime pause state", error)
                })
        })
        .await
    }

    fn check_snapshot_execution(
        &self,
        sandbox_id: SandboxId,
        fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence,
    ) -> Result<(), Status> {
        self.check_parsed_execution(&sandbox_id.to_string(), Some(fence), false)?;
        Ok(())
    }

    pub(super) async fn capture_runtime_state(
        &self,
        request: Request<StateCapture>,
    ) -> Result<Response<CaptureVmStateResponse>, Status> {
        let selected_environment = request.get_ref().environment_reference;
        let sandbox_id = request.get_ref().sandbox_id.as_str().to_owned();
        let service = self.clone();
        wait_for_capture(self.run_owned_sandbox_work(&sandbox_id, async move {
            let _vm_guard = service
                .lock_sandbox_vm_cgroup_launch(request.get_ref().sandbox_id.as_str())
                .await;
            // Reject malformed paths before any metadata IO, as on legacy nodes.
            let req = request.get_ref();
            let mut destination = match &req.storage {
                SnapshotStorage::Encrypted(mount) => {
                    trusted_paths::validate_snapshot_volume(
                        req.sandbox_id.as_str(),
                        crate::snapshot_request::path_str(mount)?,
                    )?;
                    let receipt = service
                        .open_capture_receipt(
                            req.sandbox_id.id(),
                            req.capture_directory
                                .as_ref()
                                .ok_or_else(|| {
                                    Status::invalid_argument(
                                        "encrypted capture requires its staging directory",
                                    )
                                })?
                                .as_path(),
                        )
                        .await?;
                    CaptureDestination::Encrypted(mount.clone(), receipt)
                }
                SnapshotStorage::Plain { files, directory } => CaptureDestination::Plain(
                    crate::runtime_paths::RuntimePaths::firecracker().create_snapshot_paths(
                        req.sandbox_id.as_str(),
                        files,
                        directory.as_ref(),
                    )?,
                ),
            };
            // Claim first so a fence refusal leaves durable evidence that this
            // exact attempt never reached the runtime capture operation.
            let result = async {
                let req = request.get_ref();
                service.check_snapshot_execution(
                    request.get_ref().sandbox_id.id(),
                    request.get_ref().execution_fence,
                )?;
                if service
                    .pool
                    .get(req.sandbox_id.as_str())
                    .map_err(slot_store_status)?
                    .is_none()
                {
                    return Err(Status::unavailable("snapshot lease is no longer active"));
                }
                service
                    .verify_snapshot_environment(
                        req.sandbox_id.as_str().to_owned(),
                        selected_environment,
                    )
                    .await?;
                let execution = service
                    .recorded_execution(req.sandbox_id.as_str().to_owned())
                    .await?;
                if let Some(boot_source) = snapshot_boot_source(execution.as_ref())? {
                    return service
                        .capture_cloud_hypervisor_state(request, &mut destination, boot_source)
                        .await;
                }
                service
                    .create_snapshot_inner(request, &mut destination)
                    .await
            }
            .await;
            if result.is_err()
                && let CaptureDestination::Encrypted(_, receipt) = destination
                && let Err(error) = receipt.settle_error().await
            {
                tracing::warn!(%error, "cannot persist capture error outcome; retaining existing evidence");
            }
            result
        }))
        .await
    }

    async fn create_snapshot_inner(
        &self,
        request: Request<StateCapture>,
        destination: &mut CaptureDestination,
    ) -> Result<Response<CaptureVmStateResponse>, Status> {
        let req = request.get_ref();
        let paths = match destination {
            CaptureDestination::Encrypted(mount, receipt) => {
                return self
                    .create_encrypted_snapshot(req.sandbox_id.id(), mount, receipt)
                    .await;
            }
            CaptureDestination::Plain(paths) => paths,
        };
        let dest_dir = paths.dest_dir.clone();
        let jail_root = format!("{JAIL_DIR}/firecracker/{}/root", req.sandbox_id);
        let api_socket = format!("{jail_root}/firecracker.sock");
        let client = crate::fc_api::FcApiClient::new(api_socket);
        let lease = self
            .pool
            .get(req.sandbox_id.as_str())
            .map_err(slot_store_status)?
            .ok_or_else(|| Status::unavailable("snapshot lease is no longer active"))?;
        let total_started = Instant::now();
        tracing::debug!(
            sandbox_id = %req.sandbox_id,
            snapshot_path = %paths.vmstate_name,
            mem_file_path = %paths.memory_name,
            dest_dir = ?dest_dir,
            "node snapshot create started"
        );
        let phase = Instant::now();
        client
            .create_snapshot(
                format!("/{}", paths.vmstate_name),
                format!("/{}", paths.memory_name),
            )
            .await
            .map_err(|e| Status::internal(format!("create snapshot failed: {e}")))?;
        self.adopt_validated_lease(req.sandbox_id.as_str(), &lease)?;
        tracing::debug!(
            sandbox_id = %req.sandbox_id,
            snapshot_path = %paths.vmstate_name,
            mem_file_path = %paths.memory_name,
            elapsed_ms = phase.elapsed().as_millis(),
            "node snapshot firecracker create complete"
        );

        // Reflink snapshot files from jail chroot to dest_dir (instant on btrfs).
        // The rootfs is bind-mounted from the instance subvolume — use the
        // instance path so the reflink source is always a regular file.
        // Offloaded: the reflink is a subprocess per file (with a full-copy
        // fallback) and the export chgrp is recursive — none of it belongs on
        // a runtime worker.
        if let Some(dest_dir) = paths.dest_dir.take() {
            let jail_root = jail_root.clone();
            let vmstate_name = paths.vmstate_name.clone();
            let memory_name = paths.memory_name.clone();
            let sandbox_id = req.sandbox_id.as_str().to_owned();
            self.off_runtime("export snapshot artifacts", move || {
                let destination = super::trusted_paths::create_directory(&dest_dir)?;
                for (name, src_path) in [
                    ("vmstate", format!("{jail_root}/{vmstate_name}")),
                    ("memory", format!("{jail_root}/{memory_name}")),
                    (
                        "rootfs.ext4",
                        format!("{CACHE_DIR}/{sandbox_id}/instance/rootfs.ext4"),
                    ),
                ] {
                    let copied = super::artifact_fs::copy(
                        std::path::Path::new(&src_path),
                        &destination,
                        name.as_ref(),
                    )?;
                    super::trusted_paths::set_permissions(
                        &copied,
                        Some(pplx_space_common_core::host::linux::identities::gids::SPACE_CACHE_RW),
                        0o660,
                    )?;
                }
                super::trusted_paths::set_permissions(
                    &destination,
                    Some(pplx_space_common_core::host::linux::identities::gids::SPACE_CACHE_RW),
                    0o2770,
                )?;
                destination.sync_all().map_err(super::artifact_fs::error)
            })
            .await?;
        }

        tracing::debug!(
            sandbox_id = %req.sandbox_id,
            dest_dir = ?dest_dir,
            elapsed_ms = total_started.elapsed().as_millis(),
            "node snapshot create complete"
        );
        Ok(Response::new(CaptureVmStateResponse {}))
    }

    /// Claim the pinned attempt directory off the runtime before capture admission.
    async fn open_capture_receipt(
        &self,
        sandbox: SandboxId,
        path: &Path,
    ) -> Result<capture_receipt::CaptureReceipt, Status> {
        let path = path.to_owned();
        #[cfg(test)]
        if let Some(io) = &self.encrypted_capture_io {
            let cache = io.cache.clone();
            return self
                .off_runtime("claim capture receipt", move || {
                    capture_receipt::CaptureReceipt::open_under(&cache, sandbox, &path)
                })
                .await;
        }
        self.off_runtime("claim capture receipt", move || {
            capture_receipt::CaptureReceipt::open(sandbox, &path)
        })
        .await
    }

    async fn create_encrypted_snapshot(
        &self,
        typed_sandbox_id: SandboxId,
        mount_dir: &pplx_space_common_core::types::NodeSpacePath,
        receipt: &mut capture_receipt::CaptureReceipt,
    ) -> Result<Response<CaptureVmStateResponse>, Status> {
        let sandbox_id = typed_sandbox_id.to_string();
        let jail_bind =
            format!("{JAIL_DIR}/firecracker/{sandbox_id}/root/{ENC_SNAPSHOT_JAIL_SUBDIR}");
        let lease = self
            .pool
            .get(&sandbox_id)
            .map_err(slot_store_status)?
            .ok_or_else(|| Status::unavailable("snapshot lease is no longer active"))?;

        // The outer owned capture retains the VM/work guards across caller loss.
        // This same task holds device exclusion through capture, stop, receipt
        // persistence and bind drain, so cleanup cannot interleave with writes.
        let _guard = self
            .sandbox_device_stack_lock(&typed_sandbox_id)
            .await
            .lock_owned()
            .await;
        self.write_encrypted_snapshot_files(&sandbox_id, mount_dir, receipt)
            .await?;
        self.adopt_validated_lease(&sandbox_id, &lease)?;
        stop_runtime_for_encrypted_seal(self, &sandbox_id).await?;
        receipt.commit().await?;
        self.drain_encrypted_capture_bind(sandbox_id, jail_bind)
            .await;
        Ok(Response::new(CaptureVmStateResponse {}))
    }

    async fn drain_encrypted_capture_bind(&self, sandbox_id: String, jail_bind: String) {
        #[cfg(test)]
        if let Some(io) = &self.encrypted_capture_io {
            if let Err(error) = (io.drain)().await {
                tracing::warn!(%error, "encrypted capture bind drain failed");
            }
            return;
        }
        let _ = tokio::task::spawn_blocking(move || {
            for _ in 0..MAX_ENC_SNAPSHOT_BIND_DRAIN {
                match unmount_if_present(&jail_bind) {
                    Ok(true) => continue,
                    Ok(false) => break,
                    Err(e) => {
                        tracing::warn!(
                            sandbox_id = %sandbox_id,
                            %e,
                            "failed to drain encrypted-snapshot jail bind before seal"
                        );
                        break;
                    }
                }
            }
        })
        .await;
    }

    async fn write_encrypted_snapshot_files(
        &self,
        sandbox_id: &str,
        mount_dir: &pplx_space_common_core::types::NodeSpacePath,
        receipt: &mut capture_receipt::CaptureReceipt,
    ) -> Result<(), Status> {
        #[cfg(test)]
        if let Some(io) = &self.encrypted_capture_io {
            receipt.begin_capture();
            return (io.capture)().await;
        }
        let sandbox_id = sandbox_id.to_owned();
        let mount_dir = mount_dir.clone();
        let jail_root = format!("{JAIL_DIR}/firecracker/{sandbox_id}/root");
        let jail_bind = format!("{jail_root}/{ENC_SNAPSHOT_JAIL_SUBDIR}");
        let api_socket = format!("{jail_root}/firecracker.sock");
        let fc_snapshot_path = format!("/{ENC_SNAPSHOT_JAIL_SUBDIR}/{ENC_SNAPSHOT_VMSTATE_NAME}");
        let fc_mem_file_path = format!("/{ENC_SNAPSHOT_JAIL_SUBDIR}/{ENC_SNAPSHOT_MEMORY_NAME}");
        let bind = {
            tokio::task::spawn_blocking(move || -> Result<(), Status> {
                let (volume, target) = pin_encrypted_snapshot_bind(
                    mount_dir.as_path(),
                    Path::new(&jail_root),
                    super::trusted_paths::reown,
                )?;
                super::trusted_paths::bind_pinned_files(&volume, &target)
            })
            .await
            .map_err(|e| Status::internal(format!("encrypted-snapshot bind task failed: {e}")))?
        };
        let create_result = match bind {
            // On a re-suspend, `memory` is the correctly-sized image
            // loaded by Firecracker. Firecracker 1.15 merges dirty
            // pages into that existing file, so capture work scales
            // with the dirty set while the result remains standalone.
            // The first diff after a cold boot is standalone by the
            // Firecracker contract and starts from an empty volume.
            Ok(()) => {
                receipt.begin_capture();
                crate::fc_api::FcApiClient::new(api_socket)
                    .create_snapshot(&fc_snapshot_path, &fc_mem_file_path)
                    .await
                    .map_err(|e| Status::internal(format!("create encrypted snapshot failed: {e}")))
            }
            Err(error) => Err(error),
        };
        // ALWAYS unmount the bind, so the caller can seal the volume
        // whatever happened above; idempotent if the bind never formed.
        {
            let jail_bind = jail_bind.clone();
            let sandbox_id = sandbox_id.clone();
            let _ = tokio::task::spawn_blocking(move || {
                if let Err(e) = unmount_if_present(&jail_bind) {
                    tracing::warn!(
                        sandbox_id = %sandbox_id,
                        %e,
                        "failed to unmount encrypted-snapshot jail bind"
                    );
                }
            })
            .await;
        }
        create_result
    }

    pub(super) async fn restore_runtime_state(
        &self,
        request: Request<StateRestore>,
    ) -> Result<Response<RestoreVmStateResponse>, Status> {
        let sandbox_id = request.get_ref().sandbox_id.as_str().to_owned();
        let service = self.clone();
        self.run_owned_sandbox_work(&sandbox_id, async move {
            let _vm_guard = service
                .lock_sandbox_vm_cgroup_launch(request.get_ref().sandbox_id.as_str())
                .await;
            service.check_snapshot_execution(
                request.get_ref().sandbox_id.id(),
                request.get_ref().execution_fence,
            )?;
            service.load_snapshot_inner(request).await
        })
        .await
    }

    async fn load_snapshot_inner(
        &self,
        request: Request<StateRestore>,
    ) -> Result<Response<RestoreVmStateResponse>, Status> {
        let req = request.get_ref();
        // Validate before either path stages the jail or re-owns the snapshot
        // files, and so before the load, which cannot be undone short of
        // killing the VM.
        let rootfs_limiter = req
            .rootfs_rate_limit
            .as_ref()
            .map(drive_rate_limiter)
            .transpose()?
            .map(|limiter| self.disk_io_ceiling.clamp_limiter(limiter))
            .transpose()?;
        let jail_root = format!("{JAIL_DIR}/firecracker/{}/root", req.sandbox_id);
        let _lifecycle_permit = self
            .lifecycle_limiter
            .acquire(crate::lifecycle_limiter::LifecycleMethod::RestoreVmState)
            .await?;
        let lease = self
            .pool
            .get(req.sandbox_id.as_str())
            .map_err(slot_store_status)?
            .ok_or_else(|| Status::unavailable("snapshot lease is no longer active"))?;

        // Preserve path refusals before cancellable metadata IO or jail staging.
        let source = match &req.storage {
            SnapshotStorage::Encrypted(mount) => {
                trusted_paths::validate_snapshot_volume(
                    req.sandbox_id.as_str(),
                    crate::snapshot_request::path_str(mount)?,
                )?;
                RestoreSource::Encrypted(mount.clone())
            }
            SnapshotStorage::Plain { files, directory } => RestoreSource::Plain(
                crate::runtime_paths::RuntimePaths::firecracker().load_snapshot_paths(
                    req.sandbox_id.as_str(),
                    files,
                    directory.as_ref(),
                )?,
            ),
        };

        let execution = self
            .recorded_execution(req.sandbox_id.as_str().to_owned())
            .await?;
        if let Some(boot_source) = snapshot_boot_source(execution.as_ref())? {
            return self
                .restore_cloud_hypervisor_state(
                    request,
                    &lease,
                    source,
                    boot_source,
                    rootfs_limiter,
                )
                .await;
        }

        // BYOK encrypted snapshot: bind-mount the already-mounted, decrypted
        // encrypted-snapshot volume into the jail and load vmstate+memory
        // directly from it. The volume stays mounted for the sandbox's life (FC
        // mmaps the memory file); cleanup releases it on delete. No source_dir
        // staging — the bytes never exist in plaintext on disk.
        let paths = match source {
            RestoreSource::Encrypted(mount_dir) => {
                {
                    let jail_root = jail_root.clone();
                    let mount_dir = mount_dir.clone();
                    self.off_runtime("bind encrypted snapshot volume", move || {
                        let (volume, target) = pin_encrypted_snapshot_bind(
                            mount_dir.as_path(),
                            Path::new(&jail_root),
                            reown_encrypted_snapshot_inputs,
                        )?;
                        super::trusted_paths::bind_pinned_files(&volume, &target)
                    })
                    .await?;
                }
                let fc_snapshot_path =
                    format!("/{ENC_SNAPSHOT_JAIL_SUBDIR}/{ENC_SNAPSHOT_VMSTATE_NAME}");
                let fc_mem_file_path =
                    format!("/{ENC_SNAPSHOT_JAIL_SUBDIR}/{ENC_SNAPSHOT_MEMORY_NAME}");
                self.finish_load_snapshot(
                    &lease,
                    req.sandbox_id.as_str(),
                    rootfs_limiter.as_ref(),
                    &jail_root,
                    &fc_snapshot_path,
                    &fc_mem_file_path,
                )
                .await?;
                return Ok(Response::new(RestoreVmStateResponse {}));
            }
            RestoreSource::Plain(paths) => paths,
        };
        let fc_snapshot_path = format!("/{}", paths.vmstate_name);
        let fc_mem_file_path = format!("/{}", paths.memory_name);

        // Place snapshot files in the jail chroot from source_dir.
        // vmstate: small (~11 KB), reflink-copy into jail.
        // memory: large (~1 GB), bind-mount from source to avoid copying.
        //         FC has it mmap'd; writes trigger btrfs CoW on the source (which
        //         is a reflink clone of snapshot.base), so only dirty pages cost disk.
        if let Some(source_dir) = paths.source_dir.as_ref() {
            let source_dir = source_dir.clone();
            let jail_root = jail_root.clone();
            let memory_name = paths.memory_name.clone();
            self.off_runtime("stage snapshot load inputs", move || {
                let (vmstate, memory, jail) =
                    stage_plaintext_snapshot_inputs(&source_dir, Path::new(&jail_root))?;

                // Bind-mount memory (large) — same pattern as rootfs bind mount
                let memory_dst =
                    super::trusted_paths::open_bind_target(&jail, memory_name.as_ref())?;
                let uid = jail
                    .metadata()
                    .map_err(|e| Status::internal(e.to_string()))?
                    .uid();
                super::trusted_paths::reown(&vmstate, uid)?;
                super::trusted_paths::reown(&memory, uid)?;
                super::trusted_paths::bind_pinned_files(&memory, &memory_dst)?;
                Ok(())
            })
            .await?;
        }

        self.finish_load_snapshot(
            &lease,
            req.sandbox_id.as_str(),
            rootfs_limiter.as_ref(),
            &jail_root,
            &fc_snapshot_path,
            &fc_mem_file_path,
        )
        .await?;
        Ok(Response::new(RestoreVmStateResponse {}))
    }

    /// Shared tail of the load-snapshot paths (plaintext source_dir and BYOK
    /// encrypted volume): wait for the FC API socket, load the snapshot with the
    /// given memory file, apply the sandbox's validated rootfs rate limit, then
    /// fix v.sock permissions.
    async fn finish_load_snapshot(
        &self,
        lease: &SlotInfo,
        sandbox_id: &str,
        rootfs_limiter: Option<&pplx_space_common_core::types::RateLimiter>,
        jail_root: &str,
        fc_snapshot_path: &str,
        fc_mem_file_path: &str,
    ) -> Result<(), Status> {
        let api_socket = format!("{jail_root}/firecracker.sock");
        // Wait for the API socket to appear (FC needs time to create it after launch)
        let api_socket_wait_started = Instant::now();
        loop {
            if std::path::Path::new(&api_socket).exists() {
                break;
            }
            if api_socket_wait_started.elapsed() >= FIRECRACKER_API_SOCKET_TIMEOUT {
                return Err(Status::internal(format!(
                    "FC API socket not found after 5s: {api_socket}"
                )));
            }
            tokio::time::sleep(FIRECRACKER_API_SOCKET_POLL).await;
        }
        tracing::debug!(
            sandbox_id,
            poll_interval_ms = FIRECRACKER_API_SOCKET_POLL.as_millis() as u64,
            elapsed_ms = api_socket_wait_started.elapsed().as_millis() as u64,
            "firecracker api socket observed"
        );

        let client = crate::fc_api::FcApiClient::new(api_socket);
        client
            .load_snapshot(fc_snapshot_path, fc_mem_file_path)
            .await
            .map_err(|e| Status::internal(format!("load snapshot failed: {e}")))?;
        // The loaded drive limiter is the one of the VM the snapshot was taken
        // from. Replace it with the sandbox's own while the VM is still
        // paused; a sandbox must not run under another VM's limits.
        if let Some(limiter) = rootfs_limiter
            && let Err(error) = client.update_drive_rate_limiter("rootfs", limiter).await
        {
            let status = Status::internal(format!(
                "apply rootfs rate limit after snapshot load failed: {error}"
            ));
            self.rollback_loaded_snapshot(sandbox_id, "rootfs rate limit", &status)
                .await;
            return Err(status);
        }
        // The snapshot's limiter and the saved assignment may both exceed
        // this worker's operator ceiling. Enforce it before adopting/resuming.
        self.disk_io_ceiling
            .enforce_before_adoption(&client, |status| async move {
                self.rollback_loaded_snapshot(sandbox_id, "rootfs IOPS ceiling", &status)
                    .await;
            })
            .await?;
        self.adopt_validated_lease(sandbox_id, lease)?;

        // Fix v.sock permissions so sbx-manager can connect for health checks.
        // In snapshot-restore mode, FC creates v.sock during load_snapshot (not at launch).
        let vsock_path = format!("{jail_root}/v.sock");
        if std::path::Path::new(&vsock_path).exists() {
            if let Err(status) = chgrp_path(&vsock_path, gids::SPACE_SBXMGR) {
                self.rollback_loaded_snapshot(sandbox_id, "vsock chgrp", &status)
                    .await;
                return Err(status);
            }
            if let Err(status) = chmod_path(&vsock_path, 0o770) {
                self.rollback_loaded_snapshot(sandbox_id, "vsock chmod", &status)
                    .await;
                return Err(status);
            }
        }

        Ok(())
    }

    /// Kill the VM a snapshot was loaded into when a later step of the load
    /// fails, so a retried load starts from a fresh VM.
    async fn rollback_loaded_snapshot(
        &self,
        sandbox_id: &str,
        failed_step: &'static str,
        status: &Status,
    ) {
        let stored_pid = match self.pool.get(sandbox_id) {
            Ok(info) => info.and_then(|info| info.pid),
            Err(error) => {
                tracing::error!(sandbox_id, %error, "snapshot rollback cannot read lease; checking runtime PID evidence");
                None
            }
        };
        let fc_pid =
            stored_pid.or_else(|| read_pid_from_symlink(&format!("{RUN_DIR}/{sandbox_id}/fc.pid")));

        let Some(fc_pid) = fc_pid else {
            tracing::error!(
                sandbox_id,
                failed_step,
                error = %status,
                "unable to roll back loaded snapshot because runtime pid was not found"
            );
            return;
        };

        tracing::warn!(
            sandbox_id,
            pid = fc_pid,
            failed_step,
            error = %status,
            "rolling back loaded snapshot after a failed load step"
        );
        kill_launched_process_after_pid_persist_failure(self, sandbox_id, fc_pid).await;
    }

    pub(super) async fn create_btrfs_snapshot(
        &self,
        request: Request<CreateBtrfsSnapshotRequest>,
    ) -> Result<Response<CreateBtrfsSnapshotResponse>, Status> {
        self.authorize_peer(&request, "CreateBtrfsSnapshot", ALLOW_NODE_ADMIN_OR_SBXMGR)?;
        let sandbox_id = request.get_ref().sandbox_id.clone();
        self.run_artifact_mutation(&sandbox_id, |service| async move {
            service.create_btrfs_snapshot_inner(request).await
        })
        .await
    }

    async fn create_btrfs_snapshot_inner(
        &self,
        request: Request<CreateBtrfsSnapshotRequest>,
    ) -> Result<Response<CreateBtrfsSnapshotResponse>, Status> {
        let req = request.get_ref();
        validate_sandbox_id(&req.sandbox_id)?;
        validate_snapshot_name(&req.snapshot_name)?;
        let gate = self.lock_sandbox_cgroups(&req.sandbox_id).await;
        self.check_filesystem_execution(&req.sandbox_id, req.execution_fence.as_ref())
            .await?;
        let instance_dir = format!("{CACHE_DIR}/{}/instance", req.sandbox_id);
        let snapshots_dir = format!("{CACHE_DIR}/{}/snapshots", req.sandbox_id);
        let snapshot_path = format!("{snapshots_dir}/{}", req.snapshot_name);
        let writable = req.writable;
        let sandbox_id = req
            .sandbox_id
            .parse::<SandboxId>()
            .map_err(|error| Status::invalid_argument(error.to_string()))?;
        let device_gate = self
            .sandbox_device_stack_lock(&sandbox_id)
            .await
            .lock_owned()
            .await;
        let cow_path = PathBuf::from(&instance_dir).join("cow.luks");
        let encrypted = crate::device_backing::blocking(move || {
            match crate::device_backing::Backing::open(&cow_path, false) {
                Ok(cow) => Ok(cow.is_some()),
                // The ordinary capture path owns the missing-instance error.
                // A missing parent here is not a failed encryption probe.
                Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false),
                Err(error) => Err(error),
            }
        })
        .await
        .map_err(|error| Status::unavailable(error.to_string()))?;
        let capture = self.off_runtime("create btrfs snapshot", move || {
            let name = Path::new(&snapshot_path)
                .file_name()
                .ok_or_else(|| super::artifact_fs::error("missing snapshot name"))?;
            capture_instance_snapshot(
                Path::new(&instance_dir),
                Path::new(&snapshots_dir),
                name,
                writable,
                encrypted,
            )?;
            Ok(Response::new(CreateBtrfsSnapshotResponse {}))
        });
        let result = if encrypted {
            crate::rootfs_dmcrypt::capture_quiesced(sandbox_id, capture).await
        } else {
            capture.await
        };
        drop((gate, device_gate));
        result
    }

    pub(super) async fn delete_btrfs_snapshot(
        &self,
        request: Request<DeleteBtrfsSnapshotRequest>,
    ) -> Result<Response<DeleteBtrfsSnapshotResponse>, Status> {
        self.authorize_peer(&request, "DeleteBtrfsSnapshot", ALLOW_NODE_ADMIN_OR_SBXMGR)?;
        let req = request.get_ref();
        validate_sandbox_id(&req.sandbox_id)?;
        validate_snapshot_name(&req.snapshot_name)?;
        let snapshot_path = format!(
            "{CACHE_DIR}/{}/snapshots/{}",
            req.sandbox_id, req.snapshot_name
        );

        let guard = self.lock_sandbox_artifacts(&req.sandbox_id).await?;
        let claim = self.claim_sandbox_mutation(&req.sandbox_id)?;
        self.off_runtime("delete btrfs snapshot", move || {
            let _ownership = (guard, claim);
            if !std::path::Path::new(&snapshot_path).exists() {
                return Err(Status::not_found(format!(
                    "snapshot not found: {snapshot_path}"
                )));
            }

            match delete_cache_artifact_at(&snapshot_path) {
                Ok(true) => Ok(Response::new(DeleteBtrfsSnapshotResponse {})),
                Ok(false) => Err(Status::not_found(format!(
                    "snapshot not found: {snapshot_path}"
                ))),
                Err(error) => Err(Status::internal(format!(
                    "delete snapshot {snapshot_path}: {error}"
                ))),
            }
        })
        .await
    }

    /// Generic "delete this directory under /space/cache" RPC. The
    /// path may be a btrfs subvolume (callers without
    /// `CAP_SYS_ADMIN` can't `rmdir` a subvolume root) or a plain
    /// directory; the handler tries `btrfs subvolume delete` first
    /// and falls back to `remove_dir_all`. Idempotent: a missing
    /// path returns `was_present = false` rather than `NotFound`,
    /// because every caller in practice (snapshot reaper, named-
    /// snapshot delete, `.gc-pending-*` sweep) is doing a
    /// best-effort cleanup and doesn't care which actor won the
    /// race to delete.
    pub(super) async fn delete_cache_artifact(
        &self,
        request: Request<DeleteCacheArtifactRequest>,
    ) -> Result<Response<DeleteCacheArtifactResponse>, Status> {
        self.authorize_peer(&request, "DeleteCacheArtifact", ALLOW_NODE_ADMIN_OR_SBXMGR)?;
        let req = request.get_ref();
        validate_cache_artifact_path(&req.local_path)?;
        let local_path = req.local_path.clone();
        let sandbox_id = std::path::Path::new(&local_path)
            .strip_prefix(CACHE_DIR)
            .ok()
            .and_then(|path| path.components().next())
            .and_then(|part| part.as_os_str().to_str())
            .filter(|id| id.parse::<SandboxId>().is_ok());
        // Deletion capacity is bounded where the work runs, and taken before
        // any sandbox ownership so a queued deletion never holds a healthy
        // sandbox's artifact lock while it waits. A caller that stops waiting
        // releases nothing here: the slot is held by the blocking closure
        // until the filesystem work is done.
        let permit = self
            .cache_artifact_delete_slots
            .clone()
            .acquire_owned()
            .await
            .map_err(|_| Status::unavailable("cache artifact deletion lane closed"))?;
        let guard = match sandbox_id {
            Some(id) => Some(self.lock_sandbox_artifacts(id).await?),
            None => None,
        };
        // Capture/export can outlive its caller too. Exclude shared host work,
        // not just btrfs transfers, before unlinking any sandbox artifact.
        let claim = sandbox_id
            .map(|id| self.claim_sandbox_mutation(id))
            .transpose()?;

        self.off_runtime("delete cache artifact", move || {
            let _permit = permit;
            let _guard = guard;
            let _claim = claim;
            if !std::path::Path::new(&local_path).exists() {
                return Ok(Response::new(DeleteCacheArtifactResponse {
                    was_present: false,
                }));
            }

            delete_cache_artifact_at(&local_path)
                .map(|was_present| Response::new(DeleteCacheArtifactResponse { was_present }))
                .map_err(|error| {
                    Status::internal(format!("delete cache artifact {local_path}: {error}"))
                })
        })
        .await
    }

    pub(super) async fn list_btrfs_snapshots(
        &self,
        request: Request<ListBtrfsSnapshotsRequest>,
    ) -> Result<Response<ListBtrfsSnapshotsResponse>, Status> {
        self.authorize_peer(&request, "ListBtrfsSnapshots", ALLOW_NODE_ADMIN_OR_SBXMGR)?;
        let sandbox_id: SandboxId =
            request.get_ref().sandbox_id.parse().map_err(|error| {
                Status::invalid_argument(format!("invalid sandbox id: {error}"))
            })?;
        let sandbox_root = Path::new(CACHE_DIR).join(sandbox_id.to_string());

        // Diagnostic capacity may suppress measurements, but must never reject
        // identity enumeration.
        let measurement_permit = if request.get_ref().skip_size_measurement {
            None
        } else {
            self.blocking_probe_slots.clone().try_acquire_owned().ok()
        };
        let req = request.into_inner();
        let inventory = self.snapshot_inventory.clone();
        let response = self
            .off_runtime("list btrfs snapshots", move || {
                // The blocking job retains admission even if its caller cancels.
                let permit = measurement_permit;
                let measurement_budget = if permit.is_some() {
                    DISK_MEASUREMENT_BUDGET
                } else {
                    Duration::ZERO
                };
                let (paths, next_page_token, elapsed) =
                    inventory.page(sandbox_id, &sandbox_root, &req.page_token)?;
                Ok(ListBtrfsSnapshotsResponse {
                    snapshots: measure_snapshots(
                        paths,
                        measurement_budget.saturating_sub(elapsed),
                        btrfs_du_exclusive,
                    ),
                    next_page_token,
                })
            })
            .await?;

        Ok(Response::new(response))
    }

    pub(super) async fn rollback_btrfs_snapshot(
        &self,
        request: Request<RollbackBtrfsSnapshotRequest>,
    ) -> Result<Response<RollbackBtrfsSnapshotResponse>, Status> {
        self.authorize_peer(
            &request,
            "RollbackBtrfsSnapshot",
            ALLOW_NODE_ADMIN_OR_SBXMGR,
        )?;
        let sandbox_id = request.get_ref().sandbox_id.clone();
        self.run_artifact_mutation(&sandbox_id, |service| async move {
            service.rollback_btrfs_snapshot_inner(request).await
        })
        .await
    }

    async fn rollback_btrfs_snapshot_inner(
        &self,
        request: Request<RollbackBtrfsSnapshotRequest>,
    ) -> Result<Response<RollbackBtrfsSnapshotResponse>, Status> {
        let req = request.get_ref();
        let sandbox_id = req.sandbox_id.clone();
        validate_sandbox_id(&sandbox_id)?;
        validate_snapshot_name(&req.snapshot_name)?;
        let snapshot_name = req.snapshot_name.clone();
        let gate = self.lock_sandbox_cgroups(&sandbox_id).await;
        self.check_filesystem_execution(&sandbox_id, req.execution_fence.as_ref())
            .await?;
        self.off_runtime("rollback btrfs snapshot", move || {
            let _gate = gate;
            if super::runtime::firecracker_before_jail_rebuild(&sandbox_id)?.is_some() {
                return Err(Status::failed_precondition(
                    "stop the runtime before disk rollback",
                ));
            }
            let sandbox = super::trusted_paths::open_directory(Path::new(&format!(
                "{CACHE_DIR}/{sandbox_id}"
            )))?;
            restore_instance_snapshot_in(&sandbox, &snapshot_name, gids::SPACE_CACHE_RW)?;
            let snapshots =
                super::artifact_fs::directory(&sandbox, BTRFS_SNAPSHOTS_DIR_NAME.as_ref())
                    .map_err(super::artifact_fs::error)?;

            // Prune snapshots newer than the target
            let target_name = &snapshot_name;
            let mut pruned = Vec::new();
            {
                for entry in super::artifact_fs::entries(&snapshots)? {
                    let name = entry.to_string_lossy().to_string();
                    // Compare by timestamp portion (after the prefix and tier)
                    // Handles "d-1m-20260331T140000Z", "f-1m-...", or legacy "1m-..."
                    pub(super) fn extract_timestamp(s: &str) -> &str {
                        let rest = s
                            .strip_prefix("d-")
                            .or_else(|| s.strip_prefix("f-"))
                            .unwrap_or(s);
                        rest.split_once('-').map(|(_, ts)| ts).unwrap_or(s)
                    }
                    let target_ts = extract_timestamp(target_name);
                    let snap_ts = extract_timestamp(&name);
                    if snap_ts > target_ts
                        && super::artifact_fs::remove_at(&snapshots, &entry).is_ok()
                    {
                        pruned.push(name);
                    }
                }
            }

            Ok(Response::new(RollbackBtrfsSnapshotResponse {
                pruned_snapshots: pruned,
            }))
        })
        .await
    }

    pub(super) async fn restore_instance_snapshot(
        &self,
        request: Request<RestoreInstanceSnapshotRequest>,
    ) -> Result<Response<RestoreInstanceSnapshotResponse>, Status> {
        self.authorize_peer(
            &request,
            "RestoreInstanceSnapshot",
            ALLOW_NODE_ADMIN_OR_SBXMGR,
        )?;
        let sandbox_id = request.get_ref().sandbox_id.clone();
        self.run_artifact_mutation(&sandbox_id, |service| async move {
            service.restore_instance_snapshot_inner(request).await
        })
        .await
    }

    /// [`Self::rollback_btrfs_snapshot_inner`] without its pruning, for
    /// sandbox-manager's own rollback copies. A rollback prunes every snapshot
    /// whose name sorts after the target's, which would delete a user's
    /// snapshots when the target is an internal copy.
    async fn restore_instance_snapshot_inner(
        &self,
        request: Request<RestoreInstanceSnapshotRequest>,
    ) -> Result<Response<RestoreInstanceSnapshotResponse>, Status> {
        let req = request.get_ref();
        let sandbox_id = req.sandbox_id.clone();
        validate_sandbox_id(&sandbox_id)?;
        validate_snapshot_name(&req.snapshot_name)?;
        if !req
            .snapshot_name
            .starts_with(INTERNAL_INSTANCE_SNAPSHOT_PREFIX)
        {
            return Err(Status::invalid_argument(format!(
                "instance restore accepts only {INTERNAL_INSTANCE_SNAPSHOT_PREFIX} snapshots"
            )));
        }
        let snapshot_name = req.snapshot_name.clone();
        let gate = self.lock_sandbox_cgroups(&sandbox_id).await;
        self.check_filesystem_execution(&sandbox_id, req.execution_fence.as_ref())
            .await?;
        self.off_runtime("restore instance snapshot", move || {
            let _gate = gate;
            if super::runtime::firecracker_before_jail_rebuild(&sandbox_id)?.is_some() {
                return Err(Status::failed_precondition(
                    "stop the runtime before instance restore",
                ));
            }
            let sandbox = super::trusted_paths::open_directory(Path::new(&format!(
                "{CACHE_DIR}/{sandbox_id}"
            )))?;
            restore_instance_snapshot_in(&sandbox, &snapshot_name, gids::SPACE_CACHE_RW)?;
            Ok(Response::new(RestoreInstanceSnapshotResponse {}))
        })
        .await
    }

    pub(super) async fn get_disk_usage(
        &self,
        request: Request<GetDiskUsageRequest>,
    ) -> Result<Response<GetDiskUsageResponse>, Status> {
        self.authorize_peer(&request, "GetDiskUsage", ALLOW_NODE_ADMIN_OR_SBXMGR)?;
        let ids = &request.get_ref().sandbox_ids;
        let mut usage_bytes = std::collections::HashMap::new();
        let mut missing_ids = Vec::new();
        for id in ids {
            let Ok(sandbox_id) = SandboxId::from_str(id) else {
                continue;
            };
            match self.disk_usage_cache.lookup(sandbox_id) {
                SandboxDiskUsageLookup::Usage(bytes) => {
                    usage_bytes.insert(id.clone(), bytes);
                }
                SandboxDiskUsageLookup::Missing => missing_ids.push(MissingDiskUsageId {
                    id: id.clone(),
                    reason: MissingDiskUsageReason::InstancePathMissing as i32,
                }),
                SandboxDiskUsageLookup::Unknown => {}
            }
        }
        Ok(Response::new(GetDiskUsageResponse {
            usage_bytes,
            missing_ids,
        }))
    }
}

/// Replace `sandbox`'s instance with a writable copy of one of its snapshots.
/// Every snapshot, the source included, stays where it is.
fn restore_instance_snapshot_in(
    sandbox: &std::fs::File,
    snapshot_name: &str,
    gid: u32,
) -> Result<(), Status> {
    let snapshots = super::artifact_fs::directory(sandbox, BTRFS_SNAPSHOTS_DIR_NAME.as_ref())
        .map_err(super::artifact_fs::error)?;
    let source = super::artifact_fs::directory(&snapshots, snapshot_name.as_ref())
        .map_err(super::artifact_fs::error)?;
    super::artifact_fs::rollback(sandbox, &source, gid)
}

/// Try `btrfs subvolume delete` first, fall back to
/// `remove_dir_all`. Returns `true` when the path was deleted,
/// `false` if it did not exist at delete time. Errors only when
/// both the subvolume-delete and the directory-walk fallback
/// failed to dispose of an existing path.
///
/// The two-step shape is unavoidable on a btrfs cache: a subvolume
/// root cannot be removed by an unprivileged caller via the usual
/// directory rmdir, so the fallback is needed for plain directories
/// (e.g. suspend snapshots, imported / cross-node staged dirs)
/// while the privileged path is needed for subvolumes (e.g.
/// `d-*` / `f-*` named snapshots and `d-stop-*` stopped artifacts
/// created via `create_btrfs_snapshot`).
///
/// After the leaf delete succeeds, walks up at most two parent
/// levels and removes any that have become empty as a result. The
/// walk is subvolume-aware (an empty parent that is itself a
/// subvolume root is removed via `btrfs subvolume delete`, since
/// a plain `rmdir` would EPERM on a subvol root even when empty)
/// and bounded by `SPACE_CACHE_DIR` so a runaway walk can't reach
/// `/space/cache` or any of its top-level children. The two-level
/// budget covers both shapes the reaper produces:
///
///   - suspend artifact: `<sbx>/snapshot/` → walk to `<sbx>/`.
///   - named/stopped snapshot: `<sbx>/snapshots/<name>/` →
///     `<sbx>/snapshots/` → `<sbx>/`.
///
/// If a parent is non-empty (another artifact, a live `instance/`
/// subvolume, etc.) the walk stops — the artifact deletion alone
/// is what the caller asked for and the parent's other content
/// owns the dir entry.
fn delete_cache_artifact_at(local_path: &str) -> std::io::Result<bool> {
    if !std::path::Path::new(local_path).exists() {
        return Ok(false);
    }
    delete_dir_btrfs_aware(local_path)?;
    // Best-effort empty-parent cleanup. Failures inside the walk
    // do not fail the RPC: the load-bearing contract is "leaf is
    // gone"; parent cleanup is a hygiene step. Errors at any step
    // get logged so a leaked empty dir is visible.
    prune_empty_parents(local_path, 2);
    Ok(true)
}

/// Delete through the shared bounded descriptor walk, including nested
/// subvolumes. Cache path replacements never become a recursive path fallback.
fn delete_dir_btrfs_aware(path: &str) -> std::io::Result<()> {
    super::artifact_fs::delete(std::path::Path::new(path)).map_err(std::io::Error::other)
}

/// Prune ordinary empty parents only. A failed atomic rmdir never grants
/// authority to recursively destroy a subvolume which may have new children.
fn prune_empty_parents(start_path: &str, max_levels: usize) {
    let cache_root = std::path::Path::new(CACHE_DIR);
    let mut current = std::path::Path::new(start_path);
    for _ in 0..max_levels {
        let Some(parent) = current.parent() else {
            return;
        };
        if parent == cache_root || !parent.starts_with(cache_root) {
            return;
        }
        if let Err(error) = super::artifact_fs::remove_empty_parent(parent) {
            tracing::debug!(path = %parent.display(), %error, "empty parent retained");
            return;
        }
        current = parent;
    }
}

/// Existing wire paths are restricted to snapshot artifacts and fork-copy
/// staging, not the entire cache. Containment at use time is descriptor-based.
fn validate_cache_artifact_path(local_path: &str) -> Result<(), Status> {
    let cache_root = CACHE_DIR;
    let prefix = format!("{cache_root}/");
    let Some(suffix) = local_path.strip_prefix(&prefix) else {
        return Err(Status::invalid_argument(format!(
            "local_path must start with `{prefix}`"
        )));
    };
    if suffix.is_empty() {
        return Err(Status::invalid_argument(
            "local_path must name an entry inside the cache root, not the root itself",
        ));
    }
    for component in suffix.split('/') {
        if component.is_empty() || component == "." || component == ".." {
            return Err(Status::invalid_argument(
                "local_path must not contain empty, `.`, or `..` components",
            ));
        }
    }
    let mut parts = suffix.split('/');
    let owner = parts
        .next()
        .ok_or_else(|| Status::invalid_argument("missing artifact owner"))?;
    if owner != pplx_space_common_core::host::linux::paths::FORK_COPY_STAGING_DIR_NAME {
        super::trusted_paths::canonical_sandbox_id(owner)?;
    }
    let artifact = parts.next().ok_or_else(|| {
        Status::invalid_argument("artifact path must not name a sandbox or staging root")
    })?;
    let allowed = if owner == pplx_space_common_core::host::linux::paths::FORK_COPY_STAGING_DIR_NAME
    {
        uuid::Uuid::parse_str(artifact).is_ok()
    } else {
        // GC staging sits beside the artifact it held: `.gc-pending-<id>`, or
        // `snapshot.gc-pending-<id>` from an earlier sandbox-manager.
        artifact == "snapshot"
            || artifact.starts_with(".gc-pending-")
            || artifact.starts_with("snapshot.gc-pending-")
            || (matches!(artifact, "snapshots" | "template-snapshot-staging")
                && parts.next().is_some())
    };
    if !allowed {
        return Err(Status::permission_denied(
            "path is not a snapshot or fork-copy artifact",
        ));
    }
    Ok(())
}

fn drive_rate_limiter(
    limit: &pplx_space_node_mgr_client::proto::DriveRateLimit,
) -> Result<pplx_space_common_core::types::RateLimiter, Status> {
    let bucket = |bucket: Option<&pplx_space_node_mgr_client::proto::TokenBucket>, name: &str| {
        bucket
            .map(|bucket| pplx_space_common_core::types::TokenBucket {
                size: bucket.size,
                refill_time_ms: bucket.refill_time_ms,
            })
            .ok_or_else(|| {
                Status::invalid_argument(format!("rootfs_rate_limit.{name} is required"))
            })
    };
    Ok(pplx_space_common_core::types::RateLimiter {
        bandwidth: bucket(limit.bandwidth.as_ref(), "bandwidth")?,
        operations: bucket(limit.operations.as_ref(), "operations")?,
    })
}

/// How long a Firecracker pause or resume may take. Both normally finish in
/// milliseconds; a VMM that does not answer must not keep the sandbox's VM
/// gate, which termination also takes, held indefinitely.
const VM_PAUSE_STATE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);

/// Upper bound on a requested pre-snapshot free-page-hinting wait. The pause
/// holds the sandbox's VM gate meanwhile, which termination also takes.
const MAX_BALLOON_HINT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);

/// How often a hinting run's completion is polled. Firecracker reports runs of
/// about 200 ms per GiB of guest memory.
const BALLOON_HINT_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(20);

/// Bound for each balloon read before a snapshot pause; Firecracker answers
/// them from memory without involving the guest.
const BALLOON_PROBE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(1);

fn balloon_hint_timeout(timeout_ms: u32) -> std::time::Duration {
    std::time::Duration::from_millis(u64::from(timeout_ms)).min(MAX_BALLOON_HINT_TIMEOUT)
}

/// Before a snapshot pause, record the VM's balloon statistics and, when the
/// device has free-page hinting, run one hinting pass so that free guest
/// memory is released before the snapshot captures it. This fails open: a VM
/// without a balloon is untouched beyond one read, and every balloon error or
/// timeout is logged and counted while the pause and snapshot proceed.
/// Latency is bounded by `hint_timeout` plus the fixed probe and stop bounds.
async fn prepare_balloon_for_snapshot(
    client: &crate::fc_api::FcApiClient,
    sandbox_id: &str,
    hint_timeout: std::time::Duration,
) {
    use crate::metrics::node_manager::{
        BalloonHintResult, BalloonSnapshotProbe, balloon_free_page_hint,
        balloon_snapshot_probe_failed, balloon_snapshot_statistics,
    };

    let balloon = match tokio::time::timeout(BALLOON_PROBE_TIMEOUT, client.balloon()).await {
        Ok(Ok(Some(balloon))) => balloon,
        Ok(Ok(None)) => return,
        Ok(Err(error)) => {
            tracing::warn!(sandbox_id, %error, "balloon read before snapshot pause failed");
            balloon_snapshot_probe_failed(BalloonSnapshotProbe::Device);
            return;
        }
        Err(_) => {
            tracing::warn!(sandbox_id, "balloon read before snapshot pause timed out");
            balloon_snapshot_probe_failed(BalloonSnapshotProbe::Device);
            return;
        }
    };
    let mode = balloon.mode();
    if balloon.stats_polling_interval_s > 0 {
        match tokio::time::timeout(BALLOON_PROBE_TIMEOUT, client.balloon_statistics()).await {
            Ok(Ok(stats)) => {
                let mib = |bytes: Option<u64>| bytes.map(|bytes| bytes / (1024 * 1024));
                tracing::info!(
                    sandbox_id,
                    balloon_mode = %mode,
                    target_mib = stats.target_mib,
                    actual_mib = stats.actual_mib,
                    guest_free_mib = mib(stats.free_memory_bytes),
                    guest_available_mib = mib(stats.available_memory_bytes),
                    guest_total_mib = mib(stats.total_memory_bytes),
                    guest_disk_caches_mib = mib(stats.disk_caches_bytes),
                    "balloon statistics before snapshot pause"
                );
                balloon_snapshot_statistics(mode, stats.actual_mib, mib(stats.free_memory_bytes));
            }
            Ok(Err(error)) => {
                tracing::warn!(sandbox_id, %error, "balloon statistics read before snapshot pause failed");
                balloon_snapshot_probe_failed(BalloonSnapshotProbe::Statistics);
            }
            Err(_) => {
                tracing::warn!(
                    sandbox_id,
                    "balloon statistics read before snapshot pause timed out"
                );
                balloon_snapshot_probe_failed(BalloonSnapshotProbe::Statistics);
            }
        }
    }
    if !balloon.free_page_hinting {
        return;
    }
    let started = Instant::now();
    let (result, elapsed) = match client
        .hint_free_pages(hint_timeout, BALLOON_HINT_POLL_INTERVAL)
        .await
    {
        Ok(crate::fc_api::FreePageHintOutcome::Completed { elapsed }) => {
            tracing::info!(
                sandbox_id,
                elapsed_ms = elapsed.as_millis(),
                "free-page hinting before snapshot pause completed"
            );
            (BalloonHintResult::Completed, elapsed)
        }
        Ok(crate::fc_api::FreePageHintOutcome::TimedOut { elapsed }) => {
            tracing::warn!(
                sandbox_id,
                timeout_ms = hint_timeout.as_millis(),
                elapsed_ms = elapsed.as_millis(),
                "free-page hinting before snapshot pause timed out; snapshotting anyway"
            );
            (BalloonHintResult::TimedOut, elapsed)
        }
        Err(error) => {
            tracing::warn!(
                sandbox_id,
                %error,
                "free-page hinting before snapshot pause failed; snapshotting anyway"
            );
            (BalloonHintResult::Failed, started.elapsed())
        }
    };
    balloon_free_page_hint(result, elapsed.as_secs_f64() * 1000.0);
}

/// Pause or resume the VM behind `client`, failing once `timeout` passes.
pub(super) async fn set_firecracker_paused(
    client: &crate::fc_api::FcApiClient,
    paused: bool,
    timeout: std::time::Duration,
) -> Result<(), Status> {
    let change = async {
        if paused {
            client.pause().await
        } else {
            client.resume().await
        }
    };
    match tokio::time::timeout(timeout, change).await {
        Ok(changed) => {
            changed.map_err(|error| Status::internal(format!("set VM pause state: {error}")))
        }
        Err(_) => Err(Status::deadline_exceeded(format!(
            "Firecracker did not {} the VM within {timeout:?}",
            if paused { "pause" } else { "resume" }
        ))),
    }
}

/// Stops the runtime before the encrypted seal, re-issuing the idempotent stop
/// while it fails server-side because the killed Firecracker is still exiting.
/// New supervisors report drain timeouts as `Unavailable`. Keep `Internal`
/// handling while older supervisors still report busy cgroups that way.
async fn stop_runtime_for_encrypted_seal(
    service: &NodeManagerService,
    sandbox_id: &str,
) -> Result<(), Status> {
    stop_runtime_for_encrypted_seal_retrying_after(
        service,
        sandbox_id,
        ENC_CAPTURE_STOP_RETRY_DELAY,
    )
    .await
}

/// `stop_runtime_for_encrypted_seal`, waiting `retry_delay` between attempts.
async fn stop_runtime_for_encrypted_seal_retrying_after(
    service: &NodeManagerService,
    sandbox_id: &str,
    retry_delay: std::time::Duration,
) -> Result<(), Status> {
    let mut attempt = 1;
    loop {
        match stop_sandbox_runtime_with_supervisor(service, sandbox_id, true, None).await {
            Ok(_) => return Ok(()),
            Err(error)
                if matches!(
                    error.code(),
                    tonic::Code::Internal | tonic::Code::Unavailable
                ) && attempt < ENC_CAPTURE_STOP_ATTEMPTS =>
            {
                tracing::warn!(
                    sandbox_id,
                    attempt,
                    %error,
                    "post-capture runtime stop incomplete; retrying before seal"
                );
                attempt += 1;
                tokio::time::sleep(retry_delay).await;
            }
            Err(error) => return Err(error),
        }
    }
}

/// Capture may outlive its RPC waiter. The owned task retains its node guards
/// until it actually finishes; expiry never authorizes replay over live writes.
async fn wait_for_capture<T>(
    work: impl std::future::Future<Output = Result<T, Status>>,
) -> Result<T, Status> {
    tokio::time::timeout(
        pplx_space_node_mgr_client::MEMORY_CAPTURE_WAIT_TIMEOUT,
        work,
    )
    .await
    .map_err(|_| {
        tracing::warn!(
            "memory capture waiter expired; owned capture retains node guards until completion"
        );
        Status::deadline_exceeded(
            "memory capture exceeded its wait budget; unfinished capture remains fenced",
        )
    })?
}

fn snapshot_boot_source(
    execution: Option<&pplx_space_common_core::types::ExecutionMetadata>,
) -> Result<Option<pplx_space_node_runtime_core::boot::BootSource>, Status> {
    execution
        .filter(|execution| {
            execution.runtime().kind == pplx_space_common_core::types::RuntimeKind::CloudHypervisor
        })
        .map(pplx_space_node_runtime_core::boot::BootSource::from_execution)
        .transpose()
        .map_err(|error| Status::failed_precondition(error.to_string()))
}

#[cfg(test)]
mod tests;