diff --git a/src/sources/util/net/tcp/mod.rs b/src/sources/util/net/tcp/mod.rs index 13bb464ab..b16b0a8b0 100644 --- a/src/sources/util/net/tcp/mod.rs +++ b/src/sources/util/net/tcp/mod.rs @@ -9,7 +9,7 @@ use listenfd::ListenFd; use smallvec::SmallVec; use socket2::SockRef; use tokio::{ - io::AsyncWriteExt, + io::{AsyncWrite, AsyncWriteExt}, net::{TcpListener, TcpStream}, time::sleep, }; @@ -376,11 +376,27 @@ async fn handle_stream( } } }; + // Release permit before ack write: the permit bounds in-flight + // decoded events, and that purpose is fulfilled once send_batch + // and receiver.await complete. A slow peer that never drains its + // receive window would otherwise block write_all indefinitely + // while holding the permit, starving other connections (OBE-11555). + let _ = permit.take(); if let Some(ack_bytes) = acker.build_ack(ack){ let stream = reader.get_mut().get_mut(); - if let Err(error) = stream.write_all(&ack_bytes).await { - emit!(TcpSendAckError{ error }); - break; + match write_ack(stream, &ack_bytes, ACK_WRITE_TIMEOUT).await { + AckWriteOutcome::Written => {} + AckWriteOutcome::Failed(error) => { + emit!(TcpSendAckError{ error }); + break; + } + AckWriteOutcome::TimedOut => { + warn!( + timeout_secs = ACK_WRITE_TIMEOUT.as_secs(), + "Ack write timed out; dropping connection." + ); + break; + } } } if ack != TcpSourceAck::Ack { @@ -412,6 +428,145 @@ async fn handle_stream( } } +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::AsyncReadExt; + + #[tokio::test] + async fn write_ack_succeeds_when_peer_reads() { + let (mut client, mut server) = tokio::io::duplex(64); + + let write = tokio::spawn(async move { + write_ack(&mut server, b"ack", ACK_WRITE_TIMEOUT).await + }); + + let mut buf = [0u8; 3]; + client.read_exact(&mut buf).await.unwrap(); + assert_eq!(&buf, b"ack"); + assert!(matches!(write.await.unwrap(), AckWriteOutcome::Written)); + } + + /// A peer that never drains its receive window must not block the ack write forever. Before + /// the timeout existed, this write parked indefinitely. + #[tokio::test(start_paused = true)] + async fn write_ack_times_out_against_a_peer_that_never_reads() { + // A 1-byte duplex fills immediately and `client` is never read from, so the write stalls. + let (_client, mut server) = tokio::io::duplex(1); + let payload = vec![0u8; 1024]; + + let outcome = write_ack(&mut server, &payload, ACK_WRITE_TIMEOUT).await; + assert!( + matches!(outcome, AckWriteOutcome::TimedOut), + "expected TimedOut, got {outcome:?}" + ); + } + + /// The timeout must not fire early for a peer that is merely slow rather than stuck. + #[tokio::test(start_paused = true)] + async fn write_ack_tolerates_a_slow_but_progressing_peer() { + let (mut client, mut server) = tokio::io::duplex(4); + let payload = vec![7u8; 32]; + let expected = payload.clone(); + + let write = + tokio::spawn(async move { write_ack(&mut server, &payload, ACK_WRITE_TIMEOUT).await }); + + let mut received = Vec::new(); + while received.len() < expected.len() { + // Drain in small sips, pausing well inside the timeout each round. + tokio::time::sleep(Duration::from_secs(1)).await; + let mut chunk = [0u8; 4]; + let n = client.read(&mut chunk).await.unwrap(); + received.extend_from_slice(&chunk[..n]); + } + + assert_eq!(received, expected); + assert!(matches!(write.await.unwrap(), AckWriteOutcome::Written)); + } + + #[tokio::test] + async fn write_ack_reports_failure_when_peer_hangs_up() { + let (client, mut server) = tokio::io::duplex(64); + drop(client); + + let outcome = write_ack(&mut server, &vec![0u8; 4096], ACK_WRITE_TIMEOUT).await; + assert!( + matches!(outcome, AckWriteOutcome::Failed(_)), + "expected Failed, got {outcome:?}" + ); + } + + /// The permit must be released before the ack write, so a stuck peer cannot hold a + /// `RequestLimiter` slot and starve other connections (OBE-11555). This models the ordering + /// `handle_stream` uses: take the permit, then perform the (stalling) write. + #[tokio::test(start_paused = true)] + async fn permit_is_released_before_a_stalled_ack_write() { + let limiter = RequestLimiter::new(1, 1); + // The limiter starts at its floor of 2 permits; hold every one so the next acquire blocks. + let held = limiter.acquire().await; + let mut permit = Some(limiter.acquire().await); + assert!( + tokio::time::timeout(Duration::from_millis(50), limiter.acquire()) + .await + .is_err(), + "all permits are held, so a further acquire must block" + ); + + // Ordering under test: release, then write to a peer that never reads. + let _ = permit.take(); + + let (_client, mut server) = tokio::io::duplex(1); + let write = tokio::spawn(async move { + write_ack(&mut server, &vec![0u8; 1024], ACK_WRITE_TIMEOUT).await + }); + + // While the write is stalled, another connection must still get a permit. + let second = tokio::time::timeout(Duration::from_secs(1), limiter.acquire()).await; + assert!( + second.is_ok(), + "permit must be available while the ack write is stalled" + ); + + assert!(matches!(write.await.unwrap(), AckWriteOutcome::TimedOut)); + drop(held); + } + + #[test] + fn ack_write_timeout_is_thirty_seconds() { + assert_eq!(ACK_WRITE_TIMEOUT, Duration::from_secs(30)); + } +} + +/// How long to wait for an ack to reach the peer before giving up on the connection. +const ACK_WRITE_TIMEOUT: Duration = Duration::from_secs(30); + +/// Result of attempting to write an ack back to the peer. +#[derive(Debug)] +enum AckWriteOutcome { + Written, + /// The write failed; the connection should be torn down. + Failed(std::io::Error), + /// The peer never drained its receive window within the timeout. + TimedOut, +} + +/// Writes `ack_bytes` to `stream`, bounded by `timeout`. +/// +/// Without the timeout a peer that stops reading parks this write forever. That matters because +/// the caller has already released its `RequestLimiterPermit` by this point (OBE-11555) — the +/// connection itself still needs to be reclaimed. +async fn write_ack(stream: &mut S, ack_bytes: &[u8], timeout: Duration) -> AckWriteOutcome +where + S: AsyncWrite + Unpin + ?Sized, +{ + match tokio::time::timeout(timeout, stream.write_all(ack_bytes)).await { + Ok(Ok(())) => AckWriteOutcome::Written, + Ok(Err(error)) => AckWriteOutcome::Failed(error), + Err(_elapsed) => AckWriteOutcome::TimedOut, + } +} + fn close_socket(socket: &MaybeTlsIncomingStream) -> bool { debug!("Start graceful shutdown."); // Close our write part of TCP socket to signal the other side