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

iota_network_stack/
grpc_timeout.rs
