Skip to main content

sui_rpc/client/
watchdog.rs

1//! Client-side watchdog for response bodies.
2//!
3//! Every RPC made through a [`Client`](super::Client) is multiplexed over a
4//! single HTTP/2 connection whose connection-level receive window is shared
5//! by all in-flight responses. Response data counts against that window until
6//! the application polls it out of the stream, so a streaming response that
7//! is held without being polled pins up to a full stream window of the shared
8//! budget. Enough stalled streams starve the connection, and then every RPC
9//! on the channel hangs indefinitely while TCP and HTTP/2 keepalives stay
10//! healthy. This converts into a permanent hang whenever the task that should
11//! be polling a stream is itself blocked on another call on the same starved
12//! channel: every party waits for connection window while holding connection
13//! window.
14//!
15//! The watchdog breaks both the starvation and the deadlock. Each response
16//! body is moved into a spawned task and bridged back to the caller through a
17//! bounded channel. The task's outermost await is an idle timer, so it stays
18//! responsive even when the caller has parked the response and nothing is
19//! polling it. If a whole idle period passes without a frame being delivered
20//! to the caller, the task drops the inner body, which resets the stream
21//! (RST_STREAM is a control frame and is deliverable even when the connection
22//! window is exhausted) and releases every byte of window the stream had
23//! pinned. The caller observes a `DeadlineExceeded` status on its next poll.
24//!
25//! A timer raced against the caller's own polls would not be enough: a parked
26//! stream is never polled, so a poll-driven timer freezes exactly when it is
27//! needed. The spawned task gives the timer its own polling root.
28//!
29//! The watchdog task also enforces the request's `grpc-timeout` deadline (set
30//! with [`tonic::Request::set_timeout`]) across the response body. tonic's
31//! own channel machinery only enforces that header until response headers
32//! arrive, which on a starved connection is precisely the part that still
33//! works -- headers are not flow-controlled -- so without this the deadline
34//! never bounds the actual hang. With it, `set_timeout` is a total per-call
35//! deadline: the same value is enforced locally end to end and communicated
36//! to the server through the standard header.
37
38use 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
53/// Default maximum time between response-body progress events before the
54/// watchdog resets the stream. Comfortably above the fullnode's subscription
55/// watermark cadence (~5s), so a healthy subscription always makes progress
56/// well within it.
57pub(super) const DEFAULT_BODY_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
58
59/// Per-request override for the client's idle-body watchdog.
60///
61/// Insert it into a request's extensions to raise, lower, or disable the
62/// watchdog timeout for that call only:
63///
64/// ```
65/// use std::time::Duration;
66/// use sui_rpc::client::BodyIdleTimeout;
67/// use sui_rpc::proto::sui::rpc::v2::SubscribeCheckpointsRequest;
68///
69/// let mut request = tonic::Request::new(SubscribeCheckpointsRequest::default());
70/// request
71///     .extensions_mut()
72///     .insert(BodyIdleTimeout::new(Duration::from_secs(120)));
73/// ```
74#[derive(Debug, Clone, Copy)]
75pub struct BodyIdleTimeout(Option<Duration>);
76
77impl BodyIdleTimeout {
78    /// Override the watchdog's idle timeout for this request.
79    pub fn new(timeout: Duration) -> Self {
80        Self(Some(timeout))
81    }
82
83    /// Disable the watchdog for this request. The response body can then
84    /// stall indefinitely; only use this when the caller enforces its own
85    /// bound on the call.
86    pub fn disabled() -> Self {
87        Self(None)
88    }
89}
90
91/// [`Layer`] that applies [`Watchdog`] to the client's transport.
92#[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
114/// Service that moves each response body into a watchdog task (see the
115/// module documentation).
116pub(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        // The deadline starts at request time, so the budget the header
141        // communicates to the server is the same one enforced locally.
142        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
156/// Parse the standard `grpc-timeout` request header (1-8 ASCII digits
157/// followed by a unit, per the gRPC-over-HTTP/2 spec). Malformed values are
158/// ignored, matching tonic's own leniency, since failing the call over an
159/// unenforceable header would be worse than not enforcing it.
160fn 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
179/// The caller-facing half of the bridge: receives frames from the watchdog
180/// task over a bounded channel and aborts the task if dropped, so dropping a
181/// response still resets the underlying stream promptly.
182struct WatchdogBody {
183    rx: tokio::sync::mpsc::Receiver<Result<Frame<Bytes>, Status>>,
184    /// `Some` until the task's completion has been observed. Consulted when
185    /// the channel closes without a terminal status so that a panic in the
186    /// watchdog task surfaces as an error instead of a clean end-of-stream.
187    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        // Capacity 1: the bridge only needs to decouple the two polling
197        // roots, and the tightest bound keeps the transport's own
198        // flow-control backpressure intact.
199        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                            // A clean exit, or cancellation via our own
239                            // abort-on-drop (not observable here since we
240                            // are being polled, not dropped).
241                            Ok(()) | Err(_) => Poll::Ready(None),
242                        }
243                    }
244                    Poll::Pending => Poll::Pending,
245                }
246            }
247            Poll::Pending => Poll::Pending,
248        }
249    }
250}
251
252/// The watchdog task: pump frames from the transport to the bridge channel,
253/// bounding each unit of progress with the idle timer and the whole body
254/// with the request deadline. On expiry, drop the body first (resetting the
255/// stream and releasing its pinned flow-control window in real time), then
256/// deliver the terminal status at the consumer's pace.
257async 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
280/// Returns `Some(status)` if the stream was cut by the watchdog, and `None`
281/// on natural end-of-stream, transport error, or consumer hang-up. The inner
282/// body is dropped on return in all cases -- including cancellation by the
283/// deadline in [`drive`], which drops this future and the body it owns.
284async 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        // One unit of progress: wait for bridge capacity, then receive a
291        // frame from the transport and hand it over through the permit.
292        // Reserving before polling the body means a frame is only taken from
293        // the transport once it is immediately deliverable: nothing sits in
294        // limbo between the transport and the bridge, and the frame's
295        // receive-window capacity is not released (inviting the server to
296        // send more) until the consumer can actually accept it. The timer
297        // covers both awaits, so it fires both when the consumer has parked
298        // the response without polling it (reserve never resolves) and when
299        // the transport delivers nothing (a starved or dead connection).
300        let step = async {
301            let Ok(permit) = tx.reserve().await else {
302                // Consumer dropped the response.
303                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    /// Test body driven by an unbounded channel, holding an optional beacon
345    /// whose `Drop` proves when the watchdog task has dropped the body.
346    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    /// `start_paused = true` makes `tokio::time` virtual: sleeps advance only
394    /// when the runtime is otherwise idle, so these tests run instantly and
395    /// deterministically.
396    #[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    /// The core starvation defense: when the consumer parks the response
443    /// without polling it, the watchdog must still fire and drop the inner
444    /// body (releasing its pinned flow-control window) in wall-clock time,
445    /// not on the consumer's next poll.
446    #[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        // The consumer takes one frame, then parks the response entirely.
452        tx.send(Ok(data_frame(b"first"))).unwrap();
453        assert!(next_frame(&mut watched).await.unwrap().is_ok());
454
455        // The transport keeps delivering: one frame sits in the bridge
456        // channel and the pump waits for bridge capacity before taking the
457        // next one.
458        tx.send(Ok(data_frame(b"buffered"))).unwrap();
459        tx.send(Ok(data_frame(b"blocked"))).unwrap();
460
461        // Without any consumer poll, the idle timer must drop the body.
462        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        // When the consumer returns it drains the buffered frame, then
469        // observes the watchdog cut.
470        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    /// Reserve-before-pull: the pump takes a frame from the transport only
477    /// once the bridge can accept it, so a parked consumer leaves at most one
478    /// frame buffered in the bridge and everything else in the transport,
479    /// where it remains subject to HTTP/2 flow control.
480    #[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        // The transport has plenty of frames available.
513        for _ in 0..5 {
514            tx.send(Ok(data_frame(b"tick"))).unwrap();
515        }
516
517        // The consumer takes one frame and parks. The pump may refill the
518        // one-slot bridge, but must not pull frames it cannot deliver.
519        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        // The consumer still sees the buffered frame and then the cut.
528        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    /// A panic in the watchdog task must surface as an error instead of
565    /// silently closing the bridge, which the consumer could not distinguish
566    /// from a clean end-of-stream.
567    #[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    /// The request deadline bounds the whole body even when every individual
589    /// frame arrives well within the idle timeout.
590    #[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        // Past the deadline the body must be dropped and the consumer must
604        // observe DeadlineExceeded, even though frames kept flowing.
605        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    /// The deadline is enforced even when the idle watchdog is disabled.
615    #[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    /// End-to-end through the service: the `grpc-timeout` header (what
627    /// `tonic::Request::set_timeout` produces) bounds the response body, and
628    /// a `BodyIdleTimeout::disabled()` extension does not interfere with it.
629    #[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    /// With neither an idle timeout nor a deadline, the response body must be
655    /// passed through without spawning a bridge.
656    #[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        // Nothing cuts the stream no matter how long it idles.
674        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        // Malformed values are ignored rather than failing the call.
697        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}