Skip to main content

ntex_grpc/client/
request.rs

1use std::{convert::TryFrom, fmt, ops, time};
2
3use ntex_http::{HeaderMap, HeaderName, HeaderValue, error::Error as HttpError};
4
5use crate::{client::Transport, consts, service::MethodDef};
6
7#[derive(Debug)]
8pub struct RequestContext {
9    err: Option<HttpError>,
10    headers: Vec<(HeaderName, HeaderValue)>,
11    timeout: Option<time::Duration>,
12    flags: Flags,
13}
14
15bitflags::bitflags! {
16    #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
17    struct Flags: u8 {
18        const DISCONNECT_ON_DROP = 0b0000_0001;
19    }
20}
21
22impl RequestContext {
23    /// Create new `RequestContext` instance
24    fn new() -> Self {
25        Self {
26            err: None,
27            headers: Vec::new(),
28            timeout: None,
29            flags: Flags::empty(),
30        }
31    }
32
33    /// Get request timeout
34    pub fn get_timeout(&self) -> Option<time::Duration> {
35        self.timeout
36    }
37
38    /// Set the max duration the request is allowed to take.
39    ///
40    /// The duration will be formatted according to [the spec] and use the most precise
41    /// possible.
42    ///
43    /// [the spec]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md
44    pub fn timeout<U>(&mut self, timeout: U) -> &mut Self
45    where
46        time::Duration: From<U>,
47    {
48        let to = timeout.into();
49        self.timeout = Some(to);
50        self.header(consts::GRPC_TIMEOUT, duration_to_grpc_timeout(to));
51        self
52    }
53
54    /// Disconnect connection on request drop
55    pub fn disconnect_on_drop(&mut self) -> &mut Self {
56        self.flags.insert(Flags::DISCONNECT_ON_DROP);
57        self
58    }
59
60    /// Append a header to existing headers.
61    pub fn header<K, V>(&mut self, key: K, value: V) -> &mut Self
62    where
63        HeaderName: TryFrom<K>,
64        HeaderValue: TryFrom<V>,
65        <HeaderName as TryFrom<K>>::Error: Into<HttpError>,
66        <HeaderValue as TryFrom<V>>::Error: Into<HttpError>,
67    {
68        match HeaderName::try_from(key) {
69            Ok(key) => match HeaderValue::try_from(value) {
70                Ok(value) => {
71                    if self.headers.is_empty() {
72                        self.headers.push((key, value))
73                    } else if self.headers[self.headers.len() - 1].0 == key {
74                        let idx = self.headers.len() - 1;
75                        self.headers[idx].1 = value;
76                    } else {
77                        self.headers.push((key, value))
78                    }
79                }
80                Err(e) => self.err = Some(log_error(e)),
81            },
82            Err(e) => self.err = Some(log_error(e)),
83        }
84        self
85    }
86
87    pub(crate) fn headers(&self) -> &[(HeaderName, HeaderValue)] {
88        &self.headers
89    }
90
91    pub(crate) fn get_disconnect_on_drop(&self) -> bool {
92        self.flags.contains(Flags::DISCONNECT_ON_DROP)
93    }
94}
95
96fn log_error<T: Into<HttpError>>(err: T) -> HttpError {
97    let e = err.into();
98    log::error!("Error in Grpc Request {e}");
99    e
100}
101
102pub struct Request<'a, T, M>
103where
104    T: Transport<M>,
105    T: 'a,
106    M: MethodDef,
107{
108    input: &'a M::Input,
109    transport: &'a T,
110    ctx: RequestContext,
111}
112
113impl<'a, T, M> Request<'a, T, M>
114where
115    T: Transport<M>,
116    M: MethodDef,
117{
118    pub fn new(transport: &'a T, input: &'a M::Input) -> Self {
119        Self {
120            input,
121            transport,
122            ctx: RequestContext::new(),
123        }
124    }
125
126    /// Append a header to existing headers.
127    ///
128    /// ```rust
129    /// use ntex::http::{header, Request, Response};
130    ///
131    /// fn index(req: Request) -> Response {
132    ///     Response::Ok()
133    ///         .header("X-TEST", "value")
134    ///         .header(header::CONTENT_TYPE, "application/json")
135    ///         .finish()
136    /// }
137    /// ```
138    pub fn header<K, V>(&mut self, key: K, value: V) -> &mut Self
139    where
140        HeaderName: TryFrom<K>,
141        HeaderValue: TryFrom<V>,
142        <HeaderName as TryFrom<K>>::Error: Into<HttpError>,
143        <HeaderValue as TryFrom<V>>::Error: Into<HttpError>,
144    {
145        self.ctx.header(key, value);
146        self
147    }
148
149    /// Set the max duration the request is allowed to take.
150    ///
151    /// The duration will be formatted according to [the spec] and use the most precise
152    /// possible.
153    ///
154    /// [the spec]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md
155    pub fn timeout<U>(&mut self, timeout: U) -> &mut Self
156    where
157        time::Duration: From<U>,
158    {
159        let to = timeout.into();
160        self.ctx.timeout = Some(to);
161        self.ctx
162            .header(consts::GRPC_TIMEOUT, duration_to_grpc_timeout(to));
163        self
164    }
165
166    /// Send request
167    pub async fn send(self) -> Result<Response<M>, T::Error> {
168        let Request {
169            input,
170            transport,
171            mut ctx,
172        } = self;
173
174        transport.request(input, &mut ctx).await
175    }
176}
177
178fn duration_to_grpc_timeout(duration: time::Duration) -> String {
179    fn try_format<T: Into<u128>>(
180        duration: time::Duration,
181        unit: char,
182        convert: impl FnOnce(time::Duration) -> T,
183    ) -> Option<String> {
184        // The gRPC spec specifies that the timeout most be at most 8 digits. So this is the largest a
185        // value can be before we need to use a bigger unit.
186        let max_size: u128 = 99_999_999; // exactly 8 digits
187
188        let value = convert(duration).into();
189        if value > max_size {
190            None
191        } else {
192            Some(format!("{value}{unit}"))
193        }
194    }
195
196    // pick the most precise unit that is less than or equal to 8 digits as per the gRPC spec
197    try_format(duration, 'n', |d| d.as_nanos())
198        .or_else(|| try_format(duration, 'u', |d| d.as_micros()))
199        .or_else(|| try_format(duration, 'm', |d| d.as_millis()))
200        .or_else(|| try_format(duration, 'S', |d| d.as_secs()))
201        .or_else(|| try_format(duration, 'M', |d| d.as_secs() / 60))
202        .or_else(|| {
203            try_format(duration, 'H', |d| {
204                let minutes = d.as_secs() / 60;
205                minutes / 60
206            })
207        })
208        // duration has to be more than 11_415 years for this to happen
209        .expect("duration is unrealistically large")
210}
211
212pub struct Response<T: MethodDef> {
213    pub output: T::Output,
214    pub headers: HeaderMap,
215    pub trailers: HeaderMap,
216    pub req_size: usize,
217    pub res_size: usize,
218}
219
220impl<T: MethodDef> Response<T> {
221    #[inline]
222    pub fn headers(&self) -> &HeaderMap {
223        &self.headers
224    }
225
226    #[inline]
227    pub fn trailers(&self) -> &HeaderMap {
228        &self.trailers
229    }
230
231    #[inline]
232    pub fn into_inner(self) -> T::Output {
233        self.output
234    }
235
236    #[inline]
237    pub fn into_parts(self) -> (T::Output, HeaderMap, HeaderMap) {
238        (self.output, self.headers, self.trailers)
239    }
240}
241
242impl<T: MethodDef> ops::Deref for Response<T> {
243    type Target = T::Output;
244
245    fn deref(&self) -> &Self::Target {
246        &self.output
247    }
248}
249
250impl<T: MethodDef> ops::DerefMut for Response<T> {
251    fn deref_mut(&mut self) -> &mut Self::Target {
252        &mut self.output
253    }
254}
255
256impl<T: MethodDef> fmt::Debug for Response<T>
257where
258    T::Output: fmt::Debug,
259{
260    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
261        f.debug_struct(format!("ResponseFor<{}>", T::NAME).as_str())
262            .field("output", &self.output)
263            .field("headers", &self.headers)
264            .field("translers", &self.headers)
265            .finish()
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn duration_to_grpc_timeout_less_than_second() {
275        let timeout = time::Duration::from_millis(500);
276        let value = duration_to_grpc_timeout(timeout);
277        assert_eq!(value, format!("{}u", timeout.as_micros()));
278
279        let timeout = time::Duration::from_secs(30);
280        let value = duration_to_grpc_timeout(timeout);
281        assert_eq!(value, format!("{}u", timeout.as_micros()));
282
283        let one_hour = time::Duration::from_secs(60 * 60);
284        let value = duration_to_grpc_timeout(one_hour);
285        assert_eq!(value, format!("{}m", one_hour.as_millis()));
286    }
287}