uffd-handler: the first peer dial waits for the source's capture

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()
             },
         );

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<String, Arc<PeerExport>>,
     /// 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> {
-        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.

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;
 

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<bool, PeerError>,
+) -> Result<TcpStream, PeerError> {
+    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,23 +385,47 @@ fn dial_with_socket(
     purpose: engram_migrate_proto::ConnPurpose,
     connected: impl Fn(Option<&TcpStream>) -> Result<bool, PeerError>,
 ) -> 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<SealBitmap, PeerError> {
+    write_frame(
+        stream,
         &ToSource::Hello {
             version: PROTO_VERSION,
             export_id: export_id.to_string(),
             token: token.to_string(),
             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,27 +463,30 @@ 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,
         token: String,
         expect_chunk_size: u64,
         expect_total_bytes: u64,
     ) -> Result<Self, PeerError> {
-        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);