volo_grpc/server/layer/
timeout.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5    time::Duration,
6};
7
8use http::{HeaderMap, HeaderValue};
9use metainfo::METAINFO;
10use motore::{Service, layer::Layer};
11use pin_project::pin_project;
12use tokio::time::{self, Sleep};
13
14use crate::{Request, context::ServerContext, status::Status};
15
16#[derive(Debug, Clone)]
17pub struct Timeout<S> {
18    inner: S,
19}
20
21impl<S> Timeout<S> {
22    pub fn new(inner: S) -> Self {
23        Self { inner }
24    }
25}
26
27#[derive(Clone, Default, Copy)]
28pub struct TimeoutLayer;
29
30impl TimeoutLayer {
31    pub fn new() -> Self {
32        Self
33    }
34}
35
36impl<S> Layer<S> for TimeoutLayer {
37    type Service = Timeout<S>;
38
39    fn layer(self, inner: S) -> Self::Service {
40        Timeout { inner }
41    }
42}
43
44impl<S, T> Service<ServerContext, Request<T>> for Timeout<S>
45where
46    S: Service<ServerContext, Request<T>, Error = Status> + Send + Sync,
47    T: Send + 'static,
48{
49    type Response = S::Response;
50    type Error = Status;
51
52    async fn call(
53        &self,
54        _cx: &mut ServerContext,
55        req: Request<T>,
56    ) -> Result<Self::Response, Self::Error> {
57        let client_timeout =
58            grpc_timeout_to_duration(req.metadata().headers()).unwrap_or_else(|_| {
59                tracing::trace!("Failed to parse grpc-timeout header");
60                None
61            });
62
63        // insert duration from header (if any) into METAINFO
64        if let Some(timeout_val) = client_timeout {
65            METAINFO.with(|mi| {
66                mi.borrow_mut().insert::<Duration>(timeout_val);
67            });
68        }
69
70        let sleep = client_timeout.map(time::sleep);
71        let inner = self.inner.call(_cx, req);
72
73        ResponseFuture {
74            inner,
75            sleep: sleep.map(OptionPin::Some).unwrap_or(OptionPin::None),
76        }
77        .await
78    }
79}
80
81/// Parse the timeout header in HeaderMap.
82///
83/// # Return
84///
85///  Ok(Some(duration)) => if parse success.
86///  Ok(None)           => if no success field.
87///  Err(&HeaderValue)  => if parse timeout failed or wrong format.
88fn grpc_timeout_to_duration(
89    headers: &HeaderMap<HeaderValue>,
90) -> Result<Option<Duration>, &HeaderValue> {
91    const SECONDS_HOUR: u64 = 60 * 60;
92    const SECONDS_MINUTE: u64 = 60;
93
94    match headers.get(crate::metadata::GRPC_TIMEOUT_HEADER) {
95        Some(val) => {
96            // parse the value and unit
97            let (timeout_value, timeout_unit) = val
98                .to_str()
99                .map_err(|_| val)
100                .and_then(|s| if s.is_empty() { Err(val) } else { Ok(s) })?
101                .split_at(val.len() - 1);
102            let timeout_value = timeout_value.parse::<u64>().map_err(|_| val)?;
103            // match the unit with Hour | Minute | Second | Milliseconds | Microsecond | Nanosecond
104            let duration = match timeout_unit {
105                "H" => Duration::from_secs(timeout_value * SECONDS_HOUR),
106                "M" => Duration::from_secs(timeout_value * SECONDS_MINUTE),
107                "S" => Duration::from_secs(timeout_value),
108                "m" => Duration::from_millis(timeout_value),
109                "u" => Duration::from_micros(timeout_value),
110                "n" => Duration::from_nanos(timeout_value),
111                _ => return Err(val),
112            };
113            Ok(Some(duration))
114        }
115        None => {
116            tracing::trace!("grpc-timeout header not found");
117            Ok(None)
118        }
119    }
120}
121
122#[pin_project]
123pub struct ResponseFuture<F> {
124    #[pin]
125    inner: F,
126    #[pin]
127    sleep: OptionPin<Sleep>,
128}
129
130#[pin_project(project = OptionPinProj)]
131pub enum OptionPin<T> {
132    Some(#[pin] T),
133    None,
134}
135
136impl<F, R> Future for ResponseFuture<F>
137where
138    F: Future<Output = Result<R, Status>>,
139{
140    type Output = Result<R, Status>;
141
142    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
143        let this = self.project();
144
145        if let Poll::Ready(res) = this.inner.poll(cx) {
146            return Poll::Ready(res);
147        }
148
149        if let OptionPinProj::Some(sleep) = this.sleep.project() {
150            futures_util::ready!(sleep.poll(cx));
151            let err = Status::deadline_exceeded("timeout");
152            return Poll::Ready(Err(err));
153        }
154
155        Poll::Pending
156    }
157}
158
159#[cfg(test)]
160mod tests {
161
162    use super::*;
163    use crate::metadata::GRPC_TIMEOUT_HEADER;
164
165    // init config in testing
166    fn try_set_up(val: Option<&str>) -> Result<Option<Duration>, HeaderValue> {
167        let mut hm = HeaderMap::new();
168        if let Some(v) = val {
169            let hv = HeaderValue::from_str(v).unwrap();
170            hm.insert(GRPC_TIMEOUT_HEADER, hv);
171        };
172
173        grpc_timeout_to_duration(&hm).map_err(|e| e.clone())
174    }
175
176    #[test]
177    fn test_hours() {
178        let parsed_duration = try_set_up(Some("3H")).unwrap().unwrap();
179        assert_eq!(Duration::from_secs(3 * 60 * 60), parsed_duration);
180    }
181
182    #[test]
183    fn test_minutes() {
184        let parsed_duration = try_set_up(Some("1M")).unwrap().unwrap();
185        assert_eq!(Duration::from_secs(60), parsed_duration);
186    }
187
188    #[test]
189    fn test_seconds() {
190        let parsed_duration = try_set_up(Some("42S")).unwrap().unwrap();
191        assert_eq!(Duration::from_secs(42), parsed_duration);
192    }
193
194    #[test]
195    fn test_milliseconds() {
196        let parsed_duration = try_set_up(Some("13m")).unwrap().unwrap();
197        assert_eq!(Duration::from_millis(13), parsed_duration);
198    }
199
200    #[test]
201    fn test_microseconds() {
202        let parsed_duration = try_set_up(Some("2u")).unwrap().unwrap();
203        assert_eq!(Duration::from_micros(2), parsed_duration);
204    }
205
206    #[test]
207    fn test_nanoseconds() {
208        let parsed_duration = try_set_up(Some("82n")).unwrap().unwrap();
209        assert_eq!(Duration::from_nanos(82), parsed_duration);
210    }
211
212    #[test]
213    fn test_corner_cases() {
214        // error postfix
215        let r = HeaderValue::from_str("82f").unwrap();
216        assert_eq!(try_set_up(Some("82f")), Err(r));
217
218        // error digit
219        let r = HeaderValue::from_str("abcH").unwrap();
220        assert_eq!(try_set_up(Some("abcH")), Err(r));
221    }
222}
223
224#[cfg(test)]
225mod tests_insert_and_parse {
226    use super::*;
227
228    #[tokio::test]
229    async fn test_insert_and_parse_metainfo() {
230        use std::time::Duration;
231
232        use http::HeaderValue;
233        use metainfo::{METAINFO, MetaInfo};
234
235        let mi = MetaInfo::new();
236
237        METAINFO
238            .scope(mi.into(), async {
239                // insert a Duration manually
240                METAINFO.with(|mi| {
241                    mi.borrow_mut().insert::<Duration>(Duration::from_secs(10));
242                });
243
244                // verify insertion
245                METAINFO.with(|mi| {
246                    let mi = mi.borrow();
247                    let stored = mi.get::<Duration>().expect("Duration not found");
248                    assert_eq!(*stored, Duration::from_secs(10));
249                });
250
251                // simulate parsing a grpc-timeout header and inserting
252                let mut hm = http::HeaderMap::new();
253                let hv = HeaderValue::from_str("7S").unwrap();
254                hm.insert(crate::metadata::GRPC_TIMEOUT_HEADER, hv);
255
256                // use parser function
257                let parsed = grpc_timeout_to_duration(&hm).expect("Parsing failed");
258                assert_eq!(parsed, Some(Duration::from_secs(7)));
259
260                // insert parsed duration
261                if let Some(dur) = parsed {
262                    METAINFO.with(|mi| {
263                        mi.borrow_mut().insert::<Duration>(dur);
264                    });
265                }
266
267                // check updated duration
268                METAINFO.with(|mi| {
269                    let mi = mi.borrow();
270                    let stored = mi
271                        .get::<Duration>()
272                        .expect("Duration not found after insert");
273                    assert_eq!(*stored, Duration::from_secs(7));
274                });
275            })
276            .await;
277    }
278}
volo_grpc/server/layer/timeout.rs

volo_grpc/server/layer/
timeout.rs
