AWS SDK

AWS SDK

rev. 01140188efe3be26ec6653bc30094c80ed1a1b99 (ignoring whitespace)

Files changed:

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

@@ -1,1 +142,143 @@
   10     10   
//! cell, and submits the protocol driver to the connection-owning partition.
   11     11   
   12     12   
use super::super::cell::h1::{H1CloseHandle, H1DriverGuard, H1Selection, H1Sender};
   13     13   
use super::super::cell::{EstablishmentPermit, OriginCell};
   14     14   
use super::super::connection::{
   15     15   
    CloseReason, ConnectionInfo, ConnectionIo, ConnectionState, NegotiatedProtocol,
   16     16   
};
   17     17   
use super::super::dispatch::AcquisitionContext;
   18     18   
use super::next_connection_id;
   19     19   
use crate::client::connect::BoxConn;
          20  +
use crate::client::downcast_error;
   20     21   
use aws_smithy_runtime_api::client::result::ConnectorError;
   21     22   
use aws_smithy_types::body::SdkBody;
   22     23   
use hyper_util::client::legacy::connect::Connected;
   23     24   
   24     25   
/// Handshakes, installs, and starts the owner-partition driver.
   25     26   
///
   26     27   
/// `permit` remains the sole owner of optional bounded capacity until the
   27     28   
/// connection state is created. An earlier failure returns that capacity
   28     29   
/// through the permit's fallback. A handshake failure logically closes the
   29     30   
/// new connection state, while success transfers driver ownership to the
   30     31   
/// partition runtime and returns the installed sender selection.
   31     32   
pub(in crate::client::pool) async fn establish_h1(
   32     33   
    context: AcquisitionContext,
   33     34   
    permit: EstablishmentPermit,
   34     35   
    io: BoxConn,
   35     36   
    connected: Connected,
   36     37   
) -> Result<H1Selection, ConnectorError> {
   37     38   
    let request_partition = context.partition.id();
   38     39   
    let connection_cell = context.cell.clone();
   39     40   
    tracing::debug!(
   40     41   
        request_partition = ?request_partition,
   41     42   
        connection_partition = ?connection_cell.id().partition(),
   42     43   
        origin_scheme = %connection_cell.id().origin().scheme(),
   43     44   
        origin_host = connection_cell.id().origin().host(),
   44     45   
        origin_port = ?connection_cell.id().origin().port(),
   45     46   
        "HTTP/1 connection establishment started"
   46     47   
    );
   47     48   
   48         -
    let result = handshake_and_install_h1(context, permit, io, connected).await;
          49  +
    let result = run_h1_handshake(context, permit, io, connected).await;
   49     50   
    match &result {
   50     51   
        Ok(selection) => {
   51     52   
            let connection = selection.connection();
   52     53   
            tracing::debug!(
   53     54   
                connection_id = %connection.id(),
   54     55   
                request_partition = ?request_partition,
   55     56   
                connection_partition = ?connection.owner_partition(),
   56     57   
                origin_scheme = %connection.info().origin().scheme(),
   57     58   
                origin_host = connection.info().origin().host(),
   58     59   
                origin_port = ?connection.info().origin().port(),
   59     60   
                "HTTP/1 connection established"
   60     61   
            );
   61     62   
        }
   62     63   
        Err(error) => tracing::debug!(
   63     64   
            request_partition = ?request_partition,
   64     65   
            connection_partition = ?connection_cell.id().partition(),
   65     66   
            origin_scheme = %connection_cell.id().origin().scheme(),
   66     67   
            origin_host = connection_cell.id().origin().host(),
   67     68   
            origin_port = ?connection_cell.id().origin().port(),
   68     69   
            error = ?error,
   69     70   
            "HTTP/1 connection establishment failed"
   70     71   
        ),
   71     72   
    }
   72     73   
    result
   73     74   
}
   74     75   
   75     76   
/// Handshakes and installs one already connected HTTP/1 transport.
   76         -
async fn handshake_and_install_h1(
          77  +
async fn run_h1_handshake(
   77     78   
    context: AcquisitionContext,
   78     79   
    permit: EstablishmentPermit,
   79     80   
    io: BoxConn,
   80     81   
    connected: Connected,
   81     82   
) -> Result<H1Selection, ConnectorError> {
   82     83   
    let AcquisitionContext {
   83     84   
        pool,
   84     85   
        partition,
   85     86   
        cell,
   86     87   
        absolute_uri: _,
   87     88   
        owner_spawner,
   88     89   
        connect_timeout: _,
   89     90   
    } = context;
   90     91   
   91     92   
    let id =
   92     93   
        next_connection_id(&pool).map_err(|error| ConnectorError::other(error.into(), None))?;
   93     94   
    let info = ConnectionInfo::new(
   94     95   
        id,
   95     96   
        cell.id().origin().clone(),
   96     97   
        partition.id(),
   97     98   
        NegotiatedProtocol::Http1,
   98     99   
        connected,
   99    100   
    );
  100    101   
    let (connection, physical) = ConnectionState::pending_open(info);
  101    102   
    let io = ConnectionIo::new(io, physical);
  102    103   
  103    104   
    let (sender, driver) = match hyper::client::conn::http1::Builder::new()
  104    105   
        .handshake::<_, SdkBody>(io)
  105    106   
        .await
  106    107   
    {
  107    108   
        Ok(established) => established,
  108    109   
        Err(error) => {
  109    110   
            connection.logical_close(CloseReason::ProtocolClosed);
  110         -
            return Err(super::super::super::downcast_error(Box::new(error)));
         111  +
            return Err(downcast_error(Box::new(error)));
  111    112   
        }
  112    113   
    };
  113    114   
  114    115   
    if let Err(lease) = connection.open(permit.into_lease()) {
  115    116   
        drop(lease);
  116    117   
        connection.logical_close(CloseReason::ProtocolClosed);
  117    118   
        return Err(ConnectorError::io(
  118    119   
            "HTTP/1 connection closed before installation".into(),
  119    120   
        ));
  120    121   
    }
  121    122   
  122    123   
    let selection =
  123         -
        OriginCell::install_selected_h1(&cell, connection.clone(), H1Sender::from_hyper(sender));
         124  +
        OriginCell::insert_selected_h1(&cell, connection.clone(), H1Sender::from_hyper(sender));
  124    125   
    let driver_guard = H1DriverGuard::new(H1CloseHandle::new(&cell, &connection));
  125    126   
    let driver_info = connection.info().clone();
  126    127   
    owner_spawner.spawn(Box::pin(async move {
  127    128   
        let result = driver.with_upgrades().await;
  128    129   
        if let Err(error) = result {
  129    130   
            tracing::debug!(
  130    131   
                connection_id = %driver_info.id(),
  131    132   
                connection_partition = ?driver_info.owner_partition(),
  132    133   
                origin_scheme = %driver_info.origin().scheme(),
  133    134   
                origin_host = driver_info.origin().host(),

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

@@ -1,1 +58,59 @@
    9      9   
//! accepting generation, join the current flight, or become the flight's
   10     10   
//! owner task. Only the owner performs Hyper's HTTP/2 handshake. It opens the
   11     11   
//! connection, installs the generation, submits the driver on the partition
   12     12   
//! spawner, and then disarms flight cleanup. Participants own only their
   13     13   
//! waiter entries; losing attempts close their unused transport and return
   14     14   
//! their capacity before this function reports transfer. The connector timeout
   15     15   
//! ends before this phase; no separate pool timeout covers the HTTP/2
   16     16   
//! handshake or driver submission.
   17     17   
   18     18   
use super::super::cell::h2::{
   19         -
    H2CloseHandle, H2DriverGuard, H2FlightId, H2FlightInstall, H2GenerationJoin, H2Sender,
          19  +
    H2CloseHandle, H2DriverGuard, H2FlightDecision, H2FlightId, H2GenerationJoinOutcome, H2Sender,
   20     20   
};
   21         -
use super::super::cell::{AcquisitionResult, EstablishmentPermit, OriginCell, WaiterId};
          21  +
use super::super::cell::{AcquisitionOutcome, EstablishmentPermit, OriginCell, WaiterId};
   22     22   
use super::super::connection::{
   23     23   
    CloseReason, ConnectionInfo, ConnectionIo, ConnectionState, NegotiatedProtocol,
   24     24   
};
   25     25   
use super::super::dispatch::AcquisitionContext;
   26     26   
use super::super::partition::DriverSpawner;
   27     27   
use super::{next_connection_id, EstablishmentOutcome};
   28     28   
use crate::client::connect::BoxConn;
          29  +
use crate::client::downcast_error;
   29     30   
use aws_smithy_runtime_api::box_error::BoxError;
   30     31   
use aws_smithy_runtime_api::client::result::ConnectorError;
   31     32   
use aws_smithy_types::body::SdkBody;
   32     33   
use aws_smithy_types::retry::ErrorKind;
   33     34   
use hyper::rt::Executor;
   34     35   
use hyper_util::client::legacy::connect::Connected;
   35     36   
use std::future::Future;
   36     37   
use std::sync::Arc as StdArc;
   37     38   
   38     39   
/// Hyper executor that keeps spawned HTTP/2 work on the connection partition.
@@ -76,77 +221,222 @@
   96     97   
   97     98   
/// Converges after ALPN and performs at most one HTTP/2 handshake per cell.
   98     99   
pub(super) async fn establish_h2(
   99    100   
    context: AcquisitionContext,
  100    101   
    waiter: WaiterId,
  101    102   
    permit: EstablishmentPermit,
  102    103   
    io: BoxConn,
  103    104   
    connected: Connected,
  104    105   
) -> EstablishmentOutcome {
  105    106   
    loop {
  106         -
        match context.cell.install_or_join_h2_flight(waiter) {
  107         -
            H2FlightInstall::Accepting(generation) => {
         107  +
        match context.cell.converge_h2_flight(waiter) {
         108  +
            H2FlightDecision::UseGeneration(generation) => {
  108    109   
                tracing::trace!(
  109    110   
                    request_partition = ?context.partition.id(),
  110    111   
                    connection_partition = ?context.cell.id().partition(),
  111    112   
                    origin_scheme = %context.cell.id().origin().scheme(),
  112    113   
                    origin_host = context.cell.id().origin().host(),
  113    114   
                    origin_port = ?context.cell.id().origin().port(),
  114    115   
                    h2_generation = ?generation,
  115    116   
                    "HTTP/2 establishment found an accepting generation"
  116    117   
                );
  117    118   
                match OriginCell::join_h2_generation(&context.cell, waiter, generation) {
  118         -
                    H2GenerationJoin::GenerationChanged => continue,
  119         -
                    H2GenerationJoin::Joined | H2GenerationJoin::WaiterCompleted => {
         119  +
                    H2GenerationJoinOutcome::GenerationChanged => continue,
         120  +
                    H2GenerationJoinOutcome::Joined | H2GenerationJoinOutcome::WaiterResolved => {
  120    121   
                        drop(io);
  121    122   
                        drop(permit);
  122         -
                        return EstablishmentOutcome::Transferred;
         123  +
                        return EstablishmentOutcome::WaiterCompletionTransferred;
  123    124   
                    }
  124    125   
                }
  125    126   
            }
  126         -
            H2FlightInstall::Joined => {
         127  +
            H2FlightDecision::JoinedFlight => {
  127    128   
                tracing::trace!(
  128    129   
                    request_partition = ?context.partition.id(),
  129    130   
                    connection_partition = ?context.cell.id().partition(),
  130    131   
                    origin_scheme = %context.cell.id().origin().scheme(),
  131    132   
                    origin_host = context.cell.id().origin().host(),
  132    133   
                    origin_port = ?context.cell.id().origin().port(),
  133    134   
                    "HTTP/2 establishment joined the active flight"
  134    135   
                );
  135    136   
                drop(io);
  136    137   
                drop(permit);
  137         -
                return EstablishmentOutcome::Transferred;
         138  +
                return EstablishmentOutcome::WaiterCompletionTransferred;
  138    139   
            }
  139         -
            H2FlightInstall::Driver(flight) => {
         140  +
            H2FlightDecision::RunFlight(flight) => {
  140    141   
                tracing::trace!(
  141    142   
                    request_partition = ?context.partition.id(),
  142    143   
                    connection_partition = ?context.cell.id().partition(),
  143    144   
                    origin_scheme = %context.cell.id().origin().scheme(),
  144    145   
                    origin_host = context.cell.id().origin().host(),
  145    146   
                    origin_port = ?context.cell.id().origin().port(),
  146    147   
                    h2_flight = ?flight,
  147    148   
                    "HTTP/2 establishment started a flight"
  148    149   
                );
  149    150   
                drive_flight(context, flight, permit, io, connected).await;
  150         -
                return EstablishmentOutcome::Transferred;
         151  +
                return EstablishmentOutcome::WaiterCompletionTransferred;
  151    152   
            }
  152    153   
        }
  153    154   
    }
  154    155   
}
  155    156   
  156    157   
/// Handshakes, installs, and publishes one winning flight.
  157    158   
async fn drive_flight(
  158    159   
    context: AcquisitionContext,
  159    160   
    flight: H2FlightId,
  160    161   
    permit: EstablishmentPermit,
  161    162   
    io: BoxConn,
  162    163   
    connected: Connected,
  163    164   
) {
  164    165   
    let mut completion = FlightCompletionGuard::new(context.cell.clone(), flight);
  165    166   
    let id = match next_connection_id(&context.pool) {
  166    167   
        Ok(id) => id,
  167    168   
        Err(error) => {
  168    169   
            completion.fail(ConnectorError::other(Box::new(error), None));
  169    170   
            return;
  170    171   
        }
  171    172   
    };
  172    173   
    let info = ConnectionInfo::new(
  173    174   
        id,
  174    175   
        context.cell.id().origin().clone(),
  175    176   
        context.partition.id(),
  176    177   
        NegotiatedProtocol::Http2,
  177    178   
        connected,
  178    179   
    );
  179    180   
    let (connection, physical) = ConnectionState::pending_open(info);
  180    181   
    let io = ConnectionIo::new(io, physical);
  181    182   
    let executor = PartitionExecutor {
  182    183   
        spawner: context.owner_spawner.clone(),
  183    184   
    };
  184    185   
    let (sender, driver) = match hyper::client::conn::http2::Builder::new(executor)
  185    186   
        .handshake::<_, SdkBody>(io)
  186    187   
        .await
  187    188   
    {
  188    189   
        Ok(established) => established,
  189    190   
        Err(error) => {
  190    191   
            connection.logical_close(CloseReason::ProtocolClosed);
  191         -
            completion.fail(super::super::super::downcast_error(Box::new(error)));
         192  +
            completion.fail(downcast_error(Box::new(error)));
  192    193   
            return;
  193    194   
        }
  194    195   
    };
  195    196   
  196    197   
    if let Err(lease) = connection.open(permit.into_lease()) {
  197    198   
        drop(lease);
  198    199   
        connection.logical_close(CloseReason::ProtocolClosed);
  199    200   
        completion.fail(ConnectorError::io(Box::new(std::io::Error::other(
  200    201   
            "HTTP/2 connection closed before installation",
  201    202   
        ))));
@@ -296,297 +447,448 @@
  316    317   
  317    318   
/// Fails every participant still retained by one exact flight.
  318    319   
fn fail_participants(cell: &OriginCell, flight: H2FlightId, error: ConnectorError) {
  319    320   
    let Some(participants) = cell.fail_h2_flight(flight) else {
  320    321   
        return;
  321    322   
    };
  322    323   
    let error = SharedFlightFailure::new(error);
  323    324   
    for participant in participants {
  324    325   
        cell.complete_establishment(
  325    326   
            participant,
  326         -
            AcquisitionResult::Failed(error.connector_error()),
         327  +
            AcquisitionOutcome::Failed(error.connector_error()),
  327    328   
        );
  328    329   
    }
  329    330   
}
  330    331   
  331    332   
#[cfg(all(test, not(smithy_http_client_loom)))]
  332    333   
mod tests {
  333    334   
    use super::*;
  334    335   
    use crate::client::pool::admission::ProtocolRequirement;
  335         -
    use crate::client::pool::cell::AcquisitionEvent;
         336  +
    use crate::client::pool::cell::{AcquisitionOutcome, AcquisitionStep};
  336    337   
    use crate::client::pool::origin::OriginKey;
  337    338   
    use crate::client::pool::partition::EligibilityGroup;
  338    339   
    use http_1x::uri::Scheme;
  339    340   
    use std::error::Error as _;
  340    341   
    use std::sync::atomic::{AtomicUsize, Ordering};
  341    342   
    use std::task::{Context, Poll, Waker};
  342    343   
  343    344   
    fn cell() -> crate::sync::Arc<OriginCell> {
  344    345   
        crate::sync::Arc::new(OriginCell::new(
  345    346   
            super::super::super::partition::PartitionId::from_index(1),
  346    347   
            OriginKey::from_parts(Scheme::HTTPS, "example.com", None).unwrap(),
  347    348   
            EligibilityGroup::Pool,
  348    349   
            None,
  349    350   
            None,
  350    351   
        ))
  351    352   
    }
  352    353   
  353    354   
    fn launching_waiter(cell: &crate::sync::Arc<OriginCell>) -> WaiterId {
  354    355   
        let waiter = OriginCell::register_waiter(cell, ProtocolRequirement::H2Required);
  355    356   
        let event = cell.poll_waiter(waiter, &mut Context::from_waker(Waker::noop()));
  356         -
        let Poll::Ready(AcquisitionEvent::Establish(permit)) = event else {
         357  +
        let Poll::Ready(AcquisitionStep::StartEstablishment(permit)) = event else {
  357    358   
            panic!("new H2 waiter did not receive establishment authority");
  358    359   
        };
  359    360   
        assert!(cell.start_establishment(waiter));
  360    361   
        drop(permit);
  361    362   
        waiter
  362    363   
    }
  363    364   
  364    365   
    fn failed_event(cell: &OriginCell, waiter: WaiterId) -> ConnectorError {
  365    366   
        let event = cell.poll_waiter(waiter, &mut Context::from_waker(Waker::noop()));
  366         -
        let Poll::Ready(AcquisitionEvent::Complete(AcquisitionResult::Failed(error))) = event
         367  +
        let Poll::Ready(AcquisitionStep::Resolved(AcquisitionOutcome::Failed(error))) = event
  367    368   
        else {
  368    369   
            panic!("flight participant did not receive a failure");
  369    370   
        };
  370    371   
        error
  371    372   
    }
  372    373   
  373    374   
    #[test]
  374    375   
    fn flight_failure_reaches_every_live_participant() {
  375    376   
        let cell = cell();
  376    377   
        let first = launching_waiter(&cell);
  377    378   
        let second = launching_waiter(&cell);
  378         -
        let H2FlightInstall::Driver(flight) = cell.install_or_join_h2_flight(first) else {
         379  +
        let H2FlightDecision::RunFlight(flight) = cell.converge_h2_flight(first) else {
  379    380   
            panic!("first participant did not become the flight driver");
  380    381   
        };
  381    382   
        assert!(matches!(
  382         -
            cell.install_or_join_h2_flight(second),
  383         -
            H2FlightInstall::Joined
         383  +
            cell.converge_h2_flight(second),
         384  +
            H2FlightDecision::JoinedFlight
  384    385   
        ));
  385    386   
  386    387   
        let mut completion = FlightCompletionGuard::new(cell.clone(), flight);
  387    388   
        completion.fail(ConnectorError::io(Box::new(std::io::Error::other(
  388    389   
            "synthetic HTTP/2 handshake failure",
  389    390   
        ))));
  390    391   
  391    392   
        for waiter in [first, second] {
  392    393   
            let error = failed_event(&cell, waiter);
  393    394   
            assert!(error.is_io());
  394    395   
            let source = error.source().expect("flight failure lost its source");
  395    396   
            assert_eq!("synthetic HTTP/2 handshake failure", source.to_string());
  396    397   
            assert!(
  397    398   
                source
  398    399   
                    .source()
  399    400   
                    .expect("flight failure lost its original source")
  400    401   
                    .downcast_ref::<std::io::Error>()
  401    402   
                    .is_some(),
  402    403   
                "flight failure did not preserve the original error"
  403    404   
            );
  404    405   
        }
  405    406   
    }
  406    407   
  407    408   
    #[test]
  408    409   
    fn task_drop_fails_only_participants_still_owned_by_the_flight() {
  409    410   
        let cell = cell();
  410    411   
        let live = launching_waiter(&cell);
  411    412   
        let cancelled = launching_waiter(&cell);
  412         -
        let H2FlightInstall::Driver(flight) = cell.install_or_join_h2_flight(live) else {
         413  +
        let H2FlightDecision::RunFlight(flight) = cell.converge_h2_flight(live) else {
  413    414   
            panic!("first participant did not become the flight driver");
  414    415   
        };
  415    416   
        assert!(matches!(
  416         -
            cell.install_or_join_h2_flight(cancelled),
  417         -
            H2FlightInstall::Joined
         417  +
            cell.converge_h2_flight(cancelled),
         418  +
            H2FlightDecision::JoinedFlight
  418    419   
        ));
  419    420   
        assert!(OriginCell::cancel_waiter(&cell, cancelled));
  420    421   
  421    422   
        drop(FlightCompletionGuard::new(cell.clone(), flight));
  422    423   
  423    424   
        let error = failed_event(&cell, live);
  424    425   
        assert!(error.is_io());
  425    426   
        assert_eq!(
  426    427   
            "HTTP/2 establishment task was dropped",
  427    428   
            error

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

@@ -0,1 +0,292 @@
           1  +
/*
           2  +
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
           3  +
 * SPDX-License-Identifier: Apache-2.0
           4  +
 */
           5  +
           6  +
//! Transport connection policy below HTTP protocol establishment.
           7  +
//!
           8  +
//! This module owns the boundary from a pool connection attempt to connected
           9  +
//! I/O with a negotiated protocol. [`TransportConnectContext`] keeps placement,
          10  +
//! timeout, and ALPN inputs together. A factory may delegate DNS, proxy, TLS,
          11  +
//! and ALPN to an injected connector, but it must report whether an HTTP/1-only
          12  +
//! attempt is enforceable before admission may reclaim an idle H2 connection.
          13  +
//!
          14  +
//! Cached factories retain configured connector services, not connections or
          15  +
//! DNS results. Network I/O begins only when [`TransportFactory::connect`] is
          16  +
//! called.
          17  +
          18  +
use super::super::admission::ProtocolRequirement;
          19  +
use super::super::registry::PartitionState;
          20  +
use crate::client::connect::{AsyncConn, BoxConn};
          21  +
use crate::client::timeout::{self, TimeoutKind};
          22  +
use aws_smithy_async::rt::sleep::SharedAsyncSleep;
          23  +
use aws_smithy_runtime_api::box_error::BoxError;
          24  +
use http_1x::Uri;
          25  +
use std::future::{poll_fn, Future};
          26  +
use std::pin::Pin;
          27  +
use std::sync::Arc as StdArc;
          28  +
use std::time::Duration;
          29  +
use tower::Service;
          30  +
          31  +
/// Timeout policy for one transport connection operation.
          32  +
///
          33  +
/// Connector readiness completes before this timeout starts. Keeping the
          34  +
/// duration and timer together lets request acquisition pass one complete
          35  +
/// timeout value into transport construction.
          36  +
#[derive(Clone, Debug)]
          37  +
pub(in crate::client::pool) struct TransportTimeout {
          38  +
    /// Maximum duration of the connector's connection future.
          39  +
    duration: Duration,
          40  +
    /// Runtime timer used to enforce the deadline.
          41  +
    sleep: SharedAsyncSleep,
          42  +
}
          43  +
          44  +
impl TransportTimeout {
          45  +
    /// Creates a timeout enforced by the supplied runtime timer.
          46  +
    pub(in crate::client::pool) fn new(duration: Duration, sleep: SharedAsyncSleep) -> Self {
          47  +
        Self { duration, sleep }
          48  +
    }
          49  +
}
          50  +
          51  +
/// Static ALPN protocol offer used for one transport connection.
          52  +
pub(in crate::client::pool) type AlpnProtocols = &'static [&'static [u8]];
          53  +
          54  +
/// Default offer for requests that may use HTTP/2.
          55  +
const HTTP_ALPN_PROTOCOLS: AlpnProtocols = &[b"h2", b"http/1.1"];
          56  +
          57  +
/// Narrowed offer for requests that require HTTP/1 wire semantics.
          58  +
const HTTP1_ALPN_PROTOCOLS: AlpnProtocols = &[b"http/1.1"];
          59  +
          60  +
/// Inputs needed to create one transport connection.
          61  +
///
          62  +
/// This context is the transport-attempt boundary for DNS, socket, proxy, TLS,
          63  +
/// and ALPN lifecycle observation. It carries no installed connection state
          64  +
/// because an attempt may fail before a connection exists.
          65  +
pub(in crate::client::pool) struct TransportConnectContext<'a> {
          66  +
    /// Partition whose placement policy owns the connection attempt.
          67  +
    partition: &'a PartitionState,
          68  +
    /// Absolute URI passed to the connector contract.
          69  +
    uri: Uri,
          70  +
    /// Optional deadline for the connector's connection future.
          71  +
    timeout: Option<TransportTimeout>,
          72  +
    /// Protocols the default TLS connector may advertise for this attempt.
          73  +
    alpn_protocols: AlpnProtocols,
          74  +
}
          75  +
          76  +
impl<'a> TransportConnectContext<'a> {
          77  +
    /// Creates transport inputs after request protocol classification.
          78  +
    pub(super) fn new(
          79  +
        partition: &'a PartitionState,
          80  +
        uri: Uri,
          81  +
        timeout: Option<TransportTimeout>,
          82  +
        requirement: ProtocolRequirement,
          83  +
    ) -> Self {
          84  +
        Self {
          85  +
            partition,
          86  +
            uri,
          87  +
            timeout,
          88  +
            alpn_protocols: alpn_protocols(requirement),
          89  +
        }
          90  +
    }
          91  +
}
          92  +
          93  +
/// Future returned by a type-erased transport factory.
          94  +
type TransportFuture = Pin<Box<dyn Future<Output = Result<BoxConn, BoxError>> + Send + 'static>>;
          95  +
          96  +
/// Type-erased transport construction below HTTP protocol establishment.
          97  +
///
          98  +
/// A pool retains one factory. The concrete connector remains responsible for
          99  +
/// the transport stages it implements and for returning connector metadata.
         100  +
/// Protocol establishment consumes the resulting [`BoxConn`].
         101  +
pub(in crate::client::pool) trait TransportFactory:
         102  +
    Send + Sync + 'static
         103  +
{
         104  +
    /// Returns whether an H1-required attempt is guaranteed to negotiate H1.
         105  +
    fn can_guarantee_http1(&self) -> bool;
         106  +
         107  +
    /// Creates one partition-bound transport.
         108  +
    ///
         109  +
    /// Connector readiness completes before the connection timeout starts.
         110  +
    fn connect(&self, context: TransportConnectContext<'_>) -> TransportFuture;
         111  +
}
         112  +
         113  +
/// Connector factory whose concrete service is selected by interface.
         114  +
struct ServiceTransportFactory<F> {
         115  +
    /// Builds a connector with the selected network-interface binding.
         116  +
    connector_for_interface: F,
         117  +
    /// Whether every H1-required connection is guaranteed to negotiate H1.
         118  +
    can_guarantee_http1: bool,
         119  +
}
         120  +
         121  +
impl<F, C, IO> TransportFactory for ServiceTransportFactory<F>
         122  +
where
         123  +
    F: Fn(Option<&str>) -> C + Send + Sync + 'static,
         124  +
    C: Service<Uri, Response = IO> + Send + 'static,
         125  +
    C::Error: Into<BoxError>,
         126  +
    C::Future: Send + 'static,
         127  +
    IO: AsyncConn,
         128  +
{
         129  +
    fn can_guarantee_http1(&self) -> bool {
         130  +
        self.can_guarantee_http1
         131  +
    }
         132  +
         133  +
    fn connect(&self, context: TransportConnectContext<'_>) -> TransportFuture {
         134  +
        let TransportConnectContext {
         135  +
            partition,
         136  +
            uri,
         137  +
            timeout,
         138  +
            alpn_protocols: _alpn_protocols,
         139  +
        } = context;
         140  +
        let interface = partition.interface().map(|interface| interface.as_ref());
         141  +
        let mut connector = (self.connector_for_interface)(interface);
         142  +
        Box::pin(async move {
         143  +
            // Readiness follows the existing connector contract and is not timed.
         144  +
            poll_fn(|cx| connector.poll_ready(cx))
         145  +
                .await
         146  +
                .map_err(Into::into)?;
         147  +
            let connect = connector.call(uri);
         148  +
            let io = timeout::maybe_timeout_future(
         149  +
                connect,
         150  +
                timeout.as_ref().map(|timeout| timeout.duration),
         151  +
                timeout.as_ref().map(|timeout| &timeout.sleep),
         152  +
                TimeoutKind::Connect,
         153  +
            )
         154  +
            .await?;
         155  +
            Ok(Box::new(io) as BoxConn)
         156  +
        })
         157  +
    }
         158  +
}
         159  +
         160  +
/// Erases an injected connector that owns placement and protocol negotiation.
         161  +
///
         162  +
/// The pool cannot apply its interface or ALPN inputs to this connector. It is
         163  +
/// therefore conservative about H2 reclaim for H1-required demand.
         164  +
#[cfg(any(
         165  +
    all(feature = "test-util", aws_sdk_unstable),
         166  +
    all(test, feature = "rt-tokio")
         167  +
))]
         168  +
pub(in crate::client::pool) fn from_connector<C, IO>(connector: C) -> StdArc<dyn TransportFactory>
         169  +
where
         170  +
    C: Service<Uri, Response = IO> + Clone + Send + Sync + 'static,
         171  +
    C::Error: Into<BoxError>,
         172  +
    C::Future: Send + 'static,
         173  +
    IO: AsyncConn,
         174  +
{
         175  +
    service_factory(move |_| connector.clone(), false)
         176  +
}
         177  +
         178  +
/// Erases a cleartext connector constructor that applies interface placement.
         179  +
pub(in crate::client::pool) fn from_interface_connector<F, C, IO>(
         180  +
    connector_for_interface: F,
         181  +
) -> StdArc<dyn TransportFactory>
         182  +
where
         183  +
    F: Fn(Option<&str>) -> C + Send + Sync + 'static,
         184  +
    C: Service<Uri, Response = IO> + Send + 'static,
         185  +
    C::Error: Into<BoxError>,
         186  +
    C::Future: Send + 'static,
         187  +
    IO: AsyncConn,
         188  +
{
         189  +
    service_factory(connector_for_interface, true)
         190  +
}
         191  +
         192  +
/// Erases a service connector and its HTTP/1 negotiation guarantee.
         193  +
fn service_factory<F, C, IO>(
         194  +
    connector_for_interface: F,
         195  +
    can_guarantee_http1: bool,
         196  +
) -> StdArc<dyn TransportFactory>
         197  +
where
         198  +
    F: Fn(Option<&str>) -> C + Send + Sync + 'static,
         199  +
    C: Service<Uri, Response = IO> + Send + 'static,
         200  +
    C::Error: Into<BoxError>,
         201  +
    C::Future: Send + 'static,
         202  +
    IO: AsyncConn,
         203  +
{
         204  +
    StdArc::new(ServiceTransportFactory {
         205  +
        connector_for_interface,
         206  +
        can_guarantee_http1,
         207  +
    })
         208  +
}
         209  +
         210  +
/// Erases and caches configured TLS connector services.
         211  +
///
         212  +
/// Interface placement is immutable for a partition. The first request for
         213  +
/// each `(interface, ALPN offer)` constructs its connector while holding the
         214  +
/// cache lock; later requests clone that retained connector. Provider
         215  +
/// configuration and certificate loading therefore occur at most once for
         216  +
/// each placement and offer rather than once per connection.
         217  +
#[cfg(any(feature = "__rustls", feature = "s2n-tls"))]
         218  +
pub(in crate::client::pool) fn from_cached_interface_connector<F, C, IO>(
         219  +
    connector_for_interface: F,
         220  +
    can_guarantee_http1: bool,
         221  +
) -> StdArc<dyn TransportFactory>
         222  +
where
         223  +
    F: Fn(Option<&str>, AlpnProtocols) -> C + Send + Sync + 'static,
         224  +
    C: Service<Uri, Response = IO> + Clone + Send + Sync + 'static,
         225  +
    C::Error: Into<BoxError>,
         226  +
    C::Future: Send + 'static,
         227  +
    IO: AsyncConn,
         228  +
{
         229  +
    struct CachedTransportFactory<F, C> {
         230  +
        factory: F,
         231  +
        can_guarantee_http1: bool,
         232  +
        connectors:
         233  +
            crate::sync::Mutex<std::collections::HashMap<(Option<String>, AlpnProtocols), C>>,
         234  +
    }
         235  +
         236  +
    impl<F, C, IO> TransportFactory for CachedTransportFactory<F, C>
         237  +
    where
         238  +
        F: Fn(Option<&str>, AlpnProtocols) -> C + Send + Sync + 'static,
         239  +
        C: Service<Uri, Response = IO> + Clone + Send + Sync + 'static,
         240  +
        C::Error: Into<BoxError>,
         241  +
        C::Future: Send + 'static,
         242  +
        IO: AsyncConn,
         243  +
    {
         244  +
        fn can_guarantee_http1(&self) -> bool {
         245  +
            self.can_guarantee_http1
         246  +
        }
         247  +
         248  +
        fn connect(&self, context: TransportConnectContext<'_>) -> TransportFuture {
         249  +
            let TransportConnectContext {
         250  +
                partition,
         251  +
                uri,
         252  +
                timeout,
         253  +
                alpn_protocols,
         254  +
            } = context;
         255  +
            let interface = partition.interface().map(|value| value.to_string());
         256  +
            let mut connector = self
         257  +
                .connectors
         258  +
                .lock()
         259  +
                .entry((interface.clone(), alpn_protocols))
         260  +
                .or_insert_with(|| (self.factory)(interface.as_deref(), alpn_protocols))
         261  +
                .clone();
         262  +
            Box::pin(async move {
         263  +
                poll_fn(|cx| connector.poll_ready(cx))
         264  +
                    .await
         265  +
                    .map_err(Into::into)?;
         266  +
                let connect = connector.call(uri);
         267  +
                let io = timeout::maybe_timeout_future(
         268  +
                    connect,
         269  +
                    timeout.as_ref().map(|timeout| timeout.duration),
         270  +
                    timeout.as_ref().map(|timeout| &timeout.sleep),
         271  +
                    TimeoutKind::Connect,
         272  +
                )
         273  +
                .await?;
         274  +
                Ok(Box::new(io) as BoxConn)
         275  +
            })
         276  +
        }
         277  +
    }
         278  +
         279  +
    StdArc::new(CachedTransportFactory {
         280  +
        factory: connector_for_interface,
         281  +
        can_guarantee_http1,
         282  +
        connectors: crate::sync::Mutex::new(std::collections::HashMap::new()),
         283  +
    })
         284  +
}
         285  +
         286  +
/// Selects the ALPN offer that preserves the request's wire semantics.
         287  +
fn alpn_protocols(requirement: ProtocolRequirement) -> AlpnProtocols {
         288  +
    match requirement {
         289  +
        ProtocolRequirement::H1Required => HTTP1_ALPN_PROTOCOLS,
         290  +
        ProtocolRequirement::H1Compatible | ProtocolRequirement::H2Required => HTTP_ALPN_PROTOCOLS,
         291  +
    }
         292  +
}

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

@@ -165,165 +227,227 @@
  185    185   
            .waker
  186    186   
            .as_ref()
  187    187   
            .is_none_or(|registered| !registered.will_wake(cx.waker()))
  188    188   
        {
  189    189   
            state.waker = Some(cx.waker().clone());
  190    190   
        }
  191    191   
        Poll::Pending
  192    192   
    }
  193    193   
  194    194   
    #[cfg(all(test, feature = "rt-tokio", not(smithy_http_client_loom)))]
  195         -
    fn snapshot(&self) -> MaintenanceSnapshot {
         195  +
    fn probe(&self) -> MaintenanceProbe {
  196    196   
        let state = self.state.lock();
  197         -
        MaintenanceSnapshot {
         197  +
        MaintenanceProbe {
  198    198   
            started: self.started.load(Ordering::Acquire),
  199    199   
            shutdown: state.shutdown,
  200    200   
            scheduled_deadline: state.scheduled_deadline,
  201    201   
        }
  202    202   
    }
  203    203   
}
  204    204   
  205    205   
/// Mutable scheduler state protected independently from every origin cell.
  206    206   
#[derive(Debug, Default)]
  207    207   
struct MaintenanceState {
@@ -351,351 +411,411 @@
  371    371   
            .map_or(Poll::Ready(MaintenanceWake::Revision), |current| {
  372    372   
                current
  373    373   
                    .poll_revision(observed, cx)
  374    374   
                    .map(|()| MaintenanceWake::Revision)
  375    375   
            })
  376    376   
    })
  377    377   
    .await
  378    378   
}
  379    379   
  380    380   
#[cfg(all(test, feature = "rt-tokio", not(smithy_http_client_loom)))]
  381         -
struct MaintenanceSnapshot {
         381  +
struct MaintenanceProbe {
  382    382   
    started: bool,
  383    383   
    shutdown: bool,
  384    384   
    scheduled_deadline: Option<SystemTime>,
  385    385   
}
  386    386   
  387    387   
#[cfg(all(test, smithy_http_client_loom))]
  388    388   
mod loom_tests {
  389    389   
    use super::*;
  390    390   
    use std::sync::atomic::{AtomicUsize, Ordering as StdOrdering};
  391    391   
    use std::sync::Arc as StdArc;
@@ -489,489 +569,569 @@
  509    509   
  510    510   
    #[test]
  511    511   
    fn discarded_task_releases_the_start_latch() {
  512    512   
        let (maintenance, _cell, _gate) = managed_cell(Duration::from_secs(10));
  513    513   
        let submitted = Arc::new(AtomicU64::new(0));
  514    514   
        let spawner = DroppingSpawner {
  515    515   
            submitted: submitted.clone(),
  516    516   
        };
  517    517   
  518    518   
        PartitionMaintenance::start(&maintenance, &spawner);
  519         -
        assert!(!maintenance.snapshot().started);
         519  +
        assert!(!maintenance.probe().started);
  520    520   
        PartitionMaintenance::start(&maintenance, &spawner);
  521    521   
        assert_eq!(2, submitted.load(Ordering::SeqCst));
  522         -
        assert!(!maintenance.snapshot().started);
         522  +
        assert!(!maintenance.probe().started);
  523    523   
    }
  524    524   
  525    525   
    #[test]
  526    526   
    fn disabled_idle_timeout_submits_no_maintenance_task() {
  527    527   
        let maintenance = PartitionMaintenance::new(MaintenanceConfig {
  528    528   
            idle_timeout: None,
  529    529   
            ..MaintenanceConfig::default()
  530    530   
        });
  531    531   
        let submitted = Arc::new(AtomicU64::new(0));
  532    532   
        let spawner = DroppingSpawner {
  533    533   
            submitted: submitted.clone(),
  534    534   
        };
  535    535   
  536    536   
        PartitionMaintenance::start(&maintenance, &spawner);
  537    537   
  538    538   
        assert_eq!(0, submitted.load(Ordering::SeqCst));
  539         -
        assert!(!maintenance.snapshot().started);
         539  +
        assert!(!maintenance.probe().started);
  540    540   
    }
  541    541   
  542    542   
    #[derive(Clone, Debug)]
  543    543   
    struct RewindableTimeSource(Arc<AtomicU64>);
  544    544   
  545    545   
    impl RewindableTimeSource {
  546    546   
        fn new(seconds: u64) -> Self {
  547    547   
            Self(Arc::new(AtomicU64::new(seconds)))
  548    548   
        }
  549    549   
@@ -642,642 +848,848 @@
  662    662   
            PartitionMaintenance::start(&maintenance, &spawner);
  663    663   
        }
  664    664   
        for _ in 0..10 {
  665    665   
            if spawner.active.load(Ordering::SeqCst) == 1 {
  666    666   
                break;
  667    667   
            }
  668    668   
            tokio::task::yield_now().await;
  669    669   
        }
  670    670   
        assert_eq!(1, spawner.submitted.load(Ordering::SeqCst));
  671    671   
        assert_eq!(1, spawner.active.load(Ordering::SeqCst));
  672         -
        assert!(maintenance.snapshot().started);
         672  +
        assert!(maintenance.probe().started);
  673    673   
  674    674   
        maintenance.shutdown();
  675    675   
        for _ in 0..10 {
  676    676   
            if spawner.active.load(Ordering::SeqCst) == 0 {
  677    677   
                break;
  678    678   
            }
  679    679   
            tokio::task::yield_now().await;
  680    680   
        }
  681         -
        assert!(maintenance.snapshot().shutdown);
         681  +
        assert!(maintenance.probe().shutdown);
  682    682   
        assert_eq!(0, spawner.active.load(Ordering::SeqCst));
  683    683   
    }
  684    684   
  685    685   
    #[tokio::test]
  686    686   
    async fn idle_insertion_wakes_a_task_with_no_scheduled_deadline() {
  687    687   
        let timeout = Duration::from_secs(10);
  688    688   
        let (maintenance, cell, mut gate) = managed_cell(timeout);
  689    689   
        PartitionMaintenance::start(&maintenance, &TokioDriverSpawner::current());
  690    690   
        tokio::task::yield_now().await;
  691         -
        assert_eq!(None, maintenance.snapshot().scheduled_deadline);
         691  +
        assert_eq!(None, maintenance.probe().scheduled_deadline);
  692    692   
  693         -
        OriginCell::install_idle_h1(&cell, connection(1), H1Sender::test(1));
         693  +
        OriginCell::insert_idle_h1(&cell, connection(1), H1Sender::test(1));
  694    694   
        let sleep = gate.expect_sleep().await;
  695    695   
        assert_eq!(timeout, sleep.duration());
  696         -
        assert!(maintenance.snapshot().scheduled_deadline.is_some());
         696  +
        assert!(maintenance.probe().scheduled_deadline.is_some());
  697    697   
        maintenance.shutdown();
  698    698   
    }
  699    699   
  700    700   
    #[test]
  701    701   
    fn only_an_earlier_deadline_advances_the_revision() {
  702    702   
        let (time, sleep, _gate) = controlled_time_and_sleep(UNIX_EPOCH);
  703    703   
        let maintenance = PartitionMaintenance::new(MaintenanceConfig {
  704    704   
            idle_timeout: Some(Duration::from_secs(10)),
  705    705   
            time_source: SharedTimeSource::new(time),
  706    706   
            sleep: Some(SharedAsyncSleep::new(sleep)),
  707    707   
        });
  708    708   
        let later = UNIX_EPOCH + Duration::from_secs(20);
  709    709   
        let latest = UNIX_EPOCH + Duration::from_secs(30);
  710    710   
        let earlier = UNIX_EPOCH + Duration::from_secs(10);
  711    711   
  712    712   
        maintenance.notify_deadline(Some(later));
  713    713   
        let revision = maintenance.state.lock().revision;
  714    714   
        maintenance.notify_deadline(Some(latest));
  715    715   
        assert_eq!(revision, maintenance.state.lock().revision);
  716    716   
        maintenance.notify_deadline(Some(earlier));
  717    717   
        assert!(maintenance.state.lock().revision > revision);
  718         -
        assert_eq!(Some(earlier), maintenance.snapshot().scheduled_deadline);
         718  +
        assert_eq!(Some(earlier), maintenance.probe().scheduled_deadline);
  719    719   
    }
  720    720   
  721    721   
    #[tokio::test]
  722    722   
    async fn idle_record_closes_at_its_fake_time_deadline() {
  723    723   
        let timeout = Duration::from_secs(10);
  724    724   
        let (maintenance, cell, mut gate) = managed_cell(timeout);
  725    725   
        let connection = connection(1);
  726         -
        OriginCell::install_idle_h1(&cell, connection.clone(), H1Sender::test(1));
         726  +
        OriginCell::insert_idle_h1(&cell, connection.clone(), H1Sender::test(1));
  727    727   
        PartitionMaintenance::start(&maintenance, &TokioDriverSpawner::current());
  728    728   
  729    729   
        let sleep = gate.expect_sleep().await;
  730    730   
        assert_eq!(timeout, sleep.duration());
  731         -
        assert_eq!(None, connection.snapshot().close_reason);
         731  +
        assert_eq!(None, connection.probe().close_reason);
  732    732   
        sleep.allow_progress();
  733    733   
  734    734   
        for _ in 0..10 {
  735         -
            if connection.snapshot().close_reason == Some(super::super::CloseReason::IdleTimeout) {
         735  +
            if connection.probe().close_reason == Some(super::super::CloseReason::IdleTimeout) {
  736    736   
                break;
  737    737   
            }
  738    738   
            tokio::task::yield_now().await;
  739    739   
        }
  740    740   
        assert_eq!(
  741    741   
            Some(super::super::CloseReason::IdleTimeout),
  742         -
            connection.snapshot().close_reason
         742  +
            connection.probe().close_reason
  743    743   
        );
  744    744   
    }
  745    745   
    #[tokio::test]
  746    746   
    async fn h2_generation_closes_at_its_fake_time_deadline() {
  747    747   
        let timeout = Duration::from_secs(10);
  748    748   
        let (maintenance, cell, mut gate) = managed_cell(timeout);
  749    749   
        let connection = h2_connection(1);
  750    750   
        OriginCell::install_h2_for_test(&cell, connection.clone(), 1, maintenance.idle_deadline());
  751    751   
        PartitionMaintenance::start(&maintenance, &TokioDriverSpawner::current());
  752    752   
  753    753   
        let sleep = gate.expect_sleep().await;
  754    754   
        assert_eq!(timeout, sleep.duration());
  755         -
        assert_eq!(None, connection.snapshot().close_reason);
         755  +
        assert_eq!(None, connection.probe().close_reason);
  756    756   
        sleep.allow_progress();
  757    757   
  758    758   
        for _ in 0..10 {
  759         -
            if connection.snapshot().close_reason == Some(super::super::CloseReason::IdleTimeout) {
         759  +
            if connection.probe().close_reason == Some(super::super::CloseReason::IdleTimeout) {
  760    760   
                break;
  761    761   
            }
  762    762   
            tokio::task::yield_now().await;
  763    763   
        }
  764    764   
        assert_eq!(
  765    765   
            Some(super::super::CloseReason::IdleTimeout),
  766         -
            connection.snapshot().close_reason
         766  +
            connection.probe().close_reason
  767    767   
        );
  768    768   
        assert_eq!(None, cell.accepting_h2_generation());
  769    769   
    }
  770    770   
  771    771   
    #[test]
  772    772   
    fn accepted_h2_dispatch_resets_the_idle_deadline() {
  773    773   
        let timeout = Duration::from_secs(10);
  774    774   
        let time = RewindableTimeSource::new(0);
  775    775   
        let (_unused_time, sleep, _gate) = controlled_time_and_sleep(UNIX_EPOCH);
  776    776   
        let (maintenance, cell) = cell_with_maintenance(
  777    777   
            timeout,
  778    778   
            SharedTimeSource::new(time.clone()),
  779    779   
            SharedAsyncSleep::new(sleep),
  780    780   
        );
  781    781   
        let connection = h2_connection(1);
  782    782   
        OriginCell::install_h2_for_test(&cell, connection.clone(), 1, maintenance.idle_deadline());
  783    783   
        assert_eq!(Some(UNIX_EPOCH + timeout), cell.nearest_idle_deadline());
  784    784   
  785    785   
        time.set(20);
  786    786   
        let mut activation =
  787    787   
            OriginCell::select_h2(&cell).expect("HTTP/2 generation was not selectable");
  788    788   
        let _dispatch_parts = activation.take_dispatch_parts();
  789    789   
        let dispatch = ConnectionState::try_commit_dispatch(&connection)
  790    790   
            .expect("HTTP/2 connection rejected dispatch");
  791    791   
        activation.accept(dispatch);
  792    792   
  793    793   
        assert_eq!(
  794    794   
            Some(UNIX_EPOCH + Duration::from_secs(20) + timeout),
  795    795   
            cell.nearest_idle_deadline()
  796    796   
        );
  797    797   
    }
  798    798   
  799    799   
    #[tokio::test]
  800    800   
    async fn selected_record_has_no_idle_deadline() {
  801    801   
        let timeout = Duration::from_secs(10);
  802    802   
        let (maintenance, cell, mut gate) = managed_cell(timeout);
  803    803   
        let connection = connection(1);
  804         -
        OriginCell::install_idle_h1(&cell, connection.clone(), H1Sender::test(1));
         804  +
        OriginCell::insert_idle_h1(&cell, connection.clone(), H1Sender::test(1));
  805    805   
        PartitionMaintenance::start(&maintenance, &TokioDriverSpawner::current());
  806    806   
  807    807   
        let sleep = gate.expect_sleep().await;
  808    808   
        let selection = OriginCell::select_h1(&cell).expect("idle H1 was not selected");
  809    809   
        sleep.allow_progress();
  810    810   
        for _ in 0..10 {
  811    811   
            tokio::task::yield_now().await;
  812    812   
        }
  813    813   
  814         -
        assert_eq!(None, connection.snapshot().close_reason);
         814  +
        assert_eq!(None, connection.probe().close_reason);
  815    815   
        drop(selection);
  816    816   
    }
  817    817   
  818    818   
    #[tokio::test]
  819    819   
    async fn completed_sleep_is_an_expiration_floor_after_clock_moves_backward() {
  820    820   
        let timeout = Duration::from_secs(10);
  821    821   
        let clock = RewindableTimeSource::new(100);
  822    822   
        let (_unused_time, sleep, mut gate) = controlled_time_and_sleep(UNIX_EPOCH);
  823    823   
        let (maintenance, cell) = cell_with_maintenance(
  824    824   
            timeout,
  825    825   
            SharedTimeSource::new(clock.clone()),
  826    826   
            SharedAsyncSleep::new(sleep),
  827    827   
        );
  828    828   
        let connection = connection(1);
  829         -
        OriginCell::install_idle_h1(&cell, connection.clone(), H1Sender::test(1));
         829  +
        OriginCell::insert_idle_h1(&cell, connection.clone(), H1Sender::test(1));
  830    830   
        PartitionMaintenance::start(&maintenance, &TokioDriverSpawner::current());
  831    831   
  832    832   
        let sleep = gate.expect_sleep().await;
  833    833   
        assert_eq!(timeout, sleep.duration());
  834    834   
        clock.set(90);
  835    835   
        sleep.allow_progress();
  836    836   
  837    837   
        for _ in 0..10 {
  838         -
            if connection.snapshot().close_reason == Some(super::super::CloseReason::IdleTimeout) {
         838  +
            if connection.probe().close_reason == Some(super::super::CloseReason::IdleTimeout) {
  839    839   
                break;
  840    840   
            }
  841    841   
            tokio::task::yield_now().await;
  842    842   
        }
  843    843   
        assert_eq!(
  844    844   
            Some(super::super::CloseReason::IdleTimeout),
  845         -
            connection.snapshot().close_reason
         845  +
            connection.probe().close_reason
  846    846   
        );
  847    847   
    }
  848    848   
}

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

@@ -162,162 +222,222 @@
  182    182   
    Pool,
  183    183   
}
  184    184   
  185    185   
/// Identifies the exact set of partitions allowed to share a connection.
  186    186   
///
  187    187   
/// The configured [`ConnectionReuseScope`] maps each partition to one group:
  188    188   
/// partition-local scope creates a group for one partition, network-interface
  189    189   
/// scope groups equal interface bindings, and pool scope creates one group for
  190    190   
/// the whole pool.
  191    191   
#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
  192         -
pub(crate) enum EligibilityGroup {
         192  +
pub(in crate::client::pool) enum EligibilityGroup {
  193    193   
    /// Only the partition with this identity is eligible.
  194    194   
    Partition(PartitionId),
  195    195   
    /// Partitions with this exact interface binding are eligible.
  196    196   
    NetworkInterface(Option<Arc<str>>),
  197    197   
    /// Every partition in the pool is eligible.
  198    198   
    Pool,
  199    199   
}
  200    200   
  201    201   
#[cfg(test)]
  202    202   
mod tests {

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

@@ -1,1 +307,329 @@
   21     21   
    ConnectionReuseScope, DriverSpawner, EligibilityGroup, Partition, PartitionId,
   22     22   
};
   23     23   
use crate::sync::{Arc, Mutex, RwLock};
   24     24   
use http_1x::Uri;
   25     25   
use std::collections::HashMap;
   26     26   
use std::error::Error;
   27     27   
use std::fmt;
   28     28   
use std::num::{NonZeroU16, NonZeroUsize};
   29     29   
use std::sync::{Arc as StdArc, OnceLock};
   30     30   
          31  +
/// Bounded-origin policy shared by every partition cell for that origin.
          32  +
#[derive(Clone, Copy, Debug)]
          33  +
pub(super) struct AdmissionPolicy {
          34  +
    /// Maximum number of logically open connections for one canonical origin.
          35  +
    connection_limit: NonZeroUsize,
          36  +
    /// Whether replacement establishment can be forced to HTTP/1.
          37  +
    can_reclaim_h2_for_h1: bool,
          38  +
}
          39  +
          40  +
impl AdmissionPolicy {
          41  +
    pub(super) fn new(connection_limit: NonZeroUsize, can_reclaim_h2_for_h1: bool) -> Self {
          42  +
        Self {
          43  +
            connection_limit,
          44  +
            can_reclaim_h2_for_h1,
          45  +
        }
          46  +
    }
          47  +
          48  +
    pub(super) fn connection_limit(self) -> NonZeroUsize {
          49  +
        self.connection_limit
          50  +
    }
          51  +
          52  +
    pub(super) fn can_reclaim_h2_for_h1(self) -> bool {
          53  +
        self.can_reclaim_h2_for_h1
          54  +
    }
          55  +
}
          56  +
   31     57   
/// Fixed partition set and the origin-wide admission states it shares.
   32     58   
#[derive(Debug)]
   33         -
pub(crate) struct PartitionRegistry {
          59  +
pub(in crate::client::pool) struct PartitionRegistry {
   34     60   
    /// Immutable partition identities resolved when a client is built.
   35     61   
    partitions: HashMap<PartitionId, Arc<PartitionState>>,
   36     62   
    /// Policy used to derive a new cell's eligibility group.
   37     63   
    reuse_scope: ConnectionReuseScope,
   38         -
    /// Optional origin-wide bound used when admission is first created.
   39         -
    max_connections_per_host: Option<NonZeroUsize>,
   40         -
    /// Whether admission may close idle H2 capacity for H1-required demand.
   41         -
    allow_h2_reclaim_for_h1: bool,
          64  +
    /// Bounded-origin policy, absent when connection count is unbounded.
          65  +
    admission_policy: Option<AdmissionPolicy>,
   42     66   
    /// Admission authorities retained by canonical origin.
   43         -
    bounded_origins: Mutex<HashMap<OriginKey, Arc<OriginAdmission>>>,
          67  +
    admissions: Mutex<HashMap<OriginKey, Arc<OriginAdmission>>>,
   44     68   
}
   45     69   
   46     70   
impl PartitionRegistry {
   47     71   
    /// Creates the anonymous partition or validates and retains explicit partitions.
   48         -
    pub(crate) fn new(
          72  +
    pub(in crate::client::pool) fn new(
   49     73   
        partitions: Option<Vec<Partition>>,
   50     74   
        reuse_scope: ConnectionReuseScope,
   51         -
        max_connections_per_host: Option<NonZeroUsize>,
   52         -
        allow_h2_reclaim_for_h1: bool,
          75  +
        admission_policy: Option<AdmissionPolicy>,
   53     76   
        maintenance: MaintenanceConfig,
   54     77   
    ) -> Result<Self, PartitionRegistryError> {
   55     78   
        let Some(partitions) = partitions else {
   56     79   
            let partition = Arc::new(PartitionState::anonymous(maintenance));
   57     80   
            let mut partitions = HashMap::new();
   58     81   
            partitions.insert(PartitionId::ANONYMOUS, partition);
   59     82   
            return Ok(Self {
   60     83   
                partitions,
   61     84   
                reuse_scope,
   62         -
                max_connections_per_host,
   63         -
                allow_h2_reclaim_for_h1,
   64         -
                bounded_origins: Mutex::new(HashMap::new()),
          85  +
                admission_policy,
          86  +
                admissions: Mutex::new(HashMap::new()),
   65     87   
            });
   66     88   
        };
   67     89   
   68     90   
        let mut by_id = HashMap::new();
   69     91   
        for partition in partitions {
   70     92   
            let id = partition.id();
   71     93   
            if id.is_anonymous() {
   72     94   
                return Err(PartitionRegistryError::ReservedAnonymousPartition);
   73     95   
            }
   74     96   
            if by_id
   75     97   
                .insert(
   76     98   
                    id,
   77     99   
                    Arc::new(PartitionState::explicit(partition, maintenance.clone())),
   78    100   
                )
   79    101   
                .is_some()
   80    102   
            {
   81    103   
                return Err(PartitionRegistryError::DuplicatePartition(id));
   82    104   
            }
   83    105   
        }
   84    106   
   85    107   
        if by_id.is_empty() {
   86    108   
            return Err(PartitionRegistryError::EmptyExplicitPartitionSet);
   87    109   
        }
   88    110   
   89    111   
        Ok(Self {
   90    112   
            partitions: by_id,
   91    113   
            reuse_scope,
   92         -
            max_connections_per_host,
   93         -
            allow_h2_reclaim_for_h1,
   94         -
            bounded_origins: Mutex::new(HashMap::new()),
         114  +
            admission_policy,
         115  +
            admissions: Mutex::new(HashMap::new()),
   95    116   
        })
   96    117   
    }
   97    118   
   98    119   
    /// Resolves a retained partition once during client construction.
   99         -
    pub(crate) fn partition(&self, id: PartitionId) -> Option<Arc<PartitionState>> {
         120  +
    pub(in crate::client::pool) fn partition(
         121  +
        &self,
         122  +
        id: PartitionId,
         123  +
    ) -> Option<Arc<PartitionState>> {
  100    124   
        self.partitions.get(&id).cloned()
  101    125   
    }
  102    126   
  103    127   
    /// Resolves the stable cell for an already-resolved partition and URI.
  104         -
    pub(crate) fn resolve_cell(
         128  +
    pub(in crate::client::pool) fn resolve_cell(
  105    129   
        &self,
  106    130   
        partition: &PartitionState,
  107    131   
        uri: &Uri,
  108    132   
    ) -> Result<Arc<OriginCell>, InvalidOrigin> {
  109    133   
        let lookup = OriginLookup::from_uri(uri)?;
  110    134   
        if let Some(cell) = partition.find_cell(&lookup) {
  111    135   
            return Ok(cell);
  112    136   
        }
  113    137   
  114    138   
        let eligibility_group = self.eligibility_group(partition);
  115    139   
        let origin = lookup.into_owned();
  116    140   
        let admission = self.origin_admission(&origin);
  117    141   
        let candidate = Arc::new(OriginCell::new(
  118    142   
            partition.id(),
  119    143   
            origin,
  120    144   
            eligibility_group,
  121    145   
            admission.clone(),
  122    146   
            Some(partition.maintenance.clone()),
  123    147   
        ));
  124    148   
        let cell = match admission {
  125    149   
            Some(admission) => OriginAdmission::register_cell(&admission, candidate),
  126    150   
            None => candidate,
  127    151   
        };
  128    152   
        let cell = partition.publish_cell(cell);
  129    153   
        partition.maintenance.register(&cell);
  130    154   
        Ok(cell)
  131    155   
    }
  132    156   
  133    157   
    /// Returns the shared admission authority for a bounded origin.
  134    158   
    ///
  135    159   
    /// Unbounded registries return `None`. Concurrent first use of a bounded
  136    160   
    /// origin converges on one retained authority.
  137    161   
    fn origin_admission(&self, origin: &OriginKey) -> Option<Arc<OriginAdmission>> {
  138         -
        let limit = self.max_connections_per_host?;
         162  +
        let policy = self.admission_policy?;
  139    163   
        let admission = {
  140         -
            let mut origins = self.bounded_origins.lock();
         164  +
            let mut origins = self.admissions.lock();
  141    165   
            origins
  142    166   
                .entry(origin.clone())
  143         -
                .or_insert_with(|| {
  144         -
                    OriginAdmission::new(origin.clone(), limit, self.allow_h2_reclaim_for_h1)
  145         -
                })
         167  +
                .or_insert_with(|| OriginAdmission::new(origin.clone(), policy))
  146    168   
                .clone()
  147    169   
        };
  148    170   
        Some(admission)
  149    171   
    }
  150    172   
  151    173   
    /// Derives which partitions may relieve a new cell's bounded demand.
  152    174   
    fn eligibility_group(&self, partition: &PartitionState) -> EligibilityGroup {
  153    175   
        match self.reuse_scope {
  154    176   
            ConnectionReuseScope::Partition => EligibilityGroup::Partition(partition.id()),
  155    177   
            ConnectionReuseScope::NetworkInterface => {
  156    178   
                EligibilityGroup::NetworkInterface(partition.interface().cloned())
  157    179   
            }
  158    180   
            ConnectionReuseScope::Pool => EligibilityGroup::Pool,
  159    181   
        }
  160    182   
    }
  161    183   
  162    184   
    /// Stops partition maintenance and logically closes every connection.
  163         -
    pub(crate) fn close_all(&self, reason: CloseReason) {
         185  +
    pub(in crate::client::pool) fn close_all(&self, reason: CloseReason) {
  164    186   
        for partition in self.partitions.values() {
  165    187   
            partition.shutdown_maintenance();
  166    188   
        }
  167    189   
        let cells = self
  168    190   
            .partitions
  169    191   
            .values()
  170    192   
            .flat_map(|partition| partition.cells())
  171    193   
            .collect::<Vec<_>>();
  172    194   
        for cell in cells {
  173    195   
            OriginCell::close_all(&cell, reason);
  174    196   
        }
  175    197   
    }
  176    198   
}
  177    199   
  178    200   
/// Runtime placement and retained origin cells for one partition.
  179    201   
#[derive(Debug)]
  180         -
pub(crate) struct PartitionState {
         202  +
pub(in crate::client::pool) struct PartitionState {
  181    203   
    /// Stable identity copied into every cell and owned connection.
  182    204   
    id: PartitionId,
  183    205   
    /// Configured spawner, or the first spawner published for the anonymous
  184    206   
    /// partition.
  185    207   
    spawner: OnceLock<StdArc<dyn DriverSpawner>>,
  186    208   
    /// Network interface used for placement and reuse eligibility.
  187    209   
    interface: Option<StdArc<str>>,
  188    210   
    /// Cells retained for the lifetime of this partition.
  189    211   
    origins: RwLock<OriginMap>,
  190    212   
    /// Owner-runtime idle maintenance for this partition.
  191    213   
    maintenance: Arc<PartitionMaintenance>,
  192    214   
}
  193    215   
  194    216   
impl PartitionState {
  195    217   
    /// Creates the implicit partition whose spawner is published on first use.
  196    218   
    fn anonymous(maintenance: MaintenanceConfig) -> Self {
  197    219   
        Self {
  198    220   
            id: PartitionId::ANONYMOUS,
  199    221   
            spawner: OnceLock::new(),
  200    222   
            interface: None,
  201    223   
            origins: RwLock::new(OriginMap::default()),
  202    224   
            maintenance: PartitionMaintenance::new(maintenance),
  203    225   
        }
  204    226   
    }
  205    227   
  206    228   
    /// Moves one validated explicit declaration into retained partition state.
  207    229   
    fn explicit(partition: Partition, maintenance: MaintenanceConfig) -> Self {
  208    230   
        let (id, spawner, interface) = partition.into_parts();
  209    231   
        Self {
  210    232   
            id,
  211    233   
            spawner: OnceLock::from(spawner),
  212    234   
            interface,
  213    235   
            origins: RwLock::new(OriginMap::default()),
  214    236   
            maintenance: PartitionMaintenance::new(maintenance),
  215    237   
        }
  216    238   
    }
  217    239   
  218    240   
    /// Returns this partition's stable identity.
  219         -
    pub(crate) fn id(&self) -> PartitionId {
         241  +
    pub(in crate::client::pool) fn id(&self) -> PartitionId {
  220    242   
        self.id
  221    243   
    }
  222    244   
  223    245   
    /// Returns this partition's optional network-interface binding.
  224         -
    pub(crate) fn interface(&self) -> Option<&StdArc<str>> {
         246  +
    pub(in crate::client::pool) fn interface(&self) -> Option<&StdArc<str>> {
  225    247   
        self.interface.as_ref()
  226    248   
    }
  227    249   
  228    250   
    /// Returns the runtime that owns connection drivers for this partition.
  229    251   
    ///
  230    252   
    /// Explicit partitions retain their declared spawner. The anonymous
  231    253   
    /// partition captures the first Tokio runtime on which it is used, and
  232    254   
    /// all later requests use that same runtime.
  233         -
    pub(crate) fn owner_spawner(
         255  +
    pub(in crate::client::pool) fn owner_spawner(
  234    256   
        &self,
  235    257   
    ) -> Result<StdArc<dyn DriverSpawner>, MissingAnonymousRuntime> {
  236    258   
        if let Some(spawner) = self.spawner.get() {
  237    259   
            return Ok(spawner.clone());
  238    260   
        }
  239    261   
  240    262   
        if !self.id.is_anonymous() {
  241    263   
            unreachable!("an explicit partition always has a driver spawner");
  242    264   
        }
  243    265   
  244    266   
        #[cfg(feature = "rt-tokio")]
  245    267   
        {
  246    268   
            let handle =
  247    269   
                tokio::runtime::Handle::try_current().map_err(|_| MissingAnonymousRuntime)?;
  248    270   
            Ok(self
  249    271   
                .spawner
  250    272   
                .get_or_init(|| StdArc::new(TokioDriverSpawner::from_handle(handle)))
  251    273   
                .clone())
  252    274   
        }
  253    275   
  254    276   
        #[cfg(not(feature = "rt-tokio"))]
  255    277   
        {
  256    278   
            Err(MissingAnonymousRuntime)
  257    279   
        }
  258    280   
    }
  259    281   
  260    282   
    /// Looks up a cell without materializing an owned origin key.
  261    283   
    fn find_cell(&self, lookup: &OriginLookup<'_>) -> Option<Arc<OriginCell>> {
  262    284   
        self.origins.read().get(lookup)
  263    285   
    }
  264    286   
  265    287   
    /// Publishes `cell` or returns the value that won concurrent publication.
  266    288   
    fn publish_cell(&self, cell: Arc<OriginCell>) -> Arc<OriginCell> {
  267    289   
        self.origins.write().get_or_insert(cell)
  268    290   
    }
  269    291   
  270    292   
    /// Returns the number of cells retained by this partition.
  271    293   
    #[cfg(test)]
  272         -
    pub(crate) fn cell_count(&self) -> usize {
         294  +
    pub(in crate::client::pool) fn cell_count(&self) -> usize {
  273    295   
        self.origins.read().len()
  274    296   
    }
  275    297   
  276    298   
    /// Ensures idle maintenance is running on this partition's owner runtime.
  277         -
    pub(crate) fn ensure_maintenance_started(&self, spawner: &dyn DriverSpawner) {
         299  +
    pub(in crate::client::pool) fn ensure_maintenance_started(&self, spawner: &dyn DriverSpawner) {
  278    300   
        PartitionMaintenance::start(&self.maintenance, spawner);
  279    301   
    }
  280    302   
  281    303   
    /// Stops this partition's maintenance task during pool teardown.
  282    304   
    fn shutdown_maintenance(&self) {
  283    305   
        self.maintenance.shutdown();
  284    306   
    }
  285    307   
  286    308   
    /// Snapshots retained cells before invoking any cell transition.
  287    309   
    fn cells(&self) -> Vec<Arc<OriginCell>> {
@@ -336,358 +480,504 @@
  356    378   
    fn from_lookup(lookup: &OriginLookup<'_>) -> Self {
  357    379   
        Self {
  358    380   
            scheme: lookup.scheme(),
  359    381   
            port: lookup.port(),
  360    382   
        }
  361    383   
    }
  362    384   
}
  363    385   
  364    386   
/// Error returned for an invalid explicit partition set.
  365    387   
#[derive(Clone, Debug, Eq, PartialEq)]
  366         -
pub(crate) enum PartitionRegistryError {
         388  +
pub(in crate::client::pool) enum PartitionRegistryError {
  367    389   
    /// No partitions were declared.
  368    390   
    EmptyExplicitPartitionSet,
  369    391   
    /// More than one partition used the same identity.
  370    392   
    DuplicatePartition(PartitionId),
  371    393   
    /// An explicit partition used the identity reserved for the default.
  372    394   
    ReservedAnonymousPartition,
  373    395   
}
  374    396   
  375    397   
impl fmt::Display for PartitionRegistryError {
  376    398   
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  377    399   
        match self {
  378    400   
            Self::EmptyExplicitPartitionSet => {
  379    401   
                f.write_str("explicit partition set must not be empty")
  380    402   
            }
  381    403   
            Self::DuplicatePartition(id) => write!(f, "duplicate partition identifier: {id:?}"),
  382    404   
            Self::ReservedAnonymousPartition => {
  383    405   
                f.write_str("the anonymous partition identifier is reserved")
  384    406   
            }
  385    407   
        }
  386    408   
    }
  387    409   
}
  388    410   
  389    411   
impl Error for PartitionRegistryError {}
  390    412   
  391    413   
/// Anonymous partition use requires a runtime that can own connection tasks.
  392    414   
#[derive(Debug)]
  393    415   
pub(super) struct MissingAnonymousRuntime;
  394    416   
  395    417   
impl fmt::Display for MissingAnonymousRuntime {
  396    418   
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  397    419   
        f.write_str(
  398    420   
            "the anonymous connection-pool partition requires an active Tokio runtime on first use",
  399    421   
        )
  400    422   
    }
  401    423   
}
  402    424   
  403    425   
impl Error for MissingAnonymousRuntime {}
  404    426   
  405    427   
#[cfg(all(test, not(smithy_http_client_loom)))]
  406    428   
mod tests {
  407    429   
    use super::*;
  408    430   
    use crate::client::pool::admission::ProtocolRequirement;
  409    431   
    use std::future::Future;
  410    432   
    use std::pin::Pin;
  411    433   
    use std::sync::Barrier;
  412    434   
  413    435   
    #[derive(Debug)]
  414    436   
    struct TestSpawner;
  415    437   
  416    438   
    impl DriverSpawner for TestSpawner {
  417    439   
        fn spawn(&self, driver: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) {
  418    440   
            drop(driver);
  419    441   
        }
  420    442   
    }
  421    443   
  422    444   
    fn partition(index: usize) -> Partition {
  423    445   
        Partition::new(PartitionId::from_index(index), TestSpawner)
  424    446   
    }
  425    447   
         448  +
    fn admission_policy(limit: Option<NonZeroUsize>) -> Option<AdmissionPolicy> {
         449  +
        limit.map(|limit| AdmissionPolicy::new(limit, true))
         450  +
    }
         451  +
  426    452   
    #[cfg(feature = "rt-tokio")]
  427    453   
    fn anonymous_registry(
  428    454   
        reuse_scope: ConnectionReuseScope,
  429    455   
        max_connections_per_host: Option<NonZeroUsize>,
  430    456   
    ) -> PartitionRegistry {
  431    457   
        PartitionRegistry::new(
  432    458   
            None,
  433    459   
            reuse_scope,
  434         -
            max_connections_per_host,
  435         -
            true,
         460  +
            admission_policy(max_connections_per_host),
  436    461   
            MaintenanceConfig::default(),
  437    462   
        )
  438    463   
        .unwrap()
  439    464   
    }
  440    465   
  441    466   
    fn explicit_registry(
  442    467   
        partitions: impl IntoIterator<Item = Partition>,
  443    468   
        reuse_scope: ConnectionReuseScope,
  444    469   
        max_connections_per_host: Option<NonZeroUsize>,
  445    470   
    ) -> Result<PartitionRegistry, PartitionRegistryError> {
  446    471   
        PartitionRegistry::new(
  447    472   
            Some(partitions.into_iter().collect()),
  448    473   
            reuse_scope,
  449         -
            max_connections_per_host,
  450         -
            true,
         474  +
            admission_policy(max_connections_per_host),
  451    475   
            MaintenanceConfig::default(),
  452    476   
        )
  453    477   
    }
  454    478   
  455    479   
    fn explicit_partition(registry: &PartitionRegistry, index: usize) -> Arc<PartitionState> {
  456    480   
        registry
  457    481   
            .partition(PartitionId::from_index(index))
  458    482   
            .expect("test partition was not registered")
  459    483   
    }
  460    484   
@@ -665,689 +725,749 @@
  685    709   
  686    710   
    #[test]
  687    711   
    fn unbounded_origins_construct_no_admission_state() {
  688    712   
        let registry = explicit_registry([partition(1)], ConnectionReuseScope::Pool, None).unwrap();
  689    713   
        let partition = explicit_partition(&registry, 1);
  690    714   
        let cell = registry
  691    715   
            .resolve_cell(&partition, &"https://example.com/".parse().unwrap())
  692    716   
            .unwrap();
  693    717   
  694    718   
        assert!(cell.admission().is_none());
  695         -
        assert!(registry.bounded_origins.lock().is_empty());
         719  +
        assert!(registry.admissions.lock().is_empty());
  696    720   
    }
  697    721   
  698    722   
    #[test]
  699    723   
    fn first_cell_publication_is_stable_under_contention() {
  700    724   
        const THREADS: usize = 8;
  701    725   
        let registry = Arc::new(
  702    726   
            explicit_registry(
  703    727   
                [partition(1)],
  704    728   
                ConnectionReuseScope::default(),
  705    729   
                NonZeroUsize::new(1),
@@ -731,755 +795,820 @@
  751    775   
        fn spawn(&self, driver: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) {
  752    776   
            drop(driver);
  753    777   
        }
  754    778   
    }
  755    779   
  756    780   
    fn explicit_registry(
  757    781   
        partitions: impl IntoIterator<Item = Partition>,
  758    782   
        reuse_scope: ConnectionReuseScope,
  759    783   
        max_connections_per_host: Option<NonZeroUsize>,
  760    784   
    ) -> Result<PartitionRegistry, PartitionRegistryError> {
         785  +
        let admission_policy =
         786  +
            max_connections_per_host.map(|limit| AdmissionPolicy::new(limit, true));
  761    787   
        PartitionRegistry::new(
  762    788   
            Some(partitions.into_iter().collect()),
  763    789   
            reuse_scope,
  764         -
            max_connections_per_host,
  765         -
            true,
         790  +
            admission_policy,
  766    791   
            MaintenanceConfig::default(),
  767    792   
        )
  768    793   
    }
  769    794   
  770    795   
    #[test]
  771    796   
    fn first_cell_publication_is_stable_under_contention() {
  772    797   
        loom::model(|| {
  773    798   
            let registry = Arc::new(
  774    799   
                explicit_registry(
  775    800   
                    [Partition::new(PartitionId::from_index(1), TestSpawner)],