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, PartialEq, bon::Builder)]
180#[non_exhaustive]
181pub struct RpcOptions {
182    /// Metadata attached to the RPC.
183    #[builder(default)]
184    pub metadata: RpcMetadata,
185    /// Timeout for the RPC, overriding the connection's default deadline.
186    pub timeout: Option<Duration>,
187    /// Retry behavior for the RPC, overriding the connection's retry configuration.
188    pub retry_options: Option<RetryOptions>,
189}
190
191impl Default for RpcOptions {
192    fn default() -> Self {
193        Self::builder().build()
194    }
195}
196
197impl RpcOptions {
198    pub(crate) fn apply_to<T>(&self, request: &mut tonic::Request<T>) {
199        self.metadata.apply_to(request);
200        if let Some(timeout) = self.timeout {
201            request.set_timeout(timeout);
202        }
203        if let Some(retry_options) = &self.retry_options {
204            request
205                .extensions_mut()
206                .insert(RetryConfigForCall(retry_options.clone()));
207        }
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn metadata_validates_and_exposes_values() {
217        let mut metadata = RpcMetadata::new();
218        assert_eq!(metadata.insert("trace-id", "first").unwrap(), None);
219        assert_eq!(
220            metadata.insert("trace-id", "second").unwrap(),
221            Some("first".to_owned())
222        );
223        assert_eq!(
224            metadata
225                .insert_binary("trace-data-bin", vec![0, 255])
226                .unwrap(),
227            None
228        );
229        assert_eq!(metadata.get_ascii("trace-id"), Some("second"));
230        assert_eq!(metadata.get_binary("trace-data-bin"), Some(vec![0, 255]));
231        assert!(matches!(
232            metadata.insert("bad key", "value"),
233            Err(RpcMetadataError::InvalidAsciiKey { .. })
234        ));
235        assert!(matches!(
236            metadata.insert("valid-key", "bad\nvalue"),
237            Err(RpcMetadataError::InvalidAsciiValue { .. })
238        ));
239        assert!(matches!(
240            metadata.insert_binary("missing-suffix", vec![]),
241            Err(RpcMetadataError::InvalidBinaryKey { .. })
242        ));
243        assert_eq!(metadata.remove_ascii("trace-id"), Some("second".to_owned()));
244        assert_eq!(metadata.remove_binary("trace-data-bin"), Some(vec![0, 255]));
245    }
246}