fix: prevent pipe_stream stall on multi-hop PUT forwarding (#3568)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: sanity

crates/core/src/transport/peer_connection/outbound_stream.rs

@@ -361,8 +361,51 @@ pub(super) async fn pipe_stream<S: super::super::Socket, T: TimeSource>(
     let mut fragment_number = 1u32;
     let mut pending_metadata = metadata;
 
-    while let Some(result) = stream.next().await {
-        let payload = match result {
+    // Inactivity timeout for piped streams. If no fragment arrives from the
+    // inbound buffer within this duration, the pipe fails rather than hanging
+    // indefinitely. This matches STREAM_INACTIVITY_TIMEOUT used by assemble().
+    use super::streaming::STREAM_INACTIVITY_TIMEOUT;
+    let inactivity_timeout = STREAM_INACTIVITY_TIMEOUT;
+
+    loop {
+        // Use tokio::select! with time_source.sleep() for DST compatibility.
+        // tokio::time::timeout uses real timers which don't advance in
+        // VirtualTime simulation tests.
+        let next_fragment = tokio::select! {
+            result = stream.next() => {
+                match result {
+                    Some(r) => r,
+                    None => break, // Stream complete
+                }
+            }
+            _ = time_source.sleep(inactivity_timeout) => {
+                // No fragment arrived within the inactivity timeout
+                let elapsed = time_source.now().saturating_sub(start_time);
+                tracing::warn!(
+                    stream_id = %outbound_stream_id.0,
+                    destination = %destination_addr,
+                    sent_so_far,
+                    total_bytes,
+                    fragment_number,
+                    elapsed_ms = elapsed.as_millis(),
+                    "pipe_stream stalled: no fragment received within {}s",
+                    inactivity_timeout.as_secs()
+                );
+                emit_transfer_failed(
+                    outbound_stream_id.0 as u64,
+                    destination_addr,
+                    sent_so_far,
+                    format!(
+                        "pipe stalled: no fragment for {}s (sent {sent_so_far}/{total_bytes} bytes)",
+                        inactivity_timeout.as_secs()
+                    ),
+                    elapsed.as_millis() as u64,
+                    TransferDirection::Send,
+                );
+                return Err(TransportError::ConnectionClosed(destination_addr));
+            }
+        };
+        let payload = match next_fragment {
             Ok(data) => data,
             Err(e) => {
                 let elapsed = time_source.now().saturating_sub(start_time);

crates/core/src/transport/peer_connection/streaming.rs

@@ -31,7 +31,9 @@
 
 use bytes::Bytes;
 use dashmap::DashMap;
+use event_listener::EventListener;
 use futures::Stream;
+use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll, Waker};
@@ -198,6 +200,7 @@ impl StreamHandle {
             next_fragment: 1,
             bytes_read: 0,
             auto_reclaim: false,
+            listener: None,
         }
     }
 
@@ -246,15 +249,35 @@ impl StreamHandle {
             next_fragment: 1,
             bytes_read: 0,
             auto_reclaim: true,
+            listener: None,
         }
     }
 
     /// Forks this handle, creating an independent consumer.
     ///
-    /// The forked handle shares the same underlying buffer but maintains
-    /// its own read position when used with `.stream()`.
+    /// The forked handle shares the same underlying buffer but has its own
+    /// independent `SyncState` (cancelled flag and wakers). This means:
+    /// - Cancelling the original handle does NOT cancel the fork
+    /// - Each consumer has its own waker list
+    /// - Fragment notifications reach the fork via the buffer's `Event` notifier
+    ///
+    /// This independence is critical for piped stream forwarding: when the
+    /// upstream `PeerConnection` drops and cancels its handles, the pipe task
+    /// (which reads from a forked handle) must continue forwarding data that's
+    /// already in the shared buffer.
     pub fn fork(&self) -> Self {
-        self.clone()
+        // Preserve the current cancelled state so forking an already-dead
+        // stream doesn't resurrect it. Future cancellations on the original
+        // won't propagate to the fork (independent SyncState).
+        let already_cancelled = self.sync.read().cancelled;
+        let mut sync = SyncState::new();
+        sync.cancelled = already_cancelled;
+        Self {
+            buffer: self.buffer.clone(),
+            sync: Arc::new(parking_lot::RwLock::new(sync)),
+            stream_id: self.stream_id,
+            total_bytes: self.total_bytes,
+        }
     }
 
     /// Inserts a fragment into the stream buffer.
@@ -415,6 +438,11 @@ pub struct StreamingInboundStream {
     /// If true, fragments are taken (removed) from the buffer after reading.
     /// This enables progressive memory reclamation for single-consumer scenarios.
     auto_reclaim: bool,
+    /// Listener for buffer data_available notifications.
+    /// Created when poll_next returns Pending, consumed when notified.
+    /// Uses the buffer's `event_listener::Event` which is fired by every
+    /// `buffer.insert()` — independent of which handle called `push_fragment`.
+    listener: Option<Pin<Box<EventListener>>>,
 }
 
 impl StreamingInboundStream {
@@ -460,7 +488,7 @@ impl Stream for StreamingInboundStream {
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let next_idx = self.next_fragment;
 
-        // Check cancelled state (sync lock)
+        // Check cancelled state (per-handle, independent for forks)
         if self.handle.sync.read().cancelled {
             return Poll::Ready(Some(Err(StreamError::Cancelled)));
         }
@@ -480,31 +508,83 @@ impl Stream for StreamingInboundStream {
             return Poll::Ready(None);
         }
 
-        // Try to get the next fragment (lock-free)
+        // Try to get the next fragment (lock-free fast path)
         if let Some(data) = self.try_get_fragment(next_idx) {
+            self.listener = None;
             self.next_fragment = next_idx + 1;
             self.bytes_read += data.len() as u64;
             return Poll::Ready(Some(Ok(data)));
         }
 
-        // Fragment not yet available, register waker
-        {
-            let mut sync = self.handle.sync.write();
-            // Re-check cancelled after acquiring write lock
-            if sync.cancelled {
-                return Poll::Ready(Some(Err(StreamError::Cancelled)));
-            }
-            // Re-check buffer (fragment may have arrived)
-            if let Some(data) = self.try_get_fragment(next_idx) {
-                drop(sync); // Release lock before modifying self
-                self.next_fragment = next_idx + 1;
-                self.bytes_read += data.len() as u64;
-                return Poll::Ready(Some(Ok(data)));
-            }
-            sync.wakers.push(cx.waker().clone());
+        // Fragment not yet available.
+        //
+        // Following the event_listener pattern: create the listener BEFORE
+        // re-checking the condition. This prevents a TOCTOU race where a
+        // notification fires between the check and listener creation.
+        //
+        // We use the buffer's `data_available` Event (fired by every
+        // buffer.insert()) as the primary notification mechanism. This works
+        // for both original and forked handles since buffer.insert() fires
+        // regardless of which handle called push_fragment().
+        //
+        // cancel() also fires buffer.notifier().notify(), so the listener
+        // covers both data arrival and cancellation.
+        if self.listener.is_none() {
+            self.listener = Some(Box::pin(self.handle.buffer.notifier().listen()));
+        }
+
+        // Re-check after listener creation (the event_listener pattern):
+        // if a fragment arrived between the fast-path check and listen(),
+        // we catch it here instead of missing the notification.
+        if self.handle.sync.read().cancelled {
+            self.listener = None;
+            return Poll::Ready(Some(Err(StreamError::Cancelled)));
+        }
+        if let Some(data) = self.try_get_fragment(next_idx) {
+            self.listener = None;
+            self.next_fragment = next_idx + 1;
+            self.bytes_read += data.len() as u64;
+            return Poll::Ready(Some(Ok(data)));
         }
 
-        Poll::Pending
+        // Poll the listener. The listener is guaranteed to be Some here
+        // (set above), but we use if-let to satisfy the no-unwrap rule.
+        if let Some(listener) = self.listener.as_mut() {
+            match listener.as_mut().poll(cx) {
+                Poll::Ready(()) => {
+                    // Notified — clear listener and re-check buffer inline
+                    // (avoids wake_by_ref spin-loop).
+                    self.listener = None;
+                    if self.handle.sync.read().cancelled {
+                        return Poll::Ready(Some(Err(StreamError::Cancelled)));
+                    }
+                    if let Some(data) = self.try_get_fragment(next_idx) {
+                        self.next_fragment = next_idx + 1;
+                        self.bytes_read += data.len() as u64;
+                        return Poll::Ready(Some(Ok(data)));
+                    }
+                    // Spurious notification. Create a fresh listener,
+                    // re-check, and register waker.
+                    self.listener = Some(Box::pin(self.handle.buffer.notifier().listen()));
+                    if let Some(data) = self.try_get_fragment(next_idx) {
+                        self.listener = None;
+                        self.next_fragment = next_idx + 1;
+                        self.bytes_read += data.len() as u64;
+                        return Poll::Ready(Some(Ok(data)));
+                    }
+                    if let Some(new_listener) = self.listener.as_mut() {
+                        match new_listener.as_mut().poll(cx) {
+                            Poll::Ready(()) => self.listener = None,
+                            Poll::Pending => {}
+                        }
+                    }
+                    Poll::Pending
+                }
+                Poll::Pending => Poll::Pending,
+            }
+        } else {
+            Poll::Pending
+        }
     }
 }
 
@@ -656,6 +736,127 @@ mod tests {
         assert_eq!(handle.try_assemble(), forked.try_assemble());
     }
 
+    #[test]
+    fn test_fork_independent_cancellation() {
+        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;
+
+        // Use exact fragment size so one fragment completes the stream
+        let total = FRAGMENT_PAYLOAD_SIZE as u64;
+        let data = vec![42u8; FRAGMENT_PAYLOAD_SIZE];
+        let handle = StreamHandle::new(make_stream_id(), total);
+        handle.push_fragment(1, Bytes::from(data.clone())).unwrap();
+
+        let forked = handle.fork();
+
+        // Cancel the original — simulates PeerConnection::Drop
+        handle.cancel();
+
+        // Original is cancelled
+        assert!(handle.sync.read().cancelled);
+        assert!(matches!(
+            handle.push_fragment(1, Bytes::from(vec![0u8; 10])),
+            Err(StreamError::Cancelled)
+        ));
+
+        // Fork is NOT cancelled — pipe task can continue reading
+        assert!(!forked.sync.read().cancelled);
+
+        // Fork can still read data from the shared buffer
+        assert!(forked.is_complete());
+        assert_eq!(forked.try_assemble(), Some(data));
+    }
+
+    #[tokio::test]
+    async fn test_fork_stream_reads_after_original_cancel() {
+        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;
+
+        let total = (FRAGMENT_PAYLOAD_SIZE * 2) as u64;
+        let handle = StreamHandle::new(make_stream_id(), total);
+
+        // Push both fragments via the original handle
+        handle
+            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
+            .unwrap();
+        handle
+            .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
+            .unwrap();
+
+        // Fork and create a stream from it
+        let forked = handle.fork();
+        let mut stream = forked.stream();
+
+        // Cancel the original (simulates PeerConnection drop)
+        handle.cancel();
+
+        // The forked stream should still be able to read all fragments
+        let chunk1 = stream.next().await;
+        assert!(chunk1.is_some());
+        assert_eq!(chunk1.unwrap().unwrap().len(), FRAGMENT_PAYLOAD_SIZE);
+
+        let chunk2 = stream.next().await;
+        assert!(chunk2.is_some());
+        assert_eq!(chunk2.unwrap().unwrap().len(), FRAGMENT_PAYLOAD_SIZE);
+
+        // Stream complete
+        let chunk3 = stream.next().await;
+        assert!(chunk3.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_fork_incremental_wakeup() {
+        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;
+
+        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
+        let handle = StreamHandle::new(make_stream_id(), total);
+
+        // Push only fragment #1
+        handle
+            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
+            .unwrap();
+
+        // Fork and start reading
+        let forked = handle.fork();
+        let mut stream = forked.stream();
+
+        // Read fragment #1 (available immediately)
+        let chunk1 = stream.next().await;
+        assert!(chunk1.is_some());
+        assert_eq!(chunk1.unwrap().unwrap(), vec![1u8; FRAGMENT_PAYLOAD_SIZE]);
+
+        // Fragment #2 arrives later via the ORIGINAL handle's push_fragment.
+        // The forked stream must wake up via the buffer's EventListener
+        // (not via SyncState wakers, which are on the fork's independent sync).
+        let handle_clone = handle.clone();
+        tokio::spawn(async move {
+            tokio::task::yield_now().await;
+            handle_clone
+                .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
+                .unwrap();
+        });
+
+        // The forked stream should receive fragment #2 via EventListener
+        let chunk2 = tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()).await;
+        assert!(
+            chunk2.is_ok(),
+            "fork should wake up when fragment arrives via original handle"
+        );
+        let chunk2 = chunk2.unwrap();
+        assert!(chunk2.is_some());
+        assert_eq!(chunk2.unwrap().unwrap(), vec![2u8; FRAGMENT_PAYLOAD_SIZE]);
+
+        // Push fragment #3 and verify
+        handle
+            .push_fragment(3, Bytes::from(vec![3u8; FRAGMENT_PAYLOAD_SIZE]))
+            .unwrap();
+        let chunk3 = stream.next().await;
+        assert!(chunk3.is_some());
+        assert_eq!(chunk3.unwrap().unwrap(), vec![3u8; FRAGMENT_PAYLOAD_SIZE]);
+
+        // Stream complete
+        let end = stream.next().await;
+        assert!(end.is_none());
+    }
+
     #[tokio::test]
     async fn test_streaming_inbound_stream_basic() {
         let handle = StreamHandle::new(make_stream_id(), 15);