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 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 pub fn get_timeout(&self) -> Option<time::Duration> {
36 self.0.timeout.get()
37 }
38
39 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 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 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 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 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 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 let max_size: u128 = 99_999_999; let value = convert(duration).into();
209 if value > max_size {
210 None
211 } else {
212 Some(format!("{value}{unit}"))
213 }
214 }
215
216 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 .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}