AWS SDK

AWS SDK

rev. 174400987dccd7e137fefa96b1143d21c7ddfb78

Files changed:

tmp-codegen-diff/aws-sdk/sdk/aws-smithy-http-client/src/client/pool/handshake.rs

@@ -0,1 +0,694 @@
           1  +
/*
           2  +
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
           3  +
 * SPDX-License-Identifier: Apache-2.0
           4  +
 */
           5  +
           6  +
//! Tower service adapters for hyper's HTTP protocol handshake.
           7  +
//!
           8  +
//! - `ConnectionLimit`: wraps a connector, acquires semaphore permits before
           9  +
//!   connecting, and returns an `EstablishedConnection`.
          10  +
//! - `H1ConnectAndHandshake` / `H2ConnectAndHandshake`: perform the protocol
          11  +
//!   handshake on an already-connected IO stream, producing a `ManagedConnection`.
          12  +
          13  +
use std::future::Future;
          14  +
use std::pin::Pin;
          15  +
use std::sync::Arc;
          16  +
use std::task::{Context, Poll};
          17  +
use std::time::Instant;
          18  +
          19  +
use aws_smithy_runtime_api::box_error::BoxError;
          20  +
use aws_smithy_runtime_api::client::connection::ConnectionId;
          21  +
use aws_smithy_types::body::SdkBody;
          22  +
use hyper_util::client::legacy::connect::{Connection, HttpInfo};
          23  +
use hyper_util::rt::TokioExecutor;
          24  +
use tokio::sync::Semaphore;
          25  +
use tower::Service;
          26  +
          27  +
use super::connection::{
          28  +
    Authority, ConnectCtx, ConnectionCreatedEvent, ConnectionFailedEvent, ConnectionInfo,
          29  +
    ConnectionPermit, ConnectionTiming, EstablishedConnection, ManagedConnection,
          30  +
    NegotiatedProtocol,
          31  +
};
          32  +
use super::partition::DriverSpawner;
          33  +
          34  +
/// Pool-scoped instrumentation primitives shared across layers in the
          35  +
/// pool stack. Held by `ConnectionPool` for the pool's lifetime and
          36  +
/// cloned into individual layers (the H1/H2 handshake services) at
          37  +
/// construction.
          38  +
///
          39  +
/// Cheap to clone (one `Arc` per primitive).
          40  +
#[derive(Clone)]
          41  +
pub(crate) struct PoolHooks {
          42  +
    conn_id_counter: Arc<std::sync::atomic::AtomicU64>,
          43  +
    pub(crate) listener: Option<Arc<dyn super::connection::ConnectionEventListener>>,
          44  +
}
          45  +
          46  +
impl PoolHooks {
          47  +
    pub(crate) fn new(
          48  +
        listener: Option<Arc<dyn super::connection::ConnectionEventListener>>,
          49  +
    ) -> Self {
          50  +
        Self {
          51  +
            conn_id_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
          52  +
            listener,
          53  +
        }
          54  +
    }
          55  +
          56  +
    /// Mint the next connection id. Stable for the connection's lifetime;
          57  +
    /// the underlying counter wraps at `u64::MAX`.
          58  +
    pub(crate) fn next_conn_id(&self) -> ConnectionId {
          59  +
        ConnectionId::new(
          60  +
            self.conn_id_counter
          61  +
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed),
          62  +
        )
          63  +
    }
          64  +
          65  +
    /// Fire the listener's connection-created callback, if a listener is set.
          66  +
    pub(crate) fn on_created(&self, event: &ConnectionCreatedEvent) {
          67  +
        if let Some(ref l) = self.listener {
          68  +
            l.on_created(event);
          69  +
        }
          70  +
    }
          71  +
          72  +
    /// Fire the listener's connection-failed callback, if a listener is set.
          73  +
    pub(crate) fn on_connection_failed(&self, event: &ConnectionFailedEvent) {
          74  +
        if let Some(ref l) = self.listener {
          75  +
            l.on_connection_failed(event);
          76  +
        }
          77  +
    }
          78  +
          79  +
    /// Fire the listener's connection-reused callback, if a listener is set.
          80  +
    pub(crate) fn on_reused(&self, event: &super::connection::ConnectionReusedEvent) {
          81  +
        if let Some(ref l) = self.listener {
          82  +
            l.on_reused(event);
          83  +
        }
          84  +
    }
          85  +
          86  +
    /// Fire the listener's connection-closed callback, if a listener is set.
          87  +
    pub(crate) fn on_closed(&self, event: &super::connection::ConnectionClosedEvent) {
          88  +
        if let Some(ref l) = self.listener {
          89  +
            l.on_closed(event);
          90  +
        }
          91  +
    }
          92  +
}
          93  +
          94  +
/// Wraps a connector service, acquiring semaphore permits before connecting.
          95  +
///
          96  +
/// Returns an [`EstablishedConnection`] so the permit can be stored on
          97  +
/// `ManagedConnection` and held for the connection's lifetime.
          98  +
///
          99  +
/// Target type is `ConnectCtx`: the inner TCP connector is `Service<Uri>`,
         100  +
/// so this layer extracts `ctx.uri` to pass through. Per-operation timeouts
         101  +
/// on the context are applied here: `connect_timeout` wraps the inner
         102  +
/// connector call (which includes TCP + TLS since the inner is the
         103  +
/// TLS-wrapped connector). Cache hits skip this layer entirely, so
         104  +
/// `connect_timeout` is automatically new-connection-only.
         105  +
pub(crate) struct ConnectionLimit<C> {
         106  +
    inner: C,
         107  +
    global: Option<Arc<Semaphore>>,
         108  +
    per_host: Option<Arc<Semaphore>>,
         109  +
    counters: Arc<super::stats::ConnectionCounters>,
         110  +
    /// Cross-partition active-reclaim handle. `None` when no reclaim peer
         111  +
    /// exists or the pool is mid-teardown; the cap path is then a plain
         112  +
    /// blocking acquire.
         113  +
    reclaim: Option<super::PeerReclaimHandle>,
         114  +
}
         115  +
         116  +
impl<C> ConnectionLimit<C> {
         117  +
    pub(crate) fn new(
         118  +
        inner: C,
         119  +
        global: Option<Arc<Semaphore>>,
         120  +
        per_host: Option<Arc<Semaphore>>,
         121  +
        counters: Arc<super::stats::ConnectionCounters>,
         122  +
        reclaim: Option<super::PeerReclaimHandle>,
         123  +
    ) -> Self {
         124  +
        Self {
         125  +
            inner,
         126  +
            global,
         127  +
            per_host,
         128  +
            counters,
         129  +
            reclaim,
         130  +
        }
         131  +
    }
         132  +
}
         133  +
         134  +
impl<C: Clone> Clone for ConnectionLimit<C> {
         135  +
    fn clone(&self) -> Self {
         136  +
        Self {
         137  +
            inner: self.inner.clone(),
         138  +
            global: self.global.clone(),
         139  +
            per_host: self.per_host.clone(),
         140  +
            counters: self.counters.clone(),
         141  +
            reclaim: self.reclaim.clone(),
         142  +
        }
         143  +
    }
         144  +
}
         145  +
         146  +
impl<C, IO> Service<ConnectCtx> for ConnectionLimit<C>
         147  +
where
         148  +
    C: Service<http_1x::Uri, Response = IO> + Clone + Send + 'static,
         149  +
    C::Error: Into<BoxError> + 'static,
         150  +
    C::Future: Send + 'static,
         151  +
    IO: Send + 'static,
         152  +
{
         153  +
    type Response = EstablishedConnection<IO>;
         154  +
    type Error = BoxError;
         155  +
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
         156  +
         157  +
    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         158  +
        self.inner.poll_ready(cx).map_err(Into::into)
         159  +
    }
         160  +
         161  +
    fn call(&mut self, ctx: ConnectCtx) -> Self::Future {
         162  +
        let mut inner = self.inner.clone();
         163  +
        let global = self.global.clone();
         164  +
        let per_host = self.per_host.clone();
         165  +
        let counters = self.counters.clone();
         166  +
        let reclaim = self.reclaim.clone();
         167  +
        Box::pin(async move {
         168  +
            let mode = ctx.mode;
         169  +
            // Per-host before global: never hold a global permit while
         170  +
            // waiting on a per-host permit.
         171  +
            let per_host_permit = match &per_host {
         172  +
                Some(sem) => Some(
         173  +
                    acquire_or_reclaim(sem, reclaim.as_ref(), mode, || {
         174  +
                        // `send_request` validated the URI has scheme+authority
         175  +
                        // before dispatching, so this cannot fail here.
         176  +
                        let key = super::PoolKey::from_uri(&ctx.uri)
         177  +
                            .expect("connect URI has scheme+authority");
         178  +
                        super::BindingConstraint::PerHost(key)
         179  +
                    })
         180  +
                    .await?,
         181  +
                ),
         182  +
                None => None,
         183  +
            };
         184  +
            let global_permit = match &global {
         185  +
                Some(sem) => Some(
         186  +
                    acquire_or_reclaim(sem, reclaim.as_ref(), mode, || {
         187  +
                        super::BindingConstraint::Global
         188  +
                    })
         189  +
                    .await?,
         190  +
                ),
         191  +
                None => None,
         192  +
            };
         193  +
            let permit = Arc::new(ConnectionPermit::new(global_permit, per_host_permit));
         194  +
            let establishing = super::stats::EstablishingGuard::new(counters);
         195  +
         196  +
            std::future::poll_fn(|cx| inner.poll_ready(cx))
         197  +
                .await
         198  +
                .map_err(Into::into)?;
         199  +
         200  +
            // Apply connect_timeout only around the actual connector call
         201  +
            // (TCP + TLS). If `ctx.connect_timeout` is `None`, this is a
         202  +
            // plain `inner.call(uri).await`.
         203  +
            let uri = ctx.uri;
         204  +
            let connect_fut = inner.call(uri);
         205  +
            let io = super::super::timeout::maybe_timeout_future(
         206  +
                connect_fut,
         207  +
                ctx.connect_timeout.as_ref().map(|t| t.duration),
         208  +
                ctx.connect_timeout.as_ref().map(|t| &t.sleep_impl),
         209  +
                super::super::timeout::TimeoutKind::Connect,
         210  +
            )
         211  +
            .await?;
         212  +
            Ok(EstablishedConnection {
         213  +
                io,
         214  +
                permit,
         215  +
                establishing,
         216  +
            })
         217  +
        })
         218  +
    }
         219  +
}
         220  +
         221  +
/// Acquire one owned permit. Fast path is `try_acquire_owned`. On
         222  +
/// `NoPermits`: under [`AcquireMode::NonBlocking`] return [`CapBound`]
         223  +
/// immediately (the caller will try a peer borrow); under
         224  +
/// [`AcquireMode::Blocking`] free one peer's idle connection for
         225  +
/// `constraint` (inline, best-effort) then blocking-acquire — the blocking
         226  +
/// acquire is the authoritative take regardless of whether reclaim freed a
         227  +
/// permit. `constraint` is built only on the blocking cap-bound branch.
         228  +
async fn acquire_or_reclaim(
         229  +
    sem: &Arc<Semaphore>,
         230  +
    reclaim: Option<&super::PeerReclaimHandle>,
         231  +
    mode: super::connection::AcquireMode,
         232  +
    constraint: impl FnOnce() -> super::BindingConstraint,
         233  +
) -> Result<tokio::sync::OwnedSemaphorePermit, BoxError> {
         234  +
    match sem.clone().try_acquire_owned() {
         235  +
        Ok(permit) => Ok(permit),
         236  +
        Err(tokio::sync::TryAcquireError::NoPermits) => {
         237  +
            if mode == super::connection::AcquireMode::NonBlocking {
         238  +
                return Err(super::connection::CapBound.into());
         239  +
            }
         240  +
            if let Some(reclaim) = reclaim {
         241  +
                reclaim.try_free_under_load(&constraint());
         242  +
            }
         243  +
            sem.clone()
         244  +
                .acquire_owned()
         245  +
                .await
         246  +
                .map_err(|_| "pool closed".into())
         247  +
        }
         248  +
        Err(tokio::sync::TryAcquireError::Closed) => Err("pool closed".into()),
         249  +
    }
         250  +
}
         251  +
         252  +
// ---------------------------------------------------------------------------
         253  +
// Tower Service wrappers for hyper's SendRequest
         254  +
// ---------------------------------------------------------------------------
         255  +
         256  +
/// Tower `Service` adapter for `hyper::client::conn::http1::SendRequest`.
         257  +
pub(crate) struct H1SendRequest {
         258  +
    inner: hyper::client::conn::http1::SendRequest<SdkBody>,
         259  +
}
         260  +
         261  +
impl H1SendRequest {
         262  +
    pub(crate) fn new(inner: hyper::client::conn::http1::SendRequest<SdkBody>) -> Self {
         263  +
        Self { inner }
         264  +
    }
         265  +
}
         266  +
         267  +
impl Service<http_1x::Request<SdkBody>> for H1SendRequest {
         268  +
    type Response = http_1x::Response<hyper::body::Incoming>;
         269  +
    type Error = hyper::Error;
         270  +
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
         271  +
         272  +
    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         273  +
        self.inner.poll_ready(cx)
         274  +
    }
         275  +
         276  +
    fn call(&mut self, req: http_1x::Request<SdkBody>) -> Self::Future {
         277  +
        Box::pin(self.inner.send_request(req))
         278  +
    }
         279  +
}
         280  +
         281  +
/// Tower `Service` adapter for `hyper::client::conn::http2::SendRequest`.
         282  +
///
         283  +
/// Clone is supported because HTTP/2 multiplexes requests over a single connection.
         284  +
#[derive(Clone)]
         285  +
pub(crate) struct H2SendRequest {
         286  +
    inner: hyper::client::conn::http2::SendRequest<SdkBody>,
         287  +
}
         288  +
         289  +
impl H2SendRequest {
         290  +
    pub(crate) fn new(inner: hyper::client::conn::http2::SendRequest<SdkBody>) -> Self {
         291  +
        Self { inner }
         292  +
    }
         293  +
}
         294  +
         295  +
impl Service<http_1x::Request<SdkBody>> for H2SendRequest {
         296  +
    type Response = http_1x::Response<hyper::body::Incoming>;
         297  +
    type Error = hyper::Error;
         298  +
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
         299  +
         300  +
    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         301  +
        self.inner.poll_ready(cx)
         302  +
    }
         303  +
         304  +
    fn call(&mut self, req: http_1x::Request<SdkBody>) -> Self::Future {
         305  +
        Box::pin(self.inner.send_request(req))
         306  +
    }
         307  +
}
         308  +
         309  +
// ---------------------------------------------------------------------------
         310  +
// Connect-and-handshake services
         311  +
// ---------------------------------------------------------------------------
         312  +
         313  +
/// Extract `ConnectionInfo` from a just-connected IO stream.
         314  +
fn capture_info<IO: Connection>(io: &IO, authority: Authority) -> ConnectionInfo {
         315  +
    let connected = io.connected();
         316  +
    let is_proxied = connected.is_proxied();
         317  +
    let mut extras = http_1x::Extensions::new();
         318  +
    connected.get_extras(&mut extras);
         319  +
    let http_info = extras.get::<HttpInfo>();
         320  +
    ConnectionInfo {
         321  +
        remote_addr: http_info.map(|i| i.remote_addr()),
         322  +
        local_addr: http_info.map(|i| i.local_addr()),
         323  +
        is_proxied,
         324  +
        authority,
         325  +
    }
         326  +
}
         327  +
         328  +
/// Connects and performs an HTTP/1.1 handshake, spawning the connection
         329  +
/// driver onto the partition's runtime via the captured [`DriverSpawner`].
         330  +
///
         331  +
/// The connector is expected to return an [`EstablishedConnection`], typically
         332  +
/// produced by [`ConnectionLimit`] wrapping a TCP/TLS connector.
         333  +
pub(crate) struct H1ConnectAndHandshake<C> {
         334  +
    connector: C,
         335  +
    hooks: PoolHooks,
         336  +
    spawner: Arc<dyn DriverSpawner>,
         337  +
}
         338  +
         339  +
impl<C> H1ConnectAndHandshake<C> {
         340  +
    pub(crate) fn new(connector: C, hooks: PoolHooks, spawner: Arc<dyn DriverSpawner>) -> Self {
         341  +
        Self {
         342  +
            connector,
         343  +
            hooks,
         344  +
            spawner,
         345  +
        }
         346  +
    }
         347  +
}
         348  +
         349  +
impl<C: Clone> Clone for H1ConnectAndHandshake<C> {
         350  +
    fn clone(&self) -> Self {
         351  +
        Self {
         352  +
            connector: self.connector.clone(),
         353  +
            hooks: self.hooks.clone(),
         354  +
            spawner: self.spawner.clone(),
         355  +
        }
         356  +
    }
         357  +
}
         358  +
         359  +
impl<C, IO> Service<ConnectCtx> for H1ConnectAndHandshake<C>
         360  +
where
         361  +
    C: Service<ConnectCtx, Response = EstablishedConnection<IO>>,
         362  +
    C::Error: Into<BoxError> + 'static,
         363  +
    C::Future: Send + 'static,
         364  +
    IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static,
         365  +
{
         366  +
    type Response = ManagedConnection<H1SendRequest>;
         367  +
    type Error = BoxError;
         368  +
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
         369  +
         370  +
    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         371  +
        self.connector.poll_ready(cx).map_err(Into::into)
         372  +
    }
         373  +
         374  +
    fn call(&mut self, ctx: ConnectCtx) -> Self::Future {
         375  +
        let hooks = self.hooks.clone();
         376  +
        let spawner = self.spawner.clone();
         377  +
        let authority = Authority::new(
         378  +
            ctx.uri
         379  +
                .authority()
         380  +
                .expect("request URI has authority")
         381  +
                .as_str(),
         382  +
        );
         383  +
        let fut = self.connector.call(ctx);
         384  +
        Box::pin(async move {
         385  +
            let connect_start = Instant::now();
         386  +
            let EstablishedConnection {
         387  +
                io,
         388  +
                permit,
         389  +
                establishing,
         390  +
            } = match fut.await.map_err(Into::into) {
         391  +
                Ok(v) => v,
         392  +
                Err(e) => {
         393  +
                    hooks.on_connection_failed(&ConnectionFailedEvent::new(
         394  +
                        authority,
         395  +
                        None,
         396  +
                        Into::into(e.to_string()),
         397  +
                    ));
         398  +
                    return Err(e);
         399  +
                }
         400  +
            };
         401  +
            let connect_duration = connect_start.elapsed();
         402  +
            let info = capture_info(&io, authority.clone());
         403  +
            let conn_id = hooks.next_conn_id();
         404  +
         405  +
            let (tx, conn) = match hyper::client::conn::http1::Builder::new()
         406  +
                .handshake(io)
         407  +
                .await
         408  +
            {
         409  +
                Ok(v) => v,
         410  +
                Err(e) => {
         411  +
                    let boxed: BoxError = Box::new(e);
         412  +
                    hooks.on_connection_failed(&ConnectionFailedEvent::new(
         413  +
                        authority,
         414  +
                        info.remote_addr,
         415  +
                        Into::into(boxed.to_string()),
         416  +
                    ));
         417  +
                    return Err(boxed);
         418  +
                }
         419  +
            };
         420  +
         421  +
            tracing::debug!(
         422  +
                conn_id = %conn_id,
         423  +
                protocol = "h1",
         424  +
                remote = ?info.remote_addr,
         425  +
                local = ?info.local_addr,
         426  +
                "pool: connection established"
         427  +
            );
         428  +
         429  +
            hooks.on_created(&ConnectionCreatedEvent::new(
         430  +
                conn_id,
         431  +
                authority,
         432  +
                info.remote_addr,
         433  +
                NegotiatedProtocol::Http1,
         434  +
                ConnectionTiming::new(connect_duration),
         435  +
            ));
         436  +
         437  +
            let established = establishing.promote(super::stats::PROTO_H1);
         438  +
         439  +
            spawner.spawn(Box::pin({
         440  +
                let remote_addr = info.remote_addr;
         441  +
                let local_addr = info.local_addr;
         442  +
                async move {
         443  +
                    if let Err(e) = conn.with_upgrades().await {
         444  +
                        tracing::debug!(
         445  +
                            conn_id = %conn_id,
         446  +
                            protocol = "h1",
         447  +
                            ?remote_addr,
         448  +
                            ?local_addr,
         449  +
                            error = %e,
         450  +
                            "pool: connection driver error"
         451  +
                        );
         452  +
                    }
         453  +
                }
         454  +
            }));
         455  +
         456  +
            Ok(ManagedConnection::new(
         457  +
                H1SendRequest::new(tx),
         458  +
                info,
         459  +
                conn_id,
         460  +
                permit,
         461  +
                established,
         462  +
            ))
         463  +
        })
         464  +
    }
         465  +
}
         466  +
         467  +
/// Connects and performs an HTTP/2 handshake, spawning the connection
         468  +
/// driver onto the partition's runtime via the captured [`DriverSpawner`].
         469  +
///
         470  +
/// Pinned to `Service<()>` because this service sits in the Negotiate
         471  +
/// upgrade path: the connection is already established via the shared
         472  +
/// `Inspected` slot.
         473  +
pub(crate) struct H2ConnectAndHandshake<C> {
         474  +
    connector: C,
         475  +
    h2_ref: super::connection::H2ConnectionRef,
         476  +
    hooks: PoolHooks,
         477  +
    authority: Authority,
         478  +
    spawner: Arc<dyn DriverSpawner>,
         479  +
}
         480  +
         481  +
impl<C> H2ConnectAndHandshake<C> {
         482  +
    /// Create an H2 handshake service that publishes each newly established
         483  +
    /// connection's state — `ConnectionMetadata` plus the shared
         484  +
    /// `active_streams` counter and `idle_at` timestamp — into `h2_ref`. The
         485  +
    /// same ref is held on the read side by `SingletonConnection` (clones of
         486  +
    /// the ref share the underlying slot), so the H2 checkout path can expose
         487  +
    /// connection metadata and poison support to the adapter layer, and track
         488  +
    /// stream occupancy, even though `Singled<…>` itself is opaque.
         489  +
    pub(crate) fn new(
         490  +
        connector: C,
         491  +
        h2_ref: super::connection::H2ConnectionRef,
         492  +
        hooks: PoolHooks,
         493  +
        authority: Authority,
         494  +
        spawner: Arc<dyn DriverSpawner>,
         495  +
    ) -> Self {
         496  +
        Self {
         497  +
            connector,
         498  +
            h2_ref,
         499  +
            hooks,
         500  +
            authority,
         501  +
            spawner,
         502  +
        }
         503  +
    }
         504  +
}
         505  +
         506  +
impl<C: Clone> Clone for H2ConnectAndHandshake<C> {
         507  +
    fn clone(&self) -> Self {
         508  +
        Self {
         509  +
            connector: self.connector.clone(),
         510  +
            h2_ref: self.h2_ref.clone(),
         511  +
            hooks: self.hooks.clone(),
         512  +
            authority: self.authority.clone(),
         513  +
            spawner: self.spawner.clone(),
         514  +
        }
         515  +
    }
         516  +
}
         517  +
         518  +
impl<C, IO> Service<()> for H2ConnectAndHandshake<C>
         519  +
where
         520  +
    C: Service<(), Response = EstablishedConnection<IO>>,
         521  +
    C::Error: Into<BoxError> + 'static,
         522  +
    C::Future: Send + 'static,
         523  +
    IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static,
         524  +
{
         525  +
    type Response = ManagedConnection<H2SendRequest>;
         526  +
    type Error = BoxError;
         527  +
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
         528  +
         529  +
    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         530  +
        self.connector.poll_ready(cx).map_err(Into::into)
         531  +
    }
         532  +
         533  +
    fn call(&mut self, _req: ()) -> Self::Future {
         534  +
        let fut = self.connector.call(());
         535  +
        let h2_ref = self.h2_ref.clone();
         536  +
        let hooks = self.hooks.clone();
         537  +
        let authority = self.authority.clone();
         538  +
        let spawner = self.spawner.clone();
         539  +
        Box::pin(async move {
         540  +
            let connect_start = Instant::now();
         541  +
            let EstablishedConnection {
         542  +
                io,
         543  +
                permit,
         544  +
                establishing,
         545  +
            } = match fut.await.map_err(Into::into) {
         546  +
                Ok(v) => v,
         547  +
                Err(e) => {
         548  +
                    hooks.on_connection_failed(&ConnectionFailedEvent::new(
         549  +
                        authority,
         550  +
                        None,
         551  +
                        Into::into(e.to_string()),
         552  +
                    ));
         553  +
                    return Err(e);
         554  +
                }
         555  +
            };
         556  +
            let connect_duration = connect_start.elapsed();
         557  +
            let info = capture_info(&io, authority.clone());
         558  +
            let conn_id = hooks.next_conn_id();
         559  +
         560  +
            let (tx, conn) = match hyper::client::conn::http2::Builder::new(TokioExecutor::new())
         561  +
                .handshake(io)
         562  +
                .await
         563  +
            {
         564  +
                Ok(v) => v,
         565  +
                Err(e) => {
         566  +
                    let boxed: BoxError = Box::new(e);
         567  +
                    hooks.on_connection_failed(&ConnectionFailedEvent::new(
         568  +
                        authority,
         569  +
                        info.remote_addr,
         570  +
                        Into::into(boxed.to_string()),
         571  +
                    ));
         572  +
                    return Err(boxed);
         573  +
                }
         574  +
            };
         575  +
         576  +
            tracing::debug!(
         577  +
                conn_id = %conn_id,
         578  +
                protocol = "h2",
         579  +
                remote = ?info.remote_addr,
         580  +
                local = ?info.local_addr,
         581  +
                "pool: connection established"
         582  +
            );
         583  +
         584  +
            hooks.on_created(&ConnectionCreatedEvent::new(
         585  +
                conn_id,
         586  +
                authority,
         587  +
                info.remote_addr,
         588  +
                NegotiatedProtocol::Http2,
         589  +
                ConnectionTiming::new(connect_duration),
         590  +
            ));
         591  +
         592  +
            let established = establishing.promote(super::stats::PROTO_H2);
         593  +
         594  +
            spawner.spawn(Box::pin({
         595  +
                let remote_addr = info.remote_addr;
         596  +
                let local_addr = info.local_addr;
         597  +
                async move {
         598  +
                    if let Err(e) = conn.await {
         599  +
                        tracing::debug!(
         600  +
                            conn_id = %conn_id,
         601  +
                            protocol = "h2",
         602  +
                            ?remote_addr,
         603  +
                            ?local_addr,
         604  +
                            error = %e,
         605  +
                            "pool: connection driver error"
         606  +
                        );
         607  +
                    }
         608  +
                }
         609  +
            }));
         610  +
         611  +
            let managed =
         612  +
                ManagedConnection::new(H2SendRequest::new(tx), info, conn_id, permit, established);
         613  +
            // Publish this connection's state (metadata + active_streams +
         614  +
            // idle_at refs) for the checkout side. Singleton replaces its
         615  +
            // stored connection wholesale on each new handshake, so
         616  +
            // last-writer-wins on the ref is correct.
         617  +
            h2_ref.publish(super::connection::H2ConnectionState {
         618  +
                metadata: managed.metadata(),
         619  +
                active_streams: managed.active_streams_ref(),
         620  +
                idle_at: managed.idle_at_ref(),
         621  +
            });
         622  +
            Ok(managed)
         623  +
        })
         624  +
    }
         625  +
}
         626  +
         627  +
#[cfg(test)]
         628  +
mod tests {
         629  +
    use super::*;
         630  +
    use crate::client::timeout::test::NeverConnects;
         631  +
    use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep};
         632  +
    use std::time::Duration;
         633  +
    use tower::Service as _;
         634  +
         635  +
    /// Verify `ConnectionLimit` applies `connect_timeout` from `ConnectCtx`
         636  +
    /// to the inner TCP connector. `NeverConnects` returns a connector
         637  +
    /// future that never resolves; a short connect_timeout should fire
         638  +
    /// and produce an `HTTP connect timeout occurred after …` error.
         639  +
    #[tokio::test(start_paused = true)]
         640  +
    async fn connect_timeout_fires_on_slow_connector() {
         641  +
        let mut svc = ConnectionLimit::new(
         642  +
            NeverConnects::default(),
         643  +
            None,
         644  +
            None,
         645  +
            Arc::new(super::super::stats::ConnectionCounters::default()),
         646  +
            None,
         647  +
        );
         648  +
        let sleep = SharedAsyncSleep::new(TokioSleep::new());
         649  +
        let ctx = ConnectCtx::new(
         650  +
            "http://example.com".parse().unwrap(),
         651  +
            Some(super::super::connection::TimeoutContext::new(
         652  +
                Duration::from_millis(500),
         653  +
                sleep,
         654  +
            )),
         655  +
        );
         656  +
        let err = match svc.call(ctx).await {
         657  +
            Ok(_) => panic!("connect timeout should fire against a never-resolving connector"),
         658  +
            Err(err) => err,
         659  +
        };
         660  +
        let msg = format!("{err}");
         661  +
        assert!(
         662  +
            msg.contains("HTTP connect"),
         663  +
            "expected `HTTP connect` in error, got: {msg}"
         664  +
        );
         665  +
        assert!(
         666  +
            msg.contains("500ms"),
         667  +
            "expected `500ms` in error, got: {msg}"
         668  +
        );
         669  +
    }
         670  +
         671  +
    /// Without a connect_timeout set, `ConnectionLimit` passes through the
         672  +
    /// connector call with no wrapping. Verify the absence of a timeout
         673  +
    /// means the slow connector stays pending (we just check we can start
         674  +
    /// the call; with `start_paused`, nothing advances so the future is
         675  +
    /// still pending after a short yield).
         676  +
    #[tokio::test(start_paused = true)]
         677  +
    async fn no_timeout_does_not_bound_connector() {
         678  +
        let mut svc = ConnectionLimit::new(
         679  +
            NeverConnects::default(),
         680  +
            None,
         681  +
            None,
         682  +
            Arc::new(super::super::stats::ConnectionCounters::default()),
         683  +
            None,
         684  +
        );
         685  +
        let ctx = ConnectCtx::new("http://example.com".parse().unwrap(), None);
         686  +
        let fut = svc.call(ctx);
         687  +
        // A brief tokio yield shouldn't resolve the never-connects future.
         688  +
        tokio::pin!(fut);
         689  +
        tokio::select! {
         690  +
            _ = &mut fut => panic!("future should not resolve without a timeout"),
         691  +
            _ = tokio::task::yield_now() => {}
         692  +
        }
         693  +
    }
         694  +
}

tmp-codegen-diff/aws-sdk/sdk/aws-smithy-http-client/src/client/pool/partition.rs

@@ -0,1 +0,555 @@
           1  +
/*
           2  +
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
           3  +
 * SPDX-License-Identifier: Apache-2.0
           4  +
 */
           5  +
           6  +
//! Connection pool partitioning.
           7  +
//!
           8  +
//! A partition is a label that groups connections by locality: the runtime
           9  +
//! that owns their drivers and the network interface their sockets bind to.
          10  +
//! Partition labels are assigned through client configuration; the pool
          11  +
//! indexes connections by label and respects locality at checkout.
          12  +
//!
          13  +
//! # Topologies
          14  +
//!
          15  +
//! ## Single partition
          16  +
//!
          17  +
//! No partitioning. All connections pool together. The runtime is whatever
          18  +
//! was current at client construction; no NIC binding.
          19  +
//!
          20  +
//! ```text
          21  +
//! Pool
          22  +
//! └── Partition (anonymous, runtime=tokio-mt, nic=none)
          23  +
//!     ├── conn-1
          24  +
//!     └── conn-2
          25  +
//! ```
          26  +
//!
          27  +
//! ## Per-runtime, no NIC binding
          28  +
//!
          29  +
//! N partitions, one per runtime. Each partition's connections have
          30  +
//! drivers on that partition's runtime.
          31  +
//!
          32  +
//! ```text
          33  +
//! Pool
          34  +
//! ├── Partition 0 (runtime=tokio-current R0, nic=none)
          35  +
//! │   ├── conn-1 (driver on R0)
          36  +
//! │   └── conn-2 (driver on R0)
          37  +
//! └── Partition 1 (runtime=tokio-current R1, nic=none)
          38  +
//!     └── conn-3 (driver on R1)
          39  +
//! ```
          40  +
//!
          41  +
//! ## Per-runtime, per-NIC
          42  +
//!
          43  +
//! Partitions cluster by NIC. A socket bound to one NIC cannot serve
          44  +
//! traffic on another.
          45  +
//!
          46  +
//! ```text
          47  +
//! Pool
          48  +
//! ├── Partition 0 (runtime=R0, nic=eth0) ─┐
          49  +
//! ├── Partition 1 (runtime=R1, nic=eth0) ─┴─ same NIC group
          50  +
//! ├── Partition 2 (runtime=R2, nic=eth1) ─┐
          51  +
//! └── Partition 3 (runtime=R3, nic=eth1) ─┴─ same NIC group
          52  +
//! ```
          53  +
//!
          54  +
//! ## NUMA-aware
          55  +
//!
          56  +
//! Partitions align with NUMA topology: runtimes pinned to cores on a
          57  +
//! node, NIC selected to match the node. The pool does not detect NUMA
          58  +
//! topology; it sees only `(PartitionId, runtime, nic)` and the
          59  +
//! alignment is established when clients are configured.
          60  +
//!
          61  +
//! ```text
          62  +
//! Pool
          63  +
//! ├── NUMA node 0
          64  +
//! │   ├── Partition 0 (runtime=R0 on core 0,  nic=eth0)
          65  +
//! │   ├── Partition 1 (runtime=R1 on core 1,  nic=eth0)
          66  +
//! │   └── Partition 2 (runtime=R2 on core 2,  nic=eth1)
          67  +
//! └── NUMA node 1
          68  +
//!     ├── Partition 3 (runtime=R3 on core 32, nic=eth2)
          69  +
//!     └── Partition 4 (runtime=R4 on core 33, nic=eth3)
          70  +
//! ```
          71  +
//!
          72  +
//! # Boundaries
          73  +
//!
          74  +
//! - **NIC (hard):** a connection's NIC binding is fixed at creation.
          75  +
//!   Connections form NIC groups: one group per `Some(nic)` value plus
          76  +
//!   an unbound group for `None`. The pool only returns a connection
          77  +
//!   to a checkout in the same NIC group; the unbound group is not a
          78  +
//!   wildcard.
          79  +
//! - **Runtime (soft):** a connection's driver runs on a specific
          80  +
//!   runtime. Cross-runtime checkout is feasible (the request is
          81  +
//!   dispatched through the driver's runtime via a channel) but costs
          82  +
//!   a cross-thread send. [`CrossPartitionPolicy`] controls whether
          83  +
//!   the pool crosses this boundary at checkout.
          84  +
//!
          85  +
//! # Checkout
          86  +
//!
          87  +
//! When a request arrives on partition P for authority A:
          88  +
//!
          89  +
//! 1. **Local hit:** an idle connection in P for A is reused.
          90  +
//! 2. **Local miss, under capacity:** a new connection is created
          91  +
//!    in P.
          92  +
//! 3. **Local miss, at capacity:** behavior depends on
          93  +
//!    [`CrossPartitionPolicy`].
          94  +
//!
          95  +
//! Cross-partition borrowing is a capacity-pressure fallback. Under
          96  +
//! normal load each partition creates its own connections.
          97  +
//!
          98  +
//! ## Policy: `Never`
          99  +
//!
         100  +
//! At capacity with no local idle, the request waits for a permit.
         101  +
//! Peer partitions are not consulted.
         102  +
//!
         103  +
//! ```text
         104  +
//! Capacity = 2, both in use, request arrives on P0 for A:
         105  +
//!
         106  +
//!   P0 (eth0): [active to B]   ← request for A queues here
         107  +
//!   P1 (eth0): [idle to A]     ← not consulted
         108  +
//!
         109  +
//! Outcome: request blocks until a permit is available, then either
         110  +
//! reuses a returning idle for A or creates a new connection on P0.
         111  +
//! ```
         112  +
//!
         113  +
//! ## Policy: `PreferLocal`
         114  +
//!
         115  +
//! At capacity with no local idle, the pool checks peer partitions in
         116  +
//! the same NIC group for an idle connection to the requested
         117  +
//! authority. If one is found, the request borrows it. Otherwise it
         118  +
//! waits for a permit.
         119  +
//!
         120  +
//! ```text
         121  +
//! Capacity = 2, both in use, request arrives on P0 for A:
         122  +
//!
         123  +
//!   P0 (eth0): [active to B]   ← request for A
         124  +
//!   P1 (eth0): [idle to A]     ← borrowed
         125  +
//!   P2 (eth1): [idle to A]     ← different NIC, never consulted
         126  +
//!
         127  +
//! Outcome: P1's idle connection serves the request. The connection's
         128  +
//! driver stays on P1's runtime; the request flows through P1's
         129  +
//! runtime via the connection's channel.
         130  +
//! ```
         131  +
         132  +
/// Identifier for a pool partition.
         133  +
///
         134  +
/// A partition's identity is opaque to the pool. The identifier is
         135  +
/// assigned at client construction and groups connections that share
         136  +
/// a driver spawner and network interface binding.
         137  +
///
         138  +
/// The default identifier denotes an anonymous partition used when
         139  +
/// partitioning is not configured.
         140  +
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
         141  +
pub struct PartitionId(u64);
         142  +
         143  +
impl PartitionId {
         144  +
    const ANONYMOUS: u64 = u64::MAX;
         145  +
         146  +
    /// Identifier from a numeric index.
         147  +
    pub const fn from_index(index: usize) -> Self {
         148  +
        Self(index as u64)
         149  +
    }
         150  +
         151  +
    /// Identifier from a raw value. The value `u64::MAX` is reserved
         152  +
    /// for the anonymous default partition.
         153  +
    pub const fn from_raw(id: u64) -> Self {
         154  +
        Self(id)
         155  +
    }
         156  +
         157  +
    /// Raw value of this identifier.
         158  +
    pub const fn as_u64(self) -> u64 {
         159  +
        self.0
         160  +
    }
         161  +
}
         162  +
         163  +
impl Default for PartitionId {
         164  +
    fn default() -> Self {
         165  +
        Self(Self::ANONYMOUS)
         166  +
    }
         167  +
}
         168  +
         169  +
/// Spawner for connection driver tasks.
         170  +
///
         171  +
/// A driver is the task that owns an HTTP connection's I/O state machine:
         172  +
/// reading frames, writing frames, and managing protocol-level events.
         173  +
/// Calls through a connection's request handle flow through this driver.
         174  +
/// The pool spawns one driver per established connection.
         175  +
///
         176  +
/// Different partitions may use different spawners, allowing each
         177  +
/// partition's drivers to run on a specific runtime.
         178  +
pub trait DriverSpawner: std::fmt::Debug + Send + Sync + 'static {
         179  +
    /// Spawn the connection driver future on this spawner's runtime.
         180  +
    fn spawn(
         181  +
        &self,
         182  +
        driver: std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + 'static>>,
         183  +
    );
         184  +
}
         185  +
         186  +
/// Driver spawner backed by a tokio runtime handle.
         187  +
///
         188  +
/// Spawns the connection driver via [`tokio::runtime::Handle::spawn`].
         189  +
/// The handle is captured at construction; the driver runs on the
         190  +
/// runtime the handle refers to, regardless of which runtime called
         191  +
/// [`DriverSpawner::spawn`].
         192  +
#[derive(Clone, Debug)]
         193  +
pub struct TokioDriverSpawner {
         194  +
    handle: tokio::runtime::Handle,
         195  +
}
         196  +
         197  +
impl TokioDriverSpawner {
         198  +
    /// Spawner using the current tokio runtime handle (captured eagerly).
         199  +
    ///
         200  +
    /// Panics if invoked outside a tokio runtime context.
         201  +
    pub fn current() -> Self {
         202  +
        Self::from_handle(tokio::runtime::Handle::current())
         203  +
    }
         204  +
         205  +
    /// Spawner using a specific tokio runtime handle.
         206  +
    pub fn from_handle(handle: tokio::runtime::Handle) -> Self {
         207  +
        Self { handle }
         208  +
    }
         209  +
}
         210  +
         211  +
impl DriverSpawner for TokioDriverSpawner {
         212  +
    fn spawn(
         213  +
        &self,
         214  +
        driver: std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + 'static>>,
         215  +
    ) {
         216  +
        self.handle.spawn(driver);
         217  +
    }
         218  +
}
         219  +
         220  +
/// Policy governing checkout when the local partition has no idle
         221  +
/// connection and the pool is at capacity.
         222  +
///
         223  +
/// Cross-partition borrowing applies within a NIC group only.
         224  +
/// Connections bound to different NICs are never shared regardless
         225  +
/// of policy.
         226  +
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
         227  +
#[non_exhaustive]
         228  +
pub enum CrossPartitionPolicy {
         229  +
    /// At capacity with no local idle, the request waits for a permit.
         230  +
    /// Peer partitions are not consulted.
         231  +
    #[default]
         232  +
    Never,
         233  +
    /// At capacity with no local idle, the request borrows an idle
         234  +
    /// connection from a peer partition in the same NIC group when one
         235  +
    /// is available, otherwise waits for a permit.
         236  +
    PreferLocal,
         237  +
}
         238  +
         239  +
/// A declared pool partition: a driver-spawner runtime and an optional
         240  +
/// NIC binding, identified by a caller-owned [`PartitionId`]. Declared on
         241  +
/// the pool builder via `Builder::partitions`; the pool owns the topology
         242  +
/// for its lifetime.
         243  +
#[derive(Clone, Debug)]
         244  +
pub struct Partition {
         245  +
    pub(super) id: PartitionId,
         246  +
    pub(super) spawner: std::sync::Arc<dyn DriverSpawner>,
         247  +
    pub(super) nic: Option<String>,
         248  +
}
         249  +
         250  +
impl Partition {
         251  +
    /// Declare a partition with the given id and driver spawner.
         252  +
    pub fn new<S: DriverSpawner>(id: PartitionId, spawner: S) -> Self {
         253  +
        Self {
         254  +
            id,
         255  +
            spawner: std::sync::Arc::new(spawner),
         256  +
            nic: None,
         257  +
        }
         258  +
    }
         259  +
         260  +
    /// Bind this partition's connections to a network interface.
         261  +
    pub fn interface(mut self, nic: impl Into<String>) -> Self {
         262  +
        self.nic = Some(nic.into());
         263  +
        self
         264  +
    }
         265  +
}
         266  +
         267  +
/// Pool-owned state for one declared partition. Resolved once at pool
         268  +
/// build time and referenced by [`Client`](super::Client) handles.
         269  +
pub(crate) struct PartitionState {
         270  +
    pub(crate) id: PartitionId,
         271  +
    // Captured into `make_stack` by the build factory; the field is retained
         272  +
    // on the state but read through the captured closure, not directly.
         273  +
    #[allow(dead_code)]
         274  +
    pub(crate) spawner: std::sync::Arc<dyn DriverSpawner>,
         275  +
    /// Network interface this partition's connections bind to, and the
         276  +
    /// boundary for cross-partition borrow and reclaim (peers in the same
         277  +
    /// NIC group only).
         278  +
    pub(crate) nic: Option<String>,
         279  +
    /// Per-host connection storage for this partition. Keyed by
         280  +
    /// (scheme, authority); entries built lazily on first request.
         281  +
    pub(crate) authorities:
         282  +
        std::sync::Mutex<std::collections::HashMap<super::PoolKey, Box<dyn super::PoolEntry>>>,
         283  +
    /// Builds a host entry on first touch, capturing this partition's
         284  +
    /// connector; shared budget/hooks arrive via `&SharedPoolState`.
         285  +
    pub(crate) make_stack: super::MakeStack,
         286  +
    /// Round-robin cursor over reclaim/borrow candidate peers. Advisory
         287  +
    /// (`Relaxed` `fetch_add`): rotates the starting offset into the
         288  +
    /// candidate set so concurrent cap-bound reclaims from this partition
         289  +
    /// do not all probe the lowest-numbered peer first. No correctness
         290  +
    /// invariant rides it — `try_reclaim_one` is the authoritative gate.
         291  +
    pub(crate) peer_cursor: std::sync::atomic::AtomicUsize,
         292  +
}
         293  +
         294  +
impl std::fmt::Debug for PartitionState {
         295  +
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         296  +
        f.debug_struct("PartitionState")
         297  +
            .field("id", &self.id)
         298  +
            .field("nic", &self.nic)
         299  +
            .finish_non_exhaustive()
         300  +
    }
         301  +
}
         302  +
         303  +
/// Normalize a caller-declared partition set for pool construction: when
         304  +
/// no partitions are declared, synthesize a single anonymous partition
         305  +
/// (`PartitionId::default()`, no NIC binding) using `anonymous_spawner`.
         306  +
/// Returns the caller's set unchanged when non-empty.
         307  +
///
         308  +
/// This is the one place the "no topology declared" default is decided;
         309  +
/// [`PartitionRegistry::build`] then indexes whatever set it is given.
         310  +
/// `anonymous_spawner` is a closure so the runtime handle is captured
         311  +
/// only when actually needed (e.g. `TokioDriverSpawner::current()` panics
         312  +
/// off a runtime).
         313  +
pub(crate) fn normalize_partitions(
         314  +
    partitions: Vec<Partition>,
         315  +
    anonymous_spawner: impl FnOnce() -> std::sync::Arc<dyn DriverSpawner>,
         316  +
) -> Vec<Partition> {
         317  +
    if partitions.is_empty() {
         318  +
        vec![Partition {
         319  +
            id: PartitionId::default(),
         320  +
            spawner: anonymous_spawner(),
         321  +
            nic: None,
         322  +
        }]
         323  +
    } else {
         324  +
        partitions
         325  +
    }
         326  +
}
         327  +
         328  +
/// Immutable registry of declared partitions, built once at pool
         329  +
/// construction. Maps ids and NIC groups to partition state and records
         330  +
/// the default partition used by [`Client::new`](super::Client::new).
         331  +
#[derive(Debug)]
         332  +
pub(crate) struct PartitionRegistry {
         333  +
    by_id: std::collections::HashMap<PartitionId, std::sync::Arc<PartitionState>>,
         334  +
    /// Partition ids grouped by NIC, for the cross-partition borrow and
         335  +
    /// reclaim peer walk (candidates are drawn from the requester's NIC group).
         336  +
    by_nic: std::collections::HashMap<Option<String>, Vec<PartitionId>>,
         337  +
    default_partition: PartitionId,
         338  +
}
         339  +
         340  +
impl PartitionRegistry {
         341  +
    /// Build a registry from a non-empty set of declared partitions. The
         342  +
    /// default partition is the first in the slice. Panics on a duplicate
         343  +
    /// `PartitionId`, or if `partitions` is empty (callers normalize the
         344  +
    /// no-topology case via [`normalize_partitions`] first).
         345  +
    pub(crate) fn build(
         346  +
        partitions: Vec<Partition>,
         347  +
        make_stack_for: impl Fn(&Partition) -> super::MakeStack,
         348  +
    ) -> Self {
         349  +
        assert!(
         350  +
            !partitions.is_empty(),
         351  +
            "PartitionRegistry::build requires at least one partition; \
         352  +
             normalize the empty case with normalize_partitions"
         353  +
        );
         354  +
        let default_partition = partitions[0].id;
         355  +
        let mut by_id = std::collections::HashMap::new();
         356  +
        let mut by_nic: std::collections::HashMap<Option<String>, Vec<PartitionId>> =
         357  +
            std::collections::HashMap::new();
         358  +
        for p in partitions {
         359  +
            by_nic.entry(p.nic.clone()).or_default().push(p.id);
         360  +
            let make_stack = make_stack_for(&p);
         361  +
            let state = std::sync::Arc::new(PartitionState {
         362  +
                id: p.id,
         363  +
                spawner: p.spawner,
         364  +
                nic: p.nic,
         365  +
                authorities: std::sync::Mutex::new(std::collections::HashMap::new()),
         366  +
                make_stack,
         367  +
                peer_cursor: std::sync::atomic::AtomicUsize::new(0),
         368  +
            });
         369  +
            if by_id.insert(p.id, state).is_some() {
         370  +
                panic!("duplicate PartitionId declared: {:?}", p.id);
         371  +
            }
         372  +
        }
         373  +
        Self {
         374  +
            by_id,
         375  +
            by_nic,
         376  +
            default_partition,
         377  +
        }
         378  +
    }
         379  +
         380  +
    /// Resolve the default partition (first declared or anonymous).
         381  +
    pub(crate) fn default_partition(&self) -> std::sync::Arc<PartitionState> {
         382  +
        self.by_id
         383  +
            .get(&self.default_partition)
         384  +
            .expect("default partition exists")
         385  +
            .clone()
         386  +
    }
         387  +
         388  +
    /// Resolve a declared partition by id. Panics if the id was not
         389  +
    /// declared (programming error: the caller declared the topology).
         390  +
    pub(crate) fn partition(&self, id: PartitionId) -> std::sync::Arc<PartitionState> {
         391  +
        self.by_id
         392  +
            .get(&id)
         393  +
            .unwrap_or_else(|| panic!("partition not declared: {:?}", id))
         394  +
            .clone()
         395  +
    }
         396  +
         397  +
    /// Iterate all declared partitions.
         398  +
    pub(crate) fn partitions(&self) -> impl Iterator<Item = &std::sync::Arc<PartitionState>> {
         399  +
        self.by_id.values()
         400  +
    }
         401  +
         402  +
    /// Resolve a partition by id without panicking. `None` if not declared.
         403  +
    pub(crate) fn partition_opt(&self, id: PartitionId) -> Option<&std::sync::Arc<PartitionState>> {
         404  +
        self.by_id.get(&id)
         405  +
    }
         406  +
         407  +
    /// Partition ids sharing `id`'s NIC group, excluding `id` itself.
         408  +
    ///
         409  +
    /// Reclaim is NIC-blind for the freed *permit* (P0 connects on its own
         410  +
    /// NIC), but candidate peers are still drawn from the same NIC group:
         411  +
    /// the registry only groups by NIC, and a freed permit from any
         412  +
    /// same-group peer is equivalent. Empty if `id` is alone in its group
         413  +
    /// (e.g. the single-partition default).
         414  +
    pub(crate) fn nic_group_peers(&self, id: PartitionId) -> Vec<PartitionId> {
         415  +
        let nic = match self.by_id.get(&id) {
         416  +
            Some(state) => &state.nic,
         417  +
            None => return Vec::new(),
         418  +
        };
         419  +
        self.by_nic
         420  +
            .get(nic)
         421  +
            .map(|ids| ids.iter().copied().filter(|p| *p != id).collect())
         422  +
            .unwrap_or_default()
         423  +
    }
         424  +
         425  +
    /// Attempt to reclaim one idle connection from `peer`'s entry for
         426  +
    /// `key`, freeing its permit. Returns `true` if one was freed. The idle
         427  +
    /// connection is popped under the cache lock and dropped after the lock
         428  +
    /// is released. No-op `false` if the peer or entry is absent.
         429  +
    pub(crate) fn try_reclaim_on(&self, peer: PartitionId, key: &super::PoolKey) -> bool {
         430  +
        let state = match self.by_id.get(&peer) {
         431  +
            Some(s) => s,
         432  +
            None => return false,
         433  +
        };
         434  +
        let auth = state.authorities.lock().expect("authorities poisoned");
         435  +
        match auth.get(key) {
         436  +
            Some(entry) => entry.try_reclaim_one(),
         437  +
            None => false,
         438  +
        }
         439  +
    }
         440  +
         441  +
    /// Attempt to reclaim one idle connection from *any* of `peer`'s
         442  +
    /// entries, freeing its permit. Returns `true` at the first entry that
         443  +
    /// yields. Drives the `Global` constraint, where the freed permit is
         444  +
    /// fungible across authorities — so the specific authority does not
         445  +
    /// matter, and this sidesteps reconstructing a `PoolKey` (scheme +
         446  +
    /// authority) from the authority-only stats index. No-op `false` if
         447  +
    /// the peer is absent or holds no reclaimable idle.
         448  +
    pub(crate) fn try_reclaim_any(&self, peer: PartitionId) -> bool {
         449  +
        let state = match self.by_id.get(&peer) {
         450  +
            Some(s) => s,
         451  +
            None => return false,
         452  +
        };
         453  +
        let auth = state.authorities.lock().expect("authorities poisoned");
         454  +
        auth.values().any(|entry| entry.try_reclaim_one())
         455  +
    }
         456  +
         457  +
    /// Attempt to borrow one idle connection from `peer`'s entry for
         458  +
    /// `key`, as a dispatchable handle that returns to `peer`'s pool on
         459  +
    /// drop. Returns `None` if the peer or entry is absent, or holds no
         460  +
    /// borrowable idle. The peer's `authorities` lock is held only to
         461  +
    /// check out the handle (the handle re-pools on drop, independent of
         462  +
    /// the lock); dispatch happens after the lock is released.
         463  +
    pub(crate) fn try_borrow_on(
         464  +
        &self,
         465  +
        peer: PartitionId,
         466  +
        key: &super::PoolKey,
         467  +
    ) -> Option<Box<dyn super::DispatchConn>> {
         468  +
        let state = self.by_id.get(&peer)?;
         469  +
        let auth = state.authorities.lock().expect("authorities poisoned");
         470  +
        auth.get(key)?.try_borrow_one()
         471  +
    }
         472  +
}
         473  +
         474  +
#[cfg(test)]
         475  +
mod tests {
         476  +
    use super::*;
         477  +
         478  +
    #[test]
         479  +
    fn partition_id_from_index() {
         480  +
        assert_eq!(PartitionId::from_index(0).as_u64(), 0);
         481  +
        assert_eq!(PartitionId::from_index(5).as_u64(), 5);
         482  +
        assert_eq!(PartitionId::from_index(0), PartitionId::from_index(0));
         483  +
        assert_ne!(PartitionId::from_index(0), PartitionId::from_index(1));
         484  +
    }
         485  +
         486  +
    #[test]
         487  +
    fn partition_id_default_is_anonymous() {
         488  +
        assert_eq!(PartitionId::default().as_u64(), u64::MAX);
         489  +
        assert_ne!(PartitionId::default(), PartitionId::from_index(0));
         490  +
    }
         491  +
         492  +
    #[test]
         493  +
    fn cross_partition_policy_default_is_never() {
         494  +
        assert_eq!(CrossPartitionPolicy::default(), CrossPartitionPolicy::Never);
         495  +
    }
         496  +
         497  +
    #[tokio::test]
         498  +
    async fn normalize_partitions_synthesizes_anonymous_when_empty() {
         499  +
        // Empty input → exactly one anonymous partition on the supplied spawner.
         500  +
        let parts = normalize_partitions(Vec::new(), || {
         501  +
            std::sync::Arc::new(TokioDriverSpawner::current()) as std::sync::Arc<dyn DriverSpawner>
         502  +
        });
         503  +
        assert_eq!(parts.len(), 1);
         504  +
        assert_eq!(parts[0].id, PartitionId::default());
         505  +
        assert!(parts[0].nic.is_none());
         506  +
    }
         507  +
         508  +
    #[tokio::test]
         509  +
    async fn normalize_partitions_passes_declared_set_through_untouched() {
         510  +
        // Non-empty input is returned unchanged, and the anonymous-spawner
         511  +
        // closure is never invoked.
         512  +
        let declared = vec![
         513  +
            Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()),
         514  +
            Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()),
         515  +
        ];
         516  +
        let parts = normalize_partitions(declared, || {
         517  +
            panic!("anonymous spawner must not be called when partitions are declared")
         518  +
        });
         519  +
        assert_eq!(parts.len(), 2);
         520  +
        assert_eq!(parts[0].id, PartitionId::from_index(0));
         521  +
        assert_eq!(parts[1].id, PartitionId::from_index(1));
         522  +
    }
         523  +
         524  +
    #[tokio::test]
         525  +
    async fn tokio_driver_spawner_current() {
         526  +
        let _ = TokioDriverSpawner::current();
         527  +
    }
         528  +
         529  +
    #[tokio::test]
         530  +
    async fn tokio_driver_spawner_from_handle() {
         531  +
        let h = tokio::runtime::Handle::current();
         532  +
        let _ = TokioDriverSpawner::from_handle(h);
         533  +
    }
         534  +
         535  +
    #[tokio::test]
         536  +
    async fn tokio_driver_spawner_runs_future() {
         537  +
        use std::sync::atomic::{AtomicBool, Ordering};
         538  +
        use std::sync::Arc;
         539  +
         540  +
        let sp = TokioDriverSpawner::current();
         541  +
        let flag = Arc::new(AtomicBool::new(false));
         542  +
        let f = flag.clone();
         543  +
        sp.spawn(Box::pin(async move {
         544  +
            f.store(true, Ordering::SeqCst);
         545  +
        }));
         546  +
        // Yield enough times for the spawned task to run.
         547  +
        for _ in 0..10 {
         548  +
            tokio::task::yield_now().await;
         549  +
            if flag.load(Ordering::SeqCst) {
         550  +
                break;
         551  +
            }
         552  +
        }
         553  +
        assert!(flag.load(Ordering::SeqCst), "spawned future did not run");
         554  +
    }
         555  +
}