aws_runtime/
invocation_id.rs

Help
1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use std::fmt::Debug;
7use std::sync::{Arc, Mutex};
8
9use fastrand::Rng;
10use http_02x::{HeaderName, HeaderValue};
11
12use aws_smithy_runtime_api::box_error::BoxError;
13use aws_smithy_runtime_api::client::interceptors::context::BeforeTransmitInterceptorContextMut;
14use aws_smithy_runtime_api::client::interceptors::Intercept;
15use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
16use aws_smithy_types::config_bag::{ConfigBag, Storable, StoreReplace};
17#[cfg(feature = "test-util")]
18pub use test_util::{NoInvocationIdGenerator, PredefinedInvocationIdGenerator};
19
20#[allow(clippy::declare_interior_mutable_const)] // we will never mutate this
21const AMZ_SDK_INVOCATION_ID: HeaderName = HeaderName::from_static("amz-sdk-invocation-id");
22
23/// A generator for returning new invocation IDs on demand.
24pub trait InvocationIdGenerator: Debug + Send + Sync {
25    /// Call this function to receive a new [`InvocationId`] or an error explaining why one couldn't
26    /// be provided.
27    fn generate(&self) -> Result<Option<InvocationId>, BoxError>;
28}
29
30/// Dynamic dispatch implementation of [`InvocationIdGenerator`]
31#[derive(Clone, Debug)]
32pub struct SharedInvocationIdGenerator(Arc<dyn InvocationIdGenerator>);
33
34impl SharedInvocationIdGenerator {
35    /// Creates a new [`SharedInvocationIdGenerator`].
36    pub fn new(gen: impl InvocationIdGenerator + 'static) -> Self {
37        Self(Arc::new(gen))
38    }
39}
40
41impl InvocationIdGenerator for SharedInvocationIdGenerator {
42    fn generate(&self) -> Result<Option<InvocationId>, BoxError> {
43        self.0.generate()
44    }
45}
46
47impl Storable for SharedInvocationIdGenerator {
48    type Storer = StoreReplace<Self>;
49}
50
51/// An invocation ID generator that uses random UUIDs for the invocation ID.
52#[derive(Debug, Default)]
53pub struct DefaultInvocationIdGenerator {
54    rng: Mutex<Rng>,
55}
56
57impl DefaultInvocationIdGenerator {
58    /// Creates a new [`DefaultInvocationIdGenerator`].
59    pub fn new() -> Self {
60        Default::default()
61    }
62
63    /// Creates a [`DefaultInvocationIdGenerator`] with the given seed.
64    pub fn with_seed(seed: u64) -> Self {
65        Self {
66            rng: Mutex::new(Rng::with_seed(seed)),
67        }
68    }
69}
70
71impl InvocationIdGenerator for DefaultInvocationIdGenerator {
72    fn generate(&self) -> Result<Option<InvocationId>, BoxError> {
73        let mut rng = self.rng.lock().unwrap();
74        let mut random_bytes = [0u8; 16];
75        rng.fill(&mut random_bytes);
76
77        let id = uuid::Builder::from_random_bytes(random_bytes).into_uuid();
78        Ok(Some(InvocationId::new(id.to_string())))
79    }
80}
81
82/// This interceptor generates a UUID and attaches it to all request attempts made as part of this operation.
83#[non_exhaustive]
84#[derive(Debug, Default)]
85pub struct InvocationIdInterceptor {
86    default: DefaultInvocationIdGenerator,
87}
88
89impl InvocationIdInterceptor {
90    /// Creates a new `InvocationIdInterceptor`
91    pub fn new() -> Self {
92        Self::default()
93    }
94}
95
96impl Intercept for InvocationIdInterceptor {
97    fn name(&self) -> &'static str {
98        "InvocationIdInterceptor"
99    }
100
101    fn modify_before_retry_loop(
102        &self,
103        _ctx: &mut BeforeTransmitInterceptorContextMut<'_>,
104        _runtime_components: &RuntimeComponents,
105        cfg: &mut ConfigBag,
106    ) -> Result<(), BoxError> {
107        let gen = cfg
108            .load::<SharedInvocationIdGenerator>()
109            .map(|gen| gen as &dyn InvocationIdGenerator)
110            .unwrap_or(&self.default);
111        if let Some(id) = gen.generate()? {
112            cfg.interceptor_state().store_put::<InvocationId>(id);
113        }
114
115        Ok(())
116    }
117
118    fn modify_before_transmit(
119        &self,
120        ctx: &mut BeforeTransmitInterceptorContextMut<'_>,
121        _runtime_components: &RuntimeComponents,
122        cfg: &mut ConfigBag,
123    ) -> Result<(), BoxError> {
124        let headers = ctx.request_mut().headers_mut();
125        if let Some(id) = cfg.load::<InvocationId>() {
126            headers.append(AMZ_SDK_INVOCATION_ID, id.0.clone());
127        }
128        Ok(())
129    }
130}
131
132/// InvocationId provides a consistent ID across retries
133#[derive(Debug, Clone, PartialEq, Eq)]
134pub struct InvocationId(HeaderValue);
135
136impl InvocationId {
137    /// Create an invocation ID with the given value.
138    ///
139    /// # Panics
140    /// This constructor will panic if the given invocation ID is not a valid HTTP header value.
141    pub fn new(invocation_id: String) -> Self {
142        Self(
143            HeaderValue::try_from(invocation_id)
144                .expect("invocation ID must be a valid HTTP header value"),
145        )
146    }
147}
148
149impl Storable for InvocationId {
150    type Storer = StoreReplace<Self>;
151}
152
153#[cfg(feature = "test-util")]
154mod test_util {
155    use std::sync::{Arc, Mutex};
156
157    use super::*;
158
159    impl InvocationId {
160        /// Create a new invocation ID from a `&'static str`.
161        pub fn new_from_str(uuid: &'static str) -> Self {
162            InvocationId(HeaderValue::from_static(uuid))
163        }
164    }
165
166    /// A "generator" that returns [`InvocationId`]s from a predefined list.
167    #[derive(Debug)]
168    pub struct PredefinedInvocationIdGenerator {
169        pre_generated_ids: Arc<Mutex<Vec<InvocationId>>>,
170    }
171
172    impl PredefinedInvocationIdGenerator {
173        /// Given a `Vec<InvocationId>`, create a new [`PredefinedInvocationIdGenerator`].
174        pub fn new(mut invocation_ids: Vec<InvocationId>) -> Self {
175            // We're going to pop ids off of the end of the list, so we need to reverse the list or else
176            // we'll be popping the ids in reverse order, confusing the poor test writer.
177            invocation_ids.reverse();
178
179            Self {
180                pre_generated_ids: Arc::new(Mutex::new(invocation_ids)),
181            }
182        }
183    }
184
185    impl InvocationIdGenerator for PredefinedInvocationIdGenerator {
186        fn generate(&self) -> Result<Option<InvocationId>, BoxError> {
187            Ok(Some(
188                self.pre_generated_ids
189                    .lock()
190                    .expect("this will never be under contention")
191                    .pop()
192                    .expect("testers will provide enough invocation IDs"),
193            ))
194        }
195    }
196
197    /// A "generator" that always returns `None`.
198    #[derive(Debug, Default)]
199    pub struct NoInvocationIdGenerator;
200
201    impl NoInvocationIdGenerator {
202        /// Create a new [`NoInvocationIdGenerator`].
203        pub fn new() -> Self {
204            Self
205        }
206    }
207
208    impl InvocationIdGenerator for NoInvocationIdGenerator {
209        fn generate(&self) -> Result<Option<InvocationId>, BoxError> {
210            Ok(None)
211        }
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use aws_smithy_runtime_api::client::interceptors::context::{
218        BeforeTransmitInterceptorContextMut, Input, InterceptorContext,
219    };
220    use aws_smithy_runtime_api::client::interceptors::Intercept;
221    use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
222    use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder;
223    use aws_smithy_types::config_bag::ConfigBag;
224
225    use super::*;
226
227    fn expect_header<'a>(
228        context: &'a BeforeTransmitInterceptorContextMut<'_>,
229        header_name: &str,
230    ) -> &'a str {
231        context.request().headers().get(header_name).unwrap()
232    }
233
234    #[test]
235    fn default_id_generator() {
236        let rc = RuntimeComponentsBuilder::for_tests().build().unwrap();
237        let mut ctx = InterceptorContext::new(Input::doesnt_matter());
238        ctx.enter_serialization_phase();
239        ctx.set_request(HttpRequest::empty());
240        let _ = ctx.take_input();
241        ctx.enter_before_transmit_phase();
242
243        let mut cfg = ConfigBag::base();
244        let interceptor = InvocationIdInterceptor::new();
245        let mut ctx = Into::into(&mut ctx);
246        interceptor
247            .modify_before_retry_loop(&mut ctx, &rc, &mut cfg)
248            .unwrap();
249        interceptor
250            .modify_before_transmit(&mut ctx, &rc, &mut cfg)
251            .unwrap();
252
253        let expected = cfg.load::<InvocationId>().expect("invocation ID was set");
254        let header = expect_header(&ctx, "amz-sdk-invocation-id");
255        assert_eq!(expected.0, header, "the invocation ID in the config bag must match the invocation ID in the request header");
256        // UUID should include 32 chars and 4 dashes
257        assert_eq!(header.len(), 36);
258    }
259
260    #[cfg(feature = "test-util")]
261    #[test]
262    fn custom_id_generator() {
263        use aws_smithy_types::config_bag::Layer;
264        let rc = RuntimeComponentsBuilder::for_tests().build().unwrap();
265        let mut ctx = InterceptorContext::new(Input::doesnt_matter());
266        ctx.enter_serialization_phase();
267        ctx.set_request(HttpRequest::empty());
268        let _ = ctx.take_input();
269        ctx.enter_before_transmit_phase();
270
271        let mut cfg = ConfigBag::base();
272        let mut layer = Layer::new("test");
273        layer.store_put(SharedInvocationIdGenerator::new(
274            PredefinedInvocationIdGenerator::new(vec![InvocationId::new(
275                "the-best-invocation-id".into(),
276            )]),
277        ));
278        cfg.push_layer(layer);
279
280        let interceptor = InvocationIdInterceptor::new();
281        let mut ctx = Into::into(&mut ctx);
282        interceptor
283            .modify_before_retry_loop(&mut ctx, &rc, &mut cfg)
284            .unwrap();
285        interceptor
286            .modify_before_transmit(&mut ctx, &rc, &mut cfg)
287            .unwrap();
288
289        let header = expect_header(&ctx, "amz-sdk-invocation-id");
290        assert_eq!("the-best-invocation-id", header);
291    }
292}
aws_runtime/invocation_id.rs

aws_runtime/
invocation_id.rs
