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-