temporalio_client/
rpc_options.rs1use crate::{RetryOptions, request_extensions::RetryConfigForCall};
2use std::time::Duration;
3use tonic::metadata::{
4 AsciiMetadataKey, AsciiMetadataValue, BinaryMetadataKey, BinaryMetadataValue, KeyAndValueRef,
5 MetadataMap,
6};
7
8#[derive(Clone, Debug, Default)]
12pub struct RpcMetadata {
13 inner: MetadataMap,
14}
15
16impl RpcMetadata {
17 pub fn new() -> Self {
19 Self::default()
20 }
21
22 pub fn insert(
24 &mut self,
25 key: impl Into<String>,
26 value: impl Into<String>,
27 ) -> Result<Option<String>, RpcMetadataError> {
28 let key = key.into();
29 let value = value.into();
30 let parsed_key = key
31 .parse::<AsciiMetadataKey>()
32 .map_err(|_| RpcMetadataError::InvalidAsciiKey { key: key.clone() })?;
33 let parsed_value = value.parse::<AsciiMetadataValue>().map_err(|_| {
34 RpcMetadataError::InvalidAsciiValue {
35 key,
36 value: value.clone(),
37 }
38 })?;
39 Ok(self.inner.insert(parsed_key, parsed_value).map(|previous| {
40 previous
41 .to_str()
42 .expect("ASCII RPC metadata remains valid while stored")
43 .to_owned()
44 }))
45 }
46
47 pub fn insert_binary(
49 &mut self,
50 key: impl Into<String>,
51 value: impl Into<Vec<u8>>,
52 ) -> Result<Option<Vec<u8>>, RpcMetadataError> {
53 let key = key.into();
54 let parsed_key = key
55 .parse::<BinaryMetadataKey>()
56 .map_err(|_| RpcMetadataError::InvalidBinaryKey { key })?;
57 Ok(self
58 .inner
59 .insert_bin(parsed_key, BinaryMetadataValue::from_bytes(&value.into()))
60 .map(|previous| {
61 previous
62 .to_bytes()
63 .expect("binary RPC metadata remains valid while stored")
64 .to_vec()
65 }))
66 }
67
68 pub fn get_ascii(&self, key: &str) -> Option<&str> {
70 self.inner.get(key).and_then(|value| value.to_str().ok())
71 }
72
73 pub fn get_binary(&self, key: &str) -> Option<Vec<u8>> {
75 self.inner
76 .get_bin(key)
77 .and_then(|value| value.to_bytes().ok())
78 .map(|value| value.to_vec())
79 }
80
81 pub fn remove_ascii(&mut self, key: &str) -> Option<String> {
83 self.inner.remove(key).map(|value| {
84 value
85 .to_str()
86 .expect("ASCII RPC metadata remains valid while stored")
87 .to_owned()
88 })
89 }
90
91 pub fn remove_binary(&mut self, key: &str) -> Option<Vec<u8>> {
93 self.inner.remove_bin(key).map(|value| {
94 value
95 .to_bytes()
96 .expect("binary RPC metadata remains valid while stored")
97 .to_vec()
98 })
99 }
100
101 pub fn ascii(&self) -> impl Iterator<Item = (&str, &str)> {
103 self.inner.iter().filter_map(|entry| match entry {
104 KeyAndValueRef::Ascii(key, value) => Some((
105 key.as_str(),
106 value
107 .to_str()
108 .expect("ASCII RPC metadata remains valid while stored"),
109 )),
110 KeyAndValueRef::Binary(_, _) => None,
111 })
112 }
113
114 pub fn binary(&self) -> impl Iterator<Item = (&str, Vec<u8>)> {
116 self.inner.iter().filter_map(|entry| match entry {
117 KeyAndValueRef::Ascii(_, _) => None,
118 KeyAndValueRef::Binary(key, value) => Some((
119 key.as_str(),
120 value
121 .to_bytes()
122 .expect("binary RPC metadata remains valid while stored")
123 .to_vec(),
124 )),
125 })
126 }
127
128 fn apply_to<T>(&self, request: &mut tonic::Request<T>) {
129 for entry in self.inner.iter() {
130 match entry {
131 KeyAndValueRef::Ascii(key, value) => {
132 request.metadata_mut().insert(key.clone(), value.clone());
133 }
134 KeyAndValueRef::Binary(key, value) => {
135 request
136 .metadata_mut()
137 .insert_bin(key.clone(), value.clone());
138 }
139 }
140 }
141 }
142}
143
144impl PartialEq for RpcMetadata {
145 fn eq(&self, other: &Self) -> bool {
146 self.inner.as_ref() == other.inner.as_ref()
147 }
148}
149
150impl Eq for RpcMetadata {}
151
152#[derive(Debug, thiserror::Error, PartialEq, Eq)]
154#[non_exhaustive]
155pub enum RpcMetadataError {
156 #[error("invalid ASCII RPC metadata key: {key}")]
158 InvalidAsciiKey {
159 key: String,
161 },
162 #[error("invalid ASCII RPC metadata value for key {key}: {value:?}")]
164 InvalidAsciiValue {
165 key: String,
167 value: String,
169 },
170 #[error("invalid binary RPC metadata key: {key}")]
172 InvalidBinaryKey {
173 key: String,
175 },
176}
177
178#[derive(Clone, Debug, Default)]
180#[non_exhaustive]
181pub struct RpcOptions {
182 pub metadata: RpcMetadata,
184 pub timeout: Option<Duration>,
186 pub retry_options: Option<RetryOptions>,
188}
189
190impl RpcOptions {
191 pub(crate) fn apply_to<T>(&self, request: &mut tonic::Request<T>) {
192 self.metadata.apply_to(request);
193 if let Some(timeout) = self.timeout {
194 request.set_timeout(timeout);
195 }
196 if let Some(retry_options) = &self.retry_options {
197 request
198 .extensions_mut()
199 .insert(RetryConfigForCall(retry_options.clone()));
200 }
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use super::*;
207
208 #[test]
209 fn metadata_validates_and_exposes_values() {
210 let mut metadata = RpcMetadata::new();
211 assert_eq!(metadata.insert("trace-id", "first").unwrap(), None);
212 assert_eq!(
213 metadata.insert("trace-id", "second").unwrap(),
214 Some("first".to_owned())
215 );
216 assert_eq!(
217 metadata
218 .insert_binary("trace-data-bin", vec![0, 255])
219 .unwrap(),
220 None
221 );
222 assert_eq!(metadata.get_ascii("trace-id"), Some("second"));
223 assert_eq!(metadata.get_binary("trace-data-bin"), Some(vec![0, 255]));
224 assert!(matches!(
225 metadata.insert("bad key", "value"),
226 Err(RpcMetadataError::InvalidAsciiKey { .. })
227 ));
228 assert!(matches!(
229 metadata.insert("valid-key", "bad\nvalue"),
230 Err(RpcMetadataError::InvalidAsciiValue { .. })
231 ));
232 assert!(matches!(
233 metadata.insert_binary("missing-suffix", vec![]),
234 Err(RpcMetadataError::InvalidBinaryKey { .. })
235 ));
236 assert_eq!(metadata.remove_ascii("trace-id"), Some("second".to_owned()));
237 assert_eq!(metadata.remove_binary("trace-data-bin"), Some(vec![0, 255]));
238 }
239}