diff --git a/lib/runtime/src/pipeline/network/tcp/client.rs b/lib/runtime/src/pipeline/network/tcp/client.rs index 55b3a8b52e0c..5136a57851a7 100644 --- a/lib/runtime/src/pipeline/network/tcp/client.rs +++ b/lib/runtime/src/pipeline/network/tcp/client.rs @@ -460,7 +460,14 @@ async fn handle_request_reader( } } TwoPartMessageType::DataOnly(data) => { - if bytes_tx.send(data).await.is_err() { + // Local cancellation must also interrupt a full input channel. + let sent = tokio::select! { + biased; + _ = &mut killed => break, + _ = &mut stopped => break, + sent = bytes_tx.send(data) => sent, + }; + if sent.is_err() { tracing::debug!("downstream consumer dropped; exiting request-stream reader"); break; } @@ -1816,6 +1823,175 @@ mod tests { } } + // The read notification is sent only once a whole encoded frame has reached + // FramedRead. On the current-thread test runtime, the reader then polls the + // full channel send before the waiting test task can resume. + struct FrameReadProbe { + inner: BoxRead, + remaining: usize, + read: Option>, + } + + impl tokio::io::AsyncRead for FrameReadProbe { + fn poll_read( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> std::task::Poll> { + let before = buf.filled().len(); + let result = std::pin::Pin::new(&mut self.inner).poll_read(cx, buf); + if matches!(result, std::task::Poll::Ready(Ok(()))) { + self.remaining = self.remaining.saturating_sub(buf.filled().len() - before); + if self.remaining == 0 + && let Some(read) = self.read.take() + { + let _ = read.send(()); + } + } + result + } + } + + enum FullRequestChannelAction { + Kill, + Stop, + Drop, + Drain, + RemoteKill, + } + + async fn full_request_channel(action: FullRequestChannelAction) { + use tokio_util::codec::Encoder; + let RequestReaderHarness { + mut framed_server, + framed_reader, + controller, + .. + } = request_reader_harness().await; + let (bytes_tx, mut bytes_rx) = mpsc::channel(1); + bytes_tx.send(Bytes::from_static(b"first")).await.unwrap(); + let message = TwoPartMessage::from_data(Bytes::from_static(b"second")); + let mut encoded = bytes::BytesMut::new(); + TwoPartCodec::default() + .encode(message, &mut encoded) + .unwrap(); + let (read_tx, read_rx) = oneshot::channel(); + let reader = FramedRead::new( + Box::new(FrameReadProbe { + inner: framed_reader.into_inner(), + remaining: encoded.len(), + read: Some(read_tx), + }) as BoxRead, + TwoPartCodec::default(), + ); + let counter = + IntCounter::new("full_request_channel_cancellations", "test counter").unwrap(); + let mut task = tokio::spawn(handle_request_reader( + reader, + bytes_tx, + controller.clone(), + Some(counter.clone()), + )); + framed_server + .send(TwoPartMessage::from_data(Bytes::from_static(b"second"))) + .await + .unwrap(); + time::timeout(Duration::from_secs(1), read_rx) + .await + .unwrap() + .unwrap(); + assert_eq!(bytes_rx.len(), 1); + assert!(!task.is_finished()); + + match action { + FullRequestChannelAction::Kill | FullRequestChannelAction::Stop => { + if matches!(action, FullRequestChannelAction::Kill) { + controller.kill(); + } else { + controller.stop(); + } + time::timeout(Duration::from_secs(1), &mut task) + .await + .expect("local cancellation must interrupt a full request channel") + .unwrap(); + assert!( + bytes_rx.is_closed(), + "all reader-owned senders must be dropped" + ); + assert_eq!(bytes_rx.recv().await.unwrap(), Bytes::from_static(b"first")); + assert!( + bytes_rx.recv().await.is_none(), + "pending frame must not be forwarded" + ); + assert_eq!( + counter.get(), + 0, + "local cancellation is not upstream cancellation" + ); + assert_eq!( + controller.is_killed(), + matches!(action, FullRequestChannelAction::Kill) + ); + } + FullRequestChannelAction::Drop => { + drop(bytes_rx); + time::timeout(Duration::from_secs(1), &mut task) + .await + .unwrap() + .unwrap(); + assert!(!controller.is_stopped()); + assert_eq!(counter.get(), 0); + } + FullRequestChannelAction::Drain | FullRequestChannelAction::RemoteKill => { + let control = match action { + FullRequestChannelAction::RemoteKill => ControlMessage::Kill, + _ => ControlMessage::Sentinel, + }; + // A wire control frame follows the blocked data frame. It is not + // equivalent to a local context signal and cannot bypass FIFO. + framed_server.send(control_message(&control)).await.unwrap(); + assert!(!controller.is_stopped()); + assert_eq!(bytes_rx.recv().await.unwrap(), Bytes::from_static(b"first")); + time::timeout(Duration::from_secs(1), &mut task) + .await + .unwrap() + .unwrap(); + assert_eq!( + bytes_rx.recv().await.unwrap(), + Bytes::from_static(b"second") + ); + assert!(bytes_rx.recv().await.is_none()); + let remote = !matches!(action, FullRequestChannelAction::Drain); + assert_eq!(controller.is_stopped(), remote); + assert_eq!( + controller.is_killed(), + matches!(action, FullRequestChannelAction::RemoteKill) + ); + assert_eq!(counter.get(), u64::from(remote)); + } + } + } + + #[tokio::test] + async fn test_full_request_channel_kill() { + full_request_channel(FullRequestChannelAction::Kill).await; + } + #[tokio::test] + async fn test_full_request_channel_stop() { + full_request_channel(FullRequestChannelAction::Stop).await; + } + #[tokio::test] + async fn test_full_request_channel_drop() { + full_request_channel(FullRequestChannelAction::Drop).await; + } + #[tokio::test] + async fn test_full_request_channel_drain() { + full_request_channel(FullRequestChannelAction::Drain).await; + } + #[tokio::test] + async fn test_full_request_channel_remote_kill() { + full_request_channel(FullRequestChannelAction::RemoteKill).await; + } /// Receiving Stop calls context.stop(), increments the counter, and exits. #[tokio::test] async fn test_handle_request_reader_stop_control_message() {