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 super::*;
156
157    impl InvocationId {
158        /// Create a new invocation ID from a `&'static str`.
159        pub fn new_from_str(uuid: &'static str) -> Self {
160            InvocationId(HeaderValue::from_static(uuid))
161        }
162    }
163
164    /// A "generator" that returns [`InvocationId`]s from a predefined list.
165    #[derive(Debug)]
166    pub struct PredefinedInvocationIdGenerator {
167        pre_generated_ids: Arc<Mutex<Vec<InvocationId>>>,
168    }
169
170    impl PredefinedInvocationIdGenerator {
171        /// Given a `Vec<InvocationId>`, create a new [`PredefinedInvocationIdGenerator`].
172        pub fn new(mut invocation_ids: Vec<InvocationId>) -> Self {
173            // We're going to pop ids off of the end of the list, so we need to reverse the list or else
174            // we'll be popping the ids in reverse order, confusing the poor test writer.
175            invocation_ids.reverse();
176
177            Self {
178                pre_generated_ids: Arc::new(Mutex::new(invocation_ids)),
179            }
180        }
181    }
182
183    impl InvocationIdGenerator for PredefinedInvocationIdGenerator {
184        fn generate(&self) -> Result<Option<InvocationId>, BoxError> {
185            Ok(Some(
186                self.pre_generated_ids
187                    .lock()
188                    .expect("this will never be under contention")
189                    .pop()
190                    .expect("testers will provide enough invocation IDs"),
191            ))
192        }
193    }
194
195    /// A "generator" that always returns `None`.
196    #[derive(Debug, Default)]
197    pub struct NoInvocationIdGenerator;
198
199    impl NoInvocationIdGenerator {
200        /// Create a new [`NoInvocationIdGenerator`].
201        pub fn new() -> Self {
202            Self
203        }
204    }
205
206    impl InvocationIdGenerator for NoInvocationIdGenerator {
207        fn generate(&self) -> Result<Option<InvocationId>, BoxError> {
208            Ok(None)
209        }
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use aws_smithy_runtime_api::client::interceptors::context::{
216        BeforeTransmitInterceptorContextMut, Input, InterceptorContext,
217    };
218    use aws_smithy_runtime_api::client::interceptors::Intercept;
219    use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
220    use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder;
221    use aws_smithy_types::config_bag::ConfigBag;
222
223    use super::*;
224
225    fn expect_header<'a>(
226        context: &'a BeforeTransmitInterceptorContextMut<'_>,
227        header_name: &str,
228    ) -> &'a str {
229        context.request().headers().get(header_name).unwrap()
230    }
231
232    #[test]
233    fn default_id_generator() {
234        let rc = RuntimeComponentsBuilder::for_tests().build().unwrap();
235        let mut ctx = InterceptorContext::new(Input::doesnt_matter());
236        ctx.enter_serialization_phase();
237        ctx.set_request(HttpRequest::empty());
238        let _ = ctx.take_input();
239        ctx.enter_before_transmit_phase();
240
241        let mut cfg = ConfigBag::base();
242        let interceptor = InvocationIdInterceptor::new();
243        let mut ctx = Into::into(&mut ctx);
244        interceptor
245            .modify_before_retry_loop(&mut ctx, &rc, &mut cfg)
246            .unwrap();
247        interceptor
248            .modify_before_transmit(&mut ctx, &rc, &mut cfg)
249            .unwrap();
250
251        let expected = cfg.load::<InvocationId>().expect("invocation ID was set");
252        let header = expect_header(&ctx, "amz-sdk-invocation-id");
253        assert_eq!(expected.0, header, "the invocation ID in the config bag must match the invocation ID in the request header");
254        // UUID should include 32 chars and 4 dashes
255        assert_eq!(header.len(), 36);
256    }
257
258    #[cfg(feature = "test-util")]
259    #[test]
260    fn custom_id_generator() {
261        use aws_smithy_types::config_bag::Layer;
262        let rc = RuntimeComponentsBuilder::for_tests().build().unwrap();
263        let mut ctx = InterceptorContext::new(Input::doesnt_matter());
264        ctx.enter_serialization_phase();
265        ctx.set_request(HttpRequest::empty());
266        let _ = ctx.take_input();
267        ctx.enter_before_transmit_phase();
268
269        let mut cfg = ConfigBag::base();
270        let mut layer = Layer::new("test");
271        layer.store_put(SharedInvocationIdGenerator::new(
272            PredefinedInvocationIdGenerator::new(vec![InvocationId::new(
273                "the-best-invocation-id".into(),
274            )]),
275        ));
276        cfg.push_layer(layer);
277
278        let interceptor = InvocationIdInterceptor::new();
279        let mut ctx = Into::into(&mut ctx);
280        interceptor
281            .modify_before_retry_loop(&mut ctx, &rc, &mut cfg)
282            .unwrap();
283        interceptor
284            .modify_before_transmit(&mut ctx, &rc, &mut cfg)
285            .unwrap();
286
287        let header = expect_header(&ctx, "amz-sdk-invocation-id");
288        assert_eq!("the-best-invocation-id", header);
289    }
290}
aws_runtime/invocation_id.rs

aws_runtime/
invocation_id.rs
