fix(sources): [OBE-11555] release the request-limiter permit before the ack write

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<T>(
                                             }
                                         }
                                 };
+                                // 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<T>(
     }
 }
 
+#[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<S>(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<TcpStream>) -> bool {
     debug!("Start graceful shutdown.");
     // Close our write part of TCP socket to signal the other side