aws_smithy_runtime_api/client/
connection.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Types related to connection monitoring and management.
7
8mod establishment;
9
10pub use establishment::{ConnectionEstablishmentMetadata, ConnectionEstablishmentMetadataBuilder};
11
12use aws_smithy_types::config_bag::{Storable, StoreReplace};
13use std::fmt;
14use std::net::SocketAddr;
15use std::sync::{Arc, Mutex};
16
17/// Opaque identifier for a physical HTTP connection.
18///
19/// The assigning HTTP client defines the scope of an ID. IDs from different
20/// clients or pools are not comparable.
21#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
22pub struct ConnectionId(u64);
23
24impl ConnectionId {
25    /// Creates an ID from a value assigned by an HTTP client.
26    pub const fn new(value: u64) -> Self {
27        Self(value)
28    }
29}
30
31impl fmt::Display for ConnectionId {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        self.0.fmt(f)
34    }
35}
36
37/// Metadata that tracks the state of an active connection.
38#[derive(Clone)]
39pub struct ConnectionMetadata {
40    is_proxied: bool,
41    remote_addr: Option<SocketAddr>,
42    local_addr: Option<SocketAddr>,
43    poison_fn: Arc<dyn Fn() + Send + Sync>,
44    connection_id: Option<ConnectionId>,
45    establishment: Option<ConnectionEstablishmentMetadata>,
46}
47
48impl ConnectionMetadata {
49    /// Poison this connection, ensuring that it won't be reused.
50    pub fn poison(&self) {
51        tracing::debug!(
52            see_for_more_info = "https://smithy-lang.github.io/smithy-rs/design/client/detailed_error_explanations.html",
53            "Connection encountered an issue and should not be re-used. Marking it for closure"
54        );
55        (self.poison_fn)()
56    }
57
58    /// Create a new [`ConnectionMetadata`].
59    #[deprecated(
60        since = "1.1.0",
61        note = "`ConnectionMetadata::new` is deprecated in favour of `ConnectionMetadata::builder`."
62    )]
63    pub fn new(
64        is_proxied: bool,
65        remote_addr: Option<SocketAddr>,
66        poison: impl Fn() + Send + Sync + 'static,
67    ) -> Self {
68        Self {
69            is_proxied,
70            remote_addr,
71            // need to use builder to set this field
72            local_addr: None,
73            poison_fn: Arc::new(poison),
74            connection_id: None,
75            establishment: None,
76        }
77    }
78
79    /// Builder for this connection metadata
80    pub fn builder() -> ConnectionMetadataBuilder {
81        ConnectionMetadataBuilder::new()
82    }
83
84    /// Get the remote address for this connection, if one is set.
85    pub fn remote_addr(&self) -> Option<SocketAddr> {
86        self.remote_addr
87    }
88
89    /// Get the local address for this connection, if one is set.
90    pub fn local_addr(&self) -> Option<SocketAddr> {
91        self.local_addr
92    }
93
94    /// Get the ID assigned to this connection, if the HTTP client provides one.
95    ///
96    /// Clients that do not track physical connection identity leave this unset.
97    pub fn connection_id(&self) -> Option<ConnectionId> {
98        self.connection_id
99    }
100
101    /// Get measurements from the establishment that installed this connection.
102    ///
103    /// Clients that do not expose successful establishment measurements leave
104    /// this unset.
105    pub fn establishment(&self) -> Option<&ConnectionEstablishmentMetadata> {
106        self.establishment.as_ref()
107    }
108}
109
110impl fmt::Debug for ConnectionMetadata {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        f.debug_struct("SmithyConnection")
113            .field("is_proxied", &self.is_proxied)
114            .field("remote_addr", &self.remote_addr)
115            .field("local_addr", &self.local_addr)
116            .field("connection_id", &self.connection_id)
117            .field("establishment", &self.establishment)
118            .finish()
119    }
120}
121
122/// Builder type that is used to construct a [`ConnectionMetadata`] value.
123#[derive(Default)]
124pub struct ConnectionMetadataBuilder {
125    is_proxied: Option<bool>,
126    remote_addr: Option<SocketAddr>,
127    local_addr: Option<SocketAddr>,
128    poison_fn: Option<Arc<dyn Fn() + Send + Sync>>,
129    connection_id: Option<ConnectionId>,
130    establishment: Option<ConnectionEstablishmentMetadata>,
131}
132
133impl fmt::Debug for ConnectionMetadataBuilder {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        f.debug_struct("ConnectionMetadataBuilder")
136            .field("is_proxied", &self.is_proxied)
137            .field("remote_addr", &self.remote_addr)
138            .field("local_addr", &self.local_addr)
139            .field("connection_id", &self.connection_id)
140            .field("establishment", &self.establishment)
141            .finish()
142    }
143}
144
145impl ConnectionMetadataBuilder {
146    /// Creates a new builder.
147    pub fn new() -> Self {
148        Self::default()
149    }
150
151    /// Set whether or not the associated connection is to an HTTP proxy.
152    pub fn proxied(mut self, proxied: bool) -> Self {
153        self.set_proxied(Some(proxied));
154        self
155    }
156
157    /// Set whether or not the associated connection is to an HTTP proxy.
158    pub fn set_proxied(&mut self, proxied: Option<bool>) -> &mut Self {
159        self.is_proxied = proxied;
160        self
161    }
162
163    /// Set the remote address of the connection used.
164    pub fn remote_addr(mut self, remote_addr: SocketAddr) -> Self {
165        self.set_remote_addr(Some(remote_addr));
166        self
167    }
168
169    /// Set the remote address of the connection used.
170    pub fn set_remote_addr(&mut self, remote_addr: Option<SocketAddr>) -> &mut Self {
171        self.remote_addr = remote_addr;
172        self
173    }
174
175    /// Set the local address of the connection used.
176    pub fn local_addr(mut self, local_addr: SocketAddr) -> Self {
177        self.set_local_addr(Some(local_addr));
178        self
179    }
180
181    /// Set the local address of the connection used.
182    pub fn set_local_addr(&mut self, local_addr: Option<SocketAddr>) -> &mut Self {
183        self.local_addr = local_addr;
184        self
185    }
186
187    /// Set the [`ConnectionId`] assigned by the HTTP client.
188    pub fn connection_id(mut self, connection_id: ConnectionId) -> Self {
189        self.set_connection_id(Some(connection_id));
190        self
191    }
192
193    /// Set the [`ConnectionId`] assigned by the HTTP client.
194    pub fn set_connection_id(&mut self, connection_id: Option<ConnectionId>) -> &mut Self {
195        self.connection_id = connection_id;
196        self
197    }
198
199    /// Sets measurements from the establishment that installed this connection.
200    pub fn establishment(mut self, establishment: ConnectionEstablishmentMetadata) -> Self {
201        self.set_establishment(Some(establishment));
202        self
203    }
204
205    /// Sets measurements from the establishment that installed this connection.
206    pub fn set_establishment(
207        &mut self,
208        establishment: Option<ConnectionEstablishmentMetadata>,
209    ) -> &mut Self {
210        self.establishment = establishment;
211        self
212    }
213
214    /// Set a closure which will poison the associated connection.
215    ///
216    /// A poisoned connection will not be reused for subsequent requests by the pool
217    pub fn poison_fn(mut self, poison_fn: impl Fn() + Send + Sync + 'static) -> Self {
218        self.set_poison_fn(Some(poison_fn));
219        self
220    }
221
222    /// Set a closure which will poison the associated connection.
223    ///
224    /// A poisoned connection will not be reused for subsequent requests by the pool
225    pub fn set_poison_fn(
226        &mut self,
227        poison_fn: Option<impl Fn() + Send + Sync + 'static>,
228    ) -> &mut Self {
229        self.poison_fn =
230            poison_fn.map(|poison_fn| Arc::new(poison_fn) as Arc<dyn Fn() + Send + Sync>);
231        self
232    }
233
234    /// Build a [`ConnectionMetadata`] value.
235    ///
236    /// # Panics
237    ///
238    /// If either the `is_proxied` or `poison_fn` has not been set, then this method will panic
239    pub fn build(self) -> ConnectionMetadata {
240        ConnectionMetadata {
241            is_proxied: self
242                .is_proxied
243                .expect("is_proxied should be set for ConnectionMetadata"),
244            remote_addr: self.remote_addr,
245            local_addr: self.local_addr,
246            poison_fn: self
247                .poison_fn
248                .expect("poison_fn should be set for ConnectionMetadata"),
249            connection_id: self.connection_id,
250            establishment: self.establishment,
251        }
252    }
253}
254
255type LoaderFn = dyn Fn() -> Option<ConnectionMetadata> + Send + Sync;
256
257/// State for a middleware that will monitor and manage connections.
258#[derive(Clone, Default)]
259pub struct CaptureSmithyConnection {
260    loader: Arc<Mutex<Option<Box<LoaderFn>>>>,
261}
262
263impl CaptureSmithyConnection {
264    /// Create a new connection monitor.
265    pub fn new() -> Self {
266        Self {
267            loader: Default::default(),
268        }
269    }
270
271    /// Set the retriever that will capture the `hyper` connection.
272    pub fn set_connection_retriever<F>(&self, f: F)
273    where
274        F: Fn() -> Option<ConnectionMetadata> + Send + Sync + 'static,
275    {
276        *self.loader.lock().unwrap() = Some(Box::new(f));
277    }
278
279    /// Get the associated connection metadata.
280    pub fn get(&self) -> Option<ConnectionMetadata> {
281        match self.loader.lock().unwrap().as_ref() {
282            Some(loader) => loader(),
283            None => {
284                tracing::debug!("no loader was set on the CaptureSmithyConnection");
285                None
286            }
287        }
288    }
289}
290
291impl fmt::Debug for CaptureSmithyConnection {
292    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
293        write!(f, "CaptureSmithyConnection")
294    }
295}
296
297impl Storable for CaptureSmithyConnection {
298    type Storer = StoreReplace<Self>;
299}
300
301#[cfg(test)]
302mod tests {
303    use std::{
304        net::{IpAddr, Ipv6Addr},
305        sync::Mutex,
306        time::Duration,
307    };
308
309    use super::*;
310
311    const TEST_SOCKET_ADDR: SocketAddr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 100);
312
313    #[test]
314    #[should_panic]
315    fn builder_panic_missing_proxied() {
316        ConnectionMetadataBuilder::new()
317            .poison_fn(|| {})
318            .local_addr(TEST_SOCKET_ADDR)
319            .remote_addr(TEST_SOCKET_ADDR)
320            .build();
321    }
322
323    #[test]
324    #[should_panic]
325    fn builder_panic_missing_poison_fn() {
326        ConnectionMetadataBuilder::new()
327            .proxied(true)
328            .local_addr(TEST_SOCKET_ADDR)
329            .remote_addr(TEST_SOCKET_ADDR)
330            .build();
331    }
332
333    #[test]
334    fn builder_all_fields_successful() {
335        let mutable_flag = Arc::new(Mutex::new(false));
336        let establishment = ConnectionEstablishmentMetadata::builder()
337            .total_duration(Duration::from_millis(13))
338            .transport_duration(Duration::from_millis(8))
339            .protocol_handshake_duration(Duration::from_millis(5))
340            .dns_duration(Duration::from_millis(1))
341            .socket_connect_duration(Duration::from_millis(2))
342            .proxy_duration(Duration::from_millis(3))
343            .tls_duration(Duration::from_millis(4))
344            .build();
345
346        let connection_metadata = ConnectionMetadataBuilder::new()
347            .proxied(true)
348            .local_addr(TEST_SOCKET_ADDR)
349            .remote_addr(TEST_SOCKET_ADDR)
350            .connection_id(ConnectionId::new(17))
351            .establishment(establishment.clone())
352            .poison_fn({
353                let mutable_flag = Arc::clone(&mutable_flag);
354                move || {
355                    let mut guard = mutable_flag.lock().unwrap();
356                    *guard = !*guard;
357                }
358            })
359            .build();
360
361        assert!(connection_metadata.is_proxied);
362        assert_eq!(connection_metadata.remote_addr(), Some(TEST_SOCKET_ADDR));
363        assert_eq!(connection_metadata.local_addr(), Some(TEST_SOCKET_ADDR));
364        assert_eq!(
365            connection_metadata.connection_id(),
366            Some(ConnectionId::new(17))
367        );
368        assert_eq!(connection_metadata.establishment(), Some(&establishment));
369        assert_eq!(establishment.total_duration(), Duration::from_millis(13));
370        assert_eq!(establishment.transport_duration(), Duration::from_millis(8));
371        assert_eq!(
372            establishment.protocol_handshake_duration(),
373            Some(Duration::from_millis(5))
374        );
375        assert_eq!(establishment.dns_duration(), Some(Duration::from_millis(1)));
376        assert_eq!(
377            establishment.socket_connect_duration(),
378            Some(Duration::from_millis(2))
379        );
380        assert_eq!(
381            establishment.proxy_duration(),
382            Some(Duration::from_millis(3))
383        );
384        assert_eq!(establishment.tls_duration(), Some(Duration::from_millis(4)));
385        assert_eq!("17", ConnectionId::new(17).to_string());
386        assert!(!(*mutable_flag.lock().unwrap()));
387        connection_metadata.poison();
388        assert!(*mutable_flag.lock().unwrap());
389    }
390
391    #[test]
392    fn builder_optional_fields_translate() {
393        let metadata1 = ConnectionMetadataBuilder::new()
394            .proxied(true)
395            .poison_fn(|| {})
396            .build();
397
398        assert_eq!(metadata1.local_addr(), None);
399        assert_eq!(metadata1.remote_addr(), None);
400        assert_eq!(metadata1.connection_id(), None);
401        assert_eq!(metadata1.establishment(), None);
402
403        let metadata2 = ConnectionMetadataBuilder::new()
404            .proxied(true)
405            .poison_fn(|| {})
406            .local_addr(TEST_SOCKET_ADDR)
407            .build();
408
409        assert_eq!(metadata2.local_addr(), Some(TEST_SOCKET_ADDR));
410        assert_eq!(metadata2.remote_addr(), None);
411
412        let metadata3 = ConnectionMetadataBuilder::new()
413            .proxied(true)
414            .poison_fn(|| {})
415            .remote_addr(TEST_SOCKET_ADDR)
416            .build();
417
418        assert_eq!(metadata3.local_addr(), None);
419        assert_eq!(metadata3.remote_addr(), Some(TEST_SOCKET_ADDR));
420    }
421
422    #[test]
423    #[allow(clippy::redundant_clone)]
424    fn retrieve_connection_metadata() {
425        let retriever = CaptureSmithyConnection::new();
426        let retriever_clone = retriever.clone();
427        assert!(retriever.get().is_none());
428        retriever.set_connection_retriever(|| {
429            Some(
430                ConnectionMetadata::builder()
431                    .proxied(true)
432                    .poison_fn(|| {})
433                    .build(),
434            )
435        });
436
437        assert!(retriever.get().is_some());
438        assert!(retriever_clone.get().is_some());
439    }
440}
aws_smithy_runtime_api/client/connection.rs

aws_smithy_runtime_api/client/
connection.rs
