From 00cc26519ae03ec8653e6a7f07ba772419db6409 Mon Sep 17 00:00:00 2001 From: Nikhil Unni Date: Tue, 10 Nov 2026 11:15:42 -0800 Subject: [PATCH 1/2] uffd-handler: the first peer dial waits for the source's capture The destination's handler dials the source before the source pauses the guest. The source holds that Hello until the capture registers the export, and before it pauses it waits for the destination to finish its disk precopy. The handler read HelloAck with the 5 s page-channel timeout, so a capture later than 5 s after the dial ended the handler with EAGAIN and the restore failed with postcopy-never-loaded. The Firecracker root lane showed this as an intermittent failure of the two-host live teleport tests. The first dial now waits for HelloAck for the source's park budget plus one page timeout, so the source's own answer arrives first when the export never registers. The page timeout applies again from Seal on. Redials and worker connections keep the 5 s limit: their export exists. The park budget is one constant in engram-migrate-proto, read by the server and the client. Review pass: handshake() is split from the dial so the wait is an argument of one function, not of each dial helper. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_01HNhpmz4QPYaVCE7tFna54D --- crates/engram-host-agent/src/migrate_peer.rs | 4 +- crates/engram-migrate-proto/src/lib.rs | 5 + crates/engram-uffd-handler/src/peer.rs | 117 +++++++++++++++++-- 3 files changed, 112 insertions(+), 14 deletions(-) diff --git a/crates/engram-host-agent/src/migrate_peer.rs b/crates/engram-host-agent/src/migrate_peer.rs index 5997bf5ae..54aa54a95 100644 --- a/crates/engram-host-agent/src/migrate_peer.rs +++ b/crates/engram-host-agent/src/migrate_peer.rs @@ -211,7 +211,7 @@ pub struct PeerServer { /// coordinator, which pairs it with the source's host address). port: u16, /// How long an unknown-export `Hello` parks awaiting the capture's - /// registration (export-TTL scale in production; tests shrink it). + /// registration (`EXPORT_PARK_BUDGET` in production; tests shrink it). park_budget: std::time::Duration, exports: DashMap>, /// Notified on every [`PeerServer::register`] so parked Hellos wake @@ -225,7 +225,7 @@ pub struct PeerServer { impl PeerServer { pub fn new(port: u16) -> Arc { - Self::new_with_park_budget(port, std::time::Duration::from_secs(120)) + Self::new_with_park_budget(port, engram_migrate_proto::EXPORT_PARK_BUDGET) } /// Test seam: shrink the unknown-export park window. diff --git a/crates/engram-migrate-proto/src/lib.rs b/crates/engram-migrate-proto/src/lib.rs index 8899dcee2..5485fdc56 100644 --- a/crates/engram-migrate-proto/src/lib.rs +++ b/crates/engram-migrate-proto/src/lib.rs @@ -70,6 +70,11 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; /// v6: ordered fault windows and fault-around installation accounting. pub const PROTO_VERSION: u32 = 6; +/// How long the source holds a `Hello` for an export that is not there yet. +/// The destination dials before the source pauses the guest, and the capture +/// registers the export. When the budget ends the source answers `Error`. +pub const EXPORT_PARK_BUDGET: std::time::Duration = std::time::Duration::from_secs(120); + /// Bound a demand window to limit request and response work. pub const MAX_FAULT_AROUND_CHUNKS: usize = 256; diff --git a/crates/engram-uffd-handler/src/peer.rs b/crates/engram-uffd-handler/src/peer.rs index 839461278..85ebf265d 100644 --- a/crates/engram-uffd-handler/src/peer.rs +++ b/crates/engram-uffd-handler/src/peer.rs @@ -55,6 +55,13 @@ const RECONNECT_BACKOFF: Duration = Duration::from_millis(500); pub(crate) const PEER_READ_TIMEOUT: Duration = Duration::from_secs(5); const PEER_WRITE_TIMEOUT: Duration = Duration::from_secs(5); +/// How long the first dial waits for `HelloAck`. The handler dials before +/// the source pauses the guest, and the source holds the `Hello` until the +/// capture registers the export. The source answers with an `Error` when +/// its own budget ends, so this wait is longer than that budget. +const CAPTURE_WAIT: Duration = + engram_migrate_proto::EXPORT_PARK_BUDGET.saturating_add(PEER_READ_TIMEOUT); + /// First response from a new worker: each attempt can connect, send Hello, /// read HelloAck and Seal, send NeedWindow, then read the first page. /// Include the initial attempt and every redial, plus reconnect backoff. @@ -358,6 +365,17 @@ fn connect_with_socket( Err(last_error.into()) } +/// A connected socket with the page-channel options set. +fn connect_with_sockopts( + addr: &str, + connected: &impl Fn(Option<&TcpStream>) -> Result, +) -> Result { + let stream = connect_with_socket(addr, connected)?; + stream.set_nodelay(true)?; + apply_peer_sockopts(&stream)?; + Ok(stream) +} + fn dial_with_socket( addr: &str, export_id: &str, @@ -367,15 +385,36 @@ fn dial_with_socket( purpose: engram_migrate_proto::ConnPurpose, connected: impl Fn(Option<&TcpStream>) -> Result, ) -> Result<(TcpStream, SealBitmap), PeerError> { - let mut stream = connect_with_socket(addr, &connected)?; - stream.set_nodelay(true)?; - apply_peer_sockopts(&stream)?; + let mut stream = connect_with_sockopts(addr, &connected)?; // Register worker sockets before any blocking handshake read or write. if !connected(Some(&stream))? { return Err(PeerError::Lost("connection canceled".into())); } - write_frame( + let seal = handshake( &mut stream, + export_id, + token, + expect_chunk_size, + expect_total_bytes, + purpose, + PEER_READ_TIMEOUT, + )?; + Ok((stream, seal)) +} + +/// `Hello`, `HelloAck`, `Seal`. `ack_wait` bounds the read of `HelloAck` +/// only; the source sends `Seal` directly after it. +fn handshake( + stream: &mut TcpStream, + export_id: &str, + token: &str, + expect_chunk_size: u64, + expect_total_bytes: u64, + purpose: engram_migrate_proto::ConnPurpose, + ack_wait: Duration, +) -> Result { + write_frame( + stream, &ToSource::Hello { version: PROTO_VERSION, export_id: export_id.to_string(), @@ -383,7 +422,10 @@ fn dial_with_socket( purpose, }, )?; - match read_frame::<_, FromSource>(&mut stream)? { + stream.set_read_timeout(Some(ack_wait))?; + let ack = read_frame::<_, FromSource>(stream); + stream.set_read_timeout(Some(PEER_READ_TIMEOUT))?; + match ack? { FromSource::HelloAck { version, chunk_size, @@ -408,7 +450,7 @@ fn dial_with_socket( ))); } } - let seal = match read_frame::<_, FromSource>(&mut stream)? { + let seal = match read_frame::<_, FromSource>(stream)? { FromSource::Seal { bitmap } => bitmap, FromSource::Error { message, .. } => return Err(PeerError::Server(message)), other => return Err(PeerError::Server(format!("expected Seal, got {other:?}"))), @@ -421,13 +463,14 @@ fn dial_with_socket( "seal bitmap geometry != manifest geometry".into(), )); } - Ok((stream, seal)) + Ok(seal) } impl PeerSession { - /// Connect, authenticate, and BLOCK until the source's `Seal` - /// arrives. Geometry is validated against the manifest's view - /// before any fault can be served. + /// Connect, authenticate, and block until the source's `Seal` + /// arrives, which is after its capture ([`CAPTURE_WAIT`]). Geometry + /// is validated against the manifest's view before any fault can be + /// served. pub fn connect( addr: String, export_id: String, @@ -435,13 +478,15 @@ impl PeerSession { expect_chunk_size: u64, expect_total_bytes: u64, ) -> Result { - let (stream, seal) = dial( - &addr, + let mut stream = connect_with_sockopts(&addr, &|_| Ok(true))?; + let seal = handshake( + &mut stream, &export_id, &token, expect_chunk_size, expect_total_bytes, engram_migrate_proto::ConnPurpose::Fault, + CAPTURE_WAIT, )?; tracing::info!( addr, @@ -1426,6 +1471,54 @@ pub(crate) mod tests { } } + /// A source that sends `HelloAck` and `Seal` `delay` after `Hello`. + fn slow_ack_source(delay: Duration) -> std::net::SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + std::thread::spawn(move || { + for mut s in listener.incoming().flatten() { + std::thread::spawn(move || { + let _ = read_frame::<_, ToSource>(&mut s); + std::thread::sleep(delay); + let _ = write_frame( + &mut s, + &FromSource::HelloAck { + version: PROTO_VERSION, + chunk_size: CHUNK, + total_bytes: TOTAL, + }, + ); + let bitmap = SealBitmap::new(CHUNK, TOTAL / CHUNK); + let _ = write_frame(&mut s, &FromSource::Seal { bitmap }); + // Hold the connection open for the client's checks. + let _ = read_frame::<_, ToSource>(&mut s); + }); + } + }); + addr + } + + #[test] + fn the_ack_wait_bounds_only_the_wait_for_hello_ack() { + let addr = slow_ack_source(Duration::from_millis(300)).to_string(); + let shake = |ack_wait| { + let mut stream = TcpStream::connect(&addr).unwrap(); + let purpose = engram_migrate_proto::ConnPurpose::Fault; + handshake(&mut stream, "e", "tok", CHUNK, TOTAL, purpose, ack_wait).map(|_| stream) + }; + let err = shake(Duration::from_millis(50)).expect_err("the ack is later than the wait"); + assert!(matches!(err, PeerError::Io(_)), "{err}"); + let stream = shake(Duration::from_secs(10)).expect("the ack is inside the wait"); + assert_eq!(stream.read_timeout().unwrap(), Some(PEER_READ_TIMEOUT)); + } + + /// The source answers a `Hello` whose export never registers when its + /// budget ends; the first dial must still be reading then. + #[test] + fn the_first_dial_outlasts_the_source_park_budget() { + assert!(CAPTURE_WAIT > engram_migrate_proto::EXPORT_PARK_BUDGET); + } + #[test] fn connect_validates_geometry_and_holds_seal() { let addr = fake_source(vec![0, 2], page_resp); From f0fc84f06cd19e82e24418b40a9d0fb53728fc16 Mon Sep 17 00:00:00 2001 From: Nikhil Unni Date: Tue, 10 Nov 2026 11:39:40 -0800 Subject: [PATCH 2/2] blob-client: the deadline test gives its retry one second The test set a 50 ms attempt deadline to end a hung first attempt. The same deadline applied to the retry, which reads a local file. On a loaded macOS runner that read took longer than 50 ms and the test failed with 'attempt deadline 0.05s exceeded'. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_01HNhpmz4QPYaVCE7tFna54D --- crates/engram-blob-client/src/lib.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/engram-blob-client/src/lib.rs b/crates/engram-blob-client/src/lib.rs index bcb3a4f68..70b8fa682 100644 --- a/crates/engram-blob-client/src/lib.rs +++ b/crates/engram-blob-client/src/lib.rs @@ -831,7 +831,9 @@ mod tests { let client = BlobClient::with_policy( hanging, ClientPolicy { - attempt_timeout: Duration::from_millis(50), + // The retry's read of a local file must finish inside this + // deadline on a loaded runner. + attempt_timeout: Duration::from_secs(1), ..test_policy() }, );