From 16a1f56f6ee920fb9a12adbb078f617f12da63ac Mon Sep 17 00:00:00 2001 From: Ben Pfaff Date: Wed, 5 Aug 2026 09:03:11 -0700 Subject: [PATCH] [dbsp] Fix flaky exchange wire-protocol tests under random fault injection new_message_sent_without_waiting_after_a_ping and pings_repeatedly_without_retransmitting drive the wire protocol directly against a single hand-accepted TcpStream, so they can't tolerate the random inject_fault() connection drops meant for the round-based resilience tests: a dropped connection makes the client reconnect to a fresh socket the test never accepts, leaving it reading a dead stream and panicking on an unexpected None. This fixes the problem by disabling fault injection while the tests that can tolerate it run. These tests run quickly, so other tests can still benefit from fault injection. Signed-off-by: Ben Pfaff --- .../src/operator/communication/exchange.rs | 38 ++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs index 083634f273b..cfcf96c258b 100644 --- a/crates/dbsp/src/operator/communication/exchange.rs +++ b/crates/dbsp/src/operator/communication/exchange.rs @@ -2190,10 +2190,18 @@ where Some((sender, receiver)) } +/// Count of tests currently disabling fault injection (see +/// `tests::FaultInjectionDisabled`). +#[cfg(test)] +static FAULT_INJECTION_DISABLE_COUNT: AtomicUsize = AtomicUsize::new(0); + #[cfg(test)] fn inject_fault(kind: impl Display) -> bool { use rand::Rng as _; + if FAULT_INJECTION_DISABLE_COUNT.load(Ordering::Relaxed) > 0 { + return false; + } if rand::thread_rng().gen_range(0..100) == 0 { warn!("injecting failure: {kind}"); true @@ -2261,7 +2269,7 @@ mod tests { use std::{ iter::{repeat, zip}, net::TcpListener, - sync::Arc, + sync::{Arc, atomic::Ordering}, }; /// Number of rounds for exchange. @@ -2466,6 +2474,32 @@ mod tests { test_operators_multihost(operator_circuit::); } + /// Disables the random `inject_fault` failures used by the round-based + /// resilience tests (`single_host`, `multihost`, ...) for the lifetime of + /// the guard. Tests that drive the wire protocol directly through a single + /// hand-accepted `TcpStream` have no way to recover from an injected + /// "connection failure": the client would reconnect to a fresh socket that + /// the test never accepts, leaving the test reading a dead stream. + /// + /// If we used this feature for disabling fault injection a lot, then it + /// would negate the value of fault injection entirely, but we only use it + /// in tests that run very quickly, whereas the tests that do benefit from + /// fault injection are longer running. + struct FaultInjectionDisabled; + + impl FaultInjectionDisabled { + fn new() -> Self { + super::FAULT_INJECTION_DISABLE_COUNT.fetch_add(1, Ordering::Relaxed); + Self + } + } + + impl Drop for FaultInjectionDisabled { + fn drop(&mut self) { + super::FAULT_INJECTION_DISABLE_COUNT.fetch_sub(1, Ordering::Relaxed); + } + } + // Exercises the liveness ping mechanism directly: a fake receiver // acknowledges nothing, so the sender has no way to tell a slow receiver // from a half-open connection except by asking. @@ -2478,6 +2512,7 @@ mod tests { use tokio::{io::AsyncWriteExt, net::TcpListener as AsyncTcpListener, time::timeout}; init_test_logger(); + let _guard = FaultInjectionDisabled::new(); TOKIO.block_on(async { let listener = AsyncTcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -2520,6 +2555,7 @@ mod tests { use tokio::net::TcpListener as AsyncTcpListener; init_test_logger(); + let _guard = FaultInjectionDisabled::new(); TOKIO.block_on(async { let listener = AsyncTcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap();