From 51264a13d77d70f38cc13a65912b4cda14da4c8f Mon Sep 17 00:00:00 2001 From: Nikhil Unni Date: Wed, 7 Oct 2026 16:51:39 -0700 Subject: [PATCH] feat(host): a recency hot set for every post-copy source from the kernel's idle-page bits (ADR 0045 C2) A live move of a File-mode fresh session shipped an empty hot rider, so the destination drained 8 GiB in manifest order while agentd faulted its working set on demand. The hot-first drain existed; it had nothing to order by for a source without a UFFD handler. - Presetup marks the guest's present, exclusive, anonymous pages idle through `/sys/kernel/mm/page_idle/bitmap` (compound pages by head, resolved from `/proc/kpageflags` with a linear extent cache). The capture keeps today's dirty scan only; PFN collection and the idle read run after the export is registered, while the source stays paused, so the blackout is unchanged. - Accessed pages aggregate per chunk, ordered by count then offset, mapped to chain-manifest hashes, capped at 4096 for the idle tier; the handler-trace tier stays whole and goes first. - The hint is its own optional export artifact (`MigrationItem::HotHint`), fetched by the destination only after `state.bin` lands, with backoff, never touching the export lease, and stopped by the locally observed drain. The destination drain starts at once in manifest order and reorders the remaining chunks when the hint arrives. A hint is advisory: no hint path can fail, delay or block a move. - `HostCapabilities.page_idle` is an informational probe. Wire version 35. Tests: PFN word/bit math, shared-page exclusion, aggregation and ordering, hash mapping and caps, compound-extent op count, trace-first priority; a Linux root test marks and reads a private mapping; a fake-inner test drives capture and the destination poller with hints absent, empty, delayed, malformed and failed; the two-host KVM test asserts move success and hint arrival; the drain reorder is proven deterministically on a fake peer. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01UHxL6gj4o8EvxYpgtwEWaM --- .github/workflows/ci.yml | 6 + Cargo.lock | 1 + crates/engram-coordinator/src/placement.rs | 15 + .../tests/claim_cold_base_live_pg.rs | 1 + crates/engram-core/src/types/host.rs | 16 + crates/engram-core/src/types/snapshot.rs | 9 +- crates/engram-host-agent/src/capabilities.rs | 19 + crates/engram-host-agent/src/dirty_map.rs | 80 +- crates/engram-host-agent/src/grpc_server.rs | 1 + crates/engram-host-agent/src/lib.rs | 1 + crates/engram-host-agent/src/migrate_peer.rs | 7 +- crates/engram-host-agent/src/migration.rs | 4 +- crates/engram-host-agent/src/page_idle.rs | 455 ++++++++++ .../engram-host-agent/src/pooled_backend.rs | 835 +++++++++++++++++- .../tests/two_host_live_teleport.rs | 87 +- .../engram-protocol/proto/host_service.proto | 1 + crates/engram-protocol/src/grpc_client.rs | 5 + crates/engram-protocol/src/wire.rs | 3 +- crates/engram-protocol/tests/wire_golden.rs | 30 +- crates/engram-sandbox-firecracker/src/lib.rs | 50 +- crates/engram-uffd-handler/Cargo.toml | 1 + crates/engram-uffd-handler/src/chunked.rs | 56 +- crates/engram-uffd-handler/src/main.rs | 123 +-- crates/engram-uffd-handler/src/peer.rs | 2 + crates/engram-uffd-handler/src/runtime.rs | 271 +++++- crates/engram-uffd-handler/src/working_set.rs | 43 + docs/adr/0068-capability-vector-readiness.md | 5 + 27 files changed, 1925 insertions(+), 202 deletions(-) create mode 100644 crates/engram-host-agent/src/page_idle.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cfe2b454e..9c700dcef 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1208,6 +1208,12 @@ jobs: # pushes past its window (it timed out under parallelism in CI). The # contention can't be slot-partitioned away, so this batch stays # serial; the dedup below is the win that doesn't risk a flake. + - name: Idle-page tracking and advisory drain tests (root) + run: | + sudo -E env "PATH=$PATH" cargo nextest run \ + -p engram-host-agent -p engram-uffd-handler --lib \ + -E 'test(page_idle::tests::anonymous_mapping_accesses_clear_idle_bits) | test(fake_capture_and_destination_poller_deliver_final_hot_trace) | test(peer_mode_drains_sealed_chunks_and_sweep_never_overwrites_them)' + - name: NBD chunked-disk + netlink tests (Linux + KVM + nbd, root) if: steps.nbd_setup.outputs.available == 'true' env: diff --git a/Cargo.lock b/Cargo.lock index d6e9c34db..663e4e86c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3027,6 +3027,7 @@ dependencies = [ name = "engram-uffd-handler" version = "0.1.0" dependencies = [ + "async-trait", "bytes", "chrono", "engram-blob-client", diff --git a/crates/engram-coordinator/src/placement.rs b/crates/engram-coordinator/src/placement.rs index 472aeeda7..2d2a67f01 100644 --- a/crates/engram-coordinator/src/placement.rs +++ b/crates/engram-coordinator/src/placement.rs @@ -1417,6 +1417,7 @@ mod tests { base_shm_tmpfs: base_shm, uffd_minor_shmem: uffd_minor, nbd, + page_idle: CapStatus::Unknown, bundle_stamp: bundle, fc_snapshot_version: fc_snapshot_version.map(str::to_string), wire_version: 7, @@ -1450,6 +1451,19 @@ mod tests { assert!(host_meets_capabilities(&h, &req).is_ok()); } + #[test] + fn page_idle_failure_never_gates_placement() { + let mut h = host(1); + h.capabilities = fully_ok(); + h.capabilities.page_idle = CapStatus::Failed("unavailable".into()); + let req = CapabilityRequirements { + needs_uffd_substrate: true, + needs_live_substrate: true, + fc_snapshot_version: None, + }; + assert!(host_meets_capabilities(&h, &req).is_ok()); + } + /// ADR 0123: a live post-copy DESTINATION needs the substrate to /// be present, not merely "not applicable" — the migration /// restore forces UFFD. A File-backend host that honestly reports @@ -2349,6 +2363,7 @@ mod tests { base_shm_tmpfs: engram_core::types::host::CapStatus::Ok(None), uffd_minor_shmem: engram_core::types::host::CapStatus::Ok(None), nbd: engram_core::types::host::CapStatus::Ok(None), + page_idle: engram_core::types::host::CapStatus::Unknown, bundle_stamp: engram_core::types::host::CapStatus::Ok(None), fc_snapshot_version: Some("v9.0.0".into()), wire_version: 7, diff --git a/crates/engram-coordinator/tests/claim_cold_base_live_pg.rs b/crates/engram-coordinator/tests/claim_cold_base_live_pg.rs index 3db0c2e77..4c74072b8 100644 --- a/crates/engram-coordinator/tests/claim_cold_base_live_pg.rs +++ b/crates/engram-coordinator/tests/claim_cold_base_live_pg.rs @@ -129,6 +129,7 @@ async fn seed_fc_host(meta: &Arc, fc_version: &str) -> HostId base_shm_tmpfs: CapStatus::Ok(None), uffd_minor_shmem: CapStatus::Ok(None), nbd: CapStatus::Ok(None), + page_idle: CapStatus::Unknown, bundle_stamp: CapStatus::Ok(None), fc_snapshot_version: Some(fc_version.to_string()), wire_version: engram_protocol::WIRE_VERSION, diff --git a/crates/engram-core/src/types/host.rs b/crates/engram-core/src/types/host.rs index 744dcec89..31064114a 100644 --- a/crates/engram-core/src/types/host.rs +++ b/crates/engram-core/src/types/host.rs @@ -271,6 +271,9 @@ pub struct HostCapabilities { /// dev fallback — this only makes a prod misconfiguration visible. #[serde(default)] pub nbd: CapStatus, + /// Idle-page bitmap access. Informational; never a placement requirement. + #[serde(default)] + pub page_idle: CapStatus, /// `bundles::read_stamp(bundle_dir)` returned a non-empty stamp — /// the host has *some* current RO-bundle generation staged. #[serde(default)] @@ -317,6 +320,9 @@ impl HostCapabilities { if is_failing(&self.nbd) { out.push("nbd"); } + if matches!(self.page_idle, CapStatus::Failed(_)) { + out.push("page_idle"); + } if is_failing(&self.bundle_stamp) { out.push("bundle_stamp"); } @@ -772,6 +778,7 @@ mod tests { base_shm_tmpfs: CapStatus::NotApplicable, uffd_minor_shmem: CapStatus::NotApplicable, nbd: CapStatus::NotApplicable, + page_idle: CapStatus::NotApplicable, bundle_stamp: CapStatus::Failed("no stamp".to_string()), fc_snapshot_version: None, wire_version: 7, @@ -779,6 +786,14 @@ mod tests { assert_eq!(vz_like.failing_capabilities(), vec!["bundle_stamp"]); } + #[test] + fn page_idle_failure_is_visible_but_old_payloads_still_decode() { + let mut caps: HostCapabilities = serde_json::from_str("{}").unwrap(); + assert_eq!(caps.page_idle, CapStatus::Unknown); + caps.page_idle = CapStatus::Failed("not configured".into()); + assert_eq!(caps.failing_capabilities(), vec!["page_idle"]); + } + #[test] fn failing_capabilities_sorted_deterministic_order() { let broken = HostCapabilities { @@ -788,6 +803,7 @@ mod tests { base_shm_tmpfs: CapStatus::Failed("x".into()), uffd_minor_shmem: CapStatus::Ok(None), nbd: CapStatus::Unknown, + page_idle: CapStatus::Unknown, bundle_stamp: CapStatus::Ok(None), fc_snapshot_version: None, wire_version: 7, diff --git a/crates/engram-core/src/types/snapshot.rs b/crates/engram-core/src/types/snapshot.rs index a4a4d8bc0..b7466328d 100644 --- a/crates/engram-core/src/types/snapshot.rs +++ b/crates/engram-core/src/types/snapshot.rs @@ -215,10 +215,9 @@ pub struct MigrationSourceInfo { pub source_addr: String, pub memory_manifest_json: Vec, pub memory_manifest_ref: super::manifest::ManifestRef, - /// ADR 0045 C2 (E2B fold): the source guest's hot set in fault - /// order — the destination pulls these FIRST. Best-effort rider - /// (empty when the source had no trace); serde-default keeps - /// mixed rolls safe. + /// The presetup handler trace in fault order. The source retains this + /// first-tier order and adds idle-derived hashes to the capture-time + /// root seal update. Empty when the source had no handler trace. #[serde(default)] pub hot_chunks: Vec<[u8; 32]>, /// ADR 0045 C2: this is a POST-COPY move. The destination restores @@ -365,6 +364,8 @@ pub enum MigrationItem { /// by chunk index — demand-fetched (and drained) by the dest's /// NBD backend straight out of the frozen source's RAM. DiskChunkAt(DiskRole, u64), + /// Optional advisory hot-set hint, fetched after the required VM state. + HotHint, } /// A frame of `migration_fetch`'s stream. diff --git a/crates/engram-host-agent/src/capabilities.rs b/crates/engram-host-agent/src/capabilities.rs index 24e94301d..29baf25bd 100644 --- a/crates/engram-host-agent/src/capabilities.rs +++ b/crates/engram-host-agent/src/capabilities.rs @@ -324,6 +324,24 @@ pub async fn fc_snapshot_version(firecracker_bin: Option<&Path>) -> Option CapStatus { + if !is_fc { + return CapStatus::NotApplicable; + } + #[cfg(target_os = "linux")] + { + match crate::page_idle::open_bitmap() { + Ok(_) => CapStatus::Ok(None), + Err(error) => CapStatus::Failed(error.to_string()), + } + } + #[cfg(not(target_os = "linux"))] + { + CapStatus::NotApplicable + } +} + /// Run every probe and assemble the vector this tick. `schema: 1` /// unconditionally — a `HostCapabilities::default()` (`schema: 0`) /// only ever comes from an old host-agent's absent field, never from @@ -355,6 +373,7 @@ pub async fn probe_all(inputs: &ProbeInputs) -> HostCapabilities { base_shm_tmpfs, uffd_minor_shmem, nbd, + page_idle: probe_page_idle(is_fc), bundle_stamp, fc_snapshot_version, wire_version: engram_protocol::WIRE_VERSION, diff --git a/crates/engram-host-agent/src/dirty_map.rs b/crates/engram-host-agent/src/dirty_map.rs index afb294692..98873164c 100644 --- a/crates/engram-host-agent/src/dirty_map.rs +++ b/crates/engram-host-agent/src/dirty_map.rs @@ -31,8 +31,7 @@ //! server hashes the live bytes on request and demotes matches to the //! class-2 fetch. //! -//! Everything here is blackout-critical (the scan sits between the pause -//! and the SEAL push) — reads are batched (4096 entries = 32 KiB per +//! The dirty scan is blackout-critical (between pause and the SEAL push) — reads are batched (4096 entries = 32 KiB per //! `pread`) and the caller records `scan_ms` (R6: measure, never quote). use std::io::{self, BufRead, BufReader}; @@ -149,45 +148,60 @@ pub fn scan_dirty_chunks( chunk_size: u64, total_bytes: u64, ) -> io::Result { - use std::os::unix::fs::FileExt; + let mut seal = SealBitmap::new(chunk_size, total_bytes.div_ceil(chunk_size)); + walk_pagemap_batches(pid, vmas, |base, bytes| { + let mut i = 0; + while i < bytes.len() / 8 { + let entry = u64::from_ne_bytes(bytes[i * 8..i * 8 + 8].try_into().unwrap()); + let offset = base + i as u64 * PAGE_SIZE; + if offset < total_bytes && pagemap_entry_is_dirty(entry) { + let chunk = offset / chunk_size; + seal.set(chunk); + i += (((chunk + 1) * chunk_size - offset).div_ceil(PAGE_SIZE) as usize).max(1); + } else { + i += 1; + } + } + Ok(()) + })?; + Ok(seal) +} + +/// Visit each page once, in snapshot-offset order, with batched pagemap reads. +#[cfg(target_os = "linux")] +pub(crate) fn walk_pagemap( + pid: u32, + vmas: &[GuestVma], + mut visit: impl FnMut(u64, u64) -> io::Result<()>, +) -> io::Result<()> { + walk_pagemap_batches(pid, vmas, |base, bytes| { + let (entries, _) = bytes.as_chunks::<8>(); + for (i, entry) in entries.iter().enumerate() { + visit(base + i as u64 * PAGE_SIZE, u64::from_ne_bytes(*entry))?; + } + Ok(()) + }) +} - let chunk_count = total_bytes.div_ceil(chunk_size); - let mut seal = SealBitmap::new(chunk_size, chunk_count); +#[cfg(target_os = "linux")] +fn walk_pagemap_batches( + pid: u32, + vmas: &[GuestVma], + mut visit: impl FnMut(u64, &[u8]) -> io::Result<()>, +) -> io::Result<()> { + use std::os::unix::fs::FileExt; let pagemap = std::fs::File::open(format!("/proc/{pid}/pagemap"))?; let mut buf = vec![0u8; PAGEMAP_BATCH_ENTRIES * 8]; - for vma in vmas { let mut va = vma.start; while va < vma.end { - let pages_left = ((vma.end - va) / PAGE_SIZE) as usize; - let n = pages_left.min(PAGEMAP_BATCH_ENTRIES); - let byte_len = n * 8; - pagemap.read_exact_at(&mut buf[..byte_len], (va / PAGE_SIZE) * 8)?; - - let mut i = 0usize; - while i < n { - let entry = u64::from_le_bytes(buf[i * 8..i * 8 + 8].try_into().unwrap()); - if pagemap_entry_is_dirty(entry) { - let offset = vma.file_offset + (va - vma.start) + (i as u64) * PAGE_SIZE; - if offset < total_bytes { - let chunk = offset / chunk_size; - seal.set(chunk); - // Any-page rule satisfied: skip the rest of this - // chunk (within the batch; the next batch - // re-checks cheaply via the bitmap). - let chunk_end_off = (chunk + 1) * chunk_size; - let pages_to_skip = - (chunk_end_off.saturating_sub(offset)).div_ceil(PAGE_SIZE) as usize; - i += pages_to_skip.max(1); - continue; - } - } - i += 1; - } - va += (n as u64) * PAGE_SIZE; + let n = (((vma.end - va) / PAGE_SIZE) as usize).min(PAGEMAP_BATCH_ENTRIES); + pagemap.read_exact_at(&mut buf[..n * 8], (va / PAGE_SIZE) * 8)?; + visit(vma.file_offset + va - vma.start, &buf[..n * 8])?; + va += n as u64 * PAGE_SIZE; } } - Ok(seal) + Ok(()) } #[cfg(test)] diff --git a/crates/engram-host-agent/src/grpc_server.rs b/crates/engram-host-agent/src/grpc_server.rs index 3ecd1d7be..2cfc889dd 100644 --- a/crates/engram-host-agent/src/grpc_server.rs +++ b/crates/engram-host-agent/src/grpc_server.rs @@ -434,6 +434,7 @@ impl HostService for HostServiceImpl { Err(_) => return Err(Status::invalid_argument("unknown disk role")), }; match Kind::try_from(item.kind) { + Ok(Kind::HotHint) => Ok(engram_core::types::snapshot::MigrationItem::HotHint), Ok(Kind::StateBin) => Ok(engram_core::types::snapshot::MigrationItem::StateBin), Ok(Kind::DiskSealInfo) => Ok( engram_core::types::snapshot::MigrationItem::DiskSealInfo(role), diff --git a/crates/engram-host-agent/src/lib.rs b/crates/engram-host-agent/src/lib.rs index bd7e06550..2cd04371c 100644 --- a/crates/engram-host-agent/src/lib.rs +++ b/crates/engram-host-agent/src/lib.rs @@ -43,6 +43,7 @@ pub mod harness; pub mod host_client; pub mod migrate_peer; pub mod migration; +pub mod page_idle; pub mod peer_fill; pub mod session_epochs; pub mod substrate_server; diff --git a/crates/engram-host-agent/src/migrate_peer.rs b/crates/engram-host-agent/src/migrate_peer.rs index dec690e18..0e88ecb65 100644 --- a/crates/engram-host-agent/src/migrate_peer.rs +++ b/crates/engram-host-agent/src/migrate_peer.rs @@ -701,10 +701,11 @@ impl PeerServer { /// as the proto's sync `read_frame` (4-byte big-endian length + bincode /// body, bounded by `MAX_FRAME_BYTES`). The page channel's codec is sync /// by design (the dest client must stay tokio-free), but the SOURCE -/// reads only the opening `Hello` async so an unknown-export park can +/// reads the opening `Hello` async so an unknown-export park can /// run on a tokio task instead of pinning a blocking-pool thread (issue -/// #226 (b)); everything after the export resolves uses the sync codec. -async fn read_frame_async(r: &mut R) -> std::io::Result +/// #226 (b)); the destination also uses this for cancellable control reports. +/// Page serving after the export resolves uses the sync codec. +pub(crate) async fn read_frame_async(r: &mut R) -> std::io::Result where R: tokio::io::AsyncRead + Unpin, T: serde::de::DeserializeOwned, diff --git a/crates/engram-host-agent/src/migration.rs b/crates/engram-host-agent/src/migration.rs index 3c1667dd4..05af43922 100644 --- a/crates/engram-host-agent/src/migration.rs +++ b/crates/engram-host-agent/src/migration.rs @@ -80,7 +80,7 @@ pub struct MigrationExport { /// session (`ttl_verdict` returns StayPaused, converging via the /// scanner within one cycle). pub state_served: Arc, - /// Last page/artifact-serving activity (the export TTL clock), as a + /// Last required page/artifact-serving activity (the export TTL clock), as a /// `now_mono` reading. pub last_activity: Arc>, /// The sandbox's capture lock, held for the export's lifetime — @@ -90,7 +90,7 @@ pub struct MigrationExport { } impl MigrationExport { - /// Refresh the activity clock (every artifact/page serve) off the + /// Refresh the activity clock (required artifact/page serves only) off the /// injected monotonic clock. pub fn touch(&self) { *self.last_activity.lock().expect("last_activity poisoned") = self.clock.now_mono(); diff --git a/crates/engram-host-agent/src/page_idle.rs b/crates/engram-host-agent/src/page_idle.rs new file mode 100644 index 000000000..07f439177 --- /dev/null +++ b/crates/engram-host-agent/src/page_idle.rs @@ -0,0 +1,455 @@ +//! Best-effort recency hints from Linux idle-page tracking. These hints never +//! decide which bytes must be transferred. Invariant: a hot-set hint must never +//! fail, delay, or block a move, resume, drain, or blackout. See the kernel's +//! Documentation/admin-guide/mm/idle_page_tracking.rst. + +use std::collections::{BTreeMap, HashSet}; +use std::io; + +pub const BITMAP_PATH: &str = "/sys/kernel/mm/page_idle/bitmap"; +pub const MAX_HOT_CHUNKS: usize = 4096; + +/// A native-endian bitmap word at an 8-byte aligned file offset. +pub fn pfn_word_bit(pfn: u64) -> (u64, u64) { + ((pfn / 64) * 8, 1 << (pfn % 64)) +} + +/// Return only visible PFNs of private anonymous pages. Shared/file pages +/// cannot provide per-guest hotness and must not be marked or counted. +pub fn present_pfn(entry: u64) -> Option { + let pfn = entry & ((1 << 55) - 1); + (entry & (1 << 63) != 0 && entry & (1 << 61) == 0 && entry & (1 << 56) != 0 && pfn != 0) + .then_some(pfn) +} + +/// Count accessed pages per chunk. The injected reader returns a native-endian +/// bitmap word. Equal counts keep offset order. +#[derive(Default)] +pub struct HotPages { + counts: BTreeMap, +} + +impl HotPages { + pub fn observe( + &mut self, + offset: u64, + entry: u64, + chunk_size: u64, + mut read: impl FnMut(u64) -> io::Result, + ) -> io::Result<()> { + if let Some(pfn) = present_pfn(entry) { + let (word, bit) = pfn_word_bit(pfn); + if read(word)? & bit == 0 { + *self + .counts + .entry(offset / chunk_size * chunk_size) + .or_default() += 1; + } + } + Ok(()) + } + + pub fn ordered(&self) -> Vec { + let mut counts: Vec<_> = self.counts.iter().map(|(&o, &n)| (o, n)).collect(); + counts.sort_unstable_by_key(|&(offset, count)| (std::cmp::Reverse(count), offset)); + counts.into_iter().map(|(offset, _)| offset).collect() + } +} + +/// Resolve offsets without repeated manifest scans. Deduplicate equal hashes +/// before the idle-only cap. Preserve the complete bounded handler-trace tier. +pub fn hot_hashes( + manifest: &engram_chunk_store::Manifest, + offsets: &[u64], + fallback: &[[u8; 32]], +) -> Vec<[u8; 32]> { + let by_offset: BTreeMap<_, _> = manifest + .chunks + .iter() + .map(|c| (c.offset, *c.hash.as_bytes())) + .collect(); + let mut seen = HashSet::new(); + let mut hot: Vec<_> = fallback + .iter() + .copied() + .filter(|h| seen.insert(*h)) + .collect(); + hot.extend( + offsets + .iter() + .filter_map(|o| by_offset.get(o)) + .copied() + .filter(|h| seen.insert(*h)) + .take(MAX_HOT_CHUNKS), + ); + hot +} + +#[cfg(target_os = "linux")] +pub fn open_bitmap() -> io::Result { + std::fs::OpenOptions::new() + .read(true) + .write(true) + .open(BITMAP_PATH) +} + +/// Log optional-feature failures only once per process, including read failures. +#[cfg(target_os = "linux")] +pub fn note_unavailable(error: &io::Error) { + static LOGGED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + if !LOGGED.swap(true, std::sync::atomic::Ordering::Relaxed) { + tracing::info!(%error, "page_idle unavailable; using handler working-set trace"); + } +} + +/// Set only bits for present guest pages. Writes are OR operations in the +/// kernel: no read-modify-write, and no changes to unrelated idle bits. +#[cfg(target_os = "linux")] +pub fn mark(pid: u32, vmas: &[crate::dirty_map::GuestVma]) -> io::Result<()> { + use std::os::unix::fs::FileExt; + let bitmap = open_bitmap()?; + let mut words = BTreeMap::::new(); + let mut heads = CompoundHeads::new()?; + crate::dirty_map::walk_pagemap(pid, vmas, |_, entry| { + if let Some(pfn) = present_pfn(entry) { + let (word, bit) = pfn_word_bit(heads.head(pfn)?); + *words.entry(word).or_default() |= bit; + } + Ok(()) + })?; + for (word, bits) in words { + bitmap.write_all_at(&bits.to_ne_bytes(), word)?; + } + Ok(()) +} + +/// Read the captured PFNs after export registration. The source stays paused, +/// but host reclaim can still change PFNs. These are ordering hints only. +#[cfg(target_os = "linux")] +pub fn scan(pages: &[(u64, u64)], chunk_size: u64) -> io::Result> { + use std::os::unix::fs::FileExt; + let bitmap = open_bitmap()?; + let mut words = std::collections::HashMap::new(); + let mut hot = HotPages::default(); + let mut heads = CompoundHeads::new()?; + for &(offset, pfn) in pages { + let head = heads.head(pfn)?; + hot.observe(offset, (1 << 63) | (1 << 56) | head, chunk_size, |word| { + if let Some(bits) = words.get(&word) { + return Ok(*bits); + } + let mut bytes = [0; 8]; + bitmap.read_exact_at(&mut bytes, word)?; + let bits = u64::from_ne_bytes(bytes); + words.insert(word, bits); + Ok(bits) + })?; + } + Ok(hot.ordered()) +} + +/// Idle bits exist only on compound heads. Cache kpageflags blocks so walking +/// contiguous huge pages does not issue one syscall per base page. +#[cfg(target_os = "linux")] +struct CompoundHeads { + file: std::fs::File, + blocks: std::collections::HashMap>, + heads: BTreeMap, +} + +#[cfg(target_os = "linux")] +impl CompoundHeads { + fn new() -> io::Result { + Ok(Self { + file: std::fs::File::open("/proc/kpageflags")?, + blocks: Default::default(), + heads: Default::default(), + }) + } + + fn flags(&mut self, pfn: u64) -> io::Result { + use std::os::unix::fs::FileExt; + let block = pfn / 512; + let bytes = match self.blocks.entry(block) { + std::collections::hash_map::Entry::Occupied(e) => e.into_mut(), + std::collections::hash_map::Entry::Vacant(e) => { + let mut bytes = vec![0; 4096]; + self.file.read_exact_at(&mut bytes, block * 4096)?; + e.insert(bytes) + } + }; + let pos = (pfn % 512) as usize * 8; + Ok(u64::from_ne_bytes(bytes[pos..pos + 8].try_into().unwrap())) + } + + fn head(&mut self, pfn: u64) -> io::Result { + // Split borrows: the extent cache needs only a kpageflags reader. + let mut extents = std::mem::take(&mut self.heads); + let result = compound_head(&mut extents, pfn, |p| self.flags(p)); + self.heads = extents; + result + } +} + +/// Cache complete compound extents as (head, order). Each tail is examined once, +/// then ascending tails resolve with a range lookup, without per-tail inserts. +#[cfg(any(target_os = "linux", test))] +fn compound_head( + extents: &mut BTreeMap, + pfn: u64, + mut flags: impl FnMut(u64) -> io::Result, +) -> io::Result { + const HEAD: u64 = 1 << 15; + const TAIL: u64 = 1 << 16; + if let Some((&head, &order)) = extents.range(..=pfn).next_back() { + if pfn - head < 1u64 << order { + return Ok(head); + } + } + let f = flags(pfn)?; + if f & (HEAD | TAIL) == 0 { + return Ok(pfn); + } + let mut head = pfn; + while flags(head)? & TAIL != 0 { + head = head + .checked_sub(1) + .ok_or_else(|| io::Error::other("compound tail without head"))?; + } + if flags(head)? & HEAD == 0 { + return Err(io::Error::other("invalid compound head")); + } + let mut end = head + 1; + while flags(end)? & TAIL != 0 { + end += 1; + } + let pages = end - head; + if !pages.is_power_of_two() { + return Err(io::Error::other("invalid compound extent")); + } + extents.insert(head, pages.trailing_zeros()); + Ok(head) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compound_extent_lookup_is_linear_and_validates_head() { + let mut extents = BTreeMap::new(); + let mut reads = 0; + for pfn in 513..1024 { + assert_eq!( + compound_head(&mut extents, pfn, |p| { + reads += 1; + Ok(if p == 512 { + 1 << 15 + } else if p < 1024 { + 1 << 16 + } else { + 0 + }) + }) + .unwrap(), + 512 + ); + } + assert!(reads < 520, "{reads} kpageflags operations"); + assert_eq!(extents.len(), 1); + assert!(compound_head(&mut BTreeMap::new(), 2, |p| Ok(if p == 2 { + 1 << 16 + } else { + 0 + })) + .is_err()); + } + + #[test] + fn shared_pagemap_entries_are_neither_marked_nor_counted() { + let mut hot = HotPages::default(); + for entry in [(1 << 63) | 42, (1 << 63) | (1 << 56) | (1 << 61) | 42] { + assert_eq!(present_pfn(entry), None); // mark and PFN capture use this filter + hot.observe(0, entry, 4096, |_| panic!("shared page read")) + .unwrap(); + } + assert!(hot.ordered().is_empty()); + } + + #[test] + fn pfn_alignment_presence_and_native_bits() { + assert_eq!(pfn_word_bit(63), (0, 1 << 63)); + assert_eq!(pfn_word_bit(64), (8, 1)); + assert_eq!(pfn_word_bit(129), (16, 2)); + assert_eq!(present_pfn(1 << 63), None); + assert_eq!(present_pfn((1 << 62) | 42), None); + assert_eq!(present_pfn((1 << 63) | (1 << 56) | (1 << 61) | 42), None); + } + + #[test] + fn fake_reader_counts_accessed_pages_and_breaks_ties_by_offset() { + let mut pages = HotPages::default(); + for (offset, pfn) in [(0, 64), (4096, 65), (8192, 66), (12288, 67), (16384, 68)] { + pages + .observe(offset, (1 << 63) | (1 << 56) | pfn, 8192, |word| { + assert_eq!(word, 8); + Ok(1 << 2) // page 66 alone stayed idle + }) + .unwrap(); + } + pages + .observe(24576, 0, 8192, |_| panic!("absent page read")) + .unwrap(); + assert_eq!(pages.ordered(), vec![0, 8192, 16384]); + } + + #[test] + fn hashes_keep_trace_first_skip_zero_chunks_deduplicate_and_cap() { + use engram_chunk_store::manifest::{ChunkHash, ChunkRef, ChunkSize, ManifestKind}; + let chunks: Vec<_> = (0..5000_u64) + .map(|i| { + let mut bytes = [0; 32]; + bytes[..8].copy_from_slice(&i.to_le_bytes()); + ChunkRef { + offset: i * 4096, + hash: ChunkHash::from_bytes(bytes), + } + }) + .collect(); + let manifest = engram_chunk_store::Manifest { + schema_version: 1, + kind: ManifestKind::Memory, + chunk_size: ChunkSize::bytes(4096), + total_bytes: 5001 * 4096, + chunks, + parent: None, + working_set_trace: None, + annotations: serde_json::Value::Null, + }; + let first = *manifest.chunks[5].hash.as_bytes(); + let offsets: Vec<_> = (0..=5000_u64).rev().map(|i| i * 4096).collect(); + let hot = hot_hashes(&manifest, &offsets, &[first, first]); + assert_eq!(hot.len(), MAX_HOT_CHUNKS + 1); + let whole: Vec<_> = manifest.chunks.iter().map(|c| *c.hash.as_bytes()).collect(); + assert_eq!(hot_hashes(&manifest, &[], &whole), whole); + assert_eq!(hot[0], first); + assert_eq!(hot[1], *manifest.chunks[4999].hash.as_bytes()); + assert_eq!(hot.iter().filter(|h| **h == first).count(), 1); + assert_eq!(hot_hashes(&manifest, &[], &[first]), vec![first]); + assert!(hot_hashes(&manifest, &[5000 * 4096], &[]).is_empty()); + } + + #[test] + fn idle_read_failure_is_reported() { + let mut hot = HotPages::default(); + assert!(hot + .observe(0, (1 << 63) | (1 << 56) | 1, 4096, |_| Err( + io::ErrorKind::PermissionDenied.into() + )) + .is_err()); + assert!(hot.ordered().is_empty()); + } + + #[cfg(target_os = "linux")] + #[test] + fn anonymous_mapping_accesses_clear_idle_bits() { + use std::os::unix::fs::FileExt; + // SAFETY: geteuid has no arguments or memory effects. + if unsafe { libc::geteuid() } != 0 || !std::path::Path::new(BITMAP_PATH).exists() { + eprintln!("SKIP: page_idle requires root and CONFIG_IDLE_PAGE_TRACKING"); + return; + } + let Ok(bitmap) = open_bitmap() else { + eprintln!("SKIP: page_idle is not writable"); + return; + }; + // SAFETY: anonymous mapping, owned until munmap below. + let ptr = unsafe { + libc::mmap( + std::ptr::null_mut(), + 16 * 4096, + libc::PROT_READ | libc::PROT_WRITE, + libc::MAP_PRIVATE | libc::MAP_ANONYMOUS, + -1, + 0, + ) + }; + assert_ne!(ptr, libc::MAP_FAILED); + struct Mapping(*mut libc::c_void); + impl Drop for Mapping { + fn drop(&mut self) { + // SAFETY: this guard owns the complete mapping. + unsafe { + libc::munmap(self.0, 16 * 4096); + } + } + } + let _guard = Mapping(ptr); + // SAFETY: mapping is writable and the guard outlives this slice. + let mapping = unsafe { std::slice::from_raw_parts_mut(ptr.cast::(), 16 * 4096) }; + // Avoid THP: only the compound head carries an idle flag. + // SAFETY: this is the complete live mapping and advice changes no ownership. + assert_eq!( + unsafe { + libc::madvise( + mapping.as_mut_ptr().cast(), + mapping.len(), + libc::MADV_NOHUGEPAGE, + ) + }, + 0 + ); + for i in 0..16 { + mapping[i * 4096] = 1; + } + let start = mapping.as_ptr() as u64; + let vmas = [crate::dirty_map::GuestVma { + start, + end: start + mapping.len() as u64, + file_offset: 0, + }]; + let mut pfns = Vec::new(); + crate::dirty_map::walk_pagemap(std::process::id(), &vmas, |_, entry| { + pfns.push(present_pfn(entry)); + Ok(()) + }) + .unwrap(); + if pfns.iter().any(Option::is_none) { + eprintln!("SKIP: pagemap PFNs hidden (CAP_SYS_ADMIN required)"); + return; + } + // Newly faulted pages can still be in a per-CPU LRU batch. Establish + // an idle baseline before testing access-bit changes. + let mut marked = false; + for _ in 0..100 { + mark(std::process::id(), &vmas).unwrap(); + marked = pfns.iter().flatten().all(|pfn| { + let (word, bit) = pfn_word_bit(*pfn); + let mut bytes = [0; 8]; + bitmap.read_exact_at(&mut bytes, word).unwrap(); + u64::from_ne_bytes(bytes) & bit != 0 + }); + if marked { + break; + } + std::thread::sleep(std::time::Duration::from_millis(1)); + } + assert!(marked, "anonymous pages must accept idle marks"); + for i in 0..8 { + // SAFETY: each address is within the live mapping; volatile prevents elision. + unsafe { + std::ptr::write_volatile(mapping.as_mut_ptr().add(i * 4096), 2); + } + } + for (i, pfn) in pfns.iter().flatten().enumerate() { + let (word, bit) = pfn_word_bit(*pfn); + let mut bytes = [0; 8]; + bitmap.read_exact_at(&mut bytes, word).unwrap(); + assert_eq!( + u64::from_ne_bytes(bytes) & bit == 0, + i < 8, + "access classification for page {i}" + ); + } + } +} diff --git a/crates/engram-host-agent/src/pooled_backend.rs b/crates/engram-host-agent/src/pooled_backend.rs index a87f92d22..84fd13485 100644 --- a/crates/engram-host-agent/src/pooled_backend.rs +++ b/crates/engram-host-agent/src/pooled_backend.rs @@ -184,6 +184,7 @@ fn legacy_rootfs_path(cached: &CachedImage) -> Result { #[cfg_attr(not(target_os = "linux"), allow(dead_code))] #[derive(Clone)] struct PostCopyDestPending { + hot_trace_file: String, source_addr: String, export_id: String, } @@ -296,6 +297,8 @@ struct PendingPresetup { chain_ref: engram_core::types::manifest::ManifestRef, swap_manifest_ref: Option, sidecar_json: Vec, + idle_marked: bool, + hot_chunks: Vec<[u8; 32]>, } /// Issue #202: the unwind guard for a migration / snapshot capture's @@ -818,6 +821,8 @@ pub struct PooledBackend { /// failed manifest write is retried by the next call and never turns a /// complete drain into a rewind. drained_dests: Arc>, + /// Optional hint polling stops as soon as the memory drain terminates. + hint_pollers: Arc>, /// ADR 0044 K2 (issue #224): terminal-mode flag for graceful /// shutdown. `abandon_nbd_data_planes_for_shutdown` sets this /// `true` (SeqCst) BEFORE draining `nbd_sandboxes`, turning @@ -1015,18 +1020,73 @@ fn record_warm_hook_failure_metric(kind: engram_core::types::CaptureFailureKind) .increment(1); } -/// RAII: aborts the wrapped keepalive task on drop. A leg wrapped by +/// Abort the wrapped task on drop, including hint polling after DrainDone. +/// A leg wrapped by /// [`spawn_leg_keepalive`] stops resending stale progress the moment /// the guard goes out of scope — on every path, including an early /// `?`-return, since `Drop` runs during unwind too. -struct KeepaliveGuard(tokio::task::JoinHandle<()>); +struct AbortTaskOnDrop(tokio::task::JoinHandle<()>); -impl Drop for KeepaliveGuard { +impl Drop for AbortTaskOnDrop { fn drop(&mut self) { self.0.abort(); } } +/// Observe the local handler independently of the coordinator's drain-wait RPC. +/// ControlTx replays terminal reports to each subscriber, including late ones. +fn monitor_hint_poller(poller: AbortTaskOnDrop, sock: Option) -> AbortTaskOnDrop { + AbortTaskOnDrop(tokio::spawn(async move { + let mut poller = poller; + let Some(sock) = sock else { return }; + tokio::select! { + biased; + result = observe_memory_drain(&sock) => { + if let Err(error) = result { + tracing::info!(%error, "post-copy hint observer failed; stop optional polling"); + } + } + _ = &mut poller.0 => {} + } + // Drop also cancels an in-flight fetch. Destroy/restore failure cancels + // this owner and closes the local subscription without a blocking task. + })) +} + +async fn observe_memory_drain(sock: &Path) -> std::io::Result<()> { + use engram_migrate_proto::HandlerControl; + let mut stream = tokio::net::UnixStream::connect(sock).await?; + loop { + match crate::migrate_peer::read_frame_async(&mut stream).await? { + HandlerControl::DrainDone { .. } | HandlerControl::PeerLost { .. } => return Ok(()), + HandlerControl::Sealed { .. } | HandlerControl::DrainProgress { .. } => {} + } + } +} + +/// Retry only while the advisory budget remains. The owner cancels this future +/// at the local memory drain outcome, including during a fetch or backoff sleep. +#[cfg(any(target_os = "linux", test))] +async fn poll_hot_hint(budget: std::time::Duration, mut fetch: F) +where + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + let mut delay = std::time::Duration::from_millis(100); + let poll = async { + loop { + if fetch().await { + return; + } + tokio::time::sleep(delay).await; + delay = (delay * 2).min(std::time::Duration::from_secs(2)); + } + }; + if tokio::time::timeout(budget, poll).await.is_err() { + tracing::info!("post-copy hot trace fetch timed out; drain continues"); + } +} + /// Issue #539 finding 1: `run_warm_hook`'s own >=30s keepalive only /// covers the WARM phase. The BOOT leg (cold boot + chunk materialize to /// agentd-ready) and the SNAPSHOT leg (pause/flush/chunk memory + upload @@ -1046,8 +1106,8 @@ impl Drop for KeepaliveGuard { fn spawn_leg_keepalive( progress: tokio::sync::mpsc::Sender, event: engram_core::types::CaptureProgress, -) -> KeepaliveGuard { - KeepaliveGuard(tokio::spawn(async move { +) -> AbortTaskOnDrop { + AbortTaskOnDrop(tokio::spawn(async move { let mut interval = tokio::time::interval(crate::warm_progress::capture_keepalive_secs_from_env()); interval.tick().await; // consume the immediate first tick @@ -1058,6 +1118,50 @@ fn spawn_leg_keepalive( })) } +/// Publish the optional hint without changing the required disk seal. +#[cfg(any(target_os = "linux", test))] +fn publish_hot_hint(dir: &Path, hot: &[[u8; 32]]) -> std::io::Result<()> { + let tmp = dir.join("hot-hint.json.tmp"); + std::fs::write( + &tmp, + serde_json::to_vec(&serde_json::json!({"hot_chunks": hot}))?, + )?; + std::fs::rename(tmp, dir.join("hot-hint.json")) +} + +/// Parse the optional hint independently of the required disk seal. +#[cfg(any(target_os = "linux", test))] +fn parse_hot_hint(bytes: &[u8]) -> Option> { + let result = (|| { + let value: serde_json::Value = serde_json::from_slice(bytes)?; + value + .get("hot_chunks") + .map(|hot| serde_json::from_value(hot.clone())) + .transpose() + })(); + match result { + Ok(hot) => hot, + Err(error) => { + tracing::info!(%error, "post-copy hint malformed; drain continues"); + Some(Vec::new()) + } + } +} + +/// Atomically stage an advisory hint after state.bin has opened the load gate. +#[cfg(any(target_os = "linux", test))] +fn stage_hot_trace(path: &Path, hot: &[[u8; 32]]) -> std::io::Result<()> { + let mut trace = + engram_chunk_store::working_set::WorkingSetTrace::new(0, 0, chrono::DateTime::UNIX_EPOCH); + trace.chunks = hot + .iter() + .map(|h| engram_chunk_store::manifest::ChunkHash::from_bytes(*h)) + .collect(); + let tmp = path.with_extension("json.tmp"); + std::fs::write(&tmp, serde_json::to_vec(&trace)?)?; + std::fs::rename(tmp, path) +} + impl PooledBackend { /// Compose the pre-pause export from the current device lineages. async fn prepare_migration_export( @@ -1096,6 +1200,8 @@ impl PooledBackend { chain_ref, swap_manifest_ref, sidecar_json: sidecar_json.clone(), + idle_marked: false, + hot_chunks: self.read_hot_chunks(id), }; self.pending_presetups.insert(id, pending.clone()); @@ -1108,7 +1214,7 @@ impl PooledBackend { memory_manifest_ref: chain_ref, disk_manifest_ref, swap_manifest_ref: pending.swap_manifest_ref, - hot_chunks: self.read_hot_chunks(id), + hot_chunks: pending.hot_chunks, }) } @@ -1838,8 +1944,10 @@ impl PooledBackend { // chunk peer overlay before landing state.bin (the FC load // gate), and the publisher stays fenced until the disk drain // makes the dest self-sufficient. + #[cfg(not(target_os = "linux"))] + let hint_poller: Option = None; #[cfg(target_os = "linux")] - if let Some((_, pending)) = self.postcopy_dests.remove(&metadata.id) { + let hint_poller = if let Some((_, pending)) = self.postcopy_dests.remove(&metadata.id) { let nbd = [ (DiskRole::Root, pending_nbd_state.as_ref()), (DiskRole::Swap, pending_swap_state.as_ref()), @@ -1856,12 +1964,14 @@ impl PooledBackend { }) }) .collect(); - self.spawn_postcopy_fetch_poller( + Some(AbortTaskOnDrop(self.spawn_postcopy_fetch_poller( pending, self.inner.snapshot_path_for(metadata.id), nbd, - ); - } + ))) + } else { + None + }; #[cfg(target_os = "linux")] if pending_swap_state.is_none() @@ -2068,12 +2178,18 @@ impl PooledBackend { .resolved_dirty_root() .expect("NBD state requires a dirty root"); let migration_roles = self.migration_roles.clone(); + let hint_pollers = self.hint_pollers.clone(); let join = tokio::spawn(async move { let new_id = if fresh { inner.restore_fresh(metadata, selected_mounts).await? } else { inner.restore(metadata).await? }; + // Subscribe as soon as the backend exposes the restored VM. + // Replay covers a drain that finished during restore. + let hint_poller = hint_poller.map(|poller| { + monitor_hint_poller(poller, inner.post_copy_control_sock(new_id)) + }); if post_copy_dest { migration_roles.insert(new_id, crate::migration::MigrationRole::PostCopyDest); } @@ -2159,6 +2275,9 @@ impl PooledBackend { return Err(error); } } + if let Some(poller) = hint_poller { + hint_pollers.insert(new_id, poller); + } Ok::(new_id) }); // A JoinError here means the spawned task panicked; the FC restore @@ -2172,12 +2291,19 @@ impl PooledBackend { let inner = self.inner.clone(); let migration_roles = self.migration_roles.clone(); + let hint_pollers = self.hint_pollers.clone(); tokio::spawn(async move { let new_id = if fresh { inner.restore_fresh(metadata, selected_mounts).await? } else { inner.restore(metadata).await? }; + if let Some(poller) = hint_poller { + hint_pollers.insert( + new_id, + monitor_hint_poller(poller, inner.post_copy_control_sock(new_id)), + ); + } if post_copy_dest { migration_roles.insert(new_id, crate::migration::MigrationRole::PostCopyDest); } @@ -2260,6 +2386,7 @@ impl PooledBackend { postcopy_dests: Arc::new(DashMap::new()), migration_roles: Arc::new(DashMap::new()), drained_dests: Arc::new(DashMap::new()), + hint_pollers: Arc::new(DashMap::new()), #[cfg(target_os = "linux")] abandoning: Arc::new(std::sync::atomic::AtomicBool::new(false)), captures_quiesced: Arc::new(std::sync::atomic::AtomicBool::new(false)), @@ -3140,6 +3267,10 @@ impl PooledBackend { fs::create_dir_all(&dest) .await .map_err(|e| SandboxError::Snapshot(format!("create snapshot dir: {e}")))?; + let hot_trace_file = format!( + "migration-hot-{}.json", + crate::time_source::unique_path_token() + ); fs::write(dest.join("manifest.json"), &mig.sidecar_json) .await .map_err(|e| SandboxError::Snapshot(format!("write sidecar: {e}")))?; @@ -3150,10 +3281,10 @@ impl PooledBackend { .await .map_err(|e| SandboxError::Snapshot(format!("write session manifest: {e}")))?; let peer_spec = engram_sandbox_firecracker::MigrationPeerSpec { + hot_trace_file: hot_trace_file.clone(), peer_addr: peer_addr.clone(), export_id: mig.export_id.clone(), peer_token: peer_token.clone(), - hot_chunks: mig.hot_chunks.clone(), }; fs::write( dest.join(engram_sandbox_firecracker::MIGRATION_PEER_FILE), @@ -3165,6 +3296,7 @@ impl PooledBackend { self.postcopy_dests.insert( metadata.id, PostCopyDestPending { + hot_trace_file, source_addr: mig.source_addr.clone(), export_id: mig.export_id.clone(), }, @@ -3191,7 +3323,7 @@ impl PooledBackend { Arc, std::path::PathBuf, )>, - ) { + ) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let budget = std::time::Duration::from_secs(240); let started = crate::time_source::metrics_now(); @@ -3230,7 +3362,20 @@ impl PooledBackend { }, }; let t_fetch = crate::time_source::metrics_now(); - match Self::fetch_postcopy_artifacts(client, &pending.export_id, &tmp, &roles).await + match Self::fetch_postcopy_artifacts( + client, + &pending.export_id, + &tmp, + &roles + .iter() + .copied() + .map(engram_core::types::snapshot::MigrationItem::DiskSealInfo) + .chain(std::iter::once( + engram_core::types::snapshot::MigrationItem::StateBin, + )) + .collect::>(), + ) + .await { Ok(()) => break (tmp.join("state.bin"), t_fetch.elapsed().as_millis() as u64), Err(e) => { @@ -3365,7 +3510,46 @@ impl PooledBackend { fetch_ms, "post-copy restore inputs staged (state.bin landed)", ); - }); + // Demand faults can run now. The hint task polls for the + // optional artifact, so idle reads do not extend blackout. + poll_hot_hint(budget.saturating_sub(started.elapsed()), || async { + let Some(client) = source.as_ref() else { + return true; + }; + if let Err(error) = Self::fetch_postcopy_artifacts( + client, + &pending.export_id, + &tmp, + &[engram_core::types::snapshot::MigrationItem::HotHint], + ) + .await + { + if !matches!(error, SandboxError::NotFound) { + tracing::info!(%error, "post-copy hint fetch failed; drain continues"); + } + return false; + } + let bytes = match fs::read(tmp.join("hot-hint.json")).await { + Ok(bytes) => bytes, + Err(error) => { + tracing::info!(%error, "post-copy hint read failed; drain continues"); + return true; + } + }; + let Some(hot) = parse_hot_hint(&bytes) else { + return false; + }; + let path = dest_dir.join(&pending.hot_trace_file); + // File I/O is independent of the Tokio workers serving the move. + if let result @ (Ok(Err(_)) | Err(_)) = + tokio::task::spawn_blocking(move || stage_hot_trace(&path, &hot)).await + { + tracing::info!(?result, "post-copy hint staging failed; drain continues"); + } + true + }) + .await; + }) } /// Await the disk drain's terminal outcome. `None` subscription = @@ -3420,24 +3604,17 @@ impl PooledBackend { source: &engram_protocol::grpc_client::GrpcHostClient, export_id: &str, dir: &std::path::Path, - roles: &[DiskRole], + item_specs: &[engram_core::types::snapshot::MigrationItem], ) -> Result<(), SandboxError> { use engram_core::types::snapshot::MigrationItem; - let item_specs: Vec<_> = roles - .iter() - .copied() - .map(MigrationItem::DiskSealInfo) - .chain(std::iter::once(MigrationItem::StateBin)) - .collect(); let mut stream = source .migration_fetch(export_id, item_specs.to_vec()) - .await - .map_err(|e| SandboxError::Snapshot(format!("migration fetch: {e}")))?; + .await?; use futures::StreamExt; let mut current: Vec = Vec::new(); let mut current_idx: Option = None; while let Some(frame) = stream.next().await { - let frame = frame.map_err(|e| SandboxError::Snapshot(format!("fetch frame: {e}")))?; + let frame = frame?; if current_idx != Some(frame.item_idx) { if !current.is_empty() { return Err(SandboxError::Snapshot( @@ -3451,6 +3628,7 @@ impl PooledBackend { let idx = frame.item_idx as usize; let name = match item_specs.get(idx) { Some(MigrationItem::StateBin) => "state.bin", + Some(MigrationItem::HotHint) => "hot-hint.json", Some(MigrationItem::DiskSealInfo(role)) => role.seal_file_name(), _ => { return Err(SandboxError::Snapshot( @@ -7984,8 +8162,30 @@ impl SandboxBackend for PooledBackend { let backends = self.disk_backends(id); #[cfg(not(target_os = "linux"))] let backends = Vec::new(); - self.prepare_migration_export(id, chain_ref, &chain_manifest, peer.port(), backends) - .await + let out = self + .prepare_migration_export(id, chain_ref, &chain_manifest, peer.port(), backends) + .await?; + #[cfg(target_os = "linux")] + { + let pending_presetups = self.pending_presetups.clone(); + let export_id = out.export_id.clone(); + // Marking is advisory too: never make presetup wait for procfs or sysfs. + tokio::task::spawn_blocking(move || { + let started = crate::time_source::metrics_now(); + let idle_marked = crate::dirty_map::guest_vmas(view.fc_pid, &view.memory_backing) + .and_then(|vmas| crate::page_idle::mark(view.fc_pid, &vmas)) + .inspect_err(crate::page_idle::note_unavailable) + .is_ok(); + if let Some(mut pending) = pending_presetups.get_mut(&id) { + if pending.export_id == export_id { + pending.idle_marked = idle_marked; + } + } + tracing::info!(sandbox_id = %id, idle_mark_ms = started.elapsed().as_millis() as u64, + idle_marked, "post-copy idle mark complete"); + }); + } + Ok(out) } /// ADR 0045 C2: the blackout half. Pause → NBD host-cache fsync + @@ -8193,8 +8393,9 @@ impl SandboxBackend for PooledBackend { .values() .map(|seal| seal.len() as u64) .sum::(); - for (role, info) in descriptors { - let bytes = serde_json::to_vec(&info) + // Publish immutable, hint-free disk seals before registering the export. + for (role, info) in &descriptors { + let bytes = serde_json::to_vec(info) .map_err(|e| SandboxError::Snapshot(format!("disk seal info: {e}")))?; tokio::fs::write(export_dir.join(role.seal_file_name()), bytes) .await @@ -8215,7 +8416,7 @@ impl SandboxBackend for PooledBackend { let inserted = self.migrations.insert(crate::migration::MigrationExport { export_id: export_id.to_string(), sandbox_id: id, - snapshot_dir: export_dir, + snapshot_dir: export_dir.clone(), disk_seal, clock: self.clock.clone(), state_served, @@ -8248,7 +8449,7 @@ impl SandboxBackend for PooledBackend { token: pending.peer_token, sandbox_id: id, fc_pid: view.fc_pid, - vmas, + vmas: vmas.clone(), seal, durable_at, chunk_size, @@ -8258,6 +8459,42 @@ impl SandboxBackend for PooledBackend { clock: self.clock.clone(), }); + let fallback = pending.hot_chunks; + let idle_marked = pending.idle_marked; + tokio::task::spawn_blocking(move || { + let started = crate::time_source::metrics_now(); + let mut pages = Vec::new(); + let captured = !idle_marked + || crate::dirty_map::walk_pagemap(view.fc_pid, &vmas, |offset, entry| { + if offset < total_bytes { + if let Some(pfn) = crate::page_idle::present_pfn(entry) { + pages.push((offset, pfn)); + } + } + Ok(()) + }) + .inspect_err(crate::page_idle::note_unavailable) + .is_ok(); + let pfn_capture_ms = started.elapsed().as_millis() as u64; + let offsets = if idle_marked && captured { + match crate::page_idle::scan(&pages, chunk_size) { + Ok(offsets) => offsets, + Err(error) => { + crate::page_idle::note_unavailable(&error); + Vec::new() + } + } + } else { + Vec::new() + }; + let hot = crate::page_idle::hot_hashes(&chain_manifest, &offsets, &fallback); + tracing::info!(sandbox_id = %id, idle_scan_ms = started.elapsed().as_millis() as u64, + pfn_capture_ms, hot_chunks = hot.len(), "post-copy idle scan complete after export registration"); + if let Err(error) = publish_hot_hint(&export_dir, &hot) { + tracing::info!(%error, "post-copy hot hint unavailable; drain continues"); + } + }); + tracing::info!( sandbox_id = %id, export_id, @@ -8314,6 +8551,7 @@ impl SandboxBackend for PooledBackend { .into_iter() .map(|(_, backend)| backend.postcopy_drain_subscribe()) .collect(); + let hint_pollers = self.hint_pollers.clone(); let outcome = tokio::task::spawn_blocking(move || -> Result { let mut stream = std::os::unix::net::UnixStream::connect(&sock) .map_err(|e| format!("dial control sock: {e}"))?; @@ -8335,6 +8573,7 @@ impl SandboxBackend for PooledBackend { fault_max_us, fault_around_chunks_installed, } => { + hint_pollers.remove(&id); // Restore-tail attribution: the fault-path // totals are the serial P2P cost inside the FC // load + early guest execution (cross-ref with @@ -8354,6 +8593,7 @@ impl SandboxBackend for PooledBackend { }); } engram_migrate_proto::HandlerControl::PeerLost { remaining, detail } => { + hint_pollers.remove(&id); return Ok(DrainOutcome::PeerLost { remaining, detail }); } engram_migrate_proto::HandlerControl::Sealed { .. } @@ -8424,9 +8664,15 @@ impl SandboxBackend for PooledBackend { let Some(export) = self.migrations.find_by_export_id(export_id) else { return Err(SandboxError::NotFound); }; - // Artifact serves refresh the TTL. Once state.bin is served, - // the destination may run: the TTL sweep must not resume us. - export.touch(); + // Advisory hints never extend the export lease, even when absent. + // Required artifacts and disk chunks still count as activity. + if items + .iter() + .any(|item| !matches!(item, MigrationItem::HotHint)) + { + export.touch(); + } + // Once state.bin is served, the TTL sweep must not resume us. if items.iter().any(|i| matches!(i, MigrationItem::StateBin)) { export .state_served @@ -8444,6 +8690,16 @@ impl SandboxBackend for PooledBackend { .await .map(bytes::Bytes::from) .map_err(|e| SandboxError::Snapshot(format!("read state.bin: {e}"))), + MigrationItem::HotHint => fs::read(snapshot_dir.join("hot-hint.json")) + .await + .map(bytes::Bytes::from) + .map_err(|e| { + if e.kind() == std::io::ErrorKind::NotFound { + SandboxError::NotFound + } else { + SandboxError::Snapshot(format!("read hot hint: {e}")) + } + }), MigrationItem::DiskSealInfo(role) => { fs::read(snapshot_dir.join(role.seal_file_name())) .await @@ -8800,6 +9056,7 @@ impl SandboxBackend for PooledBackend { async fn destroy(&self, id: SandboxId) -> Result<(), SandboxError> { // A held source forgets its swap policy and metadata with the VM. self.held_snapshots.remove(&id); + self.hint_pollers.remove(&id); self.teardowns .run_or_join(id, || { let session_bindings = self.session_bindings.clone(); @@ -12576,6 +12833,242 @@ mod tests { assert!(frames.iter().any(|f| f.item_idx == 1 && f.last)); } + /// The first descriptor permits resume; a later update carries the final + /// order. The handler loader must wait and preserve that order exactly. + #[test] + fn malformed_hint_is_separate_from_required_seal_fields() { + for hot in [ + serde_json::json!("invalid"), + serde_json::json!([[1, 2]]), + serde_json::Value::Null, + ] { + let bytes = serde_json::to_vec( + &serde_json::json!({"sealed_chunk_indices": [], "hot_chunks": hot}), + ) + .unwrap(); + assert_eq!(parse_hot_hint(&bytes), Some(Vec::new())); + } + assert_eq!(parse_hot_hint(br#"{"hot_chunks":[]}"#), Some(vec![])); + assert!(parse_hot_hint(br#"{}"#).is_none()); + } + + #[tokio::test(start_paused = true)] + async fn hot_hint_polling_has_bounded_requests_and_budget() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + let requests = AtomicUsize::new(0); + let started = tokio::time::Instant::now(); + poll_hot_hint(Duration::from_secs(240), || async { + requests.fetch_add(1, Ordering::SeqCst); + false // Publication failed: every request returns NotFound. + }) + .await; + assert!(requests.load(Ordering::SeqCst) <= 125); + assert!(requests.load(Ordering::SeqCst) > 5); + assert_eq!(started.elapsed(), Duration::from_secs(240)); + + // An unresponsive fetch cannot extend the remaining budget either. + let started = tokio::time::Instant::now(); + poll_hot_hint(Duration::from_millis(37), std::future::pending).await; + assert_eq!(started.elapsed(), Duration::from_millis(37)); + } + + #[tokio::test(start_paused = true)] + async fn hot_hint_polls_stop_on_local_drain_without_drain_wait() { + use engram_migrate_proto::HandlerControl; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + use tokio::io::AsyncWriteExt; + + // Socket I/O must not let Tokio auto-advance the paused clock. Only + // the explicit advances below may expire a polling backoff. + let _keep_clock_paused = AbortTaskOnDrop(tokio::spawn(async { + loop { + tokio::task::yield_now().await; + } + })); + for terminal in [ + HandlerControl::DrainDone { + pulled: 1, + alt_sourced: 0, + zero_chunks: 0, + ms: 1, + faults: 0, + fault_us: 0, + fault_max_us: 0, + fault_around_chunks_installed: 0, + }, + HandlerControl::PeerLost { + remaining: 1, + detail: "peer closed".into(), + }, + ] { + let tmp = tempfile::tempdir().unwrap(); + let sock = tmp.path().join("control.sock"); + let listener = tokio::net::UnixListener::bind(&sock).unwrap(); + let (ready_tx, ready_rx) = tokio::sync::oneshot::channel(); + let (report_tx, report_rx) = tokio::sync::oneshot::channel(); + let mut source = AbortTaskOnDrop(tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + ready_tx.send(()).unwrap(); + report_rx.await.unwrap(); + let mut bytes = Vec::new(); + engram_migrate_proto::write_frame(&mut bytes, &terminal).unwrap(); + stream.write_all(&bytes).await.unwrap(); + })); + let requests = Arc::new(AtomicUsize::new(0)); + let count = requests.clone(); + let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let (finished_tx, finished_rx) = tokio::sync::oneshot::channel(); + let (request_tx, mut request_rx) = tokio::sync::mpsc::unbounded_channel(); + let poller = AbortTaskOnDrop(tokio::spawn(async move { + start_rx.await.unwrap(); + poll_hot_hint(Duration::from_secs(240), || async { + count.fetch_add(1, Ordering::SeqCst); + request_tx.send(()).unwrap(); + false + }) + .await; + // Cancellation drops this sender instead of completing normally. + finished_tx.send(()).unwrap(); + })); + let mut observer = monitor_hint_poller(poller, Some(sock)); + ready_rx.await.unwrap(); + start_tx.send(()).unwrap(); + request_rx.recv().await.unwrap(); + assert_eq!(requests.load(Ordering::SeqCst), 1); + let intervals = [100, 200, 400]; + for (index, millis) in intervals.into_iter().enumerate() { + tokio::time::advance(Duration::from_millis(millis)).await; + request_rx.recv().await.unwrap(); + assert_eq!(requests.load(Ordering::SeqCst), index + 2); + } + // No coordinator calls migration_drain_wait. The local report + // must cancel polling while the hint remains unavailable. + report_tx.send(()).unwrap(); + (&mut source.0).await.unwrap(); + (&mut observer.0).await.unwrap(); + assert!(finished_rx.await.is_err(), "poller must be cancelled"); + let expected = 1 + intervals.len(); + assert_eq!(requests.load(Ordering::SeqCst), expected); + tokio::time::advance(Duration::from_millis(800)).await; + assert_eq!(requests.load(Ordering::SeqCst), expected); + } + } + + #[tokio::test] + async fn capture_hot_hint_leaves_initial_artifacts_unchanged() { + use engram_core::types::snapshot::MigrationItem; + use futures::StreamExt; + let tmp = tempfile::tempdir().unwrap(); + let inner = Arc::new(engram_sandbox_process::ProcessBackend::new(tmp.path())); + let pooled = PooledBackend::new(inner); + let export_dir = tmp.path().join("export"); + let dest = tmp.path().join("dest"); + std::fs::create_dir_all(&export_dir).unwrap(); + std::fs::create_dir_all(&dest).unwrap(); + let descriptor = serde_json::json!({"sealed_chunk_indices": [0], "manifest": null, "disk_ref_bincode": null}); + std::fs::write( + export_dir.join("disk-seal.json"), + serde_json::to_vec(&descriptor).unwrap(), + ) + .unwrap(); + std::fs::write(export_dir.join("state.bin"), b"state").unwrap(); + let id = SandboxId::new(); + let guard = pooled.capture_lock(id).lock_owned().await; + assert!(pooled.migrations.insert(crate::migration::MigrationExport { + export_id: "hot-export".into(), + sandbox_id: id, + snapshot_dir: export_dir.clone(), + disk_seal: Default::default(), + clock: pooled.clock.clone(), + state_served: Arc::new(std::sync::atomic::AtomicBool::new(false)), + last_activity: Arc::new(std::sync::Mutex::new(pooled.clock.now_mono())), + capture_guard: guard, + })); + let trace_path = dest.join(engram_sandbox_firecracker::MIGRATION_HOT_TRACE_FILE); + let mut loader = Box::pin(engram_uffd_handler::working_set::load_migration_trace( + &trace_path, + )); + assert!(futures::poll!(&mut loader).is_pending()); + let descriptor_bytes = std::fs::read(export_dir.join("disk-seal.json")).unwrap(); + // Even a large, already published hint cannot enlarge the initial fetch. + let hot = vec![[255; 32]; 4096]; + publish_hot_hint(&export_dir, &hot).unwrap(); + assert_eq!( + std::fs::read(export_dir.join("disk-seal.json")).unwrap(), + descriptor_bytes + ); + assert_eq!( + std::fs::metadata(export_dir.join("disk-seal.json")) + .unwrap() + .len(), + descriptor_bytes.len() as u64 + ); + let frames: Vec<_> = pooled + .migration_fetch( + "hot-export", + vec![ + MigrationItem::DiskSealInfo(DiskRole::Root), + MigrationItem::StateBin, + ], + ) + .await + .unwrap() + .collect() + .await; + assert_eq!(frames.len(), 2); + assert_eq!(frames[0].as_ref().unwrap().data.as_ref(), descriptor_bytes); + assert_eq!(frames[1].as_ref().unwrap().data.as_ref(), b"state"); + // Both missing and published hints must leave the lease untouched. + let last_activity = pooled + .migrations + .find_by_export_id("hot-export") + .unwrap() + .last_activity + .clone(); + let anchor = *last_activity.lock().unwrap(); + std::fs::remove_file(export_dir.join("hot-hint.json")).unwrap(); + for _ in 0..3 { + let frames: Vec<_> = pooled + .migration_fetch("hot-export", vec![MigrationItem::HotHint]) + .await + .unwrap() + .collect() + .await; + assert!(matches!(&frames[0], Err(SandboxError::NotFound))); + assert_eq!(*last_activity.lock().unwrap(), anchor); + } + publish_hot_hint(&export_dir, &[[2; 32], [1; 32]]).unwrap(); + let frames: Vec<_> = pooled + .migration_fetch("hot-export", vec![MigrationItem::HotHint]) + .await + .unwrap() + .collect() + .await; + let bytes: Vec<_> = frames + .into_iter() + .flat_map(|f| f.unwrap().data.to_vec()) + .collect(); + assert_eq!(*last_activity.lock().unwrap(), anchor); + let info: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + let hot: Vec<[u8; 32]> = serde_json::from_value(info["hot_chunks"].clone()).unwrap(); + stage_hot_trace(&trace_path, &hot).unwrap(); + let trace = loader.await.unwrap(); + assert_eq!( + trace + .chunks + .iter() + .map(|h| *h.as_bytes()) + .collect::>(), + vec![[2; 32], [1; 32]] + ); + assert_eq!( + std::fs::read(export_dir.join("disk-seal.json")).unwrap(), + descriptor_bytes + ); + } + #[tokio::test] async fn migration_presetup_without_chain_signals_fallback() { let tmp = tempfile::tempdir().unwrap(); @@ -12623,6 +13116,8 @@ mod tests { chain_ref: engram_core::types::manifest::ManifestRef::new(), swap_manifest_ref: None, sidecar_json: Vec::new(), + idle_marked: false, + hot_chunks: Vec::new(), }, ); @@ -12826,10 +13321,40 @@ mod tests { /// Inner backend double that records `resume` calls so the test can /// assert the unwind guard un-pauses the guest. struct ResumeSpy { + capture_root: Option, resumes: std::sync::Arc>>, } #[async_trait] impl SandboxBackend for ResumeSpy { + fn post_copy_source_view( + &self, + _: SandboxId, + ) -> Option { + Some(engram_core::traits::sandbox::PostCopySourceView { + fc_pid: std::process::id(), + memory_backing: self.capture_root.as_ref()?.join("ram"), + }) + } + fn working_set_trace_path(&self, _: SandboxId) -> Option { + Some(self.capture_root.as_ref()?.join("trace.json")) + } + async fn pause(&self, _: SandboxId) -> Result<(), SandboxError> { + Ok(()) + } + async fn snapshot_vmstate_only_package( + &self, + _: SandboxId, + _: &[u8], + ) -> Result<(engram_core::SnapshotId, PathBuf), SandboxError> { + let root = self + .capture_root + .as_ref() + .expect("capture fixture") + .join("export"); + std::fs::create_dir_all(&root).unwrap(); + std::fs::write(root.join("state.bin"), b"state").unwrap(); + Ok((engram_core::SnapshotId::new(), root)) + } fn compose_live_sidecar( &self, _: SandboxId, @@ -12868,6 +13393,246 @@ mod tests { } } + #[cfg(target_os = "linux")] + #[tokio::test] + async fn fake_capture_and_destination_poller_deliver_final_hot_trace() { + use engram_chunk_store::manifest::{ChunkHash, ChunkRef, ChunkSize, ManifestKind}; + use std::os::fd::AsRawFd; + let tmp = tempfile::tempdir().unwrap(); + let file = std::fs::File::create(tmp.path().join("ram")).unwrap(); + file.set_len(4096).unwrap(); + drop(file); + let file = std::fs::File::open(tmp.path().join("ram")).unwrap(); + // SAFETY: private file mapping owned by the guard below. + let ptr = unsafe { + libc::mmap( + std::ptr::null_mut(), + 4096, + libc::PROT_READ, + libc::MAP_PRIVATE, + file.as_raw_fd(), + 0, + ) + }; + assert_ne!(ptr, libc::MAP_FAILED); + struct Mapping(*mut libc::c_void); + impl Drop for Mapping { + fn drop(&mut self) { + // SAFETY: this guard owns the complete live mapping. + unsafe { + libc::munmap(self.0, 4096); + } + } + } + let _mapping = Mapping(ptr); + let inner = Arc::new(ResumeSpy { + capture_root: Some(tmp.path().to_path_buf()), + resumes: Default::default(), + }); + let source = Arc::new(PooledBackend::new(inner)); + source.set_migrate_peer_server(crate::migrate_peer::PeerServer::new(0)); + let mut trace = engram_chunk_store::working_set::WorkingSetTrace::new( + 0, + 0, + chrono::DateTime::UNIX_EPOCH, + ); + trace.chunks = vec![ChunkHash::from_bytes([7; 32])]; + std::fs::write( + tmp.path().join("trace.json"), + serde_json::to_vec(&trace).unwrap(), + ) + .unwrap(); + let manifest = engram_chunk_store::Manifest { + schema_version: 1, + kind: ManifestKind::Memory, + chunk_size: ChunkSize::bytes(4096), + total_bytes: 4096, + chunks: vec![ChunkRef { + offset: 0, + hash: trace.chunks[0], + }], + parent: None, + working_set_trace: None, + annotations: serde_json::Value::Null, + }; + let id = SandboxId::new(); + let chain_ref = engram_core::types::manifest::ManifestRef::new(); + let pre = source + .prepare_migration_export(id, chain_ref, &manifest, 0, vec![]) + .await + .unwrap(); + source.checkpoint_chains.insert( + id, + crate::checkpoint::CheckpointChain { + manifest_ref: chain_ref, + manifest, + }, + ); + let capture = source + .migration_capture_postcopy(id, &pre.export_id) + .await + .unwrap(); + assert_eq!(capture.total_chunks, 1); + // Wait for the independent publisher before the first initial fetch. + let seal_path = tmp.path().join("export/disk-seal.json"); + let descriptor_bytes = std::fs::read(&seal_path).unwrap(); + let hint_path = tmp.path().join("export/hot-hint.json"); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while !hint_path.exists() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + for case in [ + "published", + "fetch-failed", + "fetch-stage-failed", + "unavailable", + "empty", + "delayed", + "malformed", + "write-failed", + "rename-failed", + ] { + let _ = std::fs::remove_file(&hint_path); + let _ = std::fs::remove_dir(&hint_path); + match case { + "published" | "fetch-stage-failed" => { + publish_hot_hint(&tmp.path().join("export"), &vec![[7; 32]; 4096]).unwrap() + } + "fetch-failed" => std::fs::create_dir(&hint_path).unwrap(), + "empty" => publish_hot_hint(&tmp.path().join("export"), &[]).unwrap(), + "malformed" => std::fs::write(&hint_path, br#"{"hot_chunks":"invalid"}"#).unwrap(), + "write-failed" | "rename-failed" => { + publish_hot_hint(&tmp.path().join("export"), &[[7; 32]]).unwrap() + } + _ => {} + } + assert_eq!(std::fs::read(&seal_path).unwrap(), descriptor_bytes); + assert_eq!( + std::fs::metadata(&seal_path).unwrap().len(), + descriptor_bytes.len() as u64 + ); + let listen = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listen.local_addr().unwrap(); + drop(listen); + let client: Arc = Arc::new( + crate::host_client::LocalHostClient::with_noop_hub(source.clone()), + ); + let server = tokio::spawn(crate::grpc_server::boot( + addr, + client, + None, + crate::session_epochs::ephemeral(), + None, + )); + let dest = tmp.path().join(case); + std::fs::create_dir_all(&dest).unwrap(); + let destination = PooledBackend::new(Arc::new( + engram_sandbox_process::ProcessBackend::new(tmp.path().join("dest-inner")), + )); + let path = dest.join(engram_sandbox_firecracker::MIGRATION_HOT_TRACE_FILE); + if case == "write-failed" { + std::fs::create_dir(path.with_extension("json.tmp")).unwrap(); + } + if case == "rename-failed" { + std::fs::create_dir(&path).unwrap(); + } + if case == "fetch-stage-failed" { + std::fs::create_dir_all(dest.join("postcopy-tmp/hot-hint.json")).unwrap(); + } + let poller = destination.spawn_postcopy_fetch_poller( + PostCopyDestPending { + hot_trace_file: engram_sandbox_firecracker::MIGRATION_HOT_TRACE_FILE.into(), + source_addr: format!("http://{addr}"), + export_id: pre.export_id.clone(), + }, + dest.clone(), + vec![], + ); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while !dest.join("state.bin").exists() { + tokio::time::sleep(std::time::Duration::from_millis(2)).await; + } + }) + .await + .unwrap(); + assert_eq!(std::fs::read(dest.join("state.bin")).unwrap(), b"state"); + if matches!(case, "fetch-failed" | "fetch-stage-failed" | "unavailable") { + let client = PooledBackend::dial_migration_source(&format!("http://{addr}")) + .await + .unwrap(); + let error = PooledBackend::fetch_postcopy_artifacts( + &client, + &pre.export_id, + &dest.join("postcopy-tmp"), + &[engram_core::types::snapshot::MigrationItem::HotHint], + ) + .await + .unwrap_err(); + if case == "unavailable" { + assert!(matches!(error, SandboxError::NotFound)); + } + assert_eq!(std::fs::read(dest.join("state.bin")).unwrap(), b"state"); + } + if case == "delayed" { + assert!(!path.exists()); + publish_hot_hint(&tmp.path().join("export"), &[[7; 32]]).unwrap(); + let loaded = tokio::time::timeout( + std::time::Duration::from_secs(5), + engram_uffd_handler::working_set::load_migration_trace(&path), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(loaded.chunks, trace.chunks); + } + if case == "published" { + let loaded = tokio::time::timeout( + std::time::Duration::from_secs(5), + engram_uffd_handler::working_set::load_migration_trace(&path), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(loaded.chunks.len(), 4096); + } + assert_eq!( + std::fs::read(dest.join("postcopy-tmp/disk-seal.json")).unwrap(), + descriptor_bytes + ); + // A local terminal report stops the real fetch loop without a + // migration_drain_wait RPC, even when publication failed. + let sock = dest.join("control.sock"); + let control = engram_uffd_handler::peer::ControlTx::bind(&sock).unwrap(); + control.report(engram_migrate_proto::HandlerControl::DrainDone { + pulled: 1, + alt_sourced: 0, + zero_chunks: 0, + ms: 1, + faults: 0, + fault_us: 0, + fault_max_us: 0, + fault_around_chunks_installed: 0, + }); + let abort = poller.abort_handle(); + let mut observer = monitor_hint_poller(AbortTaskOnDrop(poller), Some(sock)); + tokio::time::timeout(std::time::Duration::from_secs(5), &mut observer.0) + .await + .unwrap() + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while !abort.is_finished() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + server.abort(); + } + } + /// Build a `ChunkedDiskBackend` over a 4-chunk sparse base with one /// dirty chunk written, plus its store. Mirrors the disk_daemon /// `build_backend` test helper. @@ -12922,6 +13687,7 @@ mod tests { let (root, store, _root_dir) = unwind_test_disk_backend().await; let (swap, _, _swap_dir) = unwind_test_disk_backend().await; let pooled = PooledBackend::new(Arc::new(ResumeSpy { + capture_root: None, resumes: Default::default(), })); let root_ref = root.manifest_ref().await; @@ -12953,6 +13719,7 @@ mod tests { swap.write(0, &[0x73; 4096]).await.unwrap(); let resumes = Arc::new(parking_lot::Mutex::new(Vec::new())); let pooled = PooledBackend::new(Arc::new(ResumeSpy { + capture_root: None, resumes: resumes.clone(), })); let id = SandboxId::new(); @@ -12997,6 +13764,7 @@ mod tests { let (backend, store, _dir) = unwind_test_disk_backend().await; let resumes = std::sync::Arc::new(parking_lot::Mutex::new(Vec::new())); let inner: Arc = Arc::new(ResumeSpy { + capture_root: None, resumes: resumes.clone(), }); let id = SandboxId::new(); @@ -13060,6 +13828,7 @@ mod tests { let (backend, _store, _dir) = unwind_test_disk_backend().await; let resumes = std::sync::Arc::new(parking_lot::Mutex::new(Vec::new())); let inner: Arc = Arc::new(ResumeSpy { + capture_root: None, resumes: resumes.clone(), }); let id = SandboxId::new(); diff --git a/crates/engram-host-agent/tests/two_host_live_teleport.rs b/crates/engram-host-agent/tests/two_host_live_teleport.rs index 2b5eaed07..9f94e4703 100644 --- a/crates/engram-host-agent/tests/two_host_live_teleport.rs +++ b/crates/engram-host-agent/tests/two_host_live_teleport.rs @@ -21,6 +21,27 @@ use futures::StreamExt; mod common; +fn idle_hint_probe() -> bool { + use std::os::unix::fs::FileExt; + let probe = || -> std::io::Result<()> { + engram_host_agent::page_idle::open_bitmap()?; + let page = Box::new([1u8; 4096]); + std::hint::black_box(&page); + let mut bytes = [0; 8]; + std::fs::File::open("/proc/self/pagemap")? + .read_exact_at(&mut bytes, (page.as_ptr() as u64 / 4096) * 8)?; + let pfn = engram_host_agent::page_idle::present_pfn(u64::from_ne_bytes(bytes)) + .ok_or_else(|| std::io::Error::other("PFN hidden or shared"))?; + std::fs::File::open("/proc/kpageflags")?.read_exact_at(&mut bytes, pfn * 8)?; + let bitmap = engram_host_agent::page_idle::open_bitmap()?; + let (word, bit) = engram_host_agent::page_idle::pfn_word_bit(pfn); + bitmap.write_all_at(&bit.to_ne_bytes(), word)?; + bitmap.read_exact_at(&mut bytes, word)?; + Ok(()) + }; + probe().is_ok() +} + struct HostStack { pooled: Arc, addr: std::net::SocketAddr, @@ -284,8 +305,33 @@ async fn two_host_live_teleport_preserves_post_checkpoint_state() { assert_eq!(alive.trim(), "alive", "harness alive on A"); // ---- The move, over the wire (the G1 downtime legs) ---- - let (moved, cap) = - common::postcopy::move_guest(&client_a, &client_b, vm, ckpt.clone(), host_a.addr).await; + let cap = client_a + .migration_presetup(vm, engram_core::traits::SessionFence::unfenced()) + .await + .expect("presetup"); + // Allocate and write private shell memory after the idle mark. + let touched = exec( + &host_a.pooled, + vm, + "buf=$(head -c 65536 /dev/zero | tr '\\000' x); echo ${#buf}", + ) + .await; + assert_eq!(touched.trim(), "65536"); + let row = common::postcopy::metadata(ckpt.clone(), &cap, host_a.addr); + let dest = client_b.clone(); + let restore = tokio::spawn(async move { + dest.restore(row, engram_core::traits::SessionFence::unfenced()) + .await + }); + client_a + .migration_capture_postcopy( + vm, + &cap.export_id, + engram_core::traits::SessionFence::unfenced(), + ) + .await + .expect("capture"); + let moved = restore.await.unwrap().expect("restore"); common::postcopy::drain(&host_b.pooled, moved).await; let stats_path = host_b .pooled @@ -306,6 +352,43 @@ async fn two_host_live_teleport_preserves_post_checkpoint_state() { .await, "handler must report the demand-window install counter" ); + if idle_hint_probe() { + let path = host_b + .pooled + .prefault_stats_path(moved) + .expect("handler stats path"); + assert!( + common::poll_until_async( + std::time::Duration::from_secs(10), + std::time::Duration::from_millis(10), + || { + let path = path.clone(); + async move { + tokio::fs::read(path) + .await + .ok() + .and_then(|bytes| { + serde_json::from_slice::(&bytes).ok() + }) + .is_some_and(|stats| stats["hot_hint_received"].as_bool().is_some()) + } + } + ) + .await, + "drain stats must be published" + ); + let stats: serde_json::Value = + serde_json::from_slice(&tokio::fs::read(path).await.unwrap()).unwrap(); + assert_eq!( + stats["hot_hint_received"].as_bool(), + Some(true), + "the hint must arrive when the idle-page probes pass" + ); + } else { + eprintln!( + "SKIP hot_hint_received assertion: page_idle, pagemap PFNs, or kpageflags unavailable" + ); + } // Commit destroys A's frozen source; the moved VM lives on B. client_a diff --git a/crates/engram-protocol/proto/host_service.proto b/crates/engram-protocol/proto/host_service.proto index 7ab267d9b..659e92f4a 100644 --- a/crates/engram-protocol/proto/host_service.proto +++ b/crates/engram-protocol/proto/host_service.proto @@ -395,6 +395,7 @@ message MigrationItem { STATE_BIN = 1; DISK_SEAL_INFO = 4; DISK_CHUNK_AT = 5; + HOT_HINT = 6; } reserved 2; reserved "hash"; diff --git a/crates/engram-protocol/src/grpc_client.rs b/crates/engram-protocol/src/grpc_client.rs index ab2685375..ba273c95c 100644 --- a/crates/engram-protocol/src/grpc_client.rs +++ b/crates/engram-protocol/src/grpc_client.rs @@ -441,6 +441,11 @@ impl GrpcHostClient { let wire_items = items .into_iter() .map(|item| match item { + engram_core::types::snapshot::MigrationItem::HotHint => MigrationItem { + kind: Kind::HotHint as i32, + chunk_idx: 0, + role: crate::grpc::DiskRole::Root as i32, + }, engram_core::types::snapshot::MigrationItem::StateBin => MigrationItem { kind: Kind::StateBin as i32, chunk_idx: 0, diff --git a/crates/engram-protocol/src/wire.rs b/crates/engram-protocol/src/wire.rs index ce72523d5..224824d3b 100644 --- a/crates/engram-protocol/src/wire.rs +++ b/crates/engram-protocol/src/wire.rs @@ -170,7 +170,8 @@ use serde::{Deserialize, Serialize}; // the same snapshot with an untyped error, which the resume verb counts toward // the Dead budget instead of recovering from disk. // v36 gates page-channel v6 fault windows for a lockstep host roll. -pub const WIRE_VERSION: u32 = 36; +// v37 appends HostCapabilities.page_idle and the optional MigrationItem::HotHint. +pub const WIRE_VERSION: u32 = 37; /// gRPC metadata (header) key carrying the caller's [`WIRE_VERSION`] on /// every coord→host request (issue #229). ASCII, lowercase — tonic diff --git a/crates/engram-protocol/tests/wire_golden.rs b/crates/engram-protocol/tests/wire_golden.rs index 6272f248f..6a7fc0101 100644 --- a/crates/engram-protocol/tests/wire_golden.rs +++ b/crates/engram-protocol/tests/wire_golden.rs @@ -446,6 +446,30 @@ fn warm_stages() -> Vec { // ---- tests ------------------------------------------------------------- +#[test] +fn migration_item_indices_stay_fixed() { + use engram_core::types::snapshot::{DiskRole, MigrationItem}; + use engram_protocol::grpc::migration_item::Kind; + for (item, index) in [ + (MigrationItem::StateBin, 0), + (MigrationItem::DiskSealInfo(DiskRole::Root), 1), + (MigrationItem::DiskChunkAt(DiskRole::Root, 9), 2), + (MigrationItem::HotHint, 3), + ] { + // Exhaustive guard: every new kind needs an explicit wire pin. + let kind = match item { + MigrationItem::StateBin => Kind::StateBin, + MigrationItem::DiskSealInfo(_) => Kind::DiskSealInfo, + MigrationItem::DiskChunkAt(_, _) => Kind::DiskChunkAt, + MigrationItem::HotHint => Kind::HotHint, + }; + assert_variant_index(&item, index, "MigrationItem"); + assert_eq!(kind as i32, [1, 4, 5, 6][index as usize]); + let bytes = bincode::serialize(&item).unwrap(); + assert_eq!(bincode::deserialize::(&bytes).unwrap(), item); + } +} + #[test] fn struct_payloads_golden() { // SandboxSpec / AgentSpec / SnapshotMetadata aren't `PartialEq`; @@ -701,8 +725,12 @@ fn wire_version_pinned() { // bump fences old hosts that refuse untyped. // 35 -> 36: page-channel v6 fault windows require a lockstep host roll. // Coordinator bincode payload bytes are unchanged. + // 36 -> 37: HostCapabilities appends page_idle. + // MigrationItem appends HotHint (bincode 3, protobuf 6). Existing item + // indices stay fixed; the optional hint is separate from immutable disk seals. + // Existing goldens embed neither payload and keep their bytes. assert_eq!( - WIRE_VERSION, 36, + WIRE_VERSION, 37, "WIRE_VERSION changed — confirm payload goldens were regenerated too" ); } diff --git a/crates/engram-sandbox-firecracker/src/lib.rs b/crates/engram-sandbox-firecracker/src/lib.rs index 8d1cca2dc..7cd048e6e 100644 --- a/crates/engram-sandbox-firecracker/src/lib.rs +++ b/crates/engram-sandbox-firecracker/src/lib.rs @@ -428,13 +428,15 @@ pub const UFFD_CONTROL_SOCK_FILE: &str = "uffd-control.sock"; /// state.bin appears late via the fetch poller). pub const MIGRATION_PEER_FILE: &str = "migration-peer.json"; -/// ADR 0045 C2 (E2B fold): the staged source-hot-set trace the spawn -/// writes into the jail for the handler's drain ordering + prefault. +/// The final source hot trace, staged by the destination poller in the +/// snapshot directory. The handler reads this advisory file concurrently with drain. pub const MIGRATION_HOT_TRACE_FILE: &str = "migration-hot-trace.json"; /// ADR 0045 C2: `migration-peer.json` content. #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub struct MigrationPeerSpec { + /// Unique to this restore attempt; stale hints need no cleanup. + pub hot_trace_file: String, /// `host:port` of the source's page server (9102). pub peer_addr: String, pub export_id: String, @@ -442,14 +444,6 @@ pub struct MigrationPeerSpec { /// on /proc/*/cmdline). The file lives root-owned in the snapshot /// staging dir for the restore's lifetime. pub peer_token: String, - /// E2B fold: the SOURCE's resume-time working set (the capture - /// rider's `hot_chunks`) — the prediction of exactly the pages - /// the dest guest will fault first. The spawn stages it as a - /// local trace file so the handler drains hot-first AND the - /// post-drain prefault walks it. Serde-default: an old spec - /// simply yields no hot ordering. - #[serde(default)] - pub hot_chunks: Vec<[u8; 32]>, } /// Host-wide knobs for `FirecrackerBackend`. The kernel image lives @@ -2130,32 +2124,16 @@ impl FirecrackerBackend { } cmd.arg("--base-shm").arg(&base); } - // Post-copy: the source's hot set beats any host-local trace — - // it is the working set of THIS guest measured minutes ago, - // not a same-template cousin's. Staged as a local file; the - // handler both orders its drain by it and prefaults from it. - let migration_hot_trace = match peer { - Some(p) if !p.hot_chunks.is_empty() => { - let trace = engram_chunk_store::working_set::WorkingSetTrace { - schema_version: 1, - captured_at: chrono::Utc::now(), - vcpu_count: 0, // unknown here; replay ignores it - capture_window_ms: 0, - chunks: p - .hot_chunks - .iter() - .map(|h| engram_chunk_store::manifest::ChunkHash::from_bytes(*h)) - .collect(), - }; - let path = jail_dir.join(MIGRATION_HOT_TRACE_FILE); - let json = serde_json::to_vec(&trace) - .map_err(|e| vm_err(format!("serialize migration hot trace: {e}")))?; - tokio::fs::write(&path, json) - .await - .map_err(|e| vm_err(format!("write migration hot trace: {e}")))?; - Some(path) - } - _ => None, + // The destination poller atomically stages the capture-time rider in + // the snapshot directory. Hint reads run independently: neither the drain + // nor the fault loop or Firecracker resume waits for this file. + let migration_hot_trace = if let Some(peer) = peer { + let dir = session_manifest_json + .and_then(Path::parent) + .ok_or_else(|| vm_err("post-copy requires a staged session manifest"))?; + Some(dir.join(&peer.hot_trace_file)) + } else { + None }; if let Some(path) = migration_hot_trace.as_ref() { cmd.arg("--prefault-trace") diff --git a/crates/engram-uffd-handler/Cargo.toml b/crates/engram-uffd-handler/Cargo.toml index 6de0fe81c..2aa5e857f 100644 --- a/crates/engram-uffd-handler/Cargo.toml +++ b/crates/engram-uffd-handler/Cargo.toml @@ -72,4 +72,5 @@ userfaultfd = { version = "0.9", features = ["linux5_13"] } [dev-dependencies] proptest.workspace = true +async-trait.workspace = true tempfile = "3" diff --git a/crates/engram-uffd-handler/src/chunked.rs b/crates/engram-uffd-handler/src/chunked.rs index 607c4ee23..ed8514d8f 100644 --- a/crates/engram-uffd-handler/src/chunked.rs +++ b/crates/engram-uffd-handler/src/chunked.rs @@ -154,6 +154,7 @@ impl PositionalManifest { pub struct ChunkedMemoryBackend { canonical: PositionalManifest, session: PositionalManifest, + session_positions: std::collections::HashMap>, /// ADR 0075: the READ-ONLY view of the shared cache dir. The /// handler cannot mutate the directory by construction. reader: ChunkCacheReader, @@ -238,9 +239,19 @@ impl ChunkedMemoryBackend { canonical.total_bytes, session.total_bytes ))); } + let mut session_positions = std::collections::HashMap::>::new(); + for (idx, hash) in session.chunks.iter().enumerate() { + if let Some(hash) = hash { + session_positions + .entry(*hash) + .or_default() + .push(idx as u64 * session.chunk_size); + } + } Ok(Self { canonical, session, + session_positions, reader, populate, store, @@ -504,20 +515,13 @@ impl ChunkedMemoryBackend { /// names a chunk hash, the runtime asks "where does the session /// place this chunk?" and pre-installs it at each position. /// - /// Linear scan over `chunks` because the chunked manifest is - /// typically small (≤8192 entries for a 4 GiB / 512 KiB layout) - /// and the prefault path runs once per restore, before vCPUs - /// unfreeze — not a hot loop. + /// Built once at load, so ordering a large hot set is linear in the + /// number of returned positions instead of trace length times RAM size. pub fn session_positions_of(&self, hash: ChunkHash) -> Vec { - self.session - .chunks - .iter() - .enumerate() - .filter_map(|(idx, entry)| match entry { - Some(h) if *h == hash => Some((idx as u64) * self.session.chunk_size), - _ => None, - }) - .collect() + self.session_positions + .get(&hash) + .cloned() + .unwrap_or_default() } } @@ -554,6 +558,32 @@ mod fallback_tests { r } + #[tokio::test] + async fn hash_index_keeps_all_positions_in_offset_order() { + let dir = tempfile::tempdir().unwrap(); + let blob: Arc = Arc::new(LocalBlobStorage::new(dir.path().to_path_buf())); + let store = ChunkStore::new(blob.clone()); + let mut manifest = mem_manifest(4 * 4096, 4096, 3); + let repeated = manifest.chunks[0].hash; + manifest.chunks[2].hash = repeated; + manifest.chunks.remove(1); // sparse zero chunk has no hash + let r = put(&store, &manifest).await; + let (backend, _) = ChunkedMemoryBackend::build_with_canonical_fallback( + r, + r, + None, + blob, + &dir.path().join("cache"), + None, + ) + .await + .unwrap(); + assert_eq!(backend.session_positions_of(repeated), vec![0, 8192]); + assert!(backend + .session_positions_of(ChunkHash::of(b"absent")) + .is_empty()); + } + /// #1066: a session whose layout predates a guest resize must fail /// open to canonical == session instead of exiting — the loop that /// stranded parked sessions for hours. diff --git a/crates/engram-uffd-handler/src/main.rs b/crates/engram-uffd-handler/src/main.rs index 6393f058f..66e5f15b3 100644 --- a/crates/engram-uffd-handler/src/main.rs +++ b/crates/engram-uffd-handler/src/main.rs @@ -517,7 +517,73 @@ mod linux { .map_err(|e| format!("build chunked backend: {e}"))?; let backend = Arc::new(backend); - let prefault = if let Some(spec) = args.prefault_trace { + // ADR 0045 C2: peer mode. Ordering is the soundness story: + // 1. bind the control sock (host-agent can subscribe NOW); + // 2. connect the peer session — BLOCKS until the source's + // capture registers the export and pushes the Seal (the + // page server parks pre-capture Hellos); + // 3. only then bind the FC-facing UDS (inside run_listener). + // FC's snapshot load waits on the UDS path, so "FC can touch + // guest memory" structurally implies "seal held" — no fault is + // ever served without a classification. + let peer_wiring = match (args.peer_addr.as_ref(), args.peer_export_id.as_ref()) { + (Some(addr), Some(export_id)) => { + let control = match args.control_sock.as_ref() { + Some(path) => { + let tx = engram_uffd_handler::peer::ControlTx::bind(path) + .map_err(|e| format!("bind --control-sock {}: {e}", path.display()))?; + Some(Arc::new(tx)) + } + None => None, + }; + let token = args.peer_token.clone().expect("validated in parse_args"); + let addr = addr.clone(); + let export_id = export_id.clone(); + let chunk_size = backend.chunk_size(); + let total_bytes = backend.total_bytes(); + // The dial blocks (server parks until capture) — do it off + // the async runtime. + let session = tokio::task::spawn_blocking(move || { + engram_uffd_handler::peer::PeerSession::connect( + addr, + export_id, + token, + chunk_size, + total_bytes, + ) + }) + .await + .map_err(|e| format!("peer connect join: {e}"))? + .map_err(|e| format!("peer connect: {e}"))?; + let session = Arc::new(session); + if let Some(control) = control.as_ref() { + control.report(engram_migrate_proto::HandlerControl::Sealed { + dirty_chunks: session.seal().count_ones(), + total_chunks: session.seal().chunk_count, + at_unix_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0), + }); + } + Some((session, control)) + } + _ => None, + }; + + // Defer file loading to the producer in peer mode. Fault service and + // FC resume must not wait for the source's idle scan. + let migration_trace = if peer_wiring.is_some() { + match args.prefault_trace.as_ref() { + Some(PrefaultTraceSpec::File(path)) => Some(path.clone()), + _ => None, + } + } else { + None + }; + let prefault = if migration_trace.is_some() { + None + } else if let Some(spec) = args.prefault_trace { let loaded = match &spec { PrefaultTraceSpec::File(path) => tokio::fs::read(path) .await @@ -586,60 +652,6 @@ mod linux { } } - // ADR 0045 C2: peer mode. Ordering is the soundness story: - // 1. bind the control sock (host-agent can subscribe NOW); - // 2. connect the peer session — BLOCKS until the source's - // capture registers the export and pushes the Seal (the - // page server parks pre-capture Hellos); - // 3. only then bind the FC-facing UDS (inside run_listener). - // FC's snapshot load waits on the UDS path, so "FC can touch - // guest memory" structurally implies "seal held" — no fault is - // ever served without a classification. - let peer_wiring = match (args.peer_addr.as_ref(), args.peer_export_id.as_ref()) { - (Some(addr), Some(export_id)) => { - let control = match args.control_sock.as_ref() { - Some(path) => { - let tx = engram_uffd_handler::peer::ControlTx::bind(path) - .map_err(|e| format!("bind --control-sock {}: {e}", path.display()))?; - Some(Arc::new(tx)) - } - None => None, - }; - let token = args.peer_token.clone().expect("validated in parse_args"); - let addr = addr.clone(); - let export_id = export_id.clone(); - let chunk_size = backend.chunk_size(); - let total_bytes = backend.total_bytes(); - // The dial blocks (server parks until capture) — do it off - // the async runtime. - let session = tokio::task::spawn_blocking(move || { - engram_uffd_handler::peer::PeerSession::connect( - addr, - export_id, - token, - chunk_size, - total_bytes, - ) - }) - .await - .map_err(|e| format!("peer connect join: {e}"))? - .map_err(|e| format!("peer connect: {e}"))?; - let session = Arc::new(session); - if let Some(control) = control.as_ref() { - control.report(engram_migrate_proto::HandlerControl::Sealed { - dirty_chunks: session.seal().count_ones(), - total_chunks: session.seal().chunk_count, - at_unix_ms: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as i64) - .unwrap_or(0), - }); - } - Some((session, control)) - } - _ => None, - }; - // ADR 0045 substrate (v2b): create + size the base shm file NOW — // before run_listener binds the UDS. The host-agent orders FC's // load (which open(O_RDONLY)s + mmaps this file) after the socket @@ -682,6 +694,7 @@ mod linux { engram_uffd_handler::runtime::RunListenerOpts { prefault_trace: prefault, fault_around_chunks: args.fault_around_chunks, + migration_trace, recorder_window: window, trace_output: trace_out, base_shm, diff --git a/crates/engram-uffd-handler/src/peer.rs b/crates/engram-uffd-handler/src/peer.rs index 2f7701900..13dd49ed3 100644 --- a/crates/engram-uffd-handler/src/peer.rs +++ b/crates/engram-uffd-handler/src/peer.rs @@ -165,6 +165,8 @@ pub fn fault_window( /// Drain accounting (rides `ToSource::DrainDone` + the control sock). #[derive(Clone, Copy, Debug, Default)] pub struct DrainStats { + pub hot_hint_received: bool, + pub hot_leading: u64, pub pulled: u64, pub alt_sourced: u64, pub zero_chunks: u64, diff --git a/crates/engram-uffd-handler/src/runtime.rs b/crates/engram-uffd-handler/src/runtime.rs index 2a280b0e7..c7c2e85ad 100644 --- a/crates/engram-uffd-handler/src/runtime.rs +++ b/crates/engram-uffd-handler/src/runtime.rs @@ -1,3 +1,6 @@ +//! Invariant: a hot-set hint must never fail, delay, or block a move, resume, +//! drain, or blackout. Hint errors mean no hint. +//! //! UFFD event loop. Linux-only — `userfaultfd(2)` is a Linux kernel //! syscall and the upstream `userfaultfd` crate fails to compile on //! other targets. @@ -130,6 +133,12 @@ pub struct PrefaultStats { /// Chunks installed by demand windows, including the faulting chunk. #[serde(default)] pub fault_around_chunks_installed: u64, + /// Sealed chunks placed first by the final migration trace. + #[serde(default)] + pub hot_leading: u64, + /// Set after drain completion: a nonempty hint arrived while requests remained. + #[serde(default)] + pub hot_hint_received: Option, } /// Filename for [`PrefaultStats`], colocated with @@ -179,13 +188,22 @@ pub fn write_prefault_stats(path: &std::path::Path, stats: &PrefaultStats) { /// would mask it. fn patch_peer_fields( existing_bytes: Option<&[u8]>, - (peer_pulled, peer_alt_sourced, peer_zero_chunks, peer_live_faults): (u64, u64, u64, u64), + ( + peer_pulled, + peer_alt_sourced, + peer_zero_chunks, + peer_live_faults, + hot_leading, + hot_hint_received, + ): (u64, u64, u64, u64, u64, bool), ) -> Option { let mut stats: PrefaultStats = serde_json::from_slice(existing_bytes?).ok()?; stats.peer_pulled = peer_pulled; stats.peer_alt_sourced = peer_alt_sourced; stats.peer_zero_chunks = peer_zero_chunks; stats.peer_live_faults = peer_live_faults; + stats.hot_leading = hot_leading; + stats.hot_hint_received = Some(hot_hint_received); Some(stats) } @@ -1850,6 +1868,15 @@ impl Runtime { &self, peer: &crate::peer::PeerSession, hot: Option<&engram_chunk_store::working_set::WorkingSetTrace>, + ) -> Result { + self.drain_with_hint(peer, &mut hot.cloned(), &mut || None) + } + + fn drain_with_hint( + &self, + peer: &crate::peer::PeerSession, + hot: &mut Option, + hint: &mut impl FnMut() -> Option, ) -> Result { use crate::peer::{DrainStats, PeerPage}; @@ -1876,7 +1903,7 @@ impl Runtime { // cold remainder keeps offset order. let mut todo: Vec = Vec::with_capacity(seal.count_ones() as usize); let mut queued = std::collections::HashSet::new(); - if let Some(trace) = hot { + if let Some(trace) = hot.as_ref() { for hash in &trace.chunks { for offset in self.backend.session_positions_of(*hash) { let idx = offset / chunk_size; @@ -1887,6 +1914,8 @@ impl Runtime { } } let hot_leading = todo.len(); + stats.hot_leading = hot_leading as u64; + stats.hot_hint_received = hot.as_ref().is_some_and(|trace| !trace.chunks.is_empty()); for i in (0..seal.chunk_count).filter(|i| seal.get(*i)) { if !queued.contains(&i) { todo.push(i * chunk_size); @@ -1898,7 +1927,7 @@ impl Runtime { let total_sealed = todo.len(); let mut in_flight: std::collections::VecDeque<(u64, u64)> = Default::default(); // (req_id, offset) - let mut iter = todo.into_iter(); + let mut todo: std::collections::VecDeque<_> = todo.into(); // Issue #227 (b): offsets the source answered with a per-request // error (`Error{req_id: Some}`, e.g. a transient process_vm_readv // EAGAIN/ENOMEM). Re-queued for a bounded re-request through the @@ -1909,6 +1938,28 @@ impl Runtime { let mut attempts: std::collections::HashMap = std::collections::HashMap::new(); let mut processed = 0usize; loop { + if let Some(trace) = hint() { + stats.hot_hint_received |= !trace.chunks.is_empty() && !todo.is_empty(); + let mut rank = std::collections::HashMap::new(); + for hash in &trace.chunks { + for offset in self.backend.session_positions_of(*hash) { + let next = rank.len(); + rank.entry(offset).or_insert(next); + } + } + stats.hot_leading += todo + .iter() + .filter(|o| { + rank.contains_key(o) && !self.chunk_installed((**o / chunk_size) as usize) + }) + .count() as u64; + todo.make_contiguous() + .sort_by_key(|o| rank.get(o).copied().unwrap_or(usize::MAX)); + // The source already merged handler and idle hints in priority + // order. Keep all hashes, including unsealed chunks, for the + // post-drain prefault. Never wait for another hint. + *hot = Some(trace); + } // Fill the pipeline — but YIELD to guest faults. A fault // request queuing behind the drain's in-flight bulk bytes // was the measured ~7 ms/fault (vs ~1 ms on a quiet wire); @@ -1925,7 +1976,7 @@ impl Runtime { // the fresh todo list. let offset = match retry.pop_front() { Some(o) => o, - None => match iter.next() { + None => match todo.pop_front() { Some(o) => o, None => break, }, @@ -2091,6 +2142,8 @@ impl Runtime { pub struct RunListenerOpts { pub prefault_trace: Option, pub fault_around_chunks: usize, + /// Advisory capture-time trace, delivered concurrently with the drain. + pub migration_trace: Option, pub recorder_window: Duration, pub trace_output: Option, pub base_shm: Option, @@ -2109,8 +2162,9 @@ pub fn run_listener( opts: RunListenerOpts, ) -> Result { let RunListenerOpts { - prefault_trace, + mut prefault_trace, fault_around_chunks, + migration_trace, recorder_window, trace_output, base_shm, @@ -2159,12 +2213,19 @@ pub fn run_listener( std::thread::Builder::new() .name("engram-uffd-prefault".to_string()) .spawn(move || { + let (hint_tx, mut hint_rx) = tokio::sync::oneshot::channel(); + let hint_task = migration_trace.map(|path| rt.handle.spawn(async move { + match crate::working_set::load_migration_trace(&path).await { + Ok(trace) => { let _ = hint_tx.send(trace); } + Err(error) => tracing::info!(%error, "migration hint unavailable; drain continues"), + } + })); // Review finding 6: (pulled, alt_sourced, zero_chunks, live // faults) from the peer drain, captured below if this // restore is peer mode — `None` otherwise (base/resume // restores never set `rt.peer`). Patched onto // `prefault-stats.json` at the end of this closure. - let mut peer_drain_stats: Option<(u64, u64, u64, u64)> = None; + let mut peer_drain_stats: Option<(u64, u64, u64, u64, u64, bool)> = None; // ADR 0045 C2: in peer mode the sealed drain runs FIRST — // sealed content exists only in the paused source, so // draining it is what releases the source, and it makes @@ -2175,7 +2236,9 @@ pub fn run_listener( // whole VM; warming a doomed guest is wasted I/O. if let Some(peer) = rt.peer.clone() { let started = std::time::Instant::now(); - match rt.drain_from_peer(&peer, prefault_trace.as_ref()) { + let outcome = rt.drain_with_hint(&peer, &mut prefault_trace, &mut || hint_rx.try_recv().ok()); + if let Some(task) = hint_task { task.abort(); } + match outcome { Ok(stats) => { // Issue #227 (a): latch drain-complete BEFORE // reporting DrainDone. The host-agent acts on @@ -2193,8 +2256,14 @@ pub fn run_listener( // `prefault-stats.json` once every other writer // in this closure has had its turn (see the // read-modify-write patch below `sweep_all`). - peer_drain_stats = - Some((stats.pulled, stats.alt_sourced, stats.zero_chunks, faults)); + peer_drain_stats = Some(( + stats.pulled, + stats.alt_sourced, + stats.zero_chunks, + faults, + stats.hot_leading, + stats.hot_hint_received, + )); tracing::info!( pulled = stats.pulled, alt_sourced = stats.alt_sourced, @@ -2328,6 +2397,8 @@ mod tests { peer_zero_chunks: 0, peer_live_faults: 0, fault_around_chunks_installed: 0, + hot_leading: 0, + hot_hint_received: None, }; write_prefault_stats(&stats_path, &stats); @@ -2382,7 +2453,8 @@ mod tests { ..Default::default() }; let bytes = serde_json::to_vec(&base).unwrap(); - let patched = patch_peer_fields(Some(&bytes), (3, 1, 4, 7)).expect("existing file parses"); + let patched = + patch_peer_fields(Some(&bytes), (3, 1, 4, 7, 2, true)).expect("existing file parses"); assert_eq!( patched, PrefaultStats { @@ -2396,6 +2468,8 @@ mod tests { peer_zero_chunks: 4, peer_live_faults: 7, fault_around_chunks_installed: 0, + hot_leading: 2, + hot_hint_received: Some(true), }, ); } @@ -2405,8 +2479,8 @@ mod tests { /// the `stats_missing` alarm finding 1 exists to protect. #[test] fn patch_peer_fields_none_when_no_prior_writer_ran() { - assert!(patch_peer_fields(None, (1, 2, 3, 4)).is_none()); - assert!(patch_peer_fields(Some(b"not json"), (1, 2, 3, 4)).is_none()); + assert!(patch_peer_fields(None, (1, 2, 3, 4, 0, false)).is_none()); + assert!(patch_peer_fields(Some(b"not json"), (1, 2, 3, 4, 0, false)).is_none()); } /// ADR 0043 P1: the prefault now runs on a background thread CONCURRENT @@ -2577,18 +2651,79 @@ mod tests { /// construction) NEVER overwritten by sweep/resolve installs. #[tokio::test(flavor = "multi_thread", worker_threads = 3)] async fn peer_mode_drains_sealed_chunks_and_sweep_never_overwrites_them() { + for hint in [ + "unavailable", + "empty", + "delayed", + "failed", + "installed", + "in-flight", + "unsealed", + "late-unsealed", + ] { + peer_drain_hint_case(hint).await; + } + } + + struct RecordingStorage { + inner: LocalBlobStorage, + fetched: std::sync::Mutex>, + } + + #[async_trait::async_trait] + impl BlobStorage for RecordingStorage { + async fn put_streaming( + &self, + key: &str, + body: engram_core::traits::ByteStream, + ) -> Result { + self.inner.put_streaming(key, body).await + } + + async fn get_streaming( + &self, + key: &str, + ) -> Result { + self.fetched.lock().unwrap().push(key.to_owned()); + self.inner.get_streaming(key).await + } + + async fn head( + &self, + key: &str, + ) -> Result { + self.inner.head(key).await + } + + async fn delete(&self, key: &str) -> Result<(), engram_core::error::BlobError> { + self.inner.delete(key).await + } + + async fn list_prefix( + &self, + prefix: &str, + ) -> Result, engram_core::error::BlobError> { + self.inner.list_prefix(prefix).await + } + } + + async fn peer_drain_hint_case(hint_kind: &str) { use engram_migrate_proto::{ read_frame, write_frame, FromSource, SealBitmap, ToSource, PROTO_VERSION, }; let page_size = 4096u64; let chunk_size = page_size; - let n_chunks = 8usize; + let n_chunks = 20usize; let total = chunk_size * n_chunks as u64; const PEER_BYTE: u8 = 0xEE; // Manifest content: chunk i = byte i+1 (the STALE durable view). let dir = tempfile::tempdir().unwrap(); - let blob: Arc = Arc::new(LocalBlobStorage::new(dir.path().to_path_buf())); + let durable = Arc::new(RecordingStorage { + inner: LocalBlobStorage::new(dir.path().to_path_buf()), + fetched: Default::default(), + }); + let blob: Arc = durable.clone(); let store = ChunkStore::new(blob); let mut entries = Vec::new(); for i in 0..n_chunks { @@ -2628,11 +2763,14 @@ mod tests { // Fake source: seals chunks {1, 3, 5}; chunk 1 = peer bytes, // chunk 3 = ZeroChunk, chunk 5 = AltSource (content matches the // durable manifest, so the dest fetches it itself). + let requested = Arc::new(std::sync::Mutex::new(Vec::new())); + let requests = requested.clone(); let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); std::thread::spawn(move || { for stream in listener.incoming() { let Ok(mut s) = stream else { break }; + let requests = requests.clone(); std::thread::spawn(move || { let Ok(ToSource::Hello { .. }) = read_frame::<_, ToSource>(&mut s) else { return; @@ -2647,7 +2785,7 @@ mod tests { ) .unwrap(); let mut bitmap = SealBitmap::new(chunk_size, n_chunks as u64); - for i in [1u64, 3, 5] { + for i in [1u64, 3, 5].into_iter().chain(8..20) { bitmap.set(i); } write_frame(&mut s, &FromSource::Seal { bitmap }).unwrap(); @@ -2656,6 +2794,7 @@ mod tests { chunk_offset, }) = read_frame::<_, ToSource>(&mut s) { + requests.lock().unwrap().push(chunk_offset / chunk_size); let resp = match chunk_offset / chunk_size { 1 => { let raw = vec![PEER_BYTE; chunk_size as usize]; @@ -2673,7 +2812,7 @@ mod tests { req_id, chunk_offset, }, - 5 => FromSource::AltSource { + 5 | 8..=19 => FromSource::AltSource { req_id, chunk_offset, durable_sha256: [0; 32], @@ -2753,11 +2892,55 @@ mod tests { let rt = Arc::new(rt); // Producer order, exactly as run_listener does it: drain, then - // sweep. Off the tokio workers (block_on inside). + // prefault, then sweep. Off the tokio workers (block_on inside). let rt_drain = Arc::clone(&rt); let sess = Arc::clone(&session); + let kind = hint_kind.to_owned(); + let mut trace = WorkingSetTrace::new(0, 0, chrono::DateTime::UNIX_EPOCH); + trace.chunks = match hint_kind { + "installed" => vec![manifest.chunks[1].hash], + "in-flight" => vec![manifest.chunks[3].hash], + "unsealed" | "late-unsealed" => { + vec![manifest.chunks[7].hash, manifest.chunks[6].hash] + } + _ => vec![manifest.chunks[19].hash, manifest.chunks[1].hash], + }; let stats = std::thread::spawn(move || { - let stats = rt_drain.drain_from_peer(&sess, None).expect("drain"); + let mut delivered = false; + let (tx, mut rx) = tokio::sync::oneshot::channel::(); + let _pending = if kind == "failed" { + drop(tx); + None + } else { + Some(tx) + }; + let mut hot = None; + let stats = rt_drain + .drain_with_hint(&sess, &mut hot, &mut || { + // An unresolved or failed delivery must not gate the first request. + if kind == "unavailable" || kind == "failed" { + return rx.try_recv().ok(); + } + let ready = if kind == "late-unsealed" { + rt_drain.sealed_uninstalled_count(&sess) == 0 + } else { + rt_drain.chunk_installed(1) + }; + if delivered || !ready { + return None; + } + delivered = true; + Some(if kind == "empty" { + WorkingSetTrace::new(0, 0, chrono::DateTime::UNIX_EPOCH) + } else { + trace.clone() + }) + }) + .expect("advisory hints cannot fail the drain"); + rt_drain.mark_drain_done(); + if let Some(trace) = hot { + rt_drain.prefault_from_trace(&trace).expect("prefault"); + } rt_drain.sweep_all().expect("sweep"); stats }) @@ -2765,10 +2948,60 @@ mod tests { .unwrap(); assert_eq!( (stats.pulled, stats.zero_chunks, stats.alt_sourced), - (1, 1, 1), + (1, 1, 13), "drain accounting" ); + assert!( + !session.is_lost(), + "hint failure must not mark the peer lost" + ); + let requests = requested.lock().unwrap(); + assert_eq!(requests[0], 1, "drain starts before hint delivery"); + assert_eq!(requests[8], if hint_kind == "delayed" { 19 } else { 13 }); + assert_eq!( + requests.len(), + 15, + "installed and in-flight chunks are not requeued" + ); + assert_eq!(stats.hot_leading, u64::from(hint_kind == "delayed")); + assert_eq!( + stats.hot_hint_received, + matches!( + hint_kind, + "delayed" | "installed" | "in-flight" | "unsealed" + ) + ); + drop(requests); + + // Only unsealed durable reads matter here: sealed AltSource reads + // happen during the drain. The hint must move late unsealed chunks + // ahead of offset zero, preserving the source's merged trace order. + let unsealed = [0, 2, 4, 6, 7]; + let fetched: Vec<_> = durable + .fetched + .lock() + .unwrap() + .iter() + .filter_map(|key| { + unsealed + .iter() + .copied() + .find(|&i| key == &manifest.chunks[i].hash.storage_key()) + }) + .collect(); + assert_eq!( + fetched, + if matches!(hint_kind, "unsealed" | "late-unsealed") { + // The uncached sweep fetches hot chunks again; installation + // remains idempotent. Their first fetch must precede the sweep. + vec![7, 6, 0, 2, 4, 6, 7] + } else { + vec![0, 2, 4, 6, 7] + }, + "unsealed durable fetch order ({hint_kind})" + ); + // A fault on an already-drained sealed chunk takes the wake // path (no peer round-trip, no error). let base = ptr as u64; diff --git a/crates/engram-uffd-handler/src/working_set.rs b/crates/engram-uffd-handler/src/working_set.rs index d51ee5aba..27eae1db6 100644 --- a/crates/engram-uffd-handler/src/working_set.rs +++ b/crates/engram-uffd-handler/src/working_set.rs @@ -113,6 +113,29 @@ impl WorkingSetRecorder { } } +/// Capture-time migration traces arrive after the peer seal. The handler must +/// never wait for them in the drain. Call only from an independent hint task. +pub async fn load_migration_trace(path: &std::path::Path) -> std::io::Result { + tokio::time::timeout(std::time::Duration::from_secs(120), async { + loop { + match tokio::fs::read(path).await { + Ok(bytes) => return serde_json::from_slice(&bytes).map_err(std::io::Error::other), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + tokio::time::sleep(std::time::Duration::from_millis(2)).await; + } + Err(error) => return Err(error), + } + } + }) + .await + .map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::TimedOut, + "capture-time trace did not arrive", + ) + })? +} + #[cfg(test)] mod tests { use super::*; @@ -121,6 +144,26 @@ mod tests { ChunkHash::of(&[byte]) } + #[tokio::test] + async fn migration_trace_reports_read_and_parse_failures_to_hint_task_only() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("hint"); + tokio::fs::write(&path, b"invalid").await.unwrap(); + assert!(load_migration_trace(&path).await.is_err()); + assert!(load_migration_trace(dir.path()).await.is_err()); + tokio::fs::remove_file(&path).await.unwrap(); + let pending = load_migration_trace(&path); + tokio::pin!(pending); + assert!(tokio::time::timeout(Duration::from_millis(5), &mut pending) + .await + .is_err()); + let trace = WorkingSetTrace::new(0, 0, chrono::DateTime::UNIX_EPOCH); + tokio::fs::write(&path, serde_json::to_vec(&trace).unwrap()) + .await + .unwrap(); + assert!(pending.await.unwrap().chunks.is_empty()); + } + #[test] fn open_window_records_in_order() { let mut r = WorkingSetRecorder::new(2, Duration::from_secs(5)); diff --git a/docs/adr/0068-capability-vector-readiness.md b/docs/adr/0068-capability-vector-readiness.md index 19a549763..27f85afba 100644 --- a/docs/adr/0068-capability-vector-readiness.md +++ b/docs/adr/0068-capability-vector-readiness.md @@ -104,6 +104,11 @@ The dead-host detector's own host-level `Ping` probe (added in `7fcc4c3c`) is un beyond a paired `engram_dead_host_probe_rescues_total` counter so both rescue paths are graphable together. +The informational `page_idle` probe reports read/write access to the kernel idle +bitmap on every Firecracker host, including File-mode hosts. Failure appears in +the fleet capability list but never gates placement. Post-copy uses this feature +for recency hints and falls back to the handler trace when it is unavailable. + ## What was deferred / deviations from the issue's literal plan - The issue sketched three PRs (host probes / coordinator gate+surface / probe-before-