Skip to main content

ntex_grpc/client/
request.rs

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