fix(relay): retain mesh receives across housekeeping ticks

Add a deterministic production-consumer regression, cancellation-safety documentation, and deadline diagnostics for mesh receive handling.

Co-authored-by: Codex <codex@openai.com>
Signed-off-by: Brad Seiler <seiler@squareup.com>
(cherry picked from commit f1674a2141)
Signed-off-by: Brad Seiler <seiler@squareup.com>
Signed-off-by: Local Builder <1c476cb1127ea270cb900bdc3891060140bdd1173f91498030ee9b4db01c2f20@buzz.block.builderlab.xyz>

crates/buzz-relay/src/api/mesh_demo.rs

@@ -336,8 +336,9 @@ mod tests {
             },
         )
         .await;
-        assert_eq!(resp.status(), StatusCode::OK);
+        let status = resp.status();
         let body = body_json(resp).await;
+        assert_eq!(status, StatusCode::OK, "forwarded response: {body}");
         assert_eq!(body["outcome"], "forwarded");
         assert_eq!(body["echoed_payload"], "mesh echo evidence");
         owner_task.abort();

crates/buzz-relay/src/mesh_boot.rs

@@ -421,24 +421,39 @@ pub(crate) async fn run_demo_echo(
     tracing::info!(%session_id, %peer, "mesh demo echo: session open");
     let mut drain_tick = tokio::time::interval(std::time::Duration::from_millis(100));
     loop {
-        let frame = tokio::select! {
-            _ = drain_tick.tick() => {
-                if shutting_down.load(Ordering::Relaxed) {
-                    if let Some(community_id) = stream.community_id() {
-                        if let Err(e) = stream.send_goodbye(community_id, GoodbyeReason::Draining).await {
-                            tracing::warn!(%session_id, "mesh demo echo: draining goodbye failed: {e}");
-                        } else {
-                            tracing::info!(%session_id, "mesh demo echo: sent draining goodbye");
+        let frame = {
+            // A receive can consume bytes before waiting on Redis validation.
+            // Keep it alive across housekeeping ticks; only shutdown may cancel it.
+            let recv = stream.recv_validated(&directory);
+            tokio::pin!(recv);
+            loop {
+                tokio::select! {
+                    _ = drain_tick.tick() => {
+                        if shutting_down.load(Ordering::Relaxed) {
+                            break None;
                         }
-                    } else {
-                        let _ = stream.finish();
-                        tracing::info!(%session_id, "mesh demo echo: drain before community latch — closing");
                     }
-                    return;
+                    frame = &mut recv => break Some(frame),
                 }
-                continue;
             }
-            frame = stream.recv_validated(&directory) => frame,
+        };
+        // The pending receive and its mutable stream borrow have ended before
+        // the drain response or next echo uses the stream.
+        let Some(frame) = frame else {
+            if let Some(community_id) = stream.community_id() {
+                if let Err(e) = stream
+                    .send_goodbye(community_id, GoodbyeReason::Draining)
+                    .await
+                {
+                    tracing::warn!(%session_id, "mesh demo echo: draining goodbye failed: {e}");
+                } else {
+                    tracing::info!(%session_id, "mesh demo echo: sent draining goodbye");
+                }
+            } else {
+                let _ = stream.finish();
+                tracing::info!(%session_id, "mesh demo echo: drain before community latch — closing");
+            }
+            return;
         };
         match frame {
             Ok(Some(ReliableFrame::Data(payload))) => {
@@ -696,6 +711,72 @@ mod tests {
         MeshStream::new(Box::new(StubSend), Box::new(StubRecv))
     }
 
+    struct PausedRecv {
+        started: Option<tokio::sync::oneshot::Sender<()>>,
+        resume: Option<tokio::sync::oneshot::Receiver<()>>,
+        calls: Arc<std::sync::atomic::AtomicUsize>,
+    }
+
+    impl StreamRecvHalf for PausedRecv {
+        fn recv_frame(&mut self) -> BoxFuture<'_, Result<Option<MeshStreamFrame>, MeshError>> {
+            self.calls.fetch_add(1, Ordering::Relaxed);
+            let started = self.started.take();
+            let resume = self.resume.take();
+            Box::pin(async move {
+                if let Some(started) = started {
+                    let _ = started.send(());
+                }
+                if let Some(resume) = resume {
+                    resume.await.expect("test resumes the in-flight receive");
+                }
+                Ok(None)
+            })
+        }
+    }
+
+    /// Housekeeping must retain a receive that already consumed transport state.
+    /// Returning to recv_frame after a tick loses that state and closes the stream.
+    #[tokio::test(start_paused = true)]
+    async fn demo_echo_retains_pending_receive_across_housekeeping_ticks() {
+        let pool = deadpool_redis::Config::from_url("redis://127.0.0.1:1")
+            .create_pool(Some(deadpool_redis::Runtime::Tokio1))
+            .expect("test pool");
+        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
+        let (resume_tx, resume_rx) = tokio::sync::oneshot::channel();
+        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        let fenced = FencedHeader {
+            session_id: uuid::Uuid::new_v4(),
+            generation: 1,
+            owner_runtime_id: rid(1),
+        };
+        let stream = MeshStream::new(
+            Box::new(StubSend),
+            Box::new(PausedRecv {
+                started: Some(started_tx),
+                resume: Some(resume_rx),
+                calls: Arc::clone(&calls),
+            }),
+        );
+        let worker = tokio::spawn(run_demo_echo(
+            SessionDirectory::new(pool),
+            ReliableInbound {
+                fenced,
+                from: rid(2),
+                stream: crate::tunnel::reliable::ReliableMeshStream::new_inbound(fenced, stream),
+            },
+            Arc::new(AtomicBool::new(false)),
+        ));
+        started_rx
+            .await
+            .expect("production consumer started receiving");
+        tokio::time::advance(std::time::Duration::from_millis(300)).await;
+        tokio::task::yield_now().await;
+        assert_eq!(calls.load(Ordering::Relaxed), 1, "tick restarted a receive");
+        assert!(!worker.is_finished(), "tick abandoned the pending receive");
+        resume_tx.send(()).expect("receive remains alive");
+        worker.await.expect("consumer completed after stream EOF");
+    }
+
     fn rid(byte: u8) -> RuntimeId {
         RuntimeId([byte; 32])
     }

crates/buzz-relay/src/tunnel/reliable.rs

@@ -329,6 +329,9 @@ impl ReliableMeshStream {
     /// pinned fenced tuple and the Redis directory. This is the reliable-stream
     /// equivalent of Dawn's hot-path media floor, but authoritative: stale or
     /// mismatched frames fail the session rather than being dropped silently.
+    /// This receive is not cancellation-safe: it can consume frame bytes before
+    /// awaiting directory validation. Preserve an in-flight receive across
+    /// housekeeping ticks; cancellation is only safe when terminating the stream.
     pub async fn recv_validated(
         &mut self,
         directory: &SessionDirectory,

crates/buzz-relay/tests/boot_lifecycle.rs

@@ -168,20 +168,38 @@ fn wait_for_relay_metrics(process: &mut RelayProcess, port: u16) -> String {
 /// [`METRICS_SCRAPE_DEADLINE`]. A relay that exits first is a failure, not a
 /// timeout, so the panic names the real cause.
 fn wait_for_scraped_metric(process: &mut RelayProcess, port: u16, needle: &str) -> String {
+    wait_for_scraped_metrics(process, port, &[needle])
+}
+
+/// Waits for one scrape to contain every requested sample, bounded by
+/// [`METRICS_SCRAPE_DEADLINE`]. A scrape can overlap the sampler's sequential
+/// metric writes, so related series may appear in different responses.
+fn wait_for_scraped_metrics(process: &mut RelayProcess, port: u16, needles: &[&str]) -> String {
+    let requested = needles.join(", ");
+    let mut last_scrape = String::new();
+    let mut last_error = None;
     let deadline = Instant::now() + METRICS_SCRAPE_DEADLINE;
     loop {
         assert!(
             process.try_wait().is_none(),
-            "relay exited before exporting {needle}"
+            "relay exited before exporting {requested}"
         );
-        if let Ok(response) = scrape_metrics(port) {
-            if response.contains(needle) {
-                return response;
+        match scrape_metrics(port) {
+            Ok(response) => {
+                if needles
+                    .iter()
+                    .all(|needle| response.lines().any(|line| line.starts_with(needle)))
+                {
+                    return response;
+                }
+                last_scrape = response;
+                last_error = None;
             }
+            Err(error) => last_error = Some(error.to_string()),
         }
         assert!(
             Instant::now() < deadline,
-            "relay did not export {needle} within {METRICS_SCRAPE_DEADLINE:?}"
+            "relay did not export {requested} within {METRICS_SCRAPE_DEADLINE:?}; last error: {last_error:?}; last scrape: {last_scrape}"
         );
         thread::sleep(Duration::from_millis(20));
     }
@@ -814,10 +832,10 @@ mod postgres_tests {
     /// fires immediately, so the wait is bounded by
     /// [`METRICS_SCRAPE_DEADLINE`] and never a fixed sleep.
     ///
-    /// The completion-timestamp gauge is the needle because `sample` writes it
-    /// last, after the cache already serves the report: observing it proves the
-    /// whole cycle ran, with no window where the counters have landed but the
-    /// timestamp has not.
+    /// The sampler writes counters and histograms before the completion gauge,
+    /// but a concurrent `/metrics` scrape is not an atomic snapshot. Wait for
+    /// the actual timestamp, counter, and duration samples together so a scrape
+    /// that overlaps those sequential writes cannot report a false failure.
     #[test]
     #[ignore = "requires PostgreSQL"]
     fn startup_owns_the_dependency_sampler() {
@@ -839,10 +857,14 @@ mod postgres_tests {
             ("REDIS_URL", &redis_url),
             ("BUZZ_GIT_CONFORMANCE_PROBE", "false"),
         ]);
-        let scrape = wait_for_scraped_metric(
+        let scrape = wait_for_scraped_metrics(
             &mut process,
             metrics_port,
-            "buzz_readiness_dependency_sample_completed_timestamp_seconds",
+            &[
+                "buzz_readiness_dependency_sample_completed_timestamp_seconds ",
+                "buzz_readiness_dependency_checks_total{",
+                "buzz_readiness_check_duration_seconds_count{check=\"overall\"}",
+            ],
         );
         let output = process.terminate();
         let logs = format!(
@@ -851,14 +873,6 @@ mod postgres_tests {
             String::from_utf8_lossy(&output.stderr)
         );
 
-        assert!(
-            scrape.contains("buzz_readiness_dependency_checks_total{"),
-            "the same cycle must publish its per-dependency outcomes: {scrape}"
-        );
-        assert!(
-            scrape.contains("buzz_readiness_check_duration_seconds"),
-            "a completed evaluation must publish its latency too: {scrape}"
-        );
         assert!(
             !scrape.contains("buzz_readiness_checks_total{"),
             "no readiness probe was sent, so the sampler alone produced this: {scrape}"