aws_smithy_runtime_api/client/
connection.rs

Help
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
8use std::fmt::{Debug, Formatter};
9use std::net::SocketAddr;
10use std::sync::Arc;
11
12/// Metadata that tracks the state of an active connection.
13#[derive(Clone)]
14pub struct ConnectionMetadata {
15    is_proxied: bool,
16    remote_addr: Option<SocketAddr>,
17    local_addr: Option<SocketAddr>,
18    poison_fn: Arc<dyn Fn() + Send + Sync>,
19}
20
21impl ConnectionMetadata {
22    /// Poison this connection, ensuring that it won't be reused.
23    pub fn poison(&self) {
24        tracing::info!(
25            see_for_more_info = "https://smithy-lang.github.io/smithy-rs/design/client/detailed_error_explanations.html",
26            "Connection encountered an issue and should not be re-used. Marking it for closure"
27        );
28        (self.poison_fn)()
29    }
30
31    /// Create a new [`ConnectionMetadata`].
32    #[deprecated(
33        since = "1.1.0",
34        note = "`ConnectionMetadata::new` is deprecated in favour of `ConnectionMetadata::builder`."
35    )]
36    pub fn new(
37        is_proxied: bool,
38        remote_addr: Option<SocketAddr>,
39        poison: impl Fn() + Send + Sync + 'static,
40    ) -> Self {
41        Self {
42            is_proxied,
43            remote_addr,
44            // need to use builder to set this field
45            local_addr: None,
46            poison_fn: Arc::new(poison),
47        }
48    }
49
50    /// Builder for this connection metadata
51    pub fn builder() -> ConnectionMetadataBuilder {
52        ConnectionMetadataBuilder::new()
53    }
54
55    /// Get the remote address for this connection, if one is set.
56    pub fn remote_addr(&self) -> Option<SocketAddr> {
57        self.remote_addr
58    }
59
60    /// Get the local address for this connection, if one is set.
61    pub fn local_addr(&self) -> Option<SocketAddr> {
62        self.local_addr
63    }
64}
65
66impl Debug for ConnectionMetadata {
67    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
68        f.debug_struct("SmithyConnection")
69            .field("is_proxied", &self.is_proxied)
70            .field("remote_addr", &self.remote_addr)
71            .field("local_addr", &self.local_addr)
72            .finish()
73    }
74}
75
76/// Builder type that is used to construct a [`ConnectionMetadata`] value.
77#[derive(Default)]
78pub struct ConnectionMetadataBuilder {
79    is_proxied: Option<bool>,
80    remote_addr: Option<SocketAddr>,
81    local_addr: Option<SocketAddr>,
82    poison_fn: Option<Arc<dyn Fn() + Send + Sync>>,
83}
84
85impl Debug for ConnectionMetadataBuilder {
86    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
87        f.debug_struct("ConnectionMetadataBuilder")
88            .field("is_proxied", &self.is_proxied)
89            .field("remote_addr", &self.remote_addr)
90            .field("local_addr", &self.local_addr)
91            .finish()
92    }
93}
94
95impl ConnectionMetadataBuilder {
96    /// Creates a new builder.
97    pub fn new() -> Self {
98        Self::default()
99    }
100
101    /// Set whether or not the associated connection is to an HTTP proxy.
102    pub fn proxied(mut self, proxied: bool) -> Self {
103        self.set_proxied(Some(proxied));
104        self
105    }
106
107    /// Set whether or not the associated connection is to an HTTP proxy.
108    pub fn set_proxied(&mut self, proxied: Option<bool>) -> &mut Self {
109        self.is_proxied = proxied;
110        self
111    }
112
113    /// Set the remote address of the connection used.
114    pub fn remote_addr(mut self, remote_addr: SocketAddr) -> Self {
115        self.set_remote_addr(Some(remote_addr));
116        self
117    }
118
119    /// Set the remote address of the connection used.
120    pub fn set_remote_addr(&mut self, remote_addr: Option<SocketAddr>) -> &mut Self {
121        self.remote_addr = remote_addr;
122        self
123    }
124
125    /// Set the local address of the connection used.
126    pub fn local_addr(mut self, local_addr: SocketAddr) -> Self {
127        self.set_local_addr(Some(local_addr));
128        self
129    }
130
131    /// Set the local address of the connection used.
132    pub fn set_local_addr(&mut self, local_addr: Option<SocketAddr>) -> &mut Self {
133        self.local_addr = local_addr;
134        self
135    }
136
137    /// Set a closure which will poison the associated connection.
138    ///
139    /// A poisoned connection will not be reused for subsequent requests by the pool
140    pub fn poison_fn(mut self, poison_fn: impl Fn() + Send + Sync + 'static) -> Self {
141        self.set_poison_fn(Some(poison_fn));
142        self
143    }
144
145    /// Set a closure which will poison the associated connection.
146    ///
147    /// A poisoned connection will not be reused for subsequent requests by the pool
148    pub fn set_poison_fn(
149        &mut self,
150        poison_fn: Option<impl Fn() + Send + Sync + 'static>,
151    ) -> &mut Self {
152        self.poison_fn =
153            poison_fn.map(|poison_fn| Arc::new(poison_fn) as Arc<dyn Fn() + Send + Sync>);
154        self
155    }
156
157    /// Build a [`ConnectionMetadata`] value.
158    ///
159    /// # Panics
160    ///
161    /// If either the `is_proxied` or `poison_fn` has not been set, then this method will panic
162    pub fn build(self) -> ConnectionMetadata {
163        ConnectionMetadata {
164            is_proxied: self
165                .is_proxied
166                .expect("is_proxied should be set for ConnectionMetadata"),
167            remote_addr: self.remote_addr,
168            local_addr: self.local_addr,
169            poison_fn: self
170                .poison_fn
171                .expect("poison_fn should be set for ConnectionMetadata"),
172        }
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use std::{
179        net::{IpAddr, Ipv6Addr},
180        sync::Mutex,
181    };
182
183    use super::*;
184
185    const TEST_SOCKET_ADDR: SocketAddr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 100);
186
187    #[test]
188    #[should_panic]
189    fn builder_panic_missing_proxied() {
190        ConnectionMetadataBuilder::new()
191            .poison_fn(|| {})
192            .local_addr(TEST_SOCKET_ADDR)
193            .remote_addr(TEST_SOCKET_ADDR)
194            .build();
195    }
196
197    #[test]
198    #[should_panic]
199    fn builder_panic_missing_poison_fn() {
200        ConnectionMetadataBuilder::new()
201            .proxied(true)
202            .local_addr(TEST_SOCKET_ADDR)
203            .remote_addr(TEST_SOCKET_ADDR)
204            .build();
205    }
206
207    #[test]
208    fn builder_all_fields_successful() {
209        let mutable_flag = Arc::new(Mutex::new(false));
210
211        let connection_metadata = ConnectionMetadataBuilder::new()
212            .proxied(true)
213            .local_addr(TEST_SOCKET_ADDR)
214            .remote_addr(TEST_SOCKET_ADDR)
215            .poison_fn({
216                let mutable_flag = Arc::clone(&mutable_flag);
217                move || {
218                    let mut guard = mutable_flag.lock().unwrap();
219                    *guard = !*guard;
220                }
221            })
222            .build();
223
224        assert!(connection_metadata.is_proxied);
225        assert_eq!(connection_metadata.remote_addr(), Some(TEST_SOCKET_ADDR));
226        assert_eq!(connection_metadata.local_addr(), Some(TEST_SOCKET_ADDR));
227        assert!(!(*mutable_flag.lock().unwrap()));
228        connection_metadata.poison();
229        assert!(*mutable_flag.lock().unwrap());
230    }
231
232    #[test]
233    fn builder_optional_fields_translate() {
234        let metadata1 = ConnectionMetadataBuilder::new()
235            .proxied(true)
236            .poison_fn(|| {})
237            .build();
238
239        assert_eq!(metadata1.local_addr(), None);
240        assert_eq!(metadata1.remote_addr(), None);
241
242        let metadata2 = ConnectionMetadataBuilder::new()
243            .proxied(true)
244            .poison_fn(|| {})
245            .local_addr(TEST_SOCKET_ADDR)
246            .build();
247
248        assert_eq!(metadata2.local_addr(), Some(TEST_SOCKET_ADDR));
249        assert_eq!(metadata2.remote_addr(), None);
250
251        let metadata3 = ConnectionMetadataBuilder::new()
252            .proxied(true)
253            .poison_fn(|| {})
254            .remote_addr(TEST_SOCKET_ADDR)
255            .build();
256
257        assert_eq!(metadata3.local_addr(), None);
258        assert_eq!(metadata3.remote_addr(), Some(TEST_SOCKET_ADDR));
259    }
260}
aws_smithy_runtime_api/client/connection.rs

aws_smithy_runtime_api/client/
connection.rs
