Skip to main content

temporalio_client/
rpc_options.rs

1use crate::{RetryOptions, request_extensions::RetryConfigForCall};
2use std::time::Duration;
3use tonic::metadata::{
4    AsciiMetadataKey, AsciiMetadataValue, BinaryMetadataKey, BinaryMetadataValue, KeyAndValueRef,
5    MetadataMap,
6};
7
8/// Metadata attached to a single high-level client RPC.
9///
10/// Per-call values take precedence over metadata configured on the connection.
11#[derive(Clone, Debug, Default)]
12pub struct RpcMetadata {
13    inner: MetadataMap,
14}
15
16impl RpcMetadata {
17    /// Create empty RPC metadata.
18    pub fn new() -> Self {
19        Self::default()
20    }
21
22    /// Insert an ASCII metadata value, returning the previous value for the key.
23    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    /// Insert a binary metadata value, returning the previous value for the key.
48    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    /// Get an ASCII metadata value.
69    pub fn get_ascii(&self, key: &str) -> Option<&str> {
70        self.inner.get(key).and_then(|value| value.to_str().ok())
71    }
72
73    /// Get a binary metadata value.
74    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    /// Remove an ASCII metadata value.
82    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    /// Remove a binary metadata value.
92    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    /// Iterate over ASCII metadata values.
102    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    /// Iterate over binary metadata values.
115    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/// An invalid key or value supplied to [`RpcMetadata`].
153#[derive(Debug, thiserror::Error, PartialEq, Eq)]
154#[non_exhaustive]
155pub enum RpcMetadataError {
156    /// The key is not valid ASCII gRPC metadata.
157    #[error("invalid ASCII RPC metadata key: {key}")]
158    InvalidAsciiKey {
159        /// The invalid key.
160        key: String,
161    },
162    /// The value is not valid ASCII gRPC metadata.
163    #[error("invalid ASCII RPC metadata value for key {key}: {value:?}")]
164    InvalidAsciiValue {
165        /// The associated key.
166        key: String,
167        /// The invalid value.
168        value: String,
169    },
170    /// The key is not valid binary gRPC metadata.
171    #[error("invalid binary RPC metadata key: {key}")]
172    InvalidBinaryKey {
173        /// The invalid key.
174        key: String,
175    },
176}
177
178/// Controls applied to a single high-level client RPC.
179#[derive(Clone, Debug, Default)]
180#[non_exhaustive]
181pub struct RpcOptions {
182    /// Metadata attached to the RPC.
183    pub metadata: RpcMetadata,
184    /// Timeout for the RPC, overriding the connection's default deadline.
185    pub timeout: Option<Duration>,
186    /// Retry behavior for the RPC, overriding the connection's retry configuration.
187    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}