iota_network_stack/
grpc_timeout.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2025 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4//
5// Ported from `tonic` crate
6// SPDX-License-Identifier: MIT
7
8use std::{
9    future::Future,
10    pin::Pin,
11    task::{Context, Poll, ready},
12    time::Duration,
13};
14
15use http::{HeaderMap, HeaderValue, Request, Response};
16use pin_project_lite::pin_project;
17use tokio::time::Sleep;
18use tonic::Status;
19use tower::Service;
20
21const GRPC_TIMEOUT_HEADER: &str = "grpc-timeout";
22
23/// Answers with the status `DEADLINE_EXCEEDED` and the default response body,
24/// which must be empty, if the inner service has not returned a response by the
25/// shorter of the client's `grpc-timeout` and `server_timeout`. Paths in
26/// `timeout_exempt_paths` skip `server_timeout`. Once the response is returned,
27/// sending its body has no deadline.
28#[derive(Debug, Clone)]
29pub struct GrpcTimeout<S> {
30    inner: S,
31    server_timeout: Option<Duration>,
32    /// Request URI paths exempt from `server_timeout` — for long-lived
33    /// server-streaming RPCs that must not be aborted by a fallback deadline.
34    /// An explicit client `grpc-timeout` is still honored for these paths.
35    timeout_exempt_paths: &'static [&'static str],
36}
37
38impl<S> GrpcTimeout<S> {
39    pub fn new(inner: S, server_timeout: Option<Duration>) -> Self {
40        Self::new_with_exempt_paths(inner, server_timeout, &[])
41    }
42
43    pub fn new_with_exempt_paths(
44        inner: S,
45        server_timeout: Option<Duration>,
46        timeout_exempt_paths: &'static [&'static str],
47    ) -> Self {
48        Self {
49            inner,
50            server_timeout,
51            timeout_exempt_paths,
52        }
53    }
54}
55
56impl<S, RequestBody, ResponseBody> Service<Request<RequestBody>> for GrpcTimeout<S>
57where
58    S: Service<Request<RequestBody>, Response = Response<ResponseBody>>,
59    ResponseBody: Default,
60{
61    type Response = Response<ResponseBody>;
62    type Error = S::Error;
63    type Future = ResponseFuture<S::Future>;
64
65    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
66        self.inner.poll_ready(cx).map_err(Into::into)
67    }
68
69    fn call(&mut self, req: Request<RequestBody>) -> Self::Future {
70        let client_timeout = try_parse_grpc_timeout(req.headers()).unwrap_or_else(|e| {
71            tracing::trace!("Error parsing `grpc-timeout` header {:?}", e);
72            None
73        });
74
75        // Exempt configured paths (long-lived server-streaming RPCs) from the
76        // server-side fallback timeout; an explicit client deadline still applies.
77        let server_timeout = if self.timeout_exempt_paths.contains(&req.uri().path()) {
78            None
79        } else {
80            self.server_timeout
81        };
82
83        // Use the shorter of the two durations, if either are set
84        let timeout_duration = match (client_timeout, server_timeout) {
85            (None, None) => None,
86            (Some(dur), None) => Some(dur),
87            (None, Some(dur)) => Some(dur),
88            (Some(header), Some(server)) => {
89                let shorter_duration = std::cmp::min(header, server);
90                Some(shorter_duration)
91            }
92        };
93
94        ResponseFuture {
95            inner: self.inner.call(req),
96            sleep: timeout_duration.map(tokio::time::sleep),
97        }
98    }
99}
100
101pin_project! {
102    pub struct ResponseFuture<F> {
103        #[pin]
104        inner: F,
105        #[pin]
106        sleep: Option<Sleep>,
107    }
108}
109
110impl<F, ResponseBody, E> Future for ResponseFuture<F>
111where
112    F: Future<Output = Result<Response<ResponseBody>, E>>,
113    ResponseBody: Default,
114{
115    type Output = F::Output;
116
117    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
118        let this = self.project();
119
120        if let Poll::Ready(result) = this.inner.poll(cx) {
121            return Poll::Ready(result);
122        }
123
124        if let Some(sleep) = this.sleep.as_pin_mut() {
125            ready!(sleep.poll(cx));
126            return Poll::Ready(Ok(Status::deadline_exceeded("Timeout expired").into_http()));
127        }
128
129        Poll::Pending
130    }
131}
132
133const SECONDS_IN_HOUR: u64 = 60 * 60;
134const SECONDS_IN_MINUTE: u64 = 60;
135
136/// Tries to parse the `grpc-timeout` header if it is present. If we fail to
137/// parse, returns the value we attempted to parse.
138///
139/// Follows the [gRPC over HTTP2 spec](https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md).
140fn try_parse_grpc_timeout(
141    headers: &HeaderMap<HeaderValue>,
142) -> Result<Option<Duration>, &HeaderValue> {
143    let Some(val) = headers.get(GRPC_TIMEOUT_HEADER) else {
144        return Ok(None);
145    };
146
147    let (timeout_value, timeout_unit) = val
148        .to_str()
149        .map_err(|_| val)
150        .and_then(|s| if s.is_empty() { Err(val) } else { Ok(s) })?
151        // `HeaderValue::to_str` only returns `Ok` if the header contains ASCII so this
152        // `split_at` will never panic from trying to split in the middle of a character.
153        // See https://docs.rs/http/0.2.4/http/header/struct.HeaderValue.html#method.to_str
154        //
155        // `len - 1` also wont panic since we just checked `s.is_empty`.
156        .split_at(val.len() - 1);
157
158    // gRPC spec specifies `TimeoutValue` will be at most 8 digits
159    // Caping this at 8 digits also prevents integer overflow from ever occurring
160    if timeout_value.len() > 8 {
161        return Err(val);
162    }
163
164    let timeout_value: u64 = timeout_value.parse().map_err(|_| val)?;
165
166    let duration = match timeout_unit {
167        // Hours
168        "H" => Duration::from_secs(timeout_value * SECONDS_IN_HOUR),
169        // Minutes
170        "M" => Duration::from_secs(timeout_value * SECONDS_IN_MINUTE),
171        // Seconds
172        "S" => Duration::from_secs(timeout_value),
173        // Milliseconds
174        "m" => Duration::from_millis(timeout_value),
175        // Microseconds
176        "u" => Duration::from_micros(timeout_value),
177        // Nanoseconds
178        "n" => Duration::from_nanos(timeout_value),
179        _ => return Err(val),
180    };
181
182    Ok(Some(duration))
183}
184
185#[cfg(test)]
186mod tests {
187    use std::{convert::Infallible, future, future::Pending, marker::PhantomData};
188
189    use http_body_util::BodyExt;
190    use tonic::body::Body;
191
192    use super::*;
193    use crate::concurrency::ReclaimableBody;
194
195    // Helper function to reduce the boiler plate of our test cases
196    fn setup_map_try_parse(val: Option<&str>) -> Result<Option<Duration>, HeaderValue> {
197        let mut hm = HeaderMap::new();
198        if let Some(v) = val {
199            let hv = HeaderValue::from_str(v).unwrap();
200            hm.insert(GRPC_TIMEOUT_HEADER, hv);
201        };
202
203        try_parse_grpc_timeout(&hm).map_err(|e| e.clone())
204    }
205
206    #[test]
207    fn test_hours() {
208        let parsed_duration = setup_map_try_parse(Some("3H")).unwrap().unwrap();
209        assert_eq!(Duration::from_secs(3 * 60 * 60), parsed_duration);
210    }
211
212    #[test]
213    fn test_minutes() {
214        let parsed_duration = setup_map_try_parse(Some("1M")).unwrap().unwrap();
215        assert_eq!(Duration::from_secs(60), parsed_duration);
216    }
217
218    #[test]
219    fn test_seconds() {
220        let parsed_duration = setup_map_try_parse(Some("42S")).unwrap().unwrap();
221        assert_eq!(Duration::from_secs(42), parsed_duration);
222    }
223
224    #[test]
225    fn test_milliseconds() {
226        let parsed_duration = setup_map_try_parse(Some("13m")).unwrap().unwrap();
227        assert_eq!(Duration::from_millis(13), parsed_duration);
228    }
229
230    #[test]
231    fn test_microseconds() {
232        let parsed_duration = setup_map_try_parse(Some("2u")).unwrap().unwrap();
233        assert_eq!(Duration::from_micros(2), parsed_duration);
234    }
235
236    #[test]
237    fn test_nanoseconds() {
238        let parsed_duration = setup_map_try_parse(Some("82n")).unwrap().unwrap();
239        assert_eq!(Duration::from_nanos(82), parsed_duration);
240    }
241
242    #[test]
243    fn test_header_not_present() {
244        let parsed_duration = setup_map_try_parse(None).unwrap();
245        assert!(parsed_duration.is_none());
246    }
247
248    #[test]
249    #[should_panic(expected = "82f")]
250    fn test_invalid_unit() {
251        // "f" is not a valid TimeoutUnit
252        setup_map_try_parse(Some("82f")).unwrap().unwrap();
253    }
254
255    #[test]
256    #[should_panic(expected = "123456789H")]
257    fn test_too_many_digits() {
258        // gRPC spec states TimeoutValue will be at most 8 digits
259        setup_map_try_parse(Some("123456789H")).unwrap().unwrap();
260    }
261
262    #[test]
263    #[should_panic(expected = "oneH")]
264    fn test_invalid_digits() {
265        // gRPC spec states TimeoutValue will be at most 8 digits
266        setup_map_try_parse(Some("oneH")).unwrap().unwrap();
267    }
268
269    /// A service whose responses never become ready, so the only way a call
270    /// resolves is via the timeout branch.
271    struct NeverResponds<B>(PhantomData<B>);
272
273    impl<B> NeverResponds<B> {
274        fn new() -> Self {
275            Self(PhantomData)
276        }
277    }
278
279    impl<B> Service<Request<()>> for NeverResponds<B> {
280        type Response = Response<B>;
281        type Error = Infallible;
282        type Future = Pending<Result<Response<B>, Infallible>>;
283
284        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
285            Poll::Ready(Ok(()))
286        }
287
288        fn call(&mut self, _req: Request<()>) -> Self::Future {
289            future::pending()
290        }
291    }
292
293    /// For each response body type that `GrpcTimeout` wraps in this repository.
294    #[tokio::test]
295    async fn a_call_past_its_deadline_ends_with_deadline_exceeded_and_an_empty_body() {
296        async fn assert_deadline_response<B>()
297        where
298            B: http_body::Body + Default,
299            B::Error: std::fmt::Debug,
300        {
301            let mut svc =
302                GrpcTimeout::new(NeverResponds::<B>::new(), Some(Duration::from_millis(10)));
303            let response = tokio::time::timeout(Duration::from_secs(5), svc.call(Request::new(())))
304                .await
305                .expect("the call must end at its deadline")
306                .unwrap();
307            assert_eq!(
308                Status::from_header_map(response.headers()).map(|status| status.code()),
309                Some(tonic::Code::DeadlineExceeded)
310            );
311            assert!(response.body().is_end_stream());
312            assert_eq!(response.body().size_hint().exact(), Some(0));
313            let body = response.into_body().collect().await.unwrap();
314            assert!(body.trailers().is_none());
315            assert!(body.to_bytes().is_empty());
316        }
317
318        assert_deadline_response::<Body>().await;
319        assert_deadline_response::<tonic::service::AxumBody>().await;
320        assert_deadline_response::<ReclaimableBody<Body, (), ()>>().await;
321    }
322
323    #[tokio::test]
324    async fn server_timeout_skips_exempt_paths() {
325        const EXEMPT: &[&str] = &["/pkg.Svc/Stream"];
326        let timeout = Duration::from_millis(50);
327
328        // Non-exempt path: the server-side fallback timeout fires.
329        let mut svc =
330            GrpcTimeout::new_with_exempt_paths(NeverResponds::<Body>::new(), Some(timeout), EXEMPT);
331        let fut = svc.call(Request::builder().uri("/pkg.Svc/Unary").body(()).unwrap());
332        assert!(
333            fut.await.is_ok(),
334            "non-exempt request should be completed by the fallback timeout"
335        );
336
337        // Exempt path: no fallback timeout, so the request is never aborted.
338        let mut svc =
339            GrpcTimeout::new_with_exempt_paths(NeverResponds::<Body>::new(), Some(timeout), EXEMPT);
340        let fut = svc.call(Request::builder().uri("/pkg.Svc/Stream").body(()).unwrap());
341        assert!(
342            tokio::time::timeout(Duration::from_millis(150), fut)
343                .await
344                .is_err(),
345            "exempt path must not be aborted by the server timeout"
346        );
347    }
348}
iota_network_stack/grpc_timeout.rs

iota_network_stack/
grpc_timeout.rs
