Skip to main content

volo_grpc/
context.rs

1use std::time::Duration;
2
3use chrono::{DateTime, Local};
4use paste::paste;
5pub use volo::context::*;
6use volo::newtype_impl_context;
7
8use crate::codec::compression::CompressionEncoding;
9
10macro_rules! stat_impl {
11    ($t: ident) => {
12        paste! {
13            #[inline]
14            pub fn $t(&self) -> Option<DateTime<Local>> {
15                self.$t
16            }
17
18            #[doc(hidden)]
19            #[inline]
20            pub fn [<set_$t>](&mut self, t: DateTime<Local>) {
21                self.$t = Some(t)
22            }
23
24            #[inline]
25            pub fn [<record_ $t>](&mut self) {
26                self.$t = Some(Local::now())
27            }
28        }
29    };
30}
31
32#[derive(Debug, Default, Clone)]
33pub struct ClientStats {
34    make_transport_start_at: Option<DateTime<Local>>,
35    make_transport_end_at: Option<DateTime<Local>>,
36}
37
38impl ClientStats {
39    stat_impl!(make_transport_start_at);
40    stat_impl!(make_transport_end_at);
41
42    #[inline]
43    pub fn reset(&mut self) {
44        self.make_transport_start_at = None;
45        self.make_transport_end_at = None;
46    }
47}
48
49#[derive(Debug, Default, Clone)]
50pub struct ServerStats {
51    process_start_at: Option<DateTime<Local>>,
52    process_end_at: Option<DateTime<Local>>,
53}
54
55impl ServerStats {
56    stat_impl!(process_start_at);
57    stat_impl!(process_end_at);
58
59    #[inline]
60    pub fn reset(&mut self) {
61        self.process_start_at = None;
62        self.process_end_at = None;
63    }
64}
65
66#[derive(Debug, Clone, Default)]
67pub struct ClientCxInner {
68    pub stats: ClientStats,
69}
70
71/// A context for client to pass information such as `RpcInfo` and `Config` between middleware
72/// during the rpc call lifecycle.
73pub struct ClientContext(pub(crate) RpcCx<ClientCxInner, Config>);
74
75newtype_impl_context!(ClientContext, Config, 0);
76
77impl ClientContext {
78    pub fn new(ri: RpcInfo<Config>) -> Self {
79        Self(RpcCx::new(
80            ri,
81            ClientCxInner {
82                stats: ClientStats::default(),
83            },
84        ))
85    }
86}
87
88impl Default for ClientContext {
89    fn default() -> Self {
90        Self::new(RpcInfo::with_role(Role::Client))
91    }
92}
93
94impl std::ops::Deref for ClientContext {
95    type Target = RpcCx<ClientCxInner, Config>;
96
97    #[inline]
98    fn deref(&self) -> &Self::Target {
99        &self.0
100    }
101}
102
103impl std::ops::DerefMut for ClientContext {
104    #[inline]
105    fn deref_mut(&mut self) -> &mut Self::Target {
106        &mut self.0
107    }
108}
109
110#[derive(Debug, Clone, Default)]
111pub struct ServerCxInner {
112    pub stats: ServerStats,
113}
114
115/// A context for server to pass information such as `RpcInfo` and `Config` between middleware
116/// during the rpc call lifecycle.
117pub struct ServerContext(pub(crate) RpcCx<ServerCxInner, Config>);
118
119newtype_impl_context!(ServerContext, Config, 0);
120
121impl Default for ServerContext {
122    fn default() -> Self {
123        Self(RpcCx::new(
124            RpcInfo::with_role(Role::Server),
125            Default::default(),
126        ))
127    }
128}
129
130impl std::ops::Deref for ServerContext {
131    type Target = RpcCx<ServerCxInner, Config>;
132
133    #[inline]
134    fn deref(&self) -> &Self::Target {
135        &self.0
136    }
137}
138
139impl std::ops::DerefMut for ServerContext {
140    #[inline]
141    fn deref_mut(&mut self) -> &mut Self::Target {
142        &mut self.0
143    }
144}
145
146const DEFAULT_RPC_TIMEOUT: Duration = Duration::from_secs(1);
147
148#[derive(Default, Debug, Clone)]
149pub struct Config {
150    pub(crate) rpc_timeout: Option<Duration>,
151    /// Amount of time to wait connecting.
152    pub(crate) connect_timeout: Option<Duration>,
153    /// Amount of time to wait reading.
154    pub(crate) read_timeout: Option<Duration>,
155    /// Amount of time to wait reading response.
156    pub(crate) write_timeout: Option<Duration>,
157
158    pub(crate) accept_compressions: Option<Vec<CompressionEncoding>>,
159    pub(crate) send_compressions: Option<Vec<CompressionEncoding>>,
160}
161
162impl Reusable for Config {
163    fn clear(&mut self) {
164        self.rpc_timeout = None;
165        self.connect_timeout = None;
166        self.read_timeout = None;
167        self.write_timeout = None;
168        if let Some(v) = self.accept_compressions.as_mut() {
169            v.clear();
170        }
171        if let Some(v) = self.send_compressions.as_mut() {
172            v.clear();
173        }
174    }
175}
176
177impl Config {
178    pub fn merge(&mut self, other: Self) {
179        if let Some(t) = other.rpc_timeout {
180            self.rpc_timeout = Some(t);
181        }
182        if let Some(t) = other.connect_timeout {
183            self.connect_timeout = Some(t);
184        }
185        if let Some(t) = other.read_timeout {
186            self.read_timeout = Some(t);
187        }
188        if let Some(t) = other.write_timeout {
189            self.write_timeout = Some(t);
190        }
191        if let Some(e) = other.accept_compressions {
192            self.accept_compressions = Some(e);
193        }
194        if let Some(e) = other.send_compressions {
195            self.send_compressions = Some(e);
196        }
197    }
198
199    #[inline]
200    pub fn rpc_timeout(&self) -> Option<Duration> {
201        self.rpc_timeout
202    }
203
204    /// Sets the rpc timeout.
205    ///
206    /// This can be set both by the client builder and the CallOpt.
207    #[inline]
208    pub fn set_rpc_timeout(&mut self, rpc_timeout: Option<Duration>) {
209        self.rpc_timeout = rpc_timeout;
210    }
211
212    #[inline]
213    pub fn rpc_timeout_or_default(&self) -> Duration {
214        self.rpc_timeout.unwrap_or(DEFAULT_RPC_TIMEOUT)
215    }
216}