1use std::future::Future;
39use std::ops::ControlFlow;
40use std::pin::Pin;
41use std::task::Context;
42use std::task::Poll;
43use std::time::Duration;
44
45use bytes::Bytes;
46use http_body::Body as _;
47use http_body::Frame;
48use tonic::Status;
49use tonic::body::Body;
50use tower::Layer;
51use tower::Service;
52
53pub(super) const DEFAULT_BODY_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
58
59#[derive(Debug, Clone, Copy)]
75pub struct BodyIdleTimeout(Option<Duration>);
76
77impl BodyIdleTimeout {
78 pub fn new(timeout: Duration) -> Self {
80 Self(Some(timeout))
81 }
82
83 pub fn disabled() -> Self {
87 Self(None)
88 }
89}
90
91#[derive(Clone)]
93pub(super) struct WatchdogLayer {
94 idle_timeout: Option<Duration>,
95}
96
97impl WatchdogLayer {
98 pub(super) fn new(idle_timeout: Option<Duration>) -> Self {
99 Self { idle_timeout }
100 }
101}
102
103impl<S> Layer<S> for WatchdogLayer {
104 type Service = Watchdog<S>;
105
106 fn layer(&self, inner: S) -> Self::Service {
107 Watchdog {
108 inner,
109 idle_timeout: self.idle_timeout,
110 }
111 }
112}
113
114pub(super) struct Watchdog<S> {
117 inner: S,
118 idle_timeout: Option<Duration>,
119}
120
121impl<S> Service<http::Request<Body>> for Watchdog<S>
122where
123 S: Service<http::Request<Body>, Response = http::Response<Body>>,
124 S::Future: Send + 'static,
125 S::Error: Send + 'static,
126{
127 type Response = http::Response<Body>;
128 type Error = S::Error;
129 type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;
130
131 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
132 self.inner.poll_ready(cx)
133 }
134
135 fn call(&mut self, request: http::Request<Body>) -> Self::Future {
136 let idle_timeout = match request.extensions().get::<BodyIdleTimeout>() {
137 Some(BodyIdleTimeout(override_timeout)) => *override_timeout,
138 None => self.idle_timeout,
139 };
140 let deadline = parse_grpc_timeout(request.headers())
143 .map(|timeout| tokio::time::Instant::now() + timeout);
144
145 let response = self.inner.call(request);
146 Box::pin(async move {
147 let response = response.await?;
148 if idle_timeout.is_none() && deadline.is_none() {
149 return Ok(response);
150 }
151 Ok(response.map(|body| Body::new(WatchdogBody::spawn(body, idle_timeout, deadline))))
152 })
153 }
154}
155
156fn parse_grpc_timeout(headers: &http::HeaderMap) -> Option<Duration> {
161 let value = headers.get("grpc-timeout")?.to_str().ok()?;
162 let (magnitude, unit) = value.split_at(value.len().checked_sub(1)?);
163 if magnitude.is_empty() || magnitude.len() > 8 || !magnitude.bytes().all(|b| b.is_ascii_digit())
164 {
165 return None;
166 }
167 let magnitude: u64 = magnitude.parse().ok()?;
168 match unit {
169 "H" => Some(Duration::from_secs(magnitude * 60 * 60)),
170 "M" => Some(Duration::from_secs(magnitude * 60)),
171 "S" => Some(Duration::from_secs(magnitude)),
172 "m" => Some(Duration::from_millis(magnitude)),
173 "u" => Some(Duration::from_micros(magnitude)),
174 "n" => Some(Duration::from_nanos(magnitude)),
175 _ => None,
176 }
177}
178
179struct WatchdogBody {
183 rx: tokio::sync::mpsc::Receiver<Result<Frame<Bytes>, Status>>,
184 task: Option<tokio::task::JoinHandle<()>>,
188}
189
190impl WatchdogBody {
191 fn spawn(
192 body: Body,
193 idle_timeout: Option<Duration>,
194 deadline: Option<tokio::time::Instant>,
195 ) -> Self {
196 let (tx, rx) = tokio::sync::mpsc::channel(1);
200 let task = tokio::spawn(drive(body, tx, idle_timeout, deadline));
201 Self {
202 rx,
203 task: Some(task),
204 }
205 }
206}
207
208impl Drop for WatchdogBody {
209 fn drop(&mut self) {
210 if let Some(task) = &self.task {
211 task.abort();
212 }
213 }
214}
215
216impl http_body::Body for WatchdogBody {
217 type Data = Bytes;
218 type Error = Status;
219
220 fn poll_frame(
221 mut self: Pin<&mut Self>,
222 cx: &mut Context<'_>,
223 ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
224 let this = &mut *self;
225 match this.rx.poll_recv(cx) {
226 Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
227 Poll::Ready(None) => {
228 let Some(task) = &mut this.task else {
229 return Poll::Ready(None);
230 };
231 match Pin::new(task).poll(cx) {
232 Poll::Ready(result) => {
233 this.task = None;
234 match result {
235 Err(e) if e.is_panic() => Poll::Ready(Some(Err(Status::internal(
236 "client idle-body watchdog task panicked",
237 )))),
238 Ok(()) | Err(_) => Poll::Ready(None),
242 }
243 }
244 Poll::Pending => Poll::Pending,
245 }
246 }
247 Poll::Pending => Poll::Pending,
248 }
249 }
250}
251
252async fn drive(
258 body: Body,
259 tx: tokio::sync::mpsc::Sender<Result<Frame<Bytes>, Status>>,
260 idle_timeout: Option<Duration>,
261 deadline: Option<tokio::time::Instant>,
262) {
263 let pump = pump(body, tx.clone(), idle_timeout);
264 let cut = match deadline {
265 Some(deadline) => tokio::time::timeout_at(deadline, pump)
266 .await
267 .unwrap_or_else(|_| {
268 Some(Status::deadline_exceeded(
269 "grpc-timeout deadline exceeded before the response body completed; \
270 stream reset by the client",
271 ))
272 }),
273 None => pump.await,
274 };
275 if let Some(status) = cut {
276 let _ = tx.send(Err(status)).await;
277 }
278}
279
280async fn pump(
285 mut body: Body,
286 tx: tokio::sync::mpsc::Sender<Result<Frame<Bytes>, Status>>,
287 idle_timeout: Option<Duration>,
288) -> Option<Status> {
289 loop {
290 let step = async {
301 let Ok(permit) = tx.reserve().await else {
302 return ControlFlow::Break(());
304 };
305 match std::future::poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await {
306 Some(Ok(frame)) => {
307 permit.send(Ok(frame));
308 ControlFlow::Continue(())
309 }
310 Some(Err(status)) => {
311 permit.send(Err(status));
312 ControlFlow::Break(())
313 }
314 None => ControlFlow::Break(()),
315 }
316 };
317
318 let flow = match idle_timeout {
319 Some(idle_timeout) => match tokio::time::timeout(idle_timeout, step).await {
320 Ok(flow) => flow,
321 Err(_) => {
322 return Some(Status::deadline_exceeded(format!(
323 "no response-body progress within {idle_timeout:?}; \
324 stream reset by the client's idle-body watchdog"
325 )));
326 }
327 },
328 None => step.await,
329 };
330 match flow {
331 ControlFlow::Continue(()) => {}
332 ControlFlow::Break(()) => return None,
333 }
334 }
335}
336
337#[cfg(test)]
338mod tests {
339 use super::*;
340 use std::sync::Arc;
341 use std::sync::atomic::AtomicBool;
342 use std::sync::atomic::Ordering;
343
344 struct TestBody {
347 rx: tokio::sync::mpsc::UnboundedReceiver<Result<Frame<Bytes>, Status>>,
348 _beacon: Option<DropBeacon>,
349 }
350
351 struct DropBeacon(Arc<AtomicBool>);
352 impl Drop for DropBeacon {
353 fn drop(&mut self) {
354 self.0.store(true, Ordering::SeqCst);
355 }
356 }
357
358 impl http_body::Body for TestBody {
359 type Data = Bytes;
360 type Error = Status;
361
362 fn poll_frame(
363 mut self: Pin<&mut Self>,
364 cx: &mut Context<'_>,
365 ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
366 self.rx.poll_recv(cx)
367 }
368 }
369
370 #[allow(clippy::type_complexity)]
371 fn test_body() -> (
372 tokio::sync::mpsc::UnboundedSender<Result<Frame<Bytes>, Status>>,
373 Arc<AtomicBool>,
374 Body,
375 ) {
376 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
377 let dropped = Arc::new(AtomicBool::new(false));
378 let body = Body::new(TestBody {
379 rx,
380 _beacon: Some(DropBeacon(dropped.clone())),
381 });
382 (tx, dropped, body)
383 }
384
385 async fn next_frame(body: &mut WatchdogBody) -> Option<Result<Frame<Bytes>, Status>> {
386 std::future::poll_fn(|cx| Pin::new(&mut *body).poll_frame(cx)).await
387 }
388
389 fn data_frame(bytes: &'static [u8]) -> Frame<Bytes> {
390 Frame::data(Bytes::from_static(bytes))
391 }
392
393 #[tokio::test(start_paused = true)]
397 async fn frames_and_trailers_pass_through() {
398 let (tx, _dropped, body) = test_body();
399 let mut watched = WatchdogBody::spawn(body, Some(Duration::from_secs(30)), None);
400
401 tx.send(Ok(data_frame(b"hello"))).unwrap();
402 let frame = next_frame(&mut watched).await.unwrap().unwrap();
403 assert_eq!(frame.into_data().unwrap(), Bytes::from_static(b"hello"));
404
405 let mut trailers = http::HeaderMap::new();
406 trailers.insert("grpc-status", "0".parse().unwrap());
407 tx.send(Ok(Frame::trailers(trailers.clone()))).unwrap();
408 let frame = next_frame(&mut watched).await.unwrap().unwrap();
409 assert_eq!(frame.into_trailers().unwrap(), trailers);
410
411 drop(tx);
412 assert!(next_frame(&mut watched).await.is_none());
413 }
414
415 #[tokio::test(start_paused = true)]
416 async fn slow_but_live_stream_is_not_cut() {
417 let (tx, _dropped, body) = test_body();
418 let mut watched = WatchdogBody::spawn(body, Some(Duration::from_secs(30)), None);
419
420 for _ in 0..5 {
421 tokio::time::sleep(Duration::from_secs(20)).await;
422 tx.send(Ok(data_frame(b"tick"))).unwrap();
423 let frame = next_frame(&mut watched).await.unwrap().unwrap();
424 assert_eq!(frame.into_data().unwrap(), Bytes::from_static(b"tick"));
425 }
426
427 drop(tx);
428 assert!(next_frame(&mut watched).await.is_none());
429 }
430
431 #[tokio::test(start_paused = true)]
432 async fn idle_transport_is_cut_and_body_dropped() {
433 let (_tx, dropped, body) = test_body();
434 let mut watched = WatchdogBody::spawn(body, Some(Duration::from_secs(30)), None);
435
436 let status = next_frame(&mut watched).await.unwrap().unwrap_err();
437 assert_eq!(status.code(), tonic::Code::DeadlineExceeded);
438 assert!(dropped.load(Ordering::SeqCst), "inner body was not dropped");
439 assert!(next_frame(&mut watched).await.is_none());
440 }
441
442 #[tokio::test(start_paused = true)]
447 async fn parked_consumer_is_cut_without_being_polled() {
448 let (tx, dropped, body) = test_body();
449 let mut watched = WatchdogBody::spawn(body, Some(Duration::from_secs(30)), None);
450
451 tx.send(Ok(data_frame(b"first"))).unwrap();
453 assert!(next_frame(&mut watched).await.unwrap().is_ok());
454
455 tx.send(Ok(data_frame(b"buffered"))).unwrap();
459 tx.send(Ok(data_frame(b"blocked"))).unwrap();
460
461 tokio::time::sleep(Duration::from_secs(60)).await;
463 assert!(
464 dropped.load(Ordering::SeqCst),
465 "inner body was not dropped while the consumer was parked"
466 );
467
468 let frame = next_frame(&mut watched).await.unwrap().unwrap();
471 assert_eq!(frame.into_data().unwrap(), Bytes::from_static(b"buffered"));
472 let status = next_frame(&mut watched).await.unwrap().unwrap_err();
473 assert_eq!(status.code(), tonic::Code::DeadlineExceeded);
474 }
475
476 #[tokio::test(start_paused = true)]
481 async fn parked_consumer_does_not_double_buffer() {
482 use std::sync::atomic::AtomicU64;
483
484 struct CountingBody {
485 rx: tokio::sync::mpsc::UnboundedReceiver<Result<Frame<Bytes>, Status>>,
486 pulled: Arc<AtomicU64>,
487 }
488 impl http_body::Body for CountingBody {
489 type Data = Bytes;
490 type Error = Status;
491
492 fn poll_frame(
493 mut self: Pin<&mut Self>,
494 cx: &mut Context<'_>,
495 ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
496 let poll = self.rx.poll_recv(cx);
497 if matches!(poll, Poll::Ready(Some(_))) {
498 self.pulled.fetch_add(1, Ordering::SeqCst);
499 }
500 poll
501 }
502 }
503
504 let pulled = Arc::new(AtomicU64::new(0));
505 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
506 let body = Body::new(CountingBody {
507 rx,
508 pulled: pulled.clone(),
509 });
510 let mut watched = WatchdogBody::spawn(body, Some(Duration::from_secs(30)), None);
511
512 for _ in 0..5 {
514 tx.send(Ok(data_frame(b"tick"))).unwrap();
515 }
516
517 assert!(next_frame(&mut watched).await.unwrap().is_ok());
520 tokio::time::sleep(Duration::from_secs(60)).await;
521 assert_eq!(
522 pulled.load(Ordering::SeqCst),
523 2,
524 "the pump pulled frames from the transport that it could not deliver"
525 );
526
527 assert!(next_frame(&mut watched).await.unwrap().is_ok());
529 let status = next_frame(&mut watched).await.unwrap().unwrap_err();
530 assert_eq!(status.code(), tonic::Code::DeadlineExceeded);
531 }
532
533 #[tokio::test(start_paused = true)]
534 async fn transport_error_passes_through() {
535 let (tx, _dropped, body) = test_body();
536 let mut watched = WatchdogBody::spawn(body, Some(Duration::from_secs(30)), None);
537
538 tx.send(Err(Status::unavailable("connection reset")))
539 .unwrap();
540 let status = next_frame(&mut watched).await.unwrap().unwrap_err();
541 assert_eq!(status.code(), tonic::Code::Unavailable);
542 assert!(next_frame(&mut watched).await.is_none());
543 }
544
545 #[tokio::test(start_paused = true)]
546 async fn dropping_response_aborts_task_and_drops_body() {
547 let (tx, dropped, body) = test_body();
548 let watched = WatchdogBody::spawn(body, Some(Duration::from_secs(30)), None);
549 tx.send(Ok(data_frame(b"unread"))).unwrap();
550
551 drop(watched);
552 for _ in 0..10 {
553 tokio::task::yield_now().await;
554 if dropped.load(Ordering::SeqCst) {
555 break;
556 }
557 }
558 assert!(
559 dropped.load(Ordering::SeqCst),
560 "inner body was not dropped after the response was dropped"
561 );
562 }
563
564 #[tokio::test(start_paused = true)]
568 async fn task_panic_surfaces_as_internal() {
569 struct PanicBody;
570 impl http_body::Body for PanicBody {
571 type Data = Bytes;
572 type Error = Status;
573
574 fn poll_frame(
575 self: Pin<&mut Self>,
576 _cx: &mut Context<'_>,
577 ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
578 panic!("boom from inner body");
579 }
580 }
581
582 let mut watched =
583 WatchdogBody::spawn(Body::new(PanicBody), Some(Duration::from_secs(30)), None);
584 let status = next_frame(&mut watched).await.unwrap().unwrap_err();
585 assert_eq!(status.code(), tonic::Code::Internal);
586 }
587
588 #[tokio::test(start_paused = true)]
591 async fn deadline_cuts_a_live_stream() {
592 let (tx, dropped, body) = test_body();
593 let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
594 let mut watched =
595 WatchdogBody::spawn(body, Some(Duration::from_secs(3600)), Some(deadline));
596
597 for _ in 0..4 {
598 tokio::time::sleep(Duration::from_secs(5)).await;
599 tx.send(Ok(data_frame(b"tick"))).unwrap();
600 assert!(next_frame(&mut watched).await.unwrap().is_ok());
601 }
602
603 tokio::time::sleep(Duration::from_secs(15)).await;
606 assert!(
607 dropped.load(Ordering::SeqCst),
608 "inner body was not dropped at the deadline"
609 );
610 let status = next_frame(&mut watched).await.unwrap().unwrap_err();
611 assert_eq!(status.code(), tonic::Code::DeadlineExceeded);
612 }
613
614 #[tokio::test(start_paused = true)]
616 async fn deadline_applies_without_idle_timeout() {
617 let (_tx, dropped, body) = test_body();
618 let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
619 let mut watched = WatchdogBody::spawn(body, None, Some(deadline));
620
621 let status = next_frame(&mut watched).await.unwrap().unwrap_err();
622 assert_eq!(status.code(), tonic::Code::DeadlineExceeded);
623 assert!(dropped.load(Ordering::SeqCst), "inner body was not dropped");
624 }
625
626 #[tokio::test(start_paused = true)]
630 async fn service_honors_grpc_timeout_header() {
631 let (_tx, dropped, body) = test_body();
632 let mut body = Some(body);
633 let mut service = WatchdogLayer::new(None).layer(tower::service_fn(
634 move |_request: http::Request<Body>| {
635 let body = body.take().unwrap();
636 async move { Ok::<_, Status>(http::Response::new(body)) }
637 },
638 ));
639
640 let mut request = http::Request::new(Body::empty());
641 request
642 .headers_mut()
643 .insert("grpc-timeout", "30S".parse().unwrap());
644 request.extensions_mut().insert(BodyIdleTimeout::disabled());
645
646 let response = service.call(request).await.unwrap();
647 let mut body = response.into_body();
648 let item = std::future::poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await;
649 let status = item.unwrap().unwrap_err();
650 assert_eq!(status.code(), tonic::Code::DeadlineExceeded);
651 assert!(dropped.load(Ordering::SeqCst), "inner body was not dropped");
652 }
653
654 #[tokio::test(start_paused = true)]
657 async fn service_passes_body_through_when_unbounded() {
658 let (tx, _dropped, body) = test_body();
659 let mut body = Some(body);
660 let mut service = WatchdogLayer::new(None).layer(tower::service_fn(
661 move |_request: http::Request<Body>| {
662 let body = body.take().unwrap();
663 async move { Ok::<_, Status>(http::Response::new(body)) }
664 },
665 ));
666
667 let response = service
668 .call(http::Request::new(Body::empty()))
669 .await
670 .unwrap();
671 let mut body = response.into_body();
672
673 tokio::time::sleep(Duration::from_secs(3600)).await;
675 tx.send(Ok(data_frame(b"still here"))).unwrap();
676 let item = std::future::poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await;
677 assert!(item.unwrap().is_ok());
678 }
679
680 #[test]
681 fn grpc_timeout_parsing() {
682 fn parse(value: &str) -> Option<Duration> {
683 let mut headers = http::HeaderMap::new();
684 headers.insert("grpc-timeout", value.parse().unwrap());
685 parse_grpc_timeout(&headers)
686 }
687
688 assert_eq!(parse("2H"), Some(Duration::from_secs(2 * 60 * 60)));
689 assert_eq!(parse("3M"), Some(Duration::from_secs(180)));
690 assert_eq!(parse("30S"), Some(Duration::from_secs(30)));
691 assert_eq!(parse("500m"), Some(Duration::from_millis(500)));
692 assert_eq!(parse("250u"), Some(Duration::from_micros(250)));
693 assert_eq!(parse("82n"), Some(Duration::from_nanos(82)));
694 assert_eq!(parse("99999999S"), Some(Duration::from_secs(99_999_999)));
695
696 assert_eq!(parse(""), None);
698 assert_eq!(parse("S"), None);
699 assert_eq!(parse("30"), None);
700 assert_eq!(parse("30f"), None);
701 assert_eq!(parse("-30S"), None);
702 assert_eq!(parse("+30S"), None);
703 assert_eq!(parse("3.5S"), None);
704 assert_eq!(parse("123456789S"), None);
705 assert_eq!(
706 parse_grpc_timeout(&http::HeaderMap::new()),
707 None,
708 "absent header"
709 );
710 }
711}