aws_smithy_http_client/client/pool/
partition.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Connection placement declarations.
7//!
8//! A [`Partition`] combines stable identity, driver placement, and an optional
9//! network-interface binding. It declares where connections live; sharing a
10//! connection never moves that connection's I/O or driver.
11
12use std::fmt;
13use std::future::Future;
14use std::pin::Pin;
15use std::sync::Arc;
16
17/// Declares where a client's connections are established and driven.
18#[derive(Clone)]
19pub struct Partition {
20    /// Stable identity used to resolve clients and index cells.
21    id: PartitionId,
22    /// Runtime placement for connection-owned tasks.
23    spawner: DriverSpawner,
24    /// Optional network-interface binding for new connections.
25    interface: Option<Arc<str>>,
26}
27
28impl Partition {
29    /// Creates a partition with a stable identity and driver spawner.
30    pub fn new(id: PartitionId, spawner: DriverSpawner) -> Self {
31        Self {
32            id,
33            spawner,
34            interface: None,
35        }
36    }
37
38    /// Binds connections established by this partition to an interface.
39    ///
40    /// The binding is applied before connect. On Linux, using this setting
41    /// sets `SO_BINDTODEVICE` and may require `CAP_NET_RAW` or root
42    /// privileges. Apple platforms, illumos, and Solaris use `IP_BOUND_IF`.
43    ///
44    /// This method is available on Linux and Android, Fuchsia, illumos,
45    /// Solaris, macOS, iOS, tvOS, visionOS, and watchOS.
46    #[cfg(any(
47        target_os = "android",
48        target_os = "fuchsia",
49        target_os = "illumos",
50        target_os = "ios",
51        target_os = "linux",
52        target_os = "macos",
53        target_os = "solaris",
54        target_os = "tvos",
55        target_os = "visionos",
56        target_os = "watchos",
57    ))]
58    #[cfg_attr(
59        docsrs,
60        doc(cfg(any(
61            target_os = "android",
62            target_os = "fuchsia",
63            target_os = "illumos",
64            target_os = "ios",
65            target_os = "linux",
66            target_os = "macos",
67            target_os = "solaris",
68            target_os = "tvos",
69            target_os = "visionos",
70            target_os = "watchos",
71        )))
72    )]
73    pub fn interface(mut self, interface: impl Into<String>) -> Self {
74        self.interface = Some(Arc::from(interface.into()));
75        self
76    }
77
78    /// Returns the configured network-interface name for validation.
79    #[cfg(any(
80        target_os = "android",
81        target_os = "fuchsia",
82        target_os = "illumos",
83        target_os = "ios",
84        target_os = "linux",
85        target_os = "macos",
86        target_os = "solaris",
87        target_os = "tvos",
88        target_os = "visionos",
89        target_os = "watchos",
90    ))]
91    pub(super) fn interface_name(&self) -> Option<&str> {
92        self.interface.as_deref()
93    }
94
95    /// Returns this partition's declared identity.
96    pub(super) fn id(&self) -> PartitionId {
97        self.id
98    }
99
100    /// Decomposes this declaration for immutable registry storage.
101    pub(super) fn into_parts(self) -> (PartitionId, DriverSpawner, Option<Arc<str>>) {
102        (self.id, self.spawner, self.interface)
103    }
104}
105
106impl fmt::Debug for Partition {
107    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        f.debug_struct("Partition")
109            .field("id", &self.id)
110            .field("spawner", &self.spawner)
111            .field("interface", &self.interface)
112            .finish()
113    }
114}
115
116/// Identifies one declared connection partition.
117#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
118pub struct PartitionId(usize);
119
120impl PartitionId {
121    /// Reserved identity for the implicit default partition.
122    pub const ANONYMOUS: Self = Self(usize::MAX);
123
124    /// Creates a partition identity from a caller-owned index.
125    ///
126    /// `usize::MAX` is reserved for [`Self::ANONYMOUS`] and is rejected when
127    /// used in an explicit partition declaration.
128    pub const fn from_index(index: usize) -> Self {
129        Self(index)
130    }
131
132    /// Returns whether this is the reserved anonymous partition identity.
133    pub const fn is_anonymous(self) -> bool {
134        self.0 == Self::ANONYMOUS.0
135    }
136}
137
138/// Runtime placement for one partition's connection-owned work.
139///
140/// Protocol drivers, establishment completion, deferred HTTP/1 readiness, and
141/// idle maintenance for a partition run through its spawner, even when another
142/// partition dispatches on one of its connections. The underlying I/O never
143/// moves to the requesting runtime.
144///
145/// A spawner is placement, not runtime ownership. It does not keep the runtime
146/// alive and does not participate in shutdown. When the owning runtime drops
147/// spawned work, the pool observes that through the connection's own close
148/// path.
149///
150/// Callers on Tokio use [`DriverSpawner::tokio`] with the handle of the runtime
151/// that should own the partition's connections.
152#[derive(Clone)]
153pub struct DriverSpawner {
154    inner: Arc<dyn Spawn>,
155}
156
157impl DriverSpawner {
158    /// Places connection-owned work on a specific Tokio runtime.
159    ///
160    /// Work is spawned on `handle` regardless of which thread invokes the
161    /// pool. Pass `tokio::runtime::Handle::current()` to use the runtime that
162    /// is constructing the partition.
163    #[cfg(feature = "rt-tokio")]
164    #[cfg_attr(docsrs, doc(cfg(feature = "rt-tokio")))]
165    pub fn tokio(handle: tokio::runtime::Handle) -> Self {
166        Self::new(TokioSpawn { handle })
167    }
168
169    /// Places connection-owned work through a caller-provided spawn function.
170    ///
171    /// `spawn` receives each connection-owned task and must run it to
172    /// completion on the runtime this partition represents. Dropping the task
173    /// instead is treated by the pool as that runtime shutting down.
174    ///
175    /// This is a test hook for observing task placement and is not a supported
176    /// extension point.
177    #[cfg(feature = "test-util")]
178    #[doc(hidden)]
179    pub fn from_fn<F>(spawn: F) -> Self
180    where
181        F: Fn(Pin<Box<dyn Future<Output = ()> + Send + 'static>>) + Send + Sync + 'static,
182    {
183        Self::new(FnSpawn(spawn))
184    }
185
186    /// Wraps one crate-internal placement implementation.
187    #[cfg_attr(
188        not(any(feature = "rt-tokio", feature = "test-util", test)),
189        allow(
190            dead_code,
191            reason = "no public constructor exists without a runtime feature"
192        )
193    )]
194    pub(in crate::client::pool) fn new(spawn: impl Spawn) -> Self {
195        Self {
196            inner: Arc::new(spawn),
197        }
198    }
199
200    /// Spawns connection-owned work on this partition's runtime.
201    pub(in crate::client::pool) fn spawn(
202        &self,
203        driver: Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
204    ) {
205        self.inner.spawn(driver);
206    }
207
208    /// Returns whether both values name the same underlying spawner.
209    #[cfg(test)]
210    pub(in crate::client::pool) fn ptr_eq(&self, other: &Self) -> bool {
211        Arc::ptr_eq(&self.inner, &other.inner)
212    }
213}
214
215impl fmt::Debug for DriverSpawner {
216    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
217        f.debug_tuple("DriverSpawner").field(&self.inner).finish()
218    }
219}
220
221/// Placement capability behind [`DriverSpawner`].
222///
223/// This trait is crate-internal so the pool can extend the placement contract
224/// without changing the public spawner type.
225pub(in crate::client::pool) trait Spawn:
226    fmt::Debug + Send + Sync + 'static
227{
228    /// Spawns connection-owned work on the partition's runtime.
229    fn spawn(&self, driver: Pin<Box<dyn Future<Output = ()> + Send + 'static>>);
230}
231
232/// Placement on a captured Tokio runtime handle.
233#[cfg(feature = "rt-tokio")]
234#[derive(Debug)]
235struct TokioSpawn {
236    /// Runtime that receives connection-owned tasks.
237    handle: tokio::runtime::Handle,
238}
239
240#[cfg(feature = "rt-tokio")]
241impl Spawn for TokioSpawn {
242    fn spawn(&self, driver: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) {
243        drop(self.handle.spawn(driver));
244    }
245}
246
247/// Placement through a caller-provided spawn function.
248#[cfg(feature = "test-util")]
249struct FnSpawn<F>(F);
250
251#[cfg(feature = "test-util")]
252impl<F> fmt::Debug for FnSpawn<F> {
253    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
254        f.write_str("FnSpawn")
255    }
256}
257
258#[cfg(feature = "test-util")]
259impl<F> Spawn for FnSpawn<F>
260where
261    F: Fn(Pin<Box<dyn Future<Output = ()> + Send + 'static>>) + Send + Sync + 'static,
262{
263    fn spawn(&self, driver: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) {
264        (self.0)(driver);
265    }
266}
267
268/// Controls which partitions may dispatch through each other's connections.
269///
270/// The default is [`NetworkInterface`](Self::NetworkInterface): partitions
271/// with the same interface binding, including all partitions with no binding,
272/// may borrow each other's idle HTTP/1 connections and route requests to each
273/// other's HTTP/2 connections. Reuse scope never changes the bounded
274/// connection budget, which is shared by every partition for one origin.
275#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
276#[non_exhaustive]
277pub enum ConnectionReuseScope {
278    /// A connection serves only requests from its owning partition.
279    Partition,
280    /// Partitions with the same interface binding may share connections.
281    #[default]
282    NetworkInterface,
283    /// Every partition in the pool may share connections.
284    Pool,
285}
286
287/// Identifies the exact set of partitions allowed to share a connection.
288///
289/// The configured [`ConnectionReuseScope`] maps each partition to one group:
290/// partition-local scope creates a group for one partition, network-interface
291/// scope groups equal interface bindings, and pool scope creates one group for
292/// the whole pool.
293#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
294pub(in crate::client::pool) enum EligibilityGroup {
295    /// Only the partition with this identity is eligible.
296    Partition(PartitionId),
297    /// Partitions with this exact interface binding are eligible.
298    NetworkInterface(Option<Arc<str>>),
299    /// Every partition in the pool is eligible.
300    Pool,
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306    #[cfg(feature = "rt-tokio")]
307    use std::sync::atomic::{AtomicBool, Ordering};
308
309    #[derive(Debug)]
310    struct TestSpawner;
311
312    impl Spawn for TestSpawner {
313        fn spawn(&self, driver: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) {
314            drop(driver);
315        }
316    }
317
318    #[test]
319    fn anonymous_identity_is_reserved() {
320        assert!(PartitionId::ANONYMOUS.is_anonymous());
321        assert!(PartitionId::from_index(usize::MAX).is_anonymous());
322        assert!(!PartitionId::from_index(0).is_anonymous());
323    }
324
325    #[test]
326    fn partition_retains_placement() {
327        let spawner = DriverSpawner::new(TestSpawner);
328        let partition = Partition::new(PartitionId::from_index(7), spawner.clone());
329        let (id, retained, interface) = partition.into_parts();
330        assert_eq!(PartitionId::from_index(7), id);
331        assert_eq!(None, interface);
332        assert!(spawner.ptr_eq(&retained));
333
334        let driver = Box::pin(async {});
335        retained.spawn(driver);
336    }
337
338    #[cfg(feature = "test-util")]
339    #[test]
340    fn from_fn_receives_each_driver() {
341        let submitted = Arc::new(std::sync::atomic::AtomicUsize::new(0));
342        let observed = submitted.clone();
343        let spawner = DriverSpawner::from_fn(move |driver| {
344            observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
345            drop(driver);
346        });
347
348        spawner.spawn(Box::pin(async {}));
349        spawner.spawn(Box::pin(async {}));
350        assert_eq!(2, submitted.load(std::sync::atomic::Ordering::SeqCst));
351        assert_eq!("DriverSpawner(FnSpawn)", format!("{spawner:?}"));
352    }
353
354    #[cfg(feature = "rt-tokio")]
355    #[test]
356    fn tokio_spawner_uses_captured_runtime() {
357        let owner = tokio::runtime::Builder::new_current_thread()
358            .build()
359            .unwrap();
360        let owner_id = owner.handle().id();
361        let ran = Arc::new(AtomicBool::new(false));
362        let task_ran = ran.clone();
363        let spawner = DriverSpawner::tokio(owner.handle().clone());
364
365        owner.block_on(async {
366            spawner.spawn(Box::pin(async move {
367                assert_eq!(owner_id, tokio::runtime::Handle::current().id());
368                task_ran.store(true, Ordering::SeqCst);
369            }));
370            for _ in 0..10 {
371                if ran.load(Ordering::SeqCst) {
372                    break;
373                }
374                tokio::task::yield_now().await;
375            }
376        });
377        assert!(ran.load(Ordering::SeqCst));
378    }
379
380    #[cfg(feature = "rt-tokio")]
381    #[test]
382    fn tokio_spawner_accepts_work_from_a_foreign_runtime() {
383        let owner = tokio::runtime::Builder::new_current_thread()
384            .build()
385            .unwrap();
386        let foreign = tokio::runtime::Builder::new_current_thread()
387            .build()
388            .unwrap();
389        let owner_id = owner.handle().id();
390        let ran = Arc::new(AtomicBool::new(false));
391        let task_ran = ran.clone();
392        let spawner = DriverSpawner::tokio(owner.handle().clone());
393
394        foreign.block_on(async {
395            spawner.spawn(Box::pin(async move {
396                assert_eq!(owner_id, tokio::runtime::Handle::current().id());
397                task_ran.store(true, Ordering::SeqCst);
398            }));
399        });
400        owner.block_on(async {
401            for _ in 0..10 {
402                if ran.load(Ordering::SeqCst) {
403                    break;
404                }
405                tokio::task::yield_now().await;
406            }
407        });
408        assert!(ran.load(Ordering::SeqCst));
409    }
410}
aws_smithy_http_client/client/pool/partition.rs

aws_smithy_http_client/client/pool/
partition.rs
