From b6fb7f4ae84706965383d1cecfe14bfbeb33cedb Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 24 Apr 2026 15:46:25 -0400 Subject: [PATCH 01/47] bootstrap low level connection harness test util --- aws/rust-runtime/aws-config/Cargo.lock | 236 ++++++------- rust-runtime/Cargo.lock | 1 + .../aws-smithy-http-client/Cargo.toml | 4 +- rust-runtime/aws-smithy-http-client/README.md | 18 + .../aws-smithy-http-client/additional-ci | 6 +- .../aws-smithy-http-client/src/client.rs | 1 + .../aws-smithy-http-client/src/client/pool.rs | 13 + .../src/test_util/wire.rs | 2 + .../src/test_util/wire/connection.rs | 333 ++++++++++++++++++ .../tests/connection_harness_test.rs | 279 +++++++++++++++ 10 files changed, 772 insertions(+), 121 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/src/client/pool.rs create mode 100644 rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs create mode 100644 rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs diff --git a/aws/rust-runtime/aws-config/Cargo.lock b/aws/rust-runtime/aws-config/Cargo.lock index 3dbae82cd5c..7d79e3f627c 100644 --- a/aws/rust-runtime/aws-config/Cargo.lock +++ b/aws/rust-runtime/aws-config/Cargo.lock @@ -35,9 +35,9 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "autocfg" -version = "1.5.1" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "aws-config" @@ -63,7 +63,7 @@ dependencies = [ "fastrand", "futures-util", "hex", - "http 1.4.2", + "http 1.4.0", "p256", "rand", "serde", @@ -92,9 +92,9 @@ dependencies = [ [[package]] name = "aws-lc-rs" -version = "1.17.0" +version = "1.16.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" +checksum = "0ec6fb3fe69024a75fa7e1bfb48aa6cf59706a101658ea01bfd33b2b248a038f" dependencies = [ "aws-lc-sys", "zeroize", @@ -102,9 +102,9 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.41.0" +version = "0.40.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a2f9779ce85b93ab6170dd940ad0169b5766ff848247aff13bb788b832fe3f4" +checksum = "f50037ee5e1e41e7b8f9d161680a725bd1626cb6f8c7e901f91f942850852fe7" dependencies = [ "cc", "cmake", @@ -114,7 +114,7 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.7.5" +version = "1.7.4" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -127,7 +127,7 @@ dependencies = [ "bytes", "bytes-utils", "fastrand", - "http 1.4.2", + "http 1.4.0", "http-body 1.0.1", "percent-encoding", "pin-project-lite", @@ -153,7 +153,7 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "regex-lite", "tracing", ] @@ -175,7 +175,7 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "regex-lite", "tracing", ] @@ -197,7 +197,7 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "regex-lite", "tracing", ] @@ -220,7 +220,7 @@ dependencies = [ "aws-types", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "regex-lite", "tracing", ] @@ -238,7 +238,7 @@ dependencies = [ "hex", "hmac 0.13.0", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "percent-encoding", "sha2 0.11.0", "time", @@ -264,7 +264,7 @@ dependencies = [ "bytes-utils", "futures-core", "futures-util", - "http 1.4.2", + "http 1.4.0", "http-body 1.0.1", "http-body-util", "percent-encoding", @@ -283,13 +283,13 @@ dependencies = [ "aws-smithy-types", "bytes", "h2 0.3.27", - "h2 0.4.14", + "h2 0.4.13", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "http-body 0.4.6", "http-body 1.0.1", "hyper 0.14.32", - "hyper 1.10.1", + "hyper 1.9.0", "hyper-rustls", "hyper-util", "indexmap", @@ -360,7 +360,7 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "http-body 0.4.6", "http-body 1.0.1", "http-body-util", @@ -380,7 +380,7 @@ dependencies = [ "aws-smithy-types", "bytes", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "pin-project-lite", "tokio", "tracing", @@ -402,18 +402,18 @@ version = "0.1.0" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", - "http 1.4.2", + "http 1.4.0", ] [[package]] name = "aws-smithy-types" -version = "1.5.0" +version = "1.4.10" dependencies = [ "base64-simd", "bytes", "bytes-utils", "http 0.2.12", - "http 1.4.2", + "http 1.4.0", "http-body 0.4.6", "http-body 1.0.1", "http-body-util", @@ -476,9 +476,9 @@ checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" [[package]] name = "bitflags" -version = "2.13.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" [[package]] name = "block-buffer" @@ -509,9 +509,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.20.3" +version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" [[package]] name = "bytes" @@ -550,9 +550,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.63" +version = "1.2.60" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "556e016178bb5662a08681bbe0f00f8e17631781a4dfc8c45e466e4b185ec27f" +checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" dependencies = [ "find-msvc-tools", "jobserver", @@ -568,9 +568,9 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "chrono" -version = "0.4.45" +version = "0.4.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" dependencies = [ "num-traits", ] @@ -613,9 +613,9 @@ dependencies = [ [[package]] name = "cmov" -version = "0.5.4" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +checksum = "3f88a43d011fc4a6876cb7344703e297c71dda42494fee094d5f7c76bf13f746" [[package]] name = "const-oid" @@ -693,9 +693,9 @@ dependencies = [ [[package]] name = "crypto-common" -version = "0.2.2" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710" dependencies = [ "hybrid-array", ] @@ -711,9 +711,9 @@ dependencies = [ [[package]] name = "data-encoding" -version = "2.11.0" +version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" [[package]] name = "der" @@ -755,21 +755,21 @@ dependencies = [ [[package]] name = "digest" -version = "0.11.3" +version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +checksum = "4850db49bf08e663084f7fb5c87d202ef91a3907271aff24a94eb97ff039153c" dependencies = [ "block-buffer 0.12.0", "const-oid 0.10.2", - "crypto-common 0.2.2", + "crypto-common 0.2.1", "ctutils", ] [[package]] name = "displaydoc" -version = "0.2.6" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", @@ -798,9 +798,9 @@ dependencies = [ [[package]] name = "either" -version = "1.16.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" [[package]] name = "elliptic-curve" @@ -923,7 +923,6 @@ dependencies = [ "futures-core", "futures-task", "pin-project-lite", - "slab", ] [[package]] @@ -1005,16 +1004,16 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.14" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "171fefbc92fe4a4de27e0698d6a5b392d6a0e333506bc49133760b3bcf948733" +checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" dependencies = [ "atomic-waker", "bytes", "fnv", "futures-core", "futures-sink", - "http 1.4.2", + "http 1.4.0", "indexmap", "slab", "tokio", @@ -1044,9 +1043,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.17.1" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" [[package]] name = "heck" @@ -1075,7 +1074,7 @@ version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" dependencies = [ - "digest 0.11.3", + "digest 0.11.2", ] [[package]] @@ -1091,9 +1090,9 @@ dependencies = [ [[package]] name = "http" -version = "1.4.2" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" dependencies = [ "bytes", "itoa", @@ -1117,7 +1116,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http 1.4.2", + "http 1.4.0", ] [[package]] @@ -1128,7 +1127,7 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http 1.4.2", + "http 1.4.0", "http-body 1.0.1", "pin-project-lite", ] @@ -1147,9 +1146,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hybrid-array" -version = "0.4.12" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" +checksum = "3944cf8cf766b40e2a1a333ee5e9b563f854d5fa49d6a8ca2764e97c6eddb214" dependencies = [ "typenum", ] @@ -1180,16 +1179,16 @@ dependencies = [ [[package]] name = "hyper" -version = "1.10.1" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +checksum = "6299f016b246a94207e63da54dbe807655bf9e00044f73ded42c3ac5305fbcca" dependencies = [ "atomic-waker", "bytes", "futures-channel", "futures-core", - "h2 0.4.14", - "http 1.4.2", + "h2 0.4.13", + "http 1.4.0", "http-body 1.0.1", "httparse", "itoa", @@ -1205,8 +1204,8 @@ version = "0.27.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ - "http 1.4.2", - "hyper 1.10.1", + "http 1.4.0", + "hyper 1.9.0", "hyper-util", "rustls", "rustls-native-certs", @@ -1225,14 +1224,14 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http 1.4.2", + "http 1.4.0", "http-body 1.0.1", - "hyper 1.10.1", + "hyper 1.9.0", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.6.3", "tokio", "tower-service", "tracing", @@ -1339,9 +1338,9 @@ dependencies = [ [[package]] name = "idna_adapter" -version = "1.2.2" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" dependencies = [ "icu_normalizer", "icu_properties", @@ -1354,7 +1353,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.17.1", + "hashbrown 0.17.0", "serde", "serde_core", ] @@ -1383,12 +1382,11 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.100" +version = "0.3.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2025f20d7a4fa7785846e7b63d10a76d3f1cee98ee5cb79ea59703f95e42162" +checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca" dependencies = [ - "cfg-if", - "futures-util", + "once_cell", "wasm-bindgen", ] @@ -1406,9 +1404,9 @@ checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" [[package]] name = "libc" -version = "0.2.186" +version = "0.2.185" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" [[package]] name = "litemap" @@ -1427,9 +1425,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.32" +version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" [[package]] name = "matchers" @@ -1442,9 +1440,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.8.1" +version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b947ae49db0d222b1dbc6b113ce7248a3fc3a6ca21b696717bfc000ba4484d8" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "minimal-lexical" @@ -1454,9 +1452,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "mio" -version = "1.2.1" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" dependencies = [ "libc", "wasi", @@ -1494,9 +1492,9 @@ dependencies = [ [[package]] name = "num-conv" -version = "0.2.2" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" +checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967" [[package]] name = "num-integer" @@ -1758,9 +1756,9 @@ checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" [[package]] name = "regex-syntax" -version = "0.8.11" +version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" [[package]] name = "rfc6979" @@ -1806,9 +1804,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.40" +version = "0.23.38" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" +checksum = "69f9466fb2c14ea04357e91413efb882e2a6d4a406e625449bc0a5d360d53a21" dependencies = [ "aws-lc-rs", "once_cell", @@ -1820,9 +1818,9 @@ dependencies = [ [[package]] name = "rustls-native-certs" -version = "0.8.4" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" dependencies = [ "openssl-probe", "rustls-pki-types", @@ -1832,9 +1830,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" dependencies = [ "zeroize", ] @@ -2001,7 +1999,7 @@ checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", "cpufeatures 0.3.0", - "digest 0.11.3", + "digest 0.11.2", ] [[package]] @@ -2015,9 +2013,9 @@ dependencies = [ [[package]] name = "shlex" -version = "2.0.1" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] name = "signal-hook-registry" @@ -2063,9 +2061,9 @@ dependencies = [ [[package]] name = "socket2" -version = "0.6.4" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", "windows-sys 0.61.2", @@ -2201,9 +2199,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.52.3" +version = "1.52.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "b67dee974fe86fd92cc45b7a95fdd2f99a36a6d7b0d431a231178d3d670bbcc6" dependencies = [ "bytes", "libc", @@ -2211,7 +2209,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2 0.6.4", + "socket2 0.6.3", "tokio-macros", "windows-sys 0.61.2", ] @@ -2375,9 +2373,9 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" [[package]] name = "typenum" -version = "1.20.1" +version = "1.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" [[package]] name = "unicode-ident" @@ -2423,9 +2421,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.23.3" +version = "1.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "144d6b123cef80b301b8f72a9e2ca4370ddec21950d0a103dd22c437006d2db7" +checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" dependencies = [ "getrandom 0.4.2", "js-sys", @@ -2485,9 +2483,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.123" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a254a4b10c19a76f09a27640e7ffbf9bc30bf67e16a3bf28aaefa4920fe81563" +checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89" dependencies = [ "cfg-if", "once_cell", @@ -2498,9 +2496,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.123" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24a40fc75b0ec6f3746ceb10d36f53a93dcd68a93b11b6445983945d79eba0dc" +checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -2508,9 +2506,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.123" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "908f34bd9b9ce3d4caf07b72dfab63d61504d156856c6bd3cd87fa350cf3985b" +checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904" dependencies = [ "bumpalo", "proc-macro2", @@ -2521,9 +2519,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.123" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7acbf7616c27b194bbb550bf77ed0c2c3e5b7fd1260a93082b95fb7f47959b92" +checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129" dependencies = [ "unicode-ident", ] @@ -2764,9 +2762,9 @@ checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" [[package]] name = "yoke" -version = "0.8.3" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" dependencies = [ "stable_deref_trait", "yoke-derive", @@ -2787,18 +2785,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.52" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.52" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" dependencies = [ "proc-macro2", "quote", @@ -2807,9 +2805,9 @@ dependencies = [ [[package]] name = "zerofrom" -version = "0.1.8" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" dependencies = [ "zerofrom-derive", ] diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock index 1b2800b4d06..c3ac2aebed4 100644 --- a/rust-runtime/Cargo.lock +++ b/rust-runtime/Cargo.lock @@ -531,6 +531,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "socket2 0.6.3", "tokio", "tokio-rustls 0.26.4", "tower 0.5.3", diff --git a/rust-runtime/aws-smithy-http-client/Cargo.toml b/rust-runtime/aws-smithy-http-client/Cargo.toml index 24893f90f30..0f30f3f166a 100644 --- a/rust-runtime/aws-smithy-http-client/Cargo.toml +++ b/rust-runtime/aws-smithy-http-client/Cargo.toml @@ -40,6 +40,7 @@ wire-mock = [ "hyper-util?/server-graceful", "tokio/macros", "dep:http-body-util", + "dep:socket2", ] test-util = [ @@ -85,7 +86,7 @@ tracing = "0.1.44" # hyper 1.x stack hyper = { version = "1.6.0", features = ["client", "http1", "http2"], optional = true } -hyper-util = { version = "0.1.16", features = ["http1", "http2"], optional = true } +hyper-util = { version = "0.1.16", features = ["http1", "http2", "client-pool"], optional = true } http-1x = { package = "http", version = "1.3.1" , optional = true } http-body-1x = { package = "http-body", version = "1.0.1", optional = true} hyper-rustls = { version = "0.27", features = ["http2", "http1", "native-tokio", "tls12"], default-features = false, optional = true } @@ -115,6 +116,7 @@ serde = { version = "1.0.228", features = ["derive"], optional = true } serde_json = { version = "1.0.146", features = ["preserve_order"], optional = true } indexmap = { version = "2.10.0", features = ["serde"], optional = true } http-body-util = { version = "0.1.3", optional = true } +socket2 = { version = "0.6", optional = true } # end test util stack [dev-dependencies] diff --git a/rust-runtime/aws-smithy-http-client/README.md b/rust-runtime/aws-smithy-http-client/README.md index 261a84b86ce..d8809c9e949 100644 --- a/rust-runtime/aws-smithy-http-client/README.md +++ b/rust-runtime/aws-smithy-http-client/README.md @@ -2,6 +2,24 @@ HTTP client abstractions for generated smithy clients. +## Testing + +### Connection pool integration tests + +```sh +cargo test --features wire-mock,default-client --test connection_harness_test +``` + +Some tests simulate multiple S3 IPs by binding to different loopback addresses (`127.0.0.1`, `127.0.0.2`, etc.) on the same port. +On Linux this works out of the box. On macOS, loopback aliases must be configured first: + +```sh +sudo ifconfig lo0 alias 127.0.0.2 +sudo ifconfig lo0 alias 127.0.0.3 +``` + +These aliases do not persist across reboots. Without them, multi-IP tests will be skipped. + This crate is part of the [AWS SDK for Rust](https://awslabs.github.io/aws-sdk-rust/) and the [smithy-rs](https://github.com/smithy-lang/smithy-rs) code generator. In most cases, it should not be used directly. diff --git a/rust-runtime/aws-smithy-http-client/additional-ci b/rust-runtime/aws-smithy-http-client/additional-ci index 21cc04a5f93..ee5eb936db7 100755 --- a/rust-runtime/aws-smithy-http-client/additional-ci +++ b/rust-runtime/aws-smithy-http-client/additional-ci @@ -4,10 +4,14 @@ # SPDX-License-Identifier: Apache-2.0 # -# This script contains additional CI checks to run for this specific package +# Additional CI checks for aws-smithy-http-client set -e echo "### Testing unstable custom rustls crypto provider" # Enabling ring because the tests use it as the "custom" provider, and it transitively enables __rustls RUSTFLAGS="--cfg aws_sdk_unstable" cargo test --features rustls-ring + + +echo "### Running connection pool integration tests" +cargo test --features wire-mock,default-client --test connection_harness_test diff --git a/rust-runtime/aws-smithy-http-client/src/client.rs b/rust-runtime/aws-smithy-http-client/src/client.rs index cc808aa6739..4ee92f2d782 100644 --- a/rust-runtime/aws-smithy-http-client/src/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client.rs @@ -11,6 +11,7 @@ mod timeout; pub mod tls; pub(crate) mod connect; +pub(crate) mod pool; use crate::cfg::cfg_tls; use crate::tls::TlsContext; diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs new file mode 100644 index 00000000000..f2a0e268ddf --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -0,0 +1,13 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Connection pool implementation for the v2 HTTP client. +//! +//! This module composes hyper-util's pool primitives (Cache, Singleton, Negotiate, Map) +//! with SDK-owned connection lifecycle management including: +//! - Connection state tracking (idle-since, remote IP, poisoned flag) +//! - Proactive health checks on checkout +//! - Connection poisoning support +//! - Queryable pool state (idle/active counts) diff --git a/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs b/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs index 8fc3172f403..92dbd3cb858 100644 --- a/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs +++ b/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs @@ -39,6 +39,8 @@ #![allow(missing_docs)] +pub mod connection; + use aws_smithy_async::future::never::Never; use aws_smithy_async::future::BoxFuture; use aws_smithy_runtime_api::client::http::SharedHttpClient; diff --git a/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs b/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs new file mode 100644 index 00000000000..0045fe76802 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs @@ -0,0 +1,333 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Test harness for connection-level behavior testing. +//! +//! Simulates multiple IPs via TCP listeners on different loopback addresses +//! (`127.0.0.1`, `127.0.0.2`, etc.) sharing the same port. + +use aws_smithy_runtime_api::client::dns::{DnsFuture, ResolveDns}; +use std::collections::{HashMap, VecDeque}; +use std::net::{IpAddr, SocketAddr}; +use std::sync::{Arc, Mutex}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; + +/// Programmable behavior for a test endpoint per accepted connection. +#[derive(Debug, Clone)] +pub enum ConnectionBehavior { + /// Accept TCP, send HTTP/1.1 response, keep connection open (reusable). + Respond { + /// HTTP status code to return. + status: u16, + /// Response body bytes. + body: &'static [u8], + }, + /// Accept TCP, immediately reset the connection (RST). + ResetOnConnect, + /// Accept TCP, hold open for duration, then close. + HoldThenClose(std::time::Duration), +} + +/// Recorded event from the test harness. +#[derive(Debug, Clone)] +pub enum ConnectionEvent { + /// TCP connection accepted at an endpoint. + TcpAccepted { + /// The address of the endpoint that accepted the connection. + endpoint_addr: SocketAddr, + }, + /// DNS lookup performed. + DnsLookup { + /// The hostname that was looked up. + hostname: String, + }, +} + +/// A TCP endpoint bound to a specific address that executes programmed behaviors. +pub struct TestEndpoint { + addr: SocketAddr, + _task: JoinHandle<()>, +} + +impl TestEndpoint { + async fn bind( + addr: &str, + behaviors: Vec, + events: Arc>>, + ) -> Self { + let listener = TcpListener::bind(addr) + .await + .unwrap_or_else(|e| panic!("failed to bind TCP listener to {addr}: {e}")); + let addr = listener + .local_addr() + .expect("failed to get local address from listener"); + let behaviors = Arc::new(Mutex::new(VecDeque::from(behaviors))); + + let task = tokio::spawn(async move { + loop { + let (stream, _) = match listener.accept().await { + Ok(conn) => conn, + Err(_) => break, + }; + events + .lock() + .expect("event lock poisoned") + .push(ConnectionEvent::TcpAccepted { + endpoint_addr: addr, + }); + let behavior = behaviors + .lock() + .expect("behavior lock poisoned") + .pop_front(); + tokio::spawn(handle_connection(stream, behavior, behaviors.clone())); + } + }); + + Self { addr, _task: task } + } + + /// The port this endpoint is listening on. + pub fn port(&self) -> u16 { + self.addr.port() + } + + /// The IP address this endpoint is listening on. + pub fn ip(&self) -> IpAddr { + self.addr.ip() + } + + /// The full socket address this endpoint is listening on. + pub fn addr(&self) -> SocketAddr { + self.addr + } +} + +async fn handle_connection( + mut stream: tokio::net::TcpStream, + first_behavior: Option, + remaining: Arc>>, +) { + let Some(behavior) = first_behavior else { + drop(stream); + return; + }; + + match behavior { + ConnectionBehavior::ResetOnConnect => { + // Set SO_LINGER to 0 to send TCP RST on close + let sock = socket2::SockRef::from(&stream); + sock.set_linger(Some(std::time::Duration::ZERO)) + .expect("failed to set SO_LINGER"); + drop(stream); + } + ConnectionBehavior::HoldThenClose(duration) => { + tokio::time::sleep(duration).await; + drop(stream); + } + ConnectionBehavior::Respond { status, body } => { + if write_response(&mut stream, status, body).await.is_err() { + return; + } + // Keep-alive loop: read next request, send next behavior's response + loop { + if read_request(&mut stream).await.is_err() { + return; + } + let next = remaining + .lock() + .expect("behavior lock poisoned") + .pop_front(); + match next { + Some(ConnectionBehavior::Respond { status, body }) => { + if write_response(&mut stream, status, body).await.is_err() { + return; + } + } + _ => return, // No more behaviors or non-Respond behavior: close + } + } + } + } +} + +async fn write_response( + stream: &mut tokio::net::TcpStream, + status: u16, + body: &[u8], +) -> Result<(), std::io::Error> { + let header = format!( + "HTTP/1.1 {status} OK\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n", + body.len() + ); + stream.write_all(header.as_bytes()).await?; + stream.write_all(body).await?; + stream.flush().await +} + +async fn read_request(stream: &mut tokio::net::TcpStream) -> Result<(), std::io::Error> { + // Read until we see \r\n\r\n (end of HTTP headers) + let mut buf = vec![0u8; 4096]; + let mut total = 0; + loop { + let n = stream.read(&mut buf[total..]).await?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "client closed connection", + )); + } + total += n; + if total >= 4 && buf[..total].windows(4).any(|w| w == b"\r\n\r\n") { + return Ok(()); + } + } +} + +/// Mock DNS resolver that returns configured IPs and logs lookups. +#[derive(Debug, Clone)] +pub struct MockDnsResolver { + responses: HashMap>, + events: Arc>>, +} + +impl MockDnsResolver { + fn new(events: Arc>>) -> Self { + Self { + responses: HashMap::new(), + events, + } + } + + /// Add a DNS entry mapping hostname to IPs. + fn with(mut self, hostname: &str, ips: Vec) -> Self { + self.responses.insert(hostname.to_string(), ips); + self + } +} + +impl ResolveDns for MockDnsResolver { + fn resolve_dns<'a>(&'a self, name: &'a str) -> DnsFuture<'a> { + let ips = self.responses.get(name).cloned().unwrap_or_default(); + self.events + .lock() + .expect("event lock poisoned") + .push(ConnectionEvent::DnsLookup { + hostname: name.to_string(), + }); + DnsFuture::ready(Ok(ips)) + } +} + +/// Test harness for connection-level behavior testing. +pub struct ConnectionTestHarness { + /// The test endpoints managed by this harness. + pub endpoints: Vec, + events: Arc>>, + dns_resolver: MockDnsResolver, +} + +impl ConnectionTestHarness { + /// Create a new builder for the test harness. + pub fn builder() -> HarnessBuilder { + HarnessBuilder { + endpoint_configs: Vec::new(), + dns_entries: Vec::new(), + } + } + + /// Clone all recorded events. + pub fn events(&self) -> Vec { + self.events.lock().expect("event lock poisoned").clone() + } + + /// Count of TCP accepted events. + pub fn tcp_accepted_count(&self) -> usize { + self.events() + .iter() + .filter(|e| matches!(e, ConnectionEvent::TcpAccepted { .. })) + .count() + } + + /// Count of TCP accepted events for a specific IP. + pub fn tcp_accepted_by(&self, ip: IpAddr) -> usize { + self.events() + .iter() + .filter(|e| matches!(e, ConnectionEvent::TcpAccepted { endpoint_addr } if endpoint_addr.ip() == ip)) + .count() + } + + /// Count of DNS lookup events. + pub fn dns_lookup_count(&self) -> usize { + self.events() + .iter() + .filter(|e| matches!(e, ConnectionEvent::DnsLookup { .. })) + .count() + } + + /// Get a clone of the mock DNS resolver. + pub fn dns_resolver(&self) -> MockDnsResolver { + self.dns_resolver.clone() + } +} + +/// Builder for [`ConnectionTestHarness`]. +pub struct HarnessBuilder { + endpoint_configs: Vec<(IpAddr, Vec)>, + dns_entries: Vec<(String, Vec)>, +} + +impl HarnessBuilder { + /// Add an endpoint with the given IP and behaviors. + pub fn endpoint(mut self, ip: IpAddr, behaviors: Vec) -> Self { + self.endpoint_configs.push((ip, behaviors)); + self + } + + /// Add a DNS entry mapping hostname to IPs. + pub fn dns(mut self, hostname: &str, ips: Vec) -> Self { + self.dns_entries.push((hostname.to_string(), ips)); + self + } + + /// Add a DNS entry mapping hostname to all configured endpoint IPs. + pub fn dns_all(mut self, hostname: &str) -> Self { + let ips: Vec = self.endpoint_configs.iter().map(|(ip, _)| *ip).collect(); + self.dns_entries.push((hostname.to_string(), ips)); + self + } + + /// Build the test harness, binding all endpoints. + pub async fn build(self) -> ConnectionTestHarness { + let events: Arc>> = Arc::new(Mutex::new(Vec::new())); + let mut endpoints = Vec::new(); + + let mut port = 0u16; + for (i, (ip, behaviors)) in self.endpoint_configs.into_iter().enumerate() { + let bind_addr = if i == 0 { + format!("{ip}:0") + } else { + format!("{ip}:{port}") + }; + let ep = TestEndpoint::bind(&bind_addr, behaviors, events.clone()).await; + if i == 0 { + port = ep.port(); + } + endpoints.push(ep); + } + + let mut resolver = MockDnsResolver::new(events.clone()); + for (hostname, ips) in self.dns_entries { + resolver = resolver.with(&hostname, ips); + } + + ConnectionTestHarness { + endpoints, + events, + dns_resolver: resolver, + } + } +} diff --git a/rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs b/rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs new file mode 100644 index 00000000000..3f07b536579 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs @@ -0,0 +1,279 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![cfg(all(feature = "wire-mock", feature = "default-client"))] + +use aws_smithy_http_client::test_util::wire::connection::{ + ConnectionBehavior, ConnectionTestHarness, +}; +use std::net::{IpAddr, Ipv4Addr}; + +const IP1: IpAddr = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)); +const IP2: IpAddr = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)); + +/// Check if a loopback address is bindable on this system. +/// On macOS, addresses other than 127.0.0.1 require explicit loopback aliases. +async fn is_bindable(ip: IpAddr) -> bool { + tokio::net::TcpListener::bind((ip, 0u16)).await.is_ok() +} + +#[tokio::test] +async fn test_harness_multi_ip_endpoints() { + if !is_bindable(IP2).await { + eprintln!("skipping test: 127.0.0.2 not bindable (loopback alias not configured)"); + return; + } + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"hello", + }], + ) + .endpoint( + IP2, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"world", + }], + ) + .dns_all("test.example.com") + .build() + .await; + + // Verify endpoints are on different IPs but same port + let eps = &harness.endpoints; + assert_ne!(eps[0].ip(), eps[1].ip()); + assert_eq!(eps[0].port(), eps[1].port()); + + // Connect to each endpoint directly to verify they work + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpStream; + + for ep in &harness.endpoints { + let mut stream = TcpStream::connect(ep.addr()).await.unwrap(); + stream + .write_all(b"GET / HTTP/1.1\r\nHost: test\r\n\r\n") + .await + .unwrap(); + let mut buf = vec![0u8; 1024]; + let n = stream.read(&mut buf).await.unwrap(); + let response = String::from_utf8_lossy(&buf[..n]); + assert!( + response.contains("200"), + "Expected 200 response, got: {}", + response + ); + } + + assert_eq!(harness.tcp_accepted_count(), 2); + assert_eq!(harness.tcp_accepted_by(IP1), 1); + assert_eq!(harness.tcp_accepted_by(IP2), 1); +} + +#[tokio::test] +async fn test_harness_reset_on_connect() { + let harness = ConnectionTestHarness::builder() + .endpoint(IP1, vec![ConnectionBehavior::ResetOnConnect]) + .build() + .await; + + use tokio::io::AsyncReadExt; + use tokio::net::TcpStream; + + let mut stream = TcpStream::connect(harness.endpoints[0].addr()) + .await + .unwrap(); + // Try to read — should get connection reset or EOF + let mut buf = vec![0u8; 1024]; + let result = stream.read(&mut buf).await; + assert!( + result.is_err() || result.unwrap() == 0, + "Expected connection reset or EOF" + ); + + assert_eq!(harness.tcp_accepted_count(), 1); +} + +#[tokio::test] +async fn test_mock_dns_resolver() { + if !is_bindable(IP2).await { + eprintln!("skipping test: 127.0.0.2 not bindable (loopback alias not configured)"); + return; + } + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }], + ) + .endpoint( + IP2, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }], + ) + .dns("s3.amazonaws.com", vec![IP1, IP2]) + .build() + .await; + + use aws_smithy_runtime_api::client::dns::ResolveDns; + let resolver = harness.dns_resolver(); + let ips = resolver.resolve_dns("s3.amazonaws.com").await.unwrap(); + assert_eq!(ips.len(), 2); + assert!(ips.contains(&IP1)); + assert!(ips.contains(&IP2)); + assert_eq!(harness.dns_lookup_count(), 1); +} + +#[tokio::test] +async fn test_harness_keep_alive_reuse() { + // Multiple responses on a single connection (keep-alive) + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::Respond { + status: 200, + body: b"first", + }, + ConnectionBehavior::Respond { + status: 201, + body: b"second", + }, + ConnectionBehavior::Respond { + status: 202, + body: b"third", + }, + ], + ) + .build() + .await; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpStream; + + let mut stream = TcpStream::connect(harness.endpoints[0].addr()) + .await + .unwrap(); + + for (i, expected_status) in ["200", "201", "202"].iter().enumerate() { + stream + .write_all(b"GET / HTTP/1.1\r\nHost: test\r\n\r\n") + .await + .unwrap(); + let mut buf = vec![0u8; 1024]; + let n = stream.read(&mut buf).await.unwrap(); + let response = String::from_utf8_lossy(&buf[..n]); + assert!( + response.contains(expected_status), + "Request {i}: expected {expected_status}, got: {response}" + ); + } + + // Only 1 TCP connection was accepted — all 3 requests reused it + assert_eq!(harness.tcp_accepted_count(), 1); +} + +#[tokio::test] +async fn test_harness_hold_then_close() { + use std::time::{Duration, Instant}; + + let hold_duration = Duration::from_millis(100); + let harness = ConnectionTestHarness::builder() + .endpoint(IP1, vec![ConnectionBehavior::HoldThenClose(hold_duration)]) + .build() + .await; + + use tokio::io::AsyncReadExt; + use tokio::net::TcpStream; + + let start = Instant::now(); + let mut stream = TcpStream::connect(harness.endpoints[0].addr()) + .await + .unwrap(); + let mut buf = vec![0u8; 1024]; + let n = stream.read(&mut buf).await.unwrap(); + let elapsed = start.elapsed(); + + // Connection held open then closed — read returns 0 (EOF) + assert_eq!(n, 0, "Expected EOF after hold-then-close"); + assert!( + elapsed >= hold_duration, + "Expected at least {hold_duration:?} hold, got {elapsed:?}" + ); + assert_eq!(harness.tcp_accepted_count(), 1); +} + +#[tokio::test] +async fn test_harness_dns_all_includes_all_endpoints() { + if !is_bindable(IP2).await { + eprintln!("skipping test: 127.0.0.2 not bindable (loopback alias not configured)"); + return; + } + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"a", + }], + ) + .endpoint( + IP2, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"b", + }], + ) + .dns_all("example.com") + .build() + .await; + + use aws_smithy_runtime_api::client::dns::ResolveDns; + let ips = harness + .dns_resolver() + .resolve_dns("example.com") + .await + .unwrap(); + assert_eq!(ips.len(), 2); + assert!(ips.contains(&IP1)); + assert!(ips.contains(&IP2)); +} + +#[tokio::test] +async fn test_mock_dns_unknown_host_returns_empty() { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }], + ) + .dns("known.com", vec![IP1]) + .build() + .await; + + use aws_smithy_runtime_api::client::dns::ResolveDns; + let ips = harness + .dns_resolver() + .resolve_dns("unknown.com") + .await + .unwrap(); + assert!(ips.is_empty(), "Unknown host should return empty IP list"); + assert_eq!( + harness.dns_lookup_count(), + 1, + "Lookup should still be recorded" + ); +} From f5036ced2c240bc84209447f273eeb93b00bb5c1 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Sun, 26 Apr 2026 09:46:22 -0400 Subject: [PATCH 02/47] refactor transient error tests to cover all stacks --- aws/rust-runtime/Cargo.lock | 2 +- .../aws-smithy-http-client/src/client.rs | 21 ++ .../src/test_util/wire.rs | 16 ++ .../tests/pool_behavior_test.rs | 237 ++++++++++++++++++ .../tests/proxy_tests.rs | 4 + .../tests/reconnect_on_transient_error.rs | 126 ++++++---- 6 files changed, 361 insertions(+), 45 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs diff --git a/aws/rust-runtime/Cargo.lock b/aws/rust-runtime/Cargo.lock index fb1ef471a11..20fa5de8665 100644 --- a/aws/rust-runtime/Cargo.lock +++ b/aws/rust-runtime/Cargo.lock @@ -254,7 +254,7 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.1.13" +version = "1.2.0" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", diff --git a/rust-runtime/aws-smithy-http-client/src/client.rs b/rust-runtime/aws-smithy-http-client/src/client.rs index 4ee92f2d782..b5ecdb679c5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client.rs @@ -975,6 +975,27 @@ impl Builder { ) } + /// Build an HTTP client (no TLS) with a custom DNS resolver. + #[doc(hidden)] + #[cfg(feature = "test-util")] + pub fn build_with_resolver( + self, + resolver: impl aws_smithy_runtime_api::client::dns::ResolveDns + Clone + 'static, + ) -> SharedHttpClient { + build_with_conn_fn( + self.client_builder, + self.pool_idle_timeout, + move |client_builder, settings, runtime_components| { + let builder = new_conn_builder(client_builder, settings, runtime_components); + use crate::client::dns::HyperUtilResolver; + let http_connector = builder.base_connector_with_resolver(HyperUtilResolver { + resolver: resolver.clone(), + }); + builder.wrap_connector(http_connector) + }, + ) + } + /// Set the TLS implementation to use pub fn tls_provider(self, provider: tls::Provider) -> Builder { Builder { diff --git a/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs b/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs index 92dbd3cb858..8564413c423 100644 --- a/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs +++ b/rust-runtime/aws-smithy-http-client/src/test_util/wire.rs @@ -411,6 +411,22 @@ impl tower::Service for InnerDnsResolver { } } +impl aws_smithy_runtime_api::client::dns::ResolveDns for LoggingDnsResolver { + fn resolve_dns<'a>( + &'a self, + name: &'a str, + ) -> aws_smithy_runtime_api::client::dns::DnsFuture<'a> { + let socket_addr = self.0.socket_addr; + let log = self.0.log.clone(); + aws_smithy_runtime_api::client::dns::DnsFuture::new(async move { + log.lock() + .unwrap() + .push(RecordedEvent::DnsLookup(name.to_string())); + Ok(vec![socket_addr.ip()]) + }) + } +} + #[cfg(all(feature = "legacy-test-util", feature = "hyper-014"))] impl hyper_0_14::service::Service for LoggingDnsResolver { type Response = Once; diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs new file mode 100644 index 00000000000..1ff181fe8e7 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -0,0 +1,237 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Pool behavior tests parameterized over HTTP client implementations. +//! +//! Each test is written once as an `async fn` that takes a `&dyn MakeClient`, +//! then invoked for each client implementation. When `Builder::new_v2()` is +//! added, a single new `MakeClient` impl covers all tests. + +#![cfg(all(feature = "wire-mock", feature = "default-client"))] + +use aws_smithy_async::time::SystemTimeSource; +use aws_smithy_http_client::test_util::wire::connection::{ + ConnectionBehavior, ConnectionTestHarness, +}; +use aws_smithy_http_client::test_util::wire::{ReplayedEvent, WireMockServer}; +use aws_smithy_http_client::{ev, match_events, Builder, Connector}; +use aws_smithy_runtime_api::client::http::{ + HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, +}; +use aws_smithy_runtime_api::client::orchestrator::HttpRequest; +use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder; +use std::borrow::Cow; +use std::net::{IpAddr, Ipv4Addr}; +use std::time::Duration; + +const IP1: IpAddr = IpAddr::V4(Ipv4Addr::LOCALHOST); + +// --------------------------------------------------------------------------- +// ClientConfig + MakeClient +// --------------------------------------------------------------------------- + +#[derive(Clone, Default)] +struct ClientConfig { + idle_timeout: Option, + max_connections: Option, +} + +impl ClientConfig { + fn with_idle_timeout(mut self, timeout: Duration) -> Self { + self.idle_timeout = Some(timeout); + self + } + + #[allow(dead_code)] + fn with_max_connections(mut self, n: usize) -> Self { + self.max_connections = Some(n); + self + } +} + +trait MakeClient: Send + Sync { + fn make(&self, config: ClientConfig) -> SharedHttpClient; +} + +/// Current hyper 1.x client via `Builder::new()` +struct V1Client; + +impl MakeClient for V1Client { + fn make(&self, config: ClientConfig) -> SharedHttpClient { + Builder::new().build_with_connector_fn(move |_settings, _components| { + let mut builder = Connector::builder(); + if let Some(timeout) = config.idle_timeout { + builder = builder.pool_idle_timeout(timeout); + } + // v1 does not support max_connections + builder.build_http() + }) + } +} + +// When Builder::new_v2() exists: +// struct V2Client; +// impl MakeClient for V2Client { +// fn make(&self, config: ClientConfig) -> SharedHttpClient { +// let mut builder = Builder::new_v2(); +// if let Some(timeout) = config.idle_timeout { +// builder = builder.pool_idle_timeout(timeout); +// } +// if let Some(n) = config.max_connections { +// builder = builder.max_connections(n); +// } +// builder.build_http() +// } +// } + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +fn runtime_components() -> aws_smithy_runtime_api::client::runtime_components::RuntimeComponents { + RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(SystemTimeSource::new())) + .build() + .expect("valid runtime components") +} + +async fn send_to( + client: &SharedHttpClient, + url: &str, +) -> Result< + aws_smithy_runtime_api::client::orchestrator::HttpResponse, + aws_smithy_runtime_api::client::result::ConnectorError, +> { + let settings = HttpConnectorSettings::builder().build(); + let components = runtime_components(); + let connector = client.http_connector(&settings, &components); + connector + .call(HttpRequest::get(url).expect("valid HTTP request")) + .await +} + +fn localhost_url(server: &WireMockServer) -> String { + let endpoint = server.endpoint_url(); + let port = endpoint + .rsplit(':') + .next() + .expect("endpoint URL should contain port"); + format!("http://127.0.0.1:{port}/") +} + +// --------------------------------------------------------------------------- +// Test implementations +// --------------------------------------------------------------------------- + +/// Connection reuse via HTTP/1.1 keep-alive: sequential requests reuse one TCP connection. +async fn connection_reuse(make: &dyn MakeClient) { + let server = WireMockServer::start(vec![ + ReplayedEvent::status(200), + ReplayedEvent::status(200), + ReplayedEvent::status(200), + ]) + .await; + + let client = make.make(ClientConfig::default()); + let url = localhost_url(&server); + + for i in 1..=3 { + let resp = send_to(&client, &url) + .await + .unwrap_or_else(|e| panic!("request {i} should succeed: {e}")); + assert_eq!(resp.status().as_u16(), 200, "request {i} should return 200"); + } + + match_events!(ev!(connect), ev!(http(200)), ev!(http(200)), ev!(http(200)))(&server.events()); +} + +/// Idle timeout eviction: a connection idle past the timeout is discarded. +async fn idle_timeout_eviction(make: &dyn MakeClient) { + let server = + WireMockServer::start(vec![ReplayedEvent::status(200), ReplayedEvent::status(200)]).await; + + let idle_timeout = Duration::from_millis(100); + let client = make.make(ClientConfig::default().with_idle_timeout(idle_timeout)); + let url = localhost_url(&server); + + let resp = send_to(&client, &url) + .await + .expect("first request should succeed"); + assert_eq!(resp.status().as_u16(), 200); + + tokio::time::sleep(idle_timeout * 2).await; + + let resp = send_to(&client, &url) + .await + .expect("second request should succeed after idle eviction"); + assert_eq!(resp.status().as_u16(), 200); + + match_events!(ev!(connect), ev!(http(200)), ev!(connect), ev!(http(200)))(&server.events()); +} + +/// A server that resets on connect should surface as a connector error. +async fn connection_reset_returns_error(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint(IP1, vec![ConnectionBehavior::ResetOnConnect]) + .build() + .await; + + let client = make.make(ClientConfig::default()); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + let result = send_to(&client, &url).await; + assert!( + result.is_err(), + "request to a reset-on-connect endpoint should return an error" + ); +} + +/// The client should report correct connector metadata. +async fn connector_metadata(make: &dyn MakeClient) { + let client = make.make(ClientConfig::default()); + let metadata = client + .connector_metadata() + .expect("connector_metadata should return Some"); + assert_eq!(metadata.name(), Cow::Borrowed("hyper")); + assert_eq!( + metadata.version(), + Some(Cow::Borrowed("1.x")), + "expected hyper 1.x connector version" + ); +} + +// --------------------------------------------------------------------------- +// v1 test runners +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn v1_connection_reuse() { + connection_reuse(&V1Client).await; +} + +#[tokio::test] +async fn v1_idle_timeout_eviction() { + idle_timeout_eviction(&V1Client).await; +} + +#[tokio::test] +async fn v1_connection_reset_returns_error() { + connection_reset_returns_error(&V1Client).await; +} + +#[tokio::test] +async fn v1_connector_metadata() { + connector_metadata(&V1Client).await; +} + +// When Builder::new_v2() exists, add: +// #[tokio::test] +// async fn v2_connection_reuse() { connection_reuse(&V2Client).await; } +// #[tokio::test] +// async fn v2_idle_timeout_eviction() { idle_timeout_eviction(&V2Client).await; } +// #[tokio::test] +// async fn v2_connection_reset_returns_error() { connection_reset_returns_error(&V2Client).await; } +// #[tokio::test] +// async fn v2_connector_metadata() { connector_metadata(&V2Client).await; } diff --git a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs index ad337624959..bbc06530978 100644 --- a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs +++ b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs @@ -791,6 +791,7 @@ async fn test_explicit_proxy_disable_overrides_environment() { /// Helper function to make HTTPS requests through proxy using TLS providers /// This is similar to make_http_request_through_proxy but uses TLS-enabled connectors +#[cfg(any(feature = "rustls-ring", feature = "s2n-tls"))] async fn make_https_request_through_proxy( proxy_config: ProxyConfig, target_url: &str, @@ -828,6 +829,7 @@ async fn make_https_request_through_proxy( /// Generic test function for HTTPS CONNECT with authentication /// Tests that HTTPS requests through HTTP proxy use CONNECT method with proper auth headers +#[cfg(any(feature = "rustls-ring", feature = "s2n-tls"))] async fn run_https_connect_with_auth_test(tls_provider: tls::Provider, provider_name: &str) { let mock_proxy = MockProxyServer::new(|req| { // For HTTPS through HTTP proxy, we should see a CONNECT request @@ -879,6 +881,7 @@ async fn run_https_connect_with_auth_test(tls_provider: tls::Provider, provider_ /// Generic test function for CONNECT without authentication (should get 407) /// Tests that HTTPS requests without auth get proper 407 response +#[cfg(any(feature = "rustls-ring", feature = "s2n-tls"))] async fn run_https_connect_auth_required_test(tls_provider: tls::Provider, provider_name: &str) { let mock_proxy = MockProxyServer::new(|req| { // For HTTPS through HTTP proxy, we should see a CONNECT request @@ -1123,6 +1126,7 @@ async fn test_uri_form_proxy_vs_direct() { /// Generic test function for CONNECT URI form validation /// Tests that CONNECT requests use the correct host:port format +#[cfg(any(feature = "rustls-ring", feature = "s2n-tls"))] async fn run_connect_uri_form_test(tls_provider: tls::Provider, provider_name: &str) { let target_host = "secure.example.com"; let target_port = 443; diff --git a/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs b/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs index 41600db4802..bde0688df6d 100644 --- a/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs +++ b/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs @@ -16,13 +16,14 @@ use aws_smithy_runtime::client::http::test_util::wire::{ use aws_smithy_runtime::client::orchestrator::operation::Operation; use aws_smithy_runtime::test_util::capture_test_logs::capture_test_logs; use aws_smithy_runtime::{ev, match_events}; +use aws_smithy_runtime_api::client::http::SharedHttpClient; use aws_smithy_runtime_api::client::interceptors::context::InterceptorContext; use aws_smithy_runtime_api::client::orchestrator::OrchestratorError; use aws_smithy_runtime_api::client::retries::classifiers::{ClassifyRetry, RetryAction}; +use aws_smithy_runtime_api::shared::IntoShared; use aws_smithy_types::body::SdkBody; use aws_smithy_types::retry::{ErrorKind, ProvideErrorKind, ReconnectMode, RetryConfig}; use aws_smithy_types::timeout::TimeoutConfig; -use hyper_0_14::client::Builder as HyperBuilder; use std::fmt; use std::time::Duration; @@ -55,9 +56,7 @@ struct TestRetryClassifier; impl ClassifyRetry for TestRetryClassifier { fn classify_retry(&self, ctx: &InterceptorContext) -> RetryAction { tracing::info!("classifying retry for {ctx:?}"); - // Check for a result let output_or_error = ctx.output_or_error(); - // Check for an error let error = match output_or_error { Some(Ok(_)) | None => return RetryAction::NoActionIndicated, Some(Err(err)) => err, @@ -86,44 +85,58 @@ impl ClassifyRetry for TestRetryClassifier { } } -async fn h1_and_h2(events: Vec, match_clause: impl Fn(&[RecordedEvent])) { - wire_level_test( - events.clone(), - |_b| {}, - ReconnectMode::ReconnectOnTransientError, - &match_clause, - ) - .await; - wire_level_test( - events, - |b| { - b.http2_only(true); - }, - ReconnectMode::ReconnectOnTransientError, - match_clause, - ) - .await; - tracing::info!("h2 ok!"); +/// MakeClient — parameterizes tests over HTTP stacks +trait MakeClient: Send + Sync { + fn make(&self, mock: &WireMockServer) -> SharedHttpClient; +} + +/// hyper 0.14 legacy stack +struct Hyper014Client; + +impl MakeClient for Hyper014Client { + fn make(&self, mock: &WireMockServer) -> SharedHttpClient { + HyperClientBuilder::new() + .build(hyper_0_14::client::HttpConnector::new_with_resolver( + mock.dns_resolver(), + )) + .into_shared() + } +} + +/// hyper 0.14 with HTTP/2 only +struct Hyper014H2Client; + +impl MakeClient for Hyper014H2Client { + fn make(&self, mock: &WireMockServer) -> SharedHttpClient { + let mut hyper_builder = hyper_0_14::Client::builder(); + hyper_builder.http2_only(true); + HyperClientBuilder::new() + .hyper_builder(hyper_builder) + .build(hyper_0_14::client::HttpConnector::new_with_resolver( + mock.dns_resolver(), + )) + .into_shared() + } +} + +/// hyper 1.x stack via public Builder API +struct Hyper1xClient; + +impl MakeClient for Hyper1xClient { + fn make(&self, mock: &WireMockServer) -> SharedHttpClient { + aws_smithy_http_client::Builder::new().build_with_resolver(mock.dns_resolver()) + } } -/// Repeatedly send test operation until `end_of_test` is received -/// -/// When the test is over, match_clause is evaluated -async fn wire_level_test( +/// Repeatedly send test operation until `end_of_test` is received, then run match_clause. +async fn run_test( + make_client: &dyn MakeClient, events: Vec, - hyper_builder_settings: impl Fn(&mut HyperBuilder), reconnect_mode: ReconnectMode, match_clause: impl Fn(&[RecordedEvent]), ) { - let mut hyper_builder = hyper_0_14::Client::builder(); - hyper_builder_settings(&mut hyper_builder); - let mock = WireMockServer::start(events).await; - let http_client = HyperClientBuilder::new() - .hyper_builder(hyper_builder) - .build(hyper_0_14::client::HttpConnector::new_with_resolver( - mock.dns_resolver(), - )); + let http_client = make_client.make(&mock); let operation = Operation::builder() .service_name("test") @@ -143,10 +156,9 @@ async fn wire_level_test( .serializer({ let endpoint_url = mock.endpoint_url(); move |_| { - let request = http_02x::Request::builder() + let request = http_1x::Request::builder() .uri(endpoint_url.clone()) - // Make the body non-replayable since we don't actually want to retry - .body(SdkBody::from_body_0_4(SdkBody::from("body"))) + .body(SdkBody::from("body")) .unwrap() .try_into() .unwrap(); @@ -193,29 +205,54 @@ async fn wire_level_test( mock.shutdown(); } +/// Run a test against all HTTP stacks +async fn all_stacks( + events: Vec, + reconnect_mode: ReconnectMode, + match_clause: impl Fn(&[RecordedEvent]), +) { + run_test( + &Hyper014Client, + events.clone(), + reconnect_mode, + &match_clause, + ) + .await; + run_test( + &Hyper014H2Client, + events.clone(), + reconnect_mode, + &match_clause, + ) + .await; + run_test(&Hyper1xClient, events, reconnect_mode, &match_clause).await; +} + #[tokio::test] async fn non_transient_errors_no_reconnect() { let _logs = capture_test_logs(); - h1_and_h2( + all_stacks( vec![ ReplayedEvent::status(400), ReplayedEvent::with_body(END_OF_TEST), ], + ReconnectMode::ReconnectOnTransientError, match_events!(ev!(dns), ev!(connect), ev!(http(400)), ev!(http(200))), ) - .await + .await; } #[tokio::test] async fn reestablish_dns_on_503() { let _logs = capture_test_logs(); - h1_and_h2( + all_stacks( vec![ ReplayedEvent::status(503), ReplayedEvent::status(503), ReplayedEvent::status(503), ReplayedEvent::with_body(END_OF_TEST), ], + ReconnectMode::ReconnectOnTransientError, match_events!( // first request ev!(dns), @@ -241,13 +278,14 @@ async fn reestablish_dns_on_503() { #[tokio::test] async fn connection_shared_on_success() { let _logs = capture_test_logs(); - h1_and_h2( + all_stacks( vec![ ReplayedEvent::ok(), ReplayedEvent::ok(), ReplayedEvent::status(503), ReplayedEvent::with_body(END_OF_TEST), ], + ReconnectMode::ReconnectOnTransientError, match_events!( ev!(dns), ev!(connect), @@ -265,12 +303,11 @@ async fn connection_shared_on_success() { #[tokio::test] async fn no_reconnect_when_disabled() { let _logs = capture_test_logs(); - wire_level_test( + all_stacks( vec![ ReplayedEvent::status(503), ReplayedEvent::with_body(END_OF_TEST), ], - |_b| {}, ReconnectMode::ReuseAllConnections, match_events!(ev!(dns), ev!(connect), ev!(http(503)), ev!(http(200))), ) @@ -280,7 +317,7 @@ async fn no_reconnect_when_disabled() { #[tokio::test] async fn connection_reestablished_after_timeout() { let _logs = capture_test_logs(); - h1_and_h2( + all_stacks( vec![ ReplayedEvent::ok(), ReplayedEvent::Timeout, @@ -288,6 +325,7 @@ async fn connection_reestablished_after_timeout() { ReplayedEvent::Timeout, ReplayedEvent::with_body(END_OF_TEST), ], + ReconnectMode::ReconnectOnTransientError, match_events!( // first connection ev!(dns), From 4653957e5a3a1286ecf6e5bb799bfc29a3b6580f Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Mon, 27 Apr 2026 16:56:02 -0400 Subject: [PATCH 03/47] bootstrap pool shape --- NOTICE | 8 + rust-runtime/Cargo.lock | 2 + .../aws-smithy-http-client/Cargo.toml | 6 + .../aws-smithy-http-client/src/client/pool.rs | 238 +++++++- .../src/client/pool/connection.rs | 317 +++++++++++ .../src/client/pool/handshake.rs | 225 ++++++++ .../src/client/pool/vendored_cache.rs | 517 ++++++++++++++++++ 7 files changed, 1306 insertions(+), 7 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs create mode 100644 rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs create mode 100644 rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs diff --git a/NOTICE b/NOTICE index 616fc588945..f1ba3e006b8 100644 --- a/NOTICE +++ b/NOTICE @@ -1 +1,9 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +This product includes software developed by Sean McArthur and contributors +to the hyper-util project (https://github.com/hyperium/hyper-util), licensed +under the MIT License. The applicable copyright notice and MIT license text +are preserved at the top of each file that incorporates hyper-util sources. +See: + + rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock index c3ac2aebed4..7aa7ec12069 100644 --- a/rust-runtime/Cargo.lock +++ b/rust-runtime/Cargo.lock @@ -506,6 +506,7 @@ dependencies = [ "aws-smithy-types 1.6.0", "base64 0.22.1", "bytes", + "futures-util", "h2 0.3.27", "h2 0.4.14", "http 0.2.12", @@ -535,6 +536,7 @@ dependencies = [ "tokio", "tokio-rustls 0.26.4", "tower 0.5.3", + "tower-test", "tracing", ] diff --git a/rust-runtime/aws-smithy-http-client/Cargo.toml b/rust-runtime/aws-smithy-http-client/Cargo.toml index 0f30f3f166a..c69529efb87 100644 --- a/rust-runtime/aws-smithy-http-client/Cargo.toml +++ b/rust-runtime/aws-smithy-http-client/Cargo.toml @@ -27,6 +27,9 @@ default-client = [ "hyper-util?/client-proxy", "dep:http-1x", "dep:tower", + "tower?/util", + "dep:futures-util", + "tokio/sync", "dep:rustls-pki-types", "dep:rustls-native-certs" ] @@ -97,6 +100,7 @@ s2n-tls-hyper = { version = "0.1.0", optional = true } s2n-tls = { version = "0.3.33", optional = true } s2n-tls-tokio = { version = "0.3.33", optional = true} tower = { version = "0.5.2", optional = true } +futures-util = { version = "0.3.32", default-features = false, features = ["std"], optional = true } rustls-pki-types = { version = "1.12.0", features = ["std"], optional = true } rustls-native-certs = { version = "0.8.1", optional = true } # end hyper 1.x stack deps @@ -131,6 +135,8 @@ rustls-pemfile = "2.2.0" rustls-pki-types = { version = "1.12.0", features = ["std"] } tokio = { version = "1.49.0", features = ["macros", "rt", "rt-multi-thread", "test-util", "full"] } tokio-rustls = "0.26.2" +# used by tests vendored from hyper-util in src/client/pool/vendored_cache.rs +tower-test = "0.4" [[example]] diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index f2a0e268ddf..2f9468cff15 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -3,11 +3,235 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Connection pool implementation for the v2 HTTP client. +//! Connection pool for the v2 HTTP client. //! -//! This module composes hyper-util's pool primitives (Cache, Singleton, Negotiate, Map) -//! with SDK-owned connection lifecycle management including: -//! - Connection state tracking (idle-since, remote IP, poisoned flag) -//! - Proactive health checks on checkout -//! - Connection poisoning support -//! - Queryable pool state (idle/active counts) +//! Manages the lifecycle of HTTP connections including: +//! - Connection state tracking (idle duration, remote address, poisoned flag) +//! - Health checks on checkout +//! - Connection poisoning +//! - Idle connection eviction + +mod connection; +mod handshake; +mod vendored_cache; + +/// Connection-caching pool layer. +/// +/// Currently re-exports [`vendored_cache`], a copy of hyper-util's +/// `pool::cache` module with SDK-specific modifications. The re-export +/// insulates the rest of the pool from the vendoring detail: call sites +/// use `cache::...` regardless of whether the implementation is vendored +/// or, later, fully owned. +mod cache { + pub(crate) use super::vendored_cache::*; +} + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Mutex; + +use aws_smithy_types::body::SdkBody; +use hyper_util::client::legacy::connect::Connection as HyperConnection; +use hyper_util::client::pool as hpool; +use hyper_util::rt::TokioExecutor; +use tower::{Service, ServiceExt}; + +use connection::ManagedConnection; +use handshake::{H1ConnectAndHandshake, H2ConnectAndHandshake}; + +pub(crate) type BoxError = Box; +type BoxFuture = Pin + Send>>; + +/// Key for per-host connection pool routing. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct PoolKey { + scheme: http_1x::uri::Scheme, + authority: http_1x::uri::Authority, +} + +impl PoolKey { + fn from_uri(uri: &http_1x::Uri) -> Option { + Some(Self { + scheme: uri.scheme()?.clone(), + authority: uri.authority()?.clone(), + }) + } +} + +/// Build a connection pool for the given connector. +pub(crate) fn build_pool(connector: C) -> ConnectionPool +where + C: Service + Clone + Send + Sync + 'static, + C::Error: Into + 'static, + C::Future: Unpin + Send + 'static, + IO: hyper::rt::Read + hyper::rt::Write + HyperConnection + Unpin + Send + 'static, +{ + let make_entry = move |_uri: &http_1x::Uri| -> Box { + let stack = hpool::negotiate::builder() + .connect(connector.clone()) + .inspect(|conn: &IO| conn.connected().is_negotiated_h2()) + .fallback(tower::layer::layer_fn(|inspector| { + cache::builder() + .executor(TokioExecutor::new()) + .build(H1ConnectAndHandshake::new(inspector)) + .map_response(CachedConnection::new) + })) + .upgrade(tower::layer::layer_fn(|inspected| { + hpool::singleton::Singleton::new(H2ConnectAndHandshake::new(inspected)) + })) + .build(); + Box::new(TypedPoolEntry(stack)) + }; + + ConnectionPool { + hosts: Mutex::new(HashMap::new()), + make_entry: Box::new(make_entry), + } +} + +/// The connection pool. +/// +/// Routes requests by (scheme, authority) to per-host pool stacks. +/// Each host gets a Negotiate stack that selects between HTTP/1.1 (Cache) +/// and HTTP/2 (Singleton) based on ALPN negotiation. +pub(crate) struct ConnectionPool { + hosts: Mutex>>, + make_entry: Box Box + Send + Sync>, +} + +impl ConnectionPool { + /// Send a request through the pool. + /// + /// Routes to the appropriate per-host pool stack, acquires a connection, + /// and sends the request. + pub(crate) async fn send_request( + &self, + uri: http_1x::Uri, + req: http_1x::Request, + ) -> Result, BoxError> { + let key = PoolKey::from_uri(&uri).ok_or("request URI must have scheme and authority")?; + + // Clone the per-host stack out of the lock. The stack is Clone + // (cheap Arc clones internally). All I/O happens outside the lock. + let fut = { + let mut hosts = self.hosts.lock().unwrap(); + if !hosts.contains_key(&key) { + hosts.insert(key.clone(), (self.make_entry)(&uri)); + } + hosts.get_mut(&key).unwrap().acquire(uri) + }; + + let conn = fut.await?; + conn.send(req).await + } +} + +/// Type-erased per-host pool entry. +/// +/// Each host has one of these, wrapping the unnameable Negotiate stack. +trait PoolEntry: Send + Sync { + /// Acquire a connection from this host's pool. + fn acquire(&mut self, uri: http_1x::Uri) -> BoxFuture>; +} + +/// A type-erased connection that can send one request. +trait Connection: Send { + fn send( + self: Box, + req: http_1x::Request, + ) -> BoxFuture, BoxError>>; +} + +type BoxedConnection = Box; + +/// Concrete PoolEntry wrapping a Negotiate stack. +struct TypedPoolEntry(S); + +impl PoolEntry for TypedPoolEntry +where + S: Service + Clone + Send + Sync + 'static, + S::Error: Into + 'static, + S::Future: Send + 'static, + Conn: Service> + Send + 'static, + Conn::Response: Into>, + Conn::Error: Into, + Conn::Future: Send + 'static, +{ + fn acquire(&mut self, uri: http_1x::Uri) -> BoxFuture> { + let mut svc = self.0.clone(); + Box::pin(async move { + // Loop past stale idle entries (poisoned or dead). Converges because: + // - Cache idle set is bounded; each post-checkout `poll_ready` Err + // flips `is_closed` so `Cached::Drop` skips reinsertion, shrinking + // the set by one. + // - Singleton clears to `Empty` on `poll_ready` Err, forcing the + // next `call` to run a fresh handshake — a handshake failure + // surfaces as an error from `svc.call` (not the inner + // `poll_ready`), which we propagate. + loop { + std::future::poll_fn(|cx| svc.poll_ready(cx)) + .await + .map_err(Into::into)?; + let mut conn = svc.call(uri.clone()).await.map_err(Into::into)?; + if std::future::poll_fn(|cx| conn.poll_ready(cx)).await.is_ok() { + return Ok(Box::new(TypedConnection(conn)) as BoxedConnection); + } + // drop `conn` → triggers pool cleanup (H1: discard via + // CachedConnection::Drop; H2: Singleton already cleared state). + } + }) + } +} + +/// Concrete Connection wrapping a Negotiated service. +struct TypedConnection(S); + +impl Connection for TypedConnection +where + S: Service> + Send + 'static, + S::Response: Into>, + S::Error: Into, + S::Future: Send + 'static, +{ + fn send( + self: Box, + req: http_1x::Request, + ) -> BoxFuture, BoxError>> { + let mut svc = self.0; + Box::pin(async move { + std::future::poll_fn(|cx| svc.poll_ready(cx)) + .await + .map_err(Into::into)?; + svc.call(req).await.map(Into::into).map_err(Into::into) + }) + } +} + +/// A connection checked out from the H1 cache. +/// +/// See `connection::CachedConnection`. +pub(crate) use connection::CachedConnection; + +/// A connection checked out from the H2 singleton. +/// +/// See `connection::SingletonConnection`. +pub(crate) use connection::SingletonConnection; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pool_key() { + let uri: http_1x::Uri = "http://example.com:8080/path".parse().unwrap(); + let key = PoolKey::from_uri(&uri).unwrap(); + assert_eq!(key.authority.as_str(), "example.com:8080"); + } + + #[test] + fn test_pool_key_missing_scheme() { + let uri: http_1x::Uri = "/path".parse().unwrap(); + assert!(PoolKey::from_uri(&uri).is_none()); + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs new file mode 100644 index 00000000000..e79b8143e50 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -0,0 +1,317 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Connection state tracking for pooled connections. + +use std::net::SocketAddr; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::Instant; + +use aws_smithy_types::body::SdkBody; +use tower::Service; + +use super::cache; +use super::handshake::H1SendRequest; +use super::BoxError; + +/// Error returned by `ManagedConnection::poll_ready` when the connection +/// has been marked poisoned and should not be reused. +#[derive(Debug)] +pub(crate) struct PoisonedError; + +impl std::fmt::Display for PoisonedError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("connection poisoned") + } +} + +impl std::error::Error for PoisonedError {} + +/// Metadata about a connection captured at establishment time. +/// +/// Captured between TLS connector output and protocol handshake — the last point +/// where the raw transport stream is accessible. +#[derive(Debug, Clone)] +pub(crate) struct ConnectionInfo { + /// Remote address of the peer. `None` when the underlying connector did + /// not attach `HttpInfo` (e.g. the s2n-tls provider today). + pub(crate) remote_addr: Option, + /// Local address of this end of the connection. + pub(crate) local_addr: Option, + // Future: tls_info, timing (tcp_connect_duration, tls_handshake_duration) +} + +/// A connection with SDK-owned lifecycle metadata. +/// +/// Wraps the inner service (typically `SendRequest`) with state needed +/// for pool management: idle tracking, poisoning, and connection identity. +/// +/// Clone is supported when the inner service is Clone (e.g., HTTP/2 multiplexed +/// connections). Clones share the same poison flag and connection metadata. +pub(crate) struct ManagedConnection { + inner: S, + info: ConnectionInfo, + created_at: Instant, + last_used: Instant, + poisoned: Arc, +} + +impl ManagedConnection { + /// Create a new managed connection wrapping the given service. + pub(crate) fn new(inner: S, info: ConnectionInfo) -> Self { + let now = Instant::now(); + Self { + inner, + info, + created_at: now, + last_used: now, + poisoned: Arc::new(AtomicBool::new(false)), + } + } + + /// Mark this connection as poisoned so it will not be reused. + /// + /// Poisoned connections are skipped on checkout and dropped on return to the pool. + pub(crate) fn poison(&self) { + self.poisoned.store(true, Ordering::Release); + } + + /// Returns a closure that poisons this connection when called. + /// + /// The returned closure can be passed to components that need to signal + /// connection failure without holding a direct reference to the connection. + pub(crate) fn poison_fn(&self) -> impl Fn() + Send + Sync + 'static { + let flag = self.poisoned.clone(); + move || flag.store(true, Ordering::Release) + } + + /// Whether this connection has been poisoned. + pub(crate) fn is_poisoned(&self) -> bool { + self.poisoned.load(Ordering::Acquire) + } + + /// How long this connection has been idle (since last request completed). + pub(crate) fn idle_duration(&self) -> std::time::Duration { + self.last_used.elapsed() + } + + /// Connection metadata. + pub(crate) fn info(&self) -> &ConnectionInfo { + &self.info + } + + /// When this connection was established. + pub(crate) fn created_at(&self) -> Instant { + self.created_at + } + + /// Mark this connection as recently used. + pub(crate) fn touch(&mut self) { + self.last_used = Instant::now(); + } + + /// Mutable access to the inner service. + pub(crate) fn inner_mut(&mut self) -> &mut S { + &mut self.inner + } +} + +impl Clone for ManagedConnection { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + info: self.info.clone(), + created_at: self.created_at, + last_used: self.last_used, + poisoned: self.poisoned.clone(), + } + } +} + +impl Service> for ManagedConnection +where + S: Service>, + S::Error: Into, + S::Future: Send + 'static, + S::Response: 'static, +{ + type Response = S::Response; + type Error = BoxError; + type Future = std::pin::Pin< + Box> + Send>, + >; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + if self.is_poisoned() { + return Poll::Ready(Err(PoisonedError.into())); + } + self.inner_mut().poll_ready(cx).map_err(Into::into) + } + + fn call(&mut self, req: http_1x::Request) -> Self::Future { + self.touch(); + let fut = self.inner_mut().call(req); + Box::pin(async move { fut.await.map_err(Into::into) }) + } +} + +/// A connection checked out from the H1 cache. +/// +/// Wraps `cache::Cached>` so that poisoned connections +/// are dropped from the pool (via `Cached::discard`) instead of being +/// returned for reuse when dropped. Healthy connections return to the pool +/// normally through `Cached::Drop`. +pub(crate) struct CachedConnection { + inner: Option>>, +} + +impl CachedConnection { + pub(crate) fn new(cached: cache::Cached>) -> Self { + Self { + inner: Some(cached), + } + } +} + +impl Service for CachedConnection +where + cache::Cached>: Service, +{ + type Response = > as Service>::Response; + type Error = > as Service>::Error; + type Future = > as Service>::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().unwrap().poll_ready(cx) + } + + fn call(&mut self, req: Req) -> Self::Future { + self.inner.as_mut().unwrap().call(req) + } +} + +impl Drop for CachedConnection { + fn drop(&mut self) { + if let Some(cached) = self.inner.take() { + if cached.inner().is_poisoned() { + cached.discard(); + } + // else: cached drops normally, returning the connection to the cache + } + } +} + +/// A connection checked out from the H2 singleton. +/// +/// Symmetric wrapper to `CachedConnection` for the H2 leg. The inner type +/// is parameterized because hyper-util's `Singled` is unnameable outside +/// its crate; we accept whatever type comes out of the singleton. Gives us +/// a named, pool-owned H2 handle we can hold as a guard, inspect, and +/// instrument later without vendoring `singleton.rs`. +/// +/// No Drop hook: dropping a Singled clone does not affect singleton state; +/// dead/poisoned H2 connections are handled reactively via +/// `Singled::poll_ready` clearing the state to `Empty`. +pub(crate) struct SingletonConnection { + inner: T, +} + +impl SingletonConnection { + pub(crate) fn new(inner: T) -> Self { + Self { inner } + } +} + +impl Service for SingletonConnection +where + T: Service, +{ + type Response = T::Response; + type Error = T::Error; + type Future = T::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Req) -> Self::Future { + self.inner.call(req) + } +} + +/// A response body that keeps the originating pool checkout alive until +/// the body is fully consumed. +/// +/// # Why this exists +/// +/// When a checked-out pool connection's `Service::call` returns, the +/// response head is available but the body may still be streaming over +/// the same underlying HTTP connection. For H1 this is critical — if the +/// `CachedConnection` drops at that point, the connection returns to the +/// pool mid-body-stream and the next checkout would be handed a still-busy +/// connection. +/// +/// `ResponseBody` holds the checkout (`CachedConnection` for H1, +/// `SingletonConnection` for H2) in its guard field until the body is +/// fully dropped. Body streaming continues through the held inner +/// `Incoming`; when the `ResponseBody` is dropped the guard drops, which +/// for H1 triggers `CachedConnection::Drop` (return-to-pool or `discard` +/// if poisoned), and for H2 simply drops the singleton clone (no-op on +/// singleton state). +/// +/// The H2 variant carries a generic type parameter because +/// `SingletonConnection`'s inner `T` is `hyper_util::client::pool:: +/// singleton::Singled<...>`, which is unnameable outside hyper-util. The +/// H1 variant is fully concrete since our vendored cache makes +/// `cache::Cached<...>` nameable. +pub(crate) struct ResponseBody { + inner: hyper::body::Incoming, + _guard: ConnectionGuard, +} + +/// What a `ResponseBody` holds alive while the body streams. +/// +/// Explicit per-leg variants so H1 vs H2 bifurcation is visible at the +/// type level and in debugger output. +pub(crate) enum ConnectionGuard { + H1(CachedConnection>), + H2(SingletonConnection), +} + +impl ResponseBody { + pub(crate) fn new(inner: hyper::body::Incoming, guard: ConnectionGuard) -> Self { + Self { + inner, + _guard: guard, + } + } +} + +impl hyper::body::Body for ResponseBody +where + H2Inner: Unpin, +{ + type Data = ::Data; + type Error = ::Error; + + fn poll_frame( + self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + // Safe: we only project to `inner`, guard is never moved. + let this = unsafe { self.get_unchecked_mut() }; + std::pin::Pin::new(&mut this.inner).poll_frame(cx) + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + + fn size_hint(&self) -> hyper::body::SizeHint { + self.inner.size_hint() + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs new file mode 100644 index 00000000000..08ee3bc65db --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -0,0 +1,225 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tower service adapters for hyper's HTTP protocol handshake. +//! +//! Provides two layers: +//! - `H1SendRequest` / `H2SendRequest`: tower `Service` wrappers around +//! hyper's `SendRequest` types. +//! - `H1ConnectAndHandshake` / `H2ConnectAndHandshake`: services that chain +//! a connector with the protocol handshake, producing a `ManagedConnection`. + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use aws_smithy_types::body::SdkBody; +use hyper_util::client::legacy::connect::{Connection, HttpInfo}; +use hyper_util::rt::TokioExecutor; +use tower::Service; + +use super::connection::{ConnectionInfo, ManagedConnection}; + +// --------------------------------------------------------------------------- +// Tower Service wrappers for hyper's SendRequest +// --------------------------------------------------------------------------- + +/// Tower `Service` adapter for `hyper::client::conn::http1::SendRequest`. +pub(crate) struct H1SendRequest { + inner: hyper::client::conn::http1::SendRequest, +} + +impl H1SendRequest { + pub(crate) fn new(inner: hyper::client::conn::http1::SendRequest) -> Self { + Self { inner } + } +} + +impl Service> for H1SendRequest { + type Response = http_1x::Response; + type Error = hyper::Error; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: http_1x::Request) -> Self::Future { + Box::pin(self.inner.send_request(req)) + } +} + +/// Tower `Service` adapter for `hyper::client::conn::http2::SendRequest`. +/// +/// Clone is supported — HTTP/2 multiplexes requests over a single connection. +#[derive(Clone)] +pub(crate) struct H2SendRequest { + inner: hyper::client::conn::http2::SendRequest, +} + +impl H2SendRequest { + pub(crate) fn new(inner: hyper::client::conn::http2::SendRequest) -> Self { + Self { inner } + } +} + +impl Service> for H2SendRequest { + type Response = http_1x::Response; + type Error = hyper::Error; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: http_1x::Request) -> Self::Future { + Box::pin(self.inner.send_request(req)) + } +} + +// --------------------------------------------------------------------------- +// Connect-and-handshake services +// --------------------------------------------------------------------------- + +use super::BoxError; + +/// Extract `ConnectionInfo` from a just-connected IO stream. +/// +/// Reads `Connected` metadata (ALPN, proxy flag, attached extras) and pulls out +/// `HttpInfo` (remote/local addr) if the connector attached it. This is the single +/// place we bridge between hyper-util's `Connection` trait and our own +/// `ConnectionInfo` type — if upstream reshapes, only this function changes. +fn capture_info(io: &IO) -> ConnectionInfo { + let connected = io.connected(); + let mut extras = http_1x::Extensions::new(); + connected.get_extras(&mut extras); + let http_info = extras.get::(); + ConnectionInfo { + remote_addr: http_info.map(|i| i.remote_addr()), + local_addr: http_info.map(|i| i.local_addr()), + } +} + +/// Connects and performs an HTTP/1.1 handshake. +/// +/// Wraps a connector service, performing the H1 protocol handshake after +/// connection and returning a `ManagedConnection`. +pub(crate) struct H1ConnectAndHandshake { + connector: C, +} + +impl H1ConnectAndHandshake { + pub(crate) fn new(connector: C) -> Self { + Self { connector } + } +} + +impl Clone for H1ConnectAndHandshake { + fn clone(&self) -> Self { + Self { + connector: self.connector.clone(), + } + } +} + +impl Service for H1ConnectAndHandshake +where + C: Service, + C::Error: Into + 'static, + C::Future: Send + 'static, + IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, +{ + type Response = ManagedConnection; + type Error = BoxError; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.connector.poll_ready(cx).map_err(Into::into) + } + + fn call(&mut self, uri: http_1x::Uri) -> Self::Future { + let fut = self.connector.call(uri); + Box::pin(async move { + let io = fut.await.map_err(Into::into)?; + let info = capture_info(&io); + + let (tx, conn) = hyper::client::conn::http1::Builder::new() + .handshake(io) + .await + .map_err(|e| Box::new(e) as BoxError)?; + + // Drive the connection I/O in the background + tokio::spawn(async move { + if let Err(e) = conn.with_upgrades().await { + tracing::debug!(error = %e, "h1 connection driver error"); + } + }); + + Ok(ManagedConnection::new(H1SendRequest::new(tx), info)) + }) + } +} + +/// Connects and performs an HTTP/2 handshake. +/// +/// Wraps a connector service, performing the H2 protocol handshake after +/// connection and returning a `ManagedConnection`. +/// +/// Generic over the request type because the Negotiate upgrade path +/// receives `()` (connection already established), not a URI. +pub(crate) struct H2ConnectAndHandshake { + connector: C, +} + +impl H2ConnectAndHandshake { + pub(crate) fn new(connector: C) -> Self { + Self { connector } + } +} + +impl Clone for H2ConnectAndHandshake { + fn clone(&self) -> Self { + Self { + connector: self.connector.clone(), + } + } +} + +impl Service for H2ConnectAndHandshake +where + C: Service, + C::Error: Into + 'static, + C::Future: Send + 'static, + IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, +{ + type Response = ManagedConnection; + type Error = BoxError; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.connector.poll_ready(cx).map_err(Into::into) + } + + fn call(&mut self, req: Req) -> Self::Future { + let fut = self.connector.call(req); + Box::pin(async move { + let io = fut.await.map_err(Into::into)?; + let info = capture_info(&io); + + let (tx, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new()) + .handshake(io) + .await + .map_err(|e| Box::new(e) as BoxError)?; + + tokio::spawn(async move { + if let Err(e) = conn.await { + tracing::debug!(error = %e, "h2 connection driver error"); + } + }); + + Ok(ManagedConnection::new(H2SendRequest::new(tx), info)) + }) + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs new file mode 100644 index 00000000000..0ecd6a56944 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs @@ -0,0 +1,517 @@ +/* + * Portions of this file are derived from hyper-util + * (https://github.com/hyperium/hyper-util), licensed under MIT: + * + * Copyright (c) 2023-2025 Sean McArthur + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files (the + * "Software"), to deal in the Software without restriction, including + * without limitation the rights to use, copy, modify, merge, publish, + * distribute, sublicense, and/or sell copies of the Software, and to + * permit persons to whom the Software is furnished to do so, subject + * to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + * + * Modifications by Amazon.com, Inc. or its affiliates: + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * The derivative work as a whole is licensed under Apache-2.0 as part of the + * smithy-rs project. The MIT notice above applies to the original portions + * as required by that license. + * + * Source: hyper-util src/client/pool/cache.rs + * Upstream: https://github.com/hyperium/hyper-util + * Commit: e1c5a6c89bfaed11fb34bd483fe9ba616f403791 + * + * Modifications from upstream: + * 1. Replaced `tower_service::Service` imports with `tower::Service` (the + * `tower` crate re-exports `Service` at the crate root, avoiding an + * extra direct dependency on `tower-service`). + * 2. Changed `pub use self::internal::builder;` to `pub(crate) use ...` + * and dropped the three `#[cfg(docsrs)] pub use` lines. The composable + * pool internals are `pub(crate)` in this crate — no types from this + * file are exposed in the smithy-rs public API. + * 3. Added `#![allow(dead_code, unreachable_pub)]` so the file can be + * kept close to upstream even when individual items aren't used yet. + * 4. Dropped the module- and struct-level rustdoc sections that reference + * "Unnameable" (that rustdoc pattern is specific to hyper-util's public + * API; these types aren't public here). + * 5. Added `Cached::discard(self)` — consumes self and prevents + * reinsertion into the pool on drop. See `// SDK MODIFICATION` marker + * below. Used to drop connections that a caller has learned are bad + * (poisoned, GOAWAY, etc.) between checkout and `poll_ready`. The + * upstream API requires the inner service's `poll_ready` to error + * in order to skip reinsertion, which overloads `poll_ready` semantics. + */ + +//! A cache of services +//! +//! The cache is a single list of cached services, bundled with a `MakeService`. +//! Calling the cache returns either an existing service, or makes a new one. +//! The returned `impl Service` can be used to send requests, and when dropped, +//! it will try to be returned back to the cache. + +#![allow(dead_code, unreachable_pub)] + +pub(crate) use self::internal::{builder, Cached}; + +mod internal { + use std::fmt; + use std::future::Future; + use std::pin::Pin; + use std::sync::{Arc, Mutex, Weak}; + use std::task::{self, ready, Poll}; + + use futures_util::future; + use tokio::sync::oneshot; + use tower::Service; + + use super::events; + + /// Start a builder to construct a `Cache` pool. + pub fn builder() -> Builder { + Builder { + events: events::Ignore, + } + } + + /// A cache pool of services from the inner make service. + #[derive(Debug)] + pub struct Cache + where + M: Service, + { + connector: M, + shared: Arc>>, + events: Ev, + } + + /// A builder to configure a `Cache`. + #[derive(Debug)] + pub struct Builder { + events: Ev, + } + + /// A cached service returned from a [`Cache`]. + /// + /// Implements `Service` by delegating to the inner service. Once dropped, + /// tries to reinsert into the `Cache`. + pub struct Cached { + is_closed: bool, + inner: Option, + shared: Weak>>, + // todo: on_idle + } + + pub enum CacheFuture + where + M: Service, + { + Racing { + shared: Arc>>, + select: future::Select, M::Future>, + events: Ev, + }, + Connecting { + // TODO: could be Weak even here... + shared: Arc>>, + future: M::Future, + }, + Cached { + svc: Option>, + }, + } + + // shouldn't be pub + #[derive(Debug)] + pub struct Shared { + services: Vec, + waiters: Vec>, + } + + // impl Builder + + impl Builder { + /// Provide a `Future` executor to be used by the `Cache`. + pub fn executor(self, exec: E) -> Builder> { + Builder { + events: events::WithExecutor(exec), + } + } + + /// Build a `Cache` pool around the `connector`. + pub fn build(self, connector: M) -> Cache + where + M: Service, + { + Cache { + connector, + events: self.events, + shared: Arc::new(Mutex::new(Shared { + services: Vec::new(), + waiters: Vec::new(), + })), + } + } + } + + // impl Cache + + impl Cache + where + M: Service, + { + /// Retain all cached services indicated by the predicate. + pub fn retain(&mut self, predicate: F) + where + F: FnMut(&mut M::Response) -> bool, + { + self.shared.lock().unwrap().services.retain_mut(predicate); + } + + /// Check whether this cache has no cached services. + pub fn is_empty(&self) -> bool { + self.shared.lock().unwrap().services.is_empty() + } + } + + impl Service for Cache + where + M: Service, + M::Future: Unpin, + M::Response: Unpin, + Ev: events::Events> + Clone + Unpin, + { + type Response = Cached; + type Error = M::Error; + type Future = CacheFuture; + + fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll> { + if !self.shared.lock().unwrap().services.is_empty() { + Poll::Ready(Ok(())) + } else { + self.connector.poll_ready(cx) + } + } + + fn call(&mut self, target: Dst) -> Self::Future { + // 1. If already cached, easy! + let waiter = { + let mut locked = self.shared.lock().unwrap(); + if let Some(found) = locked.take() { + return CacheFuture::Cached { + svc: Some(Cached::new(found, Arc::downgrade(&self.shared))), + }; + } + + let (tx, rx) = oneshot::channel(); + locked.waiters.push(tx); + rx + }; + + // 2. Otherwise, we start a new connect, and also listen for + // any newly idle. + CacheFuture::Racing { + shared: self.shared.clone(), + select: future::select(waiter, self.connector.call(target)), + events: self.events.clone(), + } + } + } + + impl Clone for Cache + where + M: Service + Clone, + Ev: Clone, + { + fn clone(&self) -> Self { + Self { + connector: self.connector.clone(), + events: self.events.clone(), + shared: self.shared.clone(), + } + } + } + + impl Future for CacheFuture + where + M: Service, + M::Future: Unpin, + M::Response: Unpin, + Ev: events::Events> + Unpin, + { + type Output = Result, M::Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll { + loop { + match &mut *self.as_mut() { + CacheFuture::Racing { + shared, + select, + events, + } => { + match ready!(Pin::new(select).poll(cx)) { + future::Either::Left((Err(_pool_closed), connecting)) => { + // pool was dropped, so we'll never get it from a waiter, + // but if this future still exists, then the user still + // wants a connection. just wait for the connecting + *self = CacheFuture::Connecting { + shared: shared.clone(), + future: connecting, + }; + } + future::Either::Left((Ok(pool_got), connecting)) => { + events.on_race_lost(BackgroundConnect { + future: connecting, + shared: Arc::downgrade(&shared), + }); + return Poll::Ready(Ok(Cached::new( + pool_got, + Arc::downgrade(&shared), + ))); + } + future::Either::Right((connected, _waiter)) => { + let inner = connected?; + return Poll::Ready(Ok(Cached::new( + inner, + Arc::downgrade(&shared), + ))); + } + } + } + CacheFuture::Connecting { shared, future } => { + let inner = ready!(Pin::new(future).poll(cx))?; + return Poll::Ready(Ok(Cached::new(inner, Arc::downgrade(&shared)))); + } + CacheFuture::Cached { svc } => { + return Poll::Ready(Ok(svc.take().unwrap())); + } + } + } + } + } + + // impl Cached + + impl Cached { + fn new(inner: S, shared: Weak>>) -> Self { + Cached { + is_closed: false, + inner: Some(inner), + shared, + } + } + + // TODO: inner()? looks like `tower` likes `get_ref()` and `get_mut()`. + + /// Get a reference to the inner service. + pub fn inner(&self) -> &S { + self.inner.as_ref().expect("inner only taken in drop") + } + + /// Get a mutable reference to the inner service. + pub fn inner_mut(&mut self) -> &mut S { + self.inner.as_mut().expect("inner only taken in drop") + } + + // SDK MODIFICATION: added `discard` so callers can prevent a bad + // connection from returning to the pool without having to cause a + // synthetic `poll_ready` error. + /// Prevent this cached service from being returned to the pool. + /// + /// Consumes self — the normal `Drop` then runs, sees `is_closed`, + /// and skips reinsertion. + pub fn discard(mut self) { + self.is_closed = true; + } + } + + impl Service for Cached + where + S: Service, + { + type Response = S::Response; + type Error = S::Error; + type Future = S::Future; + + fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll> { + self.inner.as_mut().unwrap().poll_ready(cx).map_err(|err| { + self.is_closed = true; + err + }) + } + + fn call(&mut self, req: Req) -> Self::Future { + self.inner.as_mut().unwrap().call(req) + } + } + + impl Drop for Cached { + fn drop(&mut self) { + if self.is_closed { + return; + } + if let Some(value) = self.inner.take() { + if let Some(shared) = self.shared.upgrade() { + if let Ok(mut shared) = shared.lock() { + shared.put(value); + } + } + } + } + } + + impl fmt::Debug for Cached { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("Cached") + .field(self.inner.as_ref().unwrap()) + .finish() + } + } + + // impl Shared + + impl Shared { + fn put(&mut self, val: V) { + let mut val = Some(val); + while let Some(tx) = self.waiters.pop() { + if !tx.is_closed() { + match tx.send(val.take().unwrap()) { + Ok(()) => break, + Err(v) => { + val = Some(v); + } + } + } + } + + if let Some(val) = val { + self.services.push(val); + } + } + + fn take(&mut self) -> Option { + // TODO: take in a loop + self.services.pop() + } + } + + pub struct BackgroundConnect { + future: CF, + shared: Weak>>, + } + + impl Future for BackgroundConnect + where + CF: Future> + Unpin, + { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll { + match ready!(Pin::new(&mut self.future).poll(cx)) { + Ok(svc) => { + if let Some(shared) = self.shared.upgrade() { + if let Ok(mut locked) = shared.lock() { + locked.put(svc); + } + } + Poll::Ready(()) + } + Err(_e) => Poll::Ready(()), + } + } + } +} + +mod events { + #[derive(Clone, Debug)] + #[non_exhaustive] + pub struct Ignore; + + #[derive(Clone, Debug)] + pub struct WithExecutor(pub(super) E); + + pub trait Events { + fn on_race_lost(&self, fut: CF); + } + + impl Events for Ignore { + fn on_race_lost(&self, _fut: CF) {} + } + + impl Events for WithExecutor + where + E: hyper::rt::Executor, + { + fn on_race_lost(&self, fut: CF) { + self.0.execute(fut); + } + } +} + +#[cfg(test)] +mod tests { + use futures_util::future; + use tower::Service; + use tower_test::assert_request_eq; + + #[tokio::test] + async fn test_makes_svc_when_empty() { + let (mock, mut handle) = tower_test::mock::pair(); + let mut cache = super::builder().build(mock); + handle.allow(1); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + + let f = cache.call(1); + + future::join(f, async move { + assert_request_eq!(handle, 1).send_response("one"); + }) + .await + .0 + .expect("call"); + } + + #[tokio::test] + async fn test_reuses_after_idle() { + let (mock, mut handle) = tower_test::mock::pair(); + let mut cache = super::builder().build(mock); + + // only 1 connection should ever be made + handle.allow(1); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let f = cache.call(1); + let cached = future::join(f, async { + assert_request_eq!(handle, 1).send_response("one"); + }) + .await + .0 + .expect("call"); + drop(cached); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let f = cache.call(1); + let cached = f.await.expect("call"); + drop(cached); + } +} From a7e7294539a531f26c50124bc6395469ee6b4744 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Mon, 27 Apr 2026 17:42:07 -0400 Subject: [PATCH 04/47] use pin project --- rust-runtime/Cargo.lock | 1 + .../aws-smithy-http-client/Cargo.toml | 2 + .../src/client/pool/connection.rs | 69 +++++++++---------- .../src/client/pool/vendored_cache.rs | 15 ++-- 4 files changed, 43 insertions(+), 44 deletions(-) diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock index 7aa7ec12069..06e4d56239a 100644 --- a/rust-runtime/Cargo.lock +++ b/rust-runtime/Cargo.lock @@ -536,6 +536,7 @@ dependencies = [ "tokio", "tokio-rustls 0.26.4", "tower 0.5.3", + "tower-service", "tower-test", "tracing", ] diff --git a/rust-runtime/aws-smithy-http-client/Cargo.toml b/rust-runtime/aws-smithy-http-client/Cargo.toml index c69529efb87..45748cb0890 100644 --- a/rust-runtime/aws-smithy-http-client/Cargo.toml +++ b/rust-runtime/aws-smithy-http-client/Cargo.toml @@ -27,6 +27,7 @@ default-client = [ "hyper-util?/client-proxy", "dep:http-1x", "dep:tower", + "dep:tower-service", "tower?/util", "dep:futures-util", "tokio/sync", @@ -100,6 +101,7 @@ s2n-tls-hyper = { version = "0.1.0", optional = true } s2n-tls = { version = "0.3.33", optional = true } s2n-tls-tokio = { version = "0.3.33", optional = true} tower = { version = "0.5.2", optional = true } +tower-service = { version = "0.3", optional = true } futures-util = { version = "0.3.32", default-features = false, features = ["std"], optional = true } rustls-pki-types = { version = "1.12.0", features = ["std"], optional = true } rustls-native-certs = { version = "0.8.1", optional = true } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index e79b8143e50..04dff823da0 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -12,6 +12,7 @@ use std::task::{Context, Poll}; use std::time::Instant; use aws_smithy_types::body::SdkBody; +use pin_project_lite::pin_project; use tower::Service; use super::cache; @@ -243,34 +244,37 @@ where } } -/// A response body that keeps the originating pool checkout alive until -/// the body is fully consumed. -/// -/// # Why this exists -/// -/// When a checked-out pool connection's `Service::call` returns, the -/// response head is available but the body may still be streaming over -/// the same underlying HTTP connection. For H1 this is critical — if the -/// `CachedConnection` drops at that point, the connection returns to the -/// pool mid-body-stream and the next checkout would be handed a still-busy -/// connection. -/// -/// `ResponseBody` holds the checkout (`CachedConnection` for H1, -/// `SingletonConnection` for H2) in its guard field until the body is -/// fully dropped. Body streaming continues through the held inner -/// `Incoming`; when the `ResponseBody` is dropped the guard drops, which -/// for H1 triggers `CachedConnection::Drop` (return-to-pool or `discard` -/// if poisoned), and for H2 simply drops the singleton clone (no-op on -/// singleton state). -/// -/// The H2 variant carries a generic type parameter because -/// `SingletonConnection`'s inner `T` is `hyper_util::client::pool:: -/// singleton::Singled<...>`, which is unnameable outside hyper-util. The -/// H1 variant is fully concrete since our vendored cache makes -/// `cache::Cached<...>` nameable. -pub(crate) struct ResponseBody { - inner: hyper::body::Incoming, - _guard: ConnectionGuard, +pin_project! { + /// A response body that keeps the originating pool checkout alive until + /// the body is fully consumed. + /// + /// # Why this exists + /// + /// When a checked-out pool connection's `Service::call` returns, the + /// response head is available but the body may still be streaming over + /// the same underlying HTTP connection. For H1 this is critical — if the + /// `CachedConnection` drops at that point, the connection returns to the + /// pool mid-body-stream and the next checkout would be handed a still-busy + /// connection. + /// + /// `ResponseBody` holds the checkout (`CachedConnection` for H1, + /// `SingletonConnection` for H2) in its guard field until the body is + /// fully dropped. Body streaming continues through the held inner + /// `Incoming`; when the `ResponseBody` is dropped the guard drops, which + /// for H1 triggers `CachedConnection::Drop` (return-to-pool or `discard` + /// if poisoned), and for H2 simply drops the singleton clone (no-op on + /// singleton state). + /// + /// The H2 variant carries a generic type parameter because + /// `SingletonConnection`'s inner `T` is `hyper_util::client::pool:: + /// singleton::Singled<...>`, which is unnameable outside hyper-util. The + /// H1 variant is fully concrete since our vendored cache makes + /// `cache::Cached<...>` nameable. + pub(crate) struct ResponseBody { + #[pin] + inner: hyper::body::Incoming, + _guard: ConnectionGuard, + } } /// What a `ResponseBody` holds alive while the body streams. @@ -291,10 +295,7 @@ impl ResponseBody { } } -impl hyper::body::Body for ResponseBody -where - H2Inner: Unpin, -{ +impl hyper::body::Body for ResponseBody { type Data = ::Data; type Error = ::Error; @@ -302,9 +303,7 @@ where self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll, Self::Error>>> { - // Safe: we only project to `inner`, guard is never moved. - let this = unsafe { self.get_unchecked_mut() }; - std::pin::Pin::new(&mut this.inner).poll_frame(cx) + self.project().inner.poll_frame(cx) } fn is_end_stream(&self) -> bool { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs index 0ecd6a56944..cf476ee9aa9 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs @@ -36,19 +36,16 @@ * Commit: e1c5a6c89bfaed11fb34bd483fe9ba616f403791 * * Modifications from upstream: - * 1. Replaced `tower_service::Service` imports with `tower::Service` (the - * `tower` crate re-exports `Service` at the crate root, avoiding an - * extra direct dependency on `tower-service`). - * 2. Changed `pub use self::internal::builder;` to `pub(crate) use ...` + * 1. Changed `pub use self::internal::builder;` to `pub(crate) use ...` * and dropped the three `#[cfg(docsrs)] pub use` lines. The composable * pool internals are `pub(crate)` in this crate — no types from this * file are exposed in the smithy-rs public API. - * 3. Added `#![allow(dead_code, unreachable_pub)]` so the file can be + * 2. Added `#![allow(dead_code, unreachable_pub)]` so the file can be * kept close to upstream even when individual items aren't used yet. - * 4. Dropped the module- and struct-level rustdoc sections that reference + * 3. Dropped the module- and struct-level rustdoc sections that reference * "Unnameable" (that rustdoc pattern is specific to hyper-util's public * API; these types aren't public here). - * 5. Added `Cached::discard(self)` — consumes self and prevents + * 4. Added `Cached::discard(self)` — consumes self and prevents * reinsertion into the pool on drop. See `// SDK MODIFICATION` marker * below. Used to drop connections that a caller has learned are bad * (poisoned, GOAWAY, etc.) between checkout and `poll_ready`. The @@ -76,7 +73,7 @@ mod internal { use futures_util::future; use tokio::sync::oneshot; - use tower::Service; + use tower_service::Service; use super::events; @@ -464,7 +461,7 @@ mod events { #[cfg(test)] mod tests { use futures_util::future; - use tower::Service; + use tower_service::Service; use tower_test::assert_request_eq; #[tokio::test] From 870409a0bac1ce3ad63c53bf29fbf50970e37383 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 28 Apr 2026 09:59:55 -0400 Subject: [PATCH 05/47] fix generics to be more useful --- .../aws-smithy-http-client/src/client/pool.rs | 167 +++++++++++++++--- .../src/client/pool/connection.rs | 31 ++-- 2 files changed, 161 insertions(+), 37 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 2f9468cff15..5578eb1f679 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -30,6 +30,7 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Mutex; +use std::task::Poll; use aws_smithy_types::body::SdkBody; use hyper_util::client::legacy::connect::Connection as HyperConnection; @@ -37,8 +38,11 @@ use hyper_util::client::pool as hpool; use hyper_util::rt::TokioExecutor; use tower::{Service, ServiceExt}; -use connection::ManagedConnection; -use handshake::{H1ConnectAndHandshake, H2ConnectAndHandshake}; +use connection::{ + CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, ManagedConnection, + SingletonConnection, +}; +use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake, H2SendRequest}; pub(crate) type BoxError = Box; type BoxFuture = Pin + Send>>; @@ -75,10 +79,11 @@ where cache::builder() .executor(TokioExecutor::new()) .build(H1ConnectAndHandshake::new(inspector)) - .map_response(CachedConnection::new) + .map_response(|cached| H1Checkout::new(CachedConnection::new(cached))) })) .upgrade(tower::layer::layer_fn(|inspected| { hpool::singleton::Singleton::new(H2ConnectAndHandshake::new(inspected)) + .map_response(|singled| H2Checkout::new(SingletonConnection::new(singled))) })) .build(); Box::new(TypedPoolEntry(stack)) @@ -109,7 +114,7 @@ impl ConnectionPool { &self, uri: http_1x::Uri, req: http_1x::Request, - ) -> Result, BoxError> { + ) -> Result, BoxError> { let key = PoolKey::from_uri(&uri).ok_or("request URI must have scheme and authority")?; // Clone the per-host stack out of the lock. The stack is Clone @@ -127,6 +132,114 @@ impl ConnectionPool { } } +/// Tower `Service` wrapper for an H1 pool checkout. +/// +/// Produced by the H1 fallback leg (one per checkout from the cache). +/// Its `Service::call(req)` consumes the held `CachedConnection`, runs +/// the request, and wraps the response body so the connection returns +/// to the pool only when the body is fully drained. +/// +/// # Single-use +/// +/// Each checkout represents one `CachedConnection`. `call` moves the +/// connection into the response future via `Option::take()`; calling +/// `call` twice panics (unreachable given our checkout-per-request flow). +/// +/// The `UnusedH2Phantom` generic is a phantom type parameter: the H1 path +/// never populates the H2 variant of `ConnectionGuard` and never holds an +/// H2 checkout, but both legs must produce the same `CheckoutResponse<…>` +/// so `Negotiate` can compose them uniformly. The phantom slot here is +/// whatever type the H2 leg's `SingletonConnection` holds, resolved by +/// type inference at the pool composition site — this wrapper ignores it. +pub(crate) struct H1Checkout { + conn: Option>, + _marker: std::marker::PhantomData UnusedH2Phantom>, +} + +impl H1Checkout { + pub(crate) fn new(conn: CachedConnection) -> Self { + Self { + conn: Some(conn), + _marker: std::marker::PhantomData, + } + } +} + +impl Service> for H1Checkout { + type Response = CheckoutResponse; + type Error = BoxError; + type Future = BoxFuture>; + + fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll> { + self.conn + .as_mut() + .expect("H1Checkout::poll_ready after call") + .poll_ready(cx) + } + + fn call(&mut self, req: http_1x::Request) -> Self::Future { + let mut conn = self.conn.take().expect("H1Checkout::call called twice"); + Box::pin(async move { + let resp = conn.call(req).await?; + let (parts, body) = resp.into_parts(); + let body = GuardedBody::new(body, ConnectionGuard::H1(conn)); + Ok(http_1x::Response::from_parts(parts, body)) + }) + } +} + +/// Tower `Service` wrapper for an H2 pool checkout. +/// +/// Symmetric to `H1Checkout`. The `Singleton` type parameter is the +/// `SingletonConnection` inner `T`, which at the composition site +/// resolves to the unnameable `hyper_util::client::pool::singleton:: +/// Singled>`. We carry it through +/// generics and never name it concretely. +/// +/// # Single-use +/// +/// Same contract as `H1Checkout`. +pub(crate) struct H2Checkout { + conn: Option>, +} + +impl H2Checkout { + pub(crate) fn new(conn: SingletonConnection) -> Self { + Self { conn: Some(conn) } + } +} + +impl Service> for H2Checkout +where + Singleton: Service, Response = http_1x::Response> + + Send + + 'static, + Singleton::Error: Into, + Singleton::Future: Send + 'static, +{ + type Response = CheckoutResponse; + type Error = BoxError; + type Future = BoxFuture>; + + fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll> { + self.conn + .as_mut() + .expect("H2Checkout::poll_ready after call") + .poll_ready(cx) + .map_err(Into::into) + } + + fn call(&mut self, req: http_1x::Request) -> Self::Future { + let mut conn = self.conn.take().expect("H2Checkout::call called twice"); + Box::pin(async move { + let resp = conn.call(req).await.map_err(Into::into)?; + let (parts, body) = resp.into_parts(); + let body = GuardedBody::new(body, ConnectionGuard::H2(conn)); + Ok(http_1x::Response::from_parts(parts, body)) + }) + } +} + /// Type-erased per-host pool entry. /// /// Each host has one of these, wrapping the unnameable Negotiate stack. @@ -140,23 +253,29 @@ trait Connection: Send { fn send( self: Box, req: http_1x::Request, - ) -> BoxFuture, BoxError>>; + ) -> BoxFuture, BoxError>>; } type BoxedConnection = Box; /// Concrete PoolEntry wrapping a Negotiate stack. +/// +/// `PoolUnnameable` propagates whatever pool-internal unnameable type +/// flows through the checkout services' `CheckoutResponse`; it's carried +/// through so `Conn::Response` can name `CheckoutResponse`. struct TypedPoolEntry(S); -impl PoolEntry for TypedPoolEntry +impl PoolEntry for TypedPoolEntry where S: Service + Clone + Send + Sync + 'static, S::Error: Into + 'static, S::Future: Send + 'static, - Conn: Service> + Send + 'static, - Conn::Response: Into>, + Conn: Service, Response = CheckoutResponse> + + Send + + 'static, Conn::Error: Into, Conn::Future: Send + 'static, + PoolUnnameable: Send + Sync + 'static, { fn acquire(&mut self, uri: http_1x::Uri) -> BoxFuture> { let mut svc = self.0.clone(); @@ -185,39 +304,35 @@ where } /// Concrete Connection wrapping a Negotiated service. +/// +/// Erases the `CheckoutResponse` into `SdkBody` at the +/// trait-object boundary, so `dyn Connection` can name a single concrete +/// response type. struct TypedConnection(S); -impl Connection for TypedConnection +impl Connection for TypedConnection where - S: Service> + Send + 'static, - S::Response: Into>, + S: Service, Response = CheckoutResponse> + + Send + + 'static, S::Error: Into, S::Future: Send + 'static, + PoolUnnameable: Send + Sync + 'static, { fn send( self: Box, req: http_1x::Request, - ) -> BoxFuture, BoxError>> { + ) -> BoxFuture, BoxError>> { let mut svc = self.0; Box::pin(async move { - std::future::poll_fn(|cx| svc.poll_ready(cx)) - .await - .map_err(Into::into)?; - svc.call(req).await.map(Into::into).map_err(Into::into) + let resp = svc.call(req).await.map_err(Into::into)?; + let (parts, body) = resp.into_parts(); + let body = SdkBody::from_body_1_x(body); + Ok(http_1x::Response::from_parts(parts, body)) }) } } -/// A connection checked out from the H1 cache. -/// -/// See `connection::CachedConnection`. -pub(crate) use connection::CachedConnection; - -/// A connection checked out from the H2 singleton. -/// -/// See `connection::SingletonConnection`. -pub(crate) use connection::SingletonConnection; - #[cfg(test)] mod tests { use super::*; diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 04dff823da0..6f8df3875a1 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -257,10 +257,10 @@ pin_project! { /// pool mid-body-stream and the next checkout would be handed a still-busy /// connection. /// - /// `ResponseBody` holds the checkout (`CachedConnection` for H1, + /// `GuardedBody` holds the checkout (`CachedConnection` for H1, /// `SingletonConnection` for H2) in its guard field until the body is /// fully dropped. Body streaming continues through the held inner - /// `Incoming`; when the `ResponseBody` is dropped the guard drops, which + /// `Incoming`; when the `GuardedBody` is dropped the guard drops, which /// for H1 triggers `CachedConnection::Drop` (return-to-pool or `discard` /// if poisoned), and for H2 simply drops the singleton clone (no-op on /// singleton state). @@ -270,24 +270,24 @@ pin_project! { /// singleton::Singled<...>`, which is unnameable outside hyper-util. The /// H1 variant is fully concrete since our vendored cache makes /// `cache::Cached<...>` nameable. - pub(crate) struct ResponseBody { + pub(crate) struct GuardedBody { #[pin] inner: hyper::body::Incoming, - _guard: ConnectionGuard, + _guard: ConnectionGuard, } } -/// What a `ResponseBody` holds alive while the body streams. +/// What a `GuardedBody` holds alive while the body streams. /// /// Explicit per-leg variants so H1 vs H2 bifurcation is visible at the /// type level and in debugger output. -pub(crate) enum ConnectionGuard { - H1(CachedConnection>), - H2(SingletonConnection), +pub(crate) enum ConnectionGuard { + H1(CachedConnection), + H2(SingletonConnection), } -impl ResponseBody { - pub(crate) fn new(inner: hyper::body::Incoming, guard: ConnectionGuard) -> Self { +impl GuardedBody { + pub(crate) fn new(inner: hyper::body::Incoming, guard: ConnectionGuard) -> Self { Self { inner, _guard: guard, @@ -295,7 +295,7 @@ impl ResponseBody { } } -impl hyper::body::Body for ResponseBody { +impl hyper::body::Body for GuardedBody { type Data = ::Data; type Error = ::Error; @@ -314,3 +314,12 @@ impl hyper::body::Body for ResponseBody { self.inner.size_hint() } } + +/// The response type both H1 and H2 pool checkouts produce. +/// +/// Carries `GuardedBody` because some pool-internal leg +/// holds a checkout guard whose type is unnameable outside hyper-util +/// (today: the H2 leg's `Singled<...>`). `Negotiate` requires uniform +/// response types across legs; this alias is the uniform type that +/// consumers above the Negotiate composition point work with. +pub(crate) type CheckoutResponse = http_1x::Response>; From 56e2ae4c01e4346c4321ac6d62568a8604560144 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 28 Apr 2026 11:02:23 -0400 Subject: [PATCH 06/47] collapse acquire/send --- .../aws-smithy-http-client/src/client/pool.rs | 123 +++++++++--------- 1 file changed, 62 insertions(+), 61 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 5578eb1f679..c9d52bd0633 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -108,8 +108,7 @@ pub(crate) struct ConnectionPool { impl ConnectionPool { /// Send a request through the pool. /// - /// Routes to the appropriate per-host pool stack, acquires a connection, - /// and sends the request. + /// Routes to the appropriate per-host pool stack and sends the request. pub(crate) async fn send_request( &self, uri: http_1x::Uri, @@ -117,18 +116,19 @@ impl ConnectionPool { ) -> Result, BoxError> { let key = PoolKey::from_uri(&uri).ok_or("request URI must have scheme and authority")?; - // Clone the per-host stack out of the lock. The stack is Clone - // (cheap Arc clones internally). All I/O happens outside the lock. + // Dispatch the request through the per-host entry. The lock is + // held only long enough to look up / create the entry and call + // `send` — all I/O happens in the returned future, outside the + // lock. let fut = { let mut hosts = self.hosts.lock().unwrap(); if !hosts.contains_key(&key) { hosts.insert(key.clone(), (self.make_entry)(&uri)); } - hosts.get_mut(&key).unwrap().acquire(uri) + hosts.get_mut(&key).unwrap().send(uri, req) }; - let conn = fut.await?; - conn.send(req).await + fut.await } } @@ -243,21 +243,31 @@ where /// Type-erased per-host pool entry. /// /// Each host has one of these, wrapping the unnameable Negotiate stack. +/// Dyn erasure is structural: hosts have different Negotiate compositions +/// with different unnameable inner types, so `ConnectionPool::hosts` can't +/// hold a concrete type. +/// +/// The entry is request-granular — `send` performs checkout + health +/// filter + dispatch + body-guard setup atomically for one request. This +/// folds what was previously a separate `acquire(uri) -> Connection` +/// followed by `connection.send(req)` into one call. +/// +/// # Future: splitting acquire and send +/// +/// If we later need to expose a checked-out connection handle to callers +/// (e.g. connection warming, request pipelining on a single checkout, +/// or pool-external reuse), reintroduce a separate `Connection` trait +/// and split this into `acquire(uri) -> Box` plus +/// `Connection::send(self: Box, req) -> Response`. The current +/// shape collapses that split because no caller needs it today. trait PoolEntry: Send + Sync { - /// Acquire a connection from this host's pool. - fn acquire(&mut self, uri: http_1x::Uri) -> BoxFuture>; -} - -/// A type-erased connection that can send one request. -trait Connection: Send { fn send( - self: Box, + &mut self, + uri: http_1x::Uri, req: http_1x::Request, ) -> BoxFuture, BoxError>>; } -type BoxedConnection = Box; - /// Concrete PoolEntry wrapping a Negotiate stack. /// /// `PoolUnnameable` propagates whatever pool-internal unnameable type @@ -277,58 +287,49 @@ where Conn::Future: Send + 'static, PoolUnnameable: Send + Sync + 'static, { - fn acquire(&mut self, uri: http_1x::Uri) -> BoxFuture> { + fn send( + &mut self, + uri: http_1x::Uri, + req: http_1x::Request, + ) -> BoxFuture, BoxError>> { let mut svc = self.0.clone(); Box::pin(async move { - // Loop past stale idle entries (poisoned or dead). Converges because: - // - Cache idle set is bounded; each post-checkout `poll_ready` Err - // flips `is_closed` so `Cached::Drop` skips reinsertion, shrinking - // the set by one. - // - Singleton clears to `Empty` on `poll_ready` Err, forcing the - // next `call` to run a fresh handshake — a handshake failure - // surfaces as an error from `svc.call` (not the inner - // `poll_ready`), which we propagate. - loop { + // Checkout loop. Pops past stale idle entries (poisoned or dead). + // Converges because: + // - Cache idle set is bounded; each post-checkout `poll_ready` + // Err flips `is_closed` so `Cached::Drop` skips reinsertion, + // shrinking the set by one. + // - Singleton clears to `Empty` on `poll_ready` Err, forcing + // the next `call` to run a fresh handshake — a handshake + // failure surfaces as an error from `svc.call` (not the + // inner `poll_ready`), which we propagate. + let mut conn = loop { std::future::poll_fn(|cx| svc.poll_ready(cx)) .await .map_err(Into::into)?; - let mut conn = svc.call(uri.clone()).await.map_err(Into::into)?; - if std::future::poll_fn(|cx| conn.poll_ready(cx)).await.is_ok() { - return Ok(Box::new(TypedConnection(conn)) as BoxedConnection); + let mut checkout = svc.call(uri.clone()).await.map_err(Into::into)?; + if std::future::poll_fn(|cx| checkout.poll_ready(cx)) + .await + .is_ok() + { + break checkout; } - // drop `conn` → triggers pool cleanup (H1: discard via - // CachedConnection::Drop; H2: Singleton already cleared state). - } - }) - } -} - -/// Concrete Connection wrapping a Negotiated service. -/// -/// Erases the `CheckoutResponse` into `SdkBody` at the -/// trait-object boundary, so `dyn Connection` can name a single concrete -/// response type. -struct TypedConnection(S); - -impl Connection for TypedConnection -where - S: Service, Response = CheckoutResponse> - + Send - + 'static, - S::Error: Into, - S::Future: Send + 'static, - PoolUnnameable: Send + Sync + 'static, -{ - fn send( - self: Box, - req: http_1x::Request, - ) -> BoxFuture, BoxError>> { - let mut svc = self.0; - Box::pin(async move { - let resp = svc.call(req).await.map_err(Into::into)?; + // drop `checkout` → triggers pool cleanup (H1: discard via + // CachedConnection::Drop; H2: Singleton already cleared + // state). + }; + + // Dispatch the request; convert the body into `SdkBody` at the + // boundary so consumers (and `dyn PoolEntry`) don't see our + // internal `GuardedBody<...>` type. The guard inside + // `GuardedBody` lives on inside `SdkBody` until the body is + // fully drained. + let resp = conn.call(req).await.map_err(Into::into)?; let (parts, body) = resp.into_parts(); - let body = SdkBody::from_body_1_x(body); - Ok(http_1x::Response::from_parts(parts, body)) + Ok(http_1x::Response::from_parts( + parts, + SdkBody::from_body_1_x(body), + )) }) } } From a29c2693d78000e993669b89eb7c7d8548cc7204 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 28 Apr 2026 15:20:37 -0400 Subject: [PATCH 07/47] tighten bounds on generics --- .../aws-smithy-http-client/src/client/pool.rs | 35 +++++++++++++++++-- .../src/client/pool/handshake.rs | 14 ++++---- 2 files changed, 41 insertions(+), 8 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index c9d52bd0633..3575eb3ebc6 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -47,6 +47,28 @@ use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake, H2S pub(crate) type BoxError = Box; type BoxFuture = Pin + Send>>; +/// Pool-level configuration. +/// +/// Plumbed from the eventual `Builder::new_v2()` public API through +/// `build_pool` to `ConnectionPool`. Internal type — the public surface +/// is the builder's per-knob setters, not this struct. +/// +/// All fields are `None` by default (defaults applied by the pool where +/// they take effect, not here). Fields are `pub(crate)` because +/// consumers of this struct are all within the pool module; accessor +/// methods would add noise without benefit. +#[derive(Clone, Debug, Default)] +pub(crate) struct PoolConfig { + /// Upper bound on concurrent connections (total, across all hosts). + /// Enforced via `tower::ConcurrencyLimit` at the connection + /// establishment layer. `None` = unlimited. + pub(crate) max_connections: Option, + + /// How long an idle connection may stay in the pool before being + /// evicted. `None` = no eviction (hyper-util's default behavior). + pub(crate) pool_idle_timeout: Option, +} + /// Key for per-host connection pool routing. #[derive(Clone, Debug, PartialEq, Eq, Hash)] struct PoolKey { @@ -63,8 +85,8 @@ impl PoolKey { } } -/// Build a connection pool for the given connector. -pub(crate) fn build_pool(connector: C) -> ConnectionPool +/// Build a connection pool for the given connector and configuration. +pub(crate) fn build_pool(connector: C, config: PoolConfig) -> ConnectionPool where C: Service + Clone + Send + Sync + 'static, C::Error: Into + 'static, @@ -72,6 +94,10 @@ where IO: hyper::rt::Read + hyper::rt::Write + HyperConnection + Unpin + Send + 'static, { let make_entry = move |_uri: &http_1x::Uri| -> Box { + // TODO #8: consult `config.max_connections` here (ConcurrencyLimit + // layer between cache and handshake) and `config.pool_idle_timeout` + // (retain() scheduler in Phase 4). For now the config is stored + // on the pool but not yet applied. let stack = hpool::negotiate::builder() .connect(connector.clone()) .inspect(|conn: &IO| conn.connected().is_negotiated_h2()) @@ -90,6 +116,7 @@ where }; ConnectionPool { + config, hosts: Mutex::new(HashMap::new()), make_entry: Box::new(make_entry), } @@ -101,6 +128,10 @@ where /// Each host gets a Negotiate stack that selects between HTTP/1.1 (Cache) /// and HTTP/2 (Singleton) based on ALPN negotiation. pub(crate) struct ConnectionPool { + /// Pool-wide configuration. Stored for consultation at request time + /// and for re-use if `make_entry` needs config (not yet wired). + #[allow(dead_code)] + config: PoolConfig, hosts: Mutex>>, make_entry: Box Box + Send + Sync>, } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 08ee3bc65db..e56bb5bfc03 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -167,8 +167,10 @@ where /// Wraps a connector service, performing the H2 protocol handshake after /// connection and returning a `ManagedConnection`. /// -/// Generic over the request type because the Negotiate upgrade path -/// receives `()` (connection already established), not a URI. +/// Pinned to `Service<()>` because this service sits in the Negotiate +/// upgrade path — by the time it runs, the connection is already +/// established; no target (URI) is needed. Keeping the request type +/// `()` prevents misassembly at compile time. pub(crate) struct H2ConnectAndHandshake { connector: C, } @@ -187,9 +189,9 @@ impl Clone for H2ConnectAndHandshake { } } -impl Service for H2ConnectAndHandshake +impl Service<()> for H2ConnectAndHandshake where - C: Service, + C: Service<(), Response = IO>, C::Error: Into + 'static, C::Future: Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, @@ -202,8 +204,8 @@ where self.connector.poll_ready(cx).map_err(Into::into) } - fn call(&mut self, req: Req) -> Self::Future { - let fut = self.connector.call(req); + fn call(&mut self, _req: ()) -> Self::Future { + let fut = self.connector.call(()); Box::pin(async move { let io = fut.await.map_err(Into::into)?; let info = capture_info(&io); From 201d81636df19b28a86fa0379887abc63598a2a1 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 28 Apr 2026 16:15:16 -0400 Subject: [PATCH 08/47] remove unused and cleanup logs --- .../aws-smithy-http-client/src/client/pool.rs | 9 +++--- .../src/client/pool/connection.rs | 27 +++++++--------- .../src/client/pool/handshake.rs | 32 +++++++++++++++---- 3 files changed, 42 insertions(+), 26 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 3575eb3ebc6..db536c69642 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -94,10 +94,11 @@ where IO: hyper::rt::Read + hyper::rt::Write + HyperConnection + Unpin + Send + 'static, { let make_entry = move |_uri: &http_1x::Uri| -> Box { - // TODO #8: consult `config.max_connections` here (ConcurrencyLimit - // layer between cache and handshake) and `config.pool_idle_timeout` - // (retain() scheduler in Phase 4). For now the config is stored - // on the pool but not yet applied. + // TODO: consult `config.max_connections` here by wrapping the + // per-host connector with `tower::ConcurrencyLimit`, and consult + // `config.pool_idle_timeout` by driving `retain()` on a background + // task. Currently the config is stored on the pool but not yet + // applied. let stack = hpool::negotiate::builder() .connect(connector.clone()) .inspect(|conn: &IO| conn.connected().is_negotiated_h2()) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 6f8df3875a1..4a3d8a3de3a 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -53,23 +53,30 @@ pub(crate) struct ConnectionInfo { /// /// Clone is supported when the inner service is Clone (e.g., HTTP/2 multiplexed /// connections). Clones share the same poison flag and connection metadata. +/// +/// # Missing: idle-start tracking +/// +/// A `last_used: Instant` updated on **return** to the pool is required +/// to drive `pool_idle_timeout`-based eviction. That requires a +/// re-insertion hook we do not have yet (hyper-util's `Cached::Drop` +/// gives us no callback). The hook shape will be designed alongside the +/// eviction mechanism. Until then this struct does not track idle +/// duration. `created_at` is retained because it is well-defined at +/// establishment time and needed for a future max-lifetime eviction. pub(crate) struct ManagedConnection { inner: S, info: ConnectionInfo, created_at: Instant, - last_used: Instant, poisoned: Arc, } impl ManagedConnection { /// Create a new managed connection wrapping the given service. pub(crate) fn new(inner: S, info: ConnectionInfo) -> Self { - let now = Instant::now(); Self { inner, info, - created_at: now, - last_used: now, + created_at: Instant::now(), poisoned: Arc::new(AtomicBool::new(false)), } } @@ -95,11 +102,6 @@ impl ManagedConnection { self.poisoned.load(Ordering::Acquire) } - /// How long this connection has been idle (since last request completed). - pub(crate) fn idle_duration(&self) -> std::time::Duration { - self.last_used.elapsed() - } - /// Connection metadata. pub(crate) fn info(&self) -> &ConnectionInfo { &self.info @@ -110,11 +112,6 @@ impl ManagedConnection { self.created_at } - /// Mark this connection as recently used. - pub(crate) fn touch(&mut self) { - self.last_used = Instant::now(); - } - /// Mutable access to the inner service. pub(crate) fn inner_mut(&mut self) -> &mut S { &mut self.inner @@ -127,7 +124,6 @@ impl Clone for ManagedConnection { inner: self.inner.clone(), info: self.info.clone(), created_at: self.created_at, - last_used: self.last_used, poisoned: self.poisoned.clone(), } } @@ -154,7 +150,6 @@ where } fn call(&mut self, req: http_1x::Request) -> Self::Future { - self.touch(); let fut = self.inner_mut().call(req); Box::pin(async move { fut.await.map_err(Into::into) }) } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index e56bb5bfc03..6561e457079 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -151,9 +151,19 @@ where .map_err(|e| Box::new(e) as BoxError)?; // Drive the connection I/O in the background - tokio::spawn(async move { - if let Err(e) = conn.with_upgrades().await { - tracing::debug!(error = %e, "h1 connection driver error"); + tokio::spawn({ + let remote_addr = info.remote_addr; + let local_addr = info.local_addr; + async move { + if let Err(e) = conn.with_upgrades().await { + tracing::debug!( + protocol = "h1", + ?remote_addr, + ?local_addr, + error = %e, + "connection driver error" + ); + } } }); @@ -215,9 +225,19 @@ where .await .map_err(|e| Box::new(e) as BoxError)?; - tokio::spawn(async move { - if let Err(e) = conn.await { - tracing::debug!(error = %e, "h2 connection driver error"); + tokio::spawn({ + let remote_addr = info.remote_addr; + let local_addr = info.local_addr; + async move { + if let Err(e) = conn.await { + tracing::debug!( + protocol = "h2", + ?remote_addr, + ?local_addr, + error = %e, + "connection driver error" + ); + } } }); From a4339f50c920ec4f7d8f86936833cc9f8d06ae80 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 29 Apr 2026 12:50:58 -0400 Subject: [PATCH 09/47] move v2 into own module --- .../aws-smithy-http-client/src/client.rs | 20 +- .../aws-smithy-http-client/src/client/v2.rs | 241 ++++++++++++++++++ .../aws-smithy-http-client/src/lib.rs | 4 + .../src/test_util/wire/connection.rs | 75 +++++- .../tests/pool_behavior_test.rs | 112 ++++++-- 5 files changed, 415 insertions(+), 37 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/src/client/v2.rs diff --git a/rust-runtime/aws-smithy-http-client/src/client.rs b/rust-runtime/aws-smithy-http-client/src/client.rs index b5ecdb679c5..f756d11180c 100644 --- a/rust-runtime/aws-smithy-http-client/src/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client.rs @@ -12,6 +12,8 @@ pub mod tls; pub(crate) mod connect; pub(crate) mod pool; +/// V2 HTTP client backed by composable connection pools. +pub mod v2; use crate::cfg::cfg_tls; use crate::tls::TlsContext; @@ -147,10 +149,8 @@ pub struct TlsUnset {} /// TLS implementation selected #[derive(Debug, Clone)] pub struct TlsProviderSelected { - #[allow(unused)] - provider: tls::Provider, - #[allow(unused)] - context: TlsContext, + pub(crate) provider: tls::Provider, + pub(crate) context: TlsContext, } impl ConnectorBuilder { @@ -584,7 +584,7 @@ where } /// Downcast errors coming out of hyper into an appropriate `ConnectorError` -fn downcast_error(err: BoxError) -> ConnectorError { +pub(crate) fn downcast_error(err: BoxError) -> ConnectorError { // is a `TimedOutError` (from aws_smithy_async::timeout) in the chain? if it is, this is a timeout if find_source::(err.as_ref()).is_some() { return ConnectorError::timeout(err); @@ -944,6 +944,16 @@ impl Builder { Self::default() } + /// Creates a new builder backed by the v2 composable connection pool. + /// + /// The v2 client uses hyper-util's composable pool APIs for connection + /// management, supporting per-host connection reuse with HTTP/1.1 and + /// HTTP/2 protocol negotiation. + #[doc(hidden)] + pub fn new_v2() -> v2::BuilderV2 { + v2::BuilderV2::new() + } + /// Returns a [`SharedHttpClient`] that calls the given `connector` function to select an HTTP(S) connector. #[doc(hidden)] pub fn build_with_connector_fn(self, connector_fn: F) -> SharedHttpClient diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs new file mode 100644 index 00000000000..89631ac10f1 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -0,0 +1,241 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! V2 HTTP client. +//! +//! Entry point: [`BuilderV2::new`] or [`Builder::new_v2`](super::Builder::new_v2). + +use std::borrow::Cow; +use std::sync::Arc; +use std::time::Duration; + +use aws_smithy_runtime_api::client::connection::CaptureSmithyConnection; +use aws_smithy_runtime_api::client::connection::ConnectionMetadata; +use aws_smithy_runtime_api::client::connector_metadata::ConnectorMetadata; +use aws_smithy_runtime_api::client::http::{ + HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpClient, + SharedHttpConnector, +}; +use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse}; +use aws_smithy_runtime_api::client::result::ConnectorError; +use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; +use aws_smithy_runtime_api::shared::IntoShared; +use hyper_util::client::legacy::connect::HttpConnector as HyperHttpConnector; + +use super::downcast_error; +use super::pool::{self, ConnectionPool, PoolConfig}; +use super::tls; +use super::{TlsProviderSelected, TlsUnset}; +use crate::tls::TlsContext; + +/// Builder for the v2 HTTP client. +/// +/// Provides a flat, unified configuration surface for connection management, +/// TLS, timeouts, and pool behavior — all knobs in one place. +/// +/// # Examples +/// +/// ```no_run +/// use aws_smithy_http_client::v2::BuilderV2; +/// use std::time::Duration; +/// +/// let client = BuilderV2::new() +/// .pool_idle_timeout(Duration::from_secs(20)) +/// .build_http(); +/// ``` +#[derive(Clone, Debug)] +pub struct BuilderV2 { + pool_idle_timeout: Option, + tcp_nodelay: bool, + tls: Tls, +} + +impl Default for BuilderV2 { + fn default() -> Self { + Self { + pool_idle_timeout: None, + tcp_nodelay: true, + tls: TlsUnset {}, + } + } +} + +// Methods available in any TLS state. +impl BuilderV2 { + /// Set the pool idle timeout. + /// + /// Connections idle longer than this duration are evicted from the pool. + /// Set below the server's idle timeout to avoid stale connection errors + /// (e.g., ~20s for S3). + pub fn pool_idle_timeout(mut self, timeout: Duration) -> Self { + self.pool_idle_timeout = Some(timeout); + self + } + + /// Set TCP_NODELAY on connections. Default: `true`. + pub fn tcp_nodelay(mut self, nodelay: bool) -> Self { + self.tcp_nodelay = nodelay; + self + } +} + +impl BuilderV2 { + /// Create a new v2 client builder. + pub fn new() -> Self { + Self::default() + } + + /// Set the TLS implementation. + pub fn tls_provider(self, provider: tls::Provider) -> BuilderV2 { + BuilderV2 { + pool_idle_timeout: self.pool_idle_timeout, + tcp_nodelay: self.tcp_nodelay, + tls: TlsProviderSelected { + provider, + context: TlsContext::default(), + }, + } + } + + /// Build an HTTP client without TLS. + #[doc(hidden)] + pub fn build_http(self) -> SharedHttpClient { + let mut tcp = HyperHttpConnector::new(); + tcp.set_nodelay(self.tcp_nodelay); + let config = PoolConfig { + max_connections: None, + pool_idle_timeout: self.pool_idle_timeout, + }; + let pool = pool::build_pool(tcp, config); + V2HttpClient::new(pool).into_shared() + } +} + +// TODO(3f): impl BuilderV2 { fn build_https() } + +// --------------------------------------------------------------------------- +// HttpConnector adapter +// --------------------------------------------------------------------------- + +/// Smithy `HttpConnector` backed by the v2 connection pool. +struct PooledConnector { + pool: Arc, +} + +impl std::fmt::Debug for PooledConnector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PooledConnector").finish() + } +} + +impl HttpConnector for PooledConnector { + fn call(&self, request: HttpRequest) -> HttpConnectorFuture { + let pool = self.pool.clone(); + HttpConnectorFuture::new(async move { + let mut request = request + .try_into_http1x() + .map_err(|err| ConnectorError::user(err.into()))?; + + // Save the full URI for pool routing before rewriting to origin-form. + let full_uri = request.uri().clone(); + + // Wire CaptureSmithyConnection with a no-op poison_fn so the interceptor + // pipeline doesn't panic. TODO: wire real poison from ManagedConnection. + if let Some(capture) = request.extensions().get::() { + capture.set_connection_retriever(|| { + Some( + ConnectionMetadata::builder() + .proxied(false) + .poison_fn(|| {}) + .build(), + ) + }); + } + + // Set Host header from the URI authority if not already present. + if !request.headers().contains_key(http_1x::header::HOST) { + if let Some(authority) = full_uri.authority() { + let host = match get_non_default_port(&full_uri) { + Some(port) => format!("{}:{}", authority.host(), port), + None => authority.host().to_string(), + }; + request.headers_mut().insert( + http_1x::header::HOST, + http_1x::HeaderValue::from_str(&host) + .expect("authority is valid header value"), + ); + } + } + + // Rewrite URI to origin-form for HTTP/1.1. + let origin = match request.uri().path_and_query() { + Some(pq) => { + let mut parts = http_1x::uri::Parts::default(); + parts.path_and_query = Some(pq.clone()); + http_1x::Uri::from_parts(parts).expect("path is valid uri") + } + None => http_1x::Uri::from_static("/"), + }; + *request.uri_mut() = origin; + + let response = pool + .send_request(full_uri, request) + .await + .map_err(downcast_error)?; + + HttpResponse::try_from(response).map_err(|err| ConnectorError::other(err.into(), None)) + }) + } +} + +// --------------------------------------------------------------------------- +// HttpClient impl +// --------------------------------------------------------------------------- + +/// Returns the port if it is not the default for the scheme. +fn get_non_default_port(uri: &http_1x::Uri) -> Option> { + match (uri.port().map(|p| p.as_u16()), uri.scheme()) { + (Some(443), Some(s)) if *s == http_1x::uri::Scheme::HTTPS => None, + (Some(80), Some(s)) if *s == http_1x::uri::Scheme::HTTP => None, + _ => uri.port(), + } +} + +/// Smithy `HttpClient` backed by the v2 connection pool. +struct V2HttpClient { + pool: Arc, +} + +impl std::fmt::Debug for V2HttpClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("V2HttpClient").finish() + } +} + +impl V2HttpClient { + fn new(pool: ConnectionPool) -> Self { + Self { + pool: Arc::new(pool), + } + } +} + +impl HttpClient for V2HttpClient { + fn http_connector( + &self, + _settings: &HttpConnectorSettings, + _components: &RuntimeComponents, + ) -> SharedHttpConnector { + SharedHttpConnector::new(PooledConnector { + pool: self.pool.clone(), + }) + } + + fn connector_metadata(&self) -> Option { + // TODO: consider distinguishing v1 vs v2 clients and including TLS + // provider info (e.g. "rustls"/"0.23"). Revisit alongside 3f (TLS support). + Some(ConnectorMetadata::new("hyper", Some(Cow::Borrowed("1.x")))) + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/lib.rs b/rust-runtime/aws-smithy-http-client/src/lib.rs index 85cae688018..12eb5193950 100644 --- a/rust-runtime/aws-smithy-http-client/src/lib.rs +++ b/rust-runtime/aws-smithy-http-client/src/lib.rs @@ -44,6 +44,10 @@ pub(crate) mod client; #[cfg(feature = "default-client")] pub use client::{default_connector, proxy, tls, Builder, Connector, ConnectorBuilder}; +/// V2 HTTP client backed by composable connection pools. +#[cfg(feature = "default-client")] +pub use client::v2; + #[cfg(feature = "test-util")] pub mod test_util; diff --git a/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs b/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs index 0045fe76802..b6753661e25 100644 --- a/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs @@ -40,6 +40,13 @@ pub enum ConnectionEvent { /// The address of the endpoint that accepted the connection. endpoint_addr: SocketAddr, }, + /// HTTP request received by an endpoint. + HttpRequestReceived { + /// The request-target from the request line (e.g. "/path?query"). + uri: String, + /// The Host header value, if present. + host: Option, + }, /// DNS lookup performed. DnsLookup { /// The hostname that was looked up. @@ -83,7 +90,12 @@ impl TestEndpoint { .lock() .expect("behavior lock poisoned") .pop_front(); - tokio::spawn(handle_connection(stream, behavior, behaviors.clone())); + tokio::spawn(handle_connection( + stream, + behavior, + behaviors.clone(), + events.clone(), + )); } }); @@ -110,6 +122,7 @@ async fn handle_connection( mut stream: tokio::net::TcpStream, first_behavior: Option, remaining: Arc>>, + events: Arc>>, ) { let Some(behavior) = first_behavior else { drop(stream); @@ -129,13 +142,23 @@ async fn handle_connection( drop(stream); } ConnectionBehavior::Respond { status, body } => { + // Read the first request before responding. + match read_request(&mut stream).await { + Ok(req_info) => { + events.lock().expect("event lock poisoned").push(req_info); + } + Err(_) => return, + } if write_response(&mut stream, status, body).await.is_err() { return; } // Keep-alive loop: read next request, send next behavior's response loop { - if read_request(&mut stream).await.is_err() { - return; + match read_request(&mut stream).await { + Ok(req_info) => { + events.lock().expect("event lock poisoned").push(req_info); + } + Err(_) => return, } let next = remaining .lock() @@ -159,16 +182,19 @@ async fn write_response( status: u16, body: &[u8], ) -> Result<(), std::io::Error> { - let header = format!( + let mut response = format!( "HTTP/1.1 {status} OK\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n", body.len() - ); - stream.write_all(header.as_bytes()).await?; - stream.write_all(body).await?; + ) + .into_bytes(); + response.extend_from_slice(body); + stream.write_all(&response).await?; stream.flush().await } -async fn read_request(stream: &mut tokio::net::TcpStream) -> Result<(), std::io::Error> { +async fn read_request( + stream: &mut tokio::net::TcpStream, +) -> Result { // Read until we see \r\n\r\n (end of HTTP headers) let mut buf = vec![0u8; 4096]; let mut total = 0; @@ -182,9 +208,27 @@ async fn read_request(stream: &mut tokio::net::TcpStream) -> Result<(), std::io: } total += n; if total >= 4 && buf[..total].windows(4).any(|w| w == b"\r\n\r\n") { - return Ok(()); + break; } } + let raw = String::from_utf8_lossy(&buf[..total]); + let mut lines = raw.lines(); + // Parse request-target from "GET /path HTTP/1.1" + let uri = lines + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or("/") + .to_string(); + // Find Host header + let host = lines.find_map(|line| { + let lower = line.to_ascii_lowercase(); + if lower.starts_with("host:") { + Some(line[5..].trim().to_string()) + } else { + None + } + }); + Ok(ConnectionEvent::HttpRequestReceived { uri, host }) } /// Mock DNS resolver that returns configured IPs and logs lookups. @@ -268,6 +312,19 @@ impl ConnectionTestHarness { .count() } + /// Collected HTTP request events (uri + host header). + pub fn http_requests(&self) -> Vec<(String, Option)> { + self.events() + .iter() + .filter_map(|e| match e { + ConnectionEvent::HttpRequestReceived { uri, host } => { + Some((uri.clone(), host.clone())) + } + _ => None, + }) + .collect() + } + /// Get a clone of the mock DNS resolver. pub fn dns_resolver(&self) -> MockDnsResolver { self.dns_resolver.clone() diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 1ff181fe8e7..e5755e2fe0d 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -16,6 +16,7 @@ use aws_smithy_http_client::test_util::wire::connection::{ ConnectionBehavior, ConnectionTestHarness, }; use aws_smithy_http_client::test_util::wire::{ReplayedEvent, WireMockServer}; +use aws_smithy_http_client::v2::BuilderV2; use aws_smithy_http_client::{ev, match_events, Builder, Connector}; use aws_smithy_runtime_api::client::http::{ HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, @@ -71,20 +72,19 @@ impl MakeClient for V1Client { } } -// When Builder::new_v2() exists: -// struct V2Client; -// impl MakeClient for V2Client { -// fn make(&self, config: ClientConfig) -> SharedHttpClient { -// let mut builder = Builder::new_v2(); -// if let Some(timeout) = config.idle_timeout { -// builder = builder.pool_idle_timeout(timeout); -// } -// if let Some(n) = config.max_connections { -// builder = builder.max_connections(n); -// } -// builder.build_http() -// } -// } +// V2 client backed by the composable connection pool +struct V2Client; + +impl MakeClient for V2Client { + fn make(&self, config: ClientConfig) -> SharedHttpClient { + let mut builder = BuilderV2::new(); + if let Some(timeout) = config.idle_timeout { + builder = builder.pool_idle_timeout(timeout); + } + // max_connections deferred + builder.build_http() + } +} // --------------------------------------------------------------------------- // Helpers @@ -226,12 +226,78 @@ async fn v1_connector_metadata() { connector_metadata(&V1Client).await; } -// When Builder::new_v2() exists, add: -// #[tokio::test] -// async fn v2_connection_reuse() { connection_reuse(&V2Client).await; } -// #[tokio::test] -// async fn v2_idle_timeout_eviction() { idle_timeout_eviction(&V2Client).await; } -// #[tokio::test] -// async fn v2_connection_reset_returns_error() { connection_reset_returns_error(&V2Client).await; } -// #[tokio::test] -// async fn v2_connector_metadata() { connector_metadata(&V2Client).await; } +// --------------------------------------------------------------------------- +// v2 test runners +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn v2_connection_reuse() { + connection_reuse(&V2Client).await; +} + +#[tokio::test] +async fn v2_idle_timeout_eviction() { + idle_timeout_eviction(&V2Client).await; +} + +#[tokio::test] +async fn v2_connection_reset_returns_error() { + connection_reset_returns_error(&V2Client).await; +} + +#[tokio::test] +async fn v2_connector_metadata() { + connector_metadata(&V2Client).await; +} + +// --------------------------------------------------------------------------- +// Test implementations: origin-form URI + Host header +// --------------------------------------------------------------------------- + +/// Requests should be sent with origin-form URI (just the path) and a correct Host header. +async fn origin_form_and_host_header(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }], + ) + .build() + .await; + + let client = make.make(ClientConfig::default()); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/some/path?key=val"); + + let resp = send_to(&client, &url) + .await + .expect("request should succeed"); + assert_eq!(resp.status().as_u16(), 200); + + let requests = harness.http_requests(); + assert_eq!(requests.len(), 1, "expected exactly one HTTP request"); + let (uri, host) = &requests[0]; + + // URI must be origin-form (path + query only, no scheme/authority) + assert_eq!(uri, "/some/path?key=val", "URI should be origin-form"); + + // Host header must be present with the correct authority + let host = host.as_deref().expect("Host header should be present"); + assert_eq!( + host, + format!("127.0.0.1:{port}"), + "Host header should match authority" + ); +} + +#[tokio::test] +async fn v1_origin_form_and_host_header() { + origin_form_and_host_header(&V1Client).await; +} + +#[tokio::test] +async fn v2_origin_form_and_host_header() { + origin_form_and_host_header(&V2Client).await; +} From f31feecc29bf19d9d3308dd863bd78b549323389 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 30 Apr 2026 12:00:36 -0400 Subject: [PATCH 10/47] conn limits --- .../aws-smithy-http-client/src/client/pool.rs | 35 ++- .../src/client/pool/connection.rs | 25 +- .../src/client/pool/handshake.rs | 123 +++++++-- .../aws-smithy-http-client/src/client/v2.rs | 30 ++- .../tests/pool_behavior_test.rs | 255 +++++++++++++++++- 5 files changed, 419 insertions(+), 49 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index db536c69642..8be67b7f005 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -29,6 +29,7 @@ mod cache { use std::collections::HashMap; use std::future::Future; use std::pin::Pin; +use std::sync::Arc; use std::sync::Mutex; use std::task::Poll; @@ -36,13 +37,13 @@ use aws_smithy_types::body::SdkBody; use hyper_util::client::legacy::connect::Connection as HyperConnection; use hyper_util::client::pool as hpool; use hyper_util::rt::TokioExecutor; +use tokio::sync::Semaphore; use tower::{Service, ServiceExt}; use connection::{ - CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, ManagedConnection, - SingletonConnection, + CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, SingletonConnection, }; -use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake, H2SendRequest}; +use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake}; pub(crate) type BoxError = Box; type BoxFuture = Pin + Send>>; @@ -60,10 +61,15 @@ type BoxFuture = Pin + Send>>; #[derive(Clone, Debug, Default)] pub(crate) struct PoolConfig { /// Upper bound on concurrent connections (total, across all hosts). - /// Enforced via `tower::ConcurrencyLimit` at the connection - /// establishment layer. `None` = unlimited. + /// Enforced via semaphore at the connection establishment layer. + /// `None` = unlimited. pub(crate) max_connections: Option, + /// Upper bound on concurrent connections per host. + /// Each unique (scheme, authority) pair gets an independent semaphore. + /// `None` = unlimited. + pub(crate) max_connections_per_host: Option, + /// How long an idle connection may stay in the pool before being /// evicted. `None` = no eviction (hyper-util's default behavior). pub(crate) pool_idle_timeout: Option, @@ -93,15 +99,20 @@ where C::Future: Unpin + Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + HyperConnection + Unpin + Send + 'static, { + let global_sem = config.max_connections.map(|n| Arc::new(Semaphore::new(n))); + let max_per_host = config.max_connections_per_host; + let make_entry = move |_uri: &http_1x::Uri| -> Box { - // TODO: consult `config.max_connections` here by wrapping the - // per-host connector with `tower::ConcurrencyLimit`, and consult - // `config.pool_idle_timeout` by driving `retain()` on a background - // task. Currently the config is stored on the pool but not yet - // applied. + let per_host_sem = max_per_host.map(|n| Arc::new(Semaphore::new(n))); + let limited = + handshake::ConnectionLimit::new(connector.clone(), global_sem.clone(), per_host_sem); let stack = hpool::negotiate::builder() - .connect(connector.clone()) - .inspect(|conn: &IO| conn.connected().is_negotiated_h2()) + .connect(limited) + .inspect( + |(conn, _permit): &(IO, Arc)| { + conn.connected().is_negotiated_h2() + }, + ) .fallback(tower::layer::layer_fn(|inspector| { cache::builder() .executor(TokioExecutor::new()) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 4a3d8a3de3a..c9b693eed11 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -13,12 +13,32 @@ use std::time::Instant; use aws_smithy_types::body::SdkBody; use pin_project_lite::pin_project; +use tokio::sync::OwnedSemaphorePermit; use tower::Service; use super::cache; use super::handshake::H1SendRequest; use super::BoxError; +/// Permits acquired from connection limit semaphores. +/// Held for the lifetime of the connection — dropped when the connection is dropped. +pub(crate) struct ConnectionPermit { + _global: Option, + _per_host: Option, +} + +impl ConnectionPermit { + pub(crate) fn new( + global: Option, + per_host: Option, + ) -> Self { + Self { + _global: global, + _per_host: per_host, + } + } +} + /// Error returned by `ManagedConnection::poll_ready` when the connection /// has been marked poisoned and should not be reused. #[derive(Debug)] @@ -68,16 +88,18 @@ pub(crate) struct ManagedConnection { info: ConnectionInfo, created_at: Instant, poisoned: Arc, + _permit: Arc, } impl ManagedConnection { /// Create a new managed connection wrapping the given service. - pub(crate) fn new(inner: S, info: ConnectionInfo) -> Self { + pub(crate) fn new(inner: S, info: ConnectionInfo, permit: Arc) -> Self { Self { inner, info, created_at: Instant::now(), poisoned: Arc::new(AtomicBool::new(false)), + _permit: permit, } } @@ -125,6 +147,7 @@ impl Clone for ManagedConnection { info: self.info.clone(), created_at: self.created_at, poisoned: self.poisoned.clone(), + _permit: self._permit.clone(), } } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 6561e457079..16bbdeb5ab2 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -5,22 +5,102 @@ //! Tower service adapters for hyper's HTTP protocol handshake. //! -//! Provides two layers: -//! - `H1SendRequest` / `H2SendRequest`: tower `Service` wrappers around -//! hyper's `SendRequest` types. -//! - `H1ConnectAndHandshake` / `H2ConnectAndHandshake`: services that chain -//! a connector with the protocol handshake, producing a `ManagedConnection`. +//! - `ConnectionLimit`: wraps a connector, acquires semaphore permits before +//! connecting, and returns `(IO, Arc)`. +//! - `H1ConnectAndHandshake` / `H2ConnectAndHandshake`: perform the protocol +//! handshake on an already-connected IO stream, producing a `ManagedConnection`. use std::future::Future; use std::pin::Pin; +use std::sync::Arc; use std::task::{Context, Poll}; use aws_smithy_types::body::SdkBody; use hyper_util::client::legacy::connect::{Connection, HttpInfo}; use hyper_util::rt::TokioExecutor; +use tokio::sync::Semaphore; use tower::Service; -use super::connection::{ConnectionInfo, ManagedConnection}; +use super::connection::{ConnectionInfo, ConnectionPermit, ManagedConnection}; +use super::BoxError; + +/// Wraps a connector service, acquiring semaphore permits before connecting. +/// +/// Returns `(IO, Arc)` so the permit can be stored on +/// `ManagedConnection` and held for the connection's lifetime. +pub(crate) struct ConnectionLimit { + inner: C, + global: Option>, + per_host: Option>, +} + +impl ConnectionLimit { + pub(crate) fn new( + inner: C, + global: Option>, + per_host: Option>, + ) -> Self { + Self { + inner, + global, + per_host, + } + } +} + +impl Clone for ConnectionLimit { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + global: self.global.clone(), + per_host: self.per_host.clone(), + } + } +} + +impl Service for ConnectionLimit +where + C: Service + Clone + Send + 'static, + C::Error: Into + 'static, + C::Future: Send + 'static, + IO: Send + 'static, + Req: Send + 'static, +{ + type Response = (IO, Arc); + type Error = BoxError; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx).map_err(Into::into) + } + + fn call(&mut self, req: Req) -> Self::Future { + let mut inner = self.inner.clone(); + let global = self.global.clone(); + let per_host = self.per_host.clone(); + Box::pin(async move { + // Acquire per-host permit first, then global. This ordering + // prevents deadlock: a request never holds a global permit + // while waiting on a per-host permit (which could starve + // other hosts' requests that need global permits). + let per_host_permit = match per_host { + Some(sem) => Some(sem.acquire_owned().await.map_err(|_| "pool closed")?), + None => None, + }; + let global_permit = match global { + Some(sem) => Some(sem.acquire_owned().await.map_err(|_| "pool closed")?), + None => None, + }; + let permit = Arc::new(ConnectionPermit::new(global_permit, per_host_permit)); + + std::future::poll_fn(|cx| inner.poll_ready(cx)) + .await + .map_err(Into::into)?; + let io = inner.call(req).await.map_err(Into::into)?; + Ok((io, permit)) + }) + } +} // --------------------------------------------------------------------------- // Tower Service wrappers for hyper's SendRequest @@ -83,14 +163,7 @@ impl Service> for H2SendRequest { // Connect-and-handshake services // --------------------------------------------------------------------------- -use super::BoxError; - /// Extract `ConnectionInfo` from a just-connected IO stream. -/// -/// Reads `Connected` metadata (ALPN, proxy flag, attached extras) and pulls out -/// `HttpInfo` (remote/local addr) if the connector attached it. This is the single -/// place we bridge between hyper-util's `Connection` trait and our own -/// `ConnectionInfo` type — if upstream reshapes, only this function changes. fn capture_info(io: &IO) -> ConnectionInfo { let connected = io.connected(); let mut extras = http_1x::Extensions::new(); @@ -104,8 +177,8 @@ fn capture_info(io: &IO) -> ConnectionInfo { /// Connects and performs an HTTP/1.1 handshake. /// -/// Wraps a connector service, performing the H1 protocol handshake after -/// connection and returning a `ManagedConnection`. +/// The connector is expected to return `(IO, Arc)` — typically +/// produced by [`ConnectionLimit`] wrapping a TCP/TLS connector. pub(crate) struct H1ConnectAndHandshake { connector: C, } @@ -126,7 +199,7 @@ impl Clone for H1ConnectAndHandshake { impl Service for H1ConnectAndHandshake where - C: Service, + C: Service)>, C::Error: Into + 'static, C::Future: Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, @@ -142,7 +215,7 @@ where fn call(&mut self, uri: http_1x::Uri) -> Self::Future { let fut = self.connector.call(uri); Box::pin(async move { - let io = fut.await.map_err(Into::into)?; + let (io, permit) = fut.await.map_err(Into::into)?; let info = capture_info(&io); let (tx, conn) = hyper::client::conn::http1::Builder::new() @@ -150,7 +223,6 @@ where .await .map_err(|e| Box::new(e) as BoxError)?; - // Drive the connection I/O in the background tokio::spawn({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; @@ -167,20 +239,15 @@ where } }); - Ok(ManagedConnection::new(H1SendRequest::new(tx), info)) + Ok(ManagedConnection::new(H1SendRequest::new(tx), info, permit)) }) } } /// Connects and performs an HTTP/2 handshake. /// -/// Wraps a connector service, performing the H2 protocol handshake after -/// connection and returning a `ManagedConnection`. -/// /// Pinned to `Service<()>` because this service sits in the Negotiate -/// upgrade path — by the time it runs, the connection is already -/// established; no target (URI) is needed. Keeping the request type -/// `()` prevents misassembly at compile time. +/// upgrade path — the connection is already established. pub(crate) struct H2ConnectAndHandshake { connector: C, } @@ -201,7 +268,7 @@ impl Clone for H2ConnectAndHandshake { impl Service<()> for H2ConnectAndHandshake where - C: Service<(), Response = IO>, + C: Service<(), Response = (IO, Arc)>, C::Error: Into + 'static, C::Future: Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, @@ -217,7 +284,7 @@ where fn call(&mut self, _req: ()) -> Self::Future { let fut = self.connector.call(()); Box::pin(async move { - let io = fut.await.map_err(Into::into)?; + let (io, permit) = fut.await.map_err(Into::into)?; let info = capture_info(&io); let (tx, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new()) @@ -241,7 +308,7 @@ where } }); - Ok(ManagedConnection::new(H2SendRequest::new(tx), info)) + Ok(ManagedConnection::new(H2SendRequest::new(tx), info, permit)) }) } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index 89631ac10f1..f6f3c096bf6 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -49,6 +49,8 @@ use crate::tls::TlsContext; pub struct BuilderV2 { pool_idle_timeout: Option, tcp_nodelay: bool, + max_connections: Option, + max_connections_per_host: Option, tls: Tls, } @@ -57,6 +59,8 @@ impl Default for BuilderV2 { Self { pool_idle_timeout: None, tcp_nodelay: true, + max_connections: None, + max_connections_per_host: None, tls: TlsUnset {}, } } @@ -79,6 +83,27 @@ impl BuilderV2 { self.tcp_nodelay = nodelay; self } + + /// Set the maximum number of concurrent connections. + /// + /// Limits concurrent connection handshakes across all hosts. Cached + /// (reused) connections do not count against this limit. When the limit + /// is reached, new connection attempts wait until an existing handshake + /// completes. + pub fn max_connections(mut self, n: usize) -> Self { + self.max_connections = Some(n); + self + } + + /// Set the maximum number of concurrent connections per host. + /// + /// Each unique (scheme, authority) pair has an independent connection + /// budget. This limit is independent of `max_connections` — both can + /// be set and are enforced simultaneously. + pub fn max_connections_per_host(mut self, n: usize) -> Self { + self.max_connections_per_host = Some(n); + self + } } impl BuilderV2 { @@ -92,6 +117,8 @@ impl BuilderV2 { BuilderV2 { pool_idle_timeout: self.pool_idle_timeout, tcp_nodelay: self.tcp_nodelay, + max_connections: self.max_connections, + max_connections_per_host: self.max_connections_per_host, tls: TlsProviderSelected { provider, context: TlsContext::default(), @@ -105,7 +132,8 @@ impl BuilderV2 { let mut tcp = HyperHttpConnector::new(); tcp.set_nodelay(self.tcp_nodelay); let config = PoolConfig { - max_connections: None, + max_connections: self.max_connections, + max_connections_per_host: self.max_connections_per_host, pool_idle_timeout: self.pool_idle_timeout, }; let pool = pool::build_pool(tcp, config); diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index e5755e2fe0d..8eb8a4f3450 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -37,6 +37,7 @@ const IP1: IpAddr = IpAddr::V4(Ipv4Addr::LOCALHOST); struct ClientConfig { idle_timeout: Option, max_connections: Option, + max_connections_per_host: Option, } impl ClientConfig { @@ -45,11 +46,15 @@ impl ClientConfig { self } - #[allow(dead_code)] fn with_max_connections(mut self, n: usize) -> Self { self.max_connections = Some(n); self } + + fn with_max_connections_per_host(mut self, n: usize) -> Self { + self.max_connections_per_host = Some(n); + self + } } trait MakeClient: Send + Sync { @@ -81,7 +86,12 @@ impl MakeClient for V2Client { if let Some(timeout) = config.idle_timeout { builder = builder.pool_idle_timeout(timeout); } - // max_connections deferred + if let Some(n) = config.max_connections { + builder = builder.max_connections(n); + } + if let Some(n) = config.max_connections_per_host { + builder = builder.max_connections_per_host(n); + } builder.build_http() } } @@ -181,11 +191,10 @@ async fn connection_reset_returns_error(make: &dyn MakeClient) { let client = make.make(ClientConfig::default()); let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); - let result = send_to(&client, &url).await; - assert!( - result.is_err(), - "request to a reset-on-connect endpoint should return an error" - ); + let err = send_to(&client, &url) + .await + .expect_err("request to a reset-on-connect endpoint should fail"); + assert!(err.is_io(), "expected ConnectorError::io, got: {err:?}"); } /// The client should report correct connector metadata. @@ -301,3 +310,235 @@ async fn v1_origin_form_and_host_header() { async fn v2_origin_form_and_host_header() { origin_form_and_host_header(&V2Client).await; } + +// --------------------------------------------------------------------------- +// Test implementations: max_connections +// --------------------------------------------------------------------------- + +/// With max_connections(2), concurrent requests should not open more than 2 connections, +/// and all requests should succeed (waiting requests served via connection reuse). +async fn max_connections_limits_concurrency(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..10) + .map(|_| ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }) + .collect(), + ) + .build() + .await; + + let client = make.make(ClientConfig::default().with_max_connections(2)); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + // Send 5 concurrent requests — all should succeed + let mut tasks = tokio::task::JoinSet::new(); + for _ in 0..5 { + let client = client.clone(); + let url = url.clone(); + tasks.spawn(async move { send_to(&client, &url).await }); + } + let mut success_count = 0; + while let Some(result) = tasks.join_next().await { + let resp = result + .expect("task should not panic") + .expect("request should succeed"); + assert_eq!(resp.status().as_u16(), 200); + success_count += 1; + } + assert_eq!(success_count, 5, "all 5 requests should complete"); + + let accepted = harness.tcp_accepted_count(); + assert!( + accepted <= 2, + "expected at most 2 connections with max_connections(2), got {accepted}" + ); +} + +/// Cached (reused) connections don't consume permits — sequential requests +/// with max_connections(1) all succeed on one connection. +async fn max_connections_reuse_does_not_consume_permits(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..5) + .map(|_| ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }) + .collect(), + ) + .build() + .await; + + let client = make.make(ClientConfig::default().with_max_connections(1)); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + // 5 sequential requests — all reuse the same connection + for i in 0..5 { + let resp = send_to(&client, &url) + .await + .unwrap_or_else(|e| panic!("request {i} should succeed: {e}")); + assert_eq!(resp.status().as_u16(), 200); + } + + assert_eq!( + harness.tcp_accepted_count(), + 1, + "all requests should reuse one connection" + ); +} + +#[tokio::test] +async fn v2_max_connections_limits_concurrency() { + max_connections_limits_concurrency(&V2Client).await; +} + +#[tokio::test] +async fn v2_max_connections_reuse_does_not_consume_permits() { + max_connections_reuse_does_not_consume_permits(&V2Client).await; +} + +// --------------------------------------------------------------------------- +// Test implementations: max_connections_per_host +// --------------------------------------------------------------------------- + +async fn is_bindable(ip: IpAddr) -> bool { + tokio::net::TcpListener::bind((ip, 0u16)).await.is_ok() +} + +const IP2: IpAddr = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)); + +/// Per-host limit: each host gets its own budget. +async fn max_connections_per_host_limits_per_host(make: &dyn MakeClient) { + if !is_bindable(IP2).await { + eprintln!("skipping test: 127.0.0.2 not bindable"); + return; + } + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..10) + .map(|_| ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }) + .collect(), + ) + .endpoint( + IP2, + (0..10) + .map(|_| ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }) + .collect(), + ) + .build() + .await; + + let client = make.make(ClientConfig::default().with_max_connections_per_host(2)); + let port = harness.endpoints[0].port(); + + // Send 4 concurrent requests to each host + let mut tasks = tokio::task::JoinSet::new(); + for ip in ["127.0.0.1", "127.0.0.2"] { + for _ in 0..4 { + let client = client.clone(); + let url = format!("http://{ip}:{port}/"); + tasks.spawn(async move { send_to(&client, &url).await }); + } + } + while let Some(result) = tasks.join_next().await { + result + .expect("task should not panic") + .expect("request should succeed"); + } + + let host1 = harness.tcp_accepted_by(IP1); + let host2 = harness.tcp_accepted_by(IP2); + assert!(host1 <= 2, "host 1: expected ≤2 connections, got {host1}"); + assert!(host2 <= 2, "host 2: expected ≤2 connections, got {host2}"); + // Total exceeds per-host limit (no global limit set) + assert!(host1 + host2 > 2, "total should exceed per-host limit"); +} + +/// Global and per-host limits compose: global(3) + per_host(2) with 2 hosts. +async fn max_connections_global_and_per_host_compose(make: &dyn MakeClient) { + if !is_bindable(IP2).await { + eprintln!("skipping test: 127.0.0.2 not bindable"); + return; + } + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..10) + .map(|_| ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }) + .collect(), + ) + .endpoint( + IP2, + (0..10) + .map(|_| ConnectionBehavior::Respond { + status: 200, + body: b"ok", + }) + .collect(), + ) + .build() + .await; + + let client = make.make( + ClientConfig::default() + .with_max_connections(3) + .with_max_connections_per_host(2), + ); + let port = harness.endpoints[0].port(); + + // Send 4 concurrent requests to each host (8 total) + let mut tasks = tokio::task::JoinSet::new(); + for ip in ["127.0.0.1", "127.0.0.2"] { + for _ in 0..4 { + let client = client.clone(); + let url = format!("http://{ip}:{port}/"); + tasks.spawn(async move { send_to(&client, &url).await }); + } + } + while let Some(result) = tasks.join_next().await { + result + .expect("task should not panic") + .expect("request should succeed"); + } + + let host1 = harness.tcp_accepted_by(IP1); + let host2 = harness.tcp_accepted_by(IP2); + // Per-host: each ≤2 + assert!(host1 <= 2, "host 1: expected ≤2, got {host1}"); + assert!(host2 <= 2, "host 2: expected ≤2, got {host2}"); + // Global: total ≤3 + assert!( + host1 + host2 <= 3, + "total: expected ≤3 (global limit), got {}", + host1 + host2 + ); +} + +#[tokio::test] +async fn v2_max_connections_per_host_limits_per_host() { + max_connections_per_host_limits_per_host(&V2Client).await; +} + +#[tokio::test] +async fn v2_max_connections_global_and_per_host_compose() { + max_connections_global_and_per_host_compose(&V2Client).await; +} From e8099ec79bacb4f8946b9579397aac1ae539e68b Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 30 Apr 2026 16:38:48 -0400 Subject: [PATCH 11/47] add more tests for pool behaviors --- .../src/client/pool/handshake.rs | 15 +- .../src/test_util/wire/connection.rs | 75 +++- .../tests/connection_harness_test.rs | 20 +- .../tests/pool_behavior_test.rs | 368 +++++++++++++++++- 4 files changed, 444 insertions(+), 34 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 16bbdeb5ab2..f2bf2cf000e 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use aws_smithy_types::body::SdkBody; +use hyper::rt::Executor; use hyper_util::client::legacy::connect::{Connection, HttpInfo}; use hyper_util::rt::TokioExecutor; use tokio::sync::Semaphore; @@ -175,6 +176,13 @@ fn capture_info(io: &IO) -> ConnectionInfo { } } +/// Spawn a connection driver future onto the runtime. +/// +// TODO - revisit tokio default without flags +fn spawn_driver(future: impl Future + Send + 'static) { + TokioExecutor::new().execute(future); +} + /// Connects and performs an HTTP/1.1 handshake. /// /// The connector is expected to return `(IO, Arc)` — typically @@ -223,7 +231,7 @@ where .await .map_err(|e| Box::new(e) as BoxError)?; - tokio::spawn({ + spawn_driver({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { @@ -247,7 +255,8 @@ where /// Connects and performs an HTTP/2 handshake. /// /// Pinned to `Service<()>` because this service sits in the Negotiate -/// upgrade path — the connection is already established. +/// upgrade path — the connection is already established via the shared +/// `Inspected` slot. pub(crate) struct H2ConnectAndHandshake { connector: C, } @@ -292,7 +301,7 @@ where .await .map_err(|e| Box::new(e) as BoxError)?; - tokio::spawn({ + spawn_driver({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { diff --git a/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs b/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs index b6753661e25..14408b5e558 100644 --- a/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/test_util/wire/connection.rs @@ -19,8 +19,9 @@ use tokio::task::JoinHandle; /// Programmable behavior for a test endpoint per accepted connection. #[derive(Debug, Clone)] pub enum ConnectionBehavior { - /// Accept TCP, send HTTP/1.1 response, keep connection open (reusable). - Respond { + /// Accept TCP, send HTTP/1.1 response with `Connection: keep-alive`, keep + /// connection open for subsequent requests on the same TCP stream. + RespondKeepAlive { /// HTTP status code to return. status: u16, /// Response body bytes. @@ -28,8 +29,28 @@ pub enum ConnectionBehavior { }, /// Accept TCP, immediately reset the connection (RST). ResetOnConnect, + /// Accept TCP, send HTTP/1.1 response, then close the connection. + /// Simulates a server that does not support keep-alive. + RespondThenClose { + /// HTTP status code to return. + status: u16, + /// Response body bytes. + body: &'static [u8], + }, /// Accept TCP, hold open for duration, then close. HoldThenClose(std::time::Duration), + /// Accept TCP, send HTTP/1.1 response with `Connection: keep-alive`, + /// wait for the specified duration, then close the connection. + /// Simulates a server-side idle timeout (e.g. S3 closing after ~20s). + /// The connection appears reusable to the client until the server closes it. + RespondThenIdleClose { + /// HTTP status code to return. + status: u16, + /// Response body bytes. + body: &'static [u8], + /// How long to wait after responding before closing. + idle: std::time::Duration, + }, } /// Recorded event from the test harness. @@ -137,11 +158,39 @@ async fn handle_connection( .expect("failed to set SO_LINGER"); drop(stream); } + ConnectionBehavior::RespondThenClose { status, body } => { + match read_request(&mut stream).await { + Ok(req_info) => { + events.lock().expect("event lock poisoned").push(req_info); + } + Err(_) => return, + } + let _ = write_response(&mut stream, status, body, false).await; + // Close immediately after responding — no keep-alive. + drop(stream); + } ConnectionBehavior::HoldThenClose(duration) => { tokio::time::sleep(duration).await; drop(stream); } - ConnectionBehavior::Respond { status, body } => { + ConnectionBehavior::RespondThenIdleClose { status, body, idle } => { + match read_request(&mut stream).await { + Ok(req_info) => { + events.lock().expect("event lock poisoned").push(req_info); + } + Err(_) => return, + } + if write_response(&mut stream, status, body, true) + .await + .is_err() + { + return; + } + // Wait, then close — simulates server-side idle timeout. + tokio::time::sleep(idle).await; + drop(stream); + } + ConnectionBehavior::RespondKeepAlive { status, body } => { // Read the first request before responding. match read_request(&mut stream).await { Ok(req_info) => { @@ -149,7 +198,10 @@ async fn handle_connection( } Err(_) => return, } - if write_response(&mut stream, status, body).await.is_err() { + if write_response(&mut stream, status, body, true) + .await + .is_err() + { return; } // Keep-alive loop: read next request, send next behavior's response @@ -165,11 +217,18 @@ async fn handle_connection( .expect("behavior lock poisoned") .pop_front(); match next { - Some(ConnectionBehavior::Respond { status, body }) => { - if write_response(&mut stream, status, body).await.is_err() { + Some(ConnectionBehavior::RespondKeepAlive { status, body }) => { + if write_response(&mut stream, status, body, true) + .await + .is_err() + { return; } } + Some(ConnectionBehavior::RespondThenClose { status, body }) => { + let _ = write_response(&mut stream, status, body, false).await; + return; // close after responding + } _ => return, // No more behaviors or non-Respond behavior: close } } @@ -181,9 +240,11 @@ async fn write_response( stream: &mut tokio::net::TcpStream, status: u16, body: &[u8], + keep_alive: bool, ) -> Result<(), std::io::Error> { + let conn_header = if keep_alive { "keep-alive" } else { "close" }; let mut response = format!( - "HTTP/1.1 {status} OK\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n", + "HTTP/1.1 {status} OK\r\nContent-Length: {}\r\nConnection: {conn_header}\r\n\r\n", body.len() ) .into_bytes(); diff --git a/rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs b/rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs index 3f07b536579..8c9a804a8c6 100644 --- a/rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/connection_harness_test.rs @@ -29,14 +29,14 @@ async fn test_harness_multi_ip_endpoints() { let harness = ConnectionTestHarness::builder() .endpoint( IP1, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"hello", }], ) .endpoint( IP2, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"world", }], @@ -109,14 +109,14 @@ async fn test_mock_dns_resolver() { let harness = ConnectionTestHarness::builder() .endpoint( IP1, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }], ) .endpoint( IP2, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }], @@ -141,15 +141,15 @@ async fn test_harness_keep_alive_reuse() { .endpoint( IP1, vec![ - ConnectionBehavior::Respond { + ConnectionBehavior::RespondKeepAlive { status: 200, body: b"first", }, - ConnectionBehavior::Respond { + ConnectionBehavior::RespondKeepAlive { status: 201, body: b"second", }, - ConnectionBehavior::Respond { + ConnectionBehavior::RespondKeepAlive { status: 202, body: b"third", }, @@ -223,14 +223,14 @@ async fn test_harness_dns_all_includes_all_endpoints() { let harness = ConnectionTestHarness::builder() .endpoint( IP1, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"a", }], ) .endpoint( IP2, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"b", }], @@ -255,7 +255,7 @@ async fn test_mock_dns_unknown_host_returns_empty() { let harness = ConnectionTestHarness::builder() .endpoint( IP1, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }], diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 8eb8a4f3450..3a1bf1e00b2 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -122,6 +122,27 @@ async fn send_to( .await } +/// Send a request and read the response body to completion. +/// +/// This ensures the connection's body guard is released, returning the +/// connection to the pool. Returns `(status, body_bytes)`. +async fn send_and_read_body( + client: &SharedHttpClient, + url: &str, +) -> Result<(u16, Vec), aws_smithy_runtime_api::client::result::ConnectorError> { + use http_body_util::BodyExt; + let resp = send_to(client, url).await?; + let status = resp.status().as_u16(); + let body = resp + .into_body() + .collect() + .await + .expect("body should be readable") + .to_bytes() + .to_vec(); + Ok((status, body)) +} + fn localhost_url(server: &WireMockServer) -> String { let endpoint = server.endpoint_url(); let port = endpoint @@ -166,17 +187,22 @@ async fn idle_timeout_eviction(make: &dyn MakeClient) { let client = make.make(ClientConfig::default().with_idle_timeout(idle_timeout)); let url = localhost_url(&server); - let resp = send_to(&client, &url) + let status = send_to(&client, &url) .await - .expect("first request should succeed"); - assert_eq!(resp.status().as_u16(), 200); + .expect("first request should succeed") + .status() + .as_u16(); + assert_eq!(status, 200); + // Response (and its body guard) dropped here — connection returns to pool. tokio::time::sleep(idle_timeout * 2).await; - let resp = send_to(&client, &url) + let status = send_to(&client, &url) .await - .expect("second request should succeed after idle eviction"); - assert_eq!(resp.status().as_u16(), 200); + .expect("second request should succeed after idle eviction") + .status() + .as_u16(); + assert_eq!(status, 200); match_events!(ev!(connect), ev!(http(200)), ev!(connect), ev!(http(200)))(&server.events()); } @@ -246,7 +272,20 @@ async fn v2_connection_reuse() { #[tokio::test] async fn v2_idle_timeout_eviction() { - idle_timeout_eviction(&V2Client).await; + // v2 does not yet implement idle timeout eviction (Phase 4). + // + // The v1 version of this test passes because the legacy client has + // built-in idle eviction. For v2, calling idle_timeout_eviction() + // also produces 2 connections — but NOT because of idle eviction. + // The response body is never consumed by the test (HttpResponse wraps + // a lazy SdkBody). The CachedConnection is held inside the response + // body's guard until the HttpResponse is dropped. Since the first + // response is still alive during the sleep, the connection never + // returns to Cache, and the second request opens a new connection. + // + // Real idle eviction tests (consume body → return to pool → sleep + // past timeout → verify new connection) will be added in Phase 4 + // alongside the background eviction task. } #[tokio::test] @@ -268,7 +307,7 @@ async fn origin_form_and_host_header(make: &dyn MakeClient) { let harness = ConnectionTestHarness::builder() .endpoint( IP1, - vec![ConnectionBehavior::Respond { + vec![ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }], @@ -322,7 +361,7 @@ async fn max_connections_limits_concurrency(make: &dyn MakeClient) { .endpoint( IP1, (0..10) - .map(|_| ConnectionBehavior::Respond { + .map(|_| ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }) @@ -366,7 +405,7 @@ async fn max_connections_reuse_does_not_consume_permits(make: &dyn MakeClient) { .endpoint( IP1, (0..5) - .map(|_| ConnectionBehavior::Respond { + .map(|_| ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }) @@ -425,7 +464,7 @@ async fn max_connections_per_host_limits_per_host(make: &dyn MakeClient) { .endpoint( IP1, (0..10) - .map(|_| ConnectionBehavior::Respond { + .map(|_| ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }) @@ -434,7 +473,7 @@ async fn max_connections_per_host_limits_per_host(make: &dyn MakeClient) { .endpoint( IP2, (0..10) - .map(|_| ConnectionBehavior::Respond { + .map(|_| ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }) @@ -480,7 +519,7 @@ async fn max_connections_global_and_per_host_compose(make: &dyn MakeClient) { .endpoint( IP1, (0..10) - .map(|_| ConnectionBehavior::Respond { + .map(|_| ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }) @@ -489,7 +528,7 @@ async fn max_connections_global_and_per_host_compose(make: &dyn MakeClient) { .endpoint( IP2, (0..10) - .map(|_| ConnectionBehavior::Respond { + .map(|_| ConnectionBehavior::RespondKeepAlive { status: 200, body: b"ok", }) @@ -542,3 +581,304 @@ async fn v2_max_connections_per_host_limits_per_host() { async fn v2_max_connections_global_and_per_host_compose() { max_connections_global_and_per_host_compose(&V2Client).await; } + +// --------------------------------------------------------------------------- +// Test implementations: connection lifecycle +// --------------------------------------------------------------------------- + +/// Without max_connections, connections scale freely under concurrency. +async fn no_limit_allows_unbounded_connections(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..20) + .map(|_| ConnectionBehavior::HoldThenClose(Duration::from_secs(2))) + .collect(), + ) + .build() + .await; + + let client = make.make(ClientConfig::default()); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + // Send 5 concurrent requests — server holds each connection open, + // so the client must open 5 separate connections. + let mut tasks = tokio::task::JoinSet::new(); + for _ in 0..5 { + let client = client.clone(); + let url = url.clone(); + tasks.spawn(async move { + // These will fail (server doesn't send HTTP response), but + // the point is that 5 TCP connections are opened concurrently. + let _ = send_to(&client, &url).await; + }); + } + // Give time for all connections to be established + tokio::time::sleep(Duration::from_millis(200)).await; + + let accepted = harness.tcp_accepted_count(); + assert!( + accepted >= 4, + "without max_connections, expected ≥4 concurrent connections, got {accepted}" + ); + + // Clean up tasks + tasks.shutdown().await; +} + +/// Sequential requests reuse the same connection (body fully consumed between requests). +async fn connection_reuse_after_body_consumed(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..3) + .map(|_| ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"hello world response body", + }) + .collect(), + ) + .build() + .await; + + let client = make.make(ClientConfig::default()); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + for i in 0..3 { + let (status, body) = send_and_read_body(&client, &url) + .await + .unwrap_or_else(|e| panic!("request {i} should succeed: {e}")); + assert_eq!(status, 200); + assert_eq!(body, b"hello world response body"); + } + + assert_eq!( + harness.tcp_accepted_count(), + 1, + "all requests should reuse one connection after body is consumed" + ); +} + +/// Active connections are not evicted by idle timeout. +async fn active_connection_survives_idle_timeout(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"first", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"second", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"third", + }, + ], + ) + .build() + .await; + + // Use a very short idle timeout — but sequential requests should still + // reuse the connection because the pool doesn't evict between requests + // that happen back-to-back. + let client = make.make(ClientConfig::default().with_idle_timeout(Duration::from_millis(100))); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + // Three back-to-back requests — all should reuse the same connection + let expected = [b"first".as_slice(), b"second", b"third"]; + for (i, expected_body) in expected.iter().enumerate() { + let (status, body) = send_and_read_body(&client, &url) + .await + .unwrap_or_else(|e| panic!("request {i} should succeed: {e}")); + assert_eq!(status, 200); + assert_eq!(body, *expected_body, "request {i} body mismatch"); + } + + assert_eq!( + harness.tcp_accepted_count(), + 1, + "back-to-back requests should reuse one connection even with short idle timeout" + ); +} + +#[tokio::test] +async fn v2_no_limit_allows_unbounded_connections() { + no_limit_allows_unbounded_connections(&V2Client).await; +} + +#[tokio::test] +async fn v2_connection_reuse_after_body_consumed() { + connection_reuse_after_body_consumed(&V2Client).await; +} + +#[tokio::test] +async fn v2_active_connection_survives_idle_timeout() { + active_connection_survives_idle_timeout(&V2Client).await; +} + +// --------------------------------------------------------------------------- +// Test implementations: connection lifecycle +// --------------------------------------------------------------------------- + +/// Server closes an idle connection (simulating server-side idle timeout). +/// The connection returns to the pool clean, then dies while idle. On next +/// checkout, poll_ready detects the dead connection, the checkout loop +/// discards it, and a fresh connection is created. +async fn stale_connection_detected_at_checkout(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + // First connection: respond with keep-alive, then close after 30ms + ConnectionBehavior::RespondThenIdleClose { + status: 200, + body: b"first", + idle: Duration::from_millis(30), + }, + // Second connection (after stale one is discarded) + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"second", + }, + ], + ) + .build() + .await; + + let client = make.make(ClientConfig::default()); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + // First request succeeds, body consumed, connection returns to pool + let (status, body) = send_and_read_body(&client, &url) + .await + .expect("first request"); + assert_eq!(status, 200); + assert_eq!(body, b"first"); + + // Wait for server to close the idle connection (30ms idle + margin) + tokio::time::sleep(Duration::from_millis(80)).await; + + // Second request: checkout finds stale connection, discards, creates new + let (status, body) = send_and_read_body(&client, &url) + .await + .expect("second request should succeed on fresh connection"); + assert_eq!(status, 200); + assert_eq!(body, b"second"); + + assert_eq!( + harness.tcp_accepted_count(), + 2, + "should have opened 2 connections (first died while idle in pool)" + ); +} + +/// When the response body is not consumed, the connection is held by the +/// body guard and unavailable for reuse. A concurrent request must open +/// a new connection. +async fn unconsumed_body_holds_connection(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"first", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"second", + }, + ], + ) + .build() + .await; + + let client = make.make(ClientConfig::default()); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + // First request — hold the response (body unconsumed, connection held) + let _held_resp = send_to(&client, &url).await.expect("first request"); + + // Second request while first response is still held — must open new connection + let (status, body) = send_and_read_body(&client, &url) + .await + .expect("second request"); + assert_eq!(status, 200); + assert_eq!(body, b"second"); + + assert_eq!( + harness.tcp_accepted_count(), + 2, + "second request should open a new connection while first body is held" + ); + + // Drop the held response — connection guard released + drop(_held_resp); +} + +/// Server sends Connection: close — client should not reuse the connection. +async fn connection_close_header_prevents_reuse(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondThenClose { + status: 200, + body: b"closing", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"fresh", + }, + ], + ) + .build() + .await; + + let client = make.make(ClientConfig::default()); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + + let (status, body) = send_and_read_body(&client, &url) + .await + .expect("first request"); + assert_eq!(status, 200); + assert_eq!(body, b"closing"); + + // Connection: close means the client should not attempt to reuse + let (status, body) = send_and_read_body(&client, &url) + .await + .expect("second request"); + assert_eq!(status, 200); + assert_eq!(body, b"fresh"); + + assert_eq!( + harness.tcp_accepted_count(), + 2, + "Connection: close should prevent reuse" + ); +} + +#[tokio::test] +async fn v2_stale_connection_detected_at_checkout() { + stale_connection_detected_at_checkout(&V2Client).await; +} + +#[tokio::test] +async fn v2_unconsumed_body_holds_connection() { + unconsumed_body_holds_connection(&V2Client).await; +} + +#[tokio::test] +async fn v2_connection_close_header_prevents_reuse() { + connection_close_header_prevents_reuse(&V2Client).await; +} From 32a92934cca0856ec53835b688c0d4892849edea Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 1 May 2026 13:32:27 -0400 Subject: [PATCH 12/47] wire v2 tls support --- .../aws-smithy-http-client/src/client/v2.rs | 77 ++++++++++++++- .../tests/smoke_test_clients.rs | 48 +++++++++ .../aws-smithy-http-client/tests/tls.rs | 98 +++++++++++++++++++ 3 files changed, 222 insertions(+), 1 deletion(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index f6f3c096bf6..f9c67730781 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -141,7 +141,82 @@ impl BuilderV2 { } } -// TODO(3f): impl BuilderV2 { fn build_https() } +impl BuilderV2 { + /// Set the TLS context (custom trust store, etc.). + pub fn tls_context(mut self, context: TlsContext) -> Self { + self.tls.context = context; + self + } + + /// Build an HTTPS client with the selected TLS provider. + pub fn build_https(self) -> SharedHttpClient { + let mut tcp = HyperHttpConnector::new(); + tcp.set_nodelay(self.tcp_nodelay); + tcp.enforce_http(false); + self.build_from_tcp(tcp) + } + + /// Build an HTTPS client using a custom DNS resolver. + pub fn build_with_resolver( + self, + resolver: impl aws_smithy_runtime_api::client::dns::ResolveDns + Clone + 'static, + ) -> SharedHttpClient { + use crate::client::dns::HyperUtilResolver; + let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); + tcp.set_nodelay(self.tcp_nodelay); + tcp.enforce_http(false); + self.build_from_tcp(tcp) + } + + fn build_from_tcp(self, tcp: HyperHttpConnector) -> SharedHttpClient + where + R: Clone + Send + Sync + 'static, + R: tower::Service, + R::Response: Iterator, + R::Future: Send, + R::Error: Into>, + { + use super::proxy; + + let config = PoolConfig { + max_connections: self.max_connections, + max_connections_per_host: self.max_connections_per_host, + pool_idle_timeout: self.pool_idle_timeout, + }; + + match &self.tls.provider { + #[cfg(any( + feature = "rustls-aws-lc", + feature = "rustls-aws-lc-fips", + feature = "rustls-ring" + ))] + tls::Provider::Rustls(crypto_mode) => { + let connector = tls::rustls_provider::build_connector::wrap_connector( + tcp, + crypto_mode.clone(), + &self.tls.context, + proxy::ProxyConfig::disabled(), + ); + let pool = pool::build_pool(connector, config); + V2HttpClient::new(pool).into_shared() + } + #[cfg(feature = "s2n-tls")] + tls::Provider::S2nTls => { + let connector = tls::s2n_tls_provider::build_connector::wrap_connector( + tcp, + &self.tls.context, + proxy::ProxyConfig::disabled(), + ); + let pool = pool::build_pool(connector, config); + V2HttpClient::new(pool).into_shared() + } + // Provider is #[non_exhaustive]; this arm is unreachable when any + // TLS feature is enabled (which is required to construct a Provider). + #[allow(unreachable_patterns)] + _ => unreachable!("a TLS feature must be enabled to use build_https()"), + } + } +} // --------------------------------------------------------------------------- // HttpConnector adapter diff --git a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs index 1c87cf4cd70..406c22993d0 100644 --- a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs +++ b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs @@ -125,3 +125,51 @@ async fn smoke_test_client(client: &dyn HttpClient) -> Result<(), Box .await?; Ok(()) } + +// --------------------------------------------------------------------------- +// v2 smoke tests +// --------------------------------------------------------------------------- + +use aws_smithy_http_client::v2::BuilderV2; + +#[cfg(feature = "rustls-aws-lc")] +#[tokio::test] +async fn v2_aws_lc_client() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLc, + )) + .build_https(); + smoke_test_client(&client).await.unwrap(); +} + +#[cfg(feature = "rustls-aws-lc-fips")] +#[tokio::test] +async fn v2_aws_lc_fips_client() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLcFips, + )) + .build_https(); + smoke_test_client(&client).await.unwrap(); +} + +#[cfg(feature = "rustls-ring")] +#[tokio::test] +async fn v2_ring_client() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::Ring, + )) + .build_https(); + smoke_test_client(&client).await.unwrap(); +} + +#[cfg(feature = "s2n-tls")] +#[tokio::test] +async fn v2_s2n_tls_client() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::S2nTls) + .build_https(); + smoke_test_client(&client).await.unwrap(); +} diff --git a/rust-runtime/aws-smithy-http-client/tests/tls.rs b/rust-runtime/aws-smithy-http-client/tests/tls.rs index 7f310d32d4e..4bd8d4db061 100644 --- a/rust-runtime/aws-smithy-http-client/tests/tls.rs +++ b/rust-runtime/aws-smithy-http-client/tests/tls.rs @@ -307,6 +307,104 @@ async fn test_s2n_tls_custom_ca() { run_tls_test(&client).await.unwrap() } +// --------------------------------------------------------------------------- +// v2 TLS tests +// --------------------------------------------------------------------------- + +use aws_smithy_http_client::v2::BuilderV2; + +#[cfg(feature = "rustls-aws-lc")] +#[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] +#[tokio::test] +async fn test_v2_rustls_aws_lc_native_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLc, + )) + .build_https(); + run_tls_test(&client).await.unwrap() +} + +#[cfg(feature = "rustls-aws-lc")] +#[tokio::test] +async fn test_v2_rustls_aws_lc_custom_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLc, + )) + .tls_context(tls_context_from_pem("tests/server.pem")) + .build_https(); + run_tls_test(&client).await.unwrap() +} + +#[cfg(feature = "rustls-aws-lc-fips")] +#[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] +#[tokio::test] +async fn test_v2_rustls_aws_lc_fips_native_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLcFips, + )) + .build_https(); + run_tls_test(&client).await.unwrap() +} + +#[cfg(feature = "rustls-aws-lc-fips")] +#[tokio::test] +async fn test_v2_rustls_aws_lc_fips_custom_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLcFips, + )) + .tls_context(tls_context_from_pem("tests/server.pem")) + .build_https(); + run_tls_test(&client).await.unwrap() +} + +#[cfg(feature = "rustls-ring")] +#[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] +#[tokio::test] +async fn test_v2_rustls_ring_native_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::Ring, + )) + .build_https(); + run_tls_test(&client).await.unwrap() +} + +#[cfg(feature = "rustls-ring")] +#[tokio::test] +async fn test_v2_rustls_ring_custom_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::Ring, + )) + .tls_context(tls_context_from_pem("tests/server.pem")) + .build_https(); + run_tls_test(&client).await.unwrap() +} + +#[cfg(feature = "s2n-tls")] +#[should_panic(expected = "Certificate is untrusted")] +#[tokio::test] +async fn test_v2_s2n_native_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::S2nTls) + .build_https(); + run_tls_test(&client).await.unwrap() +} + +#[cfg(feature = "s2n-tls")] +#[tokio::test] +async fn test_v2_s2n_tls_custom_ca() { + let client = BuilderV2::new() + .tls_provider(tls::Provider::S2nTls) + .tls_context(tls_context_from_pem("tests/server.pem")) + .build_https(); + run_tls_test(&client).await.unwrap() +} + async fn run_tls_test(client: &dyn HttpClient) -> Result<(), BoxError> { run_tls_test_with_idle_timeout(client, None).await } From 53b26c1245765cf7d5509e482b9282b1f631bc29 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 5 May 2026 12:53:47 -0400 Subject: [PATCH 13/47] fix poison support --- .../aws-smithy-http-client/src/client/pool.rs | 122 +++++++++---- .../src/client/pool/connection.rs | 167 ++++++++++++++---- .../src/client/pool/handshake.rs | 20 ++- .../aws-smithy-http-client/src/client/v2.rs | 25 ++- .../tests/pool_behavior_test.rs | 152 ++++++++++++++++ 5 files changed, 405 insertions(+), 81 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 8be67b7f005..85a1f403f2a 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -5,11 +5,12 @@ //! Connection pool for the v2 HTTP client. //! -//! Manages the lifecycle of HTTP connections including: -//! - Connection state tracking (idle duration, remote address, poisoned flag) -//! - Health checks on checkout -//! - Connection poisoning -//! - Idle connection eviction +//! Routes requests by `(scheme, authority)` to per-host pool stacks, +//! negotiates HTTP/1.1 vs HTTP/2 via ALPN, caches H1 connections for +//! reuse, multiplexes over a singleton H2 connection per host, enforces +//! total and per-host connection limits, and surfaces connection +//! metadata (remote/local address, poison handle) to the adapter layer +//! for the runtime's connection-poisoning interceptor. mod connection; mod handshake; @@ -17,11 +18,10 @@ mod vendored_cache; /// Connection-caching pool layer. /// -/// Currently re-exports [`vendored_cache`], a copy of hyper-util's -/// `pool::cache` module with SDK-specific modifications. The re-export -/// insulates the rest of the pool from the vendoring detail: call sites -/// use `cache::...` regardless of whether the implementation is vendored -/// or, later, fully owned. +/// Re-exports [`vendored_cache`], our vendored copy of hyper-util's +/// `pool::cache` with SDK-specific modifications. The re-export insulates +/// callers from the vendoring detail — all pool code uses `cache::…` +/// regardless of where the implementation lives. mod cache { pub(crate) use super::vendored_cache::*; } @@ -41,13 +41,54 @@ use tokio::sync::Semaphore; use tower::{Service, ServiceExt}; use connection::{ - CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, SingletonConnection, + CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, H2ConnectionRef, + SingletonConnection, }; use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake}; pub(crate) type BoxError = Box; type BoxFuture = Pin + Send>>; +/// Request-extension slot the pool checkout fills in with the +/// `ConnectionMetadata` for the selected connection. +/// +/// Read later by the adapter's `CaptureSmithyConnection` retriever, which +/// is what `ConnectionPoisoningInterceptor` uses to decide whether to call +/// `ConnectionMetadata::poison()` on a transient error. Poisoning flips the +/// shared `PoisonPill` on the actual `ManagedConnection`, so the pool +/// skips it on checkout and drops it on return. +/// +/// Write-once per request: `H{1,2}Checkout::call` sets it exactly once +/// during checkout. Subsequent sets are no-ops (the `OnceLock` guarantees +/// single-init). Retries produce fresh `HttpConnector::call` invocations +/// which create fresh capture slots — each attempt's metadata points at +/// the connection used for that attempt. +#[derive(Clone, Default)] +pub(crate) struct ConnectionMetadataCapture { + slot: Arc>, +} + +impl ConnectionMetadataCapture { + pub(crate) fn new() -> Self { + Self::default() + } + + pub(crate) fn set( + &self, + metadata: aws_smithy_runtime_api::client::connection::ConnectionMetadata, + ) { + // Silently ignore duplicate sets: single-set is the contract, extra + // sets would only happen via pool-internal bugs. + let _ = self.slot.set(metadata); + } + + pub(crate) fn get( + &self, + ) -> Option { + self.slot.get().cloned() + } +} + /// Pool-level configuration. /// /// Plumbed from the eventual `Builder::new_v2()` public API through @@ -106,6 +147,12 @@ where let per_host_sem = max_per_host.map(|n| Arc::new(Semaphore::new(n))); let limited = handshake::ConnectionLimit::new(connector.clone(), global_sem.clone(), per_host_sem); + + // Per-host bridge: `H2ConnectAndHandshake` publishes on each new + // handshake; `SingletonConnection` reads the current entry at + // checkout. Clones share the underlying slot. + let h2_ref = H2ConnectionRef::new(); + let stack = hpool::negotiate::builder() .connect(limited) .inspect( @@ -119,10 +166,21 @@ where .build(H1ConnectAndHandshake::new(inspector)) .map_response(|cached| H1Checkout::new(CachedConnection::new(cached))) })) - .upgrade(tower::layer::layer_fn(|inspected| { - hpool::singleton::Singleton::new(H2ConnectAndHandshake::new(inspected)) - .map_response(|singled| H2Checkout::new(SingletonConnection::new(singled))) - })) + .upgrade({ + let h2_ref = h2_ref.clone(); + tower::layer::layer_fn(move |inspected| { + hpool::singleton::Singleton::new(H2ConnectAndHandshake::new( + inspected, + h2_ref.clone(), + )) + .map_response({ + let h2_ref = h2_ref.clone(); + move |singled| { + H2Checkout::new(SingletonConnection::new(singled, h2_ref.clone())) + } + }) + }) + }) .build(); Box::new(TypedPoolEntry(stack)) }; @@ -140,8 +198,7 @@ where /// Each host gets a Negotiate stack that selects between HTTP/1.1 (Cache) /// and HTTP/2 (Singleton) based on ALPN negotiation. pub(crate) struct ConnectionPool { - /// Pool-wide configuration. Stored for consultation at request time - /// and for re-use if `make_entry` needs config (not yet wired). + /// Pool-wide configuration. #[allow(dead_code)] config: PoolConfig, hosts: Mutex>>, @@ -222,6 +279,12 @@ impl Service> for H1Checkout) -> Self::Future { let mut conn = self.conn.take().expect("H1Checkout::call called twice"); + // Populate the adapter-provided capture so the + // `CaptureSmithyConnection` retriever can return a live + // `ConnectionMetadata` pointing at THIS connection's poison pill. + if let Some(capture) = req.extensions().get::() { + capture.set(conn.metadata()); + } Box::pin(async move { let resp = conn.call(req).await?; let (parts, body) = resp.into_parts(); @@ -274,6 +337,16 @@ where fn call(&mut self, req: http_1x::Request) -> Self::Future { let mut conn = self.conn.take().expect("H2Checkout::call called twice"); + // Populate the adapter-provided capture. The metadata is published + // by `H2ConnectAndHandshake` when a fresh H2 connection is + // established; if a request somehow reaches us before any metadata + // is available we leave the capture empty (poison becomes a no-op + // for that request, matching the capture-absent default). + if let Some(capture) = req.extensions().get::() { + if let Some(metadata) = conn.metadata() { + capture.set(metadata); + } + } Box::pin(async move { let resp = conn.call(req).await.map_err(Into::into)?; let (parts, body) = resp.into_parts(); @@ -290,19 +363,8 @@ where /// with different unnameable inner types, so `ConnectionPool::hosts` can't /// hold a concrete type. /// -/// The entry is request-granular — `send` performs checkout + health -/// filter + dispatch + body-guard setup atomically for one request. This -/// folds what was previously a separate `acquire(uri) -> Connection` -/// followed by `connection.send(req)` into one call. -/// -/// # Future: splitting acquire and send -/// -/// If we later need to expose a checked-out connection handle to callers -/// (e.g. connection warming, request pipelining on a single checkout, -/// or pool-external reuse), reintroduce a separate `Connection` trait -/// and split this into `acquire(uri) -> Box` plus -/// `Connection::send(self: Box, req) -> Response`. The current -/// shape collapses that split because no caller needs it today. +/// `send` performs checkout, health filtering, request dispatch, and +/// body-guard setup atomically for a single request. trait PoolEntry: Send + Sync { fn send( &mut self, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index c9b693eed11..88485ffecb5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -11,6 +11,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Instant; +use aws_smithy_runtime_api::client::connection::ConnectionMetadata; use aws_smithy_types::body::SdkBody; use pin_project_lite::pin_project; use tokio::sync::OwnedSemaphorePermit; @@ -66,28 +67,49 @@ pub(crate) struct ConnectionInfo { // Future: tls_info, timing (tcp_connect_duration, tls_handshake_duration) } +/// A one-shot "this connection is dead, don't reuse it" flag. +/// +/// Shared via `Arc` — all clones observe and control the same flag. +/// `ManagedConnection` holds one and hands out clones (via `metadata()`) +/// that let the adapter layer mark the connection poisoned through +/// smithy's `ConnectionMetadata::poison_fn`. On next checkout / return, +/// the pool sees the flag set and drops the connection instead of reusing +/// it. Modeled after hyper-util's legacy `PoisonPill`, but owned by us. +#[derive(Debug, Clone, Default)] +pub(crate) struct PoisonPill { + flag: Arc, +} + +impl PoisonPill { + /// Create a fresh, non-poisoned pill. + pub(crate) fn healthy() -> Self { + Self::default() + } + + /// Mark the connection as poisoned. + pub(crate) fn poison(&self) { + self.flag.store(true, Ordering::Release); + } + + /// Whether the connection has been poisoned. + pub(crate) fn is_poisoned(&self) -> bool { + self.flag.load(Ordering::Acquire) + } +} + /// A connection with SDK-owned lifecycle metadata. /// /// Wraps the inner service (typically `SendRequest`) with state needed -/// for pool management: idle tracking, poisoning, and connection identity. +/// for pool management: poisoning, connection identity, and permit lifetime. /// /// Clone is supported when the inner service is Clone (e.g., HTTP/2 multiplexed -/// connections). Clones share the same poison flag and connection metadata. -/// -/// # Missing: idle-start tracking -/// -/// A `last_used: Instant` updated on **return** to the pool is required -/// to drive `pool_idle_timeout`-based eviction. That requires a -/// re-insertion hook we do not have yet (hyper-util's `Cached::Drop` -/// gives us no callback). The hook shape will be designed alongside the -/// eviction mechanism. Until then this struct does not track idle -/// duration. `created_at` is retained because it is well-defined at -/// establishment time and needed for a future max-lifetime eviction. +/// connections). Clones share the same poison pill and connection info, so +/// poisoning one clone poisons all of them. pub(crate) struct ManagedConnection { inner: S, info: ConnectionInfo, created_at: Instant, - poisoned: Arc, + poison: PoisonPill, _permit: Arc, } @@ -98,38 +120,24 @@ impl ManagedConnection { inner, info, created_at: Instant::now(), - poisoned: Arc::new(AtomicBool::new(false)), + poison: PoisonPill::healthy(), _permit: permit, } } - /// Mark this connection as poisoned so it will not be reused. - /// - /// Poisoned connections are skipped on checkout and dropped on return to the pool. - pub(crate) fn poison(&self) { - self.poisoned.store(true, Ordering::Release); - } - - /// Returns a closure that poisons this connection when called. - /// - /// The returned closure can be passed to components that need to signal - /// connection failure without holding a direct reference to the connection. - pub(crate) fn poison_fn(&self) -> impl Fn() + Send + Sync + 'static { - let flag = self.poisoned.clone(); - move || flag.store(true, Ordering::Release) - } - /// Whether this connection has been poisoned. pub(crate) fn is_poisoned(&self) -> bool { - self.poisoned.load(Ordering::Acquire) + self.poison.is_poisoned() } - /// Connection metadata. + /// Connection info (remote/local addresses). + #[allow(dead_code)] // future telemetry / debugging pub(crate) fn info(&self) -> &ConnectionInfo { &self.info } /// When this connection was established. + #[allow(dead_code)] // future max-lifetime eviction pub(crate) fn created_at(&self) -> Instant { self.created_at } @@ -138,6 +146,27 @@ impl ManagedConnection { pub(crate) fn inner_mut(&mut self) -> &mut S { &mut self.inner } + + /// Build a smithy `ConnectionMetadata` for this connection. + /// + /// The returned metadata captures a clone of the `PoisonPill`, so + /// calling `ConnectionMetadata::poison()` flips this connection's + /// poison flag — the same flag the pool checks on checkout/return. + /// Address fields are copied. + pub(crate) fn metadata(&self) -> ConnectionMetadata { + let poison = self.poison.clone(); + let mut builder = ConnectionMetadata::builder() + // Hardcoded `false`: v2 does not yet plumb proxy configuration + // through the pool stack. When proxy support lands, + // `ConnectionInfo` should carry an `is_proxied` flag populated + // at connection establishment and passed through here. + .proxied(false) + .poison_fn(move || poison.poison()); + builder + .set_remote_addr(self.info.remote_addr) + .set_local_addr(self.info.local_addr); + builder.build() + } } impl Clone for ManagedConnection { @@ -146,7 +175,7 @@ impl Clone for ManagedConnection { inner: self.inner.clone(), info: self.info.clone(), created_at: self.created_at, - poisoned: self.poisoned.clone(), + poison: self.poison.clone(), _permit: self._permit.clone(), } } @@ -194,6 +223,19 @@ impl CachedConnection { inner: Some(cached), } } + + /// Build a smithy `ConnectionMetadata` for the underlying H1 connection. + /// + /// See [`ManagedConnection::metadata`]. + /// + /// Panics if called after the inner cached handle has been consumed. + pub(crate) fn metadata(&self) -> ConnectionMetadata { + self.inner + .as_ref() + .expect("CachedConnection metadata after drop") + .inner() + .metadata() + } } impl Service for CachedConnection @@ -224,6 +266,48 @@ impl Drop for CachedConnection { } } +/// Per-host shared slot that bridges H2 handshake completion to H2 checkout. +/// +/// `H2ConnectAndHandshake` calls [`Self::publish`] with a fresh +/// `ConnectionMetadata` every time a new H2 connection is established. +/// `SingletonConnection::metadata` reads the current value through +/// [`Self::current`]. +/// +/// Necessary because `hyper_util::client::pool::singleton::Singled<…>` is +/// opaque to us — we can't reach through it to the underlying +/// `ManagedConnection` to build its metadata at checkout time. The +/// handshake side has direct access, so it publishes and the checkout +/// side reads. +/// +/// Last-writer-wins is correct. Singleton holds at most one live H2 +/// connection per host at a time: on `Singled::poll_ready` error (or a +/// `Singleton::retain` predicate returning false), its state transitions +/// from `Made(svc)` back to `Empty`, dropping the old service. The next +/// request's `call` runs a fresh handshake through `H2ConnectAndHandshake`, +/// which publishes new metadata — overwriting whatever was there before. +/// Clients that grabbed metadata for the previous connection retain a +/// `ConnectionMetadata` pointing at the dropped connection's `PoisonPill`; +/// poisoning it is a no-op because the connection it described is already +/// gone. +#[derive(Clone, Default)] +pub(crate) struct H2ConnectionRef { + inner: Arc>>, +} + +impl H2ConnectionRef { + pub(crate) fn new() -> Self { + Self::default() + } + + pub(crate) fn publish(&self, metadata: ConnectionMetadata) { + *self.inner.lock().unwrap() = Some(metadata); + } + + pub(crate) fn current(&self) -> Option { + self.inner.lock().unwrap().clone() + } +} + /// A connection checked out from the H2 singleton. /// /// Symmetric wrapper to `CachedConnection` for the H2 leg. The inner type @@ -235,13 +319,26 @@ impl Drop for CachedConnection { /// No Drop hook: dropping a Singled clone does not affect singleton state; /// dead/poisoned H2 connections are handled reactively via /// `Singled::poll_ready` clearing the state to `Empty`. +/// +/// `h2_ref` is a clone of the same [`H2ConnectionRef`] that the handshake +/// service writes to — it lets `SingletonConnection::metadata` expose the +/// current H2 connection's metadata to the adapter even though `Singled<…>` +/// itself is opaque. pub(crate) struct SingletonConnection { inner: T, + h2_ref: H2ConnectionRef, } impl SingletonConnection { - pub(crate) fn new(inner: T) -> Self { - Self { inner } + pub(crate) fn new(inner: T, h2_ref: H2ConnectionRef) -> Self { + Self { inner, h2_ref } + } + + /// Metadata for the current H2 connection, if a handshake has completed. + /// + /// Returns `None` before the first successful H2 handshake for this host. + pub(crate) fn metadata(&self) -> Option { + self.h2_ref.current() } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index f2bf2cf000e..e5081be880d 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -259,11 +259,18 @@ where /// `Inspected` slot. pub(crate) struct H2ConnectAndHandshake { connector: C, + h2_ref: super::connection::H2ConnectionRef, } impl H2ConnectAndHandshake { - pub(crate) fn new(connector: C) -> Self { - Self { connector } + /// Create an H2 handshake service that publishes every newly established + /// connection's `ConnectionMetadata` into `h2_ref`. The same ref is held + /// on the read side by `SingletonConnection` (clones of the ref share + /// the underlying slot), so the H2 checkout path can expose connection + /// metadata and poison support to the adapter layer even though + /// `Singled<…>` itself is opaque. + pub(crate) fn new(connector: C, h2_ref: super::connection::H2ConnectionRef) -> Self { + Self { connector, h2_ref } } } @@ -271,6 +278,7 @@ impl Clone for H2ConnectAndHandshake { fn clone(&self) -> Self { Self { connector: self.connector.clone(), + h2_ref: self.h2_ref.clone(), } } } @@ -292,6 +300,7 @@ where fn call(&mut self, _req: ()) -> Self::Future { let fut = self.connector.call(()); + let h2_ref = self.h2_ref.clone(); Box::pin(async move { let (io, permit) = fut.await.map_err(Into::into)?; let info = capture_info(&io); @@ -317,7 +326,12 @@ where } }); - Ok(ManagedConnection::new(H2SendRequest::new(tx), info, permit)) + let managed = ManagedConnection::new(H2SendRequest::new(tx), info, permit); + // Publish this connection's metadata for the checkout side. + // Singleton replaces its stored connection wholesale on each new + // handshake, so last-writer-wins on the ref is correct. + h2_ref.publish(managed.metadata()); + Ok(managed) }) } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index f9c67730781..229141ebcc3 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -12,7 +12,6 @@ use std::sync::Arc; use std::time::Duration; use aws_smithy_runtime_api::client::connection::CaptureSmithyConnection; -use aws_smithy_runtime_api::client::connection::ConnectionMetadata; use aws_smithy_runtime_api::client::connector_metadata::ConnectorMetadata; use aws_smithy_runtime_api::client::http::{ HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpClient, @@ -244,17 +243,17 @@ impl HttpConnector for PooledConnector { // Save the full URI for pool routing before rewriting to origin-form. let full_uri = request.uri().clone(); - // Wire CaptureSmithyConnection with a no-op poison_fn so the interceptor - // pipeline doesn't panic. TODO: wire real poison from ManagedConnection. - if let Some(capture) = request.extensions().get::() { - capture.set_connection_retriever(|| { - Some( - ConnectionMetadata::builder() - .proxied(false) - .poison_fn(|| {}) - .build(), - ) - }); + // Create a request-scoped capture the pool will fill in with + // the `ConnectionMetadata` for the selected connection, and + // wire `CaptureSmithyConnection` to read from it. When the + // runtime's connection-poisoning interceptor later calls + // `ConnectionMetadata::poison()`, it flips the `PoisonPill` on + // the actual `ManagedConnection` selected for this request. + if let Some(capture_smithy) = request.extensions().get::() { + let capture = pool::ConnectionMetadataCapture::new(); + let for_retriever = capture.clone(); + capture_smithy.set_connection_retriever(move || for_retriever.get()); + request.extensions_mut().insert(capture); } // Set Host header from the URI authority if not already present. @@ -338,7 +337,7 @@ impl HttpClient for V2HttpClient { fn connector_metadata(&self) -> Option { // TODO: consider distinguishing v1 vs v2 clients and including TLS - // provider info (e.g. "rustls"/"0.23"). Revisit alongside 3f (TLS support). + // provider info (e.g. "rustls"/"0.23") in the connector metadata. Some(ConnectorMetadata::new("hyper", Some(Cow::Borrowed("1.x")))) } } diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 3a1bf1e00b2..8d989bf9a5f 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -882,3 +882,155 @@ async fn v2_unconsumed_body_holds_connection() { async fn v2_connection_close_header_prevents_reuse() { connection_close_header_prevents_reuse(&V2Client).await; } + +// --------------------------------------------------------------------------- +// Test implementations: connection poisoning +// --------------------------------------------------------------------------- + +/// Send a request and return both the response body and the `ConnectionMetadata` +/// captured by a `CaptureSmithyConnection` attached to the request. This is the +/// integration surface the `ConnectionPoisoningInterceptor` uses in real SDK flows +/// — the adapter must populate a retriever that returns live metadata pointing +/// at the connection selected for this request. +async fn send_with_capture( + client: &SharedHttpClient, + url: &str, +) -> ( + u16, + Vec, + Option, +) { + use aws_smithy_runtime_api::client::connection::CaptureSmithyConnection; + use http_body_util::BodyExt; + + let settings = HttpConnectorSettings::builder().build(); + let components = runtime_components(); + let connector = client.http_connector(&settings, &components); + + let capture = CaptureSmithyConnection::new(); + let mut request = HttpRequest::get(url).expect("valid HTTP request"); + request.add_extension(capture.clone()); + + let resp = connector + .call(request) + .await + .expect("request should succeed"); + let status = resp.status().as_u16(); + let body = resp + .into_body() + .collect() + .await + .expect("body should be readable") + .to_bytes() + .to_vec(); + (status, body, capture.get()) +} + +/// Poisoning a connection prevents it from being reused. +/// +/// Mirrors the production flow: `ConnectionPoisoningInterceptor` attaches a +/// `CaptureSmithyConnection` to the request, the adapter populates it with +/// metadata pointing at the selected connection, and on a transient error +/// the interceptor calls `ConnectionMetadata::poison()`. The next request +/// for the same host must establish a new TCP connection instead of reusing +/// the poisoned one. +async fn poisoned_connection_not_reused(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"first", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"second", + }, + ], + ) + .build() + .await; + + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + let client = make.make(ClientConfig::default()); + + // First request: establish connection, capture metadata. + let (status, body, metadata) = send_with_capture(&client, &url).await; + assert_eq!(status, 200); + assert_eq!(body, b"first"); + let metadata = metadata.expect("adapter should populate CaptureSmithyConnection"); + assert_eq!( + harness.tcp_accepted_count(), + 1, + "first request opens one connection" + ); + + // Poison the connection — what the orchestrator does on a transient error. + metadata.poison(); + + // Next request to the same host must open a NEW connection; the poisoned + // one is skipped on checkout and dropped on return. + let (status, body, _) = send_with_capture(&client, &url).await; + assert_eq!(status, 200); + assert_eq!(body, b"second"); + assert_eq!( + harness.tcp_accepted_count(), + 2, + "poisoned connection must not be reused" + ); +} + +/// When no transient error occurs, the captured metadata is handed out but +/// never poisoned — the connection returns to the pool and is reused. +/// This is the non-poison control for `poisoned_connection_not_reused`. +async fn capture_without_poison_allows_reuse(make: &dyn MakeClient) { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"first", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"second", + }, + ], + ) + .build() + .await; + + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + let client = make.make(ClientConfig::default()); + + let (status, body, metadata) = send_with_capture(&client, &url).await; + assert_eq!(status, 200); + assert_eq!(body, b"first"); + assert!( + metadata.is_some(), + "adapter should populate CaptureSmithyConnection" + ); + // Deliberately not calling poison(). + drop(metadata); + + let (status, body, _) = send_with_capture(&client, &url).await; + assert_eq!(status, 200); + assert_eq!(body, b"second"); + assert_eq!( + harness.tcp_accepted_count(), + 1, + "without poison the connection should be reused" + ); +} + +#[tokio::test] +async fn v2_poisoned_connection_not_reused() { + poisoned_connection_not_reused(&V2Client).await; +} + +#[tokio::test] +async fn v2_capture_without_poison_allows_reuse() { + capture_without_poison_allows_reuse(&V2Client).await; +} From afecb70972065759519754f3c8dde68fa61f3359 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 6 May 2026 00:31:52 -0400 Subject: [PATCH 14/47] thread connection context through stack --- .../aws-smithy-http-client/src/client/pool.rs | 44 +++++++-- .../src/client/pool/connection.rs | 57 ++++++++++- .../src/client/pool/handshake.rs | 95 +++++++++++++++++-- .../src/client/timeout.rs | 46 +++++++++ .../aws-smithy-http-client/src/client/v2.rs | 55 ++++++++++- .../tests/pool_behavior_test.rs | 57 +++++++++++ 6 files changed, 330 insertions(+), 24 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 85a1f403f2a..b0e268ceb38 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -44,6 +44,7 @@ use connection::{ CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, H2ConnectionRef, SingletonConnection, }; +pub(crate) use connection::{ConnectCtx, ReadTimeoutHint, TimeoutContext}; use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake}; pub(crate) type BoxError = Box; @@ -209,12 +210,17 @@ impl ConnectionPool { /// Send a request through the pool. /// /// Routes to the appropriate per-host pool stack and sends the request. + /// `ctx` carries the routing URI plus per-operation connect-time data + /// (connect_timeout). Per-operation read_timeout, if any, must be + /// attached to `req.extensions_mut()` as [`ReadTimeoutHint`] before + /// calling — the checkout services read it from there. pub(crate) async fn send_request( &self, - uri: http_1x::Uri, + ctx: ConnectCtx, req: http_1x::Request, ) -> Result, BoxError> { - let key = PoolKey::from_uri(&uri).ok_or("request URI must have scheme and authority")?; + let key = + PoolKey::from_uri(&ctx.uri).ok_or("request URI must have scheme and authority")?; // Dispatch the request through the per-host entry. The lock is // held only long enough to look up / create the entry and call @@ -223,9 +229,9 @@ impl ConnectionPool { let fut = { let mut hosts = self.hosts.lock().unwrap(); if !hosts.contains_key(&key) { - hosts.insert(key.clone(), (self.make_entry)(&uri)); + hosts.insert(key.clone(), (self.make_entry)(&ctx.uri)); } - hosts.get_mut(&key).unwrap().send(uri, req) + hosts.get_mut(&key).unwrap().send(ctx, req) }; fut.await @@ -285,8 +291,17 @@ impl Service> for H1Checkout() { capture.set(conn.metadata()); } + // Read-timeout hint: bounds request-write + response-headers only. + let read_timeout = req.extensions().get::().cloned(); Box::pin(async move { - let resp = conn.call(req).await?; + let send_fut = conn.call(req); + let resp = super::timeout::maybe_timeout_future( + send_fut, + read_timeout.as_ref().map(|h| h.0.duration), + read_timeout.as_ref().map(|h| &h.0.sleep_impl), + super::timeout::TimeoutKind::Read, + ) + .await?; let (parts, body) = resp.into_parts(); let body = GuardedBody::new(body, ConnectionGuard::H1(conn)); Ok(http_1x::Response::from_parts(parts, body)) @@ -347,8 +362,17 @@ where capture.set(metadata); } } + // Read-timeout hint: bounds request-write + response-headers only. + let read_timeout = req.extensions().get::().cloned(); Box::pin(async move { - let resp = conn.call(req).await.map_err(Into::into)?; + let send_fut = conn.call(req); + let resp = super::timeout::maybe_timeout_future( + send_fut, + read_timeout.as_ref().map(|h| h.0.duration), + read_timeout.as_ref().map(|h| &h.0.sleep_impl), + super::timeout::TimeoutKind::Read, + ) + .await?; let (parts, body) = resp.into_parts(); let body = GuardedBody::new(body, ConnectionGuard::H2(conn)); Ok(http_1x::Response::from_parts(parts, body)) @@ -368,7 +392,7 @@ where trait PoolEntry: Send + Sync { fn send( &mut self, - uri: http_1x::Uri, + ctx: ConnectCtx, req: http_1x::Request, ) -> BoxFuture, BoxError>>; } @@ -382,7 +406,7 @@ struct TypedPoolEntry(S); impl PoolEntry for TypedPoolEntry where - S: Service + Clone + Send + Sync + 'static, + S: Service + Clone + Send + Sync + 'static, S::Error: Into + 'static, S::Future: Send + 'static, Conn: Service, Response = CheckoutResponse> @@ -394,7 +418,7 @@ where { fn send( &mut self, - uri: http_1x::Uri, + ctx: ConnectCtx, req: http_1x::Request, ) -> BoxFuture, BoxError>> { let mut svc = self.0.clone(); @@ -412,7 +436,7 @@ where std::future::poll_fn(|cx| svc.poll_ready(cx)) .await .map_err(Into::into)?; - let mut checkout = svc.call(uri.clone()).await.map_err(Into::into)?; + let mut checkout = svc.call(ctx.clone()).await.map_err(Into::into)?; if std::future::poll_fn(|cx| checkout.poll_ready(cx)) .await .is_ok() diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 88485ffecb5..acb33db4ac9 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -9,8 +9,9 @@ use std::net::SocketAddr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::task::{Context, Poll}; -use std::time::Instant; +use std::time::{Duration, Instant}; +use aws_smithy_async::rt::sleep::SharedAsyncSleep; use aws_smithy_runtime_api::client::connection::ConnectionMetadata; use aws_smithy_types::body::SdkBody; use pin_project_lite::pin_project; @@ -21,6 +22,60 @@ use super::cache; use super::handshake::H1SendRequest; use super::BoxError; +/// A duration paired with the sleep implementation used to realize it. +/// +/// This type makes "timeout without sleep impl" unrepresentable — you cannot +/// construct one without committing to a way to actually wait. Created at the +/// adapter layer from `HttpConnectorSettings` + `RuntimeComponents::sleep_impl()` +/// and passed down the pool stack where timeouts are applied. +#[derive(Clone, Debug)] +pub(crate) struct TimeoutContext { + pub(crate) duration: Duration, + pub(crate) sleep_impl: SharedAsyncSleep, +} + +impl TimeoutContext { + pub(crate) fn new(duration: Duration, sleep_impl: SharedAsyncSleep) -> Self { + Self { + duration, + sleep_impl, + } + } +} + +/// Target type for the connect portion of the pool stack. +/// +/// Replaces bare `Uri` so per-operation connect metadata flows through the +/// composable pool types (Map → Negotiate → Cache → handshake → ConnectionLimit +/// → TCP connector) to the layers that need them. +#[derive(Clone, Debug)] +pub(crate) struct ConnectCtx { + pub(crate) uri: http_1x::Uri, + /// Bounds new-connection establishment (TCP + TLS handshake). `None` + /// means no connect timeout — cached connections skip the connector + /// entirely so this is automatically a no-op on cache hit. + pub(crate) connect_timeout: Option, +} + +impl ConnectCtx { + pub(crate) fn new(uri: http_1x::Uri, connect_timeout: Option) -> Self { + Self { + uri, + connect_timeout, + } + } +} + +/// Request extension set by the adapter to hint a read timeout to the +/// checkout services (`H{1,2}Checkout`). +/// +/// `Some(...)` means: once the connection is selected (cache hit or fresh +/// handshake), wrap `conn.call(req)` with this timeout. Bounds request-write +/// + response-headers-wait only — does NOT include pool acquire or connect +/// establishment (those have their own timeouts). +#[derive(Clone, Debug)] +pub(crate) struct ReadTimeoutHint(pub(crate) TimeoutContext); + /// Permits acquired from connection limit semaphores. /// Held for the lifetime of the connection — dropped when the connection is dropped. pub(crate) struct ConnectionPermit { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index e5081be880d..b4dd4ba172c 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -22,13 +22,20 @@ use hyper_util::rt::TokioExecutor; use tokio::sync::Semaphore; use tower::Service; -use super::connection::{ConnectionInfo, ConnectionPermit, ManagedConnection}; +use super::connection::{ConnectCtx, ConnectionInfo, ConnectionPermit, ManagedConnection}; use super::BoxError; /// Wraps a connector service, acquiring semaphore permits before connecting. /// /// Returns `(IO, Arc)` so the permit can be stored on /// `ManagedConnection` and held for the connection's lifetime. +/// +/// Target type is `ConnectCtx`: the inner TCP connector is `Service`, +/// so this layer extracts `ctx.uri` to pass through. Per-operation timeouts +/// on the context are applied here — `connect_timeout` wraps the inner +/// connector call (which includes TCP + TLS since the inner is the +/// TLS-wrapped connector). Cache hits skip this layer entirely, so +/// `connect_timeout` is automatically new-connection-only. pub(crate) struct ConnectionLimit { inner: C, global: Option>, @@ -59,13 +66,12 @@ impl Clone for ConnectionLimit { } } -impl Service for ConnectionLimit +impl Service for ConnectionLimit where - C: Service + Clone + Send + 'static, + C: Service + Clone + Send + 'static, C::Error: Into + 'static, C::Future: Send + 'static, IO: Send + 'static, - Req: Send + 'static, { type Response = (IO, Arc); type Error = BoxError; @@ -75,7 +81,7 @@ where self.inner.poll_ready(cx).map_err(Into::into) } - fn call(&mut self, req: Req) -> Self::Future { + fn call(&mut self, ctx: ConnectCtx) -> Self::Future { let mut inner = self.inner.clone(); let global = self.global.clone(); let per_host = self.per_host.clone(); @@ -97,7 +103,19 @@ where std::future::poll_fn(|cx| inner.poll_ready(cx)) .await .map_err(Into::into)?; - let io = inner.call(req).await.map_err(Into::into)?; + + // Apply connect_timeout only around the actual connector call + // (TCP + TLS). If `ctx.connect_timeout` is `None`, this is a + // plain `inner.call(uri).await`. + let uri = ctx.uri; + let connect_fut = inner.call(uri); + let io = super::super::timeout::maybe_timeout_future( + connect_fut, + ctx.connect_timeout.as_ref().map(|t| t.duration), + ctx.connect_timeout.as_ref().map(|t| &t.sleep_impl), + super::super::timeout::TimeoutKind::Connect, + ) + .await?; Ok((io, permit)) }) } @@ -205,9 +223,9 @@ impl Clone for H1ConnectAndHandshake { } } -impl Service for H1ConnectAndHandshake +impl Service for H1ConnectAndHandshake where - C: Service)>, + C: Service)>, C::Error: Into + 'static, C::Future: Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, @@ -220,8 +238,8 @@ where self.connector.poll_ready(cx).map_err(Into::into) } - fn call(&mut self, uri: http_1x::Uri) -> Self::Future { - let fut = self.connector.call(uri); + fn call(&mut self, ctx: ConnectCtx) -> Self::Future { + let fut = self.connector.call(ctx); Box::pin(async move { let (io, permit) = fut.await.map_err(Into::into)?; let info = capture_info(&io); @@ -335,3 +353,60 @@ where }) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::timeout::test::NeverConnects; + use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep}; + use std::time::Duration; + use tower::Service as _; + + /// Verify `ConnectionLimit` applies `connect_timeout` from `ConnectCtx` + /// to the inner TCP connector. `NeverConnects` returns a connector + /// future that never resolves; a short connect_timeout should fire + /// and produce an `HTTP connect timeout occurred after …` error. + #[tokio::test(start_paused = true)] + async fn connect_timeout_fires_on_slow_connector() { + let mut svc = ConnectionLimit::new(NeverConnects::default(), None, None); + let sleep = SharedAsyncSleep::new(TokioSleep::new()); + let ctx = ConnectCtx::new( + "http://example.com".parse().unwrap(), + Some(super::super::connection::TimeoutContext::new( + Duration::from_millis(500), + sleep, + )), + ); + let err = match svc.call(ctx).await { + Ok(_) => panic!("connect timeout should fire against a never-resolving connector"), + Err(err) => err, + }; + let msg = format!("{err}"); + assert!( + msg.contains("HTTP connect"), + "expected `HTTP connect` in error, got: {msg}" + ); + assert!( + msg.contains("500ms"), + "expected `500ms` in error, got: {msg}" + ); + } + + /// Without a connect_timeout set, `ConnectionLimit` passes through the + /// connector call with no wrapping — verify the absence of a timeout + /// means the slow connector stays pending (we just check we can start + /// the call; with `start_paused`, nothing advances so the future is + /// still pending after a short yield). + #[tokio::test(start_paused = true)] + async fn no_timeout_does_not_bound_connector() { + let mut svc = ConnectionLimit::new(NeverConnects::default(), None, None); + let ctx = ConnectCtx::new("http://example.com".parse().unwrap(), None); + let fut = svc.call(ctx); + // A brief tokio yield shouldn't resolve the never-connects future. + tokio::pin!(fut); + tokio::select! { + _ = &mut fut => panic!("future should not resolve without a timeout"), + _ = tokio::task::yield_now() => {} + } + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/timeout.rs b/rust-runtime/aws-smithy-http-client/src/client/timeout.rs index 56c3f153d55..67f00df95fc 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/timeout.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/timeout.rs @@ -205,6 +205,52 @@ where } } +/// Which timeout label a wrapping future should carry. +/// +/// The label surfaces in the `HttpTimeoutError` produced when the timeout +/// fires; it is what users see in the error message and what the retry +/// classifier receives via `TimedOutError` downcasting. +#[derive(Clone, Copy, Debug)] +pub(crate) enum TimeoutKind { + Connect, + Read, +} + +impl TimeoutKind { + fn label(self) -> &'static str { + match self { + TimeoutKind::Connect => "HTTP connect", + TimeoutKind::Read => "HTTP read", + } + } +} + +/// Wrap `fut` in a `MaybeTimeoutFuture` if `timeout` is set. +/// +/// Used by the v2 pool layers to apply per-operation `connect_timeout` / +/// `read_timeout` without building Tower services for each. Passes +/// `sleep_impl` by reference so the helper can be used without cloning +/// the sleep impl into a dedicated service. +pub(crate) fn maybe_timeout_future( + fut: F, + timeout: Option, + sleep_impl: Option<&SharedAsyncSleep>, + kind: TimeoutKind, +) -> MaybeTimeoutFuture +where + F: Future>, + E: Into, +{ + match (timeout, sleep_impl) { + (Some(duration), Some(sleep)) => MaybeTimeoutFuture::Timeout { + timeout: Timeout::new(fut, sleep.sleep(duration)), + error_type: kind.label(), + duration, + }, + _ => MaybeTimeoutFuture::NoTimeout { future: fut }, + } +} + #[cfg(test)] pub(crate) mod test { use hyper::rt::ReadBufCursor; diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index 229141ebcc3..2af084e9348 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -222,8 +222,18 @@ impl BuilderV2 { // --------------------------------------------------------------------------- /// Smithy `HttpConnector` backed by the v2 connection pool. +/// +/// Constructed fresh per `HttpClient::http_connector` call so it can capture +/// the per-operation `HttpConnectorSettings` (connect/read timeouts). The pool +/// itself is shared across all operations via `Arc`. struct PooledConnector { pool: Arc, + connect_timeout: Option, + read_timeout: Option, + /// Present when any timeout is configured. Required to construct + /// `Timeout` futures. `V2HttpClient::http_connector` panics if + /// a timeout is set but `sleep_impl` is absent — matches v1. + sleep_impl: Option, } impl std::fmt::Debug for PooledConnector { @@ -235,6 +245,9 @@ impl std::fmt::Debug for PooledConnector { impl HttpConnector for PooledConnector { fn call(&self, request: HttpRequest) -> HttpConnectorFuture { let pool = self.pool.clone(); + let connect_timeout = self.connect_timeout; + let read_timeout = self.read_timeout; + let sleep_impl = self.sleep_impl.clone(); HttpConnectorFuture::new(async move { let mut request = request .try_into_http1x() @@ -256,6 +269,18 @@ impl HttpConnector for PooledConnector { request.extensions_mut().insert(capture); } + // Attach the per-op read timeout hint (if any) as a request + // extension. `H{1,2}Checkout::call` reads it and wraps the + // request dispatch — bounds request-write + response-headers, + // nothing else. + if let Some((duration, sleep)) = read_timeout.zip(sleep_impl.clone()) { + request + .extensions_mut() + .insert(pool::ReadTimeoutHint(pool::TimeoutContext::new( + duration, sleep, + ))); + } + // Set Host header from the URI authority if not already present. if !request.headers().contains_key(http_1x::header::HOST) { if let Some(authority) = full_uri.authority() { @@ -282,8 +307,18 @@ impl HttpConnector for PooledConnector { }; *request.uri_mut() = origin; + // Build the connect context: routing URI + connect timeout. + // Cache hits skip the connector entirely so connect_timeout is + // automatically new-connection-only. + let connect_ctx = pool::ConnectCtx::new( + full_uri, + connect_timeout + .zip(sleep_impl) + .map(|(d, s)| pool::TimeoutContext::new(d, s)), + ); + let response = pool - .send_request(full_uri, request) + .send_request(connect_ctx, request) .await .map_err(downcast_error)?; @@ -327,11 +362,25 @@ impl V2HttpClient { impl HttpClient for V2HttpClient { fn http_connector( &self, - _settings: &HttpConnectorSettings, - _components: &RuntimeComponents, + settings: &HttpConnectorSettings, + components: &RuntimeComponents, ) -> SharedHttpConnector { + let connect_timeout = settings.connect_timeout(); + let read_timeout = settings.read_timeout(); + let sleep_impl = components.sleep_impl(); + + if (connect_timeout.is_some() || read_timeout.is_some()) && sleep_impl.is_none() { + panic!( + "an async sleep impl is required to use connect/read timeouts with \ + the v2 HTTP client; provide one via `RuntimeComponents::sleep_impl`" + ); + } + SharedHttpConnector::new(PooledConnector { pool: self.pool.clone(), + connect_timeout, + read_timeout, + sleep_impl, }) } diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 8d989bf9a5f..7c8fced8527 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1034,3 +1034,60 @@ async fn v2_poisoned_connection_not_reused() { async fn v2_capture_without_poison_allows_reuse() { capture_without_poison_allows_reuse(&V2Client).await; } + +// --------------------------------------------------------------------------- +// Test implementations: per-operation timeouts +// --------------------------------------------------------------------------- + +/// Read timeout fires when the server accepts TCP but never sends a response. +/// +/// Exercises the v2 adapter's per-op timeout wrapping end-to-end: the +/// `HttpConnectorSettings::read_timeout` flows through +/// `HttpClient::http_connector` into `PooledConnector`, wraps +/// `pool.send_request`, fires because `HoldThenClose` never replies, and +/// produces a `ConnectorError::timeout` classified by `downcast_error`. +async fn v2_read_timeout_fires_on_silent_server() { + use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep}; + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::HoldThenClose(Duration::from_secs(30))], + ) + .build() + .await; + + let client = V2Client.make(ClientConfig::default()); + let components = + aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(SystemTimeSource::new())) + .with_sleep_impl(Some(SharedAsyncSleep::new(TokioSleep::new()))) + .build() + .expect("valid runtime components"); + let settings = HttpConnectorSettings::builder() + .read_timeout(Duration::from_millis(200)) + .build(); + let connector = client.http_connector(&settings, &components); + + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + let start = std::time::Instant::now(); + let err = connector + .call(HttpRequest::get(&url).expect("valid HTTP request")) + .await + .expect_err("read timeout should fire against a non-responsive server"); + let elapsed = start.elapsed(); + + assert!( + err.is_timeout(), + "expected timeout classification, got {err:?}" + ); + assert!( + elapsed < Duration::from_secs(2), + "read timeout did not fire in time (took {elapsed:?})" + ); +} + +#[tokio::test] +async fn v2_read_timeout() { + v2_read_timeout_fires_on_silent_server().await; +} From 29ad54349bf5aa27a65d713e4890982d7e916ac2 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 13 May 2026 22:20:56 -0400 Subject: [PATCH 15/47] thread pool hooks through to handshake --- rust-runtime/Cargo.lock | 1 + .../aws-smithy-http-client/Cargo.toml | 2 + .../aws-smithy-http-client/src/client/pool.rs | 527 ++++++++++++++++-- .../src/client/pool/connection.rs | 407 ++++++++++++-- .../src/client/pool/handshake.rs | 90 ++- .../aws-smithy-http-client/src/client/v2.rs | 23 + .../tests/pool_behavior_test.rs | 139 +++++ .../tests/reconnect_on_transient_error.rs | 18 +- 8 files changed, 1113 insertions(+), 94 deletions(-) diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock index 06e4d56239a..1e08424eec8 100644 --- a/rust-runtime/Cargo.lock +++ b/rust-runtime/Cargo.lock @@ -539,6 +539,7 @@ dependencies = [ "tower-service", "tower-test", "tracing", + "tracing-subscriber", ] [[package]] diff --git a/rust-runtime/aws-smithy-http-client/Cargo.toml b/rust-runtime/aws-smithy-http-client/Cargo.toml index 45748cb0890..3f9c726ab6f 100644 --- a/rust-runtime/aws-smithy-http-client/Cargo.toml +++ b/rust-runtime/aws-smithy-http-client/Cargo.toml @@ -139,6 +139,8 @@ tokio = { version = "1.49.0", features = ["macros", "rt", "rt-multi-thread", "te tokio-rustls = "0.26.2" # used by tests vendored from hyper-util in src/client/pool/vendored_cache.rs tower-test = "0.4" +# Test-only ANSI-free subscriber for capturing pool tracing events in assertions. +tracing-subscriber = { version = "0.3.19", features = ["env-filter", "fmt"] } [[example]] diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index b0e268ceb38..9ae9965e9e9 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -27,17 +27,21 @@ mod cache { } use std::collections::HashMap; +use std::convert::Infallible; use std::future::Future; use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::sync::Mutex; +use std::sync::OnceLock; use std::task::Poll; +use std::time::{Duration, Instant}; use aws_smithy_types::body::SdkBody; use hyper_util::client::legacy::connect::Connection as HyperConnection; use hyper_util::client::pool as hpool; use hyper_util::rt::TokioExecutor; -use tokio::sync::Semaphore; +use tokio::sync::{oneshot, Semaphore}; use tower::{Service, ServiceExt}; use connection::{ @@ -133,7 +137,37 @@ impl PoolKey { } } -/// Build a connection pool for the given connector and configuration. +/// Type-erased handle to one Negotiate leg's eviction interface. +/// +/// Each `TypedPoolEntry` holds at most two of these — one for its H1 cache, +/// one for its H2 singleton. The two underlying types (`Cache<…>` and +/// `Singleton<…>`) carry upstream-unnameable parameters from `Negotiate`, +/// so they're captured inside `Box` closures. The closures are +/// `Fn` (so multiple callers can invoke through `Arc`); interior +/// mutability lives in a `std::sync::Mutex` over a clone of the leg, which +/// is uncontended in practice — retain runs at most once per +/// `max(pool_idle_timeout, MIN_EVICTION_TICK)`, and the leg's own state +/// already serializes through its internal `Arc>`. +pub(crate) struct BoxedRetainer { + retain_fn: Box, + is_empty_fn: Box bool + Send + Sync + 'static>, +} + +impl BoxedRetainer { + fn retain_idle(&self, timeout: Duration) { + (self.retain_fn)(timeout) + } + + fn is_empty(&self) -> bool { + (self.is_empty_fn)() + } +} + +/// Per-host retainer registry. Bounded at two entries (one H1 cache, one +/// H2 singleton); populated lazily by the H1 fallback and H2 upgrade +/// `layer_fn`s the first time `Negotiate` constructs each leg. +type RetainerSlot = Arc>>; + pub(crate) fn build_pool(connector: C, config: PoolConfig) -> ConnectionPool where C: Service + Clone + Send + Sync + 'static, @@ -143,53 +177,179 @@ where { let global_sem = config.max_connections.map(|n| Arc::new(Semaphore::new(n))); let max_per_host = config.max_connections_per_host; - - let make_entry = move |_uri: &http_1x::Uri| -> Box { - let per_host_sem = max_per_host.map(|n| Arc::new(Semaphore::new(n))); - let limited = - handshake::ConnectionLimit::new(connector.clone(), global_sem.clone(), per_host_sem); - - // Per-host bridge: `H2ConnectAndHandshake` publishes on each new - // handshake; `SingletonConnection` reads the current entry at - // checkout. Clones share the underlying slot. - let h2_ref = H2ConnectionRef::new(); - - let stack = hpool::negotiate::builder() - .connect(limited) - .inspect( - |(conn, _permit): &(IO, Arc)| { - conn.connected().is_negotiated_h2() - }, - ) - .fallback(tower::layer::layer_fn(|inspector| { - cache::builder() - .executor(TokioExecutor::new()) - .build(H1ConnectAndHandshake::new(inspector)) - .map_response(|cached| H1Checkout::new(CachedConnection::new(cached))) - })) - .upgrade({ - let h2_ref = h2_ref.clone(); - tower::layer::layer_fn(move |inspected| { - hpool::singleton::Singleton::new(H2ConnectAndHandshake::new( - inspected, - h2_ref.clone(), - )) - .map_response({ - let h2_ref = h2_ref.clone(); - move |singled| { - H2Checkout::new(SingletonConnection::new(singled, h2_ref.clone())) - } + let pool_hooks = handshake::PoolHooks::new(); + + tracing::debug!( + max_connections = ?config.max_connections, + max_connections_per_host = ?config.max_connections_per_host, + pool_idle_timeout = ?config.pool_idle_timeout, + "v2 pool: initialized" + ); + + let make_entry = { + let pool_hooks = pool_hooks.clone(); + move |_uri: &http_1x::Uri| -> Box { + let per_host_sem = max_per_host.map(|n| Arc::new(Semaphore::new(n))); + let limited = handshake::ConnectionLimit::new( + connector.clone(), + global_sem.clone(), + per_host_sem, + ); + + // Per-host bridge: `H2ConnectAndHandshake` publishes on each new + // handshake; `SingletonConnection` reads the current entry at + // checkout. Clones share the underlying slot. + let h2_ref = H2ConnectionRef::new(); + + // Bounded at two — the H1 fallback and H2 upgrade `layer_fn`s + // each push one entry on first construction. + let retainers: RetainerSlot = Arc::new(Mutex::new(Vec::with_capacity(2))); + + let stack = hpool::negotiate::builder() + .connect(limited) + .inspect( + |(conn, _permit): &(IO, Arc)| { + conn.connected().is_negotiated_h2() + }, + ) + .fallback({ + let retainers = retainers.clone(); + let pool_hooks = pool_hooks.clone(); + tower::layer::layer_fn(move |inspector| { + let cache = cache::builder() + .executor(TokioExecutor::new()) + .build(H1ConnectAndHandshake::new(inspector, pool_hooks.clone())); + // Capture a clone of the Cache for eviction. `Cache` + // is Clone and shares state via its internal + // `Arc>`, so the clone here and the + // original-stack consumption below observe the same + // idle set. + let cache_for_retain = Arc::new(std::sync::Mutex::new(cache.clone())); + let cache_for_empty = cache_for_retain.clone(); + retainers + .lock() + .expect("retainer slot poisoned") + .push(BoxedRetainer { + retain_fn: Box::new(move |timeout| { + let now = Instant::now(); + cache_for_retain + .lock() + .expect("retain cache lock poisoned") + .retain(|managed| { + let keep = !managed.is_poisoned() + && now.saturating_duration_since(managed.idle_at()) + < timeout; + if !keep { + let reason = if managed.is_poisoned() { + "poisoned" + } else { + "idle_expired" + }; + let idle_duration = now + .saturating_duration_since(managed.idle_at()); + tracing::debug!( + conn_id = managed.conn_id(), + reason, + ?idle_duration, + "v2 pool: connection evicted" + ); + } + keep + }); + }), + is_empty_fn: Box::new(move || { + cache_for_empty + .lock() + .expect("is_empty cache lock poisoned") + .is_empty() + }), + }); + cache.map_response(|cached| H1Checkout::new(CachedConnection::new(cached))) + }) + }) + .upgrade({ + let h2_ref = h2_ref.clone(); + let retainers = retainers.clone(); + let pool_hooks = pool_hooks.clone(); + tower::layer::layer_fn(move |inspected| { + let singleton = + hpool::singleton::Singleton::new(H2ConnectAndHandshake::new( + inspected, + h2_ref.clone(), + pool_hooks.clone(), + )); + let singleton_for_retain = + Arc::new(std::sync::Mutex::new(singleton.clone())); + let singleton_for_empty = singleton_for_retain.clone(); + retainers + .lock() + .expect("retainer slot poisoned") + .push(BoxedRetainer { + retain_fn: Box::new(move |timeout| { + let now = Instant::now(); + singleton_for_retain + .lock() + .expect("retain singleton lock poisoned") + .retain(|managed| { + // For H2 the connection stays in + // Singleton while serving streams, so + // `idle_at` alone is not a valid + // idleness signal. Keep the + // connection if any stream is in + // flight; only evict when truly idle + // AND `idle_at` has exceeded the + // timeout (stamped on the 1 → 0 + // transition by + // `SingletonConnection::drop`). + let keep = !managed.is_poisoned() + && (managed.active_streams_count() > 0 + || now.saturating_duration_since( + managed.idle_at(), + ) < timeout); + if !keep { + let reason = if managed.is_poisoned() { + "poisoned" + } else { + "idle_expired" + }; + let idle_duration = now + .saturating_duration_since(managed.idle_at()); + tracing::debug!( + conn_id = managed.conn_id(), + reason, + ?idle_duration, + "v2 pool: connection evicted" + ); + } + keep + }); + }), + is_empty_fn: Box::new(move || { + singleton_for_empty + .lock() + .expect("is_empty singleton lock poisoned") + .is_empty() + }), + }); + singleton.map_response({ + let h2_ref = h2_ref.clone(); + move |singled| { + H2Checkout::new(SingletonConnection::new(singled, h2_ref.clone())) + } + }) }) }) - }) - .build(); - Box::new(TypedPoolEntry(stack)) + .build(); + Box::new(TypedPoolEntry { stack, retainers }) + } }; ConnectionPool { config, hosts: Mutex::new(HashMap::new()), make_entry: Box::new(make_entry), + eviction_spawned: AtomicBool::new(false), + drop_notifier: OnceLock::new(), } } @@ -200,12 +360,31 @@ where /// and HTTP/2 (Singleton) based on ALPN negotiation. pub(crate) struct ConnectionPool { /// Pool-wide configuration. - #[allow(dead_code)] config: PoolConfig, hosts: Mutex>>, make_entry: Box Box + Send + Sync>, + + /// Latches `true` once the eviction task has been spawned. Subsequent + /// `send_request` calls observe the latch and skip the spawn path. + /// Read only when `config.pool_idle_timeout` is set. + eviction_spawned: AtomicBool, + + /// Drop signal for the eviction task. The task awaits the matching + /// `Receiver`; we never send. When `ConnectionPool` drops, this + /// `OnceLock`'s contents drop, the sender drops, the receiver errors + /// with `Canceled`, and the task exits its `select!` loop. + /// + /// `OnceLock` because the task spawns lazily on first `send_request` + /// — at most one task per pool. The task additionally holds a + /// `Weak` as a fallback exit signal. + drop_notifier: OnceLock>, } +/// Minimum eviction tick period, matching legacy hyper-util pool's +/// `MIN_CHECK` constant. Prevents very short `pool_idle_timeout`s from +/// causing the task to spin hot. +const MIN_EVICTION_TICK: Duration = Duration::from_millis(90); + impl ConnectionPool { /// Send a request through the pool. /// @@ -214,14 +393,24 @@ impl ConnectionPool { /// (connect_timeout). Per-operation read_timeout, if any, must be /// attached to `req.extensions_mut()` as [`ReadTimeoutHint`] before /// calling — the checkout services read it from there. + /// + /// Takes `self: &Arc` so the lazy eviction task can hold a + /// `Weak` as a fallback exit signal (primary exit is the drop + /// of `drop_notifier`; the `Weak` just insulates the task against a + /// missed drop signal). pub(crate) async fn send_request( - &self, + self: &Arc, ctx: ConnectCtx, req: http_1x::Request, ) -> Result, BoxError> { let key = PoolKey::from_uri(&ctx.uri).ok_or("request URI must have scheme and authority")?; + // Lazily spawn the idle-eviction task on first use. Matches the + // legacy hyper-util pool's approach: no task if the pool is never + // used, no task if `pool_idle_timeout` is `None` or zero. + self.maybe_spawn_eviction_task(); + // Dispatch the request through the per-host entry. The lock is // held only long enough to look up / create the entry and call // `send` — all I/O happens in the returned future, outside the @@ -236,6 +425,123 @@ impl ConnectionPool { fut.await } + + /// Spawn the idle-eviction task if it hasn't been spawned yet and + /// `pool_idle_timeout` is configured. + /// + /// Idempotent and cheap after the first call (single relaxed + /// `AtomicBool` load returns early). Matches legacy + /// `spawn_idle_interval` semantics: no-op without a timeout, no-op + /// with a zero timeout, no-op if already spawned. + fn maybe_spawn_eviction_task(self: &Arc) { + let timeout = match self.config.pool_idle_timeout { + Some(d) if d > Duration::ZERO => d, + _ => return, + }; + // Fast path: already spawned. + if self.eviction_spawned.load(Ordering::Acquire) { + return; + } + // Claim the spawn slot. Only one task per pool ever. + if self + .eviction_spawned + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return; + } + let (tx, rx) = oneshot::channel::(); + // Invariant: this code path runs exactly once per pool — the + // `compare_exchange` on `eviction_spawned` above is the unique + // claim. If `drop_notifier` is already populated, that invariant + // is broken. + self.drop_notifier + .set(tx) + .expect("drop_notifier set after exclusive spawn-slot claim"); + let weak = Arc::downgrade(self); + let tick = timeout.max(MIN_EVICTION_TICK); + tracing::debug!( + interval = ?tick, + pool_idle_timeout = ?timeout, + "v2 pool: eviction task spawned" + ); + tokio::spawn(eviction_task(weak, rx, timeout)); + } + + /// Walk the host map, dropping idle connections that have exceeded + /// `timeout` and removing host entries whose retainers are empty. + /// + /// Called by the eviction task on each tick. Safe to call when the + /// host map is empty (no-op). + fn retain_idle_hosts(&self, timeout: Duration) { + let mut hosts = self.hosts.lock().unwrap(); + let before = hosts.len(); + hosts.retain(|key, entry| { + entry.retain_idle(timeout); + if entry.is_empty() { + tracing::debug!( + authority = %key.authority, + scheme = %key.scheme, + "v2 pool: host entry removed (empty after retain)" + ); + false + } else { + true + } + }); + let after = hosts.len(); + drop(hosts); + tracing::trace!( + hosts_before = before, + hosts_after = after, + "v2 pool: eviction tick complete" + ); + } +} + +/// Background loop that drops idle connections past `pool_idle_timeout`. +/// +/// On each tick: +/// 1. If the pool has been dropped (`Weak::upgrade` returns `None`), exit. +/// 2. Otherwise, walk the host map, drop connections whose `idle_at` is +/// older than `timeout`, and remove host entries whose retainers +/// report empty. +/// +/// Tick interval is `max(timeout, MIN_EVICTION_TICK)` — the floor exists +/// so a very short `pool_idle_timeout` doesn't spin the task hot. +/// Connections live on average between 1× and 2× the tick interval past +/// last use (sawtooth eviction). +/// +/// Exit paths: +/// - Primary: `drop_notifier` (the `Sender` half) drops when +/// `ConnectionPool` drops, the receiver errors with `Canceled`, the +/// `select!` left branch fires, the task exits immediately. +/// - Secondary: `Weak::upgrade` returns `None`. Belt-and-suspenders for +/// the unlikely case that the notifier didn't fire — the task will +/// exit on the next tick. +async fn eviction_task( + pool: std::sync::Weak, + mut drop_notifier: oneshot::Receiver, + timeout: Duration, +) { + let tick = timeout.max(MIN_EVICTION_TICK); + let mut interval = tokio::time::interval(tick); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + // The first tick of `tokio::time::interval` resolves immediately; + // burn it so subsequent ticks fire at `now + tick`. + interval.tick().await; + loop { + tokio::select! { + _ = &mut drop_notifier => break, + _ = interval.tick() => { + match pool.upgrade() { + Some(pool) => pool.retain_idle_hosts(timeout), + None => break, + } + } + } + } + tracing::trace!("v2 pool eviction task exiting"); } /// Tower `Service` wrapper for an H1 pool checkout. @@ -395,6 +701,19 @@ trait PoolEntry: Send + Sync { ctx: ConnectCtx, req: http_1x::Request, ) -> BoxFuture, BoxError>>; + + /// Drop idle connections that have exceeded the given timeout, and + /// drop any connection flagged as poisoned. + /// + /// Called by the background eviction task at each tick. Called on the + /// `Box` directly (no `&mut`) because the retainers + /// carry their own interior mutability (see [`BoxedRetainer`]). + fn retain_idle(&self, timeout: Duration); + + /// Whether this entry has no idle connections (H1 cache empty AND H2 + /// singleton empty). Used to drop empty host entries from the map + /// after `retain_idle`. + fn is_empty(&self) -> bool; } /// Concrete PoolEntry wrapping a Negotiate stack. @@ -402,7 +721,14 @@ trait PoolEntry: Send + Sync { /// `PoolUnnameable` propagates whatever pool-internal unnameable type /// flows through the checkout services' `CheckoutResponse`; it's carried /// through so `Conn::Response` can name `CheckoutResponse`. -struct TypedPoolEntry(S); +struct TypedPoolEntry { + stack: S, + /// Retainers for this entry's H1 cache + H2 singleton. Populated + /// lazily on first use of each leg by the `Negotiate` `layer_fn`s. + /// Empty until the first request against this host; safe to iterate + /// at any point (no leg = no-op retain, trivially empty). + retainers: RetainerSlot, +} impl PoolEntry for TypedPoolEntry where @@ -421,7 +747,7 @@ where ctx: ConnectCtx, req: http_1x::Request, ) -> BoxFuture, BoxError>> { - let mut svc = self.0.clone(); + let mut svc = self.stack.clone(); Box::pin(async move { // Checkout loop. Pops past stale idle entries (poisoned or dead). // Converges because: @@ -461,6 +787,18 @@ where )) }) } + + fn retain_idle(&self, timeout: Duration) { + let retainers = self.retainers.lock().expect("retainer slot poisoned"); + for r in retainers.iter() { + r.retain_idle(timeout); + } + } + + fn is_empty(&self) -> bool { + let retainers = self.retainers.lock().expect("retainer slot poisoned"); + retainers.iter().all(|r| r.is_empty()) + } } #[cfg(test)] @@ -479,4 +817,103 @@ mod tests { let uri: http_1x::Uri = "/path".parse().unwrap(); assert!(PoolKey::from_uri(&uri).is_none()); } + + /// A `PoolEntry` that panics if dispatched through and reports itself + /// empty. Used as a host entry for eviction-lifecycle unit tests where + /// no request is actually sent. + struct NullPoolEntry; + impl PoolEntry for NullPoolEntry { + fn send( + &mut self, + _ctx: ConnectCtx, + _req: http_1x::Request, + ) -> BoxFuture, BoxError>> { + unreachable!("NullPoolEntry::send called in a test") + } + fn retain_idle(&self, _timeout: Duration) {} + fn is_empty(&self) -> bool { + true + } + } + + fn pool_with_config(config: PoolConfig) -> Arc { + Arc::new(ConnectionPool { + config, + hosts: Mutex::new(HashMap::new()), + make_entry: Box::new(|_| Box::new(NullPoolEntry)), + eviction_spawned: AtomicBool::new(false), + drop_notifier: OnceLock::new(), + }) + } + + /// Without a `pool_idle_timeout`, `maybe_spawn_eviction_task` is a no-op. + /// Matches legacy `spawn_idle_interval` which also skips without a + /// timeout. + #[tokio::test] + async fn eviction_task_no_spawn_without_timeout() { + let pool = pool_with_config(PoolConfig { + pool_idle_timeout: None, + ..PoolConfig::default() + }); + pool.maybe_spawn_eviction_task(); + assert!(!pool.eviction_spawned.load(Ordering::Acquire)); + assert!(pool.drop_notifier.get().is_none()); + } + + /// A zero `pool_idle_timeout` is treated the same as `None`. Matches + /// legacy `spawn_idle_interval` (returns early on `Duration::ZERO`). + #[tokio::test] + async fn eviction_task_no_spawn_with_zero_timeout() { + let pool = pool_with_config(PoolConfig { + pool_idle_timeout: Some(Duration::ZERO), + ..PoolConfig::default() + }); + pool.maybe_spawn_eviction_task(); + assert!(!pool.eviction_spawned.load(Ordering::Acquire)); + } + + /// Multiple spawn calls only spawn a single task. The second call + /// observes `eviction_spawned = true` and returns without touching + /// `drop_notifier`. + #[tokio::test] + async fn eviction_task_spawn_is_idempotent() { + let pool = pool_with_config(PoolConfig { + pool_idle_timeout: Some(Duration::from_millis(100)), + ..PoolConfig::default() + }); + pool.maybe_spawn_eviction_task(); + assert!(pool.eviction_spawned.load(Ordering::Acquire)); + // drop_notifier gets set once. Grab a pointer to it so we can + // confirm the second call didn't replace it. + let first_notifier = pool.drop_notifier.get().expect("notifier set"); + let first_ptr = first_notifier as *const _; + pool.maybe_spawn_eviction_task(); + let second_notifier = pool.drop_notifier.get().expect("notifier still set"); + let second_ptr = second_notifier as *const _; + assert_eq!( + first_ptr, second_ptr, + "drop_notifier should not be replaced on repeat spawn" + ); + } + + /// `retain_idle_hosts` drops host entries whose retainers report empty. + /// Uses `NullPoolEntry` whose `is_empty` returns `true`, so the first + /// tick should remove every entry. + #[tokio::test] + async fn retain_idle_hosts_drops_empty_entries() { + let pool = pool_with_config(PoolConfig::default()); + let uri: http_1x::Uri = "http://example.com".parse().unwrap(); + let key = PoolKey::from_uri(&uri).unwrap(); + pool.hosts + .lock() + .unwrap() + .insert(key.clone(), Box::new(NullPoolEntry)); + assert_eq!(pool.hosts.lock().unwrap().len(), 1); + pool.retain_idle_hosts(Duration::from_secs(30)); + assert_eq!( + pool.hosts.lock().unwrap().len(), + 0, + "empty entry should have been evicted" + ); + } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index acb33db4ac9..46fdd4170e8 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -7,7 +7,7 @@ use std::net::SocketAddr; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; @@ -114,8 +114,8 @@ impl std::error::Error for PoisonedError {} /// where the raw transport stream is accessible. #[derive(Debug, Clone)] pub(crate) struct ConnectionInfo { - /// Remote address of the peer. `None` when the underlying connector did - /// not attach `HttpInfo` (e.g. the s2n-tls provider today). + /// Remote address of the peer. `None` when the underlying connector + /// did not attach `HttpInfo` to its `Connected` extras. pub(crate) remote_addr: Option, /// Local address of this end of the connection. pub(crate) local_addr: Option, @@ -163,23 +163,70 @@ impl PoisonPill { pub(crate) struct ManagedConnection { inner: S, info: ConnectionInfo, + /// Stable identifier for this physical connection, unique within the + /// owning pool. Shared across `Clone`s (an H2 connection's multiplexed + /// request handles all carry the same `conn_id`). Used in tracing and + /// surfaced through `ConnectionMetadata` for cross-layer correlation. + conn_id: u64, created_at: Instant, + /// Timestamp the connection last became idle (or its creation time, if + /// it has never been returned to the pool). + /// + /// - **H1**: stamped by [`CachedConnection::drop`] on the unpoisoned + /// return-to-pool path. Presence in the cache implies idle. + /// - **H2**: stamped by [`SingletonConnection::drop`] on the + /// `active_streams` 1 → 0 transition. Combined with + /// `active_streams > 0` in the H2 retain predicate, this prevents + /// eviction of a multiplexed connection that's still serving streams. + /// + /// `Arc>` because all `Clone`s of a `ManagedConnection` + /// observe the same timestamp (last-write-wins). `Mutex` over an + /// atomic-encoded `u64` because writes happen at most once per + /// request and reads at most once per eviction tick. + idle_at: Arc>, + /// In-flight request count, used by the H2 retain predicate to keep + /// actively-multiplexed connections alive regardless of `idle_at`. + /// + /// - **H1**: never incremented (always 0). H1 is serial; the cache's + /// idle set is the authoritative idleness signal. + /// - **H2**: incremented by [`SingletonConnection::call`] on dispatch, + /// decremented by [`SingletonConnection::drop`] when the body guard + /// releases. + /// + /// Arc-shared so all `Clone`s of an H2 `ManagedConnection` mutate the + /// same counter. + active_streams: Arc, poison: PoisonPill, _permit: Arc, } impl ManagedConnection { /// Create a new managed connection wrapping the given service. - pub(crate) fn new(inner: S, info: ConnectionInfo, permit: Arc) -> Self { + pub(crate) fn new( + inner: S, + info: ConnectionInfo, + conn_id: u64, + permit: Arc, + ) -> Self { + let now = Instant::now(); Self { inner, info, - created_at: Instant::now(), + conn_id, + created_at: now, + idle_at: Arc::new(Mutex::new(now)), + active_streams: Arc::new(std::sync::atomic::AtomicUsize::new(0)), poison: PoisonPill::healthy(), _permit: permit, } } + /// Stable identifier for this physical connection. Shared across + /// clones (H2 multiplexing). + pub(crate) fn conn_id(&self) -> u64 { + self.conn_id + } + /// Whether this connection has been poisoned. pub(crate) fn is_poisoned(&self) -> bool { self.poison.is_poisoned() @@ -197,6 +244,42 @@ impl ManagedConnection { self.created_at } + /// Timestamp of the last return-to-idle (or creation). + /// + /// Guaranteed meaningful for any connection sitting in the pool — the + /// initial value is `created_at`, and every return-to-pool overwrites + /// it via [`Self::mark_idle`]. + pub(crate) fn idle_at(&self) -> Instant { + *self.idle_at.lock().expect("idle_at lock poisoned") + } + + /// Stamp the return-to-idle moment. Called by [`CachedConnection::drop`] + /// on the unpoisoned path. + pub(crate) fn mark_idle(&self) { + *self.idle_at.lock().expect("idle_at lock poisoned") = Instant::now(); + } + + /// Current in-flight stream count. See [`Self::active_streams`] + /// (field doc) for semantics. + pub(crate) fn active_streams_count(&self) -> usize { + self.active_streams + .load(std::sync::atomic::Ordering::Acquire) + } + + /// Clone of the Arc'd stream counter, for publication through the H2 + /// side-channel so `SingletonConnection` can increment/decrement it + /// without reaching through the opaque `Singled<…>`. + pub(crate) fn active_streams_ref(&self) -> Arc { + self.active_streams.clone() + } + + /// Clone of the Arc'd idle timestamp, for publication through the H2 + /// side-channel so `SingletonConnection::drop` can stamp it on the + /// `active_streams` 1 → 0 transition. + pub(crate) fn idle_at_ref(&self) -> Arc> { + self.idle_at.clone() + } + /// Mutable access to the inner service. pub(crate) fn inner_mut(&mut self) -> &mut S { &mut self.inner @@ -210,13 +293,18 @@ impl ManagedConnection { /// Address fields are copied. pub(crate) fn metadata(&self) -> ConnectionMetadata { let poison = self.poison.clone(); + let conn_id = self.conn_id; + let remote = self.info.remote_addr; let mut builder = ConnectionMetadata::builder() // Hardcoded `false`: v2 does not yet plumb proxy configuration // through the pool stack. When proxy support lands, // `ConnectionInfo` should carry an `is_proxied` flag populated // at connection establishment and passed through here. .proxied(false) - .poison_fn(move || poison.poison()); + .poison_fn(move || { + tracing::debug!(conn_id, ?remote, "v2 pool: connection poisoned"); + poison.poison(); + }); builder .set_remote_addr(self.info.remote_addr) .set_local_addr(self.info.local_addr); @@ -229,7 +317,10 @@ impl Clone for ManagedConnection { Self { inner: self.inner.clone(), info: self.info.clone(), + conn_id: self.conn_id, created_at: self.created_at, + idle_at: self.idle_at.clone(), + active_streams: self.active_streams.clone(), poison: self.poison.clone(), _permit: self._permit.clone(), } @@ -313,10 +404,19 @@ where impl Drop for CachedConnection { fn drop(&mut self) { if let Some(cached) = self.inner.take() { - if cached.inner().is_poisoned() { + let managed = cached.inner(); + let conn_id = managed.conn_id; + if managed.is_poisoned() { + tracing::debug!(conn_id, "v2 pool: connection discarded (poisoned)"); cached.discard(); + } else { + // Stamp the return-to-idle moment so the eviction task + // can measure elapsed idleness. Healthy connection — the + // Cached drops normally after this, returning it to the + // cache's idle set. + managed.mark_idle(); + tracing::trace!(conn_id, "v2 pool: connection returned to idle"); } - // else: cached drops normally, returning the connection to the cache } } } @@ -340,13 +440,39 @@ impl Drop for CachedConnection { /// from `Made(svc)` back to `Empty`, dropping the old service. The next /// request's `call` runs a fresh handshake through `H2ConnectAndHandshake`, /// which publishes new metadata — overwriting whatever was there before. -/// Clients that grabbed metadata for the previous connection retain a -/// `ConnectionMetadata` pointing at the dropped connection's `PoisonPill`; -/// poisoning it is a no-op because the connection it described is already -/// gone. +/// State an H2 checkout (`SingletonConnection`) needs but cannot reach +/// through the opaque `Singled<…>` to read from the underlying +/// `ManagedConnection`. Published by the H2 handshake on each new +/// connection; consumed by `SingletonConnection::new` as a snapshot. +/// +/// `active_streams` and `idle_at` are clones of the Arcs held by the +/// `ManagedConnection`, so updates through this state are visible to the +/// retain predicate operating on the connection directly. +#[derive(Clone)] +pub(crate) struct H2ConnectionState { + pub(crate) metadata: ConnectionMetadata, + pub(crate) active_streams: Arc, + pub(crate) idle_at: Arc>, +} + +/// Per-host side-channel between H2 handshake (writer) and H2 checkout +/// (reader). +/// +/// `hyper_util::client::pool::singleton::Singled<…>` is opaque, so the +/// checkout side cannot reach through it to the underlying +/// `ManagedConnection`. This ref carries everything the checkout side +/// needs (stamped fresh on each handshake): the user-facing +/// `ConnectionMetadata` for poison and address surfacing, plus the +/// Arc-shared `active_streams` counter and `idle_at` timestamp that the +/// retain predicate consults and `SingletonConnection` mutates. +/// +/// On re-handshake (poisoning, GOAWAY, etc.), the entire state is +/// replaced last-writer-wins. A `SingletonConnection` constructed against +/// the previous handshake holds its own snapshot, so its +/// increments/decrements still target the right counter. #[derive(Clone, Default)] pub(crate) struct H2ConnectionRef { - inner: Arc>>, + inner: Arc>>, } impl H2ConnectionRef { @@ -354,46 +480,113 @@ impl H2ConnectionRef { Self::default() } - pub(crate) fn publish(&self, metadata: ConnectionMetadata) { - *self.inner.lock().unwrap() = Some(metadata); + pub(crate) fn publish(&self, state: H2ConnectionState) { + *self.inner.lock().unwrap() = Some(state); } - pub(crate) fn current(&self) -> Option { + /// Snapshot of the published state. `None` until the first + /// successful handshake for this host. + pub(crate) fn current(&self) -> Option { self.inner.lock().unwrap().clone() } } -/// A connection checked out from the H2 singleton. +/// RAII guard for one in-flight dispatch against an H2 connection. +/// +/// Existence reflects "this checkout has dispatched a request; the +/// connection is busy on its behalf." Constructed by +/// [`SingletonConnection::call`] when a request is dispatched; dropped +/// when the response body's guard releases. /// -/// Symmetric wrapper to `CachedConnection` for the H2 leg. The inner type -/// is parameterized because hyper-util's `Singled` is unnameable outside -/// its crate; we accept whatever type comes out of the singleton. Gives us -/// a named, pool-owned H2 handle we can hold as a guard, inspect, and -/// instrument later without vendoring `singleton.rs`. +/// Drop decrements `active_streams` and, on the 1 → 0 transition, stamps +/// `idle_at` — the moment the connection becomes truly idle. +struct DispatchGuard { + active_streams: Arc, + idle_at: Arc>, +} + +impl DispatchGuard { + /// Start a dispatch against the connection described by `state`. + /// Increments `active_streams`; the returned guard releases the + /// increment on drop. + fn start(state: &H2ConnectionState) -> Self { + state + .active_streams + .fetch_add(1, std::sync::atomic::Ordering::AcqRel); + Self { + active_streams: state.active_streams.clone(), + idle_at: state.idle_at.clone(), + } + } +} + +impl Drop for DispatchGuard { + fn drop(&mut self) { + let prev = self + .active_streams + .fetch_sub(1, std::sync::atomic::Ordering::AcqRel); + if prev == 1 { + if let Ok(mut idle_at) = self.idle_at.lock() { + *idle_at = Instant::now(); + } + } + } +} + +/// A checked-out H2 connection ready to dispatch requests. /// -/// No Drop hook: dropping a Singled clone does not affect singleton state; -/// dead/poisoned H2 connections are handled reactively via -/// `Singled::poll_ready` clearing the state to `Empty`. +/// H2 is multiplexed: the underlying `ManagedConnection` stays resident in +/// `Singleton` for the duration of every concurrent request, so "presence +/// in a cache" — the idleness signal we use for H1 — doesn't apply. +/// Instead, `SingletonConnection` mints a [`DispatchGuard`] on each +/// `call`; the guard's lifetime tracks one in-flight request and its +/// drop releases the corresponding `active_streams` slot, stamping +/// `idle_at` on the final release. The H2 retain predicate keeps the +/// connection alive while `active_streams > 0` regardless of `idle_at`. /// -/// `h2_ref` is a clone of the same [`H2ConnectionRef`] that the handshake -/// service writes to — it lets `SingletonConnection::metadata` expose the -/// current H2 connection's metadata to the adapter even though `Singled<…>` -/// itself is opaque. +/// Lifecycle: +/// - Constructed in the H2 upgrade `map_response` after `Singleton::call` +/// completes. Snapshots the current [`H2ConnectionState`] and holds it +/// for the checkout's lifetime. +/// - [`Self::call`] starts a [`DispatchGuard`] before delegating to +/// `Singled::call`. The guard is held in `dispatch` for the rest of +/// this `SingletonConnection`'s lifetime. +/// - On drop, the guard's drop releases the active-stream count and may +/// stamp `idle_at`. If `call` was never invoked (e.g. the surrounding +/// checkout was abandoned by the post-checkout `poll_ready` retry +/// loop), `dispatch` is `None` and drop is a no-op. pub(crate) struct SingletonConnection { inner: T, - h2_ref: H2ConnectionRef, + /// Snapshot of the H2 side-channel state taken at construction. + /// + /// Held for the lifetime of this checkout so dispatch guards always + /// target the connection this checkout was issued against, even if a + /// re-handshake replaces the underlying `H2ConnectionRef`'s published + /// state in the meantime. + state: Option, + /// Active-stream guard. `Some` between [`Self::call`] and drop. + /// `None` if `call` was never invoked or before `call` runs. + dispatch: Option, } impl SingletonConnection { pub(crate) fn new(inner: T, h2_ref: H2ConnectionRef) -> Self { - Self { inner, h2_ref } + let state = h2_ref.current(); + Self { + inner, + state, + dispatch: None, + } } - /// Metadata for the current H2 connection, if a handshake has completed. + /// Metadata for the H2 connection this checkout was issued against. /// - /// Returns `None` before the first successful H2 handshake for this host. + /// Returns `None` only if `SingletonConnection` was constructed before + /// any handshake had published — which doesn't happen in normal flow, + /// since `Singleton::call` always completes a handshake before + /// yielding the `Singled` we wrap. pub(crate) fn metadata(&self) -> Option { - self.h2_ref.current() + self.state.as_ref().map(|s| s.metadata.clone()) } } @@ -410,6 +603,9 @@ where } fn call(&mut self, req: Req) -> Self::Future { + if let Some(state) = self.state.as_ref() { + self.dispatch = Some(DispatchGuard::start(state)); + } self.inner.call(req) } } @@ -487,9 +683,144 @@ impl hyper::body::Body for GuardedBody { /// The response type both H1 and H2 pool checkouts produce. /// -/// Carries `GuardedBody` because some pool-internal leg -/// holds a checkout guard whose type is unnameable outside hyper-util -/// (today: the H2 leg's `Singled<...>`). `Negotiate` requires uniform -/// response types across legs; this alias is the uniform type that +/// Carries `GuardedBody` so the H2 leg can hold its +/// checkout guard (`Singled<…>`, type-erased through `PoolUnnameable`) +/// for the response body's lifetime. `Negotiate` requires uniform +/// response types across its legs; this alias is the uniform type that /// consumers above the Negotiate composition point work with. pub(crate) type CheckoutResponse = http_1x::Response>; + +#[cfg(test)] +mod tests { + //! Unit tests for H2 active-stream tracking. + //! + //! The test harness (`ConnectionTestHarness`) is plain HTTP only — no + //! ALPN, so we cannot exercise the full H2 pipeline end-to-end yet + //! (tracked in bosun.md as the "H2 test coverage gap"). These tests + //! verify the atomic transitions directly, which is the architectural + //! correctness property we need: an in-flight H2 connection must not be + //! evicted by the retain predicate, and the connection's `idle_at` must + //! be stamped only on the `active_streams` 1 → 0 transition. + use super::*; + use std::future::Future; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tower::Service as _; + + /// A minimal `Service<()>` that returns `Ok(())` — stand-in for + /// `Singled::call`. We only care about the wrapper's stream-count + /// side effects. + #[derive(Clone, Default)] + struct OkService; + impl Service<()> for OkService { + type Response = (); + type Error = BoxError; + type Future = std::pin::Pin> + Send>>; + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, _: ()) -> Self::Future { + Box::pin(async { Ok(()) }) + } + } + + fn publish_state(h2_ref: &H2ConnectionRef) -> (Arc, Arc>) { + use aws_smithy_runtime_api::client::connection::ConnectionMetadata; + let active = Arc::new(AtomicUsize::new(0)); + let idle_at = Arc::new(Mutex::new(Instant::now())); + h2_ref.publish(H2ConnectionState { + metadata: ConnectionMetadata::builder() + .proxied(false) + .poison_fn(|| {}) + .build(), + active_streams: active.clone(), + idle_at: idle_at.clone(), + }); + (active, idle_at) + } + + /// A round-trip through `SingletonConnection::call` mints a + /// `DispatchGuard` that releases on drop. Counter goes 0 → 1 → 0. + #[tokio::test] + async fn singleton_dispatch_round_trip_is_net_zero() { + let h2_ref = H2ConnectionRef::new(); + let (active, _) = publish_state(&h2_ref); + assert_eq!(active.load(Ordering::Acquire), 0); + + let mut sc = SingletonConnection::new(OkService, h2_ref); + let _ = sc.call(()).await; + assert_eq!( + active.load(Ordering::Acquire), + 1, + "call should mint a DispatchGuard, incrementing active_streams" + ); + drop(sc); + assert_eq!( + active.load(Ordering::Acquire), + 0, + "DispatchGuard's Drop should release the active_streams slot" + ); + } + + /// Constructing without dispatching leaves no `DispatchGuard`, so + /// dropping is a no-op. The post-checkout `poll_ready` retry loop + /// relies on this — a checkout discarded before `call` must not + /// underflow the counter. + #[tokio::test] + async fn singleton_drop_without_dispatch_is_noop() { + let h2_ref = H2ConnectionRef::new(); + let (active, _) = publish_state(&h2_ref); + active.store(5, Ordering::Release); + + let sc = SingletonConnection::new(OkService, h2_ref); + drop(sc); + assert_eq!( + active.load(Ordering::Acquire), + 5, + "uncalled SingletonConnection must not release a DispatchGuard" + ); + } + + /// Concurrent dispatches against the same H2 connection (simulating + /// multiplexed requests) all land on the same counter. `idle_at` is + /// stamped only on the LAST `DispatchGuard` drop — the 1 → 0 + /// transition. + #[tokio::test] + async fn idle_at_stamped_only_on_last_dispatch_release() { + let h2_ref = H2ConnectionRef::new(); + let (active, idle_at) = publish_state(&h2_ref); + + // Capture the stamp from construction to detect updates. + let original_idle = *idle_at.lock().unwrap(); + + let mut a = SingletonConnection::new(OkService, h2_ref.clone()); + let mut b = SingletonConnection::new(OkService, h2_ref.clone()); + let mut c = SingletonConnection::new(OkService, h2_ref); + + let _ = a.call(()).await; + let _ = b.call(()).await; + let _ = c.call(()).await; + assert_eq!(active.load(Ordering::Acquire), 3); + + // Sleep so any stamp is distinguishable from construction time. + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + + // Release two — idle_at must NOT yet be stamped (counter > 0). + drop(a); + drop(b); + assert_eq!(active.load(Ordering::Acquire), 1); + assert_eq!( + *idle_at.lock().unwrap(), + original_idle, + "idle_at must not be stamped while dispatches are in flight" + ); + + // Last release: 1 → 0 transition. idle_at must be stamped. + drop(c); + assert_eq!(active.load(Ordering::Acquire), 0); + let final_idle = *idle_at.lock().unwrap(); + assert!( + final_idle > original_idle, + "idle_at should be stamped when the last DispatchGuard releases" + ); + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index b4dd4ba172c..980f935a9f1 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -25,6 +25,32 @@ use tower::Service; use super::connection::{ConnectCtx, ConnectionInfo, ConnectionPermit, ManagedConnection}; use super::BoxError; +/// Pool-scoped instrumentation primitives shared across layers in the +/// pool stack. Held by `ConnectionPool` for the pool's lifetime and +/// cloned into individual layers (the H1/H2 handshake services) at +/// construction. +/// +/// Cheap to clone — one `Arc` per primitive. +#[derive(Clone)] +pub(crate) struct PoolHooks { + conn_id_counter: Arc, +} + +impl PoolHooks { + pub(crate) fn new() -> Self { + Self { + conn_id_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)), + } + } + + /// Mint the next connection id. Stable for the connection's lifetime; + /// the underlying counter wraps at `u64::MAX`. + pub(crate) fn next_conn_id(&self) -> u64 { + self.conn_id_counter + .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + } +} + /// Wraps a connector service, acquiring semaphore permits before connecting. /// /// Returns `(IO, Arc)` so the permit can be stored on @@ -207,11 +233,12 @@ fn spawn_driver(future: impl Future + Send + 'static) { /// produced by [`ConnectionLimit`] wrapping a TCP/TLS connector. pub(crate) struct H1ConnectAndHandshake { connector: C, + hooks: PoolHooks, } impl H1ConnectAndHandshake { - pub(crate) fn new(connector: C) -> Self { - Self { connector } + pub(crate) fn new(connector: C, hooks: PoolHooks) -> Self { + Self { connector, hooks } } } @@ -219,6 +246,7 @@ impl Clone for H1ConnectAndHandshake { fn clone(&self) -> Self { Self { connector: self.connector.clone(), + hooks: self.hooks.clone(), } } } @@ -239,22 +267,33 @@ where } fn call(&mut self, ctx: ConnectCtx) -> Self::Future { + let hooks = self.hooks.clone(); let fut = self.connector.call(ctx); Box::pin(async move { let (io, permit) = fut.await.map_err(Into::into)?; let info = capture_info(&io); + let conn_id = hooks.next_conn_id(); let (tx, conn) = hyper::client::conn::http1::Builder::new() .handshake(io) .await .map_err(|e| Box::new(e) as BoxError)?; + tracing::debug!( + conn_id, + protocol = "h1", + remote = ?info.remote_addr, + local = ?info.local_addr, + "v2 pool: connection established" + ); + spawn_driver({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { if let Err(e) = conn.with_upgrades().await { tracing::debug!( + conn_id, protocol = "h1", ?remote_addr, ?local_addr, @@ -265,7 +304,12 @@ where } }); - Ok(ManagedConnection::new(H1SendRequest::new(tx), info, permit)) + Ok(ManagedConnection::new( + H1SendRequest::new(tx), + info, + conn_id, + permit, + )) }) } } @@ -278,6 +322,7 @@ where pub(crate) struct H2ConnectAndHandshake { connector: C, h2_ref: super::connection::H2ConnectionRef, + hooks: PoolHooks, } impl H2ConnectAndHandshake { @@ -287,8 +332,16 @@ impl H2ConnectAndHandshake { /// the underlying slot), so the H2 checkout path can expose connection /// metadata and poison support to the adapter layer even though /// `Singled<…>` itself is opaque. - pub(crate) fn new(connector: C, h2_ref: super::connection::H2ConnectionRef) -> Self { - Self { connector, h2_ref } + pub(crate) fn new( + connector: C, + h2_ref: super::connection::H2ConnectionRef, + hooks: PoolHooks, + ) -> Self { + Self { + connector, + h2_ref, + hooks, + } } } @@ -297,6 +350,7 @@ impl Clone for H2ConnectAndHandshake { Self { connector: self.connector.clone(), h2_ref: self.h2_ref.clone(), + hooks: self.hooks.clone(), } } } @@ -319,21 +373,32 @@ where fn call(&mut self, _req: ()) -> Self::Future { let fut = self.connector.call(()); let h2_ref = self.h2_ref.clone(); + let hooks = self.hooks.clone(); Box::pin(async move { let (io, permit) = fut.await.map_err(Into::into)?; let info = capture_info(&io); + let conn_id = hooks.next_conn_id(); let (tx, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new()) .handshake(io) .await .map_err(|e| Box::new(e) as BoxError)?; + tracing::debug!( + conn_id, + protocol = "h2", + remote = ?info.remote_addr, + local = ?info.local_addr, + "v2 pool: connection established" + ); + spawn_driver({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { if let Err(e) = conn.await { tracing::debug!( + conn_id, protocol = "h2", ?remote_addr, ?local_addr, @@ -344,11 +409,16 @@ where } }); - let managed = ManagedConnection::new(H2SendRequest::new(tx), info, permit); - // Publish this connection's metadata for the checkout side. - // Singleton replaces its stored connection wholesale on each new - // handshake, so last-writer-wins on the ref is correct. - h2_ref.publish(managed.metadata()); + let managed = ManagedConnection::new(H2SendRequest::new(tx), info, conn_id, permit); + // Publish this connection's state (metadata + active_streams + + // idle_at refs) for the checkout side. Singleton replaces its + // stored connection wholesale on each new handshake, so + // last-writer-wins on the ref is correct. + h2_ref.publish(super::connection::H2ConnectionState { + metadata: managed.metadata(), + active_streams: managed.active_streams_ref(), + idle_at: managed.idle_at_ref(), + }); Ok(managed) }) } diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index 2af084e9348..4c9735a9fd4 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -138,6 +138,29 @@ impl BuilderV2 { let pool = pool::build_pool(tcp, config); V2HttpClient::new(pool).into_shared() } + + /// Build an HTTP client without TLS, using a custom DNS resolver. + /// + /// Mirrors `build_http` but threads a custom `ResolveDns`. Intended + /// for tests that use `WireMockServer::dns_resolver()` to redirect + /// all lookups to a loopback mock. Public API for non-TLS resolver + /// plumbing is `#[doc(hidden)]` — TLS is the common case. + #[doc(hidden)] + pub fn build_http_with_resolver( + self, + resolver: impl aws_smithy_runtime_api::client::dns::ResolveDns + Clone + 'static, + ) -> SharedHttpClient { + use crate::client::dns::HyperUtilResolver; + let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); + tcp.set_nodelay(self.tcp_nodelay); + let config = PoolConfig { + max_connections: self.max_connections, + max_connections_per_host: self.max_connections_per_host, + pool_idle_timeout: self.pool_idle_timeout, + }; + let pool = pool::build_pool(tcp, config); + V2HttpClient::new(pool).into_shared() + } } impl BuilderV2 { diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 7c8fced8527..c81fa135e72 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1091,3 +1091,142 @@ async fn v2_read_timeout_fires_on_silent_server() { async fn v2_read_timeout() { v2_read_timeout_fires_on_silent_server().await; } + +// --------------------------------------------------------------------------- +// Tracing-assertion tests for the background eviction task. +// +// These are v2-only because they assert on internal tracing events that the +// v2 pool emits. The existing `v2_idle_timeout_eviction` test proves eviction +// works end-to-end (new TCP on second request), but doesn't distinguish the +// background task from checkout-time filtering. These tests assert the +// background task specifically ran without issuing a second request. +// --------------------------------------------------------------------------- + +/// Capture `aws_smithy_http_client` tracing events at TRACE level for the +/// duration of a test. ANSI color codes are explicitly disabled so string +/// matches on fields like `reason="idle_expired"` work — the default +/// `tracing_subscriber::fmt` output embeds ANSI escapes between the field +/// name and value which breaks naive `contains` assertions. +/// +/// Returns `(guard, buffer)`. The guard resets the subscriber on drop; +/// `buffer.lock().unwrap()` yields the captured bytes. +fn capture_pool_logs() -> ( + tracing::subscriber::DefaultGuard, + std::sync::Arc>>, +) { + use std::io::Write; + use std::sync::{Arc, Mutex}; + + struct BufWriter(Arc>>); + impl Write for BufWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + let buf = Arc::new(Mutex::new(Vec::::new())); + let buf_clone = buf.clone(); + let subscriber = tracing_subscriber::fmt() + .with_ansi(false) + .with_max_level(tracing::Level::TRACE) + .with_writer(move || BufWriter(buf_clone.clone())) + .finish(); + let guard = tracing::subscriber::set_default(subscriber); + (guard, buf) +} + +fn captured_str(buf: &std::sync::Arc>>) -> String { + String::from_utf8(buf.lock().unwrap().clone()).expect("captured logs are utf-8") +} + +/// Verify the background eviction task runs and emits the "connection +/// evicted" event with `reason=idle_expired` after the idle timeout passes, +/// without any subsequent request being made. +/// +/// Also verifies the host entry is removed once its cache becomes empty. +#[tokio::test] +async fn v2_background_eviction_emits_tracing_events() { + let (_guard, logs) = capture_pool_logs(); + + let server = WireMockServer::start(vec![ReplayedEvent::status(200)]).await; + + let idle_timeout = Duration::from_millis(100); + let client = V2Client.make(ClientConfig::default().with_idle_timeout(idle_timeout)); + let url = localhost_url(&server); + + // One request — connection established, dispatched, returned to idle. + let status = send_to(&client, &url) + .await + .expect("first request should succeed") + .status() + .as_u16(); + assert_eq!(status, 200); + + // Wait past idle timeout + one tick interval. No second request here — + // any eviction must be driven by the background task. + tokio::time::sleep(idle_timeout * 3).await; + + let captured = captured_str(&logs); + assert!( + captured.contains("v2 pool: initialized"), + "expected pool init log" + ); + assert!( + captured.contains("v2 pool: eviction task spawned"), + "expected eviction task spawn log" + ); + assert!( + captured.contains("v2 pool: connection established"), + "expected connection-established log" + ); + assert!( + captured.contains("v2 pool: connection evicted"), + "expected 'connection evicted' log — background task should have evicted the idle connection" + ); + assert!( + captured.contains("reason=\"idle_expired\""), + "expected idle_expired reason. captured:\n{captured}" + ); + assert!( + captured.contains("v2 pool: host entry removed"), + "expected host-entry-removed log once cache went empty" + ); +} + +/// Verify connections carry stable `conn_id` values — first established +/// connection gets id 0, subsequent new connection (after eviction forces +/// a fresh TCP) gets id 1. +#[tokio::test] +async fn v2_conn_id_is_stable_and_monotonic() { + let (_guard, logs) = capture_pool_logs(); + + let server = + WireMockServer::start(vec![ReplayedEvent::status(200), ReplayedEvent::status(200)]).await; + + let idle_timeout = Duration::from_millis(100); + let client = V2Client.make(ClientConfig::default().with_idle_timeout(idle_timeout)); + let url = localhost_url(&server); + + send_to(&client, &url).await.expect("first request"); + tokio::time::sleep(idle_timeout * 3).await; + send_to(&client, &url).await.expect("second request"); + + let captured = captured_str(&logs); + assert!( + captured.contains("v2 pool: connection established"), + "no 'connection established' log captured" + ); + assert!( + captured.contains("conn_id=0"), + "first connection should be conn_id=0. captured:\n{captured}" + ); + assert!( + captured.contains("conn_id=1"), + "second connection (after eviction) should be conn_id=1. captured:\n{captured}" + ); + assert!(captured.contains("v2 pool: connection evicted")); +} diff --git a/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs b/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs index bde0688df6d..b3314b95b96 100644 --- a/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs +++ b/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs @@ -128,6 +128,15 @@ impl MakeClient for Hyper1xClient { } } +/// v2 HTTP client (composable pool) via `BuilderV2::new_v2()` +struct Hyper1xV2Client; + +impl MakeClient for Hyper1xV2Client { + fn make(&self, mock: &WireMockServer) -> SharedHttpClient { + aws_smithy_http_client::v2::BuilderV2::new().build_http_with_resolver(mock.dns_resolver()) + } +} + /// Repeatedly send test operation until `end_of_test` is received, then run match_clause. async fn run_test( make_client: &dyn MakeClient, @@ -225,7 +234,14 @@ async fn all_stacks( &match_clause, ) .await; - run_test(&Hyper1xClient, events, reconnect_mode, &match_clause).await; + run_test( + &Hyper1xClient, + events.clone(), + reconnect_mode, + &match_clause, + ) + .await; + run_test(&Hyper1xV2Client, events, reconnect_mode, &match_clause).await; } #[tokio::test] From df18f060b0d18470c28e4c52fd743f42e963c719 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 14 May 2026 11:29:12 -0400 Subject: [PATCH 16/47] wire proxy support --- .../aws-smithy-http-client/src/client.rs | 30 +- .../aws-smithy-http-client/src/client/pool.rs | 79 +++- .../src/client/pool/connection.rs | 76 ++-- .../src/client/pool/handshake.rs | 14 +- .../src/client/proxy.rs | 28 ++ .../aws-smithy-http-client/src/client/v2.rs | 153 +++++-- .../tests/proxy_tests.rs | 418 ++++++++++++++++++ 7 files changed, 699 insertions(+), 99 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client.rs b/rust-runtime/aws-smithy-http-client/src/client.rs index f756d11180c..25afea5ce7b 100644 --- a/rust-runtime/aws-smithy-http-client/src/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client.rs @@ -512,31 +512,13 @@ fn new_tokio_hyper_builder( } impl Adapter { - /// Add proxy authentication header to the request if needed + /// Add proxy authentication header to the request if needed. + /// + /// Delegates to the shared [`proxy::add_proxy_auth_header`] helper — + /// the v2 HTTP client uses the same logic. fn add_proxy_auth_header(&self, request: &mut http_1x::Request) { - // Only add auth for HTTP requests (not HTTPS which uses CONNECT tunneling) - if request.uri().scheme() != Some(&http_1x::uri::Scheme::HTTP) { - return; - } - - // Don't override existing proxy authorization header - if request - .headers() - .contains_key(http_1x::header::PROXY_AUTHORIZATION) - { - return; - } - - if let Some(ref matcher) = self.proxy_matcher { - if let Some(intercept) = matcher.intercept(request.uri()) { - // Add basic auth header if available - if let Some(auth_header) = intercept.basic_auth() { - request - .headers_mut() - .insert(http_1x::header::PROXY_AUTHORIZATION, auth_header.clone()); - tracing::debug!("added proxy authentication header for {}", request.uri()); - } - } + if let Some(matcher) = self.proxy_matcher.as_ref() { + proxy::add_proxy_auth_header(request, matcher); } } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 9ae9965e9e9..0b672709707 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -20,7 +20,7 @@ mod vendored_cache; /// /// Re-exports [`vendored_cache`], our vendored copy of hyper-util's /// `pool::cache` with SDK-specific modifications. The re-export insulates -/// callers from the vendoring detail — all pool code uses `cache::…` +/// callers from the vendoring detail; all pool code uses `cache::…` /// regardless of where the implementation lives. mod cache { pub(crate) use super::vendored_cache::*; @@ -66,7 +66,7 @@ type BoxFuture = Pin + Send>>; /// Write-once per request: `H{1,2}Checkout::call` sets it exactly once /// during checkout. Subsequent sets are no-ops (the `OnceLock` guarantees /// single-init). Retries produce fresh `HttpConnector::call` invocations -/// which create fresh capture slots — each attempt's metadata points at +/// which create fresh capture slots; each attempt's metadata points at /// the connection used for that attempt. #[derive(Clone, Default)] pub(crate) struct ConnectionMetadataCapture { @@ -97,7 +97,7 @@ impl ConnectionMetadataCapture { /// Pool-level configuration. /// /// Plumbed from the eventual `Builder::new_v2()` public API through -/// `build_pool` to `ConnectionPool`. Internal type — the public surface +/// `build_pool` to `ConnectionPool`. Internal type: the public surface /// is the builder's per-knob setters, not this struct. /// /// All fields are `None` by default (defaults applied by the pool where @@ -139,13 +139,13 @@ impl PoolKey { /// Type-erased handle to one Negotiate leg's eviction interface. /// -/// Each `TypedPoolEntry` holds at most two of these — one for its H1 cache, -/// one for its H2 singleton. The two underlying types (`Cache<…>` and +/// Each `TypedPoolEntry` holds at most two of these (one for its H1 cache, +/// one for its H2 singleton). The two underlying types (`Cache<…>` and /// `Singleton<…>`) carry upstream-unnameable parameters from `Negotiate`, /// so they're captured inside `Box` closures. The closures are /// `Fn` (so multiple callers can invoke through `Arc`); interior /// mutability lives in a `std::sync::Mutex` over a clone of the leg, which -/// is uncontended in practice — retain runs at most once per +/// is uncontended in practice; retain runs at most once per /// `max(pool_idle_timeout, MIN_EVICTION_TICK)`, and the leg's own state /// already serializes through its internal `Arc>`. pub(crate) struct BoxedRetainer { @@ -201,7 +201,7 @@ where // checkout. Clones share the underlying slot. let h2_ref = H2ConnectionRef::new(); - // Bounded at two — the H1 fallback and H2 upgrade `layer_fn`s + // Bounded at two: the H1 fallback and H2 upgrade `layer_fn`s // each push one entry on first construction. let retainers: RetainerSlot = Arc::new(Mutex::new(Vec::with_capacity(2))); @@ -375,7 +375,7 @@ pub(crate) struct ConnectionPool { /// with `Canceled`, and the task exits its `select!` loop. /// /// `OnceLock` because the task spawns lazily on first `send_request` - /// — at most one task per pool. The task additionally holds a + /// (at most one task per pool). The task additionally holds a /// `Weak` as a fallback exit signal. drop_notifier: OnceLock>, } @@ -392,7 +392,7 @@ impl ConnectionPool { /// `ctx` carries the routing URI plus per-operation connect-time data /// (connect_timeout). Per-operation read_timeout, if any, must be /// attached to `req.extensions_mut()` as [`ReadTimeoutHint`] before - /// calling — the checkout services read it from there. + /// calling; the checkout services read it from there. /// /// Takes `self: &Arc` so the lazy eviction task can hold a /// `Weak` as a fallback exit signal (primary exit is the drop @@ -413,7 +413,7 @@ impl ConnectionPool { // Dispatch the request through the per-host entry. The lock is // held only long enough to look up / create the entry and call - // `send` — all I/O happens in the returned future, outside the + // `send`; all I/O happens in the returned future, outside the // lock. let fut = { let mut hosts = self.hosts.lock().unwrap(); @@ -451,7 +451,7 @@ impl ConnectionPool { return; } let (tx, rx) = oneshot::channel::(); - // Invariant: this code path runs exactly once per pool — the + // Invariant: this code path runs exactly once per pool; the // `compare_exchange` on `eviction_spawned` above is the unique // claim. If `drop_notifier` is already populated, that invariant // is broken. @@ -507,7 +507,7 @@ impl ConnectionPool { /// older than `timeout`, and remove host entries whose retainers /// report empty. /// -/// Tick interval is `max(timeout, MIN_EVICTION_TICK)` — the floor exists +/// Tick interval is `max(timeout, MIN_EVICTION_TICK)`. The floor exists /// so a very short `pool_idle_timeout` doesn't spin the task hot. /// Connections live on average between 1× and 2× the tick interval past /// last use (sawtooth eviction). @@ -517,7 +517,7 @@ impl ConnectionPool { /// `ConnectionPool` drops, the receiver errors with `Canceled`, the /// `select!` left branch fires, the task exits immediately. /// - Secondary: `Weak::upgrade` returns `None`. Belt-and-suspenders for -/// the unlikely case that the notifier didn't fire — the task will +/// the unlikely case that the notifier didn't fire; the task will /// exit on the next tick. async fn eviction_task( pool: std::sync::Weak, @@ -544,6 +544,49 @@ async fn eviction_task( tracing::trace!("v2 pool eviction task exiting"); } +/// Rewrite an HTTP/1.1 request's URI to the appropriate request-target +/// form for dispatch: +/// +/// - `CONNECT` → authority-form (`host:port`) +/// - proxied non-CONNECT → absolute-form (full URL with scheme + authority) +/// - direct non-CONNECT → origin-form (path + query only) +/// +/// Mirrors `hyper-util`'s legacy `client::Client::send_request` request +/// preparation. Required because `hyper::client::conn::http1` sends the +/// URI as-is; the form selection lives in the layer above the +/// connection. +fn rewrite_h1_request_target(req: &mut http_1x::Request, is_proxied: bool) { + use http_1x::{uri::Parts, Method, Uri}; + if req.method() == Method::CONNECT { + // Authority-form: keep just the authority component. + if let Some(auth) = req.uri().authority().cloned() { + let mut parts = Parts::default(); + parts.authority = Some(auth); + *req.uri_mut() = Uri::from_parts(parts).expect("authority is valid uri"); + } + return; + } + if is_proxied { + // Absolute-form: leave the URI as-is (scheme + authority + path). + return; + } + // Origin-form: strip scheme + authority, keep path + query (defaulting + // to "/" if the URI had nothing after the authority). + let path_and_query = req + .uri() + .path_and_query() + .filter(|p| p.as_str() != "/") + .cloned(); + *req.uri_mut() = match path_and_query { + Some(pq) => { + let mut parts = Parts::default(); + parts.path_and_query = Some(pq); + Uri::from_parts(parts).expect("path-and-query is valid uri") + } + None => Uri::default(), + }; +} + /// Tower `Service` wrapper for an H1 pool checkout. /// /// Produced by the H1 fallback leg (one per checkout from the cache). @@ -562,7 +605,7 @@ async fn eviction_task( /// H2 checkout, but both legs must produce the same `CheckoutResponse<…>` /// so `Negotiate` can compose them uniformly. The phantom slot here is /// whatever type the H2 leg's `SingletonConnection` holds, resolved by -/// type inference at the pool composition site — this wrapper ignores it. +/// type inference at the pool composition site; this wrapper ignores it. pub(crate) struct H1Checkout { conn: Option>, _marker: std::marker::PhantomData UnusedH2Phantom>, @@ -599,6 +642,12 @@ impl Service> for H1Checkout().cloned(); + let mut req = req; + // HTTP/1.1 request-target form depends on whether we're talking to + // a proxy or directly to the origin. CONNECT always uses + // authority-form; otherwise proxied = absolute-form, direct = + // origin-form. Matches `hyper-util`'s legacy client behavior. + rewrite_h1_request_target(&mut req, conn.is_proxied()); Box::pin(async move { let send_fut = conn.call(req); let resp = super::timeout::maybe_timeout_future( @@ -755,7 +804,7 @@ where // Err flips `is_closed` so `Cached::Drop` skips reinsertion, // shrinking the set by one. // - Singleton clears to `Empty` on `poll_ready` Err, forcing - // the next `call` to run a fresh handshake — a handshake + // the next `call` to run a fresh handshake; a handshake // failure surfaces as an error from `svc.call` (not the // inner `poll_ready`), which we propagate. let mut conn = loop { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 46fdd4170e8..a94e4d1d7bf 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -24,7 +24,7 @@ use super::BoxError; /// A duration paired with the sleep implementation used to realize it. /// -/// This type makes "timeout without sleep impl" unrepresentable — you cannot +/// This type makes "timeout without sleep impl" unrepresentable: you cannot /// construct one without committing to a way to actually wait. Created at the /// adapter layer from `HttpConnectorSettings` + `RuntimeComponents::sleep_impl()` /// and passed down the pool stack where timeouts are applied. @@ -52,7 +52,7 @@ impl TimeoutContext { pub(crate) struct ConnectCtx { pub(crate) uri: http_1x::Uri, /// Bounds new-connection establishment (TCP + TLS handshake). `None` - /// means no connect timeout — cached connections skip the connector + /// means no connect timeout; cached connections skip the connector /// entirely so this is automatically a no-op on cache hit. pub(crate) connect_timeout: Option, } @@ -71,13 +71,13 @@ impl ConnectCtx { /// /// `Some(...)` means: once the connection is selected (cache hit or fresh /// handshake), wrap `conn.call(req)` with this timeout. Bounds request-write -/// + response-headers-wait only — does NOT include pool acquire or connect +/// + response-headers-wait only. Does NOT include pool acquire or connect /// establishment (those have their own timeouts). #[derive(Clone, Debug)] pub(crate) struct ReadTimeoutHint(pub(crate) TimeoutContext); /// Permits acquired from connection limit semaphores. -/// Held for the lifetime of the connection — dropped when the connection is dropped. +/// Held for the lifetime of the connection; dropped when the connection is dropped. pub(crate) struct ConnectionPermit { _global: Option, _per_host: Option, @@ -110,8 +110,8 @@ impl std::error::Error for PoisonedError {} /// Metadata about a connection captured at establishment time. /// -/// Captured between TLS connector output and protocol handshake — the last point -/// where the raw transport stream is accessible. +/// Captured between TLS connector output and protocol handshake (the last point +/// where the raw transport stream is accessible). #[derive(Debug, Clone)] pub(crate) struct ConnectionInfo { /// Remote address of the peer. `None` when the underlying connector @@ -119,12 +119,19 @@ pub(crate) struct ConnectionInfo { pub(crate) remote_addr: Option, /// Local address of this end of the connection. pub(crate) local_addr: Option, + /// `true` when this connection is to an HTTP proxy server (rather than + /// directly to the origin). Drives request-target form selection: H1 + /// requests dispatched on a proxied connection use absolute-form URIs; + /// direct connections use origin-form. Populated from + /// [`hyper_util::client::legacy::connect::Connected::is_proxied`] at + /// handshake. + pub(crate) is_proxied: bool, // Future: tls_info, timing (tcp_connect_duration, tls_handshake_duration) } /// A one-shot "this connection is dead, don't reuse it" flag. /// -/// Shared via `Arc` — all clones observe and control the same flag. +/// Shared via `Arc`; all clones observe and control the same flag. /// `ManagedConnection` holds one and hands out clones (via `metadata()`) /// that let the adapter layer mark the connection poisoned through /// smithy's `ConnectionMetadata::poison_fn`. On next checkout / return, @@ -246,7 +253,7 @@ impl ManagedConnection { /// Timestamp of the last return-to-idle (or creation). /// - /// Guaranteed meaningful for any connection sitting in the pool — the + /// Guaranteed meaningful for any connection sitting in the pool: the /// initial value is `created_at`, and every return-to-pool overwrites /// it via [`Self::mark_idle`]. pub(crate) fn idle_at(&self) -> Instant { @@ -289,18 +296,14 @@ impl ManagedConnection { /// /// The returned metadata captures a clone of the `PoisonPill`, so /// calling `ConnectionMetadata::poison()` flips this connection's - /// poison flag — the same flag the pool checks on checkout/return. + /// poison flag (the same flag the pool checks on checkout/return). /// Address fields are copied. pub(crate) fn metadata(&self) -> ConnectionMetadata { let poison = self.poison.clone(); let conn_id = self.conn_id; let remote = self.info.remote_addr; let mut builder = ConnectionMetadata::builder() - // Hardcoded `false`: v2 does not yet plumb proxy configuration - // through the pool stack. When proxy support lands, - // `ConnectionInfo` should carry an `is_proxied` flag populated - // at connection establishment and passed through here. - .proxied(false) + .proxied(self.info.is_proxied) .poison_fn(move || { tracing::debug!(conn_id, ?remote, "v2 pool: connection poisoned"); poison.poison(); @@ -310,6 +313,12 @@ impl ManagedConnection { .set_local_addr(self.info.local_addr); builder.build() } + + /// `true` when the underlying connection is to an HTTP proxy. Used by + /// H1 dispatch to choose absolute-form URIs over origin-form. + pub(crate) fn is_proxied(&self) -> bool { + self.info.is_proxied + } } impl Clone for ManagedConnection { @@ -382,6 +391,15 @@ impl CachedConnection { .inner() .metadata() } + + /// Forwards [`ManagedConnection::is_proxied`]. + pub(crate) fn is_proxied(&self) -> bool { + self.inner + .as_ref() + .expect("CachedConnection is_proxied after drop") + .inner() + .is_proxied() + } } impl Service for CachedConnection @@ -411,7 +429,7 @@ impl Drop for CachedConnection { cached.discard(); } else { // Stamp the return-to-idle moment so the eviction task - // can measure elapsed idleness. Healthy connection — the + // can measure elapsed idleness. Healthy connection: the // Cached drops normally after this, returning it to the // cache's idle set. managed.mark_idle(); @@ -429,7 +447,7 @@ impl Drop for CachedConnection { /// [`Self::current`]. /// /// Necessary because `hyper_util::client::pool::singleton::Singled<…>` is -/// opaque to us — we can't reach through it to the underlying +/// opaque to us; we can't reach through it to the underlying /// `ManagedConnection` to build its metadata at checkout time. The /// handshake side has direct access, so it publishes and the checkout /// side reads. @@ -439,7 +457,7 @@ impl Drop for CachedConnection { /// `Singleton::retain` predicate returning false), its state transitions /// from `Made(svc)` back to `Empty`, dropping the old service. The next /// request's `call` runs a fresh handshake through `H2ConnectAndHandshake`, -/// which publishes new metadata — overwriting whatever was there before. +/// which publishes new metadata, overwriting whatever was there before. /// State an H2 checkout (`SingletonConnection`) needs but cannot reach /// through the opaque `Singled<…>` to read from the underlying /// `ManagedConnection`. Published by the H2 handshake on each new @@ -499,7 +517,7 @@ impl H2ConnectionRef { /// when the response body's guard releases. /// /// Drop decrements `active_streams` and, on the 1 → 0 transition, stamps -/// `idle_at` — the moment the connection becomes truly idle. +/// `idle_at` (the moment the connection becomes truly idle). struct DispatchGuard { active_streams: Arc, idle_at: Arc>, @@ -537,7 +555,7 @@ impl Drop for DispatchGuard { /// /// H2 is multiplexed: the underlying `ManagedConnection` stays resident in /// `Singleton` for the duration of every concurrent request, so "presence -/// in a cache" — the idleness signal we use for H1 — doesn't apply. +/// in a cache" (the idleness signal we use for H1) doesn't apply. /// Instead, `SingletonConnection` mints a [`DispatchGuard`] on each /// `call`; the guard's lifetime tracks one in-flight request and its /// drop releases the corresponding `active_streams` slot, stamping @@ -582,9 +600,9 @@ impl SingletonConnection { /// Metadata for the H2 connection this checkout was issued against. /// /// Returns `None` only if `SingletonConnection` was constructed before - /// any handshake had published — which doesn't happen in normal flow, + /// any handshake had published (which doesn't happen in normal flow, /// since `Singleton::call` always completes a handshake before - /// yielding the `Singled` we wrap. + /// yielding the `Singled` we wrap). pub(crate) fn metadata(&self) -> Option { self.state.as_ref().map(|s| s.metadata.clone()) } @@ -618,7 +636,7 @@ pin_project! { /// /// When a checked-out pool connection's `Service::call` returns, the /// response head is available but the body may still be streaming over - /// the same underlying HTTP connection. For H1 this is critical — if the + /// the same underlying HTTP connection. For H1 this is critical: if the /// `CachedConnection` drops at that point, the connection returns to the /// pool mid-body-stream and the next checkout would be handed a still-busy /// connection. @@ -694,7 +712,7 @@ pub(crate) type CheckoutResponse = http_1x::Response` that returns `Ok(())` — stand-in for - /// `Singled::call`. We only care about the wrapper's stream-count + /// A minimal `Service<()>` that returns `Ok(())` (stand-in for + /// `Singled::call`). We only care about the wrapper's stream-count /// side effects. #[derive(Clone, Default)] struct OkService; @@ -763,7 +781,7 @@ mod tests { /// Constructing without dispatching leaves no `DispatchGuard`, so /// dropping is a no-op. The post-checkout `poll_ready` retry loop - /// relies on this — a checkout discarded before `call` must not + /// relies on this: a checkout discarded before `call` must not /// underflow the counter. #[tokio::test] async fn singleton_drop_without_dispatch_is_noop() { @@ -782,8 +800,8 @@ mod tests { /// Concurrent dispatches against the same H2 connection (simulating /// multiplexed requests) all land on the same counter. `idle_at` is - /// stamped only on the LAST `DispatchGuard` drop — the 1 → 0 - /// transition. + /// stamped only on the LAST `DispatchGuard` drop (the 1 → 0 + /// transition). #[tokio::test] async fn idle_at_stamped_only_on_last_dispatch_release() { let h2_ref = H2ConnectionRef::new(); @@ -804,7 +822,7 @@ mod tests { // Sleep so any stamp is distinguishable from construction time. tokio::time::sleep(std::time::Duration::from_millis(5)).await; - // Release two — idle_at must NOT yet be stamped (counter > 0). + // Release two; idle_at must NOT yet be stamped (counter > 0). drop(a); drop(b); assert_eq!(active.load(Ordering::Acquire), 1); diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 980f935a9f1..b00dc5600ed 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -30,7 +30,7 @@ use super::BoxError; /// cloned into individual layers (the H1/H2 handshake services) at /// construction. /// -/// Cheap to clone — one `Arc` per primitive. +/// Cheap to clone (one `Arc` per primitive). #[derive(Clone)] pub(crate) struct PoolHooks { conn_id_counter: Arc, @@ -58,7 +58,7 @@ impl PoolHooks { /// /// Target type is `ConnectCtx`: the inner TCP connector is `Service`, /// so this layer extracts `ctx.uri` to pass through. Per-operation timeouts -/// on the context are applied here — `connect_timeout` wraps the inner +/// on the context are applied here: `connect_timeout` wraps the inner /// connector call (which includes TCP + TLS since the inner is the /// TLS-wrapped connector). Cache hits skip this layer entirely, so /// `connect_timeout` is automatically new-connection-only. @@ -178,7 +178,7 @@ impl Service> for H1SendRequest { /// Tower `Service` adapter for `hyper::client::conn::http2::SendRequest`. /// -/// Clone is supported — HTTP/2 multiplexes requests over a single connection. +/// Clone is supported because HTTP/2 multiplexes requests over a single connection. #[derive(Clone)] pub(crate) struct H2SendRequest { inner: hyper::client::conn::http2::SendRequest, @@ -211,12 +211,14 @@ impl Service> for H2SendRequest { /// Extract `ConnectionInfo` from a just-connected IO stream. fn capture_info(io: &IO) -> ConnectionInfo { let connected = io.connected(); + let is_proxied = connected.is_proxied(); let mut extras = http_1x::Extensions::new(); connected.get_extras(&mut extras); let http_info = extras.get::(); ConnectionInfo { remote_addr: http_info.map(|i| i.remote_addr()), local_addr: http_info.map(|i| i.local_addr()), + is_proxied, } } @@ -229,7 +231,7 @@ fn spawn_driver(future: impl Future + Send + 'static) { /// Connects and performs an HTTP/1.1 handshake. /// -/// The connector is expected to return `(IO, Arc)` — typically +/// The connector is expected to return `(IO, Arc)`, typically /// produced by [`ConnectionLimit`] wrapping a TCP/TLS connector. pub(crate) struct H1ConnectAndHandshake { connector: C, @@ -317,7 +319,7 @@ where /// Connects and performs an HTTP/2 handshake. /// /// Pinned to `Service<()>` because this service sits in the Negotiate -/// upgrade path — the connection is already established via the shared +/// upgrade path: the connection is already established via the shared /// `Inspected` slot. pub(crate) struct H2ConnectAndHandshake { connector: C, @@ -463,7 +465,7 @@ mod tests { } /// Without a connect_timeout set, `ConnectionLimit` passes through the - /// connector call with no wrapping — verify the absence of a timeout + /// connector call with no wrapping. Verify the absence of a timeout /// means the slow connector stays pending (we just check we can start /// the call; with `start_paused`, nothing advances so the future is /// still pending after a short yield). diff --git a/rust-runtime/aws-smithy-http-client/src/client/proxy.rs b/rust-runtime/aws-smithy-http-client/src/client/proxy.rs index dfeaacfbd4e..0e69ca38135 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/proxy.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/proxy.rs @@ -519,6 +519,34 @@ impl ProxyConfig { } } +/// Inject `Proxy-Authorization` for an HTTP-through-proxy request when +/// `matcher` carries credentials for the request URI. HTTPS-through-proxy +/// uses CONNECT tunneling and authenticates during tunnel setup, so this +/// is a no-op for HTTPS URIs. Existing `Proxy-Authorization` headers are +/// preserved. +pub(crate) fn add_proxy_auth_header( + request: &mut http_1x::Request, + matcher: &Matcher, +) { + if request.uri().scheme() != Some(&http_1x::uri::Scheme::HTTP) { + return; + } + if request + .headers() + .contains_key(http_1x::header::PROXY_AUTHORIZATION) + { + return; + } + if let Some(intercept) = matcher.intercept(request.uri()) { + if let Some(auth_header) = intercept.basic_auth() { + request + .headers_mut() + .insert(http_1x::header::PROXY_AUTHORIZATION, auth_header.clone()); + tracing::debug!(uri = %request.uri(), "added proxy authentication header"); + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index 4c9735a9fd4..814b20895a2 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -32,7 +32,7 @@ use crate::tls::TlsContext; /// Builder for the v2 HTTP client. /// /// Provides a flat, unified configuration surface for connection management, -/// TLS, timeouts, and pool behavior — all knobs in one place. +/// TLS, timeouts, and pool behavior (all knobs in one place). /// /// # Examples /// @@ -50,6 +50,7 @@ pub struct BuilderV2 { tcp_nodelay: bool, max_connections: Option, max_connections_per_host: Option, + proxy_config: Option, tls: Tls, } @@ -60,6 +61,7 @@ impl Default for BuilderV2 { tcp_nodelay: true, max_connections: None, max_connections_per_host: None, + proxy_config: None, tls: TlsUnset {}, } } @@ -97,12 +99,36 @@ impl BuilderV2 { /// Set the maximum number of concurrent connections per host. /// /// Each unique (scheme, authority) pair has an independent connection - /// budget. This limit is independent of `max_connections` — both can + /// budget. This limit is independent of `max_connections`; both can /// be set and are enforced simultaneously. pub fn max_connections_per_host(mut self, n: usize) -> Self { self.max_connections_per_host = Some(n); self } + + /// Configure proxy support. + /// + /// Threads `config` through the connector chain: HTTP/HTTPS proxies + /// route requests via the proxy server; HTTPS-over-proxy uses + /// `CONNECT` tunneling. The per-host proxy decision is stable for + /// this builder's lifetime: all connections to a given authority + /// go through the same path. + /// + /// To use environment variables (`HTTP_PROXY`, `HTTPS_PROXY`, + /// `NO_PROXY`), pass `ProxyConfig::from_env()`. + pub fn proxy_config(mut self, config: super::proxy::ProxyConfig) -> Self { + self.proxy_config = Some(config); + self + } + + /// Mutable-reference variant of [`proxy_config`](Self::proxy_config). + pub fn set_proxy_config( + &mut self, + config: Option, + ) -> &mut Self { + self.proxy_config = config; + self + } } impl BuilderV2 { @@ -118,6 +144,7 @@ impl BuilderV2 { tcp_nodelay: self.tcp_nodelay, max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, + proxy_config: self.proxy_config, tls: TlsProviderSelected { provider, context: TlsContext::default(), @@ -135,8 +162,9 @@ impl BuilderV2 { max_connections_per_host: self.max_connections_per_host, pool_idle_timeout: self.pool_idle_timeout, }; - let pool = pool::build_pool(tcp, config); - V2HttpClient::new(pool).into_shared() + let proxy_matcher = proxy_matcher_from(&self.proxy_config); + let pool = build_http_pool_with_proxy(tcp, &self.proxy_config, config); + V2HttpClient::new(pool, proxy_matcher).into_shared() } /// Build an HTTP client without TLS, using a custom DNS resolver. @@ -144,7 +172,7 @@ impl BuilderV2 { /// Mirrors `build_http` but threads a custom `ResolveDns`. Intended /// for tests that use `WireMockServer::dns_resolver()` to redirect /// all lookups to a loopback mock. Public API for non-TLS resolver - /// plumbing is `#[doc(hidden)]` — TLS is the common case. + /// plumbing is `#[doc(hidden)]`; TLS is the common case. #[doc(hidden)] pub fn build_http_with_resolver( self, @@ -158,8 +186,47 @@ impl BuilderV2 { max_connections_per_host: self.max_connections_per_host, pool_idle_timeout: self.pool_idle_timeout, }; - let pool = pool::build_pool(tcp, config); - V2HttpClient::new(pool).into_shared() + let proxy_matcher = proxy_matcher_from(&self.proxy_config); + let pool = build_http_pool_with_proxy(tcp, &self.proxy_config, config); + V2HttpClient::new(pool, proxy_matcher).into_shared() + } +} + +/// Build a no-TLS pool that honors `proxy_config`. Wraps the TCP connector +/// with an HTTP proxy connector when configured. Emits a warning if an +/// HTTPS proxy is set without a TLS provider; connections to such a +/// proxy will fail at handshake time. +fn build_http_pool_with_proxy( + tcp: HyperHttpConnector, + proxy_config: &Option, + config: PoolConfig, +) -> ConnectionPool +where + R: Clone + Send + Sync + 'static, + R: tower::Service, + R::Response: Iterator, + R::Future: Send, + R::Error: Into, +{ + use super::proxy; + + let proxy_config = proxy_config + .clone() + .unwrap_or_else(proxy::ProxyConfig::disabled); + + if proxy_config.requires_tls() { + tracing::warn!( + "HTTPS proxy configured but no TLS provider set. \ + Connections to HTTPS proxy servers will fail. \ + Consider configuring a TLS provider to enable TLS support." + ); + } + + if proxy_config.is_disabled() { + pool::build_pool(tcp, config) + } else { + let connector = super::connect::HttpProxyConnector::new(tcp, proxy_config); + pool::build_pool(connector, config) } } @@ -196,7 +263,7 @@ impl BuilderV2 { R: tower::Service, R::Response: Iterator, R::Future: Send, - R::Error: Into>, + R::Error: Into, { use super::proxy; @@ -206,6 +273,12 @@ impl BuilderV2 { pool_idle_timeout: self.pool_idle_timeout, }; + let proxy_config = self + .proxy_config + .clone() + .unwrap_or_else(proxy::ProxyConfig::disabled); + let proxy_matcher = proxy_matcher_from(&self.proxy_config); + match &self.tls.provider { #[cfg(any( feature = "rustls-aws-lc", @@ -217,20 +290,20 @@ impl BuilderV2 { tcp, crypto_mode.clone(), &self.tls.context, - proxy::ProxyConfig::disabled(), + proxy_config, ); let pool = pool::build_pool(connector, config); - V2HttpClient::new(pool).into_shared() + V2HttpClient::new(pool, proxy_matcher).into_shared() } #[cfg(feature = "s2n-tls")] tls::Provider::S2nTls => { let connector = tls::s2n_tls_provider::build_connector::wrap_connector( tcp, &self.tls.context, - proxy::ProxyConfig::disabled(), + proxy_config, ); let pool = pool::build_pool(connector, config); - V2HttpClient::new(pool).into_shared() + V2HttpClient::new(pool, proxy_matcher).into_shared() } // Provider is #[non_exhaustive]; this arm is unreachable when any // TLS feature is enabled (which is required to construct a Provider). @@ -255,8 +328,11 @@ struct PooledConnector { read_timeout: Option, /// Present when any timeout is configured. Required to construct /// `Timeout` futures. `V2HttpClient::http_connector` panics if - /// a timeout is set but `sleep_impl` is absent — matches v1. + /// a timeout is set but `sleep_impl` is absent (matches v1). sleep_impl: Option, + /// Proxy URL matcher, cloned from `V2HttpClient`. Drives the + /// `Proxy-Authorization` header injection in `call`. + proxy_matcher: Option>, } impl std::fmt::Debug for PooledConnector { @@ -271,6 +347,7 @@ impl HttpConnector for PooledConnector { let connect_timeout = self.connect_timeout; let read_timeout = self.read_timeout; let sleep_impl = self.sleep_impl.clone(); + let proxy_matcher = self.proxy_matcher.clone(); HttpConnectorFuture::new(async move { let mut request = request .try_into_http1x() @@ -279,6 +356,14 @@ impl HttpConnector for PooledConnector { // Save the full URI for pool routing before rewriting to origin-form. let full_uri = request.uri().clone(); + // Inject Proxy-Authorization for HTTP-through-proxy requests. + // Shared with v1 via [`super::proxy::add_proxy_auth_header`]; + // HTTPS-through-proxy uses CONNECT tunneling and authenticates + // inside the tunnel-establishment exchange, not via this header. + if let Some(matcher) = proxy_matcher.as_ref() { + super::proxy::add_proxy_auth_header(&mut request, matcher); + } + // Create a request-scoped capture the pool will fill in with // the `ConnectionMetadata` for the selected connection, and // wire `CaptureSmithyConnection` to read from it. When the @@ -294,7 +379,7 @@ impl HttpConnector for PooledConnector { // Attach the per-op read timeout hint (if any) as a request // extension. `H{1,2}Checkout::call` reads it and wraps the - // request dispatch — bounds request-write + response-headers, + // request dispatch: bounds request-write + response-headers, // nothing else. if let Some((duration, sleep)) = read_timeout.zip(sleep_impl.clone()) { request @@ -319,16 +404,11 @@ impl HttpConnector for PooledConnector { } } - // Rewrite URI to origin-form for HTTP/1.1. - let origin = match request.uri().path_and_query() { - Some(pq) => { - let mut parts = http_1x::uri::Parts::default(); - parts.path_and_query = Some(pq.clone()); - http_1x::Uri::from_parts(parts).expect("path is valid uri") - } - None => http_1x::Uri::from_static("/"), - }; - *request.uri_mut() = origin; + // URI request-target rewriting happens in `H1Checkout::call` + // because the correct form depends on whether the connection + // landed on a proxy (absolute-form) or direct origin + // (origin-form). H2 uses pseudo-headers derived by hyper from + // the absolute URI, so it doesn't need rewriting here. // Build the connect context: routing URI + connect timeout. // Cache hits skip the connector entirely so connect_timeout is @@ -354,6 +434,18 @@ impl HttpConnector for PooledConnector { // HttpClient impl // --------------------------------------------------------------------------- +/// Build the proxy URL matcher from a `ProxyConfig`, returning `None` when +/// no proxy is configured. Wrapped in `Arc` so each `PooledConnector` +/// constructed by `http_connector` can share the (immutable) matcher +/// without per-call allocation. +fn proxy_matcher_from( + proxy_config: &Option, +) -> Option> { + proxy_config + .as_ref() + .map(|c| Arc::new(c.clone().into_hyper_util_matcher())) +} + /// Returns the port if it is not the default for the scheme. fn get_non_default_port(uri: &http_1x::Uri) -> Option> { match (uri.port().map(|p| p.as_u16()), uri.scheme()) { @@ -366,6 +458,12 @@ fn get_non_default_port(uri: &http_1x::Uri) -> Option> /// Smithy `HttpClient` backed by the v2 connection pool. struct V2HttpClient { pool: Arc, + /// Proxy URL matcher, present iff the builder was configured with a + /// `proxy_config`. Used by `PooledConnector` to inject + /// `Proxy-Authorization` headers on HTTP-through-proxy requests + /// (HTTPS-CONNECT auth is handled by the proxy connector itself + /// during tunnel setup). + proxy_matcher: Option>, } impl std::fmt::Debug for V2HttpClient { @@ -375,9 +473,13 @@ impl std::fmt::Debug for V2HttpClient { } impl V2HttpClient { - fn new(pool: ConnectionPool) -> Self { + fn new( + pool: ConnectionPool, + proxy_matcher: Option>, + ) -> Self { Self { pool: Arc::new(pool), + proxy_matcher, } } } @@ -404,6 +506,7 @@ impl HttpClient for V2HttpClient { connect_timeout, read_timeout, sleep_impl, + proxy_matcher: self.proxy_matcher.clone(), }) } diff --git a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs index bbc06530978..26bf6037947 100644 --- a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs +++ b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs @@ -1196,3 +1196,421 @@ async fn test_connect_uri_form_rustls() { async fn test_connect_uri_form_s2n_tls() { run_connect_uri_form_test(tls::Provider::S2nTls, "s2n-tls").await; } + + +// ================================================================================================ +// V2 client proxy parity tests +// ================================================================================================ + +async fn make_v2_http_request_through_proxy( + proxy_config: ProxyConfig, + target_url: &str, +) -> Result<(StatusCode, String), Box> { + use aws_smithy_http_client::v2::BuilderV2; + + let http_client = BuilderV2::new().proxy_config(proxy_config).build_http(); + let connector_settings = HttpConnectorSettings::builder().build(); + let runtime_components = RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(SystemTimeSource::new())) + .build() + .unwrap(); + let http_connector = http_client.http_connector(&connector_settings, &runtime_components); + + let request = HttpRequest::get(target_url) + .map_err(|e| Box::new(e) as Box)?; + let response = http_connector.call(request).await?; + let status = response.status(); + let body_bytes = response.into_body().collect().await?.to_bytes(); + let body_string = String::from_utf8(body_bytes.to_vec())?; + Ok((status.into(), body_string)) +} + +#[cfg(any( + feature = "rustls-aws-lc", + feature = "rustls-aws-lc-fips", + feature = "rustls-ring" +))] +async fn make_v2_https_request_through_proxy( + proxy_config: ProxyConfig, + target_url: &str, + tls_provider: tls::Provider, +) -> Result<(StatusCode, String), Box> { + use aws_smithy_http_client::v2::BuilderV2; + + let http_client = BuilderV2::new() + .tls_provider(tls_provider) + .proxy_config(proxy_config) + .build_https(); + let connector_settings = HttpConnectorSettings::builder().build(); + let runtime_components = RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(SystemTimeSource::new())) + .build() + .unwrap(); + let http_connector = http_client.http_connector(&connector_settings, &runtime_components); + + let request = HttpRequest::get(target_url) + .map_err(|e| Box::new(e) as Box)?; + let response = http_connector.call(request).await?; + let status = response.status(); + let body_bytes = response.into_body().collect().await?.to_bytes(); + let body_string = String::from_utf8(body_bytes.to_vec())?; + Ok((status.into(), body_string)) +} + +#[tokio::test] +async fn test_v2_http_proxy_basic_request() { + let mock_proxy = MockProxyServer::new(|req| { + assert_eq!(req.method, "GET"); + assert_eq!(req.uri, "http://aws.amazon.com/v2/api"); + Response::builder() + .status(StatusCode::OK) + .body("v2 proxied response".to_string()) + .unwrap() + }) + .await; + + let proxy_config = ProxyConfig::http(format!("http://{}", mock_proxy.addr())).unwrap(); + let (status, body) = + make_v2_http_request_through_proxy(proxy_config, "http://aws.amazon.com/v2/api") + .await + .expect("v2 HTTP request through proxy should succeed"); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body, "v2 proxied response"); + assert_eq!(mock_proxy.requests()[0].uri, "http://aws.amazon.com/v2/api"); +} + +#[tokio::test] +async fn test_v2_proxy_authentication() { + let mock_proxy = MockProxyServer::with_auth_validation("v2user", "v2pass").await; + + let proxy_config = ProxyConfig::http(format!("http://{}", mock_proxy.addr())) + .unwrap() + .with_basic_auth("v2user", "v2pass"); + + let (status, _) = + make_v2_http_request_through_proxy(proxy_config, "http://aws.amazon.com/auth/test") + .await + .expect("authenticated v2 proxy request should succeed"); + assert_eq!(status, StatusCode::OK); + + let requests = mock_proxy.requests(); + let auth_header = requests[0] + .headers + .get("proxy-authorization") + .expect("Proxy-Authorization header should be set"); + let expected = format!( + "Basic {}", + base64::engine::general_purpose::STANDARD.encode("v2user:v2pass") + ); + assert_eq!(auth_header, &expected); +} + +#[tokio::test] +async fn test_v2_proxy_disabled_uses_direct_connection() { + let mock_proxy = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("should never reach proxy".to_string()) + .unwrap() + }) + .await; + let direct_mock = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("direct response".to_string()) + .unwrap() + }) + .await; + + let target_url = format!("http://{}/direct", direct_mock.addr()); + let (status, body) = + make_v2_http_request_through_proxy(ProxyConfig::disabled(), &target_url) + .await + .expect("direct v2 request should succeed"); + assert_eq!(status, StatusCode::OK); + assert_eq!(body, "direct response"); + assert_eq!(mock_proxy.requests().len(), 0); + assert_eq!(direct_mock.requests().len(), 1); +} + +#[tokio::test] +async fn test_v2_proxy_from_environment_variables() { + let mock_proxy = MockProxyServer::new(|req| { + assert_eq!(req.uri, "http://aws.amazon.com/env/api"); + Response::builder() + .status(StatusCode::OK) + .body("env-proxy response".to_string()) + .unwrap() + }) + .await; + let proxy_url = format!("http://{}", mock_proxy.addr()); + + let result = with_env_vars(&[("HTTP_PROXY", &proxy_url)], || async { + let proxy_config = ProxyConfig::from_env(); + make_v2_http_request_through_proxy(proxy_config, "http://aws.amazon.com/env/api").await + }) + .await; + + let (status, body) = result.expect("env-var v2 proxy request should succeed"); + assert_eq!(status, StatusCode::OK); + assert_eq!(body, "env-proxy response"); + assert_eq!(mock_proxy.requests().len(), 1); +} + +#[tokio::test] +async fn test_v2_no_proxy_bypass_rules() { + let mock_proxy = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("proxied".to_string()) + .unwrap() + }) + .await; + let direct_mock = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("direct".to_string()) + .unwrap() + }) + .await; + + let proxy_url = format!("http://{}", mock_proxy.addr()); + let result = with_env_vars( + &[("HTTP_PROXY", &proxy_url), ("NO_PROXY", "127.0.0.1")], + || async { + let proxy_config = ProxyConfig::from_env(); + let target_url = format!("http://{}/bypassed", direct_mock.addr()); + make_v2_http_request_through_proxy(proxy_config, &target_url).await + }, + ) + .await; + + let (status, body) = result.expect("bypassed request should succeed"); + assert_eq!(status, StatusCode::OK); + assert_eq!(body, "direct"); + assert_eq!(mock_proxy.requests().len(), 0, "proxy should be bypassed"); + assert_eq!(direct_mock.requests().len(), 1); +} + +#[tokio::test] +async fn test_v2_proxy_connection_failure() { + let proxy_config = ProxyConfig::http("http://127.0.0.1:1").unwrap(); + let result = + make_v2_http_request_through_proxy(proxy_config, "http://aws.amazon.com/fail").await; + assert!(result.is_err(), "connection to non-existent proxy should fail"); +} + +#[tokio::test] +async fn test_v2_proxy_authentication_failure() { + let mock_proxy = MockProxyServer::with_auth_validation("correct", "password").await; + + let proxy_config = ProxyConfig::http(format!("http://{}", mock_proxy.addr())) + .unwrap() + .with_basic_auth("wrong", "creds"); + + let (status, _) = + make_v2_http_request_through_proxy(proxy_config, "http://aws.amazon.com/auth/fail") + .await + .expect("request should complete (proxy returns 407)"); + assert_eq!(status, StatusCode::PROXY_AUTHENTICATION_REQUIRED); +} + +#[tokio::test] +async fn test_v2_explicit_proxy_disable_overrides_environment() { + let mock_proxy = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("via proxy".to_string()) + .unwrap() + }) + .await; + let direct_mock = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("direct".to_string()) + .unwrap() + }) + .await; + + let proxy_url = format!("http://{}", mock_proxy.addr()); + let result = with_env_vars(&[("HTTP_PROXY", &proxy_url)], || async { + let proxy_config = ProxyConfig::disabled(); + let target_url = format!("http://{}/path", direct_mock.addr()); + make_v2_http_request_through_proxy(proxy_config, &target_url).await + }) + .await; + + let (status, body) = result.expect("direct request should succeed"); + assert_eq!(status, StatusCode::OK); + assert_eq!(body, "direct"); + assert_eq!(mock_proxy.requests().len(), 0); + assert_eq!(direct_mock.requests().len(), 1); +} + +#[tokio::test] +async fn test_v2_http_proxy_absolute_uri_form() { + let mock_proxy = MockProxyServer::new(|req| { + assert_eq!( + req.uri, "http://aws.amazon.com/path?query=1", + "proxy should receive absolute-form URI" + ); + Response::builder() + .status(StatusCode::OK) + .body("ok".to_string()) + .unwrap() + }) + .await; + + let proxy_config = ProxyConfig::http(format!("http://{}", mock_proxy.addr())).unwrap(); + let (status, _) = make_v2_http_request_through_proxy( + proxy_config, + "http://aws.amazon.com/path?query=1", + ) + .await + .expect("request should succeed"); + assert_eq!(status, StatusCode::OK); +} + +#[tokio::test] +async fn test_v2_direct_http_origin_uri_form() { + let direct_mock = MockProxyServer::new(|req| { + assert_eq!( + req.uri, "/path?query=1", + "direct server should receive origin-form URI" + ); + Response::builder() + .status(StatusCode::OK) + .body("ok".to_string()) + .unwrap() + }) + .await; + + let target_url = format!("http://{}/path?query=1", direct_mock.addr()); + let (status, _) = make_v2_http_request_through_proxy(ProxyConfig::disabled(), &target_url) + .await + .expect("request should succeed"); + assert_eq!(status, StatusCode::OK); +} + +#[tokio::test] +async fn test_v2_set_proxy_config_none_clears_proxy() { + use aws_smithy_http_client::v2::BuilderV2; + + let mock_proxy = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("via proxy".to_string()) + .unwrap() + }) + .await; + let direct_mock = MockProxyServer::new(|_| { + Response::builder() + .status(StatusCode::OK) + .body("direct".to_string()) + .unwrap() + }) + .await; + + let mut builder = BuilderV2::new(); + builder.set_proxy_config(Some( + ProxyConfig::http(format!("http://{}", mock_proxy.addr())).unwrap(), + )); + builder.set_proxy_config(None); + let http_client = builder.build_http(); + + let connector_settings = HttpConnectorSettings::builder().build(); + let runtime_components = RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(SystemTimeSource::new())) + .build() + .unwrap(); + let http_connector = http_client.http_connector(&connector_settings, &runtime_components); + + let target_url = format!("http://{}/path", direct_mock.addr()); + let request = HttpRequest::get(&target_url).unwrap(); + let response = http_connector.call(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK.into()); + assert_eq!(mock_proxy.requests().len(), 0); + assert_eq!(direct_mock.requests().len(), 1); +} + +#[cfg(any( + feature = "rustls-aws-lc", + feature = "rustls-aws-lc-fips", + feature = "rustls-ring" +))] +#[tokio::test] +async fn test_v2_https_connect_with_auth() { + let mock_proxy = MockProxyServer::new(|req| { + assert_eq!(req.method, "CONNECT"); + assert_eq!(req.uri, "secure.aws.amazon.com:443"); + + let expected_auth = format!( + "Basic {}", + base64::prelude::BASE64_STANDARD.encode("connectuser:connectpass") + ); + assert_eq!(req.headers.get("proxy-authorization"), Some(&expected_auth)); + + Response::builder() + .status(StatusCode::BAD_REQUEST) + .body("CONNECT tunnel setup failed".to_string()) + .unwrap() + }) + .await; + + let proxy_config = ProxyConfig::all(format!("http://{}", mock_proxy.addr())) + .unwrap() + .with_basic_auth("connectuser", "connectpass"); + + let tls_provider = tls::Provider::Rustls(tls::rustls_provider::CryptoMode::AwsLc); + let result = make_v2_https_request_through_proxy( + proxy_config, + "https://secure.aws.amazon.com/api/secure", + tls_provider, + ) + .await; + + assert!(result.is_err(), "CONNECT tunnel should fail with 400 response"); + let requests = mock_proxy.requests(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].method, "CONNECT"); + assert_eq!(requests[0].uri, "secure.aws.amazon.com:443"); +} + +#[cfg(any( + feature = "rustls-aws-lc", + feature = "rustls-aws-lc-fips", + feature = "rustls-ring" +))] +#[tokio::test] +async fn test_v2_https_connect_auth_required() { + let mock_proxy = MockProxyServer::new(|req| { + assert_eq!(req.method, "CONNECT"); + if req.headers.get("proxy-authorization").is_none() { + return Response::builder() + .status(StatusCode::PROXY_AUTHENTICATION_REQUIRED) + .body("auth required".to_string()) + .unwrap(); + } + Response::builder() + .status(StatusCode::OK) + .body(String::new()) + .unwrap() + }) + .await; + + let proxy_config = ProxyConfig::all(format!("http://{}", mock_proxy.addr())).unwrap(); + let tls_provider = tls::Provider::Rustls(tls::rustls_provider::CryptoMode::AwsLc); + let result = make_v2_https_request_through_proxy( + proxy_config, + "https://secure.aws.amazon.com/api", + tls_provider, + ) + .await; + + assert!(result.is_err(), "should fail without proxy auth"); + let requests = mock_proxy.requests(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].method, "CONNECT"); + assert!(requests[0].headers.get("proxy-authorization").is_none()); +} From 9d053b3244b1b87d45a178914d7242d50e9a4528 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 14 May 2026 13:45:16 -0400 Subject: [PATCH 17/47] add connect test for v2 --- .../aws-smithy-http-client/src/client/v2.rs | 5 +-- .../tests/pool_behavior_test.rs | 44 +++++++++++++++++++ .../tests/proxy_tests.rs | 28 ++++++------ 3 files changed, 60 insertions(+), 17 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index 814b20895a2..3c441d200ab 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -122,10 +122,7 @@ impl BuilderV2 { } /// Mutable-reference variant of [`proxy_config`](Self::proxy_config). - pub fn set_proxy_config( - &mut self, - config: Option, - ) -> &mut Self { + pub fn set_proxy_config(&mut self, config: Option) -> &mut Self { self.proxy_config = config; self } diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index c81fa135e72..3bc872eacf9 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1230,3 +1230,47 @@ async fn v2_conn_id_is_stable_and_monotonic() { ); assert!(captured.contains("v2 pool: connection evicted")); } + +// --------------------------------------------------------------------------- +// Connect timeout +// --------------------------------------------------------------------------- + +/// Proves that `connect_timeout` from `HttpConnectorSettings` fires when the +/// TCP connection cannot be established within the deadline. Uses TEST-NET-1 +/// (192.0.2.1), a non-routable address guaranteed to black-hole SYN packets. +#[tokio::test] +async fn v2_connect_timeout() { + use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep}; + + let client = V2Client.make(ClientConfig::default()); + let components = + aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(SystemTimeSource::new())) + .with_sleep_impl(Some(SharedAsyncSleep::new(TokioSleep::new()))) + .build() + .expect("valid runtime components"); + let settings = HttpConnectorSettings::builder() + .connect_timeout(Duration::from_millis(500)) + .build(); + let connector = client.http_connector(&settings, &components); + + let start = std::time::Instant::now(); + let err = connector + .call(HttpRequest::get("http://192.0.2.1:1234/unreachable").expect("valid request")) + .await + .expect_err("connect timeout should fire against non-routable address"); + let elapsed = start.elapsed(); + + assert!( + err.is_timeout(), + "expected timeout classification, got {err:?}" + ); + assert!( + elapsed >= Duration::from_millis(450), + "timeout fired too early ({elapsed:?})" + ); + assert!( + elapsed < Duration::from_secs(5), + "timeout took too long ({elapsed:?})" + ); +} diff --git a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs index 26bf6037947..7b5be9ad518 100644 --- a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs +++ b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs @@ -1197,7 +1197,6 @@ async fn test_connect_uri_form_s2n_tls() { run_connect_uri_form_test(tls::Provider::S2nTls, "s2n-tls").await; } - // ================================================================================================ // V2 client proxy parity tests // ================================================================================================ @@ -1324,10 +1323,9 @@ async fn test_v2_proxy_disabled_uses_direct_connection() { .await; let target_url = format!("http://{}/direct", direct_mock.addr()); - let (status, body) = - make_v2_http_request_through_proxy(ProxyConfig::disabled(), &target_url) - .await - .expect("direct v2 request should succeed"); + let (status, body) = make_v2_http_request_through_proxy(ProxyConfig::disabled(), &target_url) + .await + .expect("direct v2 request should succeed"); assert_eq!(status, StatusCode::OK); assert_eq!(body, "direct response"); assert_eq!(mock_proxy.requests().len(), 0); @@ -1398,7 +1396,10 @@ async fn test_v2_proxy_connection_failure() { let proxy_config = ProxyConfig::http("http://127.0.0.1:1").unwrap(); let result = make_v2_http_request_through_proxy(proxy_config, "http://aws.amazon.com/fail").await; - assert!(result.is_err(), "connection to non-existent proxy should fail"); + assert!( + result.is_err(), + "connection to non-existent proxy should fail" + ); } #[tokio::test] @@ -1463,12 +1464,10 @@ async fn test_v2_http_proxy_absolute_uri_form() { .await; let proxy_config = ProxyConfig::http(format!("http://{}", mock_proxy.addr())).unwrap(); - let (status, _) = make_v2_http_request_through_proxy( - proxy_config, - "http://aws.amazon.com/path?query=1", - ) - .await - .expect("request should succeed"); + let (status, _) = + make_v2_http_request_through_proxy(proxy_config, "http://aws.amazon.com/path?query=1") + .await + .expect("request should succeed"); assert_eq!(status, StatusCode::OK); } @@ -1570,7 +1569,10 @@ async fn test_v2_https_connect_with_auth() { ) .await; - assert!(result.is_err(), "CONNECT tunnel should fail with 400 response"); + assert!( + result.is_err(), + "CONNECT tunnel should fail with 400 response" + ); let requests = mock_proxy.requests(); assert_eq!(requests.len(), 1); assert_eq!(requests[0].method, "CONNECT"); From f39ae98bd002ab96fc8d3c5f74733dab19e1a236 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 14 May 2026 20:37:22 -0400 Subject: [PATCH 18/47] conn id --- .../aws-smithy-http-client/src/client/pool.rs | 18 +++-- .../src/client/pool/connection.rs | 22 +++++- .../src/client/pool/handshake.rs | 8 +- .../tests/pool_behavior_test.rs | 78 +++++++++++++++++-- .../src/client/connection.rs | 45 +++++++++++ 5 files changed, 149 insertions(+), 22 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 0b672709707..574c970441c 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -183,7 +183,7 @@ where max_connections = ?config.max_connections, max_connections_per_host = ?config.max_connections_per_host, pool_idle_timeout = ?config.pool_idle_timeout, - "v2 pool: initialized" + "pool: initialized" ); let make_entry = { @@ -251,7 +251,7 @@ where conn_id = managed.conn_id(), reason, ?idle_duration, - "v2 pool: connection evicted" + "pool: connection evicted" ); } keep @@ -318,7 +318,7 @@ where conn_id = managed.conn_id(), reason, ?idle_duration, - "v2 pool: connection evicted" + "pool: connection evicted" ); } keep @@ -463,7 +463,7 @@ impl ConnectionPool { tracing::debug!( interval = ?tick, pool_idle_timeout = ?timeout, - "v2 pool: eviction task spawned" + "pool: eviction task spawned" ); tokio::spawn(eviction_task(weak, rx, timeout)); } @@ -482,7 +482,7 @@ impl ConnectionPool { tracing::debug!( authority = %key.authority, scheme = %key.scheme, - "v2 pool: host entry removed (empty after retain)" + "pool: host entry removed (empty after retain)" ); false } else { @@ -494,7 +494,7 @@ impl ConnectionPool { tracing::trace!( hosts_before = before, hosts_after = after, - "v2 pool: eviction tick complete" + "pool: eviction tick complete" ); } } @@ -541,7 +541,7 @@ async fn eviction_task( } } } - tracing::trace!("v2 pool eviction task exiting"); + tracing::trace!("pool eviction task exiting"); } /// Rewrite an HTTP/1.1 request's URI to the appropriate request-target @@ -634,6 +634,7 @@ impl Service> for H1Checkout) -> Self::Future { let mut conn = self.conn.take().expect("H1Checkout::call called twice"); + tracing::trace!(conn_id = conn.conn_id(), uri = %req.uri(), "pool: dispatching request"); // Populate the adapter-provided capture so the // `CaptureSmithyConnection` retriever can return a live // `ConnectionMetadata` pointing at THIS connection's poison pill. @@ -707,6 +708,9 @@ where fn call(&mut self, req: http_1x::Request) -> Self::Future { let mut conn = self.conn.take().expect("H2Checkout::call called twice"); + if let Some(id) = conn.metadata().and_then(|m| m.connection_id()) { + tracing::trace!(conn_id = %id, uri = %req.uri(), "pool: dispatching request"); + } // Populate the adapter-provided capture. The metadata is published // by `H2ConnectAndHandshake` when a fresh H2 connection is // established; if a request somehow reaches us before any metadata diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index a94e4d1d7bf..d925824c69e 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -12,7 +12,7 @@ use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use aws_smithy_async::rt::sleep::SharedAsyncSleep; -use aws_smithy_runtime_api::client::connection::ConnectionMetadata; +use aws_smithy_runtime_api::client::connection::{ConnectionId, ConnectionMetadata}; use aws_smithy_types::body::SdkBody; use pin_project_lite::pin_project; use tokio::sync::OwnedSemaphorePermit; @@ -304,8 +304,9 @@ impl ManagedConnection { let remote = self.info.remote_addr; let mut builder = ConnectionMetadata::builder() .proxied(self.info.is_proxied) + .connection_id(ConnectionId::new(self.conn_id)) .poison_fn(move || { - tracing::debug!(conn_id, ?remote, "v2 pool: connection poisoned"); + tracing::debug!(conn_id, ?remote, "pool: connection poisoned"); poison.poison(); }); builder @@ -374,6 +375,10 @@ pub(crate) struct CachedConnection { impl CachedConnection { pub(crate) fn new(cached: cache::Cached>) -> Self { + tracing::trace!( + conn_id = cached.inner().conn_id(), + "pool: connection reused" + ); Self { inner: Some(cached), } @@ -400,6 +405,15 @@ impl CachedConnection { .inner() .is_proxied() } + + /// Forwards [`ManagedConnection::conn_id`]. + pub(crate) fn conn_id(&self) -> u64 { + self.inner + .as_ref() + .expect("CachedConnection conn_id after drop") + .inner() + .conn_id() + } } impl Service for CachedConnection @@ -425,7 +439,7 @@ impl Drop for CachedConnection { let managed = cached.inner(); let conn_id = managed.conn_id; if managed.is_poisoned() { - tracing::debug!(conn_id, "v2 pool: connection discarded (poisoned)"); + tracing::debug!(conn_id, "pool: connection discarded (poisoned)"); cached.discard(); } else { // Stamp the return-to-idle moment so the eviction task @@ -433,7 +447,7 @@ impl Drop for CachedConnection { // Cached drops normally after this, returning it to the // cache's idle set. managed.mark_idle(); - tracing::trace!(conn_id, "v2 pool: connection returned to idle"); + tracing::trace!(conn_id, "pool: connection returned to idle"); } } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index b00dc5600ed..ea8febf3721 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -286,7 +286,7 @@ where protocol = "h1", remote = ?info.remote_addr, local = ?info.local_addr, - "v2 pool: connection established" + "pool: connection established" ); spawn_driver({ @@ -300,7 +300,7 @@ where ?remote_addr, ?local_addr, error = %e, - "connection driver error" + "pool: connection driver error" ); } } @@ -391,7 +391,7 @@ where protocol = "h2", remote = ?info.remote_addr, local = ?info.local_addr, - "v2 pool: connection established" + "pool: connection established" ); spawn_driver({ @@ -405,7 +405,7 @@ where ?remote_addr, ?local_addr, error = %e, - "connection driver error" + "pool: connection driver error" ); } } diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 3bc872eacf9..18adf16579e 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1172,19 +1172,19 @@ async fn v2_background_eviction_emits_tracing_events() { let captured = captured_str(&logs); assert!( - captured.contains("v2 pool: initialized"), + captured.contains("pool: initialized"), "expected pool init log" ); assert!( - captured.contains("v2 pool: eviction task spawned"), + captured.contains("pool: eviction task spawned"), "expected eviction task spawn log" ); assert!( - captured.contains("v2 pool: connection established"), + captured.contains("pool: connection established"), "expected connection-established log" ); assert!( - captured.contains("v2 pool: connection evicted"), + captured.contains("pool: connection evicted"), "expected 'connection evicted' log — background task should have evicted the idle connection" ); assert!( @@ -1192,7 +1192,7 @@ async fn v2_background_eviction_emits_tracing_events() { "expected idle_expired reason. captured:\n{captured}" ); assert!( - captured.contains("v2 pool: host entry removed"), + captured.contains("pool: host entry removed"), "expected host-entry-removed log once cache went empty" ); } @@ -1217,7 +1217,7 @@ async fn v2_conn_id_is_stable_and_monotonic() { let captured = captured_str(&logs); assert!( - captured.contains("v2 pool: connection established"), + captured.contains("pool: connection established"), "no 'connection established' log captured" ); assert!( @@ -1228,7 +1228,7 @@ async fn v2_conn_id_is_stable_and_monotonic() { captured.contains("conn_id=1"), "second connection (after eviction) should be conn_id=1. captured:\n{captured}" ); - assert!(captured.contains("v2 pool: connection evicted")); + assert!(captured.contains("pool: connection evicted")); } // --------------------------------------------------------------------------- @@ -1274,3 +1274,67 @@ async fn v2_connect_timeout() { "timeout took too long ({elapsed:?})" ); } + +/// `ConnectionMetadata::connection_id()` surfaces the pool-assigned id +/// through `CaptureSmithyConnection`. Sequential requests on the same +/// connection share the same id; a new connection after eviction gets a +/// different id. +#[tokio::test] +async fn v2_connection_id_surfaced_through_metadata() { + use aws_smithy_runtime_api::client::connection::ConnectionId; + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"a", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"b", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"c", + }, + ], + ) + .build() + .await; + + let client = V2Client.make(ClientConfig { + idle_timeout: Some(Duration::from_millis(80)), + ..Default::default() + }); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + // First two requests reuse the same connection. + let (_, _, meta1) = send_with_capture(&client, &url).await; + let (_, _, meta2) = send_with_capture(&client, &url).await; + let id1 = meta1 + .unwrap() + .connection_id() + .expect("v2 sets connection_id"); + let id2 = meta2 + .unwrap() + .connection_id() + .expect("v2 sets connection_id"); + assert_eq!(id1, id2, "same connection should have same id"); + assert_eq!(id1, ConnectionId::new(0)); + + // Wait for eviction, then the next request gets a new connection. + tokio::time::sleep(Duration::from_millis(200)).await; + + let (_, _, meta3) = send_with_capture(&client, &url).await; + let id3 = meta3 + .unwrap() + .connection_id() + .expect("v2 sets connection_id"); + assert_ne!( + id1, id3, + "new connection after eviction should have different id" + ); + assert_eq!(id3, ConnectionId::new(1)); +} diff --git a/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs b/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs index 7d155d22813..42435f62b9b 100644 --- a/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs +++ b/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs @@ -10,6 +10,27 @@ use std::fmt; use std::net::SocketAddr; use std::sync::{Arc, Mutex}; +/// Opaque identifier for a physical connection within a pool. +/// +/// Unique for the pool's lifetime. All requests dispatched on the same +/// connection (including H2 multiplexed requests) share the same id. +/// Useful for correlating requests with connection lifecycle tracing events. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ConnectionId(u64); + +impl ConnectionId { + /// Create a new connection id. + pub fn new(id: u64) -> Self { + Self(id) + } +} + +impl fmt::Display for ConnectionId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + /// Metadata that tracks the state of an active connection. #[derive(Clone)] pub struct ConnectionMetadata { @@ -17,6 +38,7 @@ pub struct ConnectionMetadata { remote_addr: Option, local_addr: Option, poison_fn: Arc, + connection_id: Option, } impl ConnectionMetadata { @@ -45,6 +67,7 @@ impl ConnectionMetadata { // need to use builder to set this field local_addr: None, poison_fn: Arc::new(poison), + connection_id: None, } } @@ -62,6 +85,13 @@ impl ConnectionMetadata { pub fn local_addr(&self) -> Option { self.local_addr } + + /// Get the connection id, if one was assigned by the HTTP client. + /// + /// Present for the v2 HTTP client; `None` for the v1 client. + pub fn connection_id(&self) -> Option { + self.connection_id + } } impl fmt::Debug for ConnectionMetadata { @@ -70,6 +100,7 @@ impl fmt::Debug for ConnectionMetadata { .field("is_proxied", &self.is_proxied) .field("remote_addr", &self.remote_addr) .field("local_addr", &self.local_addr) + .field("connection_id", &self.connection_id) .finish() } } @@ -81,6 +112,7 @@ pub struct ConnectionMetadataBuilder { remote_addr: Option, local_addr: Option, poison_fn: Option>, + connection_id: Option, } impl fmt::Debug for ConnectionMetadataBuilder { @@ -155,6 +187,18 @@ impl ConnectionMetadataBuilder { self } + /// Set the connection id. + pub fn connection_id(mut self, id: ConnectionId) -> Self { + self.connection_id = Some(id); + self + } + + /// Set the connection id. + pub fn set_connection_id(&mut self, id: Option) -> &mut Self { + self.connection_id = id; + self + } + /// Build a [`ConnectionMetadata`] value. /// /// # Panics @@ -170,6 +214,7 @@ impl ConnectionMetadataBuilder { poison_fn: self .poison_fn .expect("poison_fn should be set for ConnectionMetadata"), + connection_id: self.connection_id, } } } From 543c56cb50aa4c2ee4b031acc9b07bd4f216fc05 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Sat, 16 May 2026 09:45:54 -0400 Subject: [PATCH 19/47] cleanup connection id and authority --- .../aws-smithy-http-client/src/client/pool.rs | 195 ++++++----- .../src/client/pool/connection.rs | 306 ++++++++++++++++-- .../src/client/pool/handshake.rs | 138 ++++++-- .../aws-smithy-http-client/src/client/v2.rs | 47 ++- .../tests/pool_behavior_test.rs | 197 +++++++++++ 5 files changed, 774 insertions(+), 109 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 574c970441c..f3d5ad65663 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -12,7 +12,7 @@ //! metadata (remote/local address, poison handle) to the adapter layer //! for the runtime's connection-poisoning interceptor. -mod connection; +pub(crate) mod connection; mod handshake; mod vendored_cache; @@ -45,8 +45,8 @@ use tokio::sync::{oneshot, Semaphore}; use tower::{Service, ServiceExt}; use connection::{ - CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, H2ConnectionRef, - SingletonConnection, + Authority, CachedConnection, CheckoutResponse, CloseReason, ConnectionClosedEvent, + ConnectionGuard, GuardedBody, H2ConnectionRef, SingletonConnection, }; pub(crate) use connection::{ConnectCtx, ReadTimeoutHint, TimeoutContext}; use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake}; @@ -104,7 +104,7 @@ impl ConnectionMetadataCapture { /// they take effect, not here). Fields are `pub(crate)` because /// consumers of this struct are all within the pool module; accessor /// methods would add noise without benefit. -#[derive(Clone, Debug, Default)] +#[derive(Clone, Default)] pub(crate) struct PoolConfig { /// Upper bound on concurrent connections (total, across all hosts). /// Enforced via semaphore at the connection establishment layer. @@ -119,6 +119,9 @@ pub(crate) struct PoolConfig { /// How long an idle connection may stay in the pool before being /// evicted. `None` = no eviction (hyper-util's default behavior). pub(crate) pool_idle_timeout: Option, + + /// Optional listener for connection lifecycle events. + pub(crate) connection_event_listener: Option>, } /// Key for per-host connection pool routing. @@ -177,7 +180,7 @@ where { let global_sem = config.max_connections.map(|n| Arc::new(Semaphore::new(n))); let max_per_host = config.max_connections_per_host; - let pool_hooks = handshake::PoolHooks::new(); + let pool_hooks = handshake::PoolHooks::new(config.connection_event_listener.clone()); tracing::debug!( max_connections = ?config.max_connections, @@ -188,7 +191,12 @@ where let make_entry = { let pool_hooks = pool_hooks.clone(); - move |_uri: &http_1x::Uri| -> Box { + move |uri: &http_1x::Uri| -> Box { + let authority = Authority::new( + uri.authority() + .expect("pool entry URI has authority") + .as_str(), + ); let per_host_sem = max_per_host.map(|n| Arc::new(Semaphore::new(n))); let limited = handshake::ConnectionLimit::new( connector.clone(), @@ -230,32 +238,51 @@ where .lock() .expect("retainer slot poisoned") .push(BoxedRetainer { - retain_fn: Box::new(move |timeout| { - let now = Instant::now(); - cache_for_retain - .lock() - .expect("retain cache lock poisoned") - .retain(|managed| { - let keep = !managed.is_poisoned() - && now.saturating_duration_since(managed.idle_at()) - < timeout; - if !keep { - let reason = if managed.is_poisoned() { - "poisoned" - } else { - "idle_expired" - }; - let idle_duration = now - .saturating_duration_since(managed.idle_at()); - tracing::debug!( - conn_id = managed.conn_id(), - reason, - ?idle_duration, - "pool: connection evicted" - ); - } - keep - }); + retain_fn: Box::new({ + let listener = pool_hooks.listener.clone(); + move |timeout| { + let now = Instant::now(); + cache_for_retain + .lock() + .expect("retain cache lock poisoned") + .retain(|managed| { + let keep = !managed.is_poisoned() + && now.saturating_duration_since( + managed.idle_at(), + ) < timeout; + if !keep { + let reason = if managed.is_poisoned() { + "poisoned" + } else { + "idle_expired" + }; + let idle_duration = now + .saturating_duration_since( + managed.idle_at(), + ); + tracing::debug!( + conn_id = %managed.conn_id(), + reason, + ?idle_duration, + "pool: connection evicted" + ); + if let Some(ref l) = listener { + l.on_closed(&ConnectionClosedEvent::new( + managed.conn_id(), + managed.info.authority.clone(), + managed.info.remote_addr, + if managed.is_poisoned() { + CloseReason::Poisoned + } else { + CloseReason::IdleTimeout + }, + None, + )); + } + } + keep + }); + } }), is_empty_fn: Box::new(move || { cache_for_empty @@ -264,7 +291,12 @@ where .is_empty() }), }); - cache.map_response(|cached| H1Checkout::new(CachedConnection::new(cached))) + cache.map_response({ + let listener = pool_hooks.listener.clone(); + move |cached| { + H1Checkout::new(CachedConnection::new(cached, listener.clone())) + } + }) }) }) .upgrade({ @@ -277,6 +309,7 @@ where inspected, h2_ref.clone(), pool_hooks.clone(), + authority.clone(), )); let singleton_for_retain = Arc::new(std::sync::Mutex::new(singleton.clone())); @@ -285,44 +318,62 @@ where .lock() .expect("retainer slot poisoned") .push(BoxedRetainer { - retain_fn: Box::new(move |timeout| { - let now = Instant::now(); - singleton_for_retain - .lock() - .expect("retain singleton lock poisoned") - .retain(|managed| { - // For H2 the connection stays in - // Singleton while serving streams, so - // `idle_at` alone is not a valid - // idleness signal. Keep the - // connection if any stream is in - // flight; only evict when truly idle - // AND `idle_at` has exceeded the - // timeout (stamped on the 1 → 0 - // transition by - // `SingletonConnection::drop`). - let keep = !managed.is_poisoned() - && (managed.active_streams_count() > 0 - || now.saturating_duration_since( - managed.idle_at(), - ) < timeout); - if !keep { - let reason = if managed.is_poisoned() { - "poisoned" - } else { - "idle_expired" - }; - let idle_duration = now - .saturating_duration_since(managed.idle_at()); - tracing::debug!( - conn_id = managed.conn_id(), - reason, - ?idle_duration, - "pool: connection evicted" - ); - } - keep - }); + retain_fn: Box::new({ + let listener = pool_hooks.listener.clone(); + move |timeout| { + let now = Instant::now(); + singleton_for_retain + .lock() + .expect("retain singleton lock poisoned") + .retain(|managed| { + // For H2 the connection stays in + // Singleton while serving streams, so + // `idle_at` alone is not a valid + // idleness signal. Keep the + // connection if any stream is in + // flight; only evict when truly idle + // AND `idle_at` has exceeded the + // timeout (stamped on the 1 → 0 + // transition by + // `SingletonConnection::drop`). + let keep = !managed.is_poisoned() + && (managed.active_streams_count() > 0 + || now.saturating_duration_since( + managed.idle_at(), + ) < timeout); + if !keep { + let reason = if managed.is_poisoned() { + "poisoned" + } else { + "idle_expired" + }; + let idle_duration = now + .saturating_duration_since( + managed.idle_at(), + ); + tracing::debug!( + conn_id = %managed.conn_id(), + reason, + ?idle_duration, + "pool: connection evicted" + ); + if let Some(ref l) = listener { + l.on_closed(&ConnectionClosedEvent::new( + managed.conn_id(), + managed.info.authority.clone(), + managed.info.remote_addr, + if managed.is_poisoned() { + CloseReason::Poisoned + } else { + CloseReason::IdleTimeout + }, + None, + )); + } + } + keep + }); + } }), is_empty_fn: Box::new(move || { singleton_for_empty @@ -634,7 +685,7 @@ impl Service> for H1Checkout) -> Self::Future { let mut conn = self.conn.take().expect("H1Checkout::call called twice"); - tracing::trace!(conn_id = conn.conn_id(), uri = %req.uri(), "pool: dispatching request"); + tracing::trace!(conn_id = %conn.conn_id(), uri = %req.uri(), "pool: dispatching request"); // Populate the adapter-provided capture so the // `CaptureSmithyConnection` retriever can return a live // `ConnectionMetadata` pointing at THIS connection's poison pill. diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index d925824c69e..0fcc67022ea 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -22,6 +22,258 @@ use super::cache; use super::handshake::H1SendRequest; use super::BoxError; +/// Opaque authority (host:port) associated with a pooled connection. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Authority(Arc); + +impl Authority { + pub(crate) fn new(s: impl Into>) -> Self { + Self(s.into()) + } + + /// The authority as a string slice. + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl AsRef for Authority { + fn as_ref(&self) -> &str { + &self.0 + } +} + +impl std::fmt::Display for Authority { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +/// Protocol negotiated for a connection. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum NegotiatedProtocol { + /// HTTP/1.1 + Http1, + /// HTTP/2 + Http2, +} + +/// Why a connection was removed from the pool. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum CloseReason { + /// Idle longer than the configured pool idle timeout. + IdleTimeout, + /// Marked as poisoned (unhealthy) by the SDK or pool. + Poisoned, + /// The pool itself was dropped. + PoolDropped, +} + +/// Emitted when a new connection is established (TCP + TLS + HTTP handshake). +#[derive(Debug)] +pub struct ConnectionCreatedEvent { + conn_id: ConnectionId, + authority: Authority, + remote_addr: Option, + protocol: NegotiatedProtocol, +} + +impl ConnectionCreatedEvent { + pub(crate) fn new( + conn_id: ConnectionId, + authority: Authority, + remote_addr: Option, + protocol: NegotiatedProtocol, + ) -> Self { + Self { + conn_id, + authority, + remote_addr, + protocol, + } + } + + /// The pool-assigned connection identifier. + pub fn conn_id(&self) -> ConnectionId { + self.conn_id + } + + /// The authority (host:port) this connection is for. + pub fn authority(&self) -> &Authority { + &self.authority + } + + /// Remote address of the peer, if known. + pub fn remote_addr(&self) -> Option { + self.remote_addr + } + + /// Negotiated protocol. + pub fn protocol(&self) -> NegotiatedProtocol { + self.protocol + } +} + +/// Emitted when an existing idle connection is checked out from the pool. +#[derive(Debug)] +pub struct ConnectionReusedEvent { + conn_id: ConnectionId, + authority: Authority, +} + +impl ConnectionReusedEvent { + pub(crate) fn new(conn_id: ConnectionId, authority: Authority) -> Self { + Self { conn_id, authority } + } + + /// The pool-assigned connection identifier. + pub fn conn_id(&self) -> ConnectionId { + self.conn_id + } + + /// The authority (host:port) this connection is for. + pub fn authority(&self) -> &Authority { + &self.authority + } +} + +/// Emitted when a connection is removed from the pool. +pub struct ConnectionClosedEvent { + conn_id: ConnectionId, + authority: Authority, + remote_addr: Option, + reason: CloseReason, + error: Option, +} + +impl std::fmt::Debug for ConnectionClosedEvent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut s = f.debug_struct("ConnectionClosedEvent"); + s.field("conn_id", &self.conn_id) + .field("authority", &self.authority) + .field("remote_addr", &self.remote_addr) + .field("reason", &self.reason); + if let Some(ref e) = self.error { + s.field("error", &format_args!("{e}")); + } + s.finish() + } +} + +impl ConnectionClosedEvent { + pub(crate) fn new( + conn_id: ConnectionId, + authority: Authority, + remote_addr: Option, + reason: CloseReason, + error: Option, + ) -> Self { + Self { + conn_id, + authority, + remote_addr, + reason, + error, + } + } + + /// The pool-assigned connection identifier. + pub fn conn_id(&self) -> ConnectionId { + self.conn_id + } + + /// The authority (host:port) this connection was for. + pub fn authority(&self) -> &Authority { + &self.authority + } + + /// Remote address of the peer, if known. + pub fn remote_addr(&self) -> Option { + self.remote_addr + } + + /// Why the connection was closed. + pub fn reason(&self) -> CloseReason { + self.reason + } + + /// The error associated with this close, if any. Present for + /// server-initiated closes; `None` for policy-driven closes + /// (idle timeout, poisoning). + pub fn error(&self) -> Option<&(dyn std::error::Error + Send + Sync)> { + self.error.as_ref().map(|e| e.as_ref()) + } +} + +/// Emitted when a connection attempt fails before completing the handshake. +pub struct ConnectionFailedEvent { + authority: Authority, + remote_addr: Option, + error: BoxError, +} + +impl std::fmt::Debug for ConnectionFailedEvent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ConnectionFailedEvent") + .field("authority", &self.authority) + .field("remote_addr", &self.remote_addr) + .field("error", &format_args!("{}", self.error)) + .finish() + } +} + +impl ConnectionFailedEvent { + pub(crate) fn new( + authority: Authority, + remote_addr: Option, + error: BoxError, + ) -> Self { + Self { + authority, + remote_addr, + error, + } + } + + /// The authority (host:port) the connection was attempting to reach. + pub fn authority(&self) -> &Authority { + &self.authority + } + + /// Remote address of the peer, if known. `None` when the failure occurred + /// before a peer address was established (e.g., DNS resolution failure or + /// connection refused before address binding). + pub fn remote_addr(&self) -> Option { + self.remote_addr + } + + /// The error that caused the connection attempt to fail. + pub fn error(&self) -> &(dyn std::error::Error + Send + Sync) { + self.error.as_ref() + } +} + +/// Callback for connection lifecycle events within the pool. +/// +/// Implementations receive notifications when connections are created, +/// reused from the pool, closed, or fail to establish. Useful for metrics, +/// DNS failure feedback, or adaptive concurrency control. +/// +/// Implementations must be cheap (nanoseconds). Expensive work should be +/// deferred to a background task. +pub trait ConnectionEventListener: Send + Sync + 'static { + /// A new connection was established. + fn on_created(&self, _event: &ConnectionCreatedEvent) {} + /// An existing idle connection was checked out from the pool. + fn on_reused(&self, _event: &ConnectionReusedEvent) {} + /// A connection was removed from the pool. + fn on_closed(&self, _event: &ConnectionClosedEvent) {} + /// A connection attempt failed before completing the handshake. + fn on_connection_failed(&self, _event: &ConnectionFailedEvent) {} +} + /// A duration paired with the sleep implementation used to realize it. /// /// This type makes "timeout without sleep impl" unrepresentable: you cannot @@ -126,6 +378,9 @@ pub(crate) struct ConnectionInfo { /// [`hyper_util::client::legacy::connect::Connected::is_proxied`] at /// handshake. pub(crate) is_proxied: bool, + /// The authority (host:port) this connection is for. Populated from the + /// URI at handshake time. + pub(crate) authority: Authority, // Future: tls_info, timing (tcp_connect_duration, tls_handshake_duration) } @@ -169,12 +424,12 @@ impl PoisonPill { /// poisoning one clone poisons all of them. pub(crate) struct ManagedConnection { inner: S, - info: ConnectionInfo, + pub(crate) info: ConnectionInfo, /// Stable identifier for this physical connection, unique within the /// owning pool. Shared across `Clone`s (an H2 connection's multiplexed /// request handles all carry the same `conn_id`). Used in tracing and /// surfaced through `ConnectionMetadata` for cross-layer correlation. - conn_id: u64, + conn_id: ConnectionId, created_at: Instant, /// Timestamp the connection last became idle (or its creation time, if /// it has never been returned to the pool). @@ -212,7 +467,7 @@ impl ManagedConnection { pub(crate) fn new( inner: S, info: ConnectionInfo, - conn_id: u64, + conn_id: ConnectionId, permit: Arc, ) -> Self { let now = Instant::now(); @@ -230,7 +485,7 @@ impl ManagedConnection { /// Stable identifier for this physical connection. Shared across /// clones (H2 multiplexing). - pub(crate) fn conn_id(&self) -> u64 { + pub(crate) fn conn_id(&self) -> ConnectionId { self.conn_id } @@ -304,9 +559,9 @@ impl ManagedConnection { let remote = self.info.remote_addr; let mut builder = ConnectionMetadata::builder() .proxied(self.info.is_proxied) - .connection_id(ConnectionId::new(self.conn_id)) + .connection_id(self.conn_id) .poison_fn(move || { - tracing::debug!(conn_id, ?remote, "pool: connection poisoned"); + tracing::debug!(conn_id = %conn_id, ?remote, "pool: connection poisoned"); poison.poison(); }); builder @@ -371,16 +626,26 @@ where /// normally through `Cached::Drop`. pub(crate) struct CachedConnection { inner: Option>>, + listener: Option>, } impl CachedConnection { - pub(crate) fn new(cached: cache::Cached>) -> Self { - tracing::trace!( - conn_id = cached.inner().conn_id(), - "pool: connection reused" - ); + pub(crate) fn new( + cached: cache::Cached>, + listener: Option>, + ) -> Self { + let is_reuse = *cached.inner().idle_at.lock().unwrap() > cached.inner().created_at; + if is_reuse { + let conn_id = cached.inner().conn_id(); + let authority = cached.inner().info.authority.clone(); + tracing::trace!(conn_id = %conn_id, "pool: connection reused"); + if let Some(ref l) = listener { + l.on_reused(&ConnectionReusedEvent::new(conn_id, authority)); + } + } Self { inner: Some(cached), + listener, } } @@ -407,7 +672,7 @@ impl CachedConnection { } /// Forwards [`ManagedConnection::conn_id`]. - pub(crate) fn conn_id(&self) -> u64 { + pub(crate) fn conn_id(&self) -> ConnectionId { self.inner .as_ref() .expect("CachedConnection conn_id after drop") @@ -439,15 +704,20 @@ impl Drop for CachedConnection { let managed = cached.inner(); let conn_id = managed.conn_id; if managed.is_poisoned() { - tracing::debug!(conn_id, "pool: connection discarded (poisoned)"); + tracing::debug!(conn_id = %conn_id, "pool: connection discarded (poisoned)"); + if let Some(ref listener) = self.listener { + listener.on_closed(&ConnectionClosedEvent::new( + managed.conn_id(), + managed.info.authority.clone(), + managed.info.remote_addr, + CloseReason::Poisoned, + None, + )); + } cached.discard(); } else { - // Stamp the return-to-idle moment so the eviction task - // can measure elapsed idleness. Healthy connection: the - // Cached drops normally after this, returning it to the - // cache's idle set. managed.mark_idle(); - tracing::trace!(conn_id, "pool: connection returned to idle"); + tracing::trace!(conn_id = %conn_id, "pool: connection returned to idle"); } } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index ea8febf3721..ce191a41939 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -15,6 +15,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use aws_smithy_runtime_api::client::connection::ConnectionId; use aws_smithy_types::body::SdkBody; use hyper::rt::Executor; use hyper_util::client::legacy::connect::{Connection, HttpInfo}; @@ -22,7 +23,10 @@ use hyper_util::rt::TokioExecutor; use tokio::sync::Semaphore; use tower::Service; -use super::connection::{ConnectCtx, ConnectionInfo, ConnectionPermit, ManagedConnection}; +use super::connection::{ + Authority, ConnectCtx, ConnectionCreatedEvent, ConnectionFailedEvent, ConnectionInfo, + ConnectionPermit, ManagedConnection, NegotiatedProtocol, +}; use super::BoxError; /// Pool-scoped instrumentation primitives shared across layers in the @@ -34,20 +38,50 @@ use super::BoxError; #[derive(Clone)] pub(crate) struct PoolHooks { conn_id_counter: Arc, + pub(crate) listener: Option>, } impl PoolHooks { - pub(crate) fn new() -> Self { + pub(crate) fn new( + listener: Option>, + ) -> Self { Self { conn_id_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)), + listener, } } /// Mint the next connection id. Stable for the connection's lifetime; /// the underlying counter wraps at `u64::MAX`. - pub(crate) fn next_conn_id(&self) -> u64 { - self.conn_id_counter - .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + pub(crate) fn next_conn_id(&self) -> ConnectionId { + ConnectionId::new( + self.conn_id_counter + .fetch_add(1, std::sync::atomic::Ordering::Relaxed), + ) + } + + pub(crate) fn on_created(&self, event: &ConnectionCreatedEvent) { + if let Some(ref l) = self.listener { + l.on_created(event); + } + } + + pub(crate) fn on_connection_failed(&self, event: &ConnectionFailedEvent) { + if let Some(ref l) = self.listener { + l.on_connection_failed(event); + } + } + + pub(crate) fn on_reused(&self, event: &super::connection::ConnectionReusedEvent) { + if let Some(ref l) = self.listener { + l.on_reused(event); + } + } + + pub(crate) fn on_closed(&self, event: &super::connection::ConnectionClosedEvent) { + if let Some(ref l) = self.listener { + l.on_closed(event); + } } } @@ -209,7 +243,7 @@ impl Service> for H2SendRequest { // --------------------------------------------------------------------------- /// Extract `ConnectionInfo` from a just-connected IO stream. -fn capture_info(io: &IO) -> ConnectionInfo { +fn capture_info(io: &IO, authority: Authority) -> ConnectionInfo { let connected = io.connected(); let is_proxied = connected.is_proxied(); let mut extras = http_1x::Extensions::new(); @@ -219,6 +253,7 @@ fn capture_info(io: &IO) -> ConnectionInfo { remote_addr: http_info.map(|i| i.remote_addr()), local_addr: http_info.map(|i| i.local_addr()), is_proxied, + authority, } } @@ -270,32 +305,66 @@ where fn call(&mut self, ctx: ConnectCtx) -> Self::Future { let hooks = self.hooks.clone(); + let authority = Authority::new( + ctx.uri + .authority() + .expect("request URI has authority") + .as_str(), + ); let fut = self.connector.call(ctx); Box::pin(async move { - let (io, permit) = fut.await.map_err(Into::into)?; - let info = capture_info(&io); + let (io, permit) = match fut.await.map_err(Into::into) { + Ok(v) => v, + Err(e) => { + hooks.on_connection_failed(&ConnectionFailedEvent::new( + authority, + None, + Into::into(e.to_string()), + )); + return Err(e); + } + }; + let info = capture_info(&io, authority.clone()); let conn_id = hooks.next_conn_id(); - let (tx, conn) = hyper::client::conn::http1::Builder::new() + let (tx, conn) = match hyper::client::conn::http1::Builder::new() .handshake(io) .await - .map_err(|e| Box::new(e) as BoxError)?; + { + Ok(v) => v, + Err(e) => { + let boxed: BoxError = Box::new(e); + hooks.on_connection_failed(&ConnectionFailedEvent::new( + authority, + info.remote_addr, + Into::into(boxed.to_string()), + )); + return Err(boxed); + } + }; tracing::debug!( - conn_id, + conn_id = %conn_id, protocol = "h1", remote = ?info.remote_addr, local = ?info.local_addr, "pool: connection established" ); + hooks.on_created(&ConnectionCreatedEvent::new( + conn_id, + authority, + info.remote_addr, + NegotiatedProtocol::Http1, + )); + spawn_driver({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { if let Err(e) = conn.with_upgrades().await { tracing::debug!( - conn_id, + conn_id = %conn_id, protocol = "h1", ?remote_addr, ?local_addr, @@ -325,6 +394,7 @@ pub(crate) struct H2ConnectAndHandshake { connector: C, h2_ref: super::connection::H2ConnectionRef, hooks: PoolHooks, + authority: Authority, } impl H2ConnectAndHandshake { @@ -338,11 +408,13 @@ impl H2ConnectAndHandshake { connector: C, h2_ref: super::connection::H2ConnectionRef, hooks: PoolHooks, + authority: Authority, ) -> Self { Self { connector, h2_ref, hooks, + authority, } } } @@ -353,6 +425,7 @@ impl Clone for H2ConnectAndHandshake { connector: self.connector.clone(), h2_ref: self.h2_ref.clone(), hooks: self.hooks.clone(), + authority: self.authority.clone(), } } } @@ -376,31 +449,60 @@ where let fut = self.connector.call(()); let h2_ref = self.h2_ref.clone(); let hooks = self.hooks.clone(); + let authority = self.authority.clone(); Box::pin(async move { - let (io, permit) = fut.await.map_err(Into::into)?; - let info = capture_info(&io); + let (io, permit) = match fut.await.map_err(Into::into) { + Ok(v) => v, + Err(e) => { + hooks.on_connection_failed(&ConnectionFailedEvent::new( + authority, + None, + Into::into(e.to_string()), + )); + return Err(e); + } + }; + let info = capture_info(&io, authority.clone()); let conn_id = hooks.next_conn_id(); - let (tx, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new()) + let (tx, conn) = match hyper::client::conn::http2::Builder::new(TokioExecutor::new()) .handshake(io) .await - .map_err(|e| Box::new(e) as BoxError)?; + { + Ok(v) => v, + Err(e) => { + let boxed: BoxError = Box::new(e); + hooks.on_connection_failed(&ConnectionFailedEvent::new( + authority, + info.remote_addr, + Into::into(boxed.to_string()), + )); + return Err(boxed); + } + }; tracing::debug!( - conn_id, + conn_id = %conn_id, protocol = "h2", remote = ?info.remote_addr, local = ?info.local_addr, "pool: connection established" ); + hooks.on_created(&ConnectionCreatedEvent::new( + conn_id, + authority, + info.remote_addr, + NegotiatedProtocol::Http2, + )); + spawn_driver({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { if let Err(e) = conn.await { tracing::debug!( - conn_id, + conn_id = %conn_id, protocol = "h2", ?remote_addr, ?local_addr, diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index 3c441d200ab..d0b0e4dc02f 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -7,6 +7,11 @@ //! //! Entry point: [`BuilderV2::new`] or [`Builder::new_v2`](super::Builder::new_v2). +pub use super::pool::connection::{ + Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, + ConnectionFailedEvent, ConnectionReusedEvent, NegotiatedProtocol, +}; + use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; @@ -44,16 +49,33 @@ use crate::tls::TlsContext; /// .pool_idle_timeout(Duration::from_secs(20)) /// .build_http(); /// ``` -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct BuilderV2 { pool_idle_timeout: Option, tcp_nodelay: bool, max_connections: Option, max_connections_per_host: Option, proxy_config: Option, + connection_event_listener: Option>, tls: Tls, } +impl std::fmt::Debug for BuilderV2 { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BuilderV2") + .field("pool_idle_timeout", &self.pool_idle_timeout) + .field("tcp_nodelay", &self.tcp_nodelay) + .field("max_connections", &self.max_connections) + .field("max_connections_per_host", &self.max_connections_per_host) + .field("proxy_config", &self.proxy_config) + .field( + "connection_event_listener", + &self.connection_event_listener.as_ref().map(|_| ".."), + ) + .finish() + } +} + impl Default for BuilderV2 { fn default() -> Self { Self { @@ -62,6 +84,7 @@ impl Default for BuilderV2 { max_connections: None, max_connections_per_host: None, proxy_config: None, + connection_event_listener: None, tls: TlsUnset {}, } } @@ -126,6 +149,24 @@ impl BuilderV2 { self.proxy_config = config; self } + + /// Set a listener for connection lifecycle events (created, reused, closed, failed). + pub fn connection_event_listener( + mut self, + listener: Arc, + ) -> Self { + self.connection_event_listener = Some(listener); + self + } + + /// Mutable-reference variant of [`connection_event_listener`](Self::connection_event_listener). + pub fn set_connection_event_listener( + &mut self, + listener: Option>, + ) -> &mut Self { + self.connection_event_listener = listener; + self + } } impl BuilderV2 { @@ -142,6 +183,7 @@ impl BuilderV2 { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, proxy_config: self.proxy_config, + connection_event_listener: self.connection_event_listener, tls: TlsProviderSelected { provider, context: TlsContext::default(), @@ -158,6 +200,7 @@ impl BuilderV2 { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, pool_idle_timeout: self.pool_idle_timeout, + connection_event_listener: self.connection_event_listener.clone(), }; let proxy_matcher = proxy_matcher_from(&self.proxy_config); let pool = build_http_pool_with_proxy(tcp, &self.proxy_config, config); @@ -182,6 +225,7 @@ impl BuilderV2 { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, pool_idle_timeout: self.pool_idle_timeout, + connection_event_listener: self.connection_event_listener.clone(), }; let proxy_matcher = proxy_matcher_from(&self.proxy_config); let pool = build_http_pool_with_proxy(tcp, &self.proxy_config, config); @@ -268,6 +312,7 @@ impl BuilderV2 { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, pool_idle_timeout: self.pool_idle_timeout, + connection_event_listener: self.connection_event_listener.clone(), }; let proxy_config = self diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 18adf16579e..5a0ae279aca 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1338,3 +1338,200 @@ async fn v2_connection_id_surfaced_through_metadata() { ); assert_eq!(id3, ConnectionId::new(1)); } + +// ---- ConnectionEventListener tests ---- + +mod listener_tests { + use super::*; + use aws_smithy_http_client::v2::{ + CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, + ConnectionFailedEvent, ConnectionReusedEvent, + }; + use std::sync::{Arc, Mutex}; + + /// Records connection lifecycle events for test assertions. + #[derive(Debug, Clone, Default)] + struct RecordingListener { + /// (conn_id, authority) for each connection created + created: Arc>>, + /// (conn_id, authority) for each connection reused from idle + reused: Arc>>, + /// (conn_id, authority, reason) for each connection closed + closed: Arc>>, + /// authority for each failed connection attempt + failed: Arc>>, + } + + impl ConnectionEventListener for RecordingListener { + fn on_created(&self, event: &ConnectionCreatedEvent) { + self.created.lock().unwrap().push(( + event.conn_id().to_string().parse().unwrap(), + event.authority().as_str().to_string(), + )); + } + fn on_reused(&self, event: &ConnectionReusedEvent) { + self.reused.lock().unwrap().push(( + event.conn_id().to_string().parse().unwrap(), + event.authority().as_str().to_string(), + )); + } + fn on_closed(&self, event: &ConnectionClosedEvent) { + self.closed.lock().unwrap().push(( + event.conn_id().to_string().parse().unwrap(), + event.authority().as_str().to_string(), + event.reason(), + )); + } + fn on_connection_failed(&self, event: &ConnectionFailedEvent) { + self.failed + .lock() + .unwrap() + .push(event.authority().as_str().to_string()); + } + } + + fn make_v2_with_listener( + harness: &ConnectionTestHarness, + idle_timeout: Option, + listener: Arc, + ) -> SharedHttpClient { + let mut builder = + aws_smithy_http_client::Builder::new_v2().connection_event_listener(listener); + if let Some(timeout) = idle_timeout { + builder = builder.pool_idle_timeout(timeout); + } + builder.build_http_with_resolver(harness.dns_resolver()) + } + + /// Listener receives created on first request, reused on second, and + /// closed(IdleTimeout) after eviction. + #[tokio::test] + async fn listener_lifecycle_created_reused_closed() { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"a", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"b", + }, + ], + ) + .build() + .await; + + let listener = Arc::new(RecordingListener::default()); + let client = + make_v2_with_listener(&harness, Some(Duration::from_millis(80)), listener.clone()); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + // First request: on_created + super::send_and_read_body(&client, &url).await; + assert_eq!(listener.created.lock().unwrap().len(), 1); + assert_eq!(listener.created.lock().unwrap()[0].0, 0); + assert!(listener.created.lock().unwrap()[0].1.contains("127.0.0.1")); + + // Second request: on_reused + super::send_and_read_body(&client, &url).await; + assert_eq!(listener.reused.lock().unwrap().len(), 1); + assert_eq!(listener.reused.lock().unwrap()[0].0, 0); + + // Wait for eviction + tokio::time::sleep(Duration::from_millis(200)).await; + + let closed = listener.closed.lock().unwrap(); + assert_eq!(closed.len(), 1); + assert_eq!(closed[0].0, 0); + assert_eq!(closed[0].2, CloseReason::IdleTimeout); + } + + /// Poisoning a connection fires on_closed with Poisoned reason. + #[tokio::test] + async fn listener_poisoned_connection_fires_on_closed() { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"ok", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"ok", + }, + ], + ) + .build() + .await; + + let listener = Arc::new(RecordingListener::default()); + let client = make_v2_with_listener(&harness, None, listener.clone()); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + // Make a request and poison the connection + let (_, _, meta) = super::send_with_capture(&client, &url).await; + meta.unwrap().poison(); + + // The poisoned connection is discarded on next checkout attempt + super::send_and_read_body(&client, &url).await; + + let closed = listener.closed.lock().unwrap(); + assert_eq!(closed.len(), 1); + assert_eq!(closed[0].2, CloseReason::Poisoned); + } + + /// Connection failure fires on_connection_failed with the authority. + #[tokio::test] + async fn listener_connection_failed() { + let listener = Arc::new(RecordingListener::default()); + + // Use a harness with no endpoints so the connection will fail + let harness = ConnectionTestHarness::builder() + .endpoint(IP1, vec![ConnectionBehavior::ResetOnConnect]) + .build() + .await; + + let client = make_v2_with_listener(&harness, None, listener.clone()); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + let settings = HttpConnectorSettings::builder().build(); + let components = runtime_components(); + let connector = client.http_connector(&settings, &components); + let request = HttpRequest::get(&url).expect("valid request"); + let _ = connector.call(request).await; + + let failed = listener.failed.lock().unwrap(); + assert_eq!(failed.len(), 1); + assert!(failed[0].contains("127.0.0.1")); + } + + /// Authority is correctly populated on events. + #[tokio::test] + async fn listener_authority_populated() { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"x", + }], + ) + .build() + .await; + + let listener = Arc::new(RecordingListener::default()); + let client = make_v2_with_listener(&harness, None, listener.clone()); + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{}/path", port); + + super::send_and_read_body(&client, &url).await; + + let created = listener.created.lock().unwrap(); + assert_eq!(created[0].1, format!("127.0.0.1:{}", port)); + } +} From a7070a930ad6b977531b3995f021fff43ca25a2a Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Sat, 16 May 2026 11:01:28 -0400 Subject: [PATCH 20/47] wire up unhealthy event --- .../src/client/pool/connection.rs | 26 +++++++++- .../tests/pool_behavior_test.rs | 47 +++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 0fcc67022ea..a21a484f24c 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -67,6 +67,10 @@ pub enum CloseReason { IdleTimeout, /// Marked as poisoned (unhealthy) by the SDK or pool. Poisoned, + /// Connection found dead at checkout (e.g., server closed the + /// connection while it was idle in the pool). The associated error + /// on the event carries specifics. + Unusable, /// The pool itself was dropped. PoolDropped, } @@ -690,7 +694,27 @@ where type Future = > as Service>::Future; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { - self.inner.as_mut().unwrap().poll_ready(cx) + match self.inner.as_mut().unwrap().poll_ready(cx) { + Poll::Ready(Err(e)) => { + // Connection is dead. If poisoned, Drop handles the event. + // If not poisoned, this is an unusable connection (server + // closed it while idle, driver died, etc.). + let managed = self.inner.as_ref().unwrap().inner(); + if !managed.is_poisoned() { + if let Some(ref l) = self.listener { + l.on_closed(&ConnectionClosedEvent::new( + managed.conn_id, + managed.info.authority.clone(), + managed.info.remote_addr, + CloseReason::Unusable, + None, + )); + } + } + Poll::Ready(Err(e)) + } + other => other, + } } fn call(&mut self, req: Req) -> Self::Future { diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 5a0ae279aca..f9671db4869 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1534,4 +1534,51 @@ mod listener_tests { let created = listener.created.lock().unwrap(); assert_eq!(created[0].1, format!("127.0.0.1:{}", port)); } + + /// A connection that the server closed while idle fires on_closed + /// with Unusable reason when detected at checkout. + #[tokio::test] + async fn listener_unusable_connection_fires_on_closed() { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + // First request succeeds, then server closes the connection. + ConnectionBehavior::RespondThenClose { + status: 200, + body: b"ok", + }, + // Second connection for the retry. + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"ok", + }, + ], + ) + .build() + .await; + + let listener = Arc::new(RecordingListener::default()); + let client = make_v2_with_listener(&harness, None, listener.clone()); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + // First request succeeds; connection returns to pool. + super::send_and_read_body(&client, &url).await; + + // Brief pause for the server-side close to propagate. + tokio::time::sleep(Duration::from_millis(50)).await; + + // Second request: pool checks out the dead connection, detects it + // in poll_ready, fires on_closed(Unusable), then creates a new one. + super::send_and_read_body(&client, &url).await; + + let closed = listener.closed.lock().unwrap(); + let unusable = closed.iter().find(|c| c.2 == CloseReason::Unusable); + assert!( + unusable.is_some(), + "expected on_closed with Unusable reason, got: {:?}", + *closed + ); + assert_eq!(unusable.unwrap().0, 0, "should be the first connection"); + } } From e0f6c8024c1b9b825da433e7620d7f522f2dd375 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Mon, 18 May 2026 11:54:08 -0400 Subject: [PATCH 21/47] add conn timing --- aws/rust-runtime/Cargo.lock | 15 -- .../aws-smithy-http-client/Cargo.toml | 20 +-- .../aws-smithy-http-client/additional-ci | 3 + .../src/client/pool/connection.rs | 145 ++++++++++++++++++ .../src/client/pool/handshake.rs | 9 +- .../src/client/tls/rustls_provider.rs | 31 +++- .../src/client/tls/s2n_tls_provider.rs | 15 +- .../aws-smithy-http-client/src/client/v2.rs | 2 +- .../tests/pool_behavior_test.rs | 117 ++++++++------ .../tests/smoke_test_clients.rs | 66 ++++++++ 10 files changed, 338 insertions(+), 85 deletions(-) diff --git a/aws/rust-runtime/Cargo.lock b/aws/rust-runtime/Cargo.lock index 20fa5de8665..a6cd6bcac17 100644 --- a/aws/rust-runtime/Cargo.lock +++ b/aws/rust-runtime/Cargo.lock @@ -652,21 +652,6 @@ dependencies = [ "libc", ] -[[package]] -name = "crc" -version = "3.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9710d3b3739c2e349eb44fe848ad0b7c8cb1e42bd87ee49371df2f7acaf3e675" -dependencies = [ - "crc-catalog", -] - -[[package]] -name = "crc-catalog" -version = "2.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" - [[package]] name = "crc-fast" version = "1.9.0" diff --git a/rust-runtime/aws-smithy-http-client/Cargo.toml b/rust-runtime/aws-smithy-http-client/Cargo.toml index 3f9c726ab6f..f15e884aef6 100644 --- a/rust-runtime/aws-smithy-http-client/Cargo.toml +++ b/rust-runtime/aws-smithy-http-client/Cargo.toml @@ -90,21 +90,21 @@ tracing = "0.1.44" # hyper 1.x stack hyper = { version = "1.6.0", features = ["client", "http1", "http2"], optional = true } -hyper-util = { version = "0.1.16", features = ["http1", "http2", "client-pool"], optional = true } +hyper-util = { version = "0.1.20", features = ["http1", "http2", "client-pool"], optional = true } http-1x = { package = "http", version = "1.3.1" , optional = true } http-body-1x = { package = "http-body", version = "1.0.1", optional = true} -hyper-rustls = { version = "0.27", features = ["http2", "http1", "native-tokio", "tls12"], default-features = false, optional = true } +hyper-rustls = { version = "0.27.9", features = ["http2", "http1", "native-tokio", "tls12"], default-features = false, optional = true } rustls = { version = "0.23.31", default-features = false, optional = true } -tokio-rustls = { version = "0.26.2", default-features = false, optional = true} +tokio-rustls = { version = "0.26.4", default-features = false, optional = true} # TODO(hyper1): add a way to enable the fips feature flag in s2n-tls s2n-tls-hyper = { version = "0.1.0", optional = true } -s2n-tls = { version = "0.3.33", optional = true } -s2n-tls-tokio = { version = "0.3.33", optional = true} +s2n-tls = { version = "0.3.36", optional = true } +s2n-tls-tokio = { version = "0.3.36", optional = true} tower = { version = "0.5.2", optional = true } tower-service = { version = "0.3", optional = true } futures-util = { version = "0.3.32", default-features = false, features = ["std"], optional = true } -rustls-pki-types = { version = "1.12.0", features = ["std"], optional = true } -rustls-native-certs = { version = "0.8.1", optional = true } +rustls-pki-types = { version = "1.14.1", features = ["std"], optional = true } +rustls-native-certs = { version = "0.8.3", optional = true } # end hyper 1.x stack deps # legacy hyper-0.14.x stack the SDK/runtime GA'd with @@ -131,12 +131,12 @@ aws-smithy-runtime-api = { path = "../aws-smithy-runtime-api", features = ["test aws-smithy-types = { path = "../aws-smithy-types", features = ["http-body-0-4-x", "test-util"] } http-body-util = { version = "0.1.3" } serial_test = "3.2" -hyper-util = { version = "0.1.16", features = ["full"] } +hyper-util = { version = "0.1.20", features = ["full"] } base64 = "0.22" rustls-pemfile = "2.2.0" -rustls-pki-types = { version = "1.12.0", features = ["std"] } +rustls-pki-types = { version = "1.14.1", features = ["std"] } tokio = { version = "1.49.0", features = ["macros", "rt", "rt-multi-thread", "test-util", "full"] } -tokio-rustls = "0.26.2" +tokio-rustls = "0.26.4" # used by tests vendored from hyper-util in src/client/pool/vendored_cache.rs tower-test = "0.4" # Test-only ANSI-free subscriber for capturing pool tracing events in assertions. diff --git a/rust-runtime/aws-smithy-http-client/additional-ci b/rust-runtime/aws-smithy-http-client/additional-ci index ee5eb936db7..260a02ec686 100755 --- a/rust-runtime/aws-smithy-http-client/additional-ci +++ b/rust-runtime/aws-smithy-http-client/additional-ci @@ -15,3 +15,6 @@ RUSTFLAGS="--cfg aws_sdk_unstable" cargo test --features rustls-ring echo "### Running connection pool integration tests" cargo test --features wire-mock,default-client --test connection_harness_test + +echo "### Running tracing-assertion tests (require single-threaded execution)" +cargo test --features wire-mock,default-client,test-util --test pool_behavior_test -- --ignored --test-threads=1 diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index a21a484f24c..57ebb333945 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -75,6 +75,25 @@ pub enum CloseReason { PoolDropped, } +/// Timing breakdown for connection establishment. +#[derive(Clone, Copy, Debug)] +pub struct ConnectionTiming { + /// TCP connect + TLS handshake combined. Measured from connector call + /// start to connected IO stream returned. + connect_duration: Duration, +} + +impl ConnectionTiming { + pub(crate) fn new(connect_duration: Duration) -> Self { + Self { connect_duration } + } + + /// Total time to establish the transport (TCP + TLS). + pub fn connect_duration(&self) -> Duration { + self.connect_duration + } +} + /// Emitted when a new connection is established (TCP + TLS + HTTP handshake). #[derive(Debug)] pub struct ConnectionCreatedEvent { @@ -82,6 +101,7 @@ pub struct ConnectionCreatedEvent { authority: Authority, remote_addr: Option, protocol: NegotiatedProtocol, + timing: ConnectionTiming, } impl ConnectionCreatedEvent { @@ -90,12 +110,14 @@ impl ConnectionCreatedEvent { authority: Authority, remote_addr: Option, protocol: NegotiatedProtocol, + timing: ConnectionTiming, ) -> Self { Self { conn_id, authority, remote_addr, protocol, + timing, } } @@ -118,6 +140,11 @@ impl ConnectionCreatedEvent { pub fn protocol(&self) -> NegotiatedProtocol { self.protocol } + + /// Timing breakdown for connection establishment. + pub fn timing(&self) -> &ConnectionTiming { + &self.timing + } } /// Emitted when an existing idle connection is checked out from the pool. @@ -1016,6 +1043,124 @@ impl hyper::body::Body for GuardedBody { /// consumers above the Negotiate composition point work with. pub(crate) type CheckoutResponse = http_1x::Response>; +/// Wire-level IO wrapper that carries TCP connect timing. +/// +/// Sits below TLS in the connector stack. Pure passthrough; never +/// modifies data or buffers. Establishes the seam for future +/// wire-level byte counting and IO error observation. +pub(crate) struct TransportIo { + inner: IO, + tcp_connect_duration: Duration, +} + +impl TransportIo { + fn new(inner: IO, tcp_connect_duration: Duration) -> Self { + Self { + inner, + tcp_connect_duration, + } + } + + /// Duration of the TCP connect (3-way handshake). + pub(crate) fn tcp_connect_duration(&self) -> Duration { + self.tcp_connect_duration + } +} + +impl Unpin for TransportIo {} + +impl hyper::rt::Read for TransportIo { + fn poll_read( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: hyper::rt::ReadBufCursor<'_>, + ) -> Poll> { + std::pin::Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +impl hyper::rt::Write for TransportIo { + fn poll_write( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &[u8], + ) -> Poll> { + std::pin::Pin::new(&mut self.inner).poll_write(cx, buf) + } + + fn poll_flush( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + std::pin::Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + std::pin::Pin::new(&mut self.inner).poll_shutdown(cx) + } +} + +impl + hyper_util::client::legacy::connect::Connection for TransportIo +{ + fn connected(&self) -> hyper_util::client::legacy::connect::Connected { + self.inner.connected() + } +} + +/// Wraps a TCP connector to stamp the connect duration on the resulting IO. +/// +/// The output `TransportIo` carries the TCP connect timing. +/// Passed to the TLS connector which layers encryption on top without +/// losing access to the transport-level timing. +pub(crate) struct TimingConnector { + inner: C, +} + +impl TimingConnector { + pub(crate) fn new(inner: C) -> Self { + Self { inner } + } +} + +impl Clone for TimingConnector { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl tower::Service for TimingConnector +where + C: tower::Service, + C::Error: Into, + C::Future: Send + 'static, + IO: Send + 'static, +{ + type Response = TransportIo; + type Error = BoxError; + type Future = std::pin::Pin< + Box> + Send>, + >; + + fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll> { + self.inner.poll_ready(cx).map_err(Into::into) + } + + fn call(&mut self, uri: http_1x::Uri) -> Self::Future { + let start = Instant::now(); + let fut = self.inner.call(uri); + Box::pin(async move { + let io = fut.await.map_err(Into::into)?; + Ok(TransportIo::new(io, start.elapsed())) + }) + } +} + #[cfg(test)] mod tests { //! Unit tests for H2 active-stream tracking. diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index ce191a41939..a70f16ece5d 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -14,6 +14,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use std::time::Instant; use aws_smithy_runtime_api::client::connection::ConnectionId; use aws_smithy_types::body::SdkBody; @@ -25,7 +26,7 @@ use tower::Service; use super::connection::{ Authority, ConnectCtx, ConnectionCreatedEvent, ConnectionFailedEvent, ConnectionInfo, - ConnectionPermit, ManagedConnection, NegotiatedProtocol, + ConnectionPermit, ConnectionTiming, ManagedConnection, NegotiatedProtocol, }; use super::BoxError; @@ -313,6 +314,7 @@ where ); let fut = self.connector.call(ctx); Box::pin(async move { + let connect_start = Instant::now(); let (io, permit) = match fut.await.map_err(Into::into) { Ok(v) => v, Err(e) => { @@ -324,6 +326,7 @@ where return Err(e); } }; + let connect_duration = connect_start.elapsed(); let info = capture_info(&io, authority.clone()); let conn_id = hooks.next_conn_id(); @@ -356,6 +359,7 @@ where authority, info.remote_addr, NegotiatedProtocol::Http1, + ConnectionTiming::new(connect_duration), )); spawn_driver({ @@ -451,6 +455,7 @@ where let hooks = self.hooks.clone(); let authority = self.authority.clone(); Box::pin(async move { + let connect_start = Instant::now(); let (io, permit) = match fut.await.map_err(Into::into) { Ok(v) => v, Err(e) => { @@ -462,6 +467,7 @@ where return Err(e); } }; + let connect_duration = connect_start.elapsed(); let info = capture_info(&io, authority.clone()); let conn_id = hooks.next_conn_id(); @@ -494,6 +500,7 @@ where authority, info.remote_addr, NegotiatedProtocol::Http2, + ConnectionTiming::new(connect_duration), )); spawn_driver({ diff --git a/rust-runtime/aws-smithy-http-client/src/client/tls/rustls_provider.rs b/rust-runtime/aws-smithy-http-client/src/client/tls/rustls_provider.rs index c043c9468f7..291a2914893 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/tls/rustls_provider.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/tls/rustls_provider.rs @@ -202,12 +202,13 @@ pub(crate) mod build_connector { ) -> super::connect::RustTlsConnector { let client_config = create_rustls_client_config(crypto_mode, tls_context); conn.enforce_http(false); + let timed = crate::client::pool::connection::TimingConnector::new(conn); let https_connector = hyper_rustls::HttpsConnectorBuilder::new() .with_tls_config(client_config.clone()) .https_or_http() .enable_http1() .enable_http2() - .wrap_connector(conn); + .wrap_connector(timed); super::connect::RustTlsConnector::new(https_connector, client_config, proxy_config) } @@ -239,14 +240,14 @@ pub(crate) mod connect { #[derive(Debug, Clone)] pub(crate) struct RustTlsConnector { - https: hyper_rustls::HttpsConnector>, + https: hyper_rustls::HttpsConnector>>, tls_config: Arc, proxy_matcher: Option>, // Pre-computed for performance } impl RustTlsConnector { pub(super) fn new( - https: hyper_rustls::HttpsConnector>, + https: hyper_rustls::HttpsConnector>>, tls_config: rustls::ClientConfig, proxy_config: ProxyConfig, ) -> Self { @@ -433,6 +434,30 @@ pub(crate) mod connect { } } } + + impl Connection + for RustTlsConn< + TokioIo< + MaybeHttpsStream< + crate::client::pool::connection::TransportIo>, + >, + >, + > + { + fn connected(&self) -> Connected { + if self.inner.inner().get_ref().1.alpn_protocol() == Some(b"h2") { + self.inner + .inner() + .get_ref() + .0 + .inner() + .connected() + .negotiated_h2() + } else { + self.inner.inner().get_ref().0.inner().connected() + } + } + } impl Read for RustTlsConn { fn poll_read( self: Pin<&mut Self>, diff --git a/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs b/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs index 888221f0de7..4d0244bbe25 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs @@ -60,10 +60,9 @@ pub(crate) mod build_connector { ) -> super::connect::S2nTlsConnector { let config = tls_context.s2n_config(); http_connector.enforce_http(false); - let mut builder = s2n_tls_hyper::connector::HttpsConnector::builder_with_http( - http_connector, - config.clone(), - ); + let timed = crate::client::pool::connection::TimingConnector::new(http_connector); + let mut builder = + s2n_tls_hyper::connector::HttpsConnector::builder_with_http(timed, config.clone()); builder.with_plaintext_http(true); let https_connector = builder.build(); @@ -89,16 +88,20 @@ pub(crate) mod connect { }; use tower::Service; + type S2nHttpsConnector = s2n_tls_hyper::connector::HttpsConnector< + crate::client::pool::connection::TimingConnector>, + >; + #[derive(Clone)] pub(crate) struct S2nTlsConnector { - https: s2n_tls_hyper::connector::HttpsConnector>, + https: S2nHttpsConnector, tls_config: s2n_tls::config::Config, proxy_matcher: Option>, // Pre-computed for performance } impl S2nTlsConnector { pub(super) fn new( - https: s2n_tls_hyper::connector::HttpsConnector>, + https: S2nHttpsConnector, tls_config: s2n_tls::config::Config, proxy_config: ProxyConfig, ) -> Self { diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index d0b0e4dc02f..ac251b192ed 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -9,7 +9,7 @@ pub use super::pool::connection::{ Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, - ConnectionFailedEvent, ConnectionReusedEvent, NegotiatedProtocol, + ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, }; use std::borrow::Cow; diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index f9671db4869..0ac96c16638 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1093,23 +1093,14 @@ async fn v2_read_timeout() { } // --------------------------------------------------------------------------- -// Tracing-assertion tests for the background eviction task. -// -// These are v2-only because they assert on internal tracing events that the -// v2 pool emits. The existing `v2_idle_timeout_eviction` test proves eviction -// works end-to-end (new TCP on second request), but doesn't distinguish the -// background task from checkout-time filtering. These tests assert the -// background task specifically ran without issuing a second request. +// Tracing output assertions // --------------------------------------------------------------------------- +// +// These tests verify the pool emits useful structured tracing events. +// They use a thread-local subscriber and `current_thread` tokio so spawned +// tasks (eviction) run on the same thread. Must be run with +// `--test-threads=1` to avoid subscriber interference from parallel tests. -/// Capture `aws_smithy_http_client` tracing events at TRACE level for the -/// duration of a test. ANSI color codes are explicitly disabled so string -/// matches on fields like `reason="idle_expired"` work — the default -/// `tracing_subscriber::fmt` output embeds ANSI escapes between the field -/// name and value which breaks naive `contains` assertions. -/// -/// Returns `(guard, buffer)`. The guard resets the subscriber on drop; -/// `buffer.lock().unwrap()` yields the captured bytes. fn capture_pool_logs() -> ( tracing::subscriber::DefaultGuard, std::sync::Arc>>, @@ -1143,12 +1134,15 @@ fn captured_str(buf: &std::sync::Arc>>) -> String { String::from_utf8(buf.lock().unwrap().clone()).expect("captured logs are utf-8") } -/// Verify the background eviction task runs and emits the "connection -/// evicted" event with `reason=idle_expired` after the idle timeout passes, -/// without any subsequent request being made. +/// Background eviction task emits structured tracing events including +/// pool init, connection established, eviction with reason, and host +/// entry removal. /// -/// Also verifies the host entry is removed once its cache becomes empty. -#[tokio::test] +/// Requires `--test-threads=1` due to thread-local subscriber; run via: +/// `cargo test --features default-client,test-util,wire-mock --test pool_behavior_test -- --ignored --test-threads=1` +#[tokio::test(flavor = "current_thread")] +#[serial_test::serial(tracing)] +#[ignore] async fn v2_background_eviction_emits_tracing_events() { let (_guard, logs) = capture_pool_logs(); @@ -1158,7 +1152,6 @@ async fn v2_background_eviction_emits_tracing_events() { let client = V2Client.make(ClientConfig::default().with_idle_timeout(idle_timeout)); let url = localhost_url(&server); - // One request — connection established, dispatched, returned to idle. let status = send_to(&client, &url) .await .expect("first request should succeed") @@ -1166,41 +1159,39 @@ async fn v2_background_eviction_emits_tracing_events() { .as_u16(); assert_eq!(status, 200); - // Wait past idle timeout + one tick interval. No second request here — - // any eviction must be driven by the background task. - tokio::time::sleep(idle_timeout * 3).await; + // Wait past idle timeout + eviction tick intervals. + tokio::time::sleep(idle_timeout * 5).await; let captured = captured_str(&logs); assert!( captured.contains("pool: initialized"), - "expected pool init log" + "expected pool init log. captured:\n{captured}" ); assert!( captured.contains("pool: eviction task spawned"), - "expected eviction task spawn log" + "expected eviction task spawn log. captured:\n{captured}" ); assert!( captured.contains("pool: connection established"), - "expected connection-established log" + "expected connection-established log. captured:\n{captured}" ); assert!( captured.contains("pool: connection evicted"), - "expected 'connection evicted' log — background task should have evicted the idle connection" - ); - assert!( - captured.contains("reason=\"idle_expired\""), - "expected idle_expired reason. captured:\n{captured}" + "expected connection-evicted log. captured:\n{captured}" ); assert!( captured.contains("pool: host entry removed"), - "expected host-entry-removed log once cache went empty" + "expected host-entry-removed log. captured:\n{captured}" ); } -/// Verify connections carry stable `conn_id` values — first established -/// connection gets id 0, subsequent new connection (after eviction forces -/// a fresh TCP) gets id 1. -#[tokio::test] +/// conn_id is stable (same across a connection's lifetime) and monotonically +/// increasing (new connection after eviction gets the next id). +/// +/// Requires `--test-threads=1`; see `v2_background_eviction_emits_tracing_events`. +#[tokio::test(flavor = "current_thread")] +#[serial_test::serial(tracing)] +#[ignore] async fn v2_conn_id_is_stable_and_monotonic() { let (_guard, logs) = capture_pool_logs(); @@ -1212,14 +1203,10 @@ async fn v2_conn_id_is_stable_and_monotonic() { let url = localhost_url(&server); send_to(&client, &url).await.expect("first request"); - tokio::time::sleep(idle_timeout * 3).await; + tokio::time::sleep(idle_timeout * 5).await; send_to(&client, &url).await.expect("second request"); let captured = captured_str(&logs); - assert!( - captured.contains("pool: connection established"), - "no 'connection established' log captured" - ); assert!( captured.contains("conn_id=0"), "first connection should be conn_id=0. captured:\n{captured}" @@ -1228,7 +1215,6 @@ async fn v2_conn_id_is_stable_and_monotonic() { captured.contains("conn_id=1"), "second connection (after eviction) should be conn_id=1. captured:\n{captured}" ); - assert!(captured.contains("pool: connection evicted")); } // --------------------------------------------------------------------------- @@ -1430,13 +1416,13 @@ mod listener_tests { let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); // First request: on_created - super::send_and_read_body(&client, &url).await; + send_and_read_body(&client, &url).await; assert_eq!(listener.created.lock().unwrap().len(), 1); assert_eq!(listener.created.lock().unwrap()[0].0, 0); assert!(listener.created.lock().unwrap()[0].1.contains("127.0.0.1")); // Second request: on_reused - super::send_and_read_body(&client, &url).await; + send_and_read_body(&client, &url).await; assert_eq!(listener.reused.lock().unwrap().len(), 1); assert_eq!(listener.reused.lock().unwrap()[0].0, 0); @@ -1478,7 +1464,7 @@ mod listener_tests { meta.unwrap().poison(); // The poisoned connection is discarded on next checkout attempt - super::send_and_read_body(&client, &url).await; + send_and_read_body(&client, &url).await; let closed = listener.closed.lock().unwrap(); assert_eq!(closed.len(), 1); @@ -1529,7 +1515,7 @@ mod listener_tests { let port = harness.endpoints[0].port(); let url = format!("http://127.0.0.1:{}/path", port); - super::send_and_read_body(&client, &url).await; + send_and_read_body(&client, &url).await; let created = listener.created.lock().unwrap(); assert_eq!(created[0].1, format!("127.0.0.1:{}", port)); @@ -1563,14 +1549,14 @@ mod listener_tests { let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); // First request succeeds; connection returns to pool. - super::send_and_read_body(&client, &url).await; + send_and_read_body(&client, &url).await; // Brief pause for the server-side close to propagate. tokio::time::sleep(Duration::from_millis(50)).await; // Second request: pool checks out the dead connection, detects it // in poll_ready, fires on_closed(Unusable), then creates a new one. - super::send_and_read_body(&client, &url).await; + send_and_read_body(&client, &url).await; let closed = listener.closed.lock().unwrap(); let unusable = closed.iter().find(|c| c.2 == CloseReason::Unusable); @@ -1581,4 +1567,37 @@ mod listener_tests { ); assert_eq!(unusable.unwrap().0, 0, "should be the first connection"); } + + /// ConnectionCreatedEvent carries non-zero connect_duration. + #[tokio::test] + async fn listener_timing_populated_on_created() { + let timing: Arc>> = Arc::new(Mutex::new(None)); + + struct TimingListener(Arc>>); + impl ConnectionEventListener for TimingListener { + fn on_created(&self, event: &ConnectionCreatedEvent) { + *self.0.lock().unwrap() = Some(event.timing().connect_duration()); + } + } + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"ok", + }], + ) + .build() + .await; + + let listener: Arc = Arc::new(TimingListener(timing.clone())); + let client = make_v2_with_listener(&harness, None, listener); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + send_and_read_body(&client, &url).await; + + let duration = timing.lock().unwrap().expect("timing should be set"); + assert!(duration > Duration::ZERO, "connect_duration should be > 0"); + } } diff --git a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs index 406c22993d0..0d8276ad66e 100644 --- a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs +++ b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs @@ -173,3 +173,69 @@ async fn v2_s2n_tls_client() { .build_https(); smoke_test_client(&client).await.unwrap(); } + +#[cfg(feature = "s2n-tls")] +#[tokio::test] +async fn v2_s2n_tls_timing_populated() { + use aws_smithy_http_client::v2::{ConnectionCreatedEvent, ConnectionEventListener}; + use std::sync::Mutex; + use std::time::Duration; + + struct TimingListener(Mutex>); + impl ConnectionEventListener for TimingListener { + fn on_created(&self, event: &ConnectionCreatedEvent) { + *self.0.lock().unwrap() = Some(event.timing().connect_duration()); + } + } + + let listener = Arc::new(TimingListener(Mutex::new(None))); + let client = BuilderV2::new() + .tls_provider(tls::Provider::S2nTls) + .connection_event_listener(listener.clone() as Arc) + .build_https(); + smoke_test_client(&client).await.unwrap(); + + let duration = listener + .0 + .lock() + .unwrap() + .expect("timing should be populated"); + assert!( + duration > Duration::ZERO, + "connect_duration should be > 0 (was {duration:?})" + ); +} + +#[cfg(feature = "rustls-aws-lc")] +#[tokio::test] +async fn v2_rustls_timing_populated() { + use aws_smithy_http_client::v2::{ConnectionCreatedEvent, ConnectionEventListener}; + use std::sync::Mutex; + use std::time::Duration; + + struct TimingListener(Mutex>); + impl ConnectionEventListener for TimingListener { + fn on_created(&self, event: &ConnectionCreatedEvent) { + *self.0.lock().unwrap() = Some(event.timing().connect_duration()); + } + } + + let listener = Arc::new(TimingListener(Mutex::new(None))); + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLc, + )) + .connection_event_listener(listener.clone() as Arc) + .build_https(); + smoke_test_client(&client).await.unwrap(); + + let duration = listener + .0 + .lock() + .unwrap() + .expect("timing should be populated"); + assert!( + duration > Duration::ZERO, + "connect_duration should be > 0 (was {duration:?})" + ); +} From 84c2c87f45ec4068c04380703fc588d287669916 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 19 May 2026 15:26:07 -0400 Subject: [PATCH 22/47] h2 ALPN tests and fix s2n when proxied --- aws/rust-runtime/Cargo.lock | 2 +- rust-runtime/Cargo.lock | 2 +- .../aws-smithy-http-client/Cargo.toml | 2 +- .../aws-smithy-http-client/additional-ci | 3 + .../src/client/tls/s2n_tls_provider.rs | 8 +- .../aws-smithy-http-client/src/client/v2.rs | 30 + .../tests/h2_pool_test.rs | 679 ++++++++++++++++++ 7 files changed, 720 insertions(+), 6 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs diff --git a/aws/rust-runtime/Cargo.lock b/aws/rust-runtime/Cargo.lock index a6cd6bcac17..dad8ae8767d 100644 --- a/aws/rust-runtime/Cargo.lock +++ b/aws/rust-runtime/Cargo.lock @@ -254,7 +254,7 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.2.0" +version = "1.3.0" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock index 1e08424eec8..77f9a408f7a 100644 --- a/rust-runtime/Cargo.lock +++ b/rust-runtime/Cargo.lock @@ -498,7 +498,7 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.2.0" +version = "1.3.0" dependencies = [ "aws-smithy-async 1.3.0", "aws-smithy-protocol-test", diff --git a/rust-runtime/aws-smithy-http-client/Cargo.toml b/rust-runtime/aws-smithy-http-client/Cargo.toml index f15e884aef6..ccfea9b06d4 100644 --- a/rust-runtime/aws-smithy-http-client/Cargo.toml +++ b/rust-runtime/aws-smithy-http-client/Cargo.toml @@ -2,7 +2,7 @@ name = "aws-smithy-http-client" authors = ["AWS Rust SDK Team "] description = "HTTP client abstractions for generated smithy clients" -version = "1.2.0" +version = "1.3.0" license = "Apache-2.0" edition = "2021" repository = "https://github.com/smithy-lang/smithy-rs" diff --git a/rust-runtime/aws-smithy-http-client/additional-ci b/rust-runtime/aws-smithy-http-client/additional-ci index 260a02ec686..bba927553ee 100755 --- a/rust-runtime/aws-smithy-http-client/additional-ci +++ b/rust-runtime/aws-smithy-http-client/additional-ci @@ -18,3 +18,6 @@ cargo test --features wire-mock,default-client --test connection_harness_test echo "### Running tracing-assertion tests (require single-threaded execution)" cargo test --features wire-mock,default-client,test-util --test pool_behavior_test -- --ignored --test-threads=1 + +echo "### Running H2 pool behavior tests (requires aws_sdk_unstable cfg)" +RUSTFLAGS='--cfg aws_sdk_unstable' cargo test --features s2n-tls,rustls-aws-lc,wire-mock,default-client,test-util --test h2_pool_test diff --git a/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs b/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs index 4d0244bbe25..2de445640fb 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/tls/s2n_tls_provider.rs @@ -262,9 +262,11 @@ pub(crate) mod connect { T: Connection + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, { fn connected(&self) -> Connected { - // For tunneled connections, we can't easily access the underlying connection info - // from s2n-tls, so we'll return a basic Connected instance - Connected::new() + let inner_connected = self.inner.inner().get_ref().connected(); + match self.inner.inner().as_ref().application_protocol() { + Some(b"h2") => inner_connected.negotiated_h2(), + _ => inner_connected, + } } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/v2.rs index ac251b192ed..e758c2e43c2 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/v2.rs @@ -231,6 +231,36 @@ impl BuilderV2 { let pool = build_http_pool_with_proxy(tcp, &self.proxy_config, config); V2HttpClient::new(pool, proxy_matcher).into_shared() } + + /// Build an HTTP client from a raw TCP-level connector. + /// + /// The connector must be a `tower::Service` producing an IO type + /// that implements hyper's `Read`, `Write`, and `Connection` traits. + /// The pool's Negotiate layer uses `Connection::connected().is_negotiated_h2()` + /// to route connections to the H2 path. + #[cfg(all(feature = "test-util", aws_sdk_unstable))] + #[doc(hidden)] + pub fn build_http_with_tcp_connector(self, connector: C) -> SharedHttpClient + where + C: tower::Service + Clone + Send + Sync + 'static, + C::Error: Into + 'static, + C::Future: Unpin + Send + 'static, + IO: hyper::rt::Read + + hyper::rt::Write + + hyper_util::client::legacy::connect::Connection + + Unpin + + Send + + 'static, + { + let config = pool::PoolConfig { + max_connections: self.max_connections, + max_connections_per_host: self.max_connections_per_host, + pool_idle_timeout: self.pool_idle_timeout, + connection_event_listener: self.connection_event_listener.clone(), + }; + let pool = pool::build_pool(connector, config); + V2HttpClient::new(pool, None).into_shared() + } } /// Build a no-TLS pool that honors `proxy_config`. Wraps the TCP connector diff --git a/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs new file mode 100644 index 00000000000..3d13b047b10 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs @@ -0,0 +1,679 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! H2 connection pool behavior tests. +//! +//! Tests the v2 pool's HTTP/2 path: multiplexing, GOAWAY handling, +//! connection poisoning, and stream limits. +//! +//! Uses plain TCP with a fake ALPN signal (`Connected::new().negotiated_h2()`) +//! so no TLS infrastructure is needed. The pool's Negotiate layer trusts +//! the `Connected` metadata to route to the H2 path. + +#![cfg(all( + feature = "wire-mock", + feature = "default-client", + feature = "test-util", + aws_sdk_unstable +))] + +use aws_smithy_http_client::v2::BuilderV2; +use aws_smithy_runtime_api::client::http::{ + HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, +}; +use aws_smithy_runtime_api::client::orchestrator::HttpRequest; +use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder; +use bytes::Bytes; +use h2::server::SendResponse; +use h2::RecvStream; +use http_body_util::BodyExt; +use hyper_util::client::legacy::connect::Connected; +use std::future::Future; +use std::net::SocketAddr; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio::net::{TcpListener, TcpStream}; +use tower::Service; + +// --------------------------------------------------------------------------- +// H2MockServer — plain TCP server speaking H2 via the h2 crate directly +// --------------------------------------------------------------------------- + +/// Handler function type for H2 requests. +type H2Handler = + Arc, SendResponse) + Send + Sync + 'static>; + +struct H2MockServer { + addr: SocketAddr, + /// Total H2 connections accepted (each connection can multiplex many streams). + connections: Arc, + /// Total streams (requests) handled across all connections. + streams: Arc, + _shutdown: tokio::sync::oneshot::Sender<()>, +} + +impl H2MockServer { + /// Start an H2 server that responds 200 with the given body to every request. + async fn start(body: &'static str) -> Self { + Self::start_with_handler(Arc::new(move |_req, mut respond| { + let response = http_1x::Response::builder().status(200).body(()).unwrap(); + let mut send_stream = respond.send_response(response, false).unwrap(); + send_stream.send_data(Bytes::from(body), true).unwrap(); + })) + .await + } + + /// Start with a custom handler for each stream. + async fn start_with_handler(handler: H2Handler) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let connections = Arc::new(AtomicUsize::new(0)); + let streams = Arc::new(AtomicUsize::new(0)); + let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + + let conns = connections.clone(); + let strms = streams.clone(); + + tokio::spawn(async move { + loop { + tokio::select! { + accept = listener.accept() => { + let (stream, _) = match accept { + Ok(v) => v, + Err(_) => break, + }; + conns.fetch_add(1, Ordering::SeqCst); + let handler = handler.clone(); + let strms = strms.clone(); + tokio::spawn(async move { + let mut conn = h2::server::Builder::new() + .handshake(stream) + .await + .unwrap(); + while let Some(result) = conn.accept().await { + let (req, respond) = result.unwrap(); + strms.fetch_add(1, Ordering::SeqCst); + handler(req, respond); + } + }); + } + _ = &mut shutdown_rx => break, + } + } + }); + + Self { + addr, + connections, + streams, + _shutdown: shutdown_tx, + } + } + + /// Start a server that sends GOAWAY after `n` streams on each connection. + async fn start_goaway_after(n: usize, body: &'static str) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let connections = Arc::new(AtomicUsize::new(0)); + let streams = Arc::new(AtomicUsize::new(0)); + let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + + let conns = connections.clone(); + let strms = streams.clone(); + + tokio::spawn(async move { + loop { + tokio::select! { + accept = listener.accept() => { + let (stream, _) = match accept { + Ok(v) => v, + Err(_) => break, + }; + conns.fetch_add(1, Ordering::SeqCst); + let strms = strms.clone(); + let per_conn_count = Arc::new(AtomicUsize::new(0)); + tokio::spawn(async move { + let mut conn = h2::server::Builder::new() + .handshake(stream) + .await + .unwrap(); + while let Some(result) = conn.accept().await { + let (req, mut respond) = result.unwrap(); + strms.fetch_add(1, Ordering::SeqCst); + let count = per_conn_count.fetch_add(1, Ordering::SeqCst) + 1; + + // Respond normally + let response = http_1x::Response::builder() + .status(200) + .body(()) + .unwrap(); + let mut send_stream = respond.send_response(response, false).unwrap(); + send_stream.send_data(Bytes::from(body), true).unwrap(); + drop(req); + + // After n streams, send GOAWAY + if count >= n { + conn.graceful_shutdown(); + } + } + }); + } + _ = &mut shutdown_rx => break, + } + } + }); + + Self { + addr, + connections, + streams, + _shutdown: shutdown_tx, + } + } + + fn connection_count(&self) -> usize { + self.connections.load(Ordering::SeqCst) + } + + fn stream_count(&self) -> usize { + self.streams.load(Ordering::SeqCst) + } + + fn url(&self) -> String { + format!("http://127.0.0.1:{}/", self.addr.port()) + } +} + +// --------------------------------------------------------------------------- +// H2Connector — connects via TCP, signals negotiated_h2() +// --------------------------------------------------------------------------- + +/// IO wrapper that signals H2 negotiation to the pool's Negotiate layer. +struct H2Io { + inner: TcpStream, +} + +impl hyper_util::client::legacy::connect::Connection for H2Io { + fn connected(&self) -> Connected { + Connected::new().negotiated_h2() + } +} + +impl AsyncRead for H2Io { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for H2Io { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} + +// hyper::rt::Read and Write are needed by the pool +impl hyper::rt::Read for H2Io { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + mut buf: hyper::rt::ReadBufCursor<'_>, + ) -> Poll> { + let n = unsafe { + let mut tbuf = ReadBuf::uninit(buf.as_mut()); + match Pin::new(&mut self.get_mut().inner).poll_read(cx, &mut tbuf) { + Poll::Ready(Ok(())) => tbuf.filled().len(), + other => return other, + } + }; + unsafe { buf.advance(n) }; + Poll::Ready(Ok(())) + } +} + +impl hyper::rt::Write for H2Io { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} + +/// Connector that establishes TCP connections and signals H2 negotiation. +#[derive(Clone)] +struct H2Connector { + addr: SocketAddr, +} + +impl Service for H2Connector { + type Response = H2Io; + type Error = Box; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _req: http_1x::Uri) -> Self::Future { + let addr = self.addr; + Box::pin(async move { + let stream = TcpStream::connect(addr).await?; + Ok(H2Io { inner: stream }) + }) + } +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +fn build_h2_client(server: &H2MockServer) -> SharedHttpClient { + BuilderV2::new().build_http_with_tcp_connector(H2Connector { addr: server.addr }) +} + +fn runtime_components() -> aws_smithy_runtime_api::client::runtime_components::RuntimeComponents { + RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(aws_smithy_async::time::SystemTimeSource::new())) + .build() + .expect("valid runtime components") +} + +async fn send_request( + client: &SharedHttpClient, + url: &str, +) -> Result<(u16, Vec), aws_smithy_runtime_api::client::result::ConnectorError> { + let settings = HttpConnectorSettings::builder().build(); + let components = runtime_components(); + let connector = client.http_connector(&settings, &components); + let resp = connector + .call(HttpRequest::get(url).expect("valid request")) + .await?; + let status = resp.status().as_u16(); + let body = resp + .into_body() + .collect() + .await + .expect("body") + .to_bytes() + .to_vec(); + Ok((status, body)) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +/// Multiple concurrent requests multiplex over a single H2 connection. +#[tokio::test] +async fn h2_multiplexing_shares_one_connection() { + let server = H2MockServer::start("ok").await; + let client = build_h2_client(&server); + let url = server.url(); + + // Warm the connection with one request first so the Singleton is populated + let (status, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + assert_eq!(server.connection_count(), 1); + + // Now send 4 concurrent requests — they should all multiplex on the existing connection + let futs: Vec<_> = (0..4).map(|_| send_request(&client, &url)).collect(); + let results = futures_util::future::join_all(futs).await; + + for (i, r) in results.iter().enumerate() { + let (status, _) = r + .as_ref() + .unwrap_or_else(|e| panic!("request {i} failed: {e}")); + assert_eq!(*status, 200); + } + + // All 5 requests (1 warm + 4 concurrent) should have used a single connection + assert_eq!( + server.connection_count(), + 1, + "H2 should multiplex on one connection" + ); + assert_eq!(server.stream_count(), 5, "should have 5 streams"); +} + +/// After GOAWAY, the pool establishes a new connection for subsequent requests. +#[tokio::test] +async fn h2_goaway_triggers_new_connection() { + // Server sends GOAWAY after 2 streams per connection + let server = H2MockServer::start_goaway_after(2, "ok").await; + let client = build_h2_client(&server); + let url = server.url(); + + // First 2 requests on connection 1 + let (s1, _) = send_request(&client, &url).await.unwrap(); + let (s2, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(s1, 200); + assert_eq!(s2, 200); + + // Give the pool time to observe the GOAWAY + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + + // Next request should go on a new connection + let (s3, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(s3, 200); + + assert!( + server.connection_count() >= 2, + "should have opened a second connection after GOAWAY, got {}", + server.connection_count() + ); +} + +/// Sequential requests reuse the same H2 connection (no unnecessary reconnects). +#[tokio::test] +async fn h2_sequential_requests_reuse_connection() { + let server = H2MockServer::start("hello").await; + let client = build_h2_client(&server); + let url = server.url(); + + for i in 0..5 { + let (status, body) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200, "request {i}"); + assert_eq!(body, b"hello", "request {i}"); + } + + assert_eq!( + server.connection_count(), + 1, + "should reuse one H2 connection" + ); + assert_eq!(server.stream_count(), 5); +} + +/// Poisoning an H2 connection forces the pool to establish a new one. +#[tokio::test] +async fn h2_poisoned_connection_not_reused() { + use aws_smithy_runtime_api::client::connection::CaptureSmithyConnection; + + let server = H2MockServer::start("ok").await; + let client = build_h2_client(&server); + let url = server.url(); + + // First request: establish H2 connection, capture metadata. + let settings = HttpConnectorSettings::builder().build(); + let components = runtime_components(); + let connector = client.http_connector(&settings, &components); + + let capture = CaptureSmithyConnection::new(); + let mut request = HttpRequest::get(&url).expect("valid request"); + request.add_extension(capture.clone()); + + let resp = connector + .call(request) + .await + .expect("request should succeed"); + let _body = resp.into_body().collect().await.expect("body"); + assert_eq!(server.connection_count(), 1); + + let metadata = capture.get().expect("adapter should populate metadata"); + metadata.poison(); + + // Give the pool a moment to observe the poison + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + + // Next request should open a NEW H2 connection + let (status, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + assert!( + server.connection_count() >= 2, + "poisoned H2 connection should not be reused, got {} connections", + server.connection_count() + ); +} + +// =========================================================================== +// TLS + ALPN — real certificate negotiation +// =========================================================================== +// +// These tests use a TLS server with self-signed certs advertising h2 via ALPN. +// They verify the v2 pool correctly routes to H2 when ALPN negotiates it +// through real TLS. + +#[cfg(feature = "rustls-aws-lc")] +mod tls_h2 { + use aws_smithy_http_client::tls; + use aws_smithy_http_client::tls::{TlsContext, TrustStore}; + use aws_smithy_http_client::v2::BuilderV2; + use aws_smithy_runtime_api::client::http::{ + HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, + }; + use aws_smithy_runtime_api::client::orchestrator::HttpRequest; + use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder; + use http_body_util::BodyExt; + use hyper_util::rt::TokioExecutor; + use std::net::SocketAddr; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use tokio::net::TcpListener; + use tokio_rustls::TlsAcceptor; + + /// TLS test server that tracks connection count and serves H2 via ALPN. + struct TlsH2Server { + addr: SocketAddr, + connections: Arc, + _shutdown: tokio::sync::oneshot::Sender<()>, + } + + impl TlsH2Server { + async fn start() -> Self { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let connections = Arc::new(AtomicUsize::new(0)); + let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + + // Load certs + let cert_pem = std::fs::read("tests/server.pem").unwrap(); + let key_pem = std::fs::read("tests/server.rsa").unwrap(); + let certs: Vec<_> = rustls_pemfile::certs(&mut &cert_pem[..]) + .collect::>() + .unwrap(); + let key = rustls_pemfile::private_key(&mut &key_pem[..]) + .unwrap() + .unwrap(); + + let mut server_config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certs, key) + .unwrap(); + // Only advertise h2 — force H2 negotiation + server_config.alpn_protocols = vec![b"h2".to_vec()]; + let tls_acceptor = TlsAcceptor::from(Arc::new(server_config)); + + let conns = connections.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + accept = listener.accept() => { + let (tcp, _) = match accept { + Ok(v) => v, + Err(_) => break, + }; + conns.fetch_add(1, Ordering::SeqCst); + let tls_acceptor = tls_acceptor.clone(); + tokio::spawn(async move { + let tls_stream = match tls_acceptor.accept(tcp).await { + Ok(s) => s, + Err(e) => { + eprintln!("TLS accept failed: {e}"); + return; + } + }; + let service = hyper::service::service_fn(|_req| async { + Ok::<_, hyper::Error>( + http_1x::Response::builder() + .status(200) + .body(http_body_util::Full::new( + bytes::Bytes::from("h2-ok"), + )) + .unwrap(), + ) + }); + let io = hyper_util::rt::TokioIo::new(tls_stream); + // Use http2 only server since we only advertise h2 + let _ = hyper_util::server::conn::auto::Builder::new(TokioExecutor::new()) + .serve_connection(io, service) + .await; + }); + } + _ = &mut shutdown_rx => break, + } + } + }); + + Self { + addr, + connections, + _shutdown: shutdown_tx, + } + } + + fn connection_count(&self) -> usize { + self.connections.load(Ordering::SeqCst) + } + + fn url(&self) -> String { + format!("https://localhost:{}/", self.addr.port()) + } + } + + fn tls_context() -> TlsContext { + let pem = std::fs::read("tests/server.pem").unwrap(); + let trust_store = TrustStore::empty().with_pem_certificate(pem); + TlsContext::builder() + .with_trust_store(trust_store) + .build() + .unwrap() + } + + fn runtime_components() -> aws_smithy_runtime_api::client::runtime_components::RuntimeComponents + { + RuntimeComponentsBuilder::for_tests() + .with_time_source(Some(aws_smithy_async::time::SystemTimeSource::new())) + .build() + .unwrap() + } + + async fn send_request( + client: &SharedHttpClient, + url: &str, + ) -> Result<(u16, Vec), aws_smithy_runtime_api::client::result::ConnectorError> { + let settings = HttpConnectorSettings::builder().build(); + let components = runtime_components(); + let connector = client.http_connector(&settings, &components); + let resp = connector + .call(HttpRequest::get(url).expect("valid request")) + .await?; + let status = resp.status().as_u16(); + let body = resp + .into_body() + .collect() + .await + .expect("body") + .to_bytes() + .to_vec(); + Ok((status, body)) + } + + /// Rustls + ALPN h2: v2 pool routes to H2, multiple requests multiplex. + #[tokio::test] + async fn rustls_alpn_h2_multiplexing() { + let server = TlsH2Server::start().await; + let client = BuilderV2::new() + .tls_provider(tls::Provider::Rustls( + tls::rustls_provider::CryptoMode::AwsLc, + )) + .tls_context(tls_context()) + .build_https(); + + let url = server.url(); + + // First request establishes connection + let (status, body) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + assert_eq!(body, b"h2-ok"); + + // Second request should multiplex on same connection + let (status, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + + // Third request too + let (status, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + + assert_eq!( + server.connection_count(), + 1, + "rustls H2: all requests should multiplex on one connection" + ); + } + + /// s2n-tls + ALPN h2: v2 pool should route to H2 (currently broken — opens + /// new connection per request because Connected never signals negotiated_h2). + #[cfg(feature = "s2n-tls")] + #[tokio::test] + async fn s2n_tls_alpn_h2_multiplexing() { + let server = TlsH2Server::start().await; + let client = BuilderV2::new() + .tls_provider(tls::Provider::S2nTls) + .tls_context(tls_context()) + .build_https(); + + let url = server.url(); + + // First request + let (status, body) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + assert_eq!(body, b"h2-ok"); + + // Second request — if H2 works, should multiplex on same connection + let (status, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + + // Third request + let (status, _) = send_request(&client, &url).await.unwrap(); + assert_eq!(status, 200); + + // This assertion will FAIL until s2n-tls signals negotiated_h2() + assert_eq!( + server.connection_count(), + 1, + "s2n-tls H2: all requests should multiplex on one connection" + ); + } +} From c9ffb257581d83fb39486c0c30da99f39db96c86 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 20 May 2026 15:39:46 -0400 Subject: [PATCH 23/47] move types around --- .../aws-smithy-http-client/src/client.rs | 8 +- .../aws-smithy-http-client/src/client/pool.rs | 14 ++- .../src/client/{v2.rs => pool/builder.rs} | 89 ++++++++++--------- .../aws-smithy-http-client/src/lib.rs | 2 +- .../tests/h2_pool_test.rs | 10 +-- .../tests/pool_behavior_test.rs | 6 +- .../tests/proxy_tests.rs | 12 +-- .../tests/smoke_test_clients.rs | 18 ++-- .../aws-smithy-http-client/tests/tls.rs | 18 ++-- 9 files changed, 93 insertions(+), 84 deletions(-) rename rust-runtime/aws-smithy-http-client/src/client/{v2.rs => pool/builder.rs} (90%) diff --git a/rust-runtime/aws-smithy-http-client/src/client.rs b/rust-runtime/aws-smithy-http-client/src/client.rs index 25afea5ce7b..35e3b8ed1e8 100644 --- a/rust-runtime/aws-smithy-http-client/src/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client.rs @@ -11,9 +11,7 @@ mod timeout; pub mod tls; pub(crate) mod connect; -pub(crate) mod pool; -/// V2 HTTP client backed by composable connection pools. -pub mod v2; +pub mod pool; use crate::cfg::cfg_tls; use crate::tls::TlsContext; @@ -932,8 +930,8 @@ impl Builder { /// management, supporting per-host connection reuse with HTTP/1.1 and /// HTTP/2 protocol negotiation. #[doc(hidden)] - pub fn new_v2() -> v2::BuilderV2 { - v2::BuilderV2::new() + pub fn new_v2() -> pool::Builder { + pool::Builder::new() } /// Returns a [`SharedHttpClient`] that calls the given `connector` function to select an HTTP(S) connector. diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index f3d5ad65663..07170585380 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -16,6 +16,16 @@ pub(crate) mod connection; mod handshake; mod vendored_cache; +pub mod builder; + +// Public re-exports — make pool::Builder, pool::ConnectionEventListener, etc. +// directly accessible without going through pool::builder. +pub use builder::{ + Authority, Builder, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, + ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, + NegotiatedProtocol, +}; + /// Connection-caching pool layer. /// /// Re-exports [`vendored_cache`], our vendored copy of hyper-util's @@ -45,8 +55,8 @@ use tokio::sync::{oneshot, Semaphore}; use tower::{Service, ServiceExt}; use connection::{ - Authority, CachedConnection, CheckoutResponse, CloseReason, ConnectionClosedEvent, - ConnectionGuard, GuardedBody, H2ConnectionRef, SingletonConnection, + CachedConnection, CheckoutResponse, ConnectionGuard, GuardedBody, H2ConnectionRef, + SingletonConnection, }; pub(crate) use connection::{ConnectCtx, ReadTimeoutHint, TimeoutContext}; use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake}; diff --git a/rust-runtime/aws-smithy-http-client/src/client/v2.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs similarity index 90% rename from rust-runtime/aws-smithy-http-client/src/client/v2.rs rename to rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index e758c2e43c2..b627f14ed05 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/v2.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -3,11 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! V2 HTTP client. +//! V2 HTTP client builder. //! -//! Entry point: [`BuilderV2::new`] or [`Builder::new_v2`](super::Builder::new_v2). +//! Entry point: [`Builder::new`] or [`crate::Builder::new_v2`]. -pub use super::pool::connection::{ +pub use super::connection::{ Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, }; @@ -28,10 +28,10 @@ use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; use aws_smithy_runtime_api::shared::IntoShared; use hyper_util::client::legacy::connect::HttpConnector as HyperHttpConnector; -use super::downcast_error; -use super::pool::{self, ConnectionPool, PoolConfig}; -use super::tls; -use super::{TlsProviderSelected, TlsUnset}; +use super::{ConnectionPool, PoolConfig}; +use crate::client::downcast_error; +use crate::client::tls; +use crate::client::{TlsProviderSelected, TlsUnset}; use crate::tls::TlsContext; /// Builder for the v2 HTTP client. @@ -42,27 +42,27 @@ use crate::tls::TlsContext; /// # Examples /// /// ```no_run -/// use aws_smithy_http_client::v2::BuilderV2; +/// use aws_smithy_http_client::pool::Builder; /// use std::time::Duration; /// -/// let client = BuilderV2::new() +/// let client = Builder::new() /// .pool_idle_timeout(Duration::from_secs(20)) /// .build_http(); /// ``` #[derive(Clone)] -pub struct BuilderV2 { +pub struct Builder { pool_idle_timeout: Option, tcp_nodelay: bool, max_connections: Option, max_connections_per_host: Option, - proxy_config: Option, - connection_event_listener: Option>, + proxy_config: Option, + connection_event_listener: Option>, tls: Tls, } -impl std::fmt::Debug for BuilderV2 { +impl std::fmt::Debug for Builder { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("BuilderV2") + f.debug_struct("Builder") .field("pool_idle_timeout", &self.pool_idle_timeout) .field("tcp_nodelay", &self.tcp_nodelay) .field("max_connections", &self.max_connections) @@ -76,7 +76,7 @@ impl std::fmt::Debug for BuilderV2 { } } -impl Default for BuilderV2 { +impl Default for Builder { fn default() -> Self { Self { pool_idle_timeout: None, @@ -91,7 +91,7 @@ impl Default for BuilderV2 { } // Methods available in any TLS state. -impl BuilderV2 { +impl Builder { /// Set the pool idle timeout. /// /// Connections idle longer than this duration are evicted from the pool. @@ -139,13 +139,16 @@ impl BuilderV2 { /// /// To use environment variables (`HTTP_PROXY`, `HTTPS_PROXY`, /// `NO_PROXY`), pass `ProxyConfig::from_env()`. - pub fn proxy_config(mut self, config: super::proxy::ProxyConfig) -> Self { + pub fn proxy_config(mut self, config: crate::client::proxy::ProxyConfig) -> Self { self.proxy_config = Some(config); self } /// Mutable-reference variant of [`proxy_config`](Self::proxy_config). - pub fn set_proxy_config(&mut self, config: Option) -> &mut Self { + pub fn set_proxy_config( + &mut self, + config: Option, + ) -> &mut Self { self.proxy_config = config; self } @@ -153,7 +156,7 @@ impl BuilderV2 { /// Set a listener for connection lifecycle events (created, reused, closed, failed). pub fn connection_event_listener( mut self, - listener: Arc, + listener: Arc, ) -> Self { self.connection_event_listener = Some(listener); self @@ -162,22 +165,22 @@ impl BuilderV2 { /// Mutable-reference variant of [`connection_event_listener`](Self::connection_event_listener). pub fn set_connection_event_listener( &mut self, - listener: Option>, + listener: Option>, ) -> &mut Self { self.connection_event_listener = listener; self } } -impl BuilderV2 { +impl Builder { /// Create a new v2 client builder. pub fn new() -> Self { Self::default() } /// Set the TLS implementation. - pub fn tls_provider(self, provider: tls::Provider) -> BuilderV2 { - BuilderV2 { + pub fn tls_provider(self, provider: tls::Provider) -> Builder { + Builder { pool_idle_timeout: self.pool_idle_timeout, tcp_nodelay: self.tcp_nodelay, max_connections: self.max_connections, @@ -252,13 +255,13 @@ impl BuilderV2 { + Send + 'static, { - let config = pool::PoolConfig { + let config = super::PoolConfig { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, pool_idle_timeout: self.pool_idle_timeout, connection_event_listener: self.connection_event_listener.clone(), }; - let pool = pool::build_pool(connector, config); + let pool = super::build_pool(connector, config); V2HttpClient::new(pool, None).into_shared() } } @@ -269,7 +272,7 @@ impl BuilderV2 { /// proxy will fail at handshake time. fn build_http_pool_with_proxy( tcp: HyperHttpConnector, - proxy_config: &Option, + proxy_config: &Option, config: PoolConfig, ) -> ConnectionPool where @@ -279,7 +282,7 @@ where R::Future: Send, R::Error: Into, { - use super::proxy; + use crate::client::proxy; let proxy_config = proxy_config .clone() @@ -294,14 +297,14 @@ where } if proxy_config.is_disabled() { - pool::build_pool(tcp, config) + super::build_pool(tcp, config) } else { - let connector = super::connect::HttpProxyConnector::new(tcp, proxy_config); - pool::build_pool(connector, config) + let connector = crate::client::connect::HttpProxyConnector::new(tcp, proxy_config); + super::build_pool(connector, config) } } -impl BuilderV2 { +impl Builder { /// Set the TLS context (custom trust store, etc.). pub fn tls_context(mut self, context: TlsContext) -> Self { self.tls.context = context; @@ -336,7 +339,7 @@ impl BuilderV2 { R::Future: Send, R::Error: Into, { - use super::proxy; + use crate::client::proxy; let config = PoolConfig { max_connections: self.max_connections, @@ -364,7 +367,7 @@ impl BuilderV2 { &self.tls.context, proxy_config, ); - let pool = pool::build_pool(connector, config); + let pool = super::build_pool(connector, config); V2HttpClient::new(pool, proxy_matcher).into_shared() } #[cfg(feature = "s2n-tls")] @@ -374,7 +377,7 @@ impl BuilderV2 { &self.tls.context, proxy_config, ); - let pool = pool::build_pool(connector, config); + let pool = super::build_pool(connector, config); V2HttpClient::new(pool, proxy_matcher).into_shared() } // Provider is #[non_exhaustive]; this arm is unreachable when any @@ -433,7 +436,7 @@ impl HttpConnector for PooledConnector { // HTTPS-through-proxy uses CONNECT tunneling and authenticates // inside the tunnel-establishment exchange, not via this header. if let Some(matcher) = proxy_matcher.as_ref() { - super::proxy::add_proxy_auth_header(&mut request, matcher); + crate::client::proxy::add_proxy_auth_header(&mut request, matcher); } // Create a request-scoped capture the pool will fill in with @@ -443,7 +446,7 @@ impl HttpConnector for PooledConnector { // `ConnectionMetadata::poison()`, it flips the `PoisonPill` on // the actual `ManagedConnection` selected for this request. if let Some(capture_smithy) = request.extensions().get::() { - let capture = pool::ConnectionMetadataCapture::new(); + let capture = super::ConnectionMetadataCapture::new(); let for_retriever = capture.clone(); capture_smithy.set_connection_retriever(move || for_retriever.get()); request.extensions_mut().insert(capture); @@ -454,11 +457,9 @@ impl HttpConnector for PooledConnector { // request dispatch: bounds request-write + response-headers, // nothing else. if let Some((duration, sleep)) = read_timeout.zip(sleep_impl.clone()) { - request - .extensions_mut() - .insert(pool::ReadTimeoutHint(pool::TimeoutContext::new( - duration, sleep, - ))); + request.extensions_mut().insert(super::ReadTimeoutHint( + super::TimeoutContext::new(duration, sleep), + )); } // Set Host header from the URI authority if not already present. @@ -485,11 +486,11 @@ impl HttpConnector for PooledConnector { // Build the connect context: routing URI + connect timeout. // Cache hits skip the connector entirely so connect_timeout is // automatically new-connection-only. - let connect_ctx = pool::ConnectCtx::new( + let connect_ctx = super::ConnectCtx::new( full_uri, connect_timeout .zip(sleep_impl) - .map(|(d, s)| pool::TimeoutContext::new(d, s)), + .map(|(d, s)| super::TimeoutContext::new(d, s)), ); let response = pool @@ -511,7 +512,7 @@ impl HttpConnector for PooledConnector { /// constructed by `http_connector` can share the (immutable) matcher /// without per-call allocation. fn proxy_matcher_from( - proxy_config: &Option, + proxy_config: &Option, ) -> Option> { proxy_config .as_ref() diff --git a/rust-runtime/aws-smithy-http-client/src/lib.rs b/rust-runtime/aws-smithy-http-client/src/lib.rs index 12eb5193950..559f149d4ff 100644 --- a/rust-runtime/aws-smithy-http-client/src/lib.rs +++ b/rust-runtime/aws-smithy-http-client/src/lib.rs @@ -46,7 +46,7 @@ pub use client::{default_connector, proxy, tls, Builder, Connector, ConnectorBui /// V2 HTTP client backed by composable connection pools. #[cfg(feature = "default-client")] -pub use client::v2; +pub use client::pool; #[cfg(feature = "test-util")] pub mod test_util; diff --git a/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs index 3d13b047b10..9c05d45bb3e 100644 --- a/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs @@ -19,7 +19,7 @@ aws_sdk_unstable ))] -use aws_smithy_http_client::v2::BuilderV2; +use aws_smithy_http_client::pool::Builder; use aws_smithy_runtime_api::client::http::{ HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, }; @@ -298,7 +298,7 @@ impl Service for H2Connector { // --------------------------------------------------------------------------- fn build_h2_client(server: &H2MockServer) -> SharedHttpClient { - BuilderV2::new().build_http_with_tcp_connector(H2Connector { addr: server.addr }) + Builder::new().build_http_with_tcp_connector(H2Connector { addr: server.addr }) } fn runtime_components() -> aws_smithy_runtime_api::client::runtime_components::RuntimeComponents { @@ -465,9 +465,9 @@ async fn h2_poisoned_connection_not_reused() { #[cfg(feature = "rustls-aws-lc")] mod tls_h2 { + use aws_smithy_http_client::pool::Builder; use aws_smithy_http_client::tls; use aws_smithy_http_client::tls::{TlsContext, TrustStore}; - use aws_smithy_http_client::v2::BuilderV2; use aws_smithy_runtime_api::client::http::{ HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, }; @@ -614,7 +614,7 @@ mod tls_h2 { #[tokio::test] async fn rustls_alpn_h2_multiplexing() { let server = TlsH2Server::start().await; - let client = BuilderV2::new() + let client = Builder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) @@ -649,7 +649,7 @@ mod tls_h2 { #[tokio::test] async fn s2n_tls_alpn_h2_multiplexing() { let server = TlsH2Server::start().await; - let client = BuilderV2::new() + let client = Builder::new() .tls_provider(tls::Provider::S2nTls) .tls_context(tls_context()) .build_https(); diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 0ac96c16638..fe424253608 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -12,11 +12,11 @@ #![cfg(all(feature = "wire-mock", feature = "default-client"))] use aws_smithy_async::time::SystemTimeSource; +use aws_smithy_http_client::pool::Builder as PoolBuilder; use aws_smithy_http_client::test_util::wire::connection::{ ConnectionBehavior, ConnectionTestHarness, }; use aws_smithy_http_client::test_util::wire::{ReplayedEvent, WireMockServer}; -use aws_smithy_http_client::v2::BuilderV2; use aws_smithy_http_client::{ev, match_events, Builder, Connector}; use aws_smithy_runtime_api::client::http::{ HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, @@ -82,7 +82,7 @@ struct V2Client; impl MakeClient for V2Client { fn make(&self, config: ClientConfig) -> SharedHttpClient { - let mut builder = BuilderV2::new(); + let mut builder = PoolBuilder::new(); if let Some(timeout) = config.idle_timeout { builder = builder.pool_idle_timeout(timeout); } @@ -1329,7 +1329,7 @@ async fn v2_connection_id_surfaced_through_metadata() { mod listener_tests { use super::*; - use aws_smithy_http_client::v2::{ + use aws_smithy_http_client::pool::{ CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, }; diff --git a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs index 7b5be9ad518..126c6a30c64 100644 --- a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs +++ b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs @@ -1205,9 +1205,9 @@ async fn make_v2_http_request_through_proxy( proxy_config: ProxyConfig, target_url: &str, ) -> Result<(StatusCode, String), Box> { - use aws_smithy_http_client::v2::BuilderV2; + use aws_smithy_http_client::pool::Builder; - let http_client = BuilderV2::new().proxy_config(proxy_config).build_http(); + let http_client = Builder::new().proxy_config(proxy_config).build_http(); let connector_settings = HttpConnectorSettings::builder().build(); let runtime_components = RuntimeComponentsBuilder::for_tests() .with_time_source(Some(SystemTimeSource::new())) @@ -1234,9 +1234,9 @@ async fn make_v2_https_request_through_proxy( target_url: &str, tls_provider: tls::Provider, ) -> Result<(StatusCode, String), Box> { - use aws_smithy_http_client::v2::BuilderV2; + use aws_smithy_http_client::pool::Builder; - let http_client = BuilderV2::new() + let http_client = Builder::new() .tls_provider(tls_provider) .proxy_config(proxy_config) .build_https(); @@ -1494,7 +1494,7 @@ async fn test_v2_direct_http_origin_uri_form() { #[tokio::test] async fn test_v2_set_proxy_config_none_clears_proxy() { - use aws_smithy_http_client::v2::BuilderV2; + use aws_smithy_http_client::pool::Builder; let mock_proxy = MockProxyServer::new(|_| { Response::builder() @@ -1511,7 +1511,7 @@ async fn test_v2_set_proxy_config_none_clears_proxy() { }) .await; - let mut builder = BuilderV2::new(); + let mut builder = Builder::new(); builder.set_proxy_config(Some( ProxyConfig::http(format!("http://{}", mock_proxy.addr())).unwrap(), )); diff --git a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs index 0d8276ad66e..3d4f6cb6981 100644 --- a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs +++ b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs @@ -130,12 +130,12 @@ async fn smoke_test_client(client: &dyn HttpClient) -> Result<(), Box // v2 smoke tests // --------------------------------------------------------------------------- -use aws_smithy_http_client::v2::BuilderV2; +use aws_smithy_http_client::pool::Builder as PoolBuilder; #[cfg(feature = "rustls-aws-lc")] #[tokio::test] async fn v2_aws_lc_client() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) @@ -146,7 +146,7 @@ async fn v2_aws_lc_client() { #[cfg(feature = "rustls-aws-lc-fips")] #[tokio::test] async fn v2_aws_lc_fips_client() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLcFips, )) @@ -157,7 +157,7 @@ async fn v2_aws_lc_fips_client() { #[cfg(feature = "rustls-ring")] #[tokio::test] async fn v2_ring_client() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::Ring, )) @@ -168,7 +168,7 @@ async fn v2_ring_client() { #[cfg(feature = "s2n-tls")] #[tokio::test] async fn v2_s2n_tls_client() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::S2nTls) .build_https(); smoke_test_client(&client).await.unwrap(); @@ -177,7 +177,7 @@ async fn v2_s2n_tls_client() { #[cfg(feature = "s2n-tls")] #[tokio::test] async fn v2_s2n_tls_timing_populated() { - use aws_smithy_http_client::v2::{ConnectionCreatedEvent, ConnectionEventListener}; + use aws_smithy_http_client::pool::{ConnectionCreatedEvent, ConnectionEventListener}; use std::sync::Mutex; use std::time::Duration; @@ -189,7 +189,7 @@ async fn v2_s2n_tls_timing_populated() { } let listener = Arc::new(TimingListener(Mutex::new(None))); - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::S2nTls) .connection_event_listener(listener.clone() as Arc) .build_https(); @@ -209,7 +209,7 @@ async fn v2_s2n_tls_timing_populated() { #[cfg(feature = "rustls-aws-lc")] #[tokio::test] async fn v2_rustls_timing_populated() { - use aws_smithy_http_client::v2::{ConnectionCreatedEvent, ConnectionEventListener}; + use aws_smithy_http_client::pool::{ConnectionCreatedEvent, ConnectionEventListener}; use std::sync::Mutex; use std::time::Duration; @@ -221,7 +221,7 @@ async fn v2_rustls_timing_populated() { } let listener = Arc::new(TimingListener(Mutex::new(None))); - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) diff --git a/rust-runtime/aws-smithy-http-client/tests/tls.rs b/rust-runtime/aws-smithy-http-client/tests/tls.rs index 4bd8d4db061..f1046a37aa5 100644 --- a/rust-runtime/aws-smithy-http-client/tests/tls.rs +++ b/rust-runtime/aws-smithy-http-client/tests/tls.rs @@ -311,13 +311,13 @@ async fn test_s2n_tls_custom_ca() { // v2 TLS tests // --------------------------------------------------------------------------- -use aws_smithy_http_client::v2::BuilderV2; +use aws_smithy_http_client::pool::Builder as PoolBuilder; #[cfg(feature = "rustls-aws-lc")] #[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] #[tokio::test] async fn test_v2_rustls_aws_lc_native_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) @@ -328,7 +328,7 @@ async fn test_v2_rustls_aws_lc_native_ca() { #[cfg(feature = "rustls-aws-lc")] #[tokio::test] async fn test_v2_rustls_aws_lc_custom_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) @@ -341,7 +341,7 @@ async fn test_v2_rustls_aws_lc_custom_ca() { #[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] #[tokio::test] async fn test_v2_rustls_aws_lc_fips_native_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLcFips, )) @@ -352,7 +352,7 @@ async fn test_v2_rustls_aws_lc_fips_native_ca() { #[cfg(feature = "rustls-aws-lc-fips")] #[tokio::test] async fn test_v2_rustls_aws_lc_fips_custom_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLcFips, )) @@ -365,7 +365,7 @@ async fn test_v2_rustls_aws_lc_fips_custom_ca() { #[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] #[tokio::test] async fn test_v2_rustls_ring_native_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::Ring, )) @@ -376,7 +376,7 @@ async fn test_v2_rustls_ring_native_ca() { #[cfg(feature = "rustls-ring")] #[tokio::test] async fn test_v2_rustls_ring_custom_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::Ring, )) @@ -389,7 +389,7 @@ async fn test_v2_rustls_ring_custom_ca() { #[should_panic(expected = "Certificate is untrusted")] #[tokio::test] async fn test_v2_s2n_native_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::S2nTls) .build_https(); run_tls_test(&client).await.unwrap() @@ -398,7 +398,7 @@ async fn test_v2_s2n_native_ca() { #[cfg(feature = "s2n-tls")] #[tokio::test] async fn test_v2_s2n_tls_custom_ca() { - let client = BuilderV2::new() + let client = PoolBuilder::new() .tls_provider(tls::Provider::S2nTls) .tls_context(tls_context_from_pem("tests/server.pem")) .build_https(); From 284ed89532c726721d783ee022007653c88a4f0f Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 20 May 2026 23:14:35 -0400 Subject: [PATCH 24/47] initial partition --- .../aws-smithy-http-client/src/client/pool.rs | 2 + .../src/client/pool/builder.rs | 27 ++ .../src/client/pool/partition.rs | 266 ++++++++++++++++++ 3 files changed, 295 insertions(+) create mode 100644 rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 07170585380..393a5af5ef3 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -17,6 +17,7 @@ mod handshake; mod vendored_cache; pub mod builder; +pub mod partition; // Public re-exports — make pool::Builder, pool::ConnectionEventListener, etc. // directly accessible without going through pool::builder. @@ -25,6 +26,7 @@ pub use builder::{ ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, }; +pub use partition::{CrossPartitionPolicy, DriverSpawner, PartitionId}; /// Connection-caching pool layer. /// diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index b627f14ed05..e4f643b95a0 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -57,6 +57,7 @@ pub struct Builder { max_connections_per_host: Option, proxy_config: Option, connection_event_listener: Option>, + cross_partition_policy: super::partition::CrossPartitionPolicy, tls: Tls, } @@ -72,6 +73,7 @@ impl std::fmt::Debug for Builder { "connection_event_listener", &self.connection_event_listener.as_ref().map(|_| ".."), ) + .field("cross_partition_policy", &self.cross_partition_policy) .finish() } } @@ -85,6 +87,7 @@ impl Default for Builder { max_connections_per_host: None, proxy_config: None, connection_event_listener: None, + cross_partition_policy: super::partition::CrossPartitionPolicy::default(), tls: TlsUnset {}, } } @@ -170,6 +173,29 @@ impl Builder { self.connection_event_listener = listener; self } + + /// Set the policy that governs checkout when the local partition has + /// no idle connection and the pool is at capacity. + /// + /// Defaults to [`CrossPartitionPolicy::Never`](super::partition::CrossPartitionPolicy::Never). + /// Has no observable effect with a single partition or when the pool + /// stays under capacity. + pub fn cross_partition_policy( + mut self, + policy: super::partition::CrossPartitionPolicy, + ) -> Self { + self.cross_partition_policy = policy; + self + } + + /// Mutable-reference variant of [`cross_partition_policy`](Self::cross_partition_policy). + pub fn set_cross_partition_policy( + &mut self, + policy: super::partition::CrossPartitionPolicy, + ) -> &mut Self { + self.cross_partition_policy = policy; + self + } } impl Builder { @@ -187,6 +213,7 @@ impl Builder { max_connections_per_host: self.max_connections_per_host, proxy_config: self.proxy_config, connection_event_listener: self.connection_event_listener, + cross_partition_policy: self.cross_partition_policy, tls: TlsProviderSelected { provider, context: TlsContext::default(), diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs new file mode 100644 index 00000000000..31d89237760 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -0,0 +1,266 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Connection pool partitioning. +//! +//! A partition is a label that groups connections by locality: the runtime +//! that owns their drivers and the network interface their sockets bind to. +//! Partition labels are assigned through client configuration; the pool +//! indexes connections by label and respects locality at checkout. +//! +//! # Topologies +//! +//! ## Single partition +//! +//! No partitioning. All connections pool together. The runtime is whatever +//! was current at client construction; no NIC binding. +//! +//! ```text +//! Pool +//! └── Partition (anonymous, runtime=tokio-mt, nic=none) +//! ├── conn-1 +//! └── conn-2 +//! ``` +//! +//! ## Per-runtime, no NIC binding +//! +//! N partitions, one per runtime. Each partition's connections have +//! drivers on that partition's runtime. +//! +//! ```text +//! Pool +//! ├── Partition 0 (runtime=tokio-current R0, nic=none) +//! │ ├── conn-1 (driver on R0) +//! │ └── conn-2 (driver on R0) +//! └── Partition 1 (runtime=tokio-current R1, nic=none) +//! └── conn-3 (driver on R1) +//! ``` +//! +//! ## Per-runtime, per-NIC +//! +//! Partitions cluster by NIC. A socket bound to one NIC cannot serve +//! traffic on another. +//! +//! ```text +//! Pool +//! ├── Partition 0 (runtime=R0, nic=eth0) ─┐ +//! ├── Partition 1 (runtime=R1, nic=eth0) ─┴─ same NIC group +//! ├── Partition 2 (runtime=R2, nic=eth1) ─┐ +//! └── Partition 3 (runtime=R3, nic=eth1) ─┴─ same NIC group +//! ``` +//! +//! ## NUMA-aware +//! +//! Partitions align with NUMA topology: runtimes pinned to cores on a +//! node, NIC selected to match the node. The pool does not detect NUMA +//! topology; it sees only `(PartitionId, runtime, nic)` and the +//! alignment is established when clients are configured. +//! +//! ```text +//! Pool +//! ├── NUMA node 0 +//! │ ├── Partition 0 (runtime=R0 on core 0, nic=eth0) +//! │ ├── Partition 1 (runtime=R1 on core 1, nic=eth0) +//! │ └── Partition 2 (runtime=R2 on core 2, nic=eth1) +//! └── NUMA node 1 +//! ├── Partition 3 (runtime=R3 on core 32, nic=eth2) +//! └── Partition 4 (runtime=R4 on core 33, nic=eth3) +//! ``` +//! +//! # Boundaries +//! +//! - **NIC (hard):** a connection's NIC binding is fixed at creation. +//! Connections form NIC groups: one group per `Some(nic)` value plus +//! an unbound group for `None`. The pool only returns a connection +//! to a checkout in the same NIC group; the unbound group is not a +//! wildcard. +//! - **Runtime (soft):** a connection's driver runs on a specific +//! runtime. Cross-runtime checkout is feasible (the request is +//! dispatched through the driver's runtime via a channel) but costs +//! a cross-thread send. [`CrossPartitionPolicy`] controls whether +//! the pool crosses this boundary at checkout. +//! +//! # Checkout +//! +//! When a request arrives on partition P for authority A: +//! +//! 1. **Local hit:** an idle connection in P for A is reused. +//! 2. **Local miss, under capacity:** a new connection is created +//! in P. +//! 3. **Local miss, at capacity:** behavior depends on +//! [`CrossPartitionPolicy`]. +//! +//! Cross-partition borrowing is a capacity-pressure fallback. Under +//! normal load each partition creates its own connections. +//! +//! ## Policy: `Never` +//! +//! At capacity with no local idle, the request waits for a permit. +//! Peer partitions are not consulted. +//! +//! ```text +//! Capacity = 2, both in use, request arrives on P0 for A: +//! +//! P0 (eth0): [active to B] ← request for A queues here +//! P1 (eth0): [idle to A] ← not consulted +//! +//! Outcome: request blocks until a permit is available, then either +//! reuses a returning idle for A or creates a new connection on P0. +//! ``` +//! +//! ## Policy: `PreferLocal` +//! +//! At capacity with no local idle, the pool checks peer partitions in +//! the same NIC group for an idle connection to the requested +//! authority. If one is found, the request borrows it. Otherwise it +//! waits for a permit. +//! +//! ```text +//! Capacity = 2, both in use, request arrives on P0 for A: +//! +//! P0 (eth0): [active to B] ← request for A +//! P1 (eth0): [idle to A] ← borrowed +//! P2 (eth1): [idle to A] ← different NIC, never consulted +//! +//! Outcome: P1's idle connection serves the request. The connection's +//! driver stays on P1's runtime; the request flows through P1's +//! runtime via the connection's channel. +//! ``` + +/// Identifier for a pool partition. +/// +/// A partition's identity is opaque to the pool. The identifier is +/// assigned at client construction; the pool validates that the same +/// identifier is associated with the same driver spawner and network +/// interface binding across all clients that reference it. +/// +/// The default identifier denotes an anonymous partition used when +/// partitioning is not configured. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct PartitionId(u64); + +impl PartitionId { + const ANONYMOUS: u64 = u64::MAX; + + /// Identifier from a numeric index. + pub const fn from_index(index: usize) -> Self { + Self(index as u64) + } + + /// Identifier from a raw value. The value `u64::MAX` is reserved + /// for the anonymous default partition. + pub const fn from_raw(id: u64) -> Self { + Self(id) + } + + /// Raw value of this identifier. + pub const fn as_u64(self) -> u64 { + self.0 + } +} + +impl Default for PartitionId { + fn default() -> Self { + Self(Self::ANONYMOUS) + } +} + +/// Spawner for connection driver tasks. +/// +/// A driver is the task that owns a connection's I/O. The spawner +/// determines which runtime executes that task. Per-runtime spawners +/// allow each connection's I/O to remain on a specific runtime. +/// +/// Construct via [`DriverSpawner::current_tokio`] (captures the +/// current tokio runtime handle) or [`DriverSpawner::from_tokio_handle`] +/// (uses a specific runtime handle). +#[derive(Clone)] +pub struct DriverSpawner { + inner: SpawnerInner, +} + +#[derive(Clone)] +enum SpawnerInner { + TokioHandle(tokio::runtime::Handle), +} + +impl DriverSpawner { + /// Spawner that uses the current tokio runtime. + /// + /// Captures the runtime handle eagerly at the call site. Panics + /// if invoked outside a tokio runtime context. + pub fn current_tokio() -> Self { + Self::from_tokio_handle(tokio::runtime::Handle::current()) + } + + /// Spawner targeting the given tokio runtime handle. + pub fn from_tokio_handle(handle: tokio::runtime::Handle) -> Self { + Self { + inner: SpawnerInner::TokioHandle(handle), + } + } +} + +impl std::fmt::Debug for DriverSpawner { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let variant = match &self.inner { + SpawnerInner::TokioHandle(_) => "TokioHandle", + }; + f.debug_tuple("DriverSpawner").field(&variant).finish() + } +} + +/// Policy governing checkout when the local partition has no idle +/// connection and the pool is at capacity. +/// +/// Cross-partition borrowing applies within a NIC group only. +/// Connections bound to different NICs are never shared regardless +/// of policy. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum CrossPartitionPolicy { + /// At capacity with no local idle, the request waits for a permit. + /// Peer partitions are not consulted. + #[default] + Never, + /// At capacity with no local idle, the request borrows an idle + /// connection from a peer partition in the same NIC group when one + /// is available, otherwise waits for a permit. + PreferLocal, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn partition_id_from_index() { + assert_eq!(PartitionId::from_index(0).as_u64(), 0); + assert_eq!(PartitionId::from_index(5).as_u64(), 5); + assert_eq!(PartitionId::from_index(0), PartitionId::from_index(0)); + assert_ne!(PartitionId::from_index(0), PartitionId::from_index(1)); + } + + #[test] + fn partition_id_default_is_anonymous() { + assert_eq!(PartitionId::default().as_u64(), u64::MAX); + assert_ne!(PartitionId::default(), PartitionId::from_index(0)); + } + + #[test] + fn cross_partition_policy_default_is_never() { + assert_eq!(CrossPartitionPolicy::default(), CrossPartitionPolicy::Never); + } + + #[tokio::test] + async fn driver_spawner_current_tokio() { + let _ = DriverSpawner::current_tokio(); + } + + #[tokio::test] + async fn driver_spawner_from_handle() { + let h = tokio::runtime::Handle::current(); + let _ = DriverSpawner::from_tokio_handle(h); + } +} From 7469abc3bacfc63cbfd58137e88001cd612a51ce Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 21 May 2026 11:00:35 -0400 Subject: [PATCH 25/47] move v2 into pool client --- .../aws-smithy-http-client/src/client.rs | 10 - .../aws-smithy-http-client/src/client/pool.rs | 71 ++-- .../src/client/pool/builder.rs | 362 +++++------------- .../src/client/pool/client.rs | 260 +++++++++++++ .../src/client/pool/connection.rs | 15 +- .../src/client/pool/partition.rs | 5 +- .../aws-smithy-http-client/src/lib.rs | 2 +- .../tests/h2_pool_test.rs | 16 +- .../tests/pool_behavior_test.rs | 19 +- .../tests/proxy_tests.rs | 19 +- .../tests/smoke_test_clients.rs | 26 +- .../aws-smithy-http-client/tests/tls.rs | 34 +- .../tests/reconnect_on_transient_error.rs | 7 +- 13 files changed, 483 insertions(+), 363 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/src/client/pool/client.rs diff --git a/rust-runtime/aws-smithy-http-client/src/client.rs b/rust-runtime/aws-smithy-http-client/src/client.rs index 35e3b8ed1e8..42ee02192c1 100644 --- a/rust-runtime/aws-smithy-http-client/src/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client.rs @@ -924,16 +924,6 @@ impl Builder { Self::default() } - /// Creates a new builder backed by the v2 composable connection pool. - /// - /// The v2 client uses hyper-util's composable pool APIs for connection - /// management, supporting per-host connection reuse with HTTP/1.1 and - /// HTTP/2 protocol negotiation. - #[doc(hidden)] - pub fn new_v2() -> pool::Builder { - pool::Builder::new() - } - /// Returns a [`SharedHttpClient`] that calls the given `connector` function to select an HTTP(S) connector. #[doc(hidden)] pub fn build_with_connector_fn(self, connector_fn: F) -> SharedHttpClient diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 393a5af5ef3..1b2961947d3 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -17,6 +17,7 @@ mod handshake; mod vendored_cache; pub mod builder; +pub mod client; pub mod partition; // Public re-exports — make pool::Builder, pool::ConnectionEventListener, etc. @@ -26,14 +27,10 @@ pub use builder::{ ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, }; +pub use client::{Client, ClientBuilder}; pub use partition::{CrossPartitionPolicy, DriverSpawner, PartitionId}; /// Connection-caching pool layer. -/// -/// Re-exports [`vendored_cache`], our vendored copy of hyper-util's -/// `pool::cache` with SDK-specific modifications. The re-export insulates -/// callers from the vendoring detail; all pool code uses `cache::…` -/// regardless of where the implementation lives. mod cache { pub(crate) use super::vendored_cache::*; } @@ -108,14 +105,8 @@ impl ConnectionMetadataCapture { /// Pool-level configuration. /// -/// Plumbed from the eventual `Builder::new_v2()` public API through -/// `build_pool` to `ConnectionPool`. Internal type: the public surface -/// is the builder's per-knob setters, not this struct. -/// -/// All fields are `None` by default (defaults applied by the pool where -/// they take effect, not here). Fields are `pub(crate)` because -/// consumers of this struct are all within the pool module; accessor -/// methods would add noise without benefit. +/// Defaults are applied at the point each setting takes effect rather +/// than in this struct. #[derive(Clone, Default)] pub(crate) struct PoolConfig { /// Upper bound on concurrent connections (total, across all hosts). @@ -136,6 +127,39 @@ pub(crate) struct PoolConfig { pub(crate) connection_event_listener: Option>, } +/// The connection pool's configuration surface. +/// +/// Owns the connection lifecycle (creation, caching, eviction, health +/// checking) and proxy routing decisions. Multiple [`Client`] instances +/// can reference one `SharedPool`, each presenting a different +/// per-partition view of the same underlying connections. +/// +/// Construct via [`SharedPool::builder`], which returns a +/// [`Builder`](crate::client::pool::builder::Builder). +/// Cloning is cheap (shared via `Arc`). +#[derive(Clone, Debug)] +pub struct SharedPool { + pub(crate) inner: Arc, +} + +pub(crate) struct SharedPoolInner { + pub(crate) pool: Arc, + pub(crate) proxy_matcher: Option>, +} + +impl std::fmt::Debug for SharedPoolInner { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SharedPoolInner").finish_non_exhaustive() + } +} + +impl SharedPool { + /// Create a [`Builder`] for configuring a new connection pool. + pub fn builder() -> builder::Builder { + builder::Builder::default() + } +} + /// Key for per-host connection pool routing. #[derive(Clone, Debug, PartialEq, Eq, Hash)] struct PoolKey { @@ -443,9 +467,8 @@ pub(crate) struct ConnectionPool { drop_notifier: OnceLock>, } -/// Minimum eviction tick period, matching legacy hyper-util pool's -/// `MIN_CHECK` constant. Prevents very short `pool_idle_timeout`s from -/// causing the task to spin hot. +/// Minimum eviction tick period. Prevents very short `pool_idle_timeout`s +/// from causing the eviction task to spin hot. const MIN_EVICTION_TICK: Duration = Duration::from_millis(90); impl ConnectionPool { @@ -493,9 +516,8 @@ impl ConnectionPool { /// `pool_idle_timeout` is configured. /// /// Idempotent and cheap after the first call (single relaxed - /// `AtomicBool` load returns early). Matches legacy - /// `spawn_idle_interval` semantics: no-op without a timeout, no-op - /// with a zero timeout, no-op if already spawned. + /// `AtomicBool` load returns early). No-op without a timeout, with + /// a zero timeout, or if already spawned. fn maybe_spawn_eviction_task(self: &Arc) { let timeout = match self.config.pool_idle_timeout { Some(d) if d > Duration::ZERO => d, @@ -614,10 +636,8 @@ async fn eviction_task( /// - proxied non-CONNECT → absolute-form (full URL with scheme + authority) /// - direct non-CONNECT → origin-form (path + query only) /// -/// Mirrors `hyper-util`'s legacy `client::Client::send_request` request -/// preparation. Required because `hyper::client::conn::http1` sends the -/// URI as-is; the form selection lives in the layer above the -/// connection. +/// The HTTP/1 connection sends request URIs as-is, so request-target +/// form selection is the caller's responsibility. fn rewrite_h1_request_target(req: &mut http_1x::Request, is_proxied: bool) { use http_1x::{uri::Parts, Method, Uri}; if req.method() == Method::CONNECT { @@ -963,8 +983,6 @@ mod tests { } /// Without a `pool_idle_timeout`, `maybe_spawn_eviction_task` is a no-op. - /// Matches legacy `spawn_idle_interval` which also skips without a - /// timeout. #[tokio::test] async fn eviction_task_no_spawn_without_timeout() { let pool = pool_with_config(PoolConfig { @@ -976,8 +994,7 @@ mod tests { assert!(pool.drop_notifier.get().is_none()); } - /// A zero `pool_idle_timeout` is treated the same as `None`. Matches - /// legacy `spawn_idle_interval` (returns early on `Duration::ZERO`). + /// A zero `pool_idle_timeout` is treated the same as `None`. #[tokio::test] async fn eviction_task_no_spawn_with_zero_timeout() { let pool = pool_with_config(PoolConfig { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index e4f643b95a0..76eee9166ab 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -3,33 +3,21 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! V2 HTTP client builder. +//! Connection pool builder. //! -//! Entry point: [`Builder::new`] or [`crate::Builder::new_v2`]. +//! Entry point: [`SharedPool::builder`](super::SharedPool::builder). pub use super::connection::{ Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, }; -use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; -use aws_smithy_runtime_api::client::connection::CaptureSmithyConnection; -use aws_smithy_runtime_api::client::connector_metadata::ConnectorMetadata; -use aws_smithy_runtime_api::client::http::{ - HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpClient, - SharedHttpConnector, -}; -use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse}; -use aws_smithy_runtime_api::client::result::ConnectorError; -use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; -use aws_smithy_runtime_api::shared::IntoShared; use hyper_util::client::legacy::connect::HttpConnector as HyperHttpConnector; use super::{ConnectionPool, PoolConfig}; -use crate::client::downcast_error; use crate::client::tls; use crate::client::{TlsProviderSelected, TlsUnset}; use crate::tls::TlsContext; @@ -42,12 +30,13 @@ use crate::tls::TlsContext; /// # Examples /// /// ```no_run -/// use aws_smithy_http_client::pool::Builder; +/// use aws_smithy_http_client::pool::{SharedPool, Client}; /// use std::time::Duration; /// -/// let client = Builder::new() +/// let pool = SharedPool::builder() /// .pool_idle_timeout(Duration::from_secs(20)) /// .build_http(); +/// let client = Client::new(&pool); /// ``` #[derive(Clone)] pub struct Builder { @@ -58,6 +47,7 @@ pub struct Builder { proxy_config: Option, connection_event_listener: Option>, cross_partition_policy: super::partition::CrossPartitionPolicy, + dns_resolver: Option, tls: Tls, } @@ -74,6 +64,7 @@ impl std::fmt::Debug for Builder { &self.connection_event_listener.as_ref().map(|_| ".."), ) .field("cross_partition_policy", &self.cross_partition_policy) + .field("dns_resolver", &self.dns_resolver.as_ref().map(|_| "..")) .finish() } } @@ -88,6 +79,7 @@ impl Default for Builder { proxy_config: None, connection_event_listener: None, cross_partition_policy: super::partition::CrossPartitionPolicy::default(), + dns_resolver: None, tls: TlsUnset {}, } } @@ -98,8 +90,7 @@ impl Builder { /// Set the pool idle timeout. /// /// Connections idle longer than this duration are evicted from the pool. - /// Set below the server's idle timeout to avoid stale connection errors - /// (e.g., ~20s for S3). + /// Set below the server's idle timeout to avoid stale connection errors. pub fn pool_idle_timeout(mut self, timeout: Duration) -> Self { self.pool_idle_timeout = Some(timeout); self @@ -111,12 +102,16 @@ impl Builder { self } - /// Set the maximum number of concurrent connections. + /// Set the global maximum number of concurrent connections. + /// + /// Caps the total live connections in the pool across all hosts, + /// including idle (cached) connections. New connection attempts wait + /// when the pool is at capacity; existing connections must be evicted + /// or closed before another can be created. /// - /// Limits concurrent connection handshakes across all hosts. Cached - /// (reused) connections do not count against this limit. When the limit - /// is reached, new connection attempts wait until an existing handshake - /// completes. + /// Should be at least [`max_connections_per_host`](Self::max_connections_per_host) + /// when both are set; the global limit applies on top of the per-host + /// limit. pub fn max_connections(mut self, n: usize) -> Self { self.max_connections = Some(n); self @@ -196,14 +191,31 @@ impl Builder { self.cross_partition_policy = policy; self } -} -impl Builder { - /// Create a new v2 client builder. - pub fn new() -> Self { - Self::default() + /// Set a custom DNS resolver. + /// + /// Connections established by this pool resolve hostnames through the + /// given resolver. Defaults to the system resolver when unset. + pub fn dns_resolver( + mut self, + resolver: impl aws_smithy_runtime_api::client::dns::ResolveDns + 'static, + ) -> Self { + use aws_smithy_runtime_api::shared::IntoShared; + self.dns_resolver = Some(resolver.into_shared()); + self } + /// Mutable-reference variant of [`dns_resolver`](Self::dns_resolver). + pub fn set_dns_resolver( + &mut self, + resolver: Option, + ) -> &mut Self { + self.dns_resolver = resolver; + self + } +} + +impl Builder { /// Set the TLS implementation. pub fn tls_provider(self, provider: tls::Provider) -> Builder { Builder { @@ -214,6 +226,7 @@ impl Builder { proxy_config: self.proxy_config, connection_event_listener: self.connection_event_listener, cross_partition_policy: self.cross_partition_policy, + dns_resolver: self.dns_resolver, tls: TlsProviderSelected { provider, context: TlsContext::default(), @@ -223,9 +236,8 @@ impl Builder { /// Build an HTTP client without TLS. #[doc(hidden)] - pub fn build_http(self) -> SharedHttpClient { - let mut tcp = HyperHttpConnector::new(); - tcp.set_nodelay(self.tcp_nodelay); + pub fn build_http(mut self) -> super::SharedPool { + let dns_resolver = self.dns_resolver.take(); let config = PoolConfig { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, @@ -233,33 +245,25 @@ impl Builder { connection_event_listener: self.connection_event_listener.clone(), }; let proxy_matcher = proxy_matcher_from(&self.proxy_config); - let pool = build_http_pool_with_proxy(tcp, &self.proxy_config, config); - V2HttpClient::new(pool, proxy_matcher).into_shared() - } - - /// Build an HTTP client without TLS, using a custom DNS resolver. - /// - /// Mirrors `build_http` but threads a custom `ResolveDns`. Intended - /// for tests that use `WireMockServer::dns_resolver()` to redirect - /// all lookups to a loopback mock. Public API for non-TLS resolver - /// plumbing is `#[doc(hidden)]`; TLS is the common case. - #[doc(hidden)] - pub fn build_http_with_resolver( - self, - resolver: impl aws_smithy_runtime_api::client::dns::ResolveDns + Clone + 'static, - ) -> SharedHttpClient { - use crate::client::dns::HyperUtilResolver; - let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); - tcp.set_nodelay(self.tcp_nodelay); - let config = PoolConfig { - max_connections: self.max_connections, - max_connections_per_host: self.max_connections_per_host, - pool_idle_timeout: self.pool_idle_timeout, - connection_event_listener: self.connection_event_listener.clone(), + let pool = match dns_resolver { + Some(resolver) => { + use crate::client::dns::HyperUtilResolver; + let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); + tcp.set_nodelay(self.tcp_nodelay); + build_http_pool_with_proxy(tcp, &self.proxy_config, config) + } + None => { + let mut tcp = HyperHttpConnector::new(); + tcp.set_nodelay(self.tcp_nodelay); + build_http_pool_with_proxy(tcp, &self.proxy_config, config) + } }; - let proxy_matcher = proxy_matcher_from(&self.proxy_config); - let pool = build_http_pool_with_proxy(tcp, &self.proxy_config, config); - V2HttpClient::new(pool, proxy_matcher).into_shared() + super::SharedPool { + inner: Arc::new(super::SharedPoolInner { + pool: Arc::new(pool), + proxy_matcher, + }), + } } /// Build an HTTP client from a raw TCP-level connector. @@ -270,7 +274,7 @@ impl Builder { /// to route connections to the H2 path. #[cfg(all(feature = "test-util", aws_sdk_unstable))] #[doc(hidden)] - pub fn build_http_with_tcp_connector(self, connector: C) -> SharedHttpClient + pub fn build_http_with_tcp_connector(self, connector: C) -> super::SharedPool where C: tower::Service + Clone + Send + Sync + 'static, C::Error: Into + 'static, @@ -289,7 +293,12 @@ impl Builder { connection_event_listener: self.connection_event_listener.clone(), }; let pool = super::build_pool(connector, config); - V2HttpClient::new(pool, None).into_shared() + super::SharedPool { + inner: Arc::new(super::SharedPoolInner { + pool: Arc::new(pool), + proxy_matcher: None, + }), + } } } @@ -339,26 +348,26 @@ impl Builder { } /// Build an HTTPS client with the selected TLS provider. - pub fn build_https(self) -> SharedHttpClient { - let mut tcp = HyperHttpConnector::new(); - tcp.set_nodelay(self.tcp_nodelay); - tcp.enforce_http(false); - self.build_from_tcp(tcp) - } - - /// Build an HTTPS client using a custom DNS resolver. - pub fn build_with_resolver( - self, - resolver: impl aws_smithy_runtime_api::client::dns::ResolveDns + Clone + 'static, - ) -> SharedHttpClient { - use crate::client::dns::HyperUtilResolver; - let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); - tcp.set_nodelay(self.tcp_nodelay); - tcp.enforce_http(false); - self.build_from_tcp(tcp) + pub fn build_https(mut self) -> super::SharedPool { + let dns_resolver = self.dns_resolver.take(); + match dns_resolver { + Some(resolver) => { + use crate::client::dns::HyperUtilResolver; + let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); + tcp.set_nodelay(self.tcp_nodelay); + tcp.enforce_http(false); + self.build_from_tcp(tcp) + } + None => { + let mut tcp = HyperHttpConnector::new(); + tcp.set_nodelay(self.tcp_nodelay); + tcp.enforce_http(false); + self.build_from_tcp(tcp) + } + } } - fn build_from_tcp(self, tcp: HyperHttpConnector) -> SharedHttpClient + fn build_from_tcp(self, tcp: HyperHttpConnector) -> super::SharedPool where R: Clone + Send + Sync + 'static, R: tower::Service, @@ -395,7 +404,12 @@ impl Builder { proxy_config, ); let pool = super::build_pool(connector, config); - V2HttpClient::new(pool, proxy_matcher).into_shared() + super::SharedPool { + inner: Arc::new(super::SharedPoolInner { + pool: Arc::new(pool), + proxy_matcher, + }), + } } #[cfg(feature = "s2n-tls")] tls::Provider::S2nTls => { @@ -405,7 +419,12 @@ impl Builder { proxy_config, ); let pool = super::build_pool(connector, config); - V2HttpClient::new(pool, proxy_matcher).into_shared() + super::SharedPool { + inner: Arc::new(super::SharedPoolInner { + pool: Arc::new(pool), + proxy_matcher, + }), + } } // Provider is #[non_exhaustive]; this arm is unreachable when any // TLS feature is enabled (which is required to construct a Provider). @@ -415,130 +434,11 @@ impl Builder { } } -// --------------------------------------------------------------------------- -// HttpConnector adapter -// --------------------------------------------------------------------------- - -/// Smithy `HttpConnector` backed by the v2 connection pool. -/// -/// Constructed fresh per `HttpClient::http_connector` call so it can capture -/// the per-operation `HttpConnectorSettings` (connect/read timeouts). The pool -/// itself is shared across all operations via `Arc`. -struct PooledConnector { - pool: Arc, - connect_timeout: Option, - read_timeout: Option, - /// Present when any timeout is configured. Required to construct - /// `Timeout` futures. `V2HttpClient::http_connector` panics if - /// a timeout is set but `sleep_impl` is absent (matches v1). - sleep_impl: Option, - /// Proxy URL matcher, cloned from `V2HttpClient`. Drives the - /// `Proxy-Authorization` header injection in `call`. - proxy_matcher: Option>, -} - -impl std::fmt::Debug for PooledConnector { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("PooledConnector").finish() - } -} - -impl HttpConnector for PooledConnector { - fn call(&self, request: HttpRequest) -> HttpConnectorFuture { - let pool = self.pool.clone(); - let connect_timeout = self.connect_timeout; - let read_timeout = self.read_timeout; - let sleep_impl = self.sleep_impl.clone(); - let proxy_matcher = self.proxy_matcher.clone(); - HttpConnectorFuture::new(async move { - let mut request = request - .try_into_http1x() - .map_err(|err| ConnectorError::user(err.into()))?; - - // Save the full URI for pool routing before rewriting to origin-form. - let full_uri = request.uri().clone(); - - // Inject Proxy-Authorization for HTTP-through-proxy requests. - // Shared with v1 via [`super::proxy::add_proxy_auth_header`]; - // HTTPS-through-proxy uses CONNECT tunneling and authenticates - // inside the tunnel-establishment exchange, not via this header. - if let Some(matcher) = proxy_matcher.as_ref() { - crate::client::proxy::add_proxy_auth_header(&mut request, matcher); - } - - // Create a request-scoped capture the pool will fill in with - // the `ConnectionMetadata` for the selected connection, and - // wire `CaptureSmithyConnection` to read from it. When the - // runtime's connection-poisoning interceptor later calls - // `ConnectionMetadata::poison()`, it flips the `PoisonPill` on - // the actual `ManagedConnection` selected for this request. - if let Some(capture_smithy) = request.extensions().get::() { - let capture = super::ConnectionMetadataCapture::new(); - let for_retriever = capture.clone(); - capture_smithy.set_connection_retriever(move || for_retriever.get()); - request.extensions_mut().insert(capture); - } - - // Attach the per-op read timeout hint (if any) as a request - // extension. `H{1,2}Checkout::call` reads it and wraps the - // request dispatch: bounds request-write + response-headers, - // nothing else. - if let Some((duration, sleep)) = read_timeout.zip(sleep_impl.clone()) { - request.extensions_mut().insert(super::ReadTimeoutHint( - super::TimeoutContext::new(duration, sleep), - )); - } - - // Set Host header from the URI authority if not already present. - if !request.headers().contains_key(http_1x::header::HOST) { - if let Some(authority) = full_uri.authority() { - let host = match get_non_default_port(&full_uri) { - Some(port) => format!("{}:{}", authority.host(), port), - None => authority.host().to_string(), - }; - request.headers_mut().insert( - http_1x::header::HOST, - http_1x::HeaderValue::from_str(&host) - .expect("authority is valid header value"), - ); - } - } - - // URI request-target rewriting happens in `H1Checkout::call` - // because the correct form depends on whether the connection - // landed on a proxy (absolute-form) or direct origin - // (origin-form). H2 uses pseudo-headers derived by hyper from - // the absolute URI, so it doesn't need rewriting here. - - // Build the connect context: routing URI + connect timeout. - // Cache hits skip the connector entirely so connect_timeout is - // automatically new-connection-only. - let connect_ctx = super::ConnectCtx::new( - full_uri, - connect_timeout - .zip(sleep_impl) - .map(|(d, s)| super::TimeoutContext::new(d, s)), - ); - - let response = pool - .send_request(connect_ctx, request) - .await - .map_err(downcast_error)?; - - HttpResponse::try_from(response).map_err(|err| ConnectorError::other(err.into(), None)) - }) - } -} - -// --------------------------------------------------------------------------- -// HttpClient impl -// --------------------------------------------------------------------------- - /// Build the proxy URL matcher from a `ProxyConfig`, returning `None` when /// no proxy is configured. Wrapped in `Arc` so each `PooledConnector` /// constructed by `http_connector` can share the (immutable) matcher /// without per-call allocation. -fn proxy_matcher_from( +pub(super) fn proxy_matcher_from( proxy_config: &Option, ) -> Option> { proxy_config @@ -547,72 +447,10 @@ fn proxy_matcher_from( } /// Returns the port if it is not the default for the scheme. -fn get_non_default_port(uri: &http_1x::Uri) -> Option> { +pub(super) fn get_non_default_port(uri: &http_1x::Uri) -> Option> { match (uri.port().map(|p| p.as_u16()), uri.scheme()) { (Some(443), Some(s)) if *s == http_1x::uri::Scheme::HTTPS => None, (Some(80), Some(s)) if *s == http_1x::uri::Scheme::HTTP => None, _ => uri.port(), } } - -/// Smithy `HttpClient` backed by the v2 connection pool. -struct V2HttpClient { - pool: Arc, - /// Proxy URL matcher, present iff the builder was configured with a - /// `proxy_config`. Used by `PooledConnector` to inject - /// `Proxy-Authorization` headers on HTTP-through-proxy requests - /// (HTTPS-CONNECT auth is handled by the proxy connector itself - /// during tunnel setup). - proxy_matcher: Option>, -} - -impl std::fmt::Debug for V2HttpClient { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("V2HttpClient").finish() - } -} - -impl V2HttpClient { - fn new( - pool: ConnectionPool, - proxy_matcher: Option>, - ) -> Self { - Self { - pool: Arc::new(pool), - proxy_matcher, - } - } -} - -impl HttpClient for V2HttpClient { - fn http_connector( - &self, - settings: &HttpConnectorSettings, - components: &RuntimeComponents, - ) -> SharedHttpConnector { - let connect_timeout = settings.connect_timeout(); - let read_timeout = settings.read_timeout(); - let sleep_impl = components.sleep_impl(); - - if (connect_timeout.is_some() || read_timeout.is_some()) && sleep_impl.is_none() { - panic!( - "an async sleep impl is required to use connect/read timeouts with \ - the v2 HTTP client; provide one via `RuntimeComponents::sleep_impl`" - ); - } - - SharedHttpConnector::new(PooledConnector { - pool: self.pool.clone(), - connect_timeout, - read_timeout, - sleep_impl, - proxy_matcher: self.proxy_matcher.clone(), - }) - } - - fn connector_metadata(&self) -> Option { - // TODO: consider distinguishing v1 vs v2 clients and including TLS - // provider info (e.g. "rustls"/"0.23") in the connector metadata. - Some(ConnectorMetadata::new("hyper", Some(Cow::Borrowed("1.x")))) - } -} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs new file mode 100644 index 00000000000..e0f647430f0 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs @@ -0,0 +1,260 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Per-partition client and its builder. + +use std::borrow::Cow; +use std::sync::Arc; +use std::time::Duration; + +use aws_smithy_runtime_api::client::connection::CaptureSmithyConnection; +use aws_smithy_runtime_api::client::connector_metadata::ConnectorMetadata; +use aws_smithy_runtime_api::client::http::{ + HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpConnector, +}; +use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse}; +use aws_smithy_runtime_api::client::result::ConnectorError; +use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; + +use super::partition::{DriverSpawner, PartitionId}; +use super::{ConnectionPool, SharedPool}; +use crate::client::downcast_error; + +/// Per-partition view of a [`SharedPool`]. +/// +/// Implements [`HttpClient`] by routing requests through the shared +/// connection pool. Multiple `Client` instances can reference the same +/// pool, each carrying a distinct [`PartitionId`], network interface +/// binding, and [`DriverSpawner`]. +/// +/// Construct via [`Client::new`] (default partition) or +/// [`Client::builder`] for per-partition configuration. +/// +/// Cloning is cheap: all fields are either `Arc`-backed or small values. +#[derive(Clone)] +pub struct Client { + pool: SharedPool, + partition: PartitionId, + interface: Option, + driver_spawner: DriverSpawner, +} + +impl std::fmt::Debug for Client { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Client") + .field("partition", &self.partition) + .field("interface", &self.interface) + .field("driver_spawner", &self.driver_spawner) + .finish_non_exhaustive() + } +} + +impl Client { + /// Construct a `Client` using the default partition. + /// + /// Equivalent to `Client::builder(pool).build()`. + pub fn new(pool: &SharedPool) -> Self { + Self::builder(pool).build() + } + + /// Return a [`ClientBuilder`] for advanced per-partition configuration. + pub fn builder(pool: &SharedPool) -> ClientBuilder { + ClientBuilder { + pool: pool.clone(), + partition: PartitionId::default(), + interface: None, + driver_spawner: None, + } + } +} + +impl HttpClient for Client { + fn http_connector( + &self, + settings: &HttpConnectorSettings, + components: &RuntimeComponents, + ) -> SharedHttpConnector { + let connect_timeout = settings.connect_timeout(); + let read_timeout = settings.read_timeout(); + let sleep_impl = components.sleep_impl(); + + if (connect_timeout.is_some() || read_timeout.is_some()) && sleep_impl.is_none() { + panic!( + "an async sleep impl is required to use connect/read timeouts with \ + the v2 HTTP client; provide one via `RuntimeComponents::sleep_impl`" + ); + } + + SharedHttpConnector::new(PooledConnector { + pool: self.pool.inner.pool.clone(), + connect_timeout, + read_timeout, + sleep_impl, + proxy_matcher: self.pool.inner.proxy_matcher.clone(), + }) + } + + fn connector_metadata(&self) -> Option { + Some(ConnectorMetadata::new("hyper", Some(Cow::Borrowed("1.x")))) + } +} + +/// Builder for configuring a [`Client`] with per-partition settings. +/// +/// Three knobs control partition behavior: +/// +/// - [`partition`](Self::partition) assigns a [`PartitionId`] (default: +/// anonymous). +/// - [`interface`](Self::interface) binds connections to a network +/// interface (default: none). +/// - [`driver_spawner`](Self::driver_spawner) selects the runtime for +/// connection driver tasks (default: captured from the current tokio +/// runtime at [`build`](Self::build) time). +pub struct ClientBuilder { + pool: SharedPool, + partition: PartitionId, + interface: Option, + driver_spawner: Option, +} + +impl ClientBuilder { + /// Set the partition identifier (consuming builder). + pub fn partition(mut self, id: PartitionId) -> Self { + self.partition = id; + self + } + + /// Set the partition identifier (mutable reference). + pub fn set_partition(&mut self, id: PartitionId) -> &mut Self { + self.partition = id; + self + } + + /// Bind connections to a network interface (consuming builder). + pub fn interface(mut self, iface: impl Into) -> Self { + self.interface = Some(iface.into()); + self + } + + /// Set or clear the network interface binding (mutable reference). + pub fn set_interface(&mut self, iface: Option) -> &mut Self { + self.interface = iface; + self + } + + /// Set the driver spawner (consuming builder). + pub fn driver_spawner(mut self, sp: DriverSpawner) -> Self { + self.driver_spawner = Some(sp); + self + } + + /// Set the driver spawner (mutable reference). + pub fn set_driver_spawner(&mut self, sp: DriverSpawner) -> &mut Self { + self.driver_spawner = Some(sp); + self + } + + /// Build the [`Client`]. + /// + /// Captures the current tokio runtime handle for the driver spawner + /// if one was not explicitly provided. + pub fn build(self) -> Client { + let driver_spawner = self + .driver_spawner + .unwrap_or_else(DriverSpawner::current_tokio); + + Client { + pool: self.pool, + partition: self.partition, + interface: self.interface, + driver_spawner, + } + } +} + +// --------------------------------------------------------------------------- +// PooledConnector (HttpConnector adapter) +// --------------------------------------------------------------------------- + +/// Smithy [`HttpConnector`] backed by the v2 connection pool. +/// +/// Constructed fresh per [`HttpClient::http_connector`] call so it can +/// capture the per-operation [`HttpConnectorSettings`] (connect/read +/// timeouts). The pool itself is shared across all operations via `Arc`. +struct PooledConnector { + pool: Arc, + connect_timeout: Option, + read_timeout: Option, + sleep_impl: Option, + proxy_matcher: Option>, +} + +impl std::fmt::Debug for PooledConnector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PooledConnector").finish() + } +} + +impl HttpConnector for PooledConnector { + fn call(&self, request: HttpRequest) -> HttpConnectorFuture { + let pool = self.pool.clone(); + let connect_timeout = self.connect_timeout; + let read_timeout = self.read_timeout; + let sleep_impl = self.sleep_impl.clone(); + let proxy_matcher = self.proxy_matcher.clone(); + HttpConnectorFuture::new(async move { + let mut request = request + .try_into_http1x() + .map_err(|err| ConnectorError::user(err.into()))?; + + let full_uri = request.uri().clone(); + + if let Some(matcher) = proxy_matcher.as_ref() { + crate::client::proxy::add_proxy_auth_header(&mut request, matcher); + } + + if let Some(capture_smithy) = request.extensions().get::() { + let capture = super::ConnectionMetadataCapture::new(); + let for_retriever = capture.clone(); + capture_smithy.set_connection_retriever(move || for_retriever.get()); + request.extensions_mut().insert(capture); + } + + if let Some((duration, sleep)) = read_timeout.zip(sleep_impl.clone()) { + request.extensions_mut().insert(super::ReadTimeoutHint( + super::TimeoutContext::new(duration, sleep), + )); + } + + if !request.headers().contains_key(http_1x::header::HOST) { + if let Some(authority) = full_uri.authority() { + let host = match super::builder::get_non_default_port(&full_uri) { + Some(port) => format!("{}:{}", authority.host(), port), + None => authority.host().to_string(), + }; + request.headers_mut().insert( + http_1x::header::HOST, + http_1x::HeaderValue::from_str(&host) + .expect("authority is valid header value"), + ); + } + } + + let connect_ctx = super::ConnectCtx::new( + full_uri, + connect_timeout + .zip(sleep_impl) + .map(|(d, s)| super::TimeoutContext::new(d, s)), + ); + + let response = pool + .send_request(connect_ctx, request) + .await + .map_err(downcast_error)?; + + HttpResponse::try_from(response).map_err(|err| ConnectorError::other(err.into(), None)) + }) + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 57ebb333945..75beab87e8d 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -420,9 +420,9 @@ pub(crate) struct ConnectionInfo { /// Shared via `Arc`; all clones observe and control the same flag. /// `ManagedConnection` holds one and hands out clones (via `metadata()`) /// that let the adapter layer mark the connection poisoned through -/// smithy's `ConnectionMetadata::poison_fn`. On next checkout / return, -/// the pool sees the flag set and drops the connection instead of reusing -/// it. Modeled after hyper-util's legacy `PoisonPill`, but owned by us. +/// smithy's `ConnectionMetadata::poison_fn`. On next checkout or return, +/// the pool sees the flag set and drops the connection instead of +/// reusing it. #[derive(Debug, Clone, Default)] pub(crate) struct PoisonPill { flag: Arc, @@ -986,9 +986,9 @@ pin_project! { /// /// The H2 variant carries a generic type parameter because /// `SingletonConnection`'s inner `T` is `hyper_util::client::pool:: - /// singleton::Singled<...>`, which is unnameable outside hyper-util. The - /// H1 variant is fully concrete since our vendored cache makes - /// `cache::Cached<...>` nameable. + /// singleton::Singled<...>`, which is unnameable outside hyper-util. + /// The H1 variant is fully concrete because `cache::Cached<...>` is + /// nameable. pub(crate) struct GuardedBody { #[pin] inner: hyper::body::Incoming, @@ -1046,8 +1046,7 @@ pub(crate) type CheckoutResponse = http_1x::Response { inner: IO, tcp_connect_duration: Duration, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 31d89237760..23e33da88fd 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -132,9 +132,8 @@ /// Identifier for a pool partition. /// /// A partition's identity is opaque to the pool. The identifier is -/// assigned at client construction; the pool validates that the same -/// identifier is associated with the same driver spawner and network -/// interface binding across all clients that reference it. +/// assigned at client construction and groups connections that share +/// a driver spawner and network interface binding. /// /// The default identifier denotes an anonymous partition used when /// partitioning is not configured. diff --git a/rust-runtime/aws-smithy-http-client/src/lib.rs b/rust-runtime/aws-smithy-http-client/src/lib.rs index 559f149d4ff..9f6c046ce92 100644 --- a/rust-runtime/aws-smithy-http-client/src/lib.rs +++ b/rust-runtime/aws-smithy-http-client/src/lib.rs @@ -44,7 +44,7 @@ pub(crate) mod client; #[cfg(feature = "default-client")] pub use client::{default_connector, proxy, tls, Builder, Connector, ConnectorBuilder}; -/// V2 HTTP client backed by composable connection pools. +/// HTTP client backed by composable connection pools. #[cfg(feature = "default-client")] pub use client::pool; diff --git a/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs index 9c05d45bb3e..a96dd0125e6 100644 --- a/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs @@ -19,12 +19,13 @@ aws_sdk_unstable ))] -use aws_smithy_http_client::pool::Builder; +use aws_smithy_http_client::pool::{Client, SharedPool}; use aws_smithy_runtime_api::client::http::{ HttpClient, HttpConnector, HttpConnectorSettings, SharedHttpClient, }; use aws_smithy_runtime_api::client::orchestrator::HttpRequest; use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder; +use aws_smithy_runtime_api::shared::IntoShared; use bytes::Bytes; use h2::server::SendResponse; use h2::RecvStream; @@ -298,7 +299,9 @@ impl Service for H2Connector { // --------------------------------------------------------------------------- fn build_h2_client(server: &H2MockServer) -> SharedHttpClient { - Builder::new().build_http_with_tcp_connector(H2Connector { addr: server.addr }) + let pool = + SharedPool::builder().build_http_with_tcp_connector(H2Connector { addr: server.addr }); + Client::new(&pool).into_shared() } fn runtime_components() -> aws_smithy_runtime_api::client::runtime_components::RuntimeComponents { @@ -465,7 +468,7 @@ async fn h2_poisoned_connection_not_reused() { #[cfg(feature = "rustls-aws-lc")] mod tls_h2 { - use aws_smithy_http_client::pool::Builder; + use aws_smithy_http_client::pool::{Client, SharedPool}; use aws_smithy_http_client::tls; use aws_smithy_http_client::tls::{TlsContext, TrustStore}; use aws_smithy_runtime_api::client::http::{ @@ -473,6 +476,7 @@ mod tls_h2 { }; use aws_smithy_runtime_api::client::orchestrator::HttpRequest; use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder; + use aws_smithy_runtime_api::shared::IntoShared; use http_body_util::BodyExt; use hyper_util::rt::TokioExecutor; use std::net::SocketAddr; @@ -614,12 +618,13 @@ mod tls_h2 { #[tokio::test] async fn rustls_alpn_h2_multiplexing() { let server = TlsH2Server::start().await; - let client = Builder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) .tls_context(tls_context()) .build_https(); + let client: SharedHttpClient = Client::new(&pool).into_shared(); let url = server.url(); @@ -649,10 +654,11 @@ mod tls_h2 { #[tokio::test] async fn s2n_tls_alpn_h2_multiplexing() { let server = TlsH2Server::start().await; - let client = Builder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::S2nTls) .tls_context(tls_context()) .build_https(); + let client: SharedHttpClient = Client::new(&pool).into_shared(); let url = server.url(); diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index fe424253608..f6807a3b4f4 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -6,13 +6,12 @@ //! Pool behavior tests parameterized over HTTP client implementations. //! //! Each test is written once as an `async fn` that takes a `&dyn MakeClient`, -//! then invoked for each client implementation. When `Builder::new_v2()` is -//! added, a single new `MakeClient` impl covers all tests. +//! then invoked for each client implementation. #![cfg(all(feature = "wire-mock", feature = "default-client"))] use aws_smithy_async::time::SystemTimeSource; -use aws_smithy_http_client::pool::Builder as PoolBuilder; +use aws_smithy_http_client::pool::{Client as PoolClient, SharedPool}; use aws_smithy_http_client::test_util::wire::connection::{ ConnectionBehavior, ConnectionTestHarness, }; @@ -23,6 +22,7 @@ use aws_smithy_runtime_api::client::http::{ }; use aws_smithy_runtime_api::client::orchestrator::HttpRequest; use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder; +use aws_smithy_runtime_api::shared::IntoShared; use std::borrow::Cow; use std::net::{IpAddr, Ipv4Addr}; use std::time::Duration; @@ -82,7 +82,7 @@ struct V2Client; impl MakeClient for V2Client { fn make(&self, config: ClientConfig) -> SharedHttpClient { - let mut builder = PoolBuilder::new(); + let mut builder = SharedPool::builder(); if let Some(timeout) = config.idle_timeout { builder = builder.pool_idle_timeout(timeout); } @@ -92,7 +92,8 @@ impl MakeClient for V2Client { if let Some(n) = config.max_connections_per_host { builder = builder.max_connections_per_host(n); } - builder.build_http() + let pool = builder.build_http(); + PoolClient::new(&pool).into_shared() } } @@ -1381,12 +1382,14 @@ mod listener_tests { idle_timeout: Option, listener: Arc, ) -> SharedHttpClient { - let mut builder = - aws_smithy_http_client::Builder::new_v2().connection_event_listener(listener); + let mut builder = aws_smithy_http_client::pool::SharedPool::builder() + .connection_event_listener(listener) + .dns_resolver(harness.dns_resolver()); if let Some(timeout) = idle_timeout { builder = builder.pool_idle_timeout(timeout); } - builder.build_http_with_resolver(harness.dns_resolver()) + let pool = builder.build_http(); + PoolClient::new(&pool).into_shared() } /// Listener receives created on first request, reused on second, and diff --git a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs index 126c6a30c64..10bb0540391 100644 --- a/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs +++ b/rust-runtime/aws-smithy-http-client/tests/proxy_tests.rs @@ -1205,9 +1205,12 @@ async fn make_v2_http_request_through_proxy( proxy_config: ProxyConfig, target_url: &str, ) -> Result<(StatusCode, String), Box> { - use aws_smithy_http_client::pool::Builder; + use aws_smithy_http_client::pool::{Client, SharedPool}; - let http_client = Builder::new().proxy_config(proxy_config).build_http(); + let pool = SharedPool::builder() + .proxy_config(proxy_config) + .build_http(); + let http_client = Client::new(&pool); let connector_settings = HttpConnectorSettings::builder().build(); let runtime_components = RuntimeComponentsBuilder::for_tests() .with_time_source(Some(SystemTimeSource::new())) @@ -1234,12 +1237,13 @@ async fn make_v2_https_request_through_proxy( target_url: &str, tls_provider: tls::Provider, ) -> Result<(StatusCode, String), Box> { - use aws_smithy_http_client::pool::Builder; + use aws_smithy_http_client::pool::{Client, SharedPool}; - let http_client = Builder::new() + let pool = SharedPool::builder() .tls_provider(tls_provider) .proxy_config(proxy_config) .build_https(); + let http_client = Client::new(&pool); let connector_settings = HttpConnectorSettings::builder().build(); let runtime_components = RuntimeComponentsBuilder::for_tests() .with_time_source(Some(SystemTimeSource::new())) @@ -1494,7 +1498,7 @@ async fn test_v2_direct_http_origin_uri_form() { #[tokio::test] async fn test_v2_set_proxy_config_none_clears_proxy() { - use aws_smithy_http_client::pool::Builder; + use aws_smithy_http_client::pool::{Client, SharedPool}; let mock_proxy = MockProxyServer::new(|_| { Response::builder() @@ -1511,12 +1515,13 @@ async fn test_v2_set_proxy_config_none_clears_proxy() { }) .await; - let mut builder = Builder::new(); + let mut builder = SharedPool::builder(); builder.set_proxy_config(Some( ProxyConfig::http(format!("http://{}", mock_proxy.addr())).unwrap(), )); builder.set_proxy_config(None); - let http_client = builder.build_http(); + let pool = builder.build_http(); + let http_client = Client::new(&pool); let connector_settings = HttpConnectorSettings::builder().build(); let runtime_components = RuntimeComponentsBuilder::for_tests() diff --git a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs index 3d4f6cb6981..92c30732abc 100644 --- a/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs +++ b/rust-runtime/aws-smithy-http-client/tests/smoke_test_clients.rs @@ -130,48 +130,48 @@ async fn smoke_test_client(client: &dyn HttpClient) -> Result<(), Box // v2 smoke tests // --------------------------------------------------------------------------- -use aws_smithy_http_client::pool::Builder as PoolBuilder; +use aws_smithy_http_client::pool::{Client as PoolClient, SharedPool}; #[cfg(feature = "rustls-aws-lc")] #[tokio::test] async fn v2_aws_lc_client() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) .build_https(); - smoke_test_client(&client).await.unwrap(); + smoke_test_client(&PoolClient::new(&pool)).await.unwrap(); } #[cfg(feature = "rustls-aws-lc-fips")] #[tokio::test] async fn v2_aws_lc_fips_client() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLcFips, )) .build_https(); - smoke_test_client(&client).await.unwrap(); + smoke_test_client(&PoolClient::new(&pool)).await.unwrap(); } #[cfg(feature = "rustls-ring")] #[tokio::test] async fn v2_ring_client() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::Ring, )) .build_https(); - smoke_test_client(&client).await.unwrap(); + smoke_test_client(&PoolClient::new(&pool)).await.unwrap(); } #[cfg(feature = "s2n-tls")] #[tokio::test] async fn v2_s2n_tls_client() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::S2nTls) .build_https(); - smoke_test_client(&client).await.unwrap(); + smoke_test_client(&PoolClient::new(&pool)).await.unwrap(); } #[cfg(feature = "s2n-tls")] @@ -189,11 +189,11 @@ async fn v2_s2n_tls_timing_populated() { } let listener = Arc::new(TimingListener(Mutex::new(None))); - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::S2nTls) .connection_event_listener(listener.clone() as Arc) .build_https(); - smoke_test_client(&client).await.unwrap(); + smoke_test_client(&PoolClient::new(&pool)).await.unwrap(); let duration = listener .0 @@ -221,13 +221,13 @@ async fn v2_rustls_timing_populated() { } let listener = Arc::new(TimingListener(Mutex::new(None))); - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) .connection_event_listener(listener.clone() as Arc) .build_https(); - smoke_test_client(&client).await.unwrap(); + smoke_test_client(&PoolClient::new(&pool)).await.unwrap(); let duration = listener .0 diff --git a/rust-runtime/aws-smithy-http-client/tests/tls.rs b/rust-runtime/aws-smithy-http-client/tests/tls.rs index f1046a37aa5..b1ca19b1306 100644 --- a/rust-runtime/aws-smithy-http-client/tests/tls.rs +++ b/rust-runtime/aws-smithy-http-client/tests/tls.rs @@ -311,98 +311,98 @@ async fn test_s2n_tls_custom_ca() { // v2 TLS tests // --------------------------------------------------------------------------- -use aws_smithy_http_client::pool::Builder as PoolBuilder; +use aws_smithy_http_client::pool::{Client as PoolClient, SharedPool}; #[cfg(feature = "rustls-aws-lc")] #[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] #[tokio::test] async fn test_v2_rustls_aws_lc_native_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } #[cfg(feature = "rustls-aws-lc")] #[tokio::test] async fn test_v2_rustls_aws_lc_custom_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLc, )) .tls_context(tls_context_from_pem("tests/server.pem")) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } #[cfg(feature = "rustls-aws-lc-fips")] #[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] #[tokio::test] async fn test_v2_rustls_aws_lc_fips_native_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLcFips, )) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } #[cfg(feature = "rustls-aws-lc-fips")] #[tokio::test] async fn test_v2_rustls_aws_lc_fips_custom_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::AwsLcFips, )) .tls_context(tls_context_from_pem("tests/server.pem")) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } #[cfg(feature = "rustls-ring")] #[should_panic(expected = "InvalidCertificate(UnknownIssuer)")] #[tokio::test] async fn test_v2_rustls_ring_native_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::Ring, )) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } #[cfg(feature = "rustls-ring")] #[tokio::test] async fn test_v2_rustls_ring_custom_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::Rustls( tls::rustls_provider::CryptoMode::Ring, )) .tls_context(tls_context_from_pem("tests/server.pem")) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } #[cfg(feature = "s2n-tls")] #[should_panic(expected = "Certificate is untrusted")] #[tokio::test] async fn test_v2_s2n_native_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::S2nTls) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } #[cfg(feature = "s2n-tls")] #[tokio::test] async fn test_v2_s2n_tls_custom_ca() { - let client = PoolBuilder::new() + let pool = SharedPool::builder() .tls_provider(tls::Provider::S2nTls) .tls_context(tls_context_from_pem("tests/server.pem")) .build_https(); - run_tls_test(&client).await.unwrap() + run_tls_test(&PoolClient::new(&pool)).await.unwrap() } async fn run_tls_test(client: &dyn HttpClient) -> Result<(), BoxError> { diff --git a/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs b/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs index b3314b95b96..24ca6793286 100644 --- a/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs +++ b/rust-runtime/aws-smithy-runtime/tests/reconnect_on_transient_error.rs @@ -128,12 +128,15 @@ impl MakeClient for Hyper1xClient { } } -/// v2 HTTP client (composable pool) via `BuilderV2::new_v2()` +/// HTTP client backed by the composable connection pool. struct Hyper1xV2Client; impl MakeClient for Hyper1xV2Client { fn make(&self, mock: &WireMockServer) -> SharedHttpClient { - aws_smithy_http_client::v2::BuilderV2::new().build_http_with_resolver(mock.dns_resolver()) + let pool = aws_smithy_http_client::pool::SharedPool::builder() + .dns_resolver(mock.dns_resolver()) + .build_http(); + aws_smithy_http_client::pool::Client::new(&pool).into_shared() } } From af9e2ef4378b8cb47c3cca1d47385c295a7bd8c0 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 21 May 2026 12:33:33 -0400 Subject: [PATCH 26/47] cleanup --- .../aws-smithy-http-client/src/client/pool.rs | 79 ++++++---- .../src/client/pool/builder.rs | 139 +++++++----------- .../src/client/pool/client.rs | 9 +- .../src/client/pool/connection.rs | 11 +- .../src/client/pool/handshake.rs | 2 +- 5 files changed, 119 insertions(+), 121 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 1b2961947d3..f7c47ecf50b 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -3,14 +3,42 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Connection pool for the v2 HTTP client. +//! HTTP connection pool. //! -//! Routes requests by `(scheme, authority)` to per-host pool stacks, -//! negotiates HTTP/1.1 vs HTTP/2 via ALPN, caches H1 connections for -//! reuse, multiplexes over a singleton H2 connection per host, enforces -//! total and per-host connection limits, and surfaces connection -//! metadata (remote/local address, poison handle) to the adapter layer -//! for the runtime's connection-poisoning interceptor. +//! Pool ownership and per-client configuration are separated: +//! +//! - [`SharedPool`] owns connection lifecycle: TLS, DNS resolution, +//! connection limits, idle eviction, proxy routing, event listening. +//! Built via [`SharedPool::builder`] which returns a [`Builder`]. +//! - [`Client`] is a per-partition view over a [`SharedPool`] that +//! implements [`HttpClient`]. Multiple [`Client`]s may share one +//! [`SharedPool`], each with a distinct [`PartitionId`], network +//! interface binding, and driver spawner. +//! +//! For partition semantics, topologies, and cross-partition checkout +//! policy, see the [`partition`] module. +//! +//! # Example +//! +//! ```no_run +//! # #[cfg(feature = "rustls-aws-lc")] +//! # { +//! use aws_smithy_http_client::pool::{Client, SharedPool}; +//! use aws_smithy_http_client::tls; +//! use std::time::Duration; +//! +//! let pool = SharedPool::builder() +//! .tls_provider(tls::Provider::Rustls( +//! tls::rustls_provider::CryptoMode::AwsLc, +//! )) +//! .max_connections(125) +//! .pool_idle_timeout(Duration::from_secs(20)) +//! .build_https(); +//! let client = Client::new(&pool); +//! # } +//! ``` +//! +//! [`HttpClient`]: aws_smithy_runtime_api::client::http::HttpClient pub(crate) mod connection; mod handshake; @@ -20,14 +48,13 @@ pub mod builder; pub mod client; pub mod partition; -// Public re-exports — make pool::Builder, pool::ConnectionEventListener, etc. -// directly accessible without going through pool::builder. -pub use builder::{ - Authority, Builder, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, - ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, - NegotiatedProtocol, -}; +// Public re-exports. +pub use builder::Builder; pub use client::{Client, ClientBuilder}; +pub use connection::{ + Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, + ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, +}; pub use partition::{CrossPartitionPolicy, DriverSpawner, PartitionId}; /// Connection-caching pool layer. @@ -46,9 +73,12 @@ use std::sync::OnceLock; use std::task::Poll; use std::time::{Duration, Instant}; +use aws_smithy_runtime_api::box_error::BoxError; +use aws_smithy_runtime_api::client::connection::ConnectionMetadata; use aws_smithy_types::body::SdkBody; use hyper_util::client::legacy::connect::Connection as HyperConnection; use hyper_util::client::pool as hpool; +use hyper_util::client::proxy::matcher::Matcher as ProxyMatcher; use hyper_util::rt::TokioExecutor; use tokio::sync::{oneshot, Semaphore}; use tower::{Service, ServiceExt}; @@ -60,7 +90,6 @@ use connection::{ pub(crate) use connection::{ConnectCtx, ReadTimeoutHint, TimeoutContext}; use handshake::{H1ConnectAndHandshake, H1SendRequest, H2ConnectAndHandshake}; -pub(crate) type BoxError = Box; type BoxFuture = Pin + Send>>; /// Request-extension slot the pool checkout fills in with the @@ -79,7 +108,7 @@ type BoxFuture = Pin + Send>>; /// the connection used for that attempt. #[derive(Clone, Default)] pub(crate) struct ConnectionMetadataCapture { - slot: Arc>, + slot: Arc>, } impl ConnectionMetadataCapture { @@ -87,18 +116,13 @@ impl ConnectionMetadataCapture { Self::default() } - pub(crate) fn set( - &self, - metadata: aws_smithy_runtime_api::client::connection::ConnectionMetadata, - ) { + pub(crate) fn set(&self, metadata: ConnectionMetadata) { // Silently ignore duplicate sets: single-set is the contract, extra // sets would only happen via pool-internal bugs. let _ = self.slot.set(metadata); } - pub(crate) fn get( - &self, - ) -> Option { + pub(crate) fn get(&self) -> Option { self.slot.get().cloned() } } @@ -134,8 +158,7 @@ pub(crate) struct PoolConfig { /// can reference one `SharedPool`, each presenting a different /// per-partition view of the same underlying connections. /// -/// Construct via [`SharedPool::builder`], which returns a -/// [`Builder`](crate::client::pool::builder::Builder). +/// Construct via [`SharedPool::builder`], which returns a [`Builder`]. /// Cloning is cheap (shared via `Arc`). #[derive(Clone, Debug)] pub struct SharedPool { @@ -144,7 +167,7 @@ pub struct SharedPool { pub(crate) struct SharedPoolInner { pub(crate) pool: Arc, - pub(crate) proxy_matcher: Option>, + pub(crate) proxy_matcher: Option>, } impl std::fmt::Debug for SharedPoolInner { @@ -155,8 +178,8 @@ impl std::fmt::Debug for SharedPoolInner { impl SharedPool { /// Create a [`Builder`] for configuring a new connection pool. - pub fn builder() -> builder::Builder { - builder::Builder::default() + pub fn builder() -> Builder { + Builder::default() } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index 76eee9166ab..8ea91578f91 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -7,47 +7,49 @@ //! //! Entry point: [`SharedPool::builder`](super::SharedPool::builder). -pub use super::connection::{ - Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, - ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, -}; - use std::sync::Arc; use std::time::Duration; +use aws_smithy_runtime_api::client::dns::{ResolveDns, SharedDnsResolver}; +use aws_smithy_runtime_api::shared::IntoShared; +use hyper_util::client::legacy::connect::dns::Name as DnsName; use hyper_util::client::legacy::connect::HttpConnector as HyperHttpConnector; +use hyper_util::client::proxy::matcher::Matcher as ProxyMatcher; -use super::{ConnectionPool, PoolConfig}; +use super::connection::ConnectionEventListener; +use super::partition::CrossPartitionPolicy; +use super::{BoxError, ConnectionPool, PoolConfig}; +use crate::client::dns::HyperUtilResolver; +use crate::client::proxy::ProxyConfig; use crate::client::tls; use crate::client::{TlsProviderSelected, TlsUnset}; use crate::tls::TlsContext; -/// Builder for the v2 HTTP client. -/// -/// Provides a flat, unified configuration surface for connection management, -/// TLS, timeouts, and pool behavior (all knobs in one place). +/// Builder for a [`SharedPool`]. /// -/// # Examples +/// Configures pool-wide settings: TLS, DNS, connection limits, idle +/// eviction, proxy, and connection event listening. Per-client settings +/// (partition, network interface, driver spawner) are configured +/// separately via [`ClientBuilder`]. /// -/// ```no_run -/// use aws_smithy_http_client::pool::{SharedPool, Client}; -/// use std::time::Duration; +/// Type-state ensures TLS is configured before [`build_https`] is +/// callable; calling [`tls_provider`] transitions the builder into the +/// state where [`build_https`] is available. /// -/// let pool = SharedPool::builder() -/// .pool_idle_timeout(Duration::from_secs(20)) -/// .build_http(); -/// let client = Client::new(&pool); -/// ``` +/// [`build_https`]: Builder::build_https +/// [`tls_provider`]: Builder::tls_provider +/// [`SharedPool`]: super::SharedPool +/// [`ClientBuilder`]: super::ClientBuilder #[derive(Clone)] pub struct Builder { pool_idle_timeout: Option, tcp_nodelay: bool, max_connections: Option, max_connections_per_host: Option, - proxy_config: Option, - connection_event_listener: Option>, - cross_partition_policy: super::partition::CrossPartitionPolicy, - dns_resolver: Option, + proxy_config: Option, + connection_event_listener: Option>, + cross_partition_policy: CrossPartitionPolicy, + dns_resolver: Option, tls: Tls, } @@ -78,7 +80,7 @@ impl Default for Builder { max_connections_per_host: None, proxy_config: None, connection_event_listener: None, - cross_partition_policy: super::partition::CrossPartitionPolicy::default(), + cross_partition_policy: CrossPartitionPolicy::default(), dns_resolver: None, tls: TlsUnset {}, } @@ -89,8 +91,9 @@ impl Default for Builder { impl Builder { /// Set the pool idle timeout. /// - /// Connections idle longer than this duration are evicted from the pool. - /// Set below the server's idle timeout to avoid stale connection errors. + /// Connections idle longer than this duration are evicted from the + /// pool. Set below the server's idle timeout to avoid stale + /// connection errors. Defaults to no eviction when unset. pub fn pool_idle_timeout(mut self, timeout: Duration) -> Self { self.pool_idle_timeout = Some(timeout); self @@ -127,35 +130,28 @@ impl Builder { self } - /// Configure proxy support. + /// Route connections through an HTTP/HTTPS/SOCKS proxy. /// - /// Threads `config` through the connector chain: HTTP/HTTPS proxies - /// route requests via the proxy server; HTTPS-over-proxy uses - /// `CONNECT` tunneling. The per-host proxy decision is stable for - /// this builder's lifetime: all connections to a given authority - /// go through the same path. + /// Per-host proxy resolution is stable for the lifetime of the + /// pool: all connections to a given authority follow the same + /// proxy decision. HTTPS through an HTTP proxy uses `CONNECT` + /// tunneling. /// - /// To use environment variables (`HTTP_PROXY`, `HTTPS_PROXY`, - /// `NO_PROXY`), pass `ProxyConfig::from_env()`. - pub fn proxy_config(mut self, config: crate::client::proxy::ProxyConfig) -> Self { + /// Pass [`ProxyConfig::from_env`] to read configuration from the + /// `HTTP_PROXY`, `HTTPS_PROXY`, and `NO_PROXY` environment variables. + pub fn proxy_config(mut self, config: ProxyConfig) -> Self { self.proxy_config = Some(config); self } /// Mutable-reference variant of [`proxy_config`](Self::proxy_config). - pub fn set_proxy_config( - &mut self, - config: Option, - ) -> &mut Self { + pub fn set_proxy_config(&mut self, config: Option) -> &mut Self { self.proxy_config = config; self } /// Set a listener for connection lifecycle events (created, reused, closed, failed). - pub fn connection_event_listener( - mut self, - listener: Arc, - ) -> Self { + pub fn connection_event_listener(mut self, listener: Arc) -> Self { self.connection_event_listener = Some(listener); self } @@ -163,7 +159,7 @@ impl Builder { /// Mutable-reference variant of [`connection_event_listener`](Self::connection_event_listener). pub fn set_connection_event_listener( &mut self, - listener: Option>, + listener: Option>, ) -> &mut Self { self.connection_event_listener = listener; self @@ -172,22 +168,16 @@ impl Builder { /// Set the policy that governs checkout when the local partition has /// no idle connection and the pool is at capacity. /// - /// Defaults to [`CrossPartitionPolicy::Never`](super::partition::CrossPartitionPolicy::Never). - /// Has no observable effect with a single partition or when the pool - /// stays under capacity. - pub fn cross_partition_policy( - mut self, - policy: super::partition::CrossPartitionPolicy, - ) -> Self { + /// Defaults to [`CrossPartitionPolicy::Never`]. Has no observable + /// effect with a single partition or when the pool stays under + /// capacity. + pub fn cross_partition_policy(mut self, policy: CrossPartitionPolicy) -> Self { self.cross_partition_policy = policy; self } /// Mutable-reference variant of [`cross_partition_policy`](Self::cross_partition_policy). - pub fn set_cross_partition_policy( - &mut self, - policy: super::partition::CrossPartitionPolicy, - ) -> &mut Self { + pub fn set_cross_partition_policy(&mut self, policy: CrossPartitionPolicy) -> &mut Self { self.cross_partition_policy = policy; self } @@ -196,20 +186,13 @@ impl Builder { /// /// Connections established by this pool resolve hostnames through the /// given resolver. Defaults to the system resolver when unset. - pub fn dns_resolver( - mut self, - resolver: impl aws_smithy_runtime_api::client::dns::ResolveDns + 'static, - ) -> Self { - use aws_smithy_runtime_api::shared::IntoShared; + pub fn dns_resolver(mut self, resolver: impl ResolveDns + 'static) -> Self { self.dns_resolver = Some(resolver.into_shared()); self } /// Mutable-reference variant of [`dns_resolver`](Self::dns_resolver). - pub fn set_dns_resolver( - &mut self, - resolver: Option, - ) -> &mut Self { + pub fn set_dns_resolver(&mut self, resolver: Option) -> &mut Self { self.dns_resolver = resolver; self } @@ -247,7 +230,6 @@ impl Builder { let proxy_matcher = proxy_matcher_from(&self.proxy_config); let pool = match dns_resolver { Some(resolver) => { - use crate::client::dns::HyperUtilResolver; let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); tcp.set_nodelay(self.tcp_nodelay); build_http_pool_with_proxy(tcp, &self.proxy_config, config) @@ -277,7 +259,7 @@ impl Builder { pub fn build_http_with_tcp_connector(self, connector: C) -> super::SharedPool where C: tower::Service + Clone + Send + Sync + 'static, - C::Error: Into + 'static, + C::Error: Into + 'static, C::Future: Unpin + Send + 'static, IO: hyper::rt::Read + hyper::rt::Write @@ -308,21 +290,17 @@ impl Builder { /// proxy will fail at handshake time. fn build_http_pool_with_proxy( tcp: HyperHttpConnector, - proxy_config: &Option, + proxy_config: &Option, config: PoolConfig, ) -> ConnectionPool where R: Clone + Send + Sync + 'static, - R: tower::Service, + R: tower::Service, R::Response: Iterator, R::Future: Send, - R::Error: Into, + R::Error: Into, { - use crate::client::proxy; - - let proxy_config = proxy_config - .clone() - .unwrap_or_else(proxy::ProxyConfig::disabled); + let proxy_config = proxy_config.clone().unwrap_or_else(ProxyConfig::disabled); if proxy_config.requires_tls() { tracing::warn!( @@ -352,7 +330,6 @@ impl Builder { let dns_resolver = self.dns_resolver.take(); match dns_resolver { Some(resolver) => { - use crate::client::dns::HyperUtilResolver; let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); tcp.set_nodelay(self.tcp_nodelay); tcp.enforce_http(false); @@ -370,13 +347,11 @@ impl Builder { fn build_from_tcp(self, tcp: HyperHttpConnector) -> super::SharedPool where R: Clone + Send + Sync + 'static, - R: tower::Service, + R: tower::Service, R::Response: Iterator, R::Future: Send, - R::Error: Into, + R::Error: Into, { - use crate::client::proxy; - let config = PoolConfig { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, @@ -387,7 +362,7 @@ impl Builder { let proxy_config = self .proxy_config .clone() - .unwrap_or_else(proxy::ProxyConfig::disabled); + .unwrap_or_else(ProxyConfig::disabled); let proxy_matcher = proxy_matcher_from(&self.proxy_config); match &self.tls.provider { @@ -438,9 +413,7 @@ impl Builder { /// no proxy is configured. Wrapped in `Arc` so each `PooledConnector` /// constructed by `http_connector` can share the (immutable) matcher /// without per-call allocation. -pub(super) fn proxy_matcher_from( - proxy_config: &Option, -) -> Option> { +pub(super) fn proxy_matcher_from(proxy_config: &Option) -> Option> { proxy_config .as_ref() .map(|c| Arc::new(c.clone().into_hyper_util_matcher())) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs index e0f647430f0..be1f6f3150d 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs @@ -9,6 +9,7 @@ use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; +use aws_smithy_async::rt::sleep::SharedAsyncSleep; use aws_smithy_runtime_api::client::connection::CaptureSmithyConnection; use aws_smithy_runtime_api::client::connector_metadata::ConnectorMetadata; use aws_smithy_runtime_api::client::http::{ @@ -17,10 +18,12 @@ use aws_smithy_runtime_api::client::http::{ use aws_smithy_runtime_api::client::orchestrator::{HttpRequest, HttpResponse}; use aws_smithy_runtime_api::client::result::ConnectorError; use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; +use hyper_util::client::proxy::matcher::Matcher as ProxyMatcher; use super::partition::{DriverSpawner, PartitionId}; use super::{ConnectionPool, SharedPool}; use crate::client::downcast_error; +use crate::client::proxy::add_proxy_auth_header; /// Per-partition view of a [`SharedPool`]. /// @@ -187,8 +190,8 @@ struct PooledConnector { pool: Arc, connect_timeout: Option, read_timeout: Option, - sleep_impl: Option, - proxy_matcher: Option>, + sleep_impl: Option, + proxy_matcher: Option>, } impl std::fmt::Debug for PooledConnector { @@ -212,7 +215,7 @@ impl HttpConnector for PooledConnector { let full_uri = request.uri().clone(); if let Some(matcher) = proxy_matcher.as_ref() { - crate::client::proxy::add_proxy_auth_header(&mut request, matcher); + add_proxy_auth_header(&mut request, matcher); } if let Some(capture_smithy) = request.extensions().get::() { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 75beab87e8d..d2f7d528b37 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -12,6 +12,7 @@ use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use aws_smithy_async::rt::sleep::SharedAsyncSleep; +use aws_smithy_runtime_api::box_error::BoxError; use aws_smithy_runtime_api::client::connection::{ConnectionId, ConnectionMetadata}; use aws_smithy_types::body::SdkBody; use pin_project_lite::pin_project; @@ -20,9 +21,8 @@ use tower::Service; use super::cache; use super::handshake::H1SendRequest; -use super::BoxError; -/// Opaque authority (host:port) associated with a pooled connection. +/// Authority of a pooled connection (host with optional port). #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct Authority(Arc); @@ -289,11 +289,10 @@ impl ConnectionFailedEvent { /// Callback for connection lifecycle events within the pool. /// /// Implementations receive notifications when connections are created, -/// reused from the pool, closed, or fail to establish. Useful for metrics, -/// DNS failure feedback, or adaptive concurrency control. +/// reused from the pool, closed, or fail to establish. /// -/// Implementations must be cheap (nanoseconds). Expensive work should be -/// deferred to a background task. +/// Implementations must be non-blocking. Defer expensive work to a +/// background task. pub trait ConnectionEventListener: Send + Sync + 'static { /// A new connection was established. fn on_created(&self, _event: &ConnectionCreatedEvent) {} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index a70f16ece5d..ff7e96a18aa 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Instant; +use aws_smithy_runtime_api::box_error::BoxError; use aws_smithy_runtime_api::client::connection::ConnectionId; use aws_smithy_types::body::SdkBody; use hyper::rt::Executor; @@ -28,7 +29,6 @@ use super::connection::{ Authority, ConnectCtx, ConnectionCreatedEvent, ConnectionFailedEvent, ConnectionInfo, ConnectionPermit, ConnectionTiming, ManagedConnection, NegotiatedProtocol, }; -use super::BoxError; /// Pool-scoped instrumentation primitives shared across layers in the /// pool stack. Held by `ConnectionPool` for the pool's lifetime and From e35c95cf21a32b04647ce028f67b5ea5341b80f0 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 22 May 2026 12:44:31 -0400 Subject: [PATCH 27/47] refactor driver spawner --- .../aws-smithy-http-client/src/client/pool.rs | 2 +- .../src/client/pool/client.rs | 18 ++-- .../src/client/pool/partition.rs | 93 ++++++++++++------- 3 files changed, 70 insertions(+), 43 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index f7c47ecf50b..26ba158d53e 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -55,7 +55,7 @@ pub use connection::{ Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, }; -pub use partition::{CrossPartitionPolicy, DriverSpawner, PartitionId}; +pub use partition::{CrossPartitionPolicy, DriverSpawner, PartitionId, TokioDriverSpawner}; /// Connection-caching pool layer. mod cache { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs index be1f6f3150d..a26344e4ffe 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs @@ -20,7 +20,7 @@ use aws_smithy_runtime_api::client::result::ConnectorError; use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; use hyper_util::client::proxy::matcher::Matcher as ProxyMatcher; -use super::partition::{DriverSpawner, PartitionId}; +use super::partition::{DriverSpawner, PartitionId, TokioDriverSpawner}; use super::{ConnectionPool, SharedPool}; use crate::client::downcast_error; use crate::client::proxy::add_proxy_auth_header; @@ -41,7 +41,7 @@ pub struct Client { pool: SharedPool, partition: PartitionId, interface: Option, - driver_spawner: DriverSpawner, + driver_spawner: Arc, } impl std::fmt::Debug for Client { @@ -119,7 +119,7 @@ pub struct ClientBuilder { pool: SharedPool, partition: PartitionId, interface: Option, - driver_spawner: Option, + driver_spawner: Option>, } impl ClientBuilder { @@ -148,14 +148,14 @@ impl ClientBuilder { } /// Set the driver spawner (consuming builder). - pub fn driver_spawner(mut self, sp: DriverSpawner) -> Self { - self.driver_spawner = Some(sp); + pub fn driver_spawner(mut self, sp: S) -> Self { + self.driver_spawner = Some(Arc::new(sp)); self } /// Set the driver spawner (mutable reference). - pub fn set_driver_spawner(&mut self, sp: DriverSpawner) -> &mut Self { - self.driver_spawner = Some(sp); + pub fn set_driver_spawner(&mut self, sp: S) -> &mut Self { + self.driver_spawner = Some(Arc::new(sp)); self } @@ -164,9 +164,9 @@ impl ClientBuilder { /// Captures the current tokio runtime handle for the driver spawner /// if one was not explicitly provided. pub fn build(self) -> Client { - let driver_spawner = self + let driver_spawner: Arc = self .driver_spawner - .unwrap_or_else(DriverSpawner::current_tokio); + .unwrap_or_else(|| Arc::new(TokioDriverSpawner::current())); Client { pool: self.pool, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 23e33da88fd..862259e7560 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -168,46 +168,52 @@ impl Default for PartitionId { /// Spawner for connection driver tasks. /// -/// A driver is the task that owns a connection's I/O. The spawner -/// determines which runtime executes that task. Per-runtime spawners -/// allow each connection's I/O to remain on a specific runtime. +/// A driver is the task that owns an HTTP connection's I/O state machine: +/// reading frames, writing frames, and managing protocol-level events. +/// Calls through a connection's request handle flow through this driver. +/// The pool spawns one driver per established connection. /// -/// Construct via [`DriverSpawner::current_tokio`] (captures the -/// current tokio runtime handle) or [`DriverSpawner::from_tokio_handle`] -/// (uses a specific runtime handle). -#[derive(Clone)] -pub struct DriverSpawner { - inner: SpawnerInner, +/// Different partitions may use different spawners, allowing each +/// partition's drivers to run on a specific runtime. +pub trait DriverSpawner: std::fmt::Debug + Send + Sync + 'static { + /// Spawn the connection driver future on this spawner's runtime. + fn spawn( + &self, + driver: std::pin::Pin + Send + 'static>>, + ); } -#[derive(Clone)] -enum SpawnerInner { - TokioHandle(tokio::runtime::Handle), +/// Driver spawner backed by a tokio runtime handle. +/// +/// Spawns the connection driver via [`tokio::runtime::Handle::spawn`]. +/// The handle is captured at construction; the driver runs on the +/// runtime the handle refers to, regardless of which runtime called +/// [`DriverSpawner::spawn`]. +#[derive(Clone, Debug)] +pub struct TokioDriverSpawner { + handle: tokio::runtime::Handle, } -impl DriverSpawner { - /// Spawner that uses the current tokio runtime. +impl TokioDriverSpawner { + /// Spawner using the current tokio runtime handle (captured eagerly). /// - /// Captures the runtime handle eagerly at the call site. Panics - /// if invoked outside a tokio runtime context. - pub fn current_tokio() -> Self { - Self::from_tokio_handle(tokio::runtime::Handle::current()) + /// Panics if invoked outside a tokio runtime context. + pub fn current() -> Self { + Self::from_handle(tokio::runtime::Handle::current()) } - /// Spawner targeting the given tokio runtime handle. - pub fn from_tokio_handle(handle: tokio::runtime::Handle) -> Self { - Self { - inner: SpawnerInner::TokioHandle(handle), - } + /// Spawner using a specific tokio runtime handle. + pub fn from_handle(handle: tokio::runtime::Handle) -> Self { + Self { handle } } } -impl std::fmt::Debug for DriverSpawner { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let variant = match &self.inner { - SpawnerInner::TokioHandle(_) => "TokioHandle", - }; - f.debug_tuple("DriverSpawner").field(&variant).finish() +impl DriverSpawner for TokioDriverSpawner { + fn spawn( + &self, + driver: std::pin::Pin + Send + 'static>>, + ) { + self.handle.spawn(driver); } } @@ -253,13 +259,34 @@ mod tests { } #[tokio::test] - async fn driver_spawner_current_tokio() { - let _ = DriverSpawner::current_tokio(); + async fn tokio_driver_spawner_current() { + let _ = TokioDriverSpawner::current(); } #[tokio::test] - async fn driver_spawner_from_handle() { + async fn tokio_driver_spawner_from_handle() { let h = tokio::runtime::Handle::current(); - let _ = DriverSpawner::from_tokio_handle(h); + let _ = TokioDriverSpawner::from_handle(h); + } + + #[tokio::test] + async fn tokio_driver_spawner_runs_future() { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + + let sp = TokioDriverSpawner::current(); + let flag = Arc::new(AtomicBool::new(false)); + let f = flag.clone(); + sp.spawn(Box::pin(async move { + f.store(true, Ordering::SeqCst); + })); + // Yield enough times for the spawned task to run. + for _ in 0..10 { + tokio::task::yield_now().await; + if flag.load(Ordering::SeqCst) { + break; + } + } + assert!(flag.load(Ordering::SeqCst), "spawned future did not run"); } } From bdc662b29a9fb9a6851f8fff789830f3ee70b192 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Sun, 24 May 2026 10:08:29 -0400 Subject: [PATCH 28/47] group connection into struct --- .../aws-smithy-http-client/src/client/pool.rs | 8 +++---- .../src/client/pool/connection.rs | 11 ++++++++++ .../src/client/pool/handshake.rs | 21 ++++++++++--------- 3 files changed, 25 insertions(+), 15 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 26ba158d53e..7745baa5bca 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -274,11 +274,9 @@ where let stack = hpool::negotiate::builder() .connect(limited) - .inspect( - |(conn, _permit): &(IO, Arc)| { - conn.connected().is_negotiated_h2() - }, - ) + .inspect(|established: &connection::EstablishedConnection| { + established.io.connected().is_negotiated_h2() + }) .fallback({ let retainers = retainers.clone(); let pool_hooks = pool_hooks.clone(); diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index d2f7d528b37..e0d166fa277 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -377,6 +377,17 @@ impl ConnectionPermit { } } +/// Outcome of establishing a new transport connection (TCP + TLS). +/// +/// Carries the IO handle and the connection permit acquired during +/// connection establishment. The IO handle drives the protocol +/// handshake; the permit holds the connection's slot in the pool's +/// limit semaphores until the connection is dropped. +pub(crate) struct EstablishedConnection { + pub(crate) io: IO, + pub(crate) permit: Arc, +} + /// Error returned by `ManagedConnection::poll_ready` when the connection /// has been marked poisoned and should not be reused. #[derive(Debug)] diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index ff7e96a18aa..7a1cafd943d 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -6,7 +6,7 @@ //! Tower service adapters for hyper's HTTP protocol handshake. //! //! - `ConnectionLimit`: wraps a connector, acquires semaphore permits before -//! connecting, and returns `(IO, Arc)`. +//! connecting, and returns an `EstablishedConnection`. //! - `H1ConnectAndHandshake` / `H2ConnectAndHandshake`: perform the protocol //! handshake on an already-connected IO stream, producing a `ManagedConnection`. @@ -27,7 +27,8 @@ use tower::Service; use super::connection::{ Authority, ConnectCtx, ConnectionCreatedEvent, ConnectionFailedEvent, ConnectionInfo, - ConnectionPermit, ConnectionTiming, ManagedConnection, NegotiatedProtocol, + ConnectionPermit, ConnectionTiming, EstablishedConnection, ManagedConnection, + NegotiatedProtocol, }; /// Pool-scoped instrumentation primitives shared across layers in the @@ -88,7 +89,7 @@ impl PoolHooks { /// Wraps a connector service, acquiring semaphore permits before connecting. /// -/// Returns `(IO, Arc)` so the permit can be stored on +/// Returns an [`EstablishedConnection`] so the permit can be stored on /// `ManagedConnection` and held for the connection's lifetime. /// /// Target type is `ConnectCtx`: the inner TCP connector is `Service`, @@ -134,7 +135,7 @@ where C::Future: Send + 'static, IO: Send + 'static, { - type Response = (IO, Arc); + type Response = EstablishedConnection; type Error = BoxError; type Future = Pin> + Send>>; @@ -177,7 +178,7 @@ where super::super::timeout::TimeoutKind::Connect, ) .await?; - Ok((io, permit)) + Ok(EstablishedConnection { io, permit }) }) } } @@ -267,7 +268,7 @@ fn spawn_driver(future: impl Future + Send + 'static) { /// Connects and performs an HTTP/1.1 handshake. /// -/// The connector is expected to return `(IO, Arc)`, typically +/// The connector is expected to return an [`EstablishedConnection`], typically /// produced by [`ConnectionLimit`] wrapping a TCP/TLS connector. pub(crate) struct H1ConnectAndHandshake { connector: C, @@ -291,7 +292,7 @@ impl Clone for H1ConnectAndHandshake { impl Service for H1ConnectAndHandshake where - C: Service)>, + C: Service>, C::Error: Into + 'static, C::Future: Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, @@ -315,7 +316,7 @@ where let fut = self.connector.call(ctx); Box::pin(async move { let connect_start = Instant::now(); - let (io, permit) = match fut.await.map_err(Into::into) { + let EstablishedConnection { io, permit } = match fut.await.map_err(Into::into) { Ok(v) => v, Err(e) => { hooks.on_connection_failed(&ConnectionFailedEvent::new( @@ -436,7 +437,7 @@ impl Clone for H2ConnectAndHandshake { impl Service<()> for H2ConnectAndHandshake where - C: Service<(), Response = (IO, Arc)>, + C: Service<(), Response = EstablishedConnection>, C::Error: Into + 'static, C::Future: Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static, @@ -456,7 +457,7 @@ where let authority = self.authority.clone(); Box::pin(async move { let connect_start = Instant::now(); - let (io, permit) = match fut.await.map_err(Into::into) { + let EstablishedConnection { io, permit } = match fut.await.map_err(Into::into) { Ok(v) => v, Err(e) => { hooks.on_connection_failed(&ConnectionFailedEvent::new( From 31c29a56c59520553ed1846555d37ace56e024e3 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 4 Jun 2026 13:25:35 -0400 Subject: [PATCH 29/47] rework partition config where pool is topology aware up front --- aws/rust-runtime/Cargo.lock | 15 ++ .../aws-smithy-http-client/src/client/pool.rs | 41 ++++- .../src/client/pool/builder.rs | 50 ++++-- .../src/client/pool/client.rs | 155 +++++++----------- .../src/client/pool/partition.rs | 118 +++++++++++++ 5 files changed, 264 insertions(+), 115 deletions(-) diff --git a/aws/rust-runtime/Cargo.lock b/aws/rust-runtime/Cargo.lock index dad8ae8767d..89cd163f4c8 100644 --- a/aws/rust-runtime/Cargo.lock +++ b/aws/rust-runtime/Cargo.lock @@ -652,6 +652,21 @@ dependencies = [ "libc", ] +[[package]] +name = "crc" +version = "3.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9710d3b3739c2e349eb44fe848ad0b7c8cb1e42bd87ee49371df2f7acaf3e675" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + [[package]] name = "crc-fast" version = "1.9.0" diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 7745baa5bca..befd233b510 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -5,15 +5,15 @@ //! HTTP connection pool. //! -//! Pool ownership and per-client configuration are separated: +//! Pool ownership and partition topology are declared at build time: //! //! - [`SharedPool`] owns connection lifecycle: TLS, DNS resolution, -//! connection limits, idle eviction, proxy routing, event listening. -//! Built via [`SharedPool::builder`] which returns a [`Builder`]. +//! connection limits, idle eviction, proxy routing, event listening, +//! and the partition registry. Built via [`SharedPool::builder`] which +//! returns a [`Builder`]. //! - [`Client`] is a per-partition view over a [`SharedPool`] that //! implements [`HttpClient`]. Multiple [`Client`]s may share one -//! [`SharedPool`], each with a distinct [`PartitionId`], network -//! interface binding, and driver spawner. +//! [`SharedPool`], each targeting a distinct declared partition. //! //! For partition semantics, topologies, and cross-partition checkout //! policy, see the [`partition`] module. @@ -50,12 +50,14 @@ pub mod partition; // Public re-exports. pub use builder::Builder; -pub use client::{Client, ClientBuilder}; +pub use client::Client; pub use connection::{ Authority, CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, ConnectionFailedEvent, ConnectionReusedEvent, ConnectionTiming, NegotiatedProtocol, }; -pub use partition::{CrossPartitionPolicy, DriverSpawner, PartitionId, TokioDriverSpawner}; +pub use partition::{ + CrossPartitionPolicy, DriverSpawner, Partition, PartitionId, TokioDriverSpawner, +}; /// Connection-caching pool layer. mod cache { @@ -230,7 +232,12 @@ impl BoxedRetainer { /// `layer_fn`s the first time `Negotiate` constructs each leg. type RetainerSlot = Arc>>; -pub(crate) fn build_pool(connector: C, config: PoolConfig) -> ConnectionPool +pub(crate) fn build_pool( + connector: C, + config: PoolConfig, + registry: Arc, + cross_partition_policy: partition::CrossPartitionPolicy, +) -> ConnectionPool where C: Service + Clone + Send + Sync + 'static, C::Error: Into + 'static, @@ -458,6 +465,8 @@ where make_entry: Box::new(make_entry), eviction_spawned: AtomicBool::new(false), drop_notifier: OnceLock::new(), + registry, + cross_partition_policy, } } @@ -472,6 +481,13 @@ pub(crate) struct ConnectionPool { hosts: Mutex>>, make_entry: Box Box + Send + Sync>, + /// Immutable partition registry, resolved at build time. + registry: Arc, + + // stored for later steps (send path unchanged this step) + #[allow(dead_code)] + cross_partition_policy: partition::CrossPartitionPolicy, + /// Latches `true` once the eviction task has been spawned. Subsequent /// `send_request` calls observe the latch and skip the spawn path. /// Read only when `config.pool_idle_timeout` is set. @@ -493,6 +509,11 @@ pub(crate) struct ConnectionPool { const MIN_EVICTION_TICK: Duration = Duration::from_millis(90); impl ConnectionPool { + /// Access the immutable partition registry. + pub(crate) fn registry(&self) -> &Arc { + &self.registry + } + /// Send a request through the pool. /// /// Routes to the appropriate per-host pool stack and sends the request. @@ -1000,6 +1021,10 @@ mod tests { make_entry: Box::new(|_| Box::new(NullPoolEntry)), eviction_spawned: AtomicBool::new(false), drop_notifier: OnceLock::new(), + registry: Arc::new(partition::PartitionRegistry::build(Vec::new(), || { + Arc::new(partition::TokioDriverSpawner::current()) + })), + cross_partition_policy: CrossPartitionPolicy::default(), }) } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index 8ea91578f91..0c87badffbe 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -17,7 +17,7 @@ use hyper_util::client::legacy::connect::HttpConnector as HyperHttpConnector; use hyper_util::client::proxy::matcher::Matcher as ProxyMatcher; use super::connection::ConnectionEventListener; -use super::partition::CrossPartitionPolicy; +use super::partition::{CrossPartitionPolicy, Partition}; use super::{BoxError, ConnectionPool, PoolConfig}; use crate::client::dns::HyperUtilResolver; use crate::client::proxy::ProxyConfig; @@ -28,9 +28,7 @@ use crate::tls::TlsContext; /// Builder for a [`SharedPool`]. /// /// Configures pool-wide settings: TLS, DNS, connection limits, idle -/// eviction, proxy, and connection event listening. Per-client settings -/// (partition, network interface, driver spawner) are configured -/// separately via [`ClientBuilder`]. +/// eviction, proxy, partition topology, and connection event listening. /// /// Type-state ensures TLS is configured before [`build_https`] is /// callable; calling [`tls_provider`] transitions the builder into the @@ -39,7 +37,6 @@ use crate::tls::TlsContext; /// [`build_https`]: Builder::build_https /// [`tls_provider`]: Builder::tls_provider /// [`SharedPool`]: super::SharedPool -/// [`ClientBuilder`]: super::ClientBuilder #[derive(Clone)] pub struct Builder { pool_idle_timeout: Option, @@ -50,6 +47,7 @@ pub struct Builder { connection_event_listener: Option>, cross_partition_policy: CrossPartitionPolicy, dns_resolver: Option, + partitions: Vec, tls: Tls, } @@ -67,6 +65,7 @@ impl std::fmt::Debug for Builder { ) .field("cross_partition_policy", &self.cross_partition_policy) .field("dns_resolver", &self.dns_resolver.as_ref().map(|_| "..")) + .field("partitions", &self.partitions.len()) .finish() } } @@ -82,6 +81,7 @@ impl Default for Builder { connection_event_listener: None, cross_partition_policy: CrossPartitionPolicy::default(), dns_resolver: None, + partitions: Vec::new(), tls: TlsUnset {}, } } @@ -196,6 +196,14 @@ impl Builder { self.dns_resolver = resolver; self } + + /// Declare the pool's partition topology. Each [`Partition`] carries a + /// driver-spawner runtime and optional NIC binding. When omitted, the + /// pool uses a single anonymous partition on the current tokio runtime. + pub fn partitions(mut self, partitions: impl IntoIterator) -> Self { + self.partitions = partitions.into_iter().collect(); + self + } } impl Builder { @@ -210,6 +218,7 @@ impl Builder { connection_event_listener: self.connection_event_listener, cross_partition_policy: self.cross_partition_policy, dns_resolver: self.dns_resolver, + partitions: self.partitions, tls: TlsProviderSelected { provider, context: TlsContext::default(), @@ -228,16 +237,21 @@ impl Builder { connection_event_listener: self.connection_event_listener.clone(), }; let proxy_matcher = proxy_matcher_from(&self.proxy_config); + let registry = Arc::new(super::partition::PartitionRegistry::build( + std::mem::take(&mut self.partitions), + || Arc::new(super::partition::TokioDriverSpawner::current()), + )); + let policy = self.cross_partition_policy; let pool = match dns_resolver { Some(resolver) => { let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); tcp.set_nodelay(self.tcp_nodelay); - build_http_pool_with_proxy(tcp, &self.proxy_config, config) + build_http_pool_with_proxy(tcp, &self.proxy_config, config, registry, policy) } None => { let mut tcp = HyperHttpConnector::new(); tcp.set_nodelay(self.tcp_nodelay); - build_http_pool_with_proxy(tcp, &self.proxy_config, config) + build_http_pool_with_proxy(tcp, &self.proxy_config, config, registry, policy) } }; super::SharedPool { @@ -274,7 +288,12 @@ impl Builder { pool_idle_timeout: self.pool_idle_timeout, connection_event_listener: self.connection_event_listener.clone(), }; - let pool = super::build_pool(connector, config); + let registry = Arc::new(super::partition::PartitionRegistry::build( + self.partitions, + || Arc::new(super::partition::TokioDriverSpawner::current()), + )); + let policy = self.cross_partition_policy; + let pool = super::build_pool(connector, config, registry, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), @@ -292,6 +311,8 @@ fn build_http_pool_with_proxy( tcp: HyperHttpConnector, proxy_config: &Option, config: PoolConfig, + registry: Arc, + policy: CrossPartitionPolicy, ) -> ConnectionPool where R: Clone + Send + Sync + 'static, @@ -311,10 +332,10 @@ where } if proxy_config.is_disabled() { - super::build_pool(tcp, config) + super::build_pool(tcp, config, registry, policy) } else { let connector = crate::client::connect::HttpProxyConnector::new(tcp, proxy_config); - super::build_pool(connector, config) + super::build_pool(connector, config, registry, policy) } } @@ -364,6 +385,11 @@ impl Builder { .clone() .unwrap_or_else(ProxyConfig::disabled); let proxy_matcher = proxy_matcher_from(&self.proxy_config); + let registry = Arc::new(super::partition::PartitionRegistry::build( + self.partitions, + || Arc::new(super::partition::TokioDriverSpawner::current()), + )); + let policy = self.cross_partition_policy; match &self.tls.provider { #[cfg(any( @@ -378,7 +404,7 @@ impl Builder { &self.tls.context, proxy_config, ); - let pool = super::build_pool(connector, config); + let pool = super::build_pool(connector, config, registry, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), @@ -393,7 +419,7 @@ impl Builder { &self.tls.context, proxy_config, ); - let pool = super::build_pool(connector, config); + let pool = super::build_pool(connector, config, registry, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs index a26344e4ffe..373a7232149 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Per-partition client and its builder. +//! Per-partition client handle. use std::borrow::Cow; use std::sync::Arc; @@ -20,7 +20,7 @@ use aws_smithy_runtime_api::client::result::ConnectorError; use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents; use hyper_util::client::proxy::matcher::Matcher as ProxyMatcher; -use super::partition::{DriverSpawner, PartitionId, TokioDriverSpawner}; +use super::partition::{PartitionId, PartitionState}; use super::{ConnectionPool, SharedPool}; use crate::client::downcast_error; use crate::client::proxy::add_proxy_auth_header; @@ -29,48 +29,54 @@ use crate::client::proxy::add_proxy_auth_header; /// /// Implements [`HttpClient`] by routing requests through the shared /// connection pool. Multiple `Client` instances can reference the same -/// pool, each carrying a distinct [`PartitionId`], network interface -/// binding, and [`DriverSpawner`]. +/// pool, each targeting a distinct declared partition. /// /// Construct via [`Client::new`] (default partition) or -/// [`Client::builder`] for per-partition configuration. +/// [`Client::from_partition`] for a specific declared partition. /// -/// Cloning is cheap: all fields are either `Arc`-backed or small values. +/// Cloning is cheap: all fields are `Arc`-backed. #[derive(Clone)] pub struct Client { pool: SharedPool, - partition: PartitionId, - interface: Option, - driver_spawner: Arc, + partition: Arc, } impl std::fmt::Debug for Client { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Client") - .field("partition", &self.partition) - .field("interface", &self.interface) - .field("driver_spawner", &self.driver_spawner) + .field("partition_id", &self.partition.id) + .field("nic", &self.partition.nic) .finish_non_exhaustive() } } impl Client { - /// Construct a `Client` using the default partition. - /// - /// Equivalent to `Client::builder(pool).build()`. + /// Construct a `Client` targeting the pool's default partition (the + /// first declared, or the anonymous partition when none were declared). pub fn new(pool: &SharedPool) -> Self { - Self::builder(pool).build() + let partition = pool.inner.pool.registry().default_partition(); + Self { + pool: pool.clone(), + partition, + } } - /// Return a [`ClientBuilder`] for advanced per-partition configuration. - pub fn builder(pool: &SharedPool) -> ClientBuilder { - ClientBuilder { + /// Construct a `Client` targeting a specific declared partition. + /// Panics if `id` was not declared on the pool builder (programming + /// error: the caller declared the topology). + pub fn from_partition(pool: &SharedPool, id: PartitionId) -> Self { + let partition = pool.inner.pool.registry().partition(id); + Self { pool: pool.clone(), - partition: PartitionId::default(), - interface: None, - driver_spawner: None, + partition, } } + + /// The partition id this client targets. + #[cfg(test)] + fn partition_id(&self) -> PartitionId { + self.partition.id + } } impl HttpClient for Client { @@ -104,79 +110,6 @@ impl HttpClient for Client { } } -/// Builder for configuring a [`Client`] with per-partition settings. -/// -/// Three knobs control partition behavior: -/// -/// - [`partition`](Self::partition) assigns a [`PartitionId`] (default: -/// anonymous). -/// - [`interface`](Self::interface) binds connections to a network -/// interface (default: none). -/// - [`driver_spawner`](Self::driver_spawner) selects the runtime for -/// connection driver tasks (default: captured from the current tokio -/// runtime at [`build`](Self::build) time). -pub struct ClientBuilder { - pool: SharedPool, - partition: PartitionId, - interface: Option, - driver_spawner: Option>, -} - -impl ClientBuilder { - /// Set the partition identifier (consuming builder). - pub fn partition(mut self, id: PartitionId) -> Self { - self.partition = id; - self - } - - /// Set the partition identifier (mutable reference). - pub fn set_partition(&mut self, id: PartitionId) -> &mut Self { - self.partition = id; - self - } - - /// Bind connections to a network interface (consuming builder). - pub fn interface(mut self, iface: impl Into) -> Self { - self.interface = Some(iface.into()); - self - } - - /// Set or clear the network interface binding (mutable reference). - pub fn set_interface(&mut self, iface: Option) -> &mut Self { - self.interface = iface; - self - } - - /// Set the driver spawner (consuming builder). - pub fn driver_spawner(mut self, sp: S) -> Self { - self.driver_spawner = Some(Arc::new(sp)); - self - } - - /// Set the driver spawner (mutable reference). - pub fn set_driver_spawner(&mut self, sp: S) -> &mut Self { - self.driver_spawner = Some(Arc::new(sp)); - self - } - - /// Build the [`Client`]. - /// - /// Captures the current tokio runtime handle for the driver spawner - /// if one was not explicitly provided. - pub fn build(self) -> Client { - let driver_spawner: Arc = self - .driver_spawner - .unwrap_or_else(|| Arc::new(TokioDriverSpawner::current())); - - Client { - pool: self.pool, - partition: self.partition, - interface: self.interface, - driver_spawner, - } - } -} - // --------------------------------------------------------------------------- // PooledConnector (HttpConnector adapter) // --------------------------------------------------------------------------- @@ -261,3 +194,35 @@ impl HttpConnector for PooledConnector { }) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::pool::partition::{Partition, TokioDriverSpawner}; + + #[tokio::test] + async fn client_new_uses_default_partition() { + let pool = SharedPool::builder().build_http(); + let client = Client::new(&pool); + assert_eq!(client.partition_id(), PartitionId::default()); + } + + #[tokio::test] + async fn from_partition_resolves_declared() { + let pool = SharedPool::builder() + .partitions([Partition::new( + PartitionId::from_index(3), + TokioDriverSpawner::current(), + )]) + .build_http(); + let client = Client::from_partition(&pool, PartitionId::from_index(3)); + assert_eq!(client.partition_id(), PartitionId::from_index(3)); + } + + #[tokio::test] + #[should_panic(expected = "partition not declared")] + async fn from_partition_unknown_panics() { + let pool = SharedPool::builder().build_http(); + Client::from_partition(&pool, PartitionId::from_index(99)); + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 862259e7560..e8977704cd8 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -235,6 +235,124 @@ pub enum CrossPartitionPolicy { PreferLocal, } +/// A declared pool partition: a driver-spawner runtime and an optional +/// NIC binding, identified by a caller-owned [`PartitionId`]. Declared on +/// the pool builder via `Builder::partitions`; the pool owns the topology +/// for its lifetime. +#[derive(Clone, Debug)] +pub struct Partition { + id: PartitionId, + spawner: std::sync::Arc, + nic: Option, +} + +impl Partition { + /// Declare a partition with the given id and driver spawner. + pub fn new(id: PartitionId, spawner: S) -> Self { + Self { + id, + spawner: std::sync::Arc::new(spawner), + nic: None, + } + } + + /// Bind this partition's connections to a network interface. + pub fn interface(mut self, nic: impl Into) -> Self { + self.nic = Some(nic.into()); + self + } + + // used by pool build validation (later step) + #[allow(dead_code)] + pub(crate) fn id(&self) -> PartitionId { + self.id + } +} + +/// Pool-owned state for one declared partition. Resolved once at pool +/// build time and referenced by [`Client`](super::Client) handles. +#[derive(Debug)] +pub(crate) struct PartitionState { + pub(crate) id: PartitionId, + // used by connection driver spawn (later step) + #[allow(dead_code)] + pub(crate) spawner: std::sync::Arc, + // used by socket bind (later step) + #[allow(dead_code)] + pub(crate) nic: Option, +} + +/// Immutable registry of declared partitions, built once at pool +/// construction. Maps ids and NIC groups to partition state and records +/// the default partition used by [`Client::new`](super::Client::new). +#[derive(Debug)] +pub(crate) struct PartitionRegistry { + by_id: std::collections::HashMap>, + // used by cross-partition borrow (later step) + #[allow(dead_code)] + by_nic: std::collections::HashMap, Vec>, + default_partition: PartitionId, +} + +impl PartitionRegistry { + /// Build a registry from declared partitions. When `partitions` is + /// empty, creates a single anonymous partition (`PartitionId::default()`) + /// with the given fallback spawner and no NIC binding. The default + /// partition is the first declared (or the anonymous one when none are + /// declared). Panics on a duplicate `PartitionId`. + pub(crate) fn build( + partitions: Vec, + anonymous_spawner: impl FnOnce() -> std::sync::Arc, + ) -> Self { + let partitions = if partitions.is_empty() { + vec![Partition { + id: PartitionId::default(), + spawner: anonymous_spawner(), + nic: None, + }] + } else { + partitions + }; + let default_partition = partitions[0].id; + let mut by_id = std::collections::HashMap::new(); + let mut by_nic: std::collections::HashMap, Vec> = + std::collections::HashMap::new(); + for p in partitions { + by_nic.entry(p.nic.clone()).or_default().push(p.id); + let state = std::sync::Arc::new(PartitionState { + id: p.id, + spawner: p.spawner, + nic: p.nic, + }); + if by_id.insert(p.id, state).is_some() { + panic!("duplicate PartitionId declared: {:?}", p.id); + } + } + Self { + by_id, + by_nic, + default_partition, + } + } + + /// Resolve the default partition (first declared or anonymous). + pub(crate) fn default_partition(&self) -> std::sync::Arc { + self.by_id + .get(&self.default_partition) + .expect("default partition exists") + .clone() + } + + /// Resolve a declared partition by id. Panics if the id was not + /// declared (programming error: the caller declared the topology). + pub(crate) fn partition(&self, id: PartitionId) -> std::sync::Arc { + self.by_id + .get(&id) + .unwrap_or_else(|| panic!("partition not declared: {:?}", id)) + .clone() + } +} + #[cfg(test)] mod tests { use super::*; From 744c44ef1dbcd429bdf708373674dcd59a369713 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 5 Jun 2026 07:28:33 -0400 Subject: [PATCH 30/47] move stack creation to partitions --- .../aws-smithy-http-client/src/client/pool.rs | 230 +++++++++++++----- .../src/client/pool/builder.rs | 30 +-- .../src/client/pool/client.rs | 5 +- .../src/client/pool/partition.rs | 27 +- 4 files changed, 203 insertions(+), 89 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index befd233b510..26eade35d11 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -187,13 +187,13 @@ impl SharedPool { /// Key for per-host connection pool routing. #[derive(Clone, Debug, PartialEq, Eq, Hash)] -struct PoolKey { +pub(crate) struct PoolKey { scheme: http_1x::uri::Scheme, authority: http_1x::uri::Authority, } impl PoolKey { - fn from_uri(uri: &http_1x::Uri) -> Option { + pub(crate) fn from_uri(uri: &http_1x::Uri) -> Option { Some(Self { scheme: uri.scheme()?.clone(), authority: uri.authority()?.clone(), @@ -232,10 +232,51 @@ impl BoxedRetainer { /// `layer_fn`s the first time `Negotiate` constructs each leg. type RetainerSlot = Arc>>; +/// Per-partition stack factory: builds a host's Negotiate(Cache,Singleton) +/// entry on first touch. Captures the partition's connector; the shared +/// budget/hooks arrive via `&SharedPoolState` at call time. +pub(crate) type MakeStack = + Arc Box + Send + Sync>; + +/// Connection-lifecycle machinery shared across every partition: event +/// hooks, the global connection budget, and the per-host budgets. One +/// instance per pool, held behind `Arc` on `ConnectionPool`. +pub(crate) struct SharedPoolState { + pub(crate) hooks: handshake::PoolHooks, + pub(crate) global_sem: Option>, + max_connections_per_host: Option, + per_host_sems: Mutex>>, +} + +impl SharedPoolState { + pub(crate) fn new(config: &PoolConfig) -> Self { + Self { + hooks: handshake::PoolHooks::new(config.connection_event_listener.clone()), + global_sem: config.max_connections.map(|n| Arc::new(Semaphore::new(n))), + max_connections_per_host: config.max_connections_per_host, + per_host_sems: Mutex::new(HashMap::new()), + } + } + + /// The shared per-host semaphore for `key`, created on first request + /// for that key. Returns `None` when no per-host limit is configured. + /// Shared across all partitions: two partitions to the same host + /// contend on the same budget. + pub(crate) fn per_host_sem(&self, key: &PoolKey) -> Option> { + let n = self.max_connections_per_host?; + let mut map = self.per_host_sems.lock().expect("per_host_sems poisoned"); + Some( + map.entry(key.clone()) + .or_insert_with(|| Arc::new(Semaphore::new(n))) + .clone(), + ) + } +} + pub(crate) fn build_pool( connector: C, config: PoolConfig, - registry: Arc, + partitions: Vec, cross_partition_policy: partition::CrossPartitionPolicy, ) -> ConnectionPool where @@ -244,9 +285,7 @@ where C::Future: Unpin + Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + HyperConnection + Unpin + Send + 'static, { - let global_sem = config.max_connections.map(|n| Arc::new(Semaphore::new(n))); - let max_per_host = config.max_connections_per_host; - let pool_hooks = handshake::PoolHooks::new(config.connection_event_listener.clone()); + let shared = Arc::new(SharedPoolState::new(&config)); tracing::debug!( max_connections = ?config.max_connections, @@ -255,21 +294,23 @@ where "pool: initialized" ); - let make_entry = { - let pool_hooks = pool_hooks.clone(); - move |uri: &http_1x::Uri| -> Box { + let make_stack: MakeStack = Arc::new( + move |uri: &http_1x::Uri, shared: &SharedPoolState| -> Box { let authority = Authority::new( uri.authority() .expect("pool entry URI has authority") .as_str(), ); - let per_host_sem = max_per_host.map(|n| Arc::new(Semaphore::new(n))); + let key = PoolKey::from_uri(uri).expect("pool entry URI has scheme+authority"); + let per_host_sem = shared.per_host_sem(&key); let limited = handshake::ConnectionLimit::new( connector.clone(), - global_sem.clone(), + shared.global_sem.clone(), per_host_sem, ); + let pool_hooks = shared.hooks.clone(); + // Per-host bridge: `H2ConnectAndHandshake` publishes on each new // handshake; `SingletonConnection` reads the current entry at // checkout. Clones share the underlying slot. @@ -456,13 +497,18 @@ where }) .build(); Box::new(TypedPoolEntry { stack, retainers }) - } - }; + }, + ); + + let registry = Arc::new(partition::PartitionRegistry::build( + partitions, + || Arc::new(partition::TokioDriverSpawner::current()), + make_stack, + )); ConnectionPool { config, - hosts: Mutex::new(HashMap::new()), - make_entry: Box::new(make_entry), + shared, eviction_spawned: AtomicBool::new(false), drop_notifier: OnceLock::new(), registry, @@ -478,13 +524,14 @@ where pub(crate) struct ConnectionPool { /// Pool-wide configuration. config: PoolConfig, - hosts: Mutex>>, - make_entry: Box Box + Send + Sync>, + + /// Shared connection-lifecycle machinery (hooks, semaphores). + shared: Arc, /// Immutable partition registry, resolved at build time. registry: Arc, - // stored for later steps (send path unchanged this step) + // used by cross-partition checkout #[allow(dead_code)] cross_partition_policy: partition::CrossPartitionPolicy, @@ -516,11 +563,12 @@ impl ConnectionPool { /// Send a request through the pool. /// - /// Routes to the appropriate per-host pool stack and sends the request. - /// `ctx` carries the routing URI plus per-operation connect-time data - /// (connect_timeout). Per-operation read_timeout, if any, must be - /// attached to `req.extensions_mut()` as [`ReadTimeoutHint`] before - /// calling; the checkout services read it from there. + /// Routes to the appropriate per-host pool stack in the given partition + /// and sends the request. `ctx` carries the routing URI plus + /// per-operation connect-time data (connect_timeout). Per-operation + /// read_timeout, if any, must be attached to `req.extensions_mut()` as + /// [`ReadTimeoutHint`] before calling; the checkout services read it + /// from there. /// /// Takes `self: &Arc` so the lazy eviction task can hold a /// `Weak` as a fallback exit signal (primary exit is the drop @@ -528,6 +576,7 @@ impl ConnectionPool { /// missed drop signal). pub(crate) async fn send_request( self: &Arc, + partition: &Arc, ctx: ConnectCtx, req: http_1x::Request, ) -> Result, BoxError> { @@ -544,11 +593,12 @@ impl ConnectionPool { // `send`; all I/O happens in the returned future, outside the // lock. let fut = { - let mut hosts = self.hosts.lock().unwrap(); - if !hosts.contains_key(&key) { - hosts.insert(key.clone(), (self.make_entry)(&ctx.uri)); + let mut auth = partition.authorities.lock().unwrap(); + if !auth.contains_key(&key) { + let entry = (partition.make_stack)(&ctx.uri, &self.shared); + auth.insert(key.clone(), entry); } - hosts.get_mut(&key).unwrap().send(ctx, req) + auth.get_mut(&key).unwrap().send(ctx, req) }; fut.await @@ -595,34 +645,24 @@ impl ConnectionPool { tokio::spawn(eviction_task(weak, rx, timeout)); } - /// Walk the host map, dropping idle connections that have exceeded + /// Walk all partitions, dropping idle connections that have exceeded /// `timeout` and removing host entries whose retainers are empty. /// - /// Called by the eviction task on each tick. Safe to call when the - /// host map is empty (no-op). - fn retain_idle_hosts(&self, timeout: Duration) { - let mut hosts = self.hosts.lock().unwrap(); - let before = hosts.len(); - hosts.retain(|key, entry| { - entry.retain_idle(timeout); - if entry.is_empty() { - tracing::debug!( - authority = %key.authority, - scheme = %key.scheme, - "pool: host entry removed (empty after retain)" - ); - false - } else { - true - } - }); - let after = hosts.len(); - drop(hosts); - tracing::trace!( - hosts_before = before, - hosts_after = after, - "pool: eviction tick complete" - ); + /// Called by the eviction task on each tick. Safe to call when + /// partitions are empty (no-op). + fn retain_idle(&self, timeout: Duration) { + for partition in self.registry.partitions() { + let mut auth = partition.authorities.lock().unwrap(); + auth.retain(|_key, entry| { + entry.retain_idle(timeout); + if entry.is_empty() { + tracing::debug!("pool: host entry removed (empty after retain)"); + false + } else { + true + } + }); + } } } @@ -662,7 +702,7 @@ async fn eviction_task( _ = &mut drop_notifier => break, _ = interval.tick() => { match pool.upgrade() { - Some(pool) => pool.retain_idle_hosts(timeout), + Some(pool) => pool.retain_idle(timeout), None => break, } } @@ -873,7 +913,7 @@ where /// /// `send` performs checkout, health filtering, request dispatch, and /// body-guard setup atomically for a single request. -trait PoolEntry: Send + Sync { +pub(crate) trait PoolEntry: Send + Sync { fn send( &mut self, ctx: ConnectCtx, @@ -1015,16 +1055,21 @@ mod tests { } fn pool_with_config(config: PoolConfig) -> Arc { + let shared = Arc::new(SharedPoolState::new(&config)); + let make_stack: MakeStack = + Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box); + let registry = Arc::new(partition::PartitionRegistry::build( + Vec::new(), + || Arc::new(partition::TokioDriverSpawner::current()), + make_stack, + )); Arc::new(ConnectionPool { config, - hosts: Mutex::new(HashMap::new()), - make_entry: Box::new(|_| Box::new(NullPoolEntry)), + shared, + registry, + cross_partition_policy: CrossPartitionPolicy::default(), eviction_spawned: AtomicBool::new(false), drop_notifier: OnceLock::new(), - registry: Arc::new(partition::PartitionRegistry::build(Vec::new(), || { - Arc::new(partition::TokioDriverSpawner::current()) - })), - cross_partition_policy: CrossPartitionPolicy::default(), }) } @@ -1075,7 +1120,7 @@ mod tests { ); } - /// `retain_idle_hosts` drops host entries whose retainers report empty. + /// `retain_idle` drops host entries whose retainers report empty. /// Uses `NullPoolEntry` whose `is_empty` returns `true`, so the first /// tick should remove every entry. #[tokio::test] @@ -1083,16 +1128,69 @@ mod tests { let pool = pool_with_config(PoolConfig::default()); let uri: http_1x::Uri = "http://example.com".parse().unwrap(); let key = PoolKey::from_uri(&uri).unwrap(); - pool.hosts + let partition = pool.registry().default_partition(); + partition + .authorities .lock() .unwrap() .insert(key.clone(), Box::new(NullPoolEntry)); - assert_eq!(pool.hosts.lock().unwrap().len(), 1); - pool.retain_idle_hosts(Duration::from_secs(30)); + assert_eq!(partition.authorities.lock().unwrap().len(), 1); + pool.retain_idle(Duration::from_secs(30)); assert_eq!( - pool.hosts.lock().unwrap().len(), + partition.authorities.lock().unwrap().len(), 0, "empty entry should have been evicted" ); } + + #[test] + fn per_host_sem_shared_by_key() { + let cfg = PoolConfig { + max_connections_per_host: Some(2), + ..PoolConfig::default() + }; + let s = SharedPoolState::new(&cfg); + let uri: http_1x::Uri = "https://example.com".parse().unwrap(); + let key = PoolKey::from_uri(&uri).unwrap(); + let a = s.per_host_sem(&key).unwrap(); + let b = s.per_host_sem(&key).unwrap(); + assert!(Arc::ptr_eq(&a, &b), "same key must share one semaphore"); + let uri2: http_1x::Uri = "https://other.com".parse().unwrap(); + let key2 = PoolKey::from_uri(&uri2).unwrap(); + let c = s.per_host_sem(&key2).unwrap(); + assert!( + !Arc::ptr_eq(&a, &c), + "different key must get a distinct semaphore" + ); + } + + #[tokio::test] + async fn two_partitions_independent_storage() { + let pool = SharedPool::builder() + .partitions([ + partition::Partition::new( + partition::PartitionId::from_index(0), + partition::TokioDriverSpawner::current(), + ), + partition::Partition::new( + partition::PartitionId::from_index(1), + partition::TokioDriverSpawner::current(), + ), + ]) + .build_http(); + let p0 = pool + .inner + .pool + .registry() + .partition(partition::PartitionId::from_index(0)); + let p1 = pool + .inner + .pool + .registry() + .partition(partition::PartitionId::from_index(1)); + assert!( + !Arc::ptr_eq(&p0, &p1), + "distinct partition ids must yield distinct PartitionState arcs" + ); + } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index 0c87badffbe..603fc462a84 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -237,21 +237,18 @@ impl Builder { connection_event_listener: self.connection_event_listener.clone(), }; let proxy_matcher = proxy_matcher_from(&self.proxy_config); - let registry = Arc::new(super::partition::PartitionRegistry::build( - std::mem::take(&mut self.partitions), - || Arc::new(super::partition::TokioDriverSpawner::current()), - )); + let partitions = std::mem::take(&mut self.partitions); let policy = self.cross_partition_policy; let pool = match dns_resolver { Some(resolver) => { let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); tcp.set_nodelay(self.tcp_nodelay); - build_http_pool_with_proxy(tcp, &self.proxy_config, config, registry, policy) + build_http_pool_with_proxy(tcp, &self.proxy_config, config, partitions, policy) } None => { let mut tcp = HyperHttpConnector::new(); tcp.set_nodelay(self.tcp_nodelay); - build_http_pool_with_proxy(tcp, &self.proxy_config, config, registry, policy) + build_http_pool_with_proxy(tcp, &self.proxy_config, config, partitions, policy) } }; super::SharedPool { @@ -288,12 +285,8 @@ impl Builder { pool_idle_timeout: self.pool_idle_timeout, connection_event_listener: self.connection_event_listener.clone(), }; - let registry = Arc::new(super::partition::PartitionRegistry::build( - self.partitions, - || Arc::new(super::partition::TokioDriverSpawner::current()), - )); let policy = self.cross_partition_policy; - let pool = super::build_pool(connector, config, registry, policy); + let pool = super::build_pool(connector, config, self.partitions, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), @@ -311,7 +304,7 @@ fn build_http_pool_with_proxy( tcp: HyperHttpConnector, proxy_config: &Option, config: PoolConfig, - registry: Arc, + partitions: Vec, policy: CrossPartitionPolicy, ) -> ConnectionPool where @@ -332,10 +325,10 @@ where } if proxy_config.is_disabled() { - super::build_pool(tcp, config, registry, policy) + super::build_pool(tcp, config, partitions, policy) } else { let connector = crate::client::connect::HttpProxyConnector::new(tcp, proxy_config); - super::build_pool(connector, config, registry, policy) + super::build_pool(connector, config, partitions, policy) } } @@ -385,10 +378,7 @@ impl Builder { .clone() .unwrap_or_else(ProxyConfig::disabled); let proxy_matcher = proxy_matcher_from(&self.proxy_config); - let registry = Arc::new(super::partition::PartitionRegistry::build( - self.partitions, - || Arc::new(super::partition::TokioDriverSpawner::current()), - )); + let partitions = self.partitions; let policy = self.cross_partition_policy; match &self.tls.provider { @@ -404,7 +394,7 @@ impl Builder { &self.tls.context, proxy_config, ); - let pool = super::build_pool(connector, config, registry, policy); + let pool = super::build_pool(connector, config, partitions, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), @@ -419,7 +409,7 @@ impl Builder { &self.tls.context, proxy_config, ); - let pool = super::build_pool(connector, config, registry, policy); + let pool = super::build_pool(connector, config, partitions, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs index 373a7232149..62848afbbc3 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/client.rs @@ -98,6 +98,7 @@ impl HttpClient for Client { SharedHttpConnector::new(PooledConnector { pool: self.pool.inner.pool.clone(), + partition: self.partition.clone(), connect_timeout, read_timeout, sleep_impl, @@ -121,6 +122,7 @@ impl HttpClient for Client { /// timeouts). The pool itself is shared across all operations via `Arc`. struct PooledConnector { pool: Arc, + partition: Arc, connect_timeout: Option, read_timeout: Option, sleep_impl: Option, @@ -136,6 +138,7 @@ impl std::fmt::Debug for PooledConnector { impl HttpConnector for PooledConnector { fn call(&self, request: HttpRequest) -> HttpConnectorFuture { let pool = self.pool.clone(); + let partition = self.partition.clone(); let connect_timeout = self.connect_timeout; let read_timeout = self.read_timeout; let sleep_impl = self.sleep_impl.clone(); @@ -186,7 +189,7 @@ impl HttpConnector for PooledConnector { ); let response = pool - .send_request(connect_ctx, request) + .send_request(&partition, connect_ctx, request) .await .map_err(downcast_error)?; diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index e8977704cd8..2aff25d0a20 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -271,15 +271,30 @@ impl Partition { /// Pool-owned state for one declared partition. Resolved once at pool /// build time and referenced by [`Client`](super::Client) handles. -#[derive(Debug)] pub(crate) struct PartitionState { pub(crate) id: PartitionId, - // used by connection driver spawn (later step) + // used by connection driver spawn #[allow(dead_code)] pub(crate) spawner: std::sync::Arc, // used by socket bind (later step) #[allow(dead_code)] pub(crate) nic: Option, + /// Per-host connection storage for this partition. Keyed by + /// (scheme, authority); entries built lazily on first request. + pub(crate) authorities: + std::sync::Mutex>>, + /// Builds a host entry on first touch, capturing this partition's + /// connector; shared budget/hooks arrive via `&SharedPoolState`. + pub(crate) make_stack: super::MakeStack, +} + +impl std::fmt::Debug for PartitionState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PartitionState") + .field("id", &self.id) + .field("nic", &self.nic) + .finish_non_exhaustive() + } } /// Immutable registry of declared partitions, built once at pool @@ -303,6 +318,7 @@ impl PartitionRegistry { pub(crate) fn build( partitions: Vec, anonymous_spawner: impl FnOnce() -> std::sync::Arc, + make_stack: super::MakeStack, ) -> Self { let partitions = if partitions.is_empty() { vec![Partition { @@ -323,6 +339,8 @@ impl PartitionRegistry { id: p.id, spawner: p.spawner, nic: p.nic, + authorities: std::sync::Mutex::new(std::collections::HashMap::new()), + make_stack: make_stack.clone(), }); if by_id.insert(p.id, state).is_some() { panic!("duplicate PartitionId declared: {:?}", p.id); @@ -351,6 +369,11 @@ impl PartitionRegistry { .unwrap_or_else(|| panic!("partition not declared: {:?}", id)) .clone() } + + /// Iterate all declared partitions. + pub(crate) fn partitions(&self) -> impl Iterator> { + self.by_id.values() + } } #[cfg(test)] From 7137a57a4de020c01305d14bce001c57205ea2d8 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 5 Jun 2026 10:04:30 -0400 Subject: [PATCH 31/47] wire driver spawner through partitions --- .../aws-smithy-http-client/src/client/pool.rs | 436 ++++++++++-------- .../src/client/pool/handshake.rs | 39 +- .../src/client/pool/partition.rs | 56 ++- .../tests/pool_behavior_test.rs | 108 +++++ 4 files changed, 400 insertions(+), 239 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 26eade35d11..4b5236c4831 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -294,216 +294,239 @@ where "pool: initialized" ); - let make_stack: MakeStack = Arc::new( - move |uri: &http_1x::Uri, shared: &SharedPoolState| -> Box { - let authority = Authority::new( - uri.authority() - .expect("pool entry URI has authority") - .as_str(), - ); - let key = PoolKey::from_uri(uri).expect("pool entry URI has scheme+authority"); - let per_host_sem = shared.per_host_sem(&key); - let limited = handshake::ConnectionLimit::new( - connector.clone(), - shared.global_sem.clone(), - per_host_sem, - ); - - let pool_hooks = shared.hooks.clone(); - - // Per-host bridge: `H2ConnectAndHandshake` publishes on each new - // handshake; `SingletonConnection` reads the current entry at - // checkout. Clones share the underlying slot. - let h2_ref = H2ConnectionRef::new(); - - // Bounded at two: the H1 fallback and H2 upgrade `layer_fn`s - // each push one entry on first construction. - let retainers: RetainerSlot = Arc::new(Mutex::new(Vec::with_capacity(2))); - - let stack = hpool::negotiate::builder() - .connect(limited) - .inspect(|established: &connection::EstablishedConnection| { - established.io.connected().is_negotiated_h2() - }) - .fallback({ - let retainers = retainers.clone(); - let pool_hooks = pool_hooks.clone(); - tower::layer::layer_fn(move |inspector| { - let cache = cache::builder() - .executor(TokioExecutor::new()) - .build(H1ConnectAndHandshake::new(inspector, pool_hooks.clone())); - // Capture a clone of the Cache for eviction. `Cache` - // is Clone and shares state via its internal - // `Arc>`, so the clone here and the - // original-stack consumption below observe the same - // idle set. - let cache_for_retain = Arc::new(std::sync::Mutex::new(cache.clone())); - let cache_for_empty = cache_for_retain.clone(); - retainers - .lock() - .expect("retainer slot poisoned") - .push(BoxedRetainer { - retain_fn: Box::new({ - let listener = pool_hooks.listener.clone(); - move |timeout| { - let now = Instant::now(); - cache_for_retain - .lock() - .expect("retain cache lock poisoned") - .retain(|managed| { - let keep = !managed.is_poisoned() - && now.saturating_duration_since( - managed.idle_at(), - ) < timeout; - if !keep { - let reason = if managed.is_poisoned() { - "poisoned" - } else { - "idle_expired" - }; - let idle_duration = now - .saturating_duration_since( + let make_stack_for = { + let connector = connector.clone(); + move |spawner: &Arc| -> MakeStack { + let connector = connector.clone(); + let spawner = spawner.clone(); + Arc::new( + move |uri: &http_1x::Uri, shared: &SharedPoolState| -> Box { + let authority = Authority::new( + uri.authority() + .expect("pool entry URI has authority") + .as_str(), + ); + let key = PoolKey::from_uri(uri).expect("pool entry URI has scheme+authority"); + let per_host_sem = shared.per_host_sem(&key); + let limited = handshake::ConnectionLimit::new( + connector.clone(), + shared.global_sem.clone(), + per_host_sem, + ); + + let pool_hooks = shared.hooks.clone(); + + // Per-host bridge: `H2ConnectAndHandshake` publishes on each new + // handshake; `SingletonConnection` reads the current entry at + // checkout. Clones share the underlying slot. + let h2_ref = H2ConnectionRef::new(); + + // Bounded at two: the H1 fallback and H2 upgrade `layer_fn`s + // each push one entry on first construction. + let retainers: RetainerSlot = Arc::new(Mutex::new(Vec::with_capacity(2))); + + let stack = + hpool::negotiate::builder() + .connect(limited) + .inspect(|established: &connection::EstablishedConnection| { + established.io.connected().is_negotiated_h2() + }) + .fallback({ + let retainers = retainers.clone(); + let pool_hooks = pool_hooks.clone(); + let spawner = spawner.clone(); + tower::layer::layer_fn(move |inspector| { + let cache = cache::builder() + .executor(TokioExecutor::new()) + .build(H1ConnectAndHandshake::new( + inspector, + pool_hooks.clone(), + spawner.clone(), + )); + // Capture a clone of the Cache for eviction. `Cache` + // is Clone and shares state via its internal + // `Arc>`, so the clone here and the + // original-stack consumption below observe the same + // idle set. + let cache_for_retain = + Arc::new(std::sync::Mutex::new(cache.clone())); + let cache_for_empty = cache_for_retain.clone(); + retainers.lock().expect("retainer slot poisoned").push( + BoxedRetainer { + retain_fn: Box::new({ + let listener = pool_hooks.listener.clone(); + move |timeout| { + let now = Instant::now(); + cache_for_retain + .lock() + .expect("retain cache lock poisoned") + .retain(|managed| { + let keep = !managed.is_poisoned() + && now.saturating_duration_since( managed.idle_at(), + ) < timeout; + if !keep { + let reason = if managed.is_poisoned() { + "poisoned" + } else { + "idle_expired" + }; + let idle_duration = now + .saturating_duration_since( + managed.idle_at(), + ); + tracing::debug!( + conn_id = %managed.conn_id(), + reason, + ?idle_duration, + "pool: connection evicted" ); - tracing::debug!( - conn_id = %managed.conn_id(), - reason, - ?idle_duration, - "pool: connection evicted" - ); - if let Some(ref l) = listener { - l.on_closed(&ConnectionClosedEvent::new( - managed.conn_id(), - managed.info.authority.clone(), - managed.info.remote_addr, - if managed.is_poisoned() { - CloseReason::Poisoned - } else { - CloseReason::IdleTimeout - }, - None, - )); + if let Some(ref l) = listener { + l.on_closed(&ConnectionClosedEvent::new( + managed.conn_id(), + managed.info.authority.clone(), + managed.info.remote_addr, + if managed.is_poisoned() { + CloseReason::Poisoned + } else { + CloseReason::IdleTimeout + }, + None, + )); + } } + keep + }); } - keep - }); - } - }), - is_empty_fn: Box::new(move || { - cache_for_empty - .lock() - .expect("is_empty cache lock poisoned") - .is_empty() - }), - }); - cache.map_response({ - let listener = pool_hooks.listener.clone(); - move |cached| { - H1Checkout::new(CachedConnection::new(cached, listener.clone())) - } - }) - }) - }) - .upgrade({ - let h2_ref = h2_ref.clone(); - let retainers = retainers.clone(); - let pool_hooks = pool_hooks.clone(); - tower::layer::layer_fn(move |inspected| { - let singleton = - hpool::singleton::Singleton::new(H2ConnectAndHandshake::new( - inspected, - h2_ref.clone(), - pool_hooks.clone(), - authority.clone(), - )); - let singleton_for_retain = - Arc::new(std::sync::Mutex::new(singleton.clone())); - let singleton_for_empty = singleton_for_retain.clone(); - retainers - .lock() - .expect("retainer slot poisoned") - .push(BoxedRetainer { - retain_fn: Box::new({ - let listener = pool_hooks.listener.clone(); - move |timeout| { - let now = Instant::now(); - singleton_for_retain - .lock() - .expect("retain singleton lock poisoned") - .retain(|managed| { - // For H2 the connection stays in - // Singleton while serving streams, so - // `idle_at` alone is not a valid - // idleness signal. Keep the - // connection if any stream is in - // flight; only evict when truly idle - // AND `idle_at` has exceeded the - // timeout (stamped on the 1 → 0 - // transition by - // `SingletonConnection::drop`). - let keep = !managed.is_poisoned() - && (managed.active_streams_count() > 0 - || now.saturating_duration_since( - managed.idle_at(), - ) < timeout); - if !keep { - let reason = if managed.is_poisoned() { - "poisoned" - } else { - "idle_expired" - }; - let idle_duration = now - .saturating_duration_since( - managed.idle_at(), + }), + is_empty_fn: Box::new(move || { + cache_for_empty + .lock() + .expect("is_empty cache lock poisoned") + .is_empty() + }), + }, + ); + cache.map_response({ + let listener = pool_hooks.listener.clone(); + move |cached| { + H1Checkout::new(CachedConnection::new( + cached, + listener.clone(), + )) + } + }) + }) + }) + .upgrade({ + let h2_ref = h2_ref.clone(); + let retainers = retainers.clone(); + let pool_hooks = pool_hooks.clone(); + let spawner = spawner.clone(); + tower::layer::layer_fn(move |inspected| { + let singleton = hpool::singleton::Singleton::new( + H2ConnectAndHandshake::new( + inspected, + h2_ref.clone(), + pool_hooks.clone(), + authority.clone(), + spawner.clone(), + ), + ); + let singleton_for_retain = + Arc::new(std::sync::Mutex::new(singleton.clone())); + let singleton_for_empty = singleton_for_retain.clone(); + retainers.lock().expect("retainer slot poisoned").push( + BoxedRetainer { + retain_fn: Box::new({ + let listener = pool_hooks.listener.clone(); + move |timeout| { + let now = Instant::now(); + singleton_for_retain + .lock() + .expect("retain singleton lock poisoned") + .retain(|managed| { + // For H2 the connection stays in + // Singleton while serving streams, so + // `idle_at` alone is not a valid + // idleness signal. Keep the + // connection if any stream is in + // flight; only evict when truly idle + // AND `idle_at` has exceeded the + // timeout (stamped on the 1 → 0 + // transition by + // `SingletonConnection::drop`). + let keep = !managed.is_poisoned() + && (managed.active_streams_count() > 0 + || now.saturating_duration_since( + managed.idle_at(), + ) < timeout); + if !keep { + let reason = if managed.is_poisoned() { + "poisoned" + } else { + "idle_expired" + }; + let idle_duration = now + .saturating_duration_since( + managed.idle_at(), + ); + tracing::debug!( + conn_id = %managed.conn_id(), + reason, + ?idle_duration, + "pool: connection evicted" ); - tracing::debug!( - conn_id = %managed.conn_id(), - reason, - ?idle_duration, - "pool: connection evicted" - ); - if let Some(ref l) = listener { - l.on_closed(&ConnectionClosedEvent::new( - managed.conn_id(), - managed.info.authority.clone(), - managed.info.remote_addr, - if managed.is_poisoned() { - CloseReason::Poisoned - } else { - CloseReason::IdleTimeout - }, - None, - )); + if let Some(ref l) = listener { + l.on_closed(&ConnectionClosedEvent::new( + managed.conn_id(), + managed.info.authority.clone(), + managed.info.remote_addr, + if managed.is_poisoned() { + CloseReason::Poisoned + } else { + CloseReason::IdleTimeout + }, + None, + )); + } } + keep + }); } - keep - }); - } - }), - is_empty_fn: Box::new(move || { - singleton_for_empty - .lock() - .expect("is_empty singleton lock poisoned") - .is_empty() - }), - }); - singleton.map_response({ - let h2_ref = h2_ref.clone(); - move |singled| { - H2Checkout::new(SingletonConnection::new(singled, h2_ref.clone())) - } - }) - }) - }) - .build(); - Box::new(TypedPoolEntry { stack, retainers }) - }, - ); + }), + is_empty_fn: Box::new(move || { + singleton_for_empty + .lock() + .expect("is_empty singleton lock poisoned") + .is_empty() + }), + }, + ); + singleton.map_response({ + let h2_ref = h2_ref.clone(); + move |singled| { + H2Checkout::new(SingletonConnection::new( + singled, + h2_ref.clone(), + )) + } + }) + }) + }) + .build(); + Box::new(TypedPoolEntry { stack, retainers }) + }, + ) + } + }; + let partitions = partition::normalize_partitions(partitions, || { + Arc::new(partition::TokioDriverSpawner::current()) + }); let registry = Arc::new(partition::PartitionRegistry::build( partitions, - || Arc::new(partition::TokioDriverSpawner::current()), - make_stack, + make_stack_for, )); ConnectionPool { @@ -1056,12 +1079,15 @@ mod tests { fn pool_with_config(config: PoolConfig) -> Arc { let shared = Arc::new(SharedPoolState::new(&config)); - let make_stack: MakeStack = - Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box); + let make_stack_for = |_spawner: &Arc| -> MakeStack { + Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box) + }; + let partitions = partition::normalize_partitions(Vec::new(), || { + Arc::new(partition::TokioDriverSpawner::current()) + }); let registry = Arc::new(partition::PartitionRegistry::build( - Vec::new(), - || Arc::new(partition::TokioDriverSpawner::current()), - make_stack, + partitions, + make_stack_for, )); Arc::new(ConnectionPool { config, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 7a1cafd943d..9a0fadb0c79 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -19,7 +19,6 @@ use std::time::Instant; use aws_smithy_runtime_api::box_error::BoxError; use aws_smithy_runtime_api::client::connection::ConnectionId; use aws_smithy_types::body::SdkBody; -use hyper::rt::Executor; use hyper_util::client::legacy::connect::{Connection, HttpInfo}; use hyper_util::rt::TokioExecutor; use tokio::sync::Semaphore; @@ -30,6 +29,7 @@ use super::connection::{ ConnectionPermit, ConnectionTiming, EstablishedConnection, ManagedConnection, NegotiatedProtocol, }; +use super::partition::DriverSpawner; /// Pool-scoped instrumentation primitives shared across layers in the /// pool stack. Held by `ConnectionPool` for the pool's lifetime and @@ -259,25 +259,24 @@ fn capture_info(io: &IO, authority: Authority) -> ConnectionInfo } } -/// Spawn a connection driver future onto the runtime. -/// -// TODO - revisit tokio default without flags -fn spawn_driver(future: impl Future + Send + 'static) { - TokioExecutor::new().execute(future); -} - -/// Connects and performs an HTTP/1.1 handshake. +/// Connects and performs an HTTP/1.1 handshake, spawning the connection +/// driver onto the partition's runtime via the captured [`DriverSpawner`]. /// /// The connector is expected to return an [`EstablishedConnection`], typically /// produced by [`ConnectionLimit`] wrapping a TCP/TLS connector. pub(crate) struct H1ConnectAndHandshake { connector: C, hooks: PoolHooks, + spawner: Arc, } impl H1ConnectAndHandshake { - pub(crate) fn new(connector: C, hooks: PoolHooks) -> Self { - Self { connector, hooks } + pub(crate) fn new(connector: C, hooks: PoolHooks, spawner: Arc) -> Self { + Self { + connector, + hooks, + spawner, + } } } @@ -286,6 +285,7 @@ impl Clone for H1ConnectAndHandshake { Self { connector: self.connector.clone(), hooks: self.hooks.clone(), + spawner: self.spawner.clone(), } } } @@ -307,6 +307,7 @@ where fn call(&mut self, ctx: ConnectCtx) -> Self::Future { let hooks = self.hooks.clone(); + let spawner = self.spawner.clone(); let authority = Authority::new( ctx.uri .authority() @@ -363,7 +364,7 @@ where ConnectionTiming::new(connect_duration), )); - spawn_driver({ + spawner.spawn(Box::pin({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { @@ -378,7 +379,7 @@ where ); } } - }); + })); Ok(ManagedConnection::new( H1SendRequest::new(tx), @@ -390,7 +391,8 @@ where } } -/// Connects and performs an HTTP/2 handshake. +/// Connects and performs an HTTP/2 handshake, spawning the connection +/// driver onto the partition's runtime via the captured [`DriverSpawner`]. /// /// Pinned to `Service<()>` because this service sits in the Negotiate /// upgrade path: the connection is already established via the shared @@ -400,6 +402,7 @@ pub(crate) struct H2ConnectAndHandshake { h2_ref: super::connection::H2ConnectionRef, hooks: PoolHooks, authority: Authority, + spawner: Arc, } impl H2ConnectAndHandshake { @@ -414,12 +417,14 @@ impl H2ConnectAndHandshake { h2_ref: super::connection::H2ConnectionRef, hooks: PoolHooks, authority: Authority, + spawner: Arc, ) -> Self { Self { connector, h2_ref, hooks, authority, + spawner, } } } @@ -431,6 +436,7 @@ impl Clone for H2ConnectAndHandshake { h2_ref: self.h2_ref.clone(), hooks: self.hooks.clone(), authority: self.authority.clone(), + spawner: self.spawner.clone(), } } } @@ -455,6 +461,7 @@ where let h2_ref = self.h2_ref.clone(); let hooks = self.hooks.clone(); let authority = self.authority.clone(); + let spawner = self.spawner.clone(); Box::pin(async move { let connect_start = Instant::now(); let EstablishedConnection { io, permit } = match fut.await.map_err(Into::into) { @@ -504,7 +511,7 @@ where ConnectionTiming::new(connect_duration), )); - spawn_driver({ + spawner.spawn(Box::pin({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; async move { @@ -519,7 +526,7 @@ where ); } } - }); + })); let managed = ManagedConnection::new(H2SendRequest::new(tx), info, conn_id, permit); // Publish this connection's state (metadata + active_streams + diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 2aff25d0a20..267d51a10f8 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -273,7 +273,7 @@ impl Partition { /// build time and referenced by [`Client`](super::Client) handles. pub(crate) struct PartitionState { pub(crate) id: PartitionId, - // used by connection driver spawn + // captured into make_stack via the factory; field retained for later NIC/runtime use #[allow(dead_code)] pub(crate) spawner: std::sync::Arc, // used by socket bind (later step) @@ -297,6 +297,31 @@ impl std::fmt::Debug for PartitionState { } } +/// Normalize a caller-declared partition set for pool construction: when +/// no partitions are declared, synthesize a single anonymous partition +/// (`PartitionId::default()`, no NIC binding) using `anonymous_spawner`. +/// Returns the caller's set unchanged when non-empty. +/// +/// This is the one place the "no topology declared" default is decided; +/// [`PartitionRegistry::build`] then indexes whatever set it is given. +/// `anonymous_spawner` is a closure so the runtime handle is captured +/// only when actually needed (e.g. `TokioDriverSpawner::current()` panics +/// off a runtime). +pub(crate) fn normalize_partitions( + partitions: Vec, + anonymous_spawner: impl FnOnce() -> std::sync::Arc, +) -> Vec { + if partitions.is_empty() { + vec![Partition { + id: PartitionId::default(), + spawner: anonymous_spawner(), + nic: None, + }] + } else { + partitions + } +} + /// Immutable registry of declared partitions, built once at pool /// construction. Maps ids and NIC groups to partition state and records /// the default partition used by [`Client::new`](super::Client::new). @@ -310,37 +335,32 @@ pub(crate) struct PartitionRegistry { } impl PartitionRegistry { - /// Build a registry from declared partitions. When `partitions` is - /// empty, creates a single anonymous partition (`PartitionId::default()`) - /// with the given fallback spawner and no NIC binding. The default - /// partition is the first declared (or the anonymous one when none are - /// declared). Panics on a duplicate `PartitionId`. + /// Build a registry from a non-empty set of declared partitions. The + /// default partition is the first in the slice. Panics on a duplicate + /// `PartitionId`, or if `partitions` is empty (callers normalize the + /// no-topology case via [`normalize_partitions`] first). pub(crate) fn build( partitions: Vec, - anonymous_spawner: impl FnOnce() -> std::sync::Arc, - make_stack: super::MakeStack, + make_stack_for: impl Fn(&std::sync::Arc) -> super::MakeStack, ) -> Self { - let partitions = if partitions.is_empty() { - vec![Partition { - id: PartitionId::default(), - spawner: anonymous_spawner(), - nic: None, - }] - } else { - partitions - }; + assert!( + !partitions.is_empty(), + "PartitionRegistry::build requires at least one partition; \ + normalize the empty case with normalize_partitions" + ); let default_partition = partitions[0].id; let mut by_id = std::collections::HashMap::new(); let mut by_nic: std::collections::HashMap, Vec> = std::collections::HashMap::new(); for p in partitions { by_nic.entry(p.nic.clone()).or_default().push(p.id); + let make_stack = make_stack_for(&p.spawner); let state = std::sync::Arc::new(PartitionState { id: p.id, spawner: p.spawner, nic: p.nic, authorities: std::sync::Mutex::new(std::collections::HashMap::new()), - make_stack: make_stack.clone(), + make_stack, }); if by_id.insert(p.id, state).is_some() { panic!("duplicate PartitionId declared: {:?}", p.id); diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index f6807a3b4f4..2d18bf4248f 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1604,3 +1604,111 @@ mod listener_tests { assert!(duration > Duration::ZERO, "connect_duration should be > 0"); } } + +/// Prove that a partition's declared spawner is used to spawn the connection +/// driver, not the free `TokioExecutor::new().execute(…)` helper that would +/// target whatever runtime is current at connect time. +#[tokio::test] +async fn partition_spawners_are_isolated() { + use aws_smithy_http_client::pool::{DriverSpawner, Partition, PartitionId, TokioDriverSpawner}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + /// A spawner that records how many times `spawn` was called, then + /// delegates to an inner `TokioDriverSpawner` so the driver runs. + #[derive(Debug)] + struct RecordingSpawner { + inner: TokioDriverSpawner, + spawned: Arc, + } + + impl RecordingSpawner { + fn new(spawned: Arc) -> Self { + Self { + inner: TokioDriverSpawner::current(), + spawned, + } + } + } + + impl DriverSpawner for RecordingSpawner { + fn spawn( + &self, + driver: std::pin::Pin + Send + 'static>>, + ) { + self.spawned.fetch_add(1, Ordering::Relaxed); + self.inner.spawn(driver); + } + } + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"hello", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"hello", + }, + ], + ) + .build() + .await; + + // Two partitions, each with its own recording spawner. + let count0 = Arc::new(AtomicUsize::new(0)); + let count1 = Arc::new(AtomicUsize::new(0)); + let p0 = Partition::new( + PartitionId::from_index(0), + RecordingSpawner::new(count0.clone()), + ); + let p1 = Partition::new( + PartitionId::from_index(1), + RecordingSpawner::new(count1.clone()), + ); + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .partitions([p0, p1]) + .build_http(); + + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + // A request on partition 0 must drive its connection through spawner 0 + // only — spawner 1 stays untouched (no cross-partition leakage). + let client0 = PoolClient::from_partition(&pool, PartitionId::from_index(0)).into_shared(); + let (status, _) = send_and_read_body(&client0, &url) + .await + .expect("p0 request should succeed"); + assert_eq!(status, 200); + assert!( + count0.load(Ordering::Relaxed) >= 1, + "partition 0's spawner should have driven its connection" + ); + assert_eq!( + count1.load(Ordering::Relaxed), + 0, + "partition 1's spawner must not be touched by a partition 0 request" + ); + + // A request on partition 1 drives through spawner 1; spawner 0's count + // does not change (each partition opens its own connection). + let count0_before = count0.load(Ordering::Relaxed); + let client1 = PoolClient::from_partition(&pool, PartitionId::from_index(1)).into_shared(); + let (status, _) = send_and_read_body(&client1, &url) + .await + .expect("p1 request should succeed"); + assert_eq!(status, 200); + assert!( + count1.load(Ordering::Relaxed) >= 1, + "partition 1's spawner should have driven its connection" + ); + assert_eq!( + count0.load(Ordering::Relaxed), + count0_before, + "partition 0's spawner must not be touched by a partition 1 request" + ); +} From db9b1fed6344ea3754f628089979e2d5e84eeaf2 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Sun, 7 Jun 2026 10:03:10 -0400 Subject: [PATCH 32/47] wire stats index --- .../aws-smithy-http-client/src/client/pool.rs | 144 ++++- .../src/client/pool/connection.rs | 69 ++- .../src/client/pool/handshake.rs | 46 +- .../src/client/pool/partition.rs | 4 +- .../src/client/pool/stats.rs | 537 ++++++++++++++++++ .../tests/pool_behavior_test.rs | 208 +++++++ 6 files changed, 979 insertions(+), 29 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 4b5236c4831..dd71c0c41fe 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -42,6 +42,7 @@ pub(crate) mod connection; mod handshake; +pub(crate) mod stats; mod vendored_cache; pub mod builder; @@ -58,6 +59,9 @@ pub use connection::{ pub use partition::{ CrossPartitionPolicy, DriverSpawner, Partition, PartitionId, TokioDriverSpawner, }; +pub use stats::{AuthorityStats, PartitionStats}; + +pub(crate) use stats::{ConnectionCounters, StatsIndex}; /// Connection-caching pool layer. mod cache { @@ -183,6 +187,20 @@ impl SharedPool { pub fn builder() -> Builder { Builder::default() } + + /// Point-in-time, per-partition snapshot of connection counts for `authority`. + /// + /// Relaxed atomics; values may be stale. Sparse — only partitions that have + /// opened a connection to this authority appear. + /// + /// The `authority` must match the form the pool keys on (see + /// [`Authority::from_host`]); a value that matches no keyed authority + /// yields empty stats. + /// + /// [`Authority::from_host`]: crate::client::pool::Authority::from_host + pub fn stats(&self, authority: &Authority) -> stats::AuthorityStats { + self.inner.pool.shared.stats_index().snapshot(authority) + } } /// Key for per-host connection pool routing. @@ -246,6 +264,7 @@ pub(crate) struct SharedPoolState { pub(crate) global_sem: Option>, max_connections_per_host: Option, per_host_sems: Mutex>>, + stats_index: StatsIndex, } impl SharedPoolState { @@ -255,6 +274,7 @@ impl SharedPoolState { global_sem: config.max_connections.map(|n| Arc::new(Semaphore::new(n))), max_connections_per_host: config.max_connections_per_host, per_host_sems: Mutex::new(HashMap::new()), + stats_index: StatsIndex::default(), } } @@ -271,6 +291,11 @@ impl SharedPoolState { .clone(), ) } + + /// Pool-level inverted index for stats reads. + pub(crate) fn stats_index(&self) -> &StatsIndex { + &self.stats_index + } } pub(crate) fn build_pool( @@ -296,7 +321,9 @@ where let make_stack_for = { let connector = connector.clone(); - move |spawner: &Arc| -> MakeStack { + move |partition_id: partition::PartitionId, + spawner: &Arc| + -> MakeStack { let connector = connector.clone(); let spawner = spawner.clone(); Arc::new( @@ -306,12 +333,19 @@ where .expect("pool entry URI has authority") .as_str(), ); + + let counters = Arc::new(ConnectionCounters::default()); + shared + .stats_index() + .register(authority.clone(), partition_id, &counters); + let key = PoolKey::from_uri(uri).expect("pool entry URI has scheme+authority"); let per_host_sem = shared.per_host_sem(&key); let limited = handshake::ConnectionLimit::new( connector.clone(), shared.global_sem.clone(), per_host_sem, + counters.clone(), ); let pool_hooks = shared.hooks.clone(); @@ -335,6 +369,7 @@ where let retainers = retainers.clone(); let pool_hooks = pool_hooks.clone(); let spawner = spawner.clone(); + let counters = counters.clone(); tower::layer::layer_fn(move |inspector| { let cache = cache::builder() .executor(TokioExecutor::new()) @@ -409,10 +444,12 @@ where ); cache.map_response({ let listener = pool_hooks.listener.clone(); + let counters = counters.clone(); move |cached| { H1Checkout::new(CachedConnection::new( cached, listener.clone(), + counters.clone(), )) } }) @@ -423,6 +460,7 @@ where let retainers = retainers.clone(); let pool_hooks = pool_hooks.clone(); let spawner = spawner.clone(); + let counters = counters.clone(); tower::layer::layer_fn(move |inspected| { let singleton = hpool::singleton::Singleton::new( H2ConnectAndHandshake::new( @@ -505,17 +543,23 @@ where ); singleton.map_response({ let h2_ref = h2_ref.clone(); + let counters = counters.clone(); move |singled| { H2Checkout::new(SingletonConnection::new( singled, h2_ref.clone(), + counters.clone(), )) } }) }) }) .build(); - Box::new(TypedPoolEntry { stack, retainers }) + Box::new(TypedPoolEntry { + stack, + retainers, + counters, + }) }, ) } @@ -675,16 +719,30 @@ impl ConnectionPool { /// partitions are empty (no-op). fn retain_idle(&self, timeout: Duration) { for partition in self.registry.partitions() { - let mut auth = partition.authorities.lock().unwrap(); - auth.retain(|_key, entry| { - entry.retain_idle(timeout); - if entry.is_empty() { - tracing::debug!("pool: host entry removed (empty after retain)"); - false - } else { - true - } - }); + let partition_id = partition.id; + let mut evicted: Vec = Vec::new(); + { + let mut auth = partition.authorities.lock().unwrap(); + auth.retain(|key, entry| { + entry.retain_idle(timeout); + if entry.is_empty() { + tracing::debug!("pool: host entry removed (empty after retain)"); + evicted.push(Authority::new(key.authority.as_str())); + false + } else { + true + } + }); + // `auth` guard drops here, dropping the removed entries (and + // the strong `Arc` they held). Only then + // can the index prune observe a zero strong count — unless a + // checkout is still in flight, in which case the cell stays. + } + for authority in evicted { + self.shared + .stats_index() + .prune_if_dead(&authority, partition_id); + } } } } @@ -969,6 +1027,9 @@ struct TypedPoolEntry { /// Empty until the first request against this host; safe to iterate /// at any point (no leg = no-op retain, trivially empty). retainers: RetainerSlot, + // active write-handle; read API added in a later change + #[allow(dead_code)] + counters: Arc, } impl PoolEntry for TypedPoolEntry @@ -1079,7 +1140,9 @@ mod tests { fn pool_with_config(config: PoolConfig) -> Arc { let shared = Arc::new(SharedPoolState::new(&config)); - let make_stack_for = |_spawner: &Arc| -> MakeStack { + let make_stack_for = |_partition_id: partition::PartitionId, + _spawner: &Arc| + -> MakeStack { Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box) }; let partitions = partition::normalize_partitions(Vec::new(), || { @@ -1219,4 +1282,59 @@ mod tests { "distinct partition ids must yield distinct PartitionState arcs" ); } + + /// End-to-end lifecycle: `EstablishingGuard::new` → `promote` → + /// `ManagedConnection` holds `EstablishedGuard` → drop decrements. + /// Verifies the wiring from handshake through connection lifetime. + #[test] + fn established_set_after_handshake() { + use connection::{Authority, ConnectionInfo, ConnectionPermit, ManagedConnection}; + use stats::{ConnectionCounters, EstablishingGuard}; + + let counters = Arc::new(ConnectionCounters::default()); + let authority = Authority::new("example.com:443"); + let partition_id = partition::PartitionId::from_index(0); + + // Register in StatsIndex so we can read through that path too + let index = StatsIndex::default(); + index.register(authority.clone(), partition_id, &counters); + + // Simulate: ConnectionLimit creates guard post-permit + let establishing = EstablishingGuard::new(counters.clone()); + assert_eq!(index.establishing_for(&authority, partition_id), 1); + assert_eq!(index.established_for(&authority, partition_id), 0); + + // Simulate: handshake succeeds, promote + let established = establishing.promote(stats::PROTO_H1); + assert_eq!(index.establishing_for(&authority, partition_id), 0); + assert_eq!(index.established_for(&authority, partition_id), 1); + + // Simulate: ManagedConnection holds the EstablishedGuard + let permit = Arc::new(ConnectionPermit::new(None, None)); + let info = ConnectionInfo { + remote_addr: None, + local_addr: None, + is_proxied: false, + authority: authority.clone(), + }; + let conn: ManagedConnection<()> = ManagedConnection::new( + (), + info, + aws_smithy_runtime_api::client::connection::ConnectionId::new(0), + permit, + established, + ); + + // H2-style: clone shares the same Arc + let conn2 = conn.clone(); + assert_eq!(index.established_for(&authority, partition_id), 1); + + // Drop one clone — guard still alive via the other + drop(conn); + assert_eq!(index.established_for(&authority, partition_id), 1); + + // Drop last clone — guard fires, established decrements + drop(conn2); + assert_eq!(index.established_for(&authority, partition_id), 0); + } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index e0d166fa277..a85197dd932 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -31,6 +31,22 @@ impl Authority { Self(s.into()) } + /// Construct an authority key for a host, as `host` or `host:port`. + /// + /// Used to query connection counts via [`SharedPool::stats`]. The + /// pool keys connection state by the authority component of each + /// request's URI, compared as an exact byte string. The value passed + /// here must match that form for a lookup to hit: notably, a port is + /// present only when the URI carried a non-default port (an HTTPS URI + /// to `example.com` keys as `example.com`, not `example.com:443`), and + /// the host is matched case-sensitively. A value that does not match + /// any keyed authority yields empty stats rather than an error. + /// + /// [`SharedPool::stats`]: crate::client::pool::SharedPool::stats + pub fn from_host(host: impl Into>) -> Self { + Self(host.into()) + } + /// The authority as a string slice. pub fn as_str(&self) -> &str { &self.0 @@ -386,6 +402,7 @@ impl ConnectionPermit { pub(crate) struct EstablishedConnection { pub(crate) io: IO, pub(crate) permit: Arc, + pub(crate) establishing: super::stats::EstablishingGuard, } /// Error returned by `ManagedConnection::poll_ready` when the connection @@ -501,6 +518,9 @@ pub(crate) struct ManagedConnection { active_streams: Arc, poison: PoisonPill, _permit: Arc, + /// Fires `established--` when the last clone drops. Arc-shared so H2 + /// clones share one guard that drops once. + _established: Arc, } impl ManagedConnection { @@ -510,6 +530,7 @@ impl ManagedConnection { info: ConnectionInfo, conn_id: ConnectionId, permit: Arc, + established: super::stats::EstablishedGuard, ) -> Self { let now = Instant::now(); Self { @@ -521,6 +542,7 @@ impl ManagedConnection { active_streams: Arc::new(std::sync::atomic::AtomicUsize::new(0)), poison: PoisonPill::healthy(), _permit: permit, + _established: Arc::new(established), } } @@ -629,6 +651,7 @@ impl Clone for ManagedConnection { active_streams: self.active_streams.clone(), poison: self.poison.clone(), _permit: self._permit.clone(), + _established: self._established.clone(), } } } @@ -668,13 +691,16 @@ where pub(crate) struct CachedConnection { inner: Option>>, listener: Option>, + counters: Arc, } impl CachedConnection { pub(crate) fn new( cached: cache::Cached>, listener: Option>, + counters: Arc, ) -> Self { + counters.incr_active(); let is_reuse = *cached.inner().idle_at.lock().unwrap() > cached.inner().created_at; if is_reuse { let conn_id = cached.inner().conn_id(); @@ -687,6 +713,7 @@ impl CachedConnection { Self { inner: Some(cached), listener, + counters, } } @@ -761,6 +788,7 @@ where impl Drop for CachedConnection { fn drop(&mut self) { + self.counters.decr_active(); if let Some(cached) = self.inner.take() { let managed = cached.inner(); let conn_id = managed.conn_id; @@ -866,25 +894,29 @@ impl H2ConnectionRef { struct DispatchGuard { active_streams: Arc, idle_at: Arc>, + counters: Arc, } impl DispatchGuard { /// Start a dispatch against the connection described by `state`. /// Increments `active_streams`; the returned guard releases the /// increment on drop. - fn start(state: &H2ConnectionState) -> Self { + fn start(state: &H2ConnectionState, counters: Arc) -> Self { state .active_streams .fetch_add(1, std::sync::atomic::Ordering::AcqRel); + counters.incr_active(); Self { active_streams: state.active_streams.clone(), idle_at: state.idle_at.clone(), + counters, } } } impl Drop for DispatchGuard { fn drop(&mut self) { + self.counters.decr_active(); let prev = self .active_streams .fetch_sub(1, std::sync::atomic::Ordering::AcqRel); @@ -930,15 +962,22 @@ pub(crate) struct SingletonConnection { /// Active-stream guard. `Some` between [`Self::call`] and drop. /// `None` if `call` was never invoked or before `call` runs. dispatch: Option, + /// Pool-level warmth counters for this (partition, authority) cell. + counters: Arc, } impl SingletonConnection { - pub(crate) fn new(inner: T, h2_ref: H2ConnectionRef) -> Self { + pub(crate) fn new( + inner: T, + h2_ref: H2ConnectionRef, + counters: Arc, + ) -> Self { let state = h2_ref.current(); Self { inner, state, dispatch: None, + counters, } } @@ -967,7 +1006,7 @@ where fn call(&mut self, req: Req) -> Self::Future { if let Some(state) = self.state.as_ref() { - self.dispatch = Some(DispatchGuard::start(state)); + self.dispatch = Some(DispatchGuard::start(state, self.counters.clone())); } self.inner.call(req) } @@ -1226,7 +1265,8 @@ mod tests { let (active, _) = publish_state(&h2_ref); assert_eq!(active.load(Ordering::Acquire), 0); - let mut sc = SingletonConnection::new(OkService, h2_ref); + let counters = Arc::new(super::super::stats::ConnectionCounters::default()); + let mut sc = SingletonConnection::new(OkService, h2_ref, counters); let _ = sc.call(()).await; assert_eq!( active.load(Ordering::Acquire), @@ -1251,7 +1291,8 @@ mod tests { let (active, _) = publish_state(&h2_ref); active.store(5, Ordering::Release); - let sc = SingletonConnection::new(OkService, h2_ref); + let counters = Arc::new(super::super::stats::ConnectionCounters::default()); + let sc = SingletonConnection::new(OkService, h2_ref, counters); drop(sc); assert_eq!( active.load(Ordering::Acquire), @@ -1272,9 +1313,21 @@ mod tests { // Capture the stamp from construction to detect updates. let original_idle = *idle_at.lock().unwrap(); - let mut a = SingletonConnection::new(OkService, h2_ref.clone()); - let mut b = SingletonConnection::new(OkService, h2_ref.clone()); - let mut c = SingletonConnection::new(OkService, h2_ref); + let mut a = SingletonConnection::new( + OkService, + h2_ref.clone(), + Arc::new(super::super::stats::ConnectionCounters::default()), + ); + let mut b = SingletonConnection::new( + OkService, + h2_ref.clone(), + Arc::new(super::super::stats::ConnectionCounters::default()), + ); + let mut c = SingletonConnection::new( + OkService, + h2_ref, + Arc::new(super::super::stats::ConnectionCounters::default()), + ); let _ = a.call(()).await; let _ = b.call(()).await; diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 9a0fadb0c79..a4338824312 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -102,6 +102,7 @@ pub(crate) struct ConnectionLimit { inner: C, global: Option>, per_host: Option>, + counters: Arc, } impl ConnectionLimit { @@ -109,11 +110,13 @@ impl ConnectionLimit { inner: C, global: Option>, per_host: Option>, + counters: Arc, ) -> Self { Self { inner, global, per_host, + counters, } } } @@ -124,6 +127,7 @@ impl Clone for ConnectionLimit { inner: self.inner.clone(), global: self.global.clone(), per_host: self.per_host.clone(), + counters: self.counters.clone(), } } } @@ -147,6 +151,7 @@ where let mut inner = self.inner.clone(); let global = self.global.clone(); let per_host = self.per_host.clone(); + let counters = self.counters.clone(); Box::pin(async move { // Acquire per-host permit first, then global. This ordering // prevents deadlock: a request never holds a global permit @@ -161,6 +166,7 @@ where None => None, }; let permit = Arc::new(ConnectionPermit::new(global_permit, per_host_permit)); + let establishing = super::stats::EstablishingGuard::new(counters); std::future::poll_fn(|cx| inner.poll_ready(cx)) .await @@ -178,7 +184,11 @@ where super::super::timeout::TimeoutKind::Connect, ) .await?; - Ok(EstablishedConnection { io, permit }) + Ok(EstablishedConnection { + io, + permit, + establishing, + }) }) } } @@ -317,7 +327,11 @@ where let fut = self.connector.call(ctx); Box::pin(async move { let connect_start = Instant::now(); - let EstablishedConnection { io, permit } = match fut.await.map_err(Into::into) { + let EstablishedConnection { + io, + permit, + establishing, + } = match fut.await.map_err(Into::into) { Ok(v) => v, Err(e) => { hooks.on_connection_failed(&ConnectionFailedEvent::new( @@ -364,6 +378,8 @@ where ConnectionTiming::new(connect_duration), )); + let established = establishing.promote(super::stats::PROTO_H1); + spawner.spawn(Box::pin({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; @@ -386,6 +402,7 @@ where info, conn_id, permit, + established, )) }) } @@ -464,7 +481,11 @@ where let spawner = self.spawner.clone(); Box::pin(async move { let connect_start = Instant::now(); - let EstablishedConnection { io, permit } = match fut.await.map_err(Into::into) { + let EstablishedConnection { + io, + permit, + establishing, + } = match fut.await.map_err(Into::into) { Ok(v) => v, Err(e) => { hooks.on_connection_failed(&ConnectionFailedEvent::new( @@ -511,6 +532,8 @@ where ConnectionTiming::new(connect_duration), )); + let established = establishing.promote(super::stats::PROTO_H2); + spawner.spawn(Box::pin({ let remote_addr = info.remote_addr; let local_addr = info.local_addr; @@ -528,7 +551,8 @@ where } })); - let managed = ManagedConnection::new(H2SendRequest::new(tx), info, conn_id, permit); + let managed = + ManagedConnection::new(H2SendRequest::new(tx), info, conn_id, permit, established); // Publish this connection's state (metadata + active_streams + // idle_at refs) for the checkout side. Singleton replaces its // stored connection wholesale on each new handshake, so @@ -557,7 +581,12 @@ mod tests { /// and produce an `HTTP connect timeout occurred after …` error. #[tokio::test(start_paused = true)] async fn connect_timeout_fires_on_slow_connector() { - let mut svc = ConnectionLimit::new(NeverConnects::default(), None, None); + let mut svc = ConnectionLimit::new( + NeverConnects::default(), + None, + None, + Arc::new(super::super::stats::ConnectionCounters::default()), + ); let sleep = SharedAsyncSleep::new(TokioSleep::new()); let ctx = ConnectCtx::new( "http://example.com".parse().unwrap(), @@ -588,7 +617,12 @@ mod tests { /// still pending after a short yield). #[tokio::test(start_paused = true)] async fn no_timeout_does_not_bound_connector() { - let mut svc = ConnectionLimit::new(NeverConnects::default(), None, None); + let mut svc = ConnectionLimit::new( + NeverConnects::default(), + None, + None, + Arc::new(super::super::stats::ConnectionCounters::default()), + ); let ctx = ConnectCtx::new("http://example.com".parse().unwrap(), None); let fut = svc.call(ctx); // A brief tokio yield shouldn't resolve the never-connects future. diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 267d51a10f8..dc3964f23b5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -341,7 +341,7 @@ impl PartitionRegistry { /// no-topology case via [`normalize_partitions`] first). pub(crate) fn build( partitions: Vec, - make_stack_for: impl Fn(&std::sync::Arc) -> super::MakeStack, + make_stack_for: impl Fn(PartitionId, &std::sync::Arc) -> super::MakeStack, ) -> Self { assert!( !partitions.is_empty(), @@ -354,7 +354,7 @@ impl PartitionRegistry { std::collections::HashMap::new(); for p in partitions { by_nic.entry(p.nic.clone()).or_default().push(p.id); - let make_stack = make_stack_for(&p.spawner); + let make_stack = make_stack_for(p.id, &p.spawner); let state = std::sync::Arc::new(PartitionState { id: p.id, spawner: p.spawner, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs new file mode 100644 index 00000000000..b03d598398a --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs @@ -0,0 +1,537 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Per-(partition, authority) connection counters and a pool-level inverted index. +//! +//! Counters are maintained with `Relaxed` atomics. A snapshot observes each counter +//! independently; transient mutual inconsistency is expected (e.g. `active` may +//! briefly exceed `established`). The index never blocks the connection hot path. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; +use std::sync::{Arc, Weak}; + +use super::connection::Authority; +use super::partition::PartitionId; + +// Protocol-tag encoding for the cell. A cell is per-(partition, authority); +// protocol is per-connection, so a cell CAN observe both over its life +// (multi-endpoint authority, server reconfig, proxy). The tag latches to +// MIXED once it sees two different protocols and never leaves it, so +// `capacity_hint` only answers `Some` for a uniformly-H1 cell and never +// overstates reusable capacity. +pub(crate) const PROTO_UNSET: u8 = 0; +pub(crate) const PROTO_H1: u8 = 1; +pub(crate) const PROTO_H2: u8 = 2; +const PROTO_MIXED: u8 = 3; + +/// Per-(partition, authority) connection counts maintained with `Relaxed` atomics. +/// +/// Each counter is loaded independently; concurrent reads may observe transiently +/// inconsistent combinations (e.g. `active` > `established`). Intended for heuristic +/// reads that tolerate stale or momentarily inconsistent values. +#[derive(Debug, Default)] +pub(crate) struct ConnectionCounters { + /// Connections that have completed handshake and exist (idle + active). + pub(crate) established: AtomicUsize, + /// Handshakes in flight. + pub(crate) establishing: AtomicUsize, + /// Connections/streams currently checked out. + pub(crate) active: AtomicUsize, + /// Protocol tag for this cell (monotonic toward MIXED). + protocol: AtomicU8, +} + +/// Tracks a connection committed to handshaking (TCP + TLS + protocol). +/// +/// Construction increments `establishing`. [`promote`](Self::promote) transitions to +/// an [`EstablishedGuard`] (established++ then establishing--) on success; any other +/// drop path (failure, cancel, panic) decrements `establishing`. Exactly-once by +/// construction: `promote` consumes `self`. +pub(crate) struct EstablishingGuard { + counters: Arc, + promoted: bool, +} + +impl EstablishingGuard { + pub(crate) fn new(counters: Arc) -> Self { + counters.establishing.fetch_add(1, Ordering::Relaxed); + Self { + counters, + promoted: false, + } + } + + /// Handshake succeeded: transition establishing → established. + /// + /// `established` is incremented BEFORE `establishing` is decremented so a + /// concurrent reader may observe a transient overcount. The overcount is in the + /// direction of over-reporting readiness, never under-reporting. + pub(crate) fn promote(mut self, proto: u8) -> EstablishedGuard { + self.counters.observe_protocol(proto); + self.counters.established.fetch_add(1, Ordering::Relaxed); + self.counters.establishing.fetch_sub(1, Ordering::Relaxed); + self.promoted = true; + EstablishedGuard { + counters: self.counters.clone(), + } + } +} + +impl Drop for EstablishingGuard { + fn drop(&mut self) { + if !self.promoted { + self.counters.establishing.fetch_sub(1, Ordering::Relaxed); + } + } +} + +/// Owns one connection's contribution to `established`. Held on the +/// Arc-shared inner of a `ManagedConnection`, so for H2 (N clones share +/// one connection) it fires `established--` exactly once — when the last +/// clone drops. Non-`Clone` by design: single ownership is compiler- +/// enforced. +pub(crate) struct EstablishedGuard { + counters: Arc, +} + +impl Drop for EstablishedGuard { + fn drop(&mut self) { + self.counters.established.fetch_sub(1, Ordering::Relaxed); + } +} + +impl ConnectionCounters { + pub(crate) fn incr_active(&self) { + self.active.fetch_add(1, Ordering::Relaxed); + } + + /// Decrement the active count. + /// + /// Every `incr_active` is paired with exactly one `decr_active` via RAII checkout + /// guards, so `active` is non-negative by construction. A saturating sub here + /// would mask a broken-pairing bug; saturation belongs only on the cross-atomic + /// READ (`idle() = established.saturating_sub(active)`) where two independent + /// relaxed loads may transiently cross. + pub(crate) fn decr_active(&self) { + let prev = self.active.fetch_sub(1, Ordering::Relaxed); + debug_assert!(prev > 0, "active underflow: decr without matching incr"); + } + + /// Record the negotiated protocol for a connection in this cell. + /// + /// Latches monotonically toward `MIXED`: `UNSET` → observed protocol; same + /// protocol → no-op; different protocol → `MIXED` (terminal). A cell reaches + /// `MIXED` when one authority negotiates differently across connections + /// (multi-endpoint DNS, server reconfiguration, an intermediary). + /// + /// The sole purpose of this tag is to keep `capacity_hint` from overstating: + /// `capacity_hint` returns `Some` only for a cell known to be uniformly one + /// protocol it can reason about (HTTP/1). Accurate per-connection protocol and + /// stream-limit accounting (per-connection indexing) is future work that would + /// make this cell-level latch unnecessary. + pub(crate) fn observe_protocol(&self, proto: u8) { + let mut cur = self.protocol.load(Ordering::Relaxed); + loop { + let next = match cur { + PROTO_UNSET => proto, + c if c == proto => return, + PROTO_MIXED => return, + _ => PROTO_MIXED, + }; + match self.protocol.compare_exchange_weak( + cur, + next, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(actual) => cur = actual, + } + } + } + + pub(crate) fn protocol(&self) -> u8 { + self.protocol.load(Ordering::Relaxed) + } +} + +/// One authority's row in the inverted index: per-partition weak references to counters. +#[derive(Debug, Default)] +pub(crate) struct AuthorityCounters { + pub(crate) by_partition: HashMap>, +} + +/// Pool-level inverted index: authority → per-partition counters. +/// +/// Non-owning projection: holds `Weak` references. Strong owners +/// (pool entries, live checkouts) keep cells alive; once all strong references drop, +/// the `Weak` goes dead and is pruned on the next `snapshot`. The index never extends +/// connection lifetime and never blocks the per-request hot path. +#[derive(Debug, Default)] +pub(crate) struct StatsIndex { + inner: std::sync::Mutex>, +} + +impl StatsIndex { + /// Register a cell's counters at first-touch of (authority, partition). + /// Stores a `Weak` reference; the caller retains the strong `Arc`. + /// Idempotent: re-registration overwrites the previous entry. + pub(crate) fn register( + &self, + authority: Authority, + partition: PartitionId, + counters: &Arc, + ) { + let mut idx = self.inner.lock().expect("stats index poisoned"); + idx.entry(authority) + .or_default() + .by_partition + .insert(partition, Arc::downgrade(counters)); + } + + /// Drop the `(authority, partition)` cell if its counters are no longer + /// referenced, removing the authority entry when its last partition is + /// pruned. A no-op if the cell is still strongly held (e.g. a checkout + /// is in flight when the host's idle connections are evicted), so it is + /// safe to call from the eviction path. Reconstructs nothing the caller + /// does not already hold. + pub(crate) fn prune_if_dead(&self, authority: &Authority, partition: PartitionId) { + let mut idx = self.inner.lock().expect("stats index poisoned"); + if let Some(a) = idx.get_mut(authority) { + if a.by_partition + .get(&partition) + .is_some_and(|w| w.strong_count() == 0) + { + a.by_partition.remove(&partition); + } + if a.by_partition.is_empty() { + idx.remove(authority); + } + } + } + + #[cfg(test)] + pub(crate) fn len(&self) -> usize { + self.inner.lock().expect("stats index poisoned").len() + } + + #[cfg(test)] + pub(crate) fn established_for(&self, authority: &Authority, partition: PartitionId) -> usize { + let idx = self.inner.lock().expect("stats index poisoned"); + idx.get(authority) + .and_then(|a| a.by_partition.get(&partition)) + .and_then(|w| w.upgrade()) + .map(|c| c.established.load(Ordering::Relaxed)) + .unwrap_or(0) + } + + #[cfg(test)] + pub(crate) fn establishing_for(&self, authority: &Authority, partition: PartitionId) -> usize { + let idx = self.inner.lock().expect("stats index poisoned"); + idx.get(authority) + .and_then(|a| a.by_partition.get(&partition)) + .and_then(|w| w.upgrade()) + .map(|c| c.establishing.load(Ordering::Relaxed)) + .unwrap_or(0) + } + + /// Snapshot one authority's per-partition counters. + /// + /// Under the lock: upgrades each `Weak`, prunes dead entries (removing the + /// authority entirely if all its partitions are dead). Releases the lock before + /// loading atomics into the returned snapshot. + pub(crate) fn snapshot(&self, authority: &Authority) -> AuthorityStats { + let handles: Vec<(PartitionId, Arc)> = { + let mut idx = self.inner.lock().expect("stats index poisoned"); + match idx.get_mut(authority) { + Some(a) => { + let mut live = Vec::new(); + a.by_partition.retain(|partition, weak| { + if let Some(strong) = weak.upgrade() { + live.push((*partition, strong)); + true + } else { + false + } + }); + if a.by_partition.is_empty() { + idx.remove(authority); + } + live + } + None => Vec::new(), + } + }; // lock released here + let by_partition = handles + .into_iter() + .map(|(p, c)| { + ( + p, + PartitionStats { + established: c.established.load(Ordering::Relaxed), + establishing: c.establishing.load(Ordering::Relaxed), + active: c.active.load(Ordering::Relaxed), + protocol: c.protocol(), + }, + ) + }) + .collect(); + AuthorityStats { by_partition } + } +} + +/// Point-in-time snapshot of one (partition, authority) cell's connection counts. +/// +/// Plain `usize` values from relaxed loads; cheap and `Copy`. Each field is loaded +/// independently and may be transiently inconsistent with the others. +/// `#[non_exhaustive]` permits future additions (e.g. per-connection stream depth). +#[non_exhaustive] +#[derive(Debug, Clone, Copy)] +pub struct PartitionStats { + /// Connections that have completed handshake (idle + active). + pub established: usize, + /// Handshakes in flight. + pub establishing: usize, + /// Connections/streams currently checked out. + pub active: usize, + // Private: protocol tag for capacity_hint. Not public — implementation + // detail of the hint that would otherwise freeze an internal encoding + // into the API. + protocol: u8, +} + +impl PartitionStats { + /// Connections not currently checked out: `established.saturating_sub(active)`. + /// + /// Exact for HTTP/1 (one stream per connection). For HTTP/2 a positive value + /// means connections with no active streams; saturated multiplexed connections + /// contribute 0. The saturating subtraction handles transient inconsistency + /// between the two independently-loaded atomics. + pub fn idle(&self) -> usize { + self.established.saturating_sub(self.active) + } + + /// Spare stream capacity, if determinable from current state. + /// + /// - Uniformly HTTP/1 cell: `Some(idle)` (one stream per connection). + /// - HTTP/2, mixed-protocol, or not-yet-handshaken cell: `None` (per-connection + /// stream limits are not indexed). + /// + /// When `Some`, still a relaxed snapshot — treat as a hint. + pub fn capacity_hint(&self) -> Option { + match self.protocol { + PROTO_H1 => Some(self.idle()), + _ => None, + } + } +} + +/// Point-in-time, per-partition snapshot of connection counts for one authority. +/// +/// Sparse: only partitions that have opened a connection to the authority appear. +/// Returned by [`super::SharedPool::stats`]. +pub struct AuthorityStats { + by_partition: Vec<(PartitionId, PartitionStats)>, +} + +impl AuthorityStats { + /// Stats for a specific partition, if it has opened a connection to this authority. + pub fn get(&self, partition: PartitionId) -> Option { + self.by_partition + .iter() + .find(|(p, _)| *p == partition) + .map(|(_, s)| *s) + } + + /// Iterate (partition, stats) pairs. + pub fn iter(&self) -> impl Iterator + '_ { + self.by_partition.iter().copied() + } + + /// Number of partitions that have opened a connection to this authority. + pub fn len(&self) -> usize { + self.by_partition.len() + } + + /// Whether no partition has connected to this authority. + pub fn is_empty(&self) -> bool { + self.by_partition.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn active_incremented_on_checkout_decremented_on_drop() { + let counters = Arc::new(ConnectionCounters::default()); + assert_eq!(counters.active.load(Ordering::Relaxed), 0); + + counters.incr_active(); + assert_eq!(counters.active.load(Ordering::Relaxed), 1); + + counters.decr_active(); + assert_eq!(counters.active.load(Ordering::Relaxed), 0); + } + + #[test] + fn establishing_guard_promote_increments_established_and_clears_establishing() { + let counters = Arc::new(ConnectionCounters::default()); + + let guard = EstablishingGuard::new(counters.clone()); + assert_eq!(counters.establishing.load(Ordering::Relaxed), 1); + assert_eq!(counters.established.load(Ordering::Relaxed), 0); + + let established = guard.promote(PROTO_H1); + assert_eq!(counters.establishing.load(Ordering::Relaxed), 0); + assert_eq!(counters.established.load(Ordering::Relaxed), 1); + + drop(established); + assert_eq!(counters.established.load(Ordering::Relaxed), 0); + } + + #[test] + fn establishing_guard_drop_without_promote_decrements() { + let counters = Arc::new(ConnectionCounters::default()); + + let guard = EstablishingGuard::new(counters.clone()); + assert_eq!(counters.establishing.load(Ordering::Relaxed), 1); + + drop(guard); + assert_eq!(counters.establishing.load(Ordering::Relaxed), 0); + assert_eq!(counters.established.load(Ordering::Relaxed), 0); + } + + #[test] + fn stats_index_registers_counters_arc() { + let index = StatsIndex::default(); + assert_eq!(index.len(), 0); + + let counters = Arc::new(ConnectionCounters::default()); + let authority_a = Authority::new("a.example.com:443"); + let partition = PartitionId::from_index(0); + index.register(authority_a, partition, &counters); + assert_eq!(index.len(), 1); + + // Different authority → new entry + let authority_b = Authority::new("b.example.com:443"); + let counters2 = Arc::new(ConnectionCounters::default()); + index.register(authority_b, partition, &counters2); + assert_eq!(index.len(), 2); + + // Same authority, different partition → same entry (len unchanged) + let authority_a2 = Authority::new("a.example.com:443"); + let counters3 = Arc::new(ConnectionCounters::default()); + let partition2 = PartitionId::from_index(1); + index.register(authority_a2, partition2, &counters3); + assert_eq!(index.len(), 2); + } + + #[test] + fn stats_index_prunes_dead_cells() { + let index = StatsIndex::default(); + let authority = Authority::new("ephemeral.example.com:443"); + let partition = PartitionId::from_index(0); + + let counters = Arc::new(ConnectionCounters::default()); + index.register(authority.clone(), partition, &counters); + assert_eq!(index.len(), 1); + + // Drop the only strong reference — the Weak in the index is now dead + drop(counters); + + // Snapshot triggers pruning; dead cell is removed + let snap = index.snapshot(&authority); + assert!(snap.is_empty()); + assert_eq!(index.len(), 0); + } + + #[test] + fn prune_if_dead_keeps_live_cell_removes_dead_cell() { + let index = StatsIndex::default(); + let authority = Authority::new("host.example.com:443"); + let partition = PartitionId::from_index(0); + let counters = Arc::new(ConnectionCounters::default()); + index.register(authority.clone(), partition, &counters); + assert_eq!(index.len(), 1); + + // Cell is still strongly held (mirrors a checkout in flight when the + // host's idle connections are evicted): prune is a no-op. + index.prune_if_dead(&authority, partition); + assert_eq!(index.len(), 1); + + // Strong ref gone (entry + checkouts dropped): prune removes the cell + // and, as its last partition, the authority entry. + drop(counters); + index.prune_if_dead(&authority, partition); + assert_eq!(index.len(), 0); + } + + #[test] + fn capacity_hint_h1_some_h2_none_mixed_none() { + // H1-only cell: capacity_hint == Some(idle) + let s = PartitionStats { + established: 5, + establishing: 0, + active: 2, + protocol: PROTO_H1, + }; + assert_eq!(s.capacity_hint(), Some(3)); + + // H2-only cell: None + let s = PartitionStats { + established: 5, + establishing: 0, + active: 2, + protocol: PROTO_H2, + }; + assert_eq!(s.capacity_hint(), None); + + // UNSET cell: None + let s = PartitionStats { + established: 0, + establishing: 1, + active: 0, + protocol: PROTO_UNSET, + }; + assert_eq!(s.capacity_hint(), None); + + // Mixed cell via observe_protocol transitions: None + let counters = Arc::new(ConnectionCounters::default()); + counters.observe_protocol(PROTO_H1); + assert_eq!(counters.protocol(), PROTO_H1); + counters.observe_protocol(PROTO_H2); + assert_eq!(counters.protocol(), PROTO_MIXED); + let s = PartitionStats { + established: 5, + establishing: 0, + active: 2, + protocol: counters.protocol(), + }; + assert_eq!(s.capacity_hint(), None); + } + + #[test] + fn idle_is_established_minus_active_saturating() { + let s = PartitionStats { + established: 3, + establishing: 0, + active: 1, + protocol: PROTO_H1, + }; + assert_eq!(s.idle(), 2); + + // H2 over-subscription: active > established (multiple streams per conn) + let s = PartitionStats { + established: 1, + establishing: 0, + active: 5, + protocol: PROTO_H2, + }; + assert_eq!(s.idle(), 0); + } +} diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 2d18bf4248f..d8140944d8f 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1712,3 +1712,211 @@ async fn partition_spawners_are_isolated() { "partition 0's spawner must not be touched by a partition 1 request" ); } + +// --------------------------------------------------------------------------- +// Test implementations: stats read API +// --------------------------------------------------------------------------- + +/// SharedPool::stats returns sparse per-partition snapshots: only partitions +/// that have touched an authority appear. After a request completes and the +/// body is consumed, counters reflect the idle state. +#[tokio::test] +async fn v2_stats_reports_per_partition_sparse() { + use aws_smithy_http_client::pool::{Authority, Partition, PartitionId, TokioDriverSpawner}; + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"stats-test", + }], + ) + .build() + .await; + + let p0 = Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()); + let p1 = Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()); + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .partitions([p0, p1]) + .build_http(); + + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + let authority = Authority::from_host(format!("127.0.0.1:{port}")); + + // Before any request, stats are empty for this authority. + let stats = pool.stats(&authority); + assert!( + stats.is_empty(), + "no partition should have touched this authority yet" + ); + + // Send a request on partition 0, consume the body so the connection idles. + let client0 = PoolClient::from_partition(&pool, PartitionId::from_index(0)).into_shared(); + let (status, body) = send_and_read_body(&client0, &url) + .await + .expect("p0 request should succeed"); + assert_eq!(status, 200); + assert_eq!(body, b"stats-test"); + + // After the request idles, partition 0 should appear, partition 1 should not. + let stats = pool.stats(&authority); + assert_eq!(stats.len(), 1, "only partition 0 should appear (sparse)"); + + let p0_stats = stats + .get(PartitionId::from_index(0)) + .expect("partition 0 should have stats"); + assert_eq!(p0_stats.established, 1, "one connection established"); + assert_eq!(p0_stats.establishing, 0, "no handshakes in flight"); + assert_eq!(p0_stats.active, 0, "connection is idle after body consumed"); + assert_eq!(p0_stats.idle(), 1, "one idle connection"); + // H1 cell: capacity_hint is Some(idle) + assert_eq!(p0_stats.capacity_hint(), Some(1)); + + assert!( + stats.get(PartitionId::from_index(1)).is_none(), + "partition 1 has not touched this authority" + ); +} + +/// `active` tracks an in-flight request end-to-end. The H1 checkout guard +/// rides the response body (`GuardedBody`): it is held while the response +/// value is alive and releases when the body is consumed/dropped, returning +/// the connection to the pool. So `active == 1` is observable for as long as +/// the caller holds the response, and drops to 0 once the body is drained. +#[tokio::test] +async fn v2_stats_active_tracks_in_flight_request() { + use aws_smithy_http_client::pool::{Authority, PartitionId}; + use http_body_util::BodyExt; + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"in-flight", + }], + ) + .build() + .await; + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .build_http(); + + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + let authority = Authority::from_host(format!("127.0.0.1:{port}")); + let partition = PartitionId::default(); + + let client = PoolClient::new(&pool).into_shared(); + + // Issue the request but hold the response without draining the body. The + // connection is checked out: its guard is alive inside `resp`'s body. + let resp = send_to(&client, &url) + .await + .expect("request should succeed"); + assert_eq!(resp.status().as_u16(), 200); + + let stats = pool.stats(&authority); + let in_flight = stats + .get(partition) + .expect("partition should have touched this authority"); + assert_eq!(in_flight.established, 1, "one connection established"); + assert_eq!( + in_flight.active, 1, + "request is in flight, connection checked out" + ); + assert_eq!(in_flight.idle(), 0, "no idle connection while in flight"); + + // Drain the body: the GuardedBody drops, CachedConnection::Drop fires, + // active decrements and the connection returns to the pool as idle. + let body = resp + .into_body() + .collect() + .await + .expect("body should be readable") + .to_bytes() + .to_vec(); + assert_eq!(body, b"in-flight"); + + let stats = pool.stats(&authority); + let idle = stats + .get(partition) + .expect("partition still present after request completes"); + assert_eq!(idle.established, 1, "connection still established (idle)"); + assert_eq!(idle.active, 0, "no in-flight request after body drained"); + assert_eq!(idle.idle(), 1, "connection is idle and reusable"); +} + +/// Eviction decrements `established` and prunes the stats-index cell through +/// its real trigger — the background eviction task — not a direct prune call. +/// After an idle connection is evicted, the host entry is removed and the +/// eviction-triggered prune drops the now-dead cell from the index, so +/// `stats()` reports the authority as untracked. +#[tokio::test] +async fn v2_stats_pruned_after_eviction() { + use aws_smithy_http_client::pool::{Authority, PartitionId}; + + let idle_timeout = Duration::from_millis(100); + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"evict-me", + }, + // A second connection is available if eviction forces a reconnect; + // the test asserts on stats, not connection count. + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"evict-me", + }, + ], + ) + .build() + .await; + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .pool_idle_timeout(idle_timeout) + .build_http(); + + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + let authority = Authority::from_host(format!("127.0.0.1:{port}")); + + let client = PoolClient::new(&pool).into_shared(); + + // Request completes and the connection idles. This also lazily spawns the + // eviction task (pool_idle_timeout is set). + let (status, _) = send_and_read_body(&client, &url) + .await + .expect("request should succeed"); + assert_eq!(status, 200); + + let stats = pool.stats(&authority); + assert_eq!( + stats + .get(PartitionId::default()) + .expect("partition present after request") + .established, + 1, + "one established idle connection before eviction" + ); + + // Wait past the idle timeout: the eviction task drops the idle connection, + // removes the host entry, and prunes the now-dead index cell. + tokio::time::sleep(idle_timeout * 3).await; + + let stats = pool.stats(&authority); + assert!( + stats.is_empty(), + "eviction should have decremented established and pruned the cell" + ); +} From 92d7819fb3e671b7c68932858ba18617882b33bf Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Mon, 8 Jun 2026 13:22:54 -0400 Subject: [PATCH 33/47] try reclaim --- .../aws-smithy-http-client/src/client/pool.rs | 76 +++++++++++++++++++ .../src/client/pool/connection.rs | 5 ++ .../src/client/pool/vendored_cache.rs | 15 ++++ 3 files changed, 96 insertions(+) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index dd71c0c41fe..4108d321c75 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -233,6 +233,12 @@ impl PoolKey { pub(crate) struct BoxedRetainer { retain_fn: Box, is_empty_fn: Box bool + Send + Sync + 'static>, + /// Pop one idle connection and drop it to free its permit (active + /// reclaim). Returns `true` if a connection was freed. The H1 cache + /// leg pops from its idle Vec; the H2 singleton leg has no + /// simply-reclaimable idle (an idle H2 connection still multiplexes), + /// so its `reclaim_fn` is a no-op returning `false`. + reclaim_fn: Box bool + Send + Sync + 'static>, } impl BoxedRetainer { @@ -243,6 +249,10 @@ impl BoxedRetainer { fn is_empty(&self) -> bool { (self.is_empty_fn)() } + + fn reclaim_one(&self) -> bool { + (self.reclaim_fn)() + } } /// Per-host retainer registry. Bounded at two entries (one H1 cache, one @@ -386,6 +396,7 @@ where let cache_for_retain = Arc::new(std::sync::Mutex::new(cache.clone())); let cache_for_empty = cache_for_retain.clone(); + let cache_for_reclaim = cache_for_retain.clone(); retainers.lock().expect("retainer slot poisoned").push( BoxedRetainer { retain_fn: Box::new({ @@ -440,6 +451,47 @@ where .expect("is_empty cache lock poisoned") .is_empty() }), + reclaim_fn: Box::new({ + let listener = pool_hooks.listener.clone(); + move || { + // Pop under the cache lock, RELEASE the + // lock, THEN drop the connection (never + // drop while holding the cache Mutex). + let managed = cache_for_reclaim + .lock() + .expect("reclaim cache lock poisoned") + .try_pop_idle(); + match managed { + Some(managed) => { + tracing::debug!( + conn_id = %managed.conn_id(), + "pool: connection reclaimed" + ); + if let Some(ref l) = listener { + l.on_closed( + &ConnectionClosedEvent::new( + managed.conn_id(), + managed + .info + .authority + .clone(), + managed.info.remote_addr, + CloseReason::Reclaimed, + None, + ), + ); + } + // `managed` drops here, after the + // cache lock is released: its + // `ConnectionPermit` returns a + // permit to the shared semaphore. + drop(managed); + true + } + None => false, + } + } + }), }, ); cache.map_response({ @@ -539,6 +591,10 @@ where .expect("is_empty singleton lock poisoned") .is_empty() }), + // An idle H2 connection still multiplexes; it is + // not simply-reclaimable like an H1 cache entry. + // Reclaim skips the H2 leg. + reclaim_fn: Box::new(|| false), }, ); singleton.map_response({ @@ -1013,6 +1069,15 @@ pub(crate) trait PoolEntry: Send + Sync { /// singleton empty). Used to drop empty host entries from the map /// after `retain_idle`. fn is_empty(&self) -> bool; + + /// Pop one idle connection and drop it to free its permit, returning + /// `true` if one was freed. Drives cross-partition active reclaim: a + /// starved partition frees an over-supplied peer's idle capacity at + /// the cap-bound point. The dropped connection's `ConnectionPermit` + /// releases back to the shared semaphore. Fires `on_closed` with + /// [`CloseReason::Reclaimed`] for the freed connection. Best-effort: + /// returns `false` if the entry has no reclaimable idle. + fn try_reclaim_one(&self) -> bool; } /// Concrete PoolEntry wrapping a Negotiate stack. @@ -1101,6 +1166,14 @@ where let retainers = self.retainers.lock().expect("retainer slot poisoned"); retainers.iter().all(|r| r.is_empty()) } + + fn try_reclaim_one(&self) -> bool { + let retainers = self.retainers.lock().expect("retainer slot poisoned"); + // Stop at the first leg that frees a connection. The H1 cache leg + // pops an idle connection; the H2 leg is a no-op (no + // simply-reclaimable idle). + retainers.iter().any(|r| r.reclaim_one()) + } } #[cfg(test)] @@ -1136,6 +1209,9 @@ mod tests { fn is_empty(&self) -> bool { true } + fn try_reclaim_one(&self) -> bool { + false + } } fn pool_with_config(config: PoolConfig) -> Arc { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index a85197dd932..fa230a816d0 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -89,6 +89,11 @@ pub enum CloseReason { Unusable, /// The pool itself was dropped. PoolDropped, + /// Dropped to free a connection permit for another partition under + /// cap pressure (cross-partition active reclaim). Distinct from + /// `IdleTimeout`: the connection was still within its idle window but + /// was reclaimed because a starved partition needed the capacity. + Reclaimed, } /// Timing breakdown for connection establishment. diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs index cf476ee9aa9..83157653d70 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs @@ -182,6 +182,21 @@ mod internal { pub fn is_empty(&self) -> bool { self.shared.lock().unwrap().services.is_empty() } + + // SDK MODIFICATION: added `try_pop_idle` so a caller can remove an + // idle cached service to free whatever resource it holds, instead + // of waiting for it to be re-handed-out or evicted. + /// Remove and return one idle cached service, if any. + /// + /// Unlike [`Service::call`], which wraps the taken service in a + /// [`Cached`] that returns to the pool on drop, this hands back the + /// raw service with no return-to-pool wrapper: dropping it drops the + /// service outright. Serialized against [`Self::retain`] by the + /// shared `Mutex`, so a popped service is removed before a retain + /// pass can observe it. + pub fn try_pop_idle(&self) -> Option { + self.shared.lock().unwrap().services.pop() + } } impl Service for Cache From 8e8e8f3c5cfde25b154e073ee56d820684ec6e3f Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 9 Jun 2026 11:01:36 -0400 Subject: [PATCH 34/47] Add binding constraint and peer reclaim --- .../aws-smithy-http-client/src/client/pool.rs | 145 +++++++++++++++++- .../src/client/pool/partition.rs | 63 ++++++++ .../src/client/pool/stats.rs | 69 +++++++++ 3 files changed, 275 insertions(+), 2 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 4108d321c75..04da1d67dc5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -266,6 +266,114 @@ type RetainerSlot = Arc>>; pub(crate) type MakeStack = Arc Box + Send + Sync>; +/// Which semaphore bound a connect attempt at the cap, identifying what a +/// reclaim must free to relieve it. Acquire order is per-host then global +/// (§8), so a per-host failure yields `PerHost`, and a failure on the +/// global semaphore while already holding the per-host permit yields +/// `Global`. +#[derive(Debug, Clone)] +pub(crate) enum BindingConstraint { + /// The global `max_connections` semaphore is exhausted. Any + /// over-supplied peer's idle connection frees a fungible permit. + Global, + /// The per-host `max_connections_per_host` semaphore for this key is + /// exhausted. Only a peer's idle connection *to the same host* frees + /// the relevant permit. + PerHost(PoolKey), +} + +/// Handle for cross-partition active reclaim, held by `ConnectionLimit`. +/// +/// At a cap-bound connect, the requesting partition uses this to free one +/// over-supplied peer's idle connection (dropping it returns its permit to +/// the bound semaphore) before blocking-acquiring. Connection-shaped and +/// NIC-blind: the freed *permit* is what matters; P0 then connects on its +/// own NIC. Candidates are narrowed by the (advisory) stats index and +/// confirmed by the authoritative cache pop. +#[derive(Clone)] +pub(crate) struct PeerReclaimHandle { + /// `Weak` to avoid a cycle: the registry transitively owns the + /// partitions whose `ConnectionLimit`s hold this handle. + registry: std::sync::Weak, + /// Owned (Arc) so the handle does not borrow `SharedPoolState`. + stats_index: Arc, + /// The requesting partition — excluded from its own candidate walk. + self_partition: PartitionId, +} + +impl PeerReclaimHandle { + /// Free one over-supplied peer's idle connection to relieve + /// `constraint`, returning `true` if a permit was freed. Best-effort: + /// `false` if the pool is gone, there are no NIC-group peers, or no + /// peer holds reclaimable idle (P0 then just blocks on the permit). + pub(crate) fn try_free_under_load(&self, constraint: &BindingConstraint) -> bool { + let registry = match self.registry.upgrade() { + Some(r) => r, + None => return false, // pool dropped → nothing to reclaim + }; + let peers = registry.nic_group_peers(self.self_partition); + if peers.is_empty() { + return false; // alone in the NIC group (e.g. single-partition default) + } + match constraint { + BindingConstraint::PerHost(key) => { + // Candidates = peers with idle to this authority (index + // narrows; cache pop confirms). Round-robin start offset. + let authority = Authority::new(key.authority.as_str()); + let mut candidates: Vec = self + .stats_index + .idle_partitions_for(&authority) + .into_iter() + .map(|(p, _idle)| p) + .filter(|p| *p != self.self_partition && peers.contains(p)) + .collect(); + self.rotate(&mut candidates); + candidates + .into_iter() + .any(|peer| registry.try_reclaim_on(peer, key)) + } + BindingConstraint::Global => { + // Fungible permit: any same-NIC-group peer's idle (any + // authority) relieves the global cap. Narrow to peers that + // the index shows holding idle. + let mut candidates: Vec = self + .stats_index + .idle_cells() + .into_iter() + .map(|(_authority, p)| p) + .filter(|p| *p != self.self_partition && peers.contains(p)) + .collect(); + candidates.sort_unstable_by_key(|p| p.as_u64()); + candidates.dedup(); + self.rotate(&mut candidates); + candidates + .into_iter() + .any(|peer| registry.try_reclaim_any(peer)) + } + } + } + + /// Rotate the candidate vec by the partition's advisory `peer_cursor`, + /// so concurrent reclaims from this partition do not all probe the + /// lowest-numbered candidate first. No-op if the registry or partition + /// is gone, or the candidate set is empty. + fn rotate(&self, candidates: &mut [PartitionId]) { + if candidates.len() < 2 { + return; + } + if let Some(registry) = self.registry.upgrade() { + if let Some(state) = registry.partition_opt(self.self_partition) { + let n = candidates.len(); + let start = state + .peer_cursor + .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + % n; + candidates.rotate_left(start); + } + } + } +} + /// Connection-lifecycle machinery shared across every partition: event /// hooks, the global connection budget, and the per-host budgets. One /// instance per pool, held behind `Arc` on `ConnectionPool`. @@ -274,7 +382,15 @@ pub(crate) struct SharedPoolState { pub(crate) global_sem: Option>, max_connections_per_host: Option, per_host_sems: Mutex>>, - stats_index: StatsIndex, + stats_index: Arc, + /// Late-bound back-reference to the registry, for cross-partition + /// reclaim. `Weak` (not `Arc`): the registry transitively owns the + /// partitions whose `ConnectionLimit`s reach back through here — a + /// strong reference would be a cycle (pool never drops). Set once, + /// immediately after the registry is built (`build_pool`), because the + /// registry does not exist when `SharedPoolState` is constructed. + /// Upgrade-or-skip: a dropped pool makes reclaim a correct no-op. + registry: OnceLock>, } impl SharedPoolState { @@ -284,10 +400,30 @@ impl SharedPoolState { global_sem: config.max_connections.map(|n| Arc::new(Semaphore::new(n))), max_connections_per_host: config.max_connections_per_host, per_host_sems: Mutex::new(HashMap::new()), - stats_index: StatsIndex::default(), + stats_index: Arc::new(StatsIndex::default()), + registry: OnceLock::new(), } } + /// Bind the registry back-reference. Called once in `build_pool` + /// right after the registry is constructed. Idempotent-safe: a second + /// call is ignored (the first binding wins). + pub(crate) fn set_registry(&self, registry: &Arc) { + let _ = self.registry.set(Arc::downgrade(registry)); + } + + /// A reclaim handle for the given requesting partition, if the + /// registry is bound and still alive. `None` when no registry is set + /// or the pool has been dropped (reclaim is then a no-op). + pub(crate) fn reclaim_handle(&self, self_partition: PartitionId) -> Option { + let registry = self.registry.get()?.clone(); + Some(PeerReclaimHandle { + registry, + stats_index: self.stats_index.clone(), + self_partition, + }) + } + /// The shared per-host semaphore for `key`, created on first request /// for that key. Returns `None` when no per-host limit is configured. /// Shared across all partitions: two partitions to the same host @@ -629,6 +765,11 @@ where make_stack_for, )); + // Late-bind the registry back-reference for cross-partition reclaim. + // Done here, after the registry exists, because `SharedPoolState` is + // built before it (the `make_stack` closures capture `shared`). + shared.set_registry(®istry); + ConnectionPool { config, shared, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index dc3964f23b5..a25351fb5a9 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -286,6 +286,12 @@ pub(crate) struct PartitionState { /// Builds a host entry on first touch, capturing this partition's /// connector; shared budget/hooks arrive via `&SharedPoolState`. pub(crate) make_stack: super::MakeStack, + /// Round-robin cursor over reclaim/borrow candidate peers. Advisory + /// (`Relaxed` `fetch_add`): rotates the starting offset into the + /// candidate set so concurrent cap-bound reclaims from this partition + /// do not all probe the lowest-numbered peer first. No correctness + /// invariant rides it — `try_reclaim_one` is the authoritative gate. + pub(crate) peer_cursor: std::sync::atomic::AtomicUsize, } impl std::fmt::Debug for PartitionState { @@ -361,6 +367,7 @@ impl PartitionRegistry { nic: p.nic, authorities: std::sync::Mutex::new(std::collections::HashMap::new()), make_stack, + peer_cursor: std::sync::atomic::AtomicUsize::new(0), }); if by_id.insert(p.id, state).is_some() { panic!("duplicate PartitionId declared: {:?}", p.id); @@ -394,6 +401,62 @@ impl PartitionRegistry { pub(crate) fn partitions(&self) -> impl Iterator> { self.by_id.values() } + + /// Resolve a partition by id without panicking. `None` if not declared. + pub(crate) fn partition_opt(&self, id: PartitionId) -> Option<&std::sync::Arc> { + self.by_id.get(&id) + } + + /// Partition ids sharing `id`'s NIC group, excluding `id` itself. + /// + /// Reclaim is NIC-blind for the freed *permit* (P0 connects on its own + /// NIC), but candidate peers are still drawn from the same NIC group: + /// the registry only groups by NIC, and a freed permit from any + /// same-group peer is equivalent. Empty if `id` is alone in its group + /// (e.g. the single-partition default). + pub(crate) fn nic_group_peers(&self, id: PartitionId) -> Vec { + let nic = match self.by_id.get(&id) { + Some(state) => &state.nic, + None => return Vec::new(), + }; + self.by_nic + .get(nic) + .map(|ids| ids.iter().copied().filter(|p| *p != id).collect()) + .unwrap_or_default() + } + + /// Attempt to reclaim one idle connection from `peer`'s entry for + /// `key`, freeing its permit. Returns `true` if one was freed. The + /// peer's `authorities` lock is held only for the `try_reclaim_one` + /// call (which pops + drops under the cache lock, releasing it before + /// the drop). No-op `false` if the peer or entry is absent. + pub(crate) fn try_reclaim_on(&self, peer: PartitionId, key: &super::PoolKey) -> bool { + let state = match self.by_id.get(&peer) { + Some(s) => s, + None => return false, + }; + let auth = state.authorities.lock().expect("authorities poisoned"); + match auth.get(key) { + Some(entry) => entry.try_reclaim_one(), + None => false, + } + } + + /// Attempt to reclaim one idle connection from *any* of `peer`'s + /// entries, freeing its permit. Returns `true` at the first entry that + /// yields. Drives the `Global` constraint, where the freed permit is + /// fungible across authorities — so the specific authority does not + /// matter, and this sidesteps reconstructing a `PoolKey` (scheme + + /// authority) from the authority-only stats index. No-op `false` if + /// the peer is absent or holds no reclaimable idle. + pub(crate) fn try_reclaim_any(&self, peer: PartitionId) -> bool { + let state = match self.by_id.get(&peer) { + Some(s) => s, + None => return false, + }; + let auth = state.authorities.lock().expect("authorities poisoned"); + auth.values().any(|entry| entry.try_reclaim_one()) + } } #[cfg(test)] diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs index b03d598398a..e1e71f307ad 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs @@ -238,6 +238,75 @@ impl StatsIndex { .unwrap_or(0) } + /// Partitions with at least one idle connection to `authority`, as + /// `(partition, idle_count)`. Advisory: a relaxed snapshot that + /// *narrows* reclaim/borrow candidates — the cache pop is the + /// authoritative confirmation. Prunes dead `Weak`s under the lock, + /// same as [`Self::snapshot`]. + pub(crate) fn idle_partitions_for(&self, authority: &Authority) -> Vec<(PartitionId, usize)> { + let handles: Vec<(PartitionId, Arc)> = { + let mut idx = self.inner.lock().expect("stats index poisoned"); + match idx.get_mut(authority) { + Some(a) => { + let mut live = Vec::new(); + a.by_partition.retain(|partition, weak| { + if let Some(strong) = weak.upgrade() { + live.push((*partition, strong)); + true + } else { + false + } + }); + if a.by_partition.is_empty() { + idx.remove(authority); + } + live + } + None => Vec::new(), + } + }; // lock released before loading atomics + handles + .into_iter() + .filter_map(|(p, c)| { + let established = c.established.load(Ordering::Relaxed); + let active = c.active.load(Ordering::Relaxed); + let idle = established.saturating_sub(active); + (idle > 0).then_some((p, idle)) + }) + .collect() + } + + /// All `(authority, partition)` cells with at least one idle + /// connection. Drives `Global`-constraint reclaim, where a freed + /// permit is fungible across authorities. Advisory/narrowing, same + /// contract as [`Self::idle_partitions_for`]. + pub(crate) fn idle_cells(&self) -> Vec<(Authority, PartitionId)> { + let handles: Vec<(Authority, PartitionId, Arc)> = { + let mut idx = self.inner.lock().expect("stats index poisoned"); + let mut live = Vec::new(); + idx.retain(|authority, a| { + a.by_partition.retain(|partition, weak| { + if let Some(strong) = weak.upgrade() { + live.push((authority.clone(), *partition, strong)); + true + } else { + false + } + }); + !a.by_partition.is_empty() + }); + live + }; // lock released before loading atomics + handles + .into_iter() + .filter_map(|(authority, p, c)| { + let established = c.established.load(Ordering::Relaxed); + let active = c.active.load(Ordering::Relaxed); + (established.saturating_sub(active) > 0).then_some((authority, p)) + }) + .collect() + } + /// Snapshot one authority's per-partition counters. /// /// Under the lock: upgrades each `Weak`, prunes dead entries (removing the From a437080619a3a5e8d37cea401b946674f645cf42 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 9 Jun 2026 12:04:00 -0400 Subject: [PATCH 35/47] reclaim on cap --- .../aws-smithy-http-client/src/client/pool.rs | 3 +- .../src/client/pool/handshake.rs | 61 +++++- .../tests/pool_behavior_test.rs | 173 ++++++++++++++++++ 3 files changed, 228 insertions(+), 9 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 04da1d67dc5..3348cf3b948 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -268,7 +268,7 @@ pub(crate) type MakeStack = /// Which semaphore bound a connect attempt at the cap, identifying what a /// reclaim must free to relieve it. Acquire order is per-host then global -/// (§8), so a per-host failure yields `PerHost`, and a failure on the +/// so a per-host failure yields `PerHost`, and a failure on the /// global semaphore while already holding the per-host permit yields /// `Global`. #[derive(Debug, Clone)] @@ -492,6 +492,7 @@ where shared.global_sem.clone(), per_host_sem, counters.clone(), + shared.reclaim_handle(partition_id), ); let pool_hooks = shared.hooks.clone(); diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index a4338824312..24d2688aac5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -103,6 +103,10 @@ pub(crate) struct ConnectionLimit { global: Option>, per_host: Option>, counters: Arc, + /// Cross-partition active-reclaim handle. `None` when no reclaim peer + /// exists or the pool is mid-teardown; the cap path is then a plain + /// blocking acquire. + reclaim: Option, } impl ConnectionLimit { @@ -111,12 +115,14 @@ impl ConnectionLimit { global: Option>, per_host: Option>, counters: Arc, + reclaim: Option, ) -> Self { Self { inner, global, per_host, counters, + reclaim, } } } @@ -128,6 +134,7 @@ impl Clone for ConnectionLimit { global: self.global.clone(), per_host: self.per_host.clone(), counters: self.counters.clone(), + reclaim: self.reclaim.clone(), } } } @@ -152,17 +159,28 @@ where let global = self.global.clone(); let per_host = self.per_host.clone(); let counters = self.counters.clone(); + let reclaim = self.reclaim.clone(); Box::pin(async move { - // Acquire per-host permit first, then global. This ordering - // prevents deadlock: a request never holds a global permit - // while waiting on a per-host permit (which could starve - // other hosts' requests that need global permits). - let per_host_permit = match per_host { - Some(sem) => Some(sem.acquire_owned().await.map_err(|_| "pool closed")?), + // Per-host before global: never hold a global permit while + // waiting on a per-host permit. + let per_host_permit = match &per_host { + Some(sem) => Some( + acquire_or_reclaim(sem, reclaim.as_ref(), || { + // `send_request` validated the URI has scheme+authority + // before dispatching, so this cannot fail here. + let key = super::PoolKey::from_uri(&ctx.uri) + .expect("connect URI has scheme+authority"); + super::BindingConstraint::PerHost(key) + }) + .await?, + ), None => None, }; - let global_permit = match global { - Some(sem) => Some(sem.acquire_owned().await.map_err(|_| "pool closed")?), + let global_permit = match &global { + Some(sem) => Some( + acquire_or_reclaim(sem, reclaim.as_ref(), || super::BindingConstraint::Global) + .await?, + ), None => None, }; let permit = Arc::new(ConnectionPermit::new(global_permit, per_host_permit)); @@ -193,6 +211,31 @@ where } } +/// Acquire one owned permit. Fast path is `try_acquire_owned`. On +/// `NoPermits`, free one peer's idle connection for `constraint` (inline, +/// best-effort) then blocking-acquire — the blocking acquire is the +/// authoritative take regardless of whether reclaim freed a permit. +/// `constraint` is built only on the cap-bound branch. +async fn acquire_or_reclaim( + sem: &Arc, + reclaim: Option<&super::PeerReclaimHandle>, + constraint: impl FnOnce() -> super::BindingConstraint, +) -> Result { + match sem.clone().try_acquire_owned() { + Ok(permit) => Ok(permit), + Err(tokio::sync::TryAcquireError::NoPermits) => { + if let Some(reclaim) = reclaim { + reclaim.try_free_under_load(&constraint()); + } + sem.clone() + .acquire_owned() + .await + .map_err(|_| "pool closed".into()) + } + Err(tokio::sync::TryAcquireError::Closed) => Err("pool closed".into()), + } +} + // --------------------------------------------------------------------------- // Tower Service wrappers for hyper's SendRequest // --------------------------------------------------------------------------- @@ -586,6 +629,7 @@ mod tests { None, None, Arc::new(super::super::stats::ConnectionCounters::default()), + None, ); let sleep = SharedAsyncSleep::new(TokioSleep::new()); let ctx = ConnectCtx::new( @@ -622,6 +666,7 @@ mod tests { None, None, Arc::new(super::super::stats::ConnectionCounters::default()), + None, ); let ctx = ConnectCtx::new("http://example.com".parse().unwrap(), None); let fut = svc.call(ctx); diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index d8140944d8f..5819a383f5c 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1713,6 +1713,179 @@ async fn partition_spawners_are_isolated() { ); } +// --------------------------------------------------------------------------- +// Test implementations: cross-partition active reclaim (Never policy) +// --------------------------------------------------------------------------- + +/// Under cap pressure, a starved partition reclaims an over-supplied peer's +/// idle connection (freeing its permit) and connects locally — rather than +/// waiting for passive idle eviction. +/// +/// Setup distinguishes *active* reclaim from passive eviction: the pool's +/// idle timeout is set very long, so the eviction tick will not fire during +/// the test. P1 establishes an idle connection holding the single global +/// permit; P0 (same authority, same NIC group) then requests, finds the cap +/// bound, reclaims P1's *non-expired* idle, and connects on a fresh +/// connection. Proven via the `ConnectionEventListener`: +/// - P1's connection closes with `CloseReason::Reclaimed` (not +/// `IdleTimeout` — the tick never ran), +/// - P0's connection is freshly created (distinct `conn_id`), +/// - the server accepted two connections (P1's reclaimed, P0's new), +/// - P0 completed well under the idle timeout. +#[tokio::test] +async fn v2_cross_partition_reclaim_frees_peer_idle() { + use aws_smithy_http_client::pool::{ + CloseReason, ConnectionClosedEvent, ConnectionCreatedEvent, ConnectionEventListener, + ConnectionFailedEvent, ConnectionReusedEvent, Partition, PartitionId, TokioDriverSpawner, + }; + use std::sync::{Arc, Mutex}; + + #[derive(Debug, Clone, Default)] + struct RecordingListener { + created: Arc>>, + closed: Arc>>, + } + impl ConnectionEventListener for RecordingListener { + fn on_created(&self, event: &ConnectionCreatedEvent) { + self.created.lock().unwrap().push(( + event.conn_id().to_string().parse().unwrap(), + event.authority().as_str().to_string(), + )); + } + fn on_reused(&self, _event: &ConnectionReusedEvent) {} + fn on_closed(&self, event: &ConnectionClosedEvent) { + self.closed.lock().unwrap().push(( + event.conn_id().to_string().parse().unwrap(), + event.authority().as_str().to_string(), + event.reason(), + )); + } + fn on_connection_failed(&self, _event: &ConnectionFailedEvent) {} + } + + // One endpoint, two connections total: P1's (later reclaimed) and P0's. + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p1", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p0", + }, + ], + ) + .build() + .await; + + let listener = Arc::new(RecordingListener::default()); + + // Two partitions sharing one NIC group (so they are reclaim peers), + // global cap of 1 (so the second partition is cap-bound), and a long + // idle timeout so passive eviction never fires during the test. + let p0 = Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()) + .interface("eth-test"); + let p1 = Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()) + .interface("eth-test"); + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .connection_event_listener(listener.clone() as Arc) + .max_connections(1) + .pool_idle_timeout(Duration::from_secs(3600)) + .partitions([p0, p1]) + .build_http(); + + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + let client0 = PoolClient::from_partition(&pool, PartitionId::from_index(0)).into_shared(); + let client1 = PoolClient::from_partition(&pool, PartitionId::from_index(1)).into_shared(); + + // P1 establishes a connection and returns it idle to its cache, holding + // the single global permit. + let (_, _, meta1) = send_with_capture(&client1, &url).await; + let p1_conn_id: u64 = meta1 + .expect("p1 metadata") + .connection_id() + .expect("p1 conn id") + .to_string() + .parse() + .unwrap(); + + // P0 requests the same authority. The global permit is held by P1's + // idle connection, so P0 is cap-bound; it reclaims P1's idle inline, + // freeing the permit, then connects locally. Bounded so a hang (failed + // reclaim → indefinite block) surfaces as a test timeout rather than + // a stall. + let started = std::time::Instant::now(); + let (status, body, meta0) = + tokio::time::timeout(Duration::from_secs(5), send_with_capture(&client0, &url)) + .await + .expect("p0 must not block indefinitely — reclaim should free a permit"); + assert_eq!(status, 200); + assert_eq!(body, b"p0"); + let p0_conn_id: u64 = meta0 + .expect("p0 metadata") + .connection_id() + .expect("p0 conn id") + .to_string() + .parse() + .unwrap(); + + // P0 completed promptly — far under the 3600s idle timeout, so this was + // active reclaim, not passive eviction. + assert!( + started.elapsed() < Duration::from_secs(60), + "p0 should complete promptly via reclaim, took {:?}", + started.elapsed() + ); + + // P0 ran on a *fresh local* connection, not P1's (that is reclaim, not + // borrow — borrow would reuse P1's exact connection). + assert_ne!( + p0_conn_id, p1_conn_id, + "reclaim gives P0 its own fresh connection, distinct from P1's" + ); + + // The server accepted two connections: P1's (reclaimed) and P0's (new). + assert_eq!( + harness.tcp_accepted_count(), + 2, + "expected P1's connection + P0's fresh connection" + ); + + // The listener proves the mechanism directly: P1's connection closed + // with `Reclaimed` (not `IdleTimeout` — the tick never ran), and P0's + // connection was created fresh. + let closed = listener.closed.lock().unwrap(); + let reclaimed: Vec<_> = closed + .iter() + .filter(|(_, _, reason)| matches!(reason, CloseReason::Reclaimed)) + .collect(); + assert_eq!( + reclaimed.len(), + 1, + "exactly one connection should close with Reclaimed, got: {closed:?}" + ); + assert_eq!( + reclaimed[0].0, p1_conn_id, + "the reclaimed connection must be P1's" + ); + assert!( + closed + .iter() + .all(|(_, _, reason)| !matches!(reason, CloseReason::IdleTimeout)), + "no IdleTimeout close — the long idle timeout means reclaim, not passive eviction" + ); + + let created = listener.created.lock().unwrap(); + assert!( + created.iter().any(|(id, _)| *id == p0_conn_id), + "P0's fresh connection should fire on_created" + ); +} + // --------------------------------------------------------------------------- // Test implementations: stats read API // --------------------------------------------------------------------------- From d49918b920517d08810b010fe0774e66f71a916c Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 10 Jun 2026 10:17:18 -0400 Subject: [PATCH 36/47] add peer borrow walk --- .../aws-smithy-http-client/src/client/pool.rs | 340 ++++++++++++++++-- .../src/client/pool/connection.rs | 55 +++ .../src/client/pool/handshake.rs | 23 +- .../src/client/pool/partition.rs | 16 + .../src/client/pool/vendored_cache.rs | 19 + .../tests/pool_behavior_test.rs | 189 ++++++++++ 6 files changed, 607 insertions(+), 35 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 3348cf3b948..cb88686d87c 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -219,6 +219,29 @@ impl PoolKey { } } +/// Type-erased, single-use handle to a checked-out connection that can +/// dispatch one request. +/// +/// A borrowed peer connection crosses the peer's `dyn PoolEntry` boundary +/// as one of these: the concrete checkout type (`H1Checkout`, holding the +/// peer's `CachedConnection`) carries upstream-unnameable parameters, so +/// it is boxed. Dispatching consumes the handle; the response body holds +/// the underlying connection alive and returns it to *the peer's* pool on +/// drop (it was never the borrower's). Boxed only on the actual-borrow +/// path — the common local path dispatches through the concrete checkout +/// with no erasure. +pub(crate) trait DispatchConn: Send { + /// Liveness check before dispatch. `Err` means the connection is dead; + /// the caller drops the handle and falls back. + fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll>; + + /// Dispatch one request, erasing the response body to `SdkBody`. + fn dispatch( + self: Box, + req: http_1x::Request, + ) -> BoxFuture, BoxError>>; +} + /// Type-erased handle to one Negotiate leg's eviction interface. /// /// Each `TypedPoolEntry` holds at most two of these (one for its H1 cache, @@ -239,6 +262,13 @@ pub(crate) struct BoxedRetainer { /// simply-reclaimable idle (an idle H2 connection still multiplexes), /// so its `reclaim_fn` is a no-op returning `false`. reclaim_fn: Box bool + Send + Sync + 'static>, + /// Take one idle connection wrapped as a dispatchable handle that + /// returns to *this* (the owner's) pool on drop — cross-partition + /// borrow. Returns `None` if no idle connection is available. The H1 + /// cache leg checks out an idle connection; the H2 singleton leg is a + /// no-op returning `None` (H2 borrow is not supported — an idle H2 + /// connection still multiplexes on its owner). + borrow_fn: Box Option> + Send + Sync + 'static>, } impl BoxedRetainer { @@ -253,6 +283,10 @@ impl BoxedRetainer { fn reclaim_one(&self) -> bool { (self.reclaim_fn)() } + + fn borrow_one(&self) -> Option> { + (self.borrow_fn)() + } } /// Per-host retainer registry. Bounded at two entries (one H1 cache, one @@ -374,6 +408,75 @@ impl PeerReclaimHandle { } } +/// Handle for cross-partition borrow (`PreferLocal`), held by the pool +/// entry under a binding cap. +/// +/// At a cap-bound local checkout (`AcquireMode::NonBlocking` returned +/// `CapBound`), the requesting partition uses this to borrow one peer's +/// idle connection and dispatch its request through it — no permit moves, +/// no cold start. Response-shaped and NIC-*bounded*: the borrowed +/// connection physically lives on the peer's NIC, so candidates are drawn +/// only from the same NIC group (unlike reclaim, which frees a fungible +/// permit and is NIC-blind). Always keyed to the requested authority (the +/// borrowed connection must already be connected to that host). +#[derive(Clone)] +pub(crate) struct PeerBorrowHandle { + /// `Weak` to avoid a cycle: the registry transitively owns the + /// partitions whose entries hold this handle. + registry: std::sync::Weak, + /// Owned (Arc) so the handle does not borrow `SharedPoolState`. + stats_index: Arc, + /// The requesting partition — excluded from its own candidate walk. + self_partition: PartitionId, +} + +impl PeerBorrowHandle { + /// Borrow one same-NIC-group peer's idle connection to `key`'s + /// authority, as a dispatchable handle. `None` if the pool is gone, + /// there are no NIC-group peers, or no peer holds idle to this + /// authority (the caller then falls back to a blocking local acquire). + pub(crate) fn try_borrow(&self, key: &PoolKey) -> Option> { + let registry = self.registry.upgrade()?; + let peers = registry.nic_group_peers(self.self_partition); + if peers.is_empty() { + return None; // alone in the NIC group (e.g. single-partition default) + } + // Candidates = peers with idle to this authority (index narrows; + // the cache checkout confirms). Round-robin start offset. + let authority = Authority::new(key.authority.as_str()); + let mut candidates: Vec = self + .stats_index + .idle_partitions_for(&authority) + .into_iter() + .map(|(p, _idle)| p) + .filter(|p| *p != self.self_partition && peers.contains(p)) + .collect(); + self.rotate(&mut candidates); + candidates + .into_iter() + .find_map(|peer| registry.try_borrow_on(peer, key)) + } + + /// Rotate the candidate vec by the partition's advisory `peer_cursor` + /// (shared with reclaim), so concurrent borrows from this partition do + /// not all probe the lowest-numbered candidate first. + fn rotate(&self, candidates: &mut [PartitionId]) { + if candidates.len() < 2 { + return; + } + if let Some(registry) = self.registry.upgrade() { + if let Some(state) = registry.partition_opt(self.self_partition) { + let n = candidates.len(); + let start = state + .peer_cursor + .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + % n; + candidates.rotate_left(start); + } + } + } +} + /// Connection-lifecycle machinery shared across every partition: event /// hooks, the global connection budget, and the per-host budgets. One /// instance per pool, held behind `Arc` on `ConnectionPool`. @@ -424,6 +527,19 @@ impl SharedPoolState { }) } + /// A borrow handle for the given requesting partition, if the registry + /// is bound and still alive. `None` when no registry is set or the + /// pool has been dropped (borrow is then unavailable; the caller + /// blocks locally). + pub(crate) fn borrow_handle(&self, self_partition: PartitionId) -> Option { + let registry = self.registry.get()?.clone(); + Some(PeerBorrowHandle { + registry, + stats_index: self.stats_index.clone(), + self_partition, + }) + } + /// The shared per-host semaphore for `key`, created on first request /// for that key. Returns `None` when no per-host limit is configured. /// Shared across all partitions: two partitions to the same host @@ -472,6 +588,7 @@ where -> MakeStack { let connector = connector.clone(); let spawner = spawner.clone(); + let cross_partition_policy = cross_partition_policy; Arc::new( move |uri: &http_1x::Uri, shared: &SharedPoolState| -> Box { let authority = Authority::new( @@ -534,6 +651,7 @@ where Arc::new(std::sync::Mutex::new(cache.clone())); let cache_for_empty = cache_for_retain.clone(); let cache_for_reclaim = cache_for_retain.clone(); + let cache_for_borrow = cache_for_retain.clone(); retainers.lock().expect("retainer slot poisoned").push( BoxedRetainer { retain_fn: Box::new({ @@ -629,6 +747,33 @@ where } } }), + borrow_fn: Box::new({ + let listener = pool_hooks.listener.clone(); + let counters = counters.clone(); + move || { + // Take an idle connection wrapped so it + // returns to THIS (the owner's) cache on + // drop — the borrower dispatches one + // request through it but never owns it. + // `None` if no idle connection is + // available (borrower falls through to + // the next peer or to a blocking local + // acquire). + let cached = cache_for_borrow + .lock() + .expect("borrow cache lock poisoned") + .try_checkout_idle()?; + let checkout = H1Checkout::<()>::new( + CachedConnection::new( + cached, + listener.clone(), + counters.clone(), + ), + ); + Some(Box::new(checkout) + as Box) + } + }), }, ); cache.map_response({ @@ -732,6 +877,11 @@ where // not simply-reclaimable like an H1 cache entry. // Reclaim skips the H2 leg. reclaim_fn: Box::new(|| false), + // H2 borrow is unsupported: an idle H2 + // connection still multiplexes on its owner, + // so there is no idle connection to hand to a + // borrower. Borrow skips the H2 leg. + borrow_fn: Box::new(|| None), }, ); singleton.map_response({ @@ -752,6 +902,12 @@ where stack, retainers, counters, + borrow: match cross_partition_policy { + partition::CrossPartitionPolicy::PreferLocal => { + shared.borrow_handle(partition_id) + } + partition::CrossPartitionPolicy::Never => None, + }, }) }, ) @@ -1109,6 +1265,35 @@ impl Service> for H1Checkout DispatchConn for H1Checkout +where + UnusedH2Phantom: Send + Sync + 'static, +{ + fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll> { + Service::poll_ready(self, cx) + } + + fn dispatch( + mut self: Box, + req: http_1x::Request, + ) -> BoxFuture, BoxError>> { + // `H1Checkout::call` produces `CheckoutResponse<…>` whose body is + // the internal `GuardedBody`; erase it to `SdkBody` at this + // boundary, exactly as `TypedPoolEntry::send` does for the local + // path. The guard inside lives on in `SdkBody` until the body is + // drained, returning the connection to the peer's pool. + let fut = Service::call(&mut *self, req); + Box::pin(async move { + let resp = fut.await?; + let (parts, body) = resp.into_parts(); + Ok(http_1x::Response::from_parts( + parts, + SdkBody::from_body_1_x(body), + )) + }) + } +} + /// Tower `Service` wrapper for an H2 pool checkout. /// /// Symmetric to `H1Checkout`. The `Singleton` type parameter is the @@ -1220,6 +1405,14 @@ pub(crate) trait PoolEntry: Send + Sync { /// [`CloseReason::Reclaimed`] for the freed connection. Best-effort: /// returns `false` if the entry has no reclaimable idle. fn try_reclaim_one(&self) -> bool; + + /// Take one idle connection as a dispatchable handle that returns to + /// this entry's pool on drop, for cross-partition borrow. The borrower + /// dispatches one request through it; the connection stays this + /// entry's (no permit moves). Returns `None` if no idle connection is + /// available. Best-effort — only the H1 cache leg yields a handle (H2 + /// idle still multiplexes on its owner, so the H2 leg returns `None`). + fn try_borrow_one(&self) -> Option>; } /// Concrete PoolEntry wrapping a Negotiate stack. @@ -1237,6 +1430,11 @@ struct TypedPoolEntry { // active write-handle; read API added in a later change #[allow(dead_code)] counters: Arc, + /// Cross-partition borrow handle. `Some` only under + /// `CrossPartitionPolicy::PreferLocal`; `None` under `Never` (and for + /// the single-partition default, where it would have no NIC-group + /// peers anyway). Drives the `send` cap-bound borrow branch. + borrow: Option, } impl PoolEntry for TypedPoolEntry @@ -1257,9 +1455,15 @@ where req: http_1x::Request, ) -> BoxFuture, BoxError>> { let mut svc = self.stack.clone(); + let borrow = self.borrow.clone(); Box::pin(async move { - // Checkout loop. Pops past stale idle entries (poisoned or dead). - // Converges because: + // Local checkout. The post-checkout `poll_ready` is the + // reactive health check: the composable pool has no proactive + // checkout-time health check, so a popped idle connection may + // be dead (server closed it, driver gone). `poll_ready` Err + // identifies that; the loop discards it and pops the next, + // until a live connection is checked out or the connect path + // is reached. Converges because: // - Cache idle set is bounded; each post-checkout `poll_ready` // Err flips `is_closed` so `Cached::Drop` skips reinsertion, // shrinking the set by one. @@ -1267,33 +1471,102 @@ where // the next `call` to run a fresh handshake; a handshake // failure surfaces as an error from `svc.call` (not the // inner `poll_ready`), which we propagate. - let mut conn = loop { - std::future::poll_fn(|cx| svc.poll_ready(cx)) - .await - .map_err(Into::into)?; - let mut checkout = svc.call(ctx.clone()).await.map_err(Into::into)?; - if std::future::poll_fn(|cx| checkout.poll_ready(cx)) - .await - .is_ok() - { - break checkout; + // Under `AcquireMode::NonBlocking`, a cache miss at the connect + // path returns `CapBound` instead of blocking — surfaced here + // so the caller can try a peer borrow before committing to wait. + async fn local_checkout( + svc: &mut S, + ctx: ConnectCtx, + ) -> Result + where + S: Service, + S::Error: Into, + Conn: + Service, Response = CheckoutResponse>, + Conn::Error: Into, + { + loop { + std::future::poll_fn(|cx| svc.poll_ready(cx)) + .await + .map_err(Into::into)?; + let mut checkout = svc.call(ctx.clone()).await.map_err(Into::into)?; + if std::future::poll_fn(|cx| checkout.poll_ready(cx)) + .await + .is_ok() + { + return Ok(checkout); + } + // drop `checkout` → pool cleanup (H1: discard via + // CachedConnection::Drop; H2: Singleton already cleared). } - // drop `checkout` → triggers pool cleanup (H1: discard via - // CachedConnection::Drop; H2: Singleton already cleared - // state). - }; - - // Dispatch the request; convert the body into `SdkBody` at the - // boundary so consumers (and `dyn PoolEntry`) don't see our - // internal `GuardedBody<...>` type. The guard inside - // `GuardedBody` lives on inside `SdkBody` until the body is - // fully drained. - let resp = conn.call(req).await.map_err(Into::into)?; - let (parts, body) = resp.into_parts(); - Ok(http_1x::Response::from_parts( - parts, - SdkBody::from_body_1_x(body), - )) + } + + // Erase a checkout's `CheckoutResponse` body to `SdkBody` at the + // `dyn PoolEntry` boundary so consumers never see the internal + // `GuardedBody<...>`. The guard inside lives on in `SdkBody` + // until the body drains, returning the connection to its pool. + fn erase_body( + resp: CheckoutResponse, + ) -> http_1x::Response + where + PoolUnnameable: Send + Sync + 'static, + { + let (parts, body) = resp.into_parts(); + http_1x::Response::from_parts(parts, SdkBody::from_body_1_x(body)) + } + + match &borrow { + // PreferLocal: probe locally without blocking; on a + // cap-bound miss, borrow a peer's idle connection before + // falling back to a blocking local acquire. + Some(handle) => { + let probe = ctx.clone().with_mode(connection::AcquireMode::NonBlocking); + match local_checkout(&mut svc, probe).await { + Ok(mut conn) => { + let resp = conn.call(req).await.map_err(Into::into)?; + Ok(erase_body(resp)) + } + Err(err) if connection::CapBound::is(&*err) => { + // Cap-bound locally. Try to borrow a same-NIC + // peer's idle connection to this authority. + let key = PoolKey::from_uri(&ctx.uri) + .ok_or("request URI must have scheme and authority")?; + if let Some(mut borrowed) = handle.try_borrow(&key) { + // Single-shot liveness check. On death, + // fall through to the authoritative local + // acquire rather than walk further peers. + if std::future::poll_fn(|cx| borrowed.poll_ready(cx)) + .await + .is_ok() + { + // Dispatch through the peer's connection; + // its driver stays on the peer's runtime, + // and the connection returns to the peer's + // pool when the body drains. + return borrowed.dispatch(req).await; + } + // Dead borrowed connection drops here: + // `CachedConnection::Drop` discards it from + // the peer's pool and balances the peer's + // `active`. + } + // Borrow miss → authoritative blocking local + // acquire (today's wait; includes reclaim). + let mut conn = local_checkout(&mut svc, ctx).await?; + let resp = conn.call(req).await.map_err(Into::into)?; + Ok(erase_body(resp)) + } + Err(err) => Err(err), + } + } + // Never (and single-partition default): blocking local + // acquire, unchanged. + None => { + let mut conn = local_checkout(&mut svc, ctx).await?; + let resp = conn.call(req).await.map_err(Into::into)?; + Ok(erase_body(resp)) + } + } }) } @@ -1316,6 +1589,14 @@ where // simply-reclaimable idle). retainers.iter().any(|r| r.reclaim_one()) } + + fn try_borrow_one(&self) -> Option> { + let retainers = self.retainers.lock().expect("retainer slot poisoned"); + // Return the first leg that yields a borrowable handle. The H1 + // cache leg checks out an idle connection; the H2 leg is a no-op + // returning `None`. + retainers.iter().find_map(|r| r.borrow_one()) + } } #[cfg(test)] @@ -1354,6 +1635,9 @@ mod tests { fn try_reclaim_one(&self) -> bool { false } + fn try_borrow_one(&self) -> Option> { + None + } } fn pool_with_config(config: PoolConfig) -> Arc { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index fa230a816d0..159c70c6c1a 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -358,6 +358,54 @@ pub(crate) struct ConnectCtx { /// means no connect timeout; cached connections skip the connector /// entirely so this is automatically a no-op on cache hit. pub(crate) connect_timeout: Option, + /// How the connect path behaves when a permit cannot be acquired. + pub(crate) mode: AcquireMode, +} + +/// Behavior of the connect path when the connection cap is reached. +/// +/// Selected per request (carried on [`ConnectCtx`]) because the same +/// partition stack is exercised twice under `PreferLocal`: once +/// `NonBlocking` to probe for local capacity, then `Blocking` as the +/// fallback after a peer-borrow miss. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) enum AcquireMode { + /// Block on the semaphore until a permit is free (active reclaim, then + /// FIFO wait). The authoritative take. + #[default] + Blocking, + /// Return [`CapBound`] immediately on `NoPermits` instead of blocking, + /// so the caller can try borrowing a peer's connection first. + NonBlocking, +} + +/// Sentinel error from the connect path under [`AcquireMode::NonBlocking`]: +/// the connection cap is reached and no permit is available. Distinct from +/// any real connect failure, and distinct from negotiate's internal +/// `UseOther` sentinel, so it propagates verbatim up the stack. +#[derive(Debug)] +pub(crate) struct CapBound; + +impl std::fmt::Display for CapBound { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("connection cap reached (non-blocking acquire)") + } +} + +impl std::error::Error for CapBound {} + +impl CapBound { + /// Whether `err` is a `CapBound` sentinel anywhere in its chain. + pub(crate) fn is(err: &(dyn std::error::Error + 'static)) -> bool { + let mut e: Option<&(dyn std::error::Error + 'static)> = Some(err); + while let Some(cur) = e { + if cur.is::() { + return true; + } + e = cur.source(); + } + false + } } impl ConnectCtx { @@ -365,8 +413,15 @@ impl ConnectCtx { Self { uri, connect_timeout, + mode: AcquireMode::Blocking, } } + + /// Set the acquire mode (defaults to [`AcquireMode::Blocking`]). + pub(crate) fn with_mode(mut self, mode: AcquireMode) -> Self { + self.mode = mode; + self + } } /// Request extension set by the adapter to hint a read timeout to the diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 24d2688aac5..202b1d7b6db 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -161,11 +161,12 @@ where let counters = self.counters.clone(); let reclaim = self.reclaim.clone(); Box::pin(async move { + let mode = ctx.mode; // Per-host before global: never hold a global permit while // waiting on a per-host permit. let per_host_permit = match &per_host { Some(sem) => Some( - acquire_or_reclaim(sem, reclaim.as_ref(), || { + acquire_or_reclaim(sem, reclaim.as_ref(), mode, || { // `send_request` validated the URI has scheme+authority // before dispatching, so this cannot fail here. let key = super::PoolKey::from_uri(&ctx.uri) @@ -178,8 +179,10 @@ where }; let global_permit = match &global { Some(sem) => Some( - acquire_or_reclaim(sem, reclaim.as_ref(), || super::BindingConstraint::Global) - .await?, + acquire_or_reclaim(sem, reclaim.as_ref(), mode, || { + super::BindingConstraint::Global + }) + .await?, ), None => None, }; @@ -212,18 +215,24 @@ where } /// Acquire one owned permit. Fast path is `try_acquire_owned`. On -/// `NoPermits`, free one peer's idle connection for `constraint` (inline, -/// best-effort) then blocking-acquire — the blocking acquire is the -/// authoritative take regardless of whether reclaim freed a permit. -/// `constraint` is built only on the cap-bound branch. +/// `NoPermits`: under [`AcquireMode::NonBlocking`] return [`CapBound`] +/// immediately (the caller will try a peer borrow); under +/// [`AcquireMode::Blocking`] free one peer's idle connection for +/// `constraint` (inline, best-effort) then blocking-acquire — the blocking +/// acquire is the authoritative take regardless of whether reclaim freed a +/// permit. `constraint` is built only on the blocking cap-bound branch. async fn acquire_or_reclaim( sem: &Arc, reclaim: Option<&super::PeerReclaimHandle>, + mode: super::connection::AcquireMode, constraint: impl FnOnce() -> super::BindingConstraint, ) -> Result { match sem.clone().try_acquire_owned() { Ok(permit) => Ok(permit), Err(tokio::sync::TryAcquireError::NoPermits) => { + if mode == super::connection::AcquireMode::NonBlocking { + return Err(super::connection::CapBound.into()); + } if let Some(reclaim) = reclaim { reclaim.try_free_under_load(&constraint()); } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index a25351fb5a9..98461f233b3 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -457,6 +457,22 @@ impl PartitionRegistry { let auth = state.authorities.lock().expect("authorities poisoned"); auth.values().any(|entry| entry.try_reclaim_one()) } + + /// Attempt to borrow one idle connection from `peer`'s entry for + /// `key`, as a dispatchable handle that returns to `peer`'s pool on + /// drop. Returns `None` if the peer or entry is absent, or holds no + /// borrowable idle. The peer's `authorities` lock is held only to + /// check out the handle (the handle re-pools on drop, independent of + /// the lock); dispatch happens after the lock is released. + pub(crate) fn try_borrow_on( + &self, + peer: PartitionId, + key: &super::PoolKey, + ) -> Option> { + let state = self.by_id.get(&peer)?; + let auth = state.authorities.lock().expect("authorities poisoned"); + auth.get(key)?.try_borrow_one() + } } #[cfg(test)] diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs index 83157653d70..65a9f06c4ac 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs @@ -197,6 +197,25 @@ mod internal { pub fn try_pop_idle(&self) -> Option { self.shared.lock().unwrap().services.pop() } + + // SDK MODIFICATION: added `try_checkout_idle` so a caller can take an + // idle cached service for one use and have it return to the pool on + // drop, without going through `Service::call` (which may also start a + // new connection when none is idle). + /// Take one idle cached service, if any, wrapped so it returns to the + /// pool on drop. + /// + /// Unlike [`Self::try_pop_idle`], which hands back the raw service + /// (dropping it drops the service), this returns the same + /// [`Cached`] wrapper [`Service::call`] produces: dropping it + /// re-inserts the service into the pool. Unlike [`Service::call`], + /// it never starts a new connection — it returns `None` when no + /// service is idle. Serialized against [`Self::retain`] and + /// [`Self::try_pop_idle`] by the shared `Mutex`. + pub fn try_checkout_idle(&self) -> Option> { + let inner = self.shared.lock().unwrap().services.pop()?; + Some(Cached::new(inner, Arc::downgrade(&self.shared))) + } } impl Service for Cache diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 5819a383f5c..f45be7158b5 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1886,6 +1886,195 @@ async fn v2_cross_partition_reclaim_frees_peer_idle() { ); } +// --------------------------------------------------------------------------- +// Test implementations: cross-partition borrow (PreferLocal policy) +// --------------------------------------------------------------------------- + +/// Under cap pressure with `PreferLocal`, a starved partition borrows a +/// same-NIC peer's idle connection and dispatches its request through it — +/// no new connection, no permit. This is the opposite disposition from +/// reclaim: the connection stays the peer's (P0 runs on P1's exact +/// connection), proven by `p0_conn_id == p1_conn_id` and a single TCP +/// accept (P1's connection serves both requests via keep-alive). +#[tokio::test] +async fn v2_cross_partition_borrow_reuses_peer_connection() { + use aws_smithy_http_client::pool::{ + CrossPartitionPolicy, Partition, PartitionId, TokioDriverSpawner, + }; + + // One endpoint, ONE connection serving TWO requests (P1's, then P0's + // borrowed request via the keep-alive loop). + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p1", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p0", + }, + ], + ) + .build() + .await; + + // Two partitions, same NIC group (borrow peers), global cap 1, long + // idle timeout so the eviction tick never interferes. + let p0 = Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()) + .interface("eth-test"); + let p1 = Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()) + .interface("eth-test"); + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .cross_partition_policy(CrossPartitionPolicy::PreferLocal) + .max_connections(1) + .pool_idle_timeout(Duration::from_secs(3600)) + .partitions([p0, p1]) + .build_http(); + + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + let client0 = PoolClient::from_partition(&pool, PartitionId::from_index(0)).into_shared(); + let client1 = PoolClient::from_partition(&pool, PartitionId::from_index(1)).into_shared(); + + // P1 establishes a connection, drains it, and returns it idle to its + // cache holding the single global permit. + let (_, _, meta1) = send_with_capture(&client1, &url).await; + let p1_conn_id: u64 = meta1 + .expect("p1 metadata") + .connection_id() + .expect("p1 conn id") + .to_string() + .parse() + .unwrap(); + + // P0 requests the same authority. The permit is held by P1's idle + // connection, so P0 is cap-bound; under PreferLocal it borrows P1's + // connection and dispatches through it. Bounded so a hang surfaces as + // a failure. + let (status, body, meta0) = + tokio::time::timeout(Duration::from_secs(5), send_with_capture(&client0, &url)) + .await + .expect("p0 must not block — borrow should reuse P1's connection"); + assert_eq!(status, 200); + assert_eq!(body, b"p0"); + let p0_conn_id: u64 = meta0 + .expect("p0 metadata") + .connection_id() + .expect("p0 conn id") + .to_string() + .parse() + .unwrap(); + + // The direct proof of borrow: P0 ran on P1's *exact* connection. + assert_eq!( + p0_conn_id, p1_conn_id, + "PreferLocal borrow dispatches P0's request through P1's connection" + ); + + // One TCP accept total — P1's connection served both requests. A fresh + // local connection for P0 (reclaim, or no borrow) would be 2. + assert_eq!( + harness.tcp_accepted_count(), + 1, + "borrow reuses the peer's connection; no new connection is opened" + ); +} + +/// Borrow is NIC-bounded: a peer on a different NIC is not a borrow +/// candidate. With `PreferLocal` but P0 and P1 on different NICs, P0 +/// cannot borrow P1's connection — nor reclaim its permit (reclaim +/// candidates are also drawn from the NIC group). P0's cap-bound wait is +/// instead released when P1's idle connection is evicted (scenario B), and +/// P0 then connects locally. Proven by `p0_conn_id != p1_conn_id` and two +/// TCP accepts — the opposite of the same-NIC borrow case, which reuses +/// P1's exact connection. +#[tokio::test] +async fn v2_cross_partition_borrow_respects_nic_boundary() { + use aws_smithy_http_client::pool::{ + CrossPartitionPolicy, Partition, PartitionId, TokioDriverSpawner, + }; + + // One endpoint, two connections: P1's, then P0's fresh local one. + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p1", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p0", + }, + ], + ) + .build() + .await; + + // Two partitions on DIFFERENT NICs — not borrow peers (and not reclaim + // peers). A short idle timeout lets P0's cap-bound wait be released by + // eviction of P1's idle connection, rather than hanging. + let p0 = Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()) + .interface("eth-zero"); + let p1 = Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()) + .interface("eth-one"); + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .cross_partition_policy(CrossPartitionPolicy::PreferLocal) + .max_connections(1) + .pool_idle_timeout(Duration::from_millis(150)) + .partitions([p0, p1]) + .build_http(); + + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + let client0 = PoolClient::from_partition(&pool, PartitionId::from_index(0)).into_shared(); + let client1 = PoolClient::from_partition(&pool, PartitionId::from_index(1)).into_shared(); + + let (_, _, meta1) = send_with_capture(&client1, &url).await; + let p1_conn_id: u64 = meta1 + .expect("p1 metadata") + .connection_id() + .expect("p1 conn id") + .to_string() + .parse() + .unwrap(); + + // P0 is cap-bound and cannot borrow across the NIC boundary; it waits + // for P1's idle to be evicted, then connects locally. Bounded so a true + // hang still surfaces as a failure. + let (status, body, meta0) = + tokio::time::timeout(Duration::from_secs(5), send_with_capture(&client0, &url)) + .await + .expect("p0 should proceed once P1's idle is evicted"); + assert_eq!(status, 200); + assert_eq!(body, b"p0"); + let p0_conn_id: u64 = meta0 + .expect("p0 metadata") + .connection_id() + .expect("p0 conn id") + .to_string() + .parse() + .unwrap(); + + // P1 is on a different NIC → not a borrow candidate. P0 does NOT run on + // P1's connection; it gets its own. + assert_ne!( + p0_conn_id, p1_conn_id, + "no borrow across NICs — P0 runs on its own connection" + ); + + // Two TCP accepts: P1's connection plus P0's fresh local one. + assert_eq!( + harness.tcp_accepted_count(), + 2, + "NIC boundary blocks borrow; P0 opens its own connection" + ); +} + // --------------------------------------------------------------------------- // Test implementations: stats read API // --------------------------------------------------------------------------- From 22d16960f0b020700a450e1e619575490ddfcc40 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 10 Jun 2026 12:18:17 -0400 Subject: [PATCH 37/47] add tsan and stress test. Workaround upstream issue --- .../aws-smithy-http-client/additional-ci | 16 ++++ .../aws-smithy-http-client/src/client/pool.rs | 44 ++++++++++- .../src/client/pool/connection.rs | 12 +-- .../tests/pool_behavior_test.rs | 78 +++++++++++++++++++ 4 files changed, 141 insertions(+), 9 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/additional-ci b/rust-runtime/aws-smithy-http-client/additional-ci index bba927553ee..a3540ce0ab3 100755 --- a/rust-runtime/aws-smithy-http-client/additional-ci +++ b/rust-runtime/aws-smithy-http-client/additional-ci @@ -21,3 +21,19 @@ cargo test --features wire-mock,default-client,test-util --test pool_behavior_te echo "### Running H2 pool behavior tests (requires aws_sdk_unstable cfg)" RUSTFLAGS='--cfg aws_sdk_unstable' cargo test --features s2n-tls,rustls-aws-lc,wire-mock,default-client,test-util --test h2_pool_test + +# ThreadSanitizer over the multi-partition concurrency stress test. The pool +# composes std synchronization primitives across threads; TSan checks the +# assembled concurrent system for data races. Requires a nightly toolchain +# with the rust-src component for instrumented std (`-Zbuild-std`): +# rustup toolchain install nightly && rustup component add rust-src --toolchain nightly +echo "### Running concurrency stress test under ThreadSanitizer (nightly + build-std)" +TSAN_TARGET="$(rustc -vV | awk '/^host:/{print $2}')" +RUSTFLAGS="-Zsanitizer=thread --cfg aws_sdk_unstable" \ +RUSTDOCFLAGS="-Zsanitizer=thread" \ + cargo +nightly test \ + -Zbuild-std \ + --target "${TSAN_TARGET}" \ + --features s2n-tls,rustls-aws-lc,wire-mock,default-client,test-util \ + --test pool_behavior_test \ + v2_cross_partition_concurrency_stress diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index cb88686d87c..3c6a03adcc5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -1155,6 +1155,25 @@ async fn eviction_task( /// /// The HTTP/1 connection sends request URIs as-is, so request-target /// form selection is the caller's responsibility. +/// Recognize hyper-util's `Singleton` coalescing-cancel sentinel by its +/// `Display`. The error type (`SingletonError`) and its inner `Canceled` +/// are upstream-private and not downcastable, so the chain is matched by +/// string. Deliberately narrow: only the exact "singleton connection +/// canceled" message re-enters a checkout; any other error propagates. +/// The string match is exercised by the multi-threaded concurrency test, +/// which fails if the upstream message changes; if upstream stops +/// producing the sentinel, this becomes inert. +fn is_singleton_canceled(err: &(dyn std::error::Error + 'static)) -> bool { + let mut e = Some(err); + while let Some(cur) = e { + if cur.to_string() == "singleton connection canceled" { + return true; + } + e = cur.source(); + } + false +} + fn rewrite_h1_request_target(req: &mut http_1x::Request, is_proxied: bool) { use http_1x::{uri::Parts, Method, Uri}; if req.method() == Method::CONNECT { @@ -1489,7 +1508,22 @@ where std::future::poll_fn(|cx| svc.poll_ready(cx)) .await .map_err(Into::into)?; - let mut checkout = svc.call(ctx.clone()).await.map_err(Into::into)?; + let mut checkout = match svc.call(ctx.clone()).await { + Ok(c) => c, + Err(e) => { + let e: BoxError = e.into(); + // A coalesced waiter whose `Singleton` maker + // bounced to the H1 fallback resolves `Canceled` + // without being reused or connected. Nothing was + // established, so re-enter the checkout (the + // retry becomes a maker or reuses the made + // connection) rather than fail. + if is_singleton_canceled(&*e) { + continue; + } + return Err(e); + } + }; if std::future::poll_fn(|cx| checkout.poll_ready(cx)) .await .is_ok() @@ -1526,9 +1560,13 @@ where let resp = conn.call(req).await.map_err(Into::into)?; Ok(erase_body(resp)) } + // Cap-bound: the cap is full (no connect was + // attempted under `NonBlocking`). Borrow a peer's + // idle connection, else fall back to the + // authoritative blocking acquire. A genuine connect + // error (a permit was free but the connect failed) + // is not `CapBound` and propagates. Err(err) if connection::CapBound::is(&*err) => { - // Cap-bound locally. Try to borrow a same-NIC - // peer's idle connection to this authority. let key = PoolKey::from_uri(&ctx.uri) .ok_or("request URI must have scheme and authority")?; if let Some(mut borrowed) = handle.try_borrow(&key) { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 159c70c6c1a..cb1c9acb5d5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -1274,12 +1274,12 @@ mod tests { //! Unit tests for H2 active-stream tracking. //! //! The test harness (`ConnectionTestHarness`) is plain HTTP only, without - //! ALPN, so we cannot exercise the full H2 pipeline end-to-end yet - //! (tracked in bosun.md as the "H2 test coverage gap"). These tests - //! verify the atomic transitions directly, which is the architectural - //! correctness property we need: an in-flight H2 connection must not be - //! evicted by the retain predicate, and the connection's `idle_at` must - //! be stamped only on the `active_streams` 1 → 0 transition. + //! ALPN, so the full H2 pipeline is not exercised end-to-end here. These + //! tests verify the atomic transitions directly, which is the + //! architectural correctness property we need: an in-flight H2 connection + //! must not be evicted by the retain predicate, and the connection's + //! `idle_at` must be stamped only on the `active_streams` 1 → 0 + //! transition. use super::*; use std::future::Future; use std::sync::atomic::{AtomicUsize, Ordering}; diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index f45be7158b5..c664e7dbf3f 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -2075,6 +2075,84 @@ async fn v2_cross_partition_borrow_respects_nic_boundary() { ); } +// --------------------------------------------------------------------------- +// Test implementations: cross-partition concurrency stress (TSan target) +// --------------------------------------------------------------------------- + +/// Drives concurrent cross-partition borrow, reclaim, and eviction against +/// a shared authority under a binding cap. The assertion is intentionally +/// weak (every request completes); the value is the interleaving: +/// concurrent `try_borrow_on` / `try_reclaim_on` touch a peer's +/// `authorities` and cache locks while that peer serves its own requests +/// and the eviction task runs `retain` on the same caches. `additional-ci` +/// additionally runs it under ThreadSanitizer. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn v2_cross_partition_concurrency_stress() { + use aws_smithy_http_client::pool::{ + CrossPartitionPolicy, Partition, PartitionId, TokioDriverSpawner, + }; + + const PARTITIONS: usize = 4; + const ROUNDS: usize = 8; + const REQUESTS_PER_ROUND: usize = 16; + + // Single loopback endpoint serving many keep-alive requests. The + // cross-partition contention comes from multiple partitions sharing + // one authority under a binding cap, not from IP spread. + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..512) + .map(|_| ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"ok", + }) + .collect(), + ) + .build() + .await; + + // No NIC binding (the common case): all partitions land in the single + // implicit NIC group, so they are borrow + reclaim peers. A small + // global cap forces cross-partition contention; a short idle timeout + // makes the eviction task churn concurrently with borrow/reclaim. + let parts = (0..PARTITIONS) + .map(|i| Partition::new(PartitionId::from_index(i), TokioDriverSpawner::current())); + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .cross_partition_policy(CrossPartitionPolicy::PreferLocal) + .max_connections(PARTITIONS) // bind the cap below the offered load + .pool_idle_timeout(Duration::from_millis(20)) + .partitions(parts) + .build_http(); + + let clients: Vec = (0..PARTITIONS) + .map(|i| PoolClient::from_partition(&pool, PartitionId::from_index(i)).into_shared()) + .collect(); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + for _ in 0..ROUNDS { + let mut tasks = tokio::task::JoinSet::new(); + for r in 0..REQUESTS_PER_ROUND { + // Spread requests across partitions so borrow/reclaim peers + // are all live at once. + let client = clients[r % PARTITIONS].clone(); + let url = url.clone(); + tasks.spawn(async move { send_and_read_body(&client, &url).await }); + } + while let Some(result) = tasks.join_next().await { + let (status, _) = result + .expect("task should not panic") + .expect("request should succeed under cross-partition contention"); + assert_eq!(status, 200); + } + // Let the eviction tick fire between rounds so the next round + // races fresh connects against reclaim/borrow on partly-evicted + // caches. + tokio::time::sleep(Duration::from_millis(30)).await; + } +} + // --------------------------------------------------------------------------- // Test implementations: stats read API // --------------------------------------------------------------------------- From 8be288b7afc087c5aca8855847a308e8e3766dcf Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 10 Jun 2026 14:22:00 -0400 Subject: [PATCH 38/47] extend tests --- .../tests/pool_behavior_test.rs | 141 ++++++++++++++++++ 1 file changed, 141 insertions(+) diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index c664e7dbf3f..ed0509e184a 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -2360,3 +2360,144 @@ async fn v2_stats_pruned_after_eviction() { "eviction should have decremented established and pruned the cell" ); } + +// --------------------------------------------------------------------------- +// Test implementations: at-the-limit scenarios +// --------------------------------------------------------------------------- + +/// A request blocked at the connection cap proceeds once an in-flight +/// request releases its permit. With `max_connections(1)`, the first +/// request pins the only permit by holding its response; a second request +/// cannot acquire and stays pending until the first is released, then +/// completes. +#[tokio::test] +async fn v2_cap_bound_request_waits_then_proceeds() { + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..2) + .map(|_| ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"ok", + }) + .collect(), + ) + .build() + .await; + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .max_connections(1) + .build_http(); + let client = PoolClient::new(&pool).into_shared(); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + // First request holds the only permit: the response is kept, body + // undrained, so the connection stays checked out. + let held = send_to(&client, &url) + .await + .expect("first request succeeds"); + assert_eq!(held.status().as_u16(), 200); + + // Second request cannot acquire a permit; it must not complete while the + // first is held. + let mut second = tokio::spawn({ + let client = client.clone(); + let url = url.clone(); + async move { send_and_read_body(&client, &url).await } + }); + assert!( + tokio::time::timeout(Duration::from_millis(200), &mut second) + .await + .is_err(), + "second request must stay pending while the cap is held" + ); + + // Release the permit by dropping the first response (its body guard + // drops, returning the connection to the pool). + drop(held); + + // The second request now acquires the freed permit and completes. + let (status, body) = tokio::time::timeout(Duration::from_secs(5), second) + .await + .expect("second request completes after permit release") + .expect("spawned task does not panic") + .expect("second request succeeds"); + assert_eq!(status, 200); + assert_eq!(body, b"ok"); + + // At most two connections: dropping the first response with its body + // undrained closes that H1 connection (an unconsumed body cannot be + // reused), so the second request may open a fresh one. The scenario + // under test is the permit wait-then-proceed, not connection reuse. + assert!(harness.tcp_accepted_count() <= 2); +} + +/// A host saturated at its per-host cap does not block requests to a +/// different host. With `max_connections_per_host(1)` and global headroom, +/// holding host X's only permit leaves a request to X pending while a +/// request to host Y proceeds — the per-host-before-global acquire order. +#[tokio::test] +async fn v2_per_host_cap_isolates_hosts() { + // Two distinct authorities, same loopback endpoint (the per-host cap is + // keyed by authority, so distinct hostnames are distinct hosts even on + // one listener). + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + (0..4) + .map(|_| ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"ok", + }) + .collect(), + ) + .dns("host-x.test", vec![IP1]) + .dns("host-y.test", vec![IP1]) + .build() + .await; + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .max_connections_per_host(1) + .build_http(); + let client = PoolClient::new(&pool).into_shared(); + let port = harness.endpoints[0].port(); + let url_x = format!("http://host-x.test:{port}/"); + let url_y = format!("http://host-y.test:{port}/"); + + // Pin host X at its per-host cap by holding the response. + let held_x = send_to(&client, &url_x).await.expect("x request succeeds"); + assert_eq!(held_x.status().as_u16(), 200); + + // A second request to X is blocked on X's per-host permit. + let mut x2 = tokio::spawn({ + let client = client.clone(); + let url_x = url_x.clone(); + async move { send_and_read_body(&client, &url_x).await } + }); + assert!( + tokio::time::timeout(Duration::from_millis(200), &mut x2) + .await + .is_err(), + "second X request must wait on X's saturated per-host cap" + ); + + // A request to host Y proceeds: Y's per-host cap is independent and + // global has headroom. + let (status, body) = + tokio::time::timeout(Duration::from_secs(5), send_and_read_body(&client, &url_y)) + .await + .expect("Y request must not be blocked by X's saturation") + .expect("y request succeeds"); + assert_eq!(status, 200); + assert_eq!(body, b"ok"); + + // Release X and let its waiter finish so the spawned task is not leaked. + drop(held_x); + let _ = tokio::time::timeout(Duration::from_secs(5), x2) + .await + .expect("x2 completes after release") + .expect("x2 task does not panic") + .expect("x2 succeeds"); +} From 5880017e9f6712cb620b2847d5efd6c2e771aee2 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 10 Jun 2026 14:36:04 -0400 Subject: [PATCH 39/47] additional tests --- .../src/client/pool/stats.rs | 30 +++++++ .../tests/pool_behavior_test.rs | 90 +++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs index e1e71f307ad..a472114c199 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs @@ -584,6 +584,36 @@ mod tests { assert_eq!(s.capacity_hint(), None); } + #[test] + fn observe_protocol_latches_monotonically_to_mixed() { + // First observation sets the protocol from UNSET. + let c = Arc::new(ConnectionCounters::default()); + assert_eq!(c.protocol(), PROTO_UNSET); + c.observe_protocol(PROTO_H1); + assert_eq!(c.protocol(), PROTO_H1); + + // Same protocol again is a no-op (stays H1, does not advance to MIXED). + c.observe_protocol(PROTO_H1); + assert_eq!(c.protocol(), PROTO_H1); + + // A different protocol latches to MIXED. + c.observe_protocol(PROTO_H2); + assert_eq!(c.protocol(), PROTO_MIXED); + + // MIXED is terminal: observing either protocol again cannot un-mix it. + c.observe_protocol(PROTO_H1); + assert_eq!(c.protocol(), PROTO_MIXED); + c.observe_protocol(PROTO_H2); + assert_eq!(c.protocol(), PROTO_MIXED); + + // The symmetric first-observation path: H2 first, then H1 → MIXED. + let c = Arc::new(ConnectionCounters::default()); + c.observe_protocol(PROTO_H2); + assert_eq!(c.protocol(), PROTO_H2); + c.observe_protocol(PROTO_H1); + assert_eq!(c.protocol(), PROTO_MIXED); + } + #[test] fn idle_is_established_minus_active_saturating() { let s = PartitionStats { diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index ed0509e184a..b3e1eb28966 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1983,6 +1983,96 @@ async fn v2_cross_partition_borrow_reuses_peer_connection() { ); } +/// Storage residency under borrow: a borrowed connection never changes +/// partitions. When P0 borrows P1's connection under `PreferLocal`, the +/// connection stays resident in P1's storage — `stats()` shows P1 owning the +/// one established connection and P0 owning zero, even though P0's request +/// was served. Borrow dispatches *through* the peer's connection; it does +/// not migrate it (connections never move). This is the wiring counterpart +/// to the borrow test's `conn_id` equality: the id matches because the +/// connection is P1's, and the residency confirms it stayed P1's. +#[tokio::test] +async fn v2_borrowed_connection_stays_resident_in_peer() { + use aws_smithy_http_client::pool::{ + Authority, CrossPartitionPolicy, Partition, PartitionId, TokioDriverSpawner, + }; + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p1", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"p0", + }, + ], + ) + .build() + .await; + + // Same setup as the borrow test: two same-NIC partitions, cap 1, long + // idle timeout. The borrow is forced by P1 holding the only permit. + let p0 = Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()) + .interface("eth-test"); + let p1 = Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()) + .interface("eth-test"); + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .cross_partition_policy(CrossPartitionPolicy::PreferLocal) + .max_connections(1) + .pool_idle_timeout(Duration::from_secs(3600)) + .partitions([p0, p1]) + .build_http(); + + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + let authority = Authority::from_host(format!("127.0.0.1:{port}")); + let client0 = PoolClient::from_partition(&pool, PartitionId::from_index(0)).into_shared(); + let client1 = PoolClient::from_partition(&pool, PartitionId::from_index(1)).into_shared(); + + // P1 establishes and returns its connection idle (holding the permit). + let _ = send_with_capture(&client1, &url).await; + + // P0 borrows P1's connection (cap-bound, PreferLocal) and its request + // completes through it. send_with_capture drains the body, so all active + // counts settle before we read stats. + let (status, _, _) = + tokio::time::timeout(Duration::from_secs(5), send_with_capture(&client0, &url)) + .await + .expect("p0 must not block — borrow should reuse P1's connection"); + assert_eq!(status, 200); + + // Residency: the single established connection lives in P1's cell. P0 + // borrowed rather than created, so its cell holds zero established. + let stats = pool.stats(&authority); + let p1_stats = stats + .get(PartitionId::from_index(1)) + .expect("P1 owns the established connection"); + assert_eq!( + p1_stats.established, 1, + "the connection is resident in P1's storage" + ); + assert_eq!( + stats + .get(PartitionId::from_index(0)) + .map(|s| s.established) + .unwrap_or(0), + 0, + "P0 borrowed P1's connection; it added nothing to P0's established" + ); + + // Exactly one connection exists across both partitions. + let total_established: usize = stats.iter().map(|(_, s)| s.established).sum(); + assert_eq!( + total_established, 1, + "one connection total — borrow does not create or duplicate" + ); +} + /// Borrow is NIC-bounded: a peer on a different NIC is not a borrow /// candidate. With `PreferLocal` but P0 and P1 on different NICs, P0 /// cannot borrow P1's connection — nor reclaim its permit (reclaim From a3040f7dc7044e595cdd39428b95a04faef12f28 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 10 Jun 2026 15:16:43 -0400 Subject: [PATCH 40/47] fill in test --- .../tests/h2_pool_test.rs | 8 ++- .../tests/pool_behavior_test.rs | 65 +++++++++++++++---- 2 files changed, 56 insertions(+), 17 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs index a96dd0125e6..1e3cafb752a 100644 --- a/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/h2_pool_test.rs @@ -648,8 +648,10 @@ mod tls_h2 { ); } - /// s2n-tls + ALPN h2: v2 pool should route to H2 (currently broken — opens - /// new connection per request because Connected never signals negotiated_h2). + /// s2n-tls + ALPN h2: the v2 pool routes to H2 and multiplexes. Three + /// sequential requests share one connection because `S2nTlsConn::connected()` + /// reports the negotiated protocol from the TLS `application_protocol`, so the + /// Negotiate layer selects the H2 leg. #[cfg(feature = "s2n-tls")] #[tokio::test] async fn s2n_tls_alpn_h2_multiplexing() { @@ -675,7 +677,7 @@ mod tls_h2 { let (status, _) = send_request(&client, &url).await.unwrap(); assert_eq!(status, 200); - // This assertion will FAIL until s2n-tls signals negotiated_h2() + // All three requests multiplex on one H2 connection. assert_eq!( server.connection_count(), 1, diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index b3e1eb28966..7dfbdd218d6 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -273,20 +273,57 @@ async fn v2_connection_reuse() { #[tokio::test] async fn v2_idle_timeout_eviction() { - // v2 does not yet implement idle timeout eviction (Phase 4). - // - // The v1 version of this test passes because the legacy client has - // built-in idle eviction. For v2, calling idle_timeout_eviction() - // also produces 2 connections — but NOT because of idle eviction. - // The response body is never consumed by the test (HttpResponse wraps - // a lazy SdkBody). The CachedConnection is held inside the response - // body's guard until the HttpResponse is dropped. Since the first - // response is still alive during the sleep, the connection never - // returns to Cache, and the second request opens a new connection. - // - // Real idle eviction tests (consume body → return to pool → sleep - // past timeout → verify new connection) will be added in Phase 4 - // alongside the background eviction task. + // An idle connection is evicted after the pool idle timeout, so a request + // arriving after the timeout opens a fresh connection. The body of the + // first response is fully consumed, so its connection returns to the pool + // and idles (the eviction target) rather than staying checked out. + let idle_timeout = Duration::from_millis(100); + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"first", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"second", + }, + ], + ) + .build() + .await; + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .pool_idle_timeout(idle_timeout) + .build_http(); + let client = PoolClient::new(&pool).into_shared(); + let url = format!("http://127.0.0.1:{}/", harness.endpoints[0].port()); + + // First request: body consumed, connection returns to the pool idle. This + // also lazily spawns the eviction task (pool_idle_timeout is set). + let (status, _) = send_and_read_body(&client, &url) + .await + .expect("first request should succeed"); + assert_eq!(status, 200); + assert_eq!(harness.tcp_accepted_count(), 1); + + // Wait past the idle timeout so the eviction task drops the idle connection. + tokio::time::sleep(idle_timeout * 3).await; + + // The next request cannot reuse the evicted connection and opens a new one. + let (status, _) = send_and_read_body(&client, &url) + .await + .expect("second request should succeed after idle eviction"); + assert_eq!(status, 200); + assert_eq!( + harness.tcp_accepted_count(), + 2, + "the idle connection was evicted, so the second request reconnects" + ); } #[tokio::test] From 7de1bd25784613f9202ca74bfa017f9a6e10d0a3 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Wed, 10 Jun 2026 20:31:45 -0400 Subject: [PATCH 41/47] fill in some basic builder tests --- .../aws-smithy-http-client/src/client/pool.rs | 67 +++++++++++++++++ .../src/client/pool/builder.rs | 72 +++++++++++++++++++ .../src/client/pool/connection.rs | 5 ++ .../src/client/pool/partition.rs | 28 ++++++++ 4 files changed, 172 insertions(+) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 3c6a03adcc5..685bcae4216 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -1702,6 +1702,73 @@ mod tests { }) } + /// A `make_stack_for` stub that produces `NullPoolEntry` stacks — enough + /// to exercise registry indexing without standing up real connectors. + fn null_make_stack_for( + ) -> impl Fn(partition::PartitionId, &Arc) -> MakeStack { + |_id, _spawner| Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box) + } + + /// The registry's default partition is the first declared. + #[tokio::test] + async fn registry_default_partition_is_first_declared() { + let partitions = vec![ + partition::Partition::new( + partition::PartitionId::from_index(7), + partition::TokioDriverSpawner::current(), + ), + partition::Partition::new( + partition::PartitionId::from_index(3), + partition::TokioDriverSpawner::current(), + ), + ]; + let registry = partition::PartitionRegistry::build(partitions, null_make_stack_for()); + assert_eq!( + registry.default_partition().id, + partition::PartitionId::from_index(7), + "first declared partition is the default" + ); + // Both declared partitions resolve. + assert!(registry + .partition_opt(partition::PartitionId::from_index(3)) + .is_some()); + assert!(registry + .partition_opt(partition::PartitionId::from_index(99)) + .is_none()); + } + + /// Declaring the same `PartitionId` twice is a programming error and panics. + #[tokio::test] + #[should_panic(expected = "duplicate PartitionId")] + async fn registry_build_panics_on_duplicate_partition_id() { + let partitions = vec![ + partition::Partition::new( + partition::PartitionId::from_index(0), + partition::TokioDriverSpawner::current(), + ), + partition::Partition::new( + partition::PartitionId::from_index(0), + partition::TokioDriverSpawner::current(), + ), + ]; + let _ = partition::PartitionRegistry::build(partitions, null_make_stack_for()); + } + + /// The default (no-topology) path normalizes to a single anonymous + /// partition that the registry indexes and resolves via `Client::new`. + #[tokio::test] + async fn registry_anonymous_default_when_no_partitions_declared() { + let partitions = partition::normalize_partitions(Vec::new(), || { + Arc::new(partition::TokioDriverSpawner::current()) + }); + let registry = partition::PartitionRegistry::build(partitions, null_make_stack_for()); + assert_eq!( + registry.default_partition().id, + partition::PartitionId::default(), + "no-topology default is the anonymous partition" + ); + } + /// Without a `pool_idle_timeout`, `maybe_spawn_eviction_task` is a no-op. #[tokio::test] async fn eviction_task_no_spawn_without_timeout() { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index 603fc462a84..011ecc5083d 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -443,3 +443,75 @@ pub(super) fn get_non_default_port(uri: &http_1x::Uri) -> Option uri.port(), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::pool::PartitionId; + + #[test] + fn builder_defaults() { + let b = Builder::default(); + assert_eq!(b.pool_idle_timeout, None); + assert!(b.tcp_nodelay, "tcp_nodelay defaults to true"); + assert_eq!(b.max_connections, None); + assert_eq!(b.max_connections_per_host, None); + assert!(b.proxy_config.is_none()); + assert!(b.connection_event_listener.is_none()); + assert_eq!(b.cross_partition_policy, CrossPartitionPolicy::Never); + assert!(b.dns_resolver.is_none()); + assert!(b.partitions.is_empty()); + } + + /// `tls_provider` transitions the type-state while preserving every + /// configured field. Guards the hand-written field-by-field move in + /// `tls_provider` against a dropped field on a future edit. + #[test] + fn tls_provider_preserves_all_config() { + let b = Builder::default() + .pool_idle_timeout(Duration::from_secs(30)) + .tcp_nodelay(false) + .max_connections(100) + .max_connections_per_host(10) + .cross_partition_policy(CrossPartitionPolicy::PreferLocal) + .partitions([Partition::new( + PartitionId::from_index(0), + crate::client::pool::TokioDriverSpawner::from_handle( + // a handle is only needed to construct the spawner; no + // runtime work happens here. + tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() + .handle() + .clone(), + ), + )]); + + let provider = tls::Provider::Rustls(tls::rustls_provider::CryptoMode::AwsLc); + let b = b.tls_provider(provider); + + assert_eq!(b.pool_idle_timeout, Some(Duration::from_secs(30))); + assert!(!b.tcp_nodelay); + assert_eq!(b.max_connections, Some(100)); + assert_eq!(b.max_connections_per_host, Some(10)); + assert_eq!(b.cross_partition_policy, CrossPartitionPolicy::PreferLocal); + assert_eq!(b.partitions.len(), 1); + } + + #[test] + fn non_default_port_elided_per_scheme() { + let cases = [ + ("https://example.com/", None), // 443 elided + ("http://example.com/", None), // 80 elided + ("https://example.com:8443/", Some(8443)), + ("http://example.com:8080/", Some(8080)), + ("https://example.com:80/", Some(80)), // 80 is non-default for https + ("http://example.com:443/", Some(443)), // 443 is non-default for http + ]; + for (uri, expected) in cases { + let uri: http_1x::Uri = uri.parse().unwrap(); + let port = get_non_default_port(&uri).map(|p| p.as_u16()); + assert_eq!(port, expected, "uri = {uri}"); + } + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index cb1c9acb5d5..31530d53a76 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -98,6 +98,7 @@ pub enum CloseReason { /// Timing breakdown for connection establishment. #[derive(Clone, Copy, Debug)] +#[non_exhaustive] pub struct ConnectionTiming { /// TCP connect + TLS handshake combined. Measured from connector call /// start to connected IO stream returned. @@ -117,6 +118,7 @@ impl ConnectionTiming { /// Emitted when a new connection is established (TCP + TLS + HTTP handshake). #[derive(Debug)] +#[non_exhaustive] pub struct ConnectionCreatedEvent { conn_id: ConnectionId, authority: Authority, @@ -170,6 +172,7 @@ impl ConnectionCreatedEvent { /// Emitted when an existing idle connection is checked out from the pool. #[derive(Debug)] +#[non_exhaustive] pub struct ConnectionReusedEvent { conn_id: ConnectionId, authority: Authority, @@ -192,6 +195,7 @@ impl ConnectionReusedEvent { } /// Emitted when a connection is removed from the pool. +#[non_exhaustive] pub struct ConnectionClosedEvent { conn_id: ConnectionId, authority: Authority, @@ -260,6 +264,7 @@ impl ConnectionClosedEvent { } /// Emitted when a connection attempt fails before completing the handshake. +#[non_exhaustive] pub struct ConnectionFailedEvent { authority: Authority, remote_addr: Option, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 98461f233b3..e553a5ebc70 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -224,6 +224,7 @@ impl DriverSpawner for TokioDriverSpawner { /// Connections bound to different NICs are never shared regardless /// of policy. #[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[non_exhaustive] pub enum CrossPartitionPolicy { /// At capacity with no local idle, the request waits for a permit. /// Peer partitions are not consulted. @@ -498,6 +499,33 @@ mod tests { assert_eq!(CrossPartitionPolicy::default(), CrossPartitionPolicy::Never); } + #[tokio::test] + async fn normalize_partitions_synthesizes_anonymous_when_empty() { + // Empty input → exactly one anonymous partition on the supplied spawner. + let parts = normalize_partitions(Vec::new(), || { + std::sync::Arc::new(TokioDriverSpawner::current()) as std::sync::Arc + }); + assert_eq!(parts.len(), 1); + assert_eq!(parts[0].id, PartitionId::default()); + assert!(parts[0].nic.is_none()); + } + + #[tokio::test] + async fn normalize_partitions_passes_declared_set_through_untouched() { + // Non-empty input is returned unchanged, and the anonymous-spawner + // closure is never invoked. + let declared = vec![ + Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()), + Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()), + ]; + let parts = normalize_partitions(declared, || { + panic!("anonymous spawner must not be called when partitions are declared") + }); + assert_eq!(parts.len(), 2); + assert_eq!(parts[0].id, PartitionId::from_index(0)); + assert_eq!(parts[1].id, PartitionId::from_index(1)); + } + #[tokio::test] async fn tokio_driver_spawner_current() { let _ = TokioDriverSpawner::current(); From b73c00cc1aca434b32490618b246b5400f5ea430 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Thu, 11 Jun 2026 06:57:45 -0400 Subject: [PATCH 42/47] cleanup dead code --- .../src/client/pool/connection.rs | 2 +- .../src/client/pool/partition.rs | 18 +++++++----------- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 31530d53a76..6de4536cd10 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -1095,7 +1095,7 @@ pin_project! { /// fully dropped. Body streaming continues through the held inner /// `Incoming`; when the `GuardedBody` is dropped the guard drops, which /// for H1 triggers `CachedConnection::Drop` (return-to-pool or `discard` - /// if poisoned), and for H2 simply drops the singleton clone (no-op on + /// if poisoned), and for H2 drops the singleton clone (no-op on /// singleton state). /// /// The H2 variant carries a generic type parameter because diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index e553a5ebc70..16884d8ccf2 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -262,23 +262,19 @@ impl Partition { self.nic = Some(nic.into()); self } - - // used by pool build validation (later step) - #[allow(dead_code)] - pub(crate) fn id(&self) -> PartitionId { - self.id - } } /// Pool-owned state for one declared partition. Resolved once at pool /// build time and referenced by [`Client`](super::Client) handles. pub(crate) struct PartitionState { pub(crate) id: PartitionId, - // captured into make_stack via the factory; field retained for later NIC/runtime use + // Captured into `make_stack` by the build factory; the field is retained + // on the state but read through the captured closure, not directly. #[allow(dead_code)] pub(crate) spawner: std::sync::Arc, - // used by socket bind (later step) - #[allow(dead_code)] + /// Network interface this partition's connections bind to, and the + /// boundary for cross-partition borrow and reclaim (peers in the same + /// NIC group only). pub(crate) nic: Option, /// Per-host connection storage for this partition. Keyed by /// (scheme, authority); entries built lazily on first request. @@ -335,8 +331,8 @@ pub(crate) fn normalize_partitions( #[derive(Debug)] pub(crate) struct PartitionRegistry { by_id: std::collections::HashMap>, - // used by cross-partition borrow (later step) - #[allow(dead_code)] + /// Partition ids grouped by NIC, for the cross-partition borrow and + /// reclaim peer walk (candidates are drawn from the requester's NIC group). by_nic: std::collections::HashMap, Vec>, default_partition: PartitionId, } From 8dfe39e1662eaa87f3605924433855f11664947a Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 12 Jun 2026 08:42:34 -0400 Subject: [PATCH 43/47] add examples and cleanup docs --- .../aws-smithy-http-client/Cargo.toml | 10 +++ .../examples/pool-basic.rs | 32 +++++++++ .../examples/pool-partitioned.rs | 62 ++++++++++++++++ .../aws-smithy-http-client/src/client.rs | 4 +- .../aws-smithy-http-client/src/client/pool.rs | 71 ++++++++++--------- .../src/client/pool/builder.rs | 5 +- .../src/client/pool/connection.rs | 64 ++++++----------- .../src/client/pool/handshake.rs | 17 +++-- .../src/client/pool/partition.rs | 7 +- .../src/client/pool/stats.rs | 12 ++-- .../src/client/pool/vendored_cache.rs | 4 +- .../src/client/timeout.rs | 7 +- .../tests/pool_behavior_test.rs | 3 +- .../aws-smithy-http-client/tests/tls.rs | 2 +- .../src/client/connection.rs | 15 ++-- 15 files changed, 206 insertions(+), 109 deletions(-) create mode 100644 rust-runtime/aws-smithy-http-client/examples/pool-basic.rs create mode 100644 rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs diff --git a/rust-runtime/aws-smithy-http-client/Cargo.toml b/rust-runtime/aws-smithy-http-client/Cargo.toml index ccfea9b06d4..76cb549a22a 100644 --- a/rust-runtime/aws-smithy-http-client/Cargo.toml +++ b/rust-runtime/aws-smithy-http-client/Cargo.toml @@ -163,6 +163,16 @@ name = "custom-dns" required-features = ["rustls-ring"] doc-scrape-examples = true +[[example]] +name = "pool-basic" +required-features = ["rustls-aws-lc"] +doc-scrape-examples = true + +[[example]] +name = "pool-partitioned" +required-features = ["rustls-aws-lc"] +doc-scrape-examples = true + [package.metadata.smithy-rs-release-tooling] stable = true diff --git a/rust-runtime/aws-smithy-http-client/examples/pool-basic.rs b/rust-runtime/aws-smithy-http-client/examples/pool-basic.rs new file mode 100644 index 00000000000..2ec8ec94250 --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/examples/pool-basic.rs @@ -0,0 +1,32 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Basic use of the connection pool: a single shared pool and a client. +//! +//! `SharedPool` is the configuration surface (TLS, limits, idle eviction, +//! events); `Client` is a lightweight handle that implements `HttpClient` +//! and can be passed to an SDK client config. With no partitions declared, +//! the pool uses a single anonymous partition on the current runtime. + +use aws_smithy_http_client::pool::{Client, SharedPool}; +use aws_smithy_http_client::tls::{self, rustls_provider::CryptoMode}; +use std::time::Duration; + +#[tokio::main] +async fn main() { + // Configure the pool once. + let pool = SharedPool::builder() + .tls_provider(tls::Provider::Rustls(CryptoMode::AwsLc)) + .max_connections(125) + .pool_idle_timeout(Duration::from_secs(20)) + .build_https(); + + // A `Client` is a cheap handle over the pool; it implements `HttpClient`, + // so it can be handed to an SDK client config via `.http_client(...)`. + let _client = Client::new(&pool); + + // Cloning the pool is cheap (an `Arc` bump); additional clients can share it. + let _another = Client::new(&pool.clone()); +} diff --git a/rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs b/rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs new file mode 100644 index 00000000000..c67bab75ade --- /dev/null +++ b/rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs @@ -0,0 +1,62 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Splitting a pool into partitions. +//! +//! A partition is a group of connections that share a driver runtime and an +//! optional network interface. A connection's driver task is pinned to the +//! runtime that created it, so partitioning keeps each connection's I/O on +//! its owning runtime while one shared pool enforces a single global +//! connection budget. +//! +//! This example declares two interface-bound partitions, builds a `Client` +//! for one, and reads per-partition connection state for an authority. + +use aws_smithy_http_client::pool::{ + Authority, Client, CrossPartitionPolicy, Partition, PartitionId, SharedPool, TokioDriverSpawner, +}; +use aws_smithy_http_client::tls::{self, rustls_provider::CryptoMode}; + +#[tokio::main] +async fn main() { + // Declare one partition per interface. `PartitionId` is a caller-owned + // label; the driver spawner sets the runtime a partition's connection + // drivers run on (here, the current runtime). + let pool = SharedPool::builder() + .tls_provider(tls::Provider::Rustls(CryptoMode::AwsLc)) + // A single ceiling shared across all partitions. + .max_connections(2_000) + // `Never` (the default) keeps each partition's requests on its own + // connections. `PreferLocal` lets a cap-bound partition borrow a + // same-interface peer's idle connection instead of waiting. + .cross_partition_policy(CrossPartitionPolicy::Never) + .partitions([ + Partition::new(PartitionId::from_index(0), TokioDriverSpawner::current()) + .interface("eth0"), + Partition::new(PartitionId::from_index(1), TokioDriverSpawner::current()) + .interface("eth1"), + ]) + .build_https(); + + // `Client::new` targets the default partition; `from_partition` targets a + // declared one. The client is a cheap handle and implements `HttpClient`. + let _client = Client::from_partition(&pool, PartitionId::from_index(0)); + + // `stats` reports per-partition connection counts for an authority. Only + // partitions that have opened a connection to it appear. The counts are a + // lock-free, point-in-time read. + let authority = Authority::from_host("example.com"); + for (partition, s) in pool.stats(&authority).iter() { + // established: connections that exist; idle(): those not in use; + // establishing: handshakes in flight. + println!( + "partition {:?}: established={} idle={} establishing={}", + partition, + s.established, + s.idle(), + s.establishing, + ); + } +} diff --git a/rust-runtime/aws-smithy-http-client/src/client.rs b/rust-runtime/aws-smithy-http-client/src/client.rs index 42ee02192c1..5b3296c3543 100644 --- a/rust-runtime/aws-smithy-http-client/src/client.rs +++ b/rust-runtime/aws-smithy-http-client/src/client.rs @@ -512,8 +512,8 @@ fn new_tokio_hyper_builder( impl Adapter { /// Add proxy authentication header to the request if needed. /// - /// Delegates to the shared [`proxy::add_proxy_auth_header`] helper — - /// the v2 HTTP client uses the same logic. + /// Inject the `Proxy-Authorization` header when proxy credentials apply, + /// via the shared [`proxy::add_proxy_auth_header`] helper. fn add_proxy_auth_header(&self, request: &mut http_1x::Request) { if let Some(matcher) = self.proxy_matcher.as_ref() { proxy::add_proxy_auth_header(request, matcher); diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 685bcae4216..58aeaf3edce 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -150,7 +150,7 @@ pub(crate) struct PoolConfig { pub(crate) max_connections_per_host: Option, /// How long an idle connection may stay in the pool before being - /// evicted. `None` = no eviction (hyper-util's default behavior). + /// evicted. `None` = no eviction. pub(crate) pool_idle_timeout: Option, /// Optional listener for connection lifecycle events. @@ -171,6 +171,8 @@ pub struct SharedPool { pub(crate) inner: Arc, } +/// Interior of [`SharedPool`]: the connection pool plus the optional proxy +/// matcher consulted per request to decide proxy vs. direct routing. pub(crate) struct SharedPoolInner { pub(crate) pool: Arc, pub(crate) proxy_matcher: Option>, @@ -259,7 +261,7 @@ pub(crate) struct BoxedRetainer { /// Pop one idle connection and drop it to free its permit (active /// reclaim). Returns `true` if a connection was freed. The H1 cache /// leg pops from its idle Vec; the H2 singleton leg has no - /// simply-reclaimable idle (an idle H2 connection still multiplexes), + /// reclaimable idle (an idle H2 connection still multiplexes), /// so its `reclaim_fn` is a no-op returning `false`. reclaim_fn: Box bool + Send + Sync + 'static>, /// Take one idle connection wrapped as a dispatchable handle that @@ -339,7 +341,7 @@ impl PeerReclaimHandle { /// Free one over-supplied peer's idle connection to relieve /// `constraint`, returning `true` if a permit was freed. Best-effort: /// `false` if the pool is gone, there are no NIC-group peers, or no - /// peer holds reclaimable idle (P0 then just blocks on the permit). + /// peer holds reclaimable idle (P0 then blocks on the permit). pub(crate) fn try_free_under_load(&self, constraint: &BindingConstraint) -> bool { let registry = match self.registry.upgrade() { Some(r) => r, @@ -560,6 +562,15 @@ impl SharedPoolState { } } +/// Assemble a [`ConnectionPool`] from a TCP/TLS connector, pool +/// configuration, the declared partitions, and the cross-partition policy. +/// +/// Builds the shared state (budget semaphores, hooks, stats index) and a +/// per-partition stack factory: on first touch of an authority, a partition +/// lazily builds its `Negotiate(ConnectionLimit, Cache/Singleton)` stack, +/// registering that cell's counters into the stats index and binding it to +/// the shared budget. The connector is captured per partition so each can +/// carry its own NIC binding. pub(crate) fn build_pool( connector: C, config: PoolConfig, @@ -874,7 +885,7 @@ where .is_empty() }), // An idle H2 connection still multiplexes; it is - // not simply-reclaimable like an H1 cache entry. + // not directly reclaimable like an H1 cache entry. // Reclaim skips the H2 leg. reclaim_fn: Box::new(|| false), // H2 borrow is unsupported: an idle H2 @@ -933,7 +944,6 @@ where eviction_spawned: AtomicBool::new(false), drop_notifier: OnceLock::new(), registry, - cross_partition_policy, } } @@ -952,10 +962,6 @@ pub(crate) struct ConnectionPool { /// Immutable partition registry, resolved at build time. registry: Arc, - // used by cross-partition checkout - #[allow(dead_code)] - cross_partition_policy: partition::CrossPartitionPolicy, - /// Latches `true` once the eviction task has been spawned. Subsequent /// `send_request` calls observe the latch and skip the spawn path. /// Read only when `config.pool_idle_timeout` is set. @@ -993,7 +999,7 @@ impl ConnectionPool { /// /// Takes `self: &Arc` so the lazy eviction task can hold a /// `Weak` as a fallback exit signal (primary exit is the drop - /// of `drop_notifier`; the `Weak` just insulates the task against a + /// of `drop_notifier`; the `Weak` insulates the task against a /// missed drop signal). pub(crate) async fn send_request( self: &Arc, @@ -1004,9 +1010,8 @@ impl ConnectionPool { let key = PoolKey::from_uri(&ctx.uri).ok_or("request URI must have scheme and authority")?; - // Lazily spawn the idle-eviction task on first use. Matches the - // legacy hyper-util pool's approach: no task if the pool is never - // used, no task if `pool_idle_timeout` is `None` or zero. + // Lazily spawn the idle-eviction task on first use: no task if the + // pool is never used, and no task if `pool_idle_timeout` is `None` or zero. self.maybe_spawn_eviction_task(); // Dispatch the request through the per-host entry. The lock is @@ -1146,15 +1151,6 @@ async fn eviction_task( tracing::trace!("pool eviction task exiting"); } -/// Rewrite an HTTP/1.1 request's URI to the appropriate request-target -/// form for dispatch: -/// -/// - `CONNECT` → authority-form (`host:port`) -/// - proxied non-CONNECT → absolute-form (full URL with scheme + authority) -/// - direct non-CONNECT → origin-form (path + query only) -/// -/// The HTTP/1 connection sends request URIs as-is, so request-target -/// form selection is the caller's responsibility. /// Recognize hyper-util's `Singleton` coalescing-cancel sentinel by its /// `Display`. The error type (`SingletonError`) and its inner `Canceled` /// are upstream-private and not downcastable, so the chain is matched by @@ -1174,10 +1170,19 @@ fn is_singleton_canceled(err: &(dyn std::error::Error + 'static)) -> bool { false } +/// Rewrite an HTTP/1.1 request's URI to the appropriate request-target +/// form for dispatch: +/// +/// - `CONNECT` → authority-form (`host:port`) +/// - proxied non-CONNECT → absolute-form (full URL with scheme + authority) +/// - direct non-CONNECT → origin-form (path + query only) +/// +/// The HTTP/1 connection sends request URIs as-is, so request-target +/// form selection is the caller's responsibility. fn rewrite_h1_request_target(req: &mut http_1x::Request, is_proxied: bool) { use http_1x::{uri::Parts, Method, Uri}; if req.method() == Method::CONNECT { - // Authority-form: keep just the authority component. + // Authority-form: retain only the authority component. if let Some(auth) = req.uri().authority().cloned() { let mut parts = Parts::default(); parts.authority = Some(auth); @@ -1266,7 +1271,7 @@ impl Service> for H1Checkout { /// Empty until the first request against this host; safe to iterate /// at any point (no leg = no-op retain, trivially empty). retainers: RetainerSlot, - // active write-handle; read API added in a later change + // Write handle for this cell's counters; the read path is `StatsIndex`. #[allow(dead_code)] counters: Arc, /// Cross-partition borrow handle. `Some` only under @@ -1589,7 +1595,7 @@ where // `active`. } // Borrow miss → authoritative blocking local - // acquire (today's wait; includes reclaim). + // acquire (blocks on the permit; includes reclaim). let mut conn = local_checkout(&mut svc, ctx).await?; let resp = conn.call(req).await.map_err(Into::into)?; Ok(erase_body(resp)) @@ -1624,7 +1630,7 @@ where let retainers = self.retainers.lock().expect("retainer slot poisoned"); // Stop at the first leg that frees a connection. The H1 cache leg // pops an idle connection; the H2 leg is a no-op (no - // simply-reclaimable idle). + // reclaimable idle). retainers.iter().any(|r| r.reclaim_one()) } @@ -1696,7 +1702,6 @@ mod tests { config, shared, registry, - cross_partition_policy: CrossPartitionPolicy::default(), eviction_spawned: AtomicBool::new(false), drop_notifier: OnceLock::new(), }) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index 011ecc5083d..5bd8293f672 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -426,9 +426,8 @@ impl Builder { } /// Build the proxy URL matcher from a `ProxyConfig`, returning `None` when -/// no proxy is configured. Wrapped in `Arc` so each `PooledConnector` -/// constructed by `http_connector` can share the (immutable) matcher -/// without per-call allocation. +/// no proxy is configured. Wrapped in `Arc` for shared ownership: the pool +/// stores it once and each `PooledConnector` clones the handle. pub(super) fn proxy_matcher_from(proxy_config: &Option) -> Option> { proxy_config .as_ref() diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs index 6de4536cd10..ecbd58e07d5 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/connection.rs @@ -358,6 +358,8 @@ impl TimeoutContext { /// → TCP connector) to the layers that need them. #[derive(Clone, Debug)] pub(crate) struct ConnectCtx { + /// Target URI: its authority is the pool key, and the full URI is + /// passed to the TCP connector. pub(crate) uri: http_1x::Uri, /// Bounds new-connection establishment (TCP + TLS handshake). `None` /// means no connect timeout; cached connections skip the connector @@ -504,7 +506,6 @@ pub(crate) struct ConnectionInfo { /// The authority (host:port) this connection is for. Populated from the /// URI at handshake time. pub(crate) authority: Authority, - // Future: tls_info, timing (tcp_connect_duration, tls_handshake_duration) } /// A one-shot "this connection is dead, don't reuse it" flag. @@ -623,13 +624,13 @@ impl ManagedConnection { } /// Connection info (remote/local addresses). - #[allow(dead_code)] // future telemetry / debugging + #[allow(dead_code)] // accessor for telemetry/debugging consumers pub(crate) fn info(&self) -> &ConnectionInfo { &self.info } /// When this connection was established. - #[allow(dead_code)] // future max-lifetime eviction + #[allow(dead_code)] // accessor for connection age (telemetry/debugging) pub(crate) fn created_at(&self) -> Instant { self.created_at } @@ -877,25 +878,6 @@ impl Drop for CachedConnection { } } -/// Per-host shared slot that bridges H2 handshake completion to H2 checkout. -/// -/// `H2ConnectAndHandshake` calls [`Self::publish`] with a fresh -/// `ConnectionMetadata` every time a new H2 connection is established. -/// `SingletonConnection::metadata` reads the current value through -/// [`Self::current`]. -/// -/// Necessary because `hyper_util::client::pool::singleton::Singled<…>` is -/// opaque to us; we can't reach through it to the underlying -/// `ManagedConnection` to build its metadata at checkout time. The -/// handshake side has direct access, so it publishes and the checkout -/// side reads. -/// -/// Last-writer-wins is correct. Singleton holds at most one live H2 -/// connection per host at a time: on `Singled::poll_ready` error (or a -/// `Singleton::retain` predicate returning false), its state transitions -/// from `Made(svc)` back to `Empty`, dropping the old service. The next -/// request's `call` runs a fresh handshake through `H2ConnectAndHandshake`, -/// which publishes new metadata, overwriting whatever was there before. /// State an H2 checkout (`SingletonConnection`) needs but cannot reach /// through the opaque `Singled<…>` to read from the underlying /// `ManagedConnection`. Published by the H2 handshake on each new @@ -936,6 +918,8 @@ impl H2ConnectionRef { Self::default() } + /// Replace the current state with a freshly-handshaked connection's + /// state (last-writer-wins). pub(crate) fn publish(&self, state: H2ConnectionState) { *self.inner.lock().unwrap() = Some(state); } @@ -1095,8 +1079,9 @@ pin_project! { /// fully dropped. Body streaming continues through the held inner /// `Incoming`; when the `GuardedBody` is dropped the guard drops, which /// for H1 triggers `CachedConnection::Drop` (return-to-pool or `discard` - /// if poisoned), and for H2 drops the singleton clone (no-op on - /// singleton state). + /// if poisoned), and for H2 releases the stream's `DispatchGuard` + /// (decrementing `active_streams`, and stamping `idle_at` on the final + /// release). /// /// The H2 variant carries a generic type parameter because /// `SingletonConnection`'s inner `T` is `hyper_util::client::pool:: @@ -1157,26 +1142,19 @@ impl hyper::body::Body for GuardedBody { /// consumers above the Negotiate composition point work with. pub(crate) type CheckoutResponse = http_1x::Response>; -/// Wire-level IO wrapper that carries TCP connect timing. +/// Wire-level IO wrapper sitting below TLS in the connector stack. Pure +/// passthrough; never modifies data or buffers. /// -/// Sits below TLS in the connector stack. Pure passthrough; never -/// modifies data or buffers. +/// TODO(pool): instrument TCP-connect vs. TLS-handshake timing separately +/// here (this wrapper is the seam below TLS where the TCP-only duration is +/// observable); surface the split on `ConnectionTiming`. pub(crate) struct TransportIo { inner: IO, - tcp_connect_duration: Duration, } impl TransportIo { - fn new(inner: IO, tcp_connect_duration: Duration) -> Self { - Self { - inner, - tcp_connect_duration, - } - } - - /// Duration of the TCP connect (3-way handshake). - pub(crate) fn tcp_connect_duration(&self) -> Duration { - self.tcp_connect_duration + fn new(inner: IO) -> Self { + Self { inner } } } @@ -1224,11 +1202,10 @@ impl } } -/// Wraps a TCP connector to stamp the connect duration on the resulting IO. +/// Wraps a TCP connector with the [`TransportIo`] seam below the TLS layer. /// -/// The output `TransportIo` carries the TCP connect timing. -/// Passed to the TLS connector which layers encryption on top without -/// losing access to the transport-level timing. +/// The wrapper is the point at which transport-level (TCP) timing and byte +/// accounting can be observed independently of the TLS handshake above it. pub(crate) struct TimingConnector { inner: C, } @@ -1265,11 +1242,10 @@ where } fn call(&mut self, uri: http_1x::Uri) -> Self::Future { - let start = Instant::now(); let fut = self.inner.call(uri); Box::pin(async move { let io = fut.await.map_err(Into::into)?; - Ok(TransportIo::new(io, start.elapsed())) + Ok(TransportIo::new(io)) }) } } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs index 202b1d7b6db..1e61b9feda7 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/handshake.rs @@ -62,24 +62,28 @@ impl PoolHooks { ) } + /// Fire the listener's connection-created callback, if a listener is set. pub(crate) fn on_created(&self, event: &ConnectionCreatedEvent) { if let Some(ref l) = self.listener { l.on_created(event); } } + /// Fire the listener's connection-failed callback, if a listener is set. pub(crate) fn on_connection_failed(&self, event: &ConnectionFailedEvent) { if let Some(ref l) = self.listener { l.on_connection_failed(event); } } + /// Fire the listener's connection-reused callback, if a listener is set. pub(crate) fn on_reused(&self, event: &super::connection::ConnectionReusedEvent) { if let Some(ref l) = self.listener { l.on_reused(event); } } + /// Fire the listener's connection-closed callback, if a listener is set. pub(crate) fn on_closed(&self, event: &super::connection::ConnectionClosedEvent) { if let Some(ref l) = self.listener { l.on_closed(event); @@ -475,12 +479,13 @@ pub(crate) struct H2ConnectAndHandshake { } impl H2ConnectAndHandshake { - /// Create an H2 handshake service that publishes every newly established - /// connection's `ConnectionMetadata` into `h2_ref`. The same ref is held - /// on the read side by `SingletonConnection` (clones of the ref share - /// the underlying slot), so the H2 checkout path can expose connection - /// metadata and poison support to the adapter layer even though - /// `Singled<…>` itself is opaque. + /// Create an H2 handshake service that publishes each newly established + /// connection's state — `ConnectionMetadata` plus the shared + /// `active_streams` counter and `idle_at` timestamp — into `h2_ref`. The + /// same ref is held on the read side by `SingletonConnection` (clones of + /// the ref share the underlying slot), so the H2 checkout path can expose + /// connection metadata and poison support to the adapter layer, and track + /// stream occupancy, even though `Singled<…>` itself is opaque. pub(crate) fn new( connector: C, h2_ref: super::connection::H2ConnectionRef, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 16884d8ccf2..36df5702f10 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -423,10 +423,9 @@ impl PartitionRegistry { } /// Attempt to reclaim one idle connection from `peer`'s entry for - /// `key`, freeing its permit. Returns `true` if one was freed. The - /// peer's `authorities` lock is held only for the `try_reclaim_one` - /// call (which pops + drops under the cache lock, releasing it before - /// the drop). No-op `false` if the peer or entry is absent. + /// `key`, freeing its permit. Returns `true` if one was freed. The idle + /// connection is popped under the cache lock and dropped after the lock + /// is released. No-op `false` if the peer or entry is absent. pub(crate) fn try_reclaim_on(&self, peer: PartitionId, key: &super::PoolKey) -> bool { let state = match self.by_id.get(&peer) { Some(s) => s, diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs index a472114c199..790d9cf4598 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/stats.rs @@ -104,6 +104,8 @@ impl Drop for EstablishedGuard { } impl ConnectionCounters { + /// Increment the active (checked-out) count. Paired with exactly one + /// [`decr_active`](Self::decr_active) via an RAII checkout guard. pub(crate) fn incr_active(&self) { self.active.fetch_add(1, Ordering::Relaxed); } @@ -129,9 +131,8 @@ impl ConnectionCounters { /// /// The sole purpose of this tag is to keep `capacity_hint` from overstating: /// `capacity_hint` returns `Some` only for a cell known to be uniformly one - /// protocol it can reason about (HTTP/1). Accurate per-connection protocol and - /// stream-limit accounting (per-connection indexing) is future work that would - /// make this cell-level latch unnecessary. + /// protocol it can reason about (HTTP/1). Per-connection protocol and + /// stream-limit accounting would supersede this cell-level latch. pub(crate) fn observe_protocol(&self, proto: u8) { let mut cur = self.protocol.load(Ordering::Relaxed); loop { @@ -153,6 +154,9 @@ impl ConnectionCounters { } } + /// Current protocol tag for this cell: `PROTO_UNSET` before the first + /// handshake, `PROTO_H1`/`PROTO_H2` for a uniform cell, or the internal + /// mixed marker once a cell has seen both. pub(crate) fn protocol(&self) -> u8 { self.protocol.load(Ordering::Relaxed) } @@ -356,7 +360,7 @@ impl StatsIndex { /// /// Plain `usize` values from relaxed loads; cheap and `Copy`. Each field is loaded /// independently and may be transiently inconsistent with the others. -/// `#[non_exhaustive]` permits future additions (e.g. per-connection stream depth). +/// `#[non_exhaustive]` allows fields to be added without a breaking change. #[non_exhaustive] #[derive(Debug, Clone, Copy)] pub struct PartitionStats { diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs index 65a9f06c4ac..587240a2e04 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs @@ -362,8 +362,8 @@ mod internal { // synthetic `poll_ready` error. /// Prevent this cached service from being returned to the pool. /// - /// Consumes self — the normal `Drop` then runs, sees `is_closed`, - /// and skips reinsertion. + /// Consumes `self`; the inner service is dropped without + /// reinsertion, regardless of whether it is still healthy. pub fn discard(mut self) { self.is_closed = true; } diff --git a/rust-runtime/aws-smithy-http-client/src/client/timeout.rs b/rust-runtime/aws-smithy-http-client/src/client/timeout.rs index 67f00df95fc..a6c1bddf031 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/timeout.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/timeout.rs @@ -227,10 +227,9 @@ impl TimeoutKind { /// Wrap `fut` in a `MaybeTimeoutFuture` if `timeout` is set. /// -/// Used by the v2 pool layers to apply per-operation `connect_timeout` / -/// `read_timeout` without building Tower services for each. Passes -/// `sleep_impl` by reference so the helper can be used without cloning -/// the sleep impl into a dedicated service. +/// Applies per-operation `connect_timeout` / `read_timeout` without +/// requiring a dedicated Tower service per request. Takes `sleep_impl` by +/// reference to avoid cloning it into a dedicated service. pub(crate) fn maybe_timeout_future( fut: F, timeout: Option, diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index 7dfbdd218d6..dc6eeb6bca9 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -699,7 +699,8 @@ async fn connection_reuse_after_body_consumed(make: &dyn MakeClient) { ); } -/// Active connections are not evicted by idle timeout. +/// Back-to-back requests reuse one connection despite a short idle timeout: +/// each re-checkout beats the eviction tick, so no reconnect occurs. async fn active_connection_survives_idle_timeout(make: &dyn MakeClient) { let harness = ConnectionTestHarness::builder() .endpoint( diff --git a/rust-runtime/aws-smithy-http-client/tests/tls.rs b/rust-runtime/aws-smithy-http-client/tests/tls.rs index b1ca19b1306..90b79ed9767 100644 --- a/rust-runtime/aws-smithy-http-client/tests/tls.rs +++ b/rust-runtime/aws-smithy-http-client/tests/tls.rs @@ -39,7 +39,7 @@ struct TestServer { impl TestServer { /// Return the number of active connections to this server fn conn_count(&self) -> usize { - // 1 reference for the struct MockProxyServer, 1 reference for the + // 1 reference for the struct TestServer, 1 reference for the // socket task. Arc::strong_count(&self.conn_count) .checked_sub(2) diff --git a/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs b/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs index 42435f62b9b..f7d0d35f38b 100644 --- a/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs +++ b/rust-runtime/aws-smithy-runtime-api/src/client/connection.rs @@ -12,14 +12,18 @@ use std::sync::{Arc, Mutex}; /// Opaque identifier for a physical connection within a pool. /// -/// Unique for the pool's lifetime. All requests dispatched on the same +/// Distinct within a pool's lifetime under normal operation (backed by a +/// monotonic 64-bit counter). All requests dispatched on the same /// connection (including H2 multiplexed requests) share the same id. /// Useful for correlating requests with connection lifecycle tracing events. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub struct ConnectionId(u64); impl ConnectionId { - /// Create a new connection id. + /// Create a connection id from a raw numeric value. + /// + /// For HTTP client/pool implementations that assign connection + /// identifiers; the value should be distinct per live connection. pub fn new(id: u64) -> Self { Self(id) } @@ -88,7 +92,8 @@ impl ConnectionMetadata { /// Get the connection id, if one was assigned by the HTTP client. /// - /// Present for the v2 HTTP client; `None` for the v1 client. + /// `Some` when the HTTP client assigns pool-level connection identifiers; + /// `None` for clients that do not track connection identity. pub fn connection_id(&self) -> Option { self.connection_id } @@ -187,13 +192,13 @@ impl ConnectionMetadataBuilder { self } - /// Set the connection id. + /// Set the [`ConnectionId`] the HTTP client assigned to this connection. pub fn connection_id(mut self, id: ConnectionId) -> Self { self.connection_id = Some(id); self } - /// Set the connection id. + /// Set the [`ConnectionId`] the HTTP client assigned to this connection. pub fn set_connection_id(&mut self, id: Option) -> &mut Self { self.connection_id = id; self From a6a94d93f9129b885cb64cfb421626b8a796ab02 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 16 Jun 2026 16:28:30 -0400 Subject: [PATCH 44/47] expose more of config --- .../src/client/pool/builder.rs | 186 ++++++++++++++++-- 1 file changed, 173 insertions(+), 13 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index 5bd8293f672..e2938e28308 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -25,6 +25,10 @@ use crate::client::tls; use crate::client::{TlsProviderSelected, TlsUnset}; use crate::tls::TlsContext; +/// Default idle-connection eviction timeout, applied when the caller does +/// not configure one. +const DEFAULT_POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(60); + /// Builder for a [`SharedPool`]. /// /// Configures pool-wide settings: TLS, DNS, connection limits, idle @@ -39,8 +43,9 @@ use crate::tls::TlsContext; /// [`SharedPool`]: super::SharedPool #[derive(Clone)] pub struct Builder { - pool_idle_timeout: Option, + pool_idle_timeout: Option>, tcp_nodelay: bool, + tcp_keepalive: Option>, max_connections: Option, max_connections_per_host: Option, proxy_config: Option, @@ -56,6 +61,7 @@ impl std::fmt::Debug for Builder { f.debug_struct("Builder") .field("pool_idle_timeout", &self.pool_idle_timeout) .field("tcp_nodelay", &self.tcp_nodelay) + .field("tcp_keepalive", &self.tcp_keepalive) .field("max_connections", &self.max_connections) .field("max_connections_per_host", &self.max_connections_per_host) .field("proxy_config", &self.proxy_config) @@ -75,6 +81,7 @@ impl Default for Builder { Self { pool_idle_timeout: None, tcp_nodelay: true, + tcp_keepalive: None, max_connections: None, max_connections_per_host: None, proxy_config: None, @@ -92,10 +99,25 @@ impl Builder { /// Set the pool idle timeout. /// /// Connections idle longer than this duration are evicted from the - /// pool. Set below the server's idle timeout to avoid stale - /// connection errors. Defaults to no eviction when unset. - pub fn pool_idle_timeout(mut self, timeout: Duration) -> Self { - self.pool_idle_timeout = Some(timeout); + /// pool. Set below the server's idle timeout to avoid dispatching on a + /// connection the server has already closed. + /// + /// Unset, the pool uses a default of 60 seconds. Pass `Some(duration)` + /// to override, or `None` to disable idle eviction entirely. + pub fn pool_idle_timeout(mut self, timeout: D) -> Self + where + D: Into>, + { + self.pool_idle_timeout = Some(timeout.into()); + self + } + + /// This is the mutable version of [`pool_idle_timeout`](Self::pool_idle_timeout). + /// + /// The outer `None` selects the default; `Some(None)` disables idle + /// eviction; `Some(Some(d))` sets the timeout to `d`. + pub fn set_pool_idle_timeout(&mut self, timeout: Option>) -> &mut Self { + self.pool_idle_timeout = timeout; self } @@ -105,6 +127,37 @@ impl Builder { self } + /// This is the mutable version of [`tcp_nodelay`](Self::tcp_nodelay). + pub fn set_tcp_nodelay(&mut self, nodelay: bool) -> &mut Self { + self.tcp_nodelay = nodelay; + self + } + + /// Set the TCP keepalive idle time. + /// + /// Enables `SO_KEEPALIVE` with the given idle time before the first + /// probe. Keepalive detects dead peers faster than idle eviction + /// alone, notably for long-lived H2 connections. + /// + /// Keepalive is disabled by default. Pass `Some(duration)` to enable + /// it with that idle time; `None` leaves it disabled. + pub fn tcp_keepalive(mut self, time: D) -> Self + where + D: Into>, + { + self.tcp_keepalive = Some(time.into()); + self + } + + /// This is the mutable version of [`tcp_keepalive`](Self::tcp_keepalive). + /// + /// The outer `None` selects the default; `Some(None)` disables + /// keepalive; `Some(Some(d))` sets the idle time to `d`. + pub fn set_tcp_keepalive(&mut self, time: Option>) -> &mut Self { + self.tcp_keepalive = time; + self + } + /// Set the global maximum number of concurrent connections. /// /// Caps the total live connections in the pool across all hosts, @@ -120,6 +173,14 @@ impl Builder { self } + /// This is the mutable version of [`max_connections`](Self::max_connections). + /// + /// `None` leaves the global limit unset (unbounded). + pub fn set_max_connections(&mut self, n: Option) -> &mut Self { + self.max_connections = n; + self + } + /// Set the maximum number of concurrent connections per host. /// /// Each unique (scheme, authority) pair has an independent connection @@ -130,6 +191,15 @@ impl Builder { self } + /// This is the mutable version of + /// [`max_connections_per_host`](Self::max_connections_per_host). + /// + /// `None` leaves the per-host limit unset (unbounded). + pub fn set_max_connections_per_host(&mut self, n: Option) -> &mut Self { + self.max_connections_per_host = n; + self + } + /// Route connections through an HTTP/HTTPS/SOCKS proxy. /// /// Per-host proxy resolution is stable for the lifetime of the @@ -144,7 +214,7 @@ impl Builder { self } - /// Mutable-reference variant of [`proxy_config`](Self::proxy_config). + /// This is the mutable version of [`proxy_config`](Self::proxy_config). pub fn set_proxy_config(&mut self, config: Option) -> &mut Self { self.proxy_config = config; self @@ -156,7 +226,7 @@ impl Builder { self } - /// Mutable-reference variant of [`connection_event_listener`](Self::connection_event_listener). + /// This is the mutable version of [`connection_event_listener`](Self::connection_event_listener). pub fn set_connection_event_listener( &mut self, listener: Option>, @@ -176,7 +246,7 @@ impl Builder { self } - /// Mutable-reference variant of [`cross_partition_policy`](Self::cross_partition_policy). + /// This is the mutable version of [`cross_partition_policy`](Self::cross_partition_policy). pub fn set_cross_partition_policy(&mut self, policy: CrossPartitionPolicy) -> &mut Self { self.cross_partition_policy = policy; self @@ -191,7 +261,7 @@ impl Builder { self } - /// Mutable-reference variant of [`dns_resolver`](Self::dns_resolver). + /// This is the mutable version of [`dns_resolver`](Self::dns_resolver). pub fn set_dns_resolver(&mut self, resolver: Option) -> &mut Self { self.dns_resolver = resolver; self @@ -204,6 +274,15 @@ impl Builder { self.partitions = partitions.into_iter().collect(); self } + + /// This is the mutable version of [`partitions`](Self::partitions). + /// + /// `None` clears any declared topology; the pool then uses a single + /// anonymous partition on the current tokio runtime. + pub fn set_partitions(&mut self, partitions: Option>) -> &mut Self { + self.partitions = partitions.unwrap_or_default(); + self + } } impl Builder { @@ -212,6 +291,7 @@ impl Builder { Builder { pool_idle_timeout: self.pool_idle_timeout, tcp_nodelay: self.tcp_nodelay, + tcp_keepalive: self.tcp_keepalive, max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, proxy_config: self.proxy_config, @@ -233,9 +313,10 @@ impl Builder { let config = PoolConfig { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, - pool_idle_timeout: self.pool_idle_timeout, + pool_idle_timeout: resolve_pool_idle_timeout(self.pool_idle_timeout), connection_event_listener: self.connection_event_listener.clone(), }; + let keepalive = resolve_tcp_keepalive(self.tcp_keepalive); let proxy_matcher = proxy_matcher_from(&self.proxy_config); let partitions = std::mem::take(&mut self.partitions); let policy = self.cross_partition_policy; @@ -243,11 +324,13 @@ impl Builder { Some(resolver) => { let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); tcp.set_nodelay(self.tcp_nodelay); + tcp.set_keepalive(keepalive); build_http_pool_with_proxy(tcp, &self.proxy_config, config, partitions, policy) } None => { let mut tcp = HyperHttpConnector::new(); tcp.set_nodelay(self.tcp_nodelay); + tcp.set_keepalive(keepalive); build_http_pool_with_proxy(tcp, &self.proxy_config, config, partitions, policy) } }; @@ -282,7 +365,7 @@ impl Builder { let config = super::PoolConfig { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, - pool_idle_timeout: self.pool_idle_timeout, + pool_idle_timeout: resolve_pool_idle_timeout(self.pool_idle_timeout), connection_event_listener: self.connection_event_listener.clone(), }; let policy = self.cross_partition_policy; @@ -339,19 +422,28 @@ impl Builder { self } + /// This is the mutable version of [`tls_context`](Self::tls_context). + pub fn set_tls_context(&mut self, context: TlsContext) -> &mut Self { + self.tls.context = context; + self + } + /// Build an HTTPS client with the selected TLS provider. pub fn build_https(mut self) -> super::SharedPool { let dns_resolver = self.dns_resolver.take(); + let keepalive = resolve_tcp_keepalive(self.tcp_keepalive); match dns_resolver { Some(resolver) => { let mut tcp = HyperHttpConnector::new_with_resolver(HyperUtilResolver { resolver }); tcp.set_nodelay(self.tcp_nodelay); + tcp.set_keepalive(keepalive); tcp.enforce_http(false); self.build_from_tcp(tcp) } None => { let mut tcp = HyperHttpConnector::new(); tcp.set_nodelay(self.tcp_nodelay); + tcp.set_keepalive(keepalive); tcp.enforce_http(false); self.build_from_tcp(tcp) } @@ -369,7 +461,7 @@ impl Builder { let config = PoolConfig { max_connections: self.max_connections, max_connections_per_host: self.max_connections_per_host, - pool_idle_timeout: self.pool_idle_timeout, + pool_idle_timeout: resolve_pool_idle_timeout(self.pool_idle_timeout), connection_event_listener: self.connection_event_listener.clone(), }; @@ -425,6 +517,22 @@ impl Builder { } } +/// Resolve the configured pool idle timeout. Outer `None` applies the +/// default; `Some(None)` disables eviction; `Some(Some(d))` uses `d`. +fn resolve_pool_idle_timeout(configured: Option>) -> Option { + match configured { + None => Some(DEFAULT_POOL_IDLE_TIMEOUT), + Some(inner) => inner, + } +} + +/// Resolve the configured TCP keepalive idle time. Keepalive is disabled +/// by default, so the outer `None` resolves to off. `Some(None)` is also +/// off; `Some(Some(d))` enables keepalive with idle time `d`. +fn resolve_tcp_keepalive(configured: Option>) -> Option { + configured.flatten() +} + /// Build the proxy URL matcher from a `ProxyConfig`, returning `None` when /// no proxy is configured. Wrapped in `Arc` for shared ownership: the pool /// stores it once and each `PooledConnector` clones the handle. @@ -451,7 +559,9 @@ mod tests { #[test] fn builder_defaults() { let b = Builder::default(); + // Outer None = "not configured"; the default is applied at build time. assert_eq!(b.pool_idle_timeout, None); + assert_eq!(b.tcp_keepalive, None); assert!(b.tcp_nodelay, "tcp_nodelay defaults to true"); assert_eq!(b.max_connections, None); assert_eq!(b.max_connections_per_host, None); @@ -462,6 +572,54 @@ mod tests { assert!(b.partitions.is_empty()); } + #[test] + fn idle_timeout_resolution() { + // Unconfigured → default applied. + assert_eq!( + resolve_pool_idle_timeout(None), + Some(DEFAULT_POOL_IDLE_TIMEOUT) + ); + // Some(None) → explicitly disabled. + assert_eq!(resolve_pool_idle_timeout(Some(None)), None); + // Some(Some(d)) → overridden. + let d = Duration::from_secs(5); + assert_eq!(resolve_pool_idle_timeout(Some(Some(d))), Some(d)); + } + + #[test] + fn keepalive_resolution() { + // Unset → off (keepalive is disabled by default). + assert_eq!(resolve_tcp_keepalive(None), None); + assert_eq!(resolve_tcp_keepalive(Some(None)), None); + let d = Duration::from_secs(45); + assert_eq!(resolve_tcp_keepalive(Some(Some(d))), Some(d)); + } + + #[test] + fn idle_timeout_setters_set_three_states() { + // chaining setter: Duration → Some(Some(d)) + let b = Builder::default().pool_idle_timeout(Duration::from_secs(7)); + assert_eq!(b.pool_idle_timeout, Some(Some(Duration::from_secs(7)))); + // chaining setter: None → Some(None) (disable) + let b = Builder::default().pool_idle_timeout(None); + assert_eq!(b.pool_idle_timeout, Some(None)); + // mutable setter passes through verbatim + let mut b = Builder::default(); + b.set_pool_idle_timeout(Some(None)); + assert_eq!(b.pool_idle_timeout, Some(None)); + } + + #[test] + fn keepalive_setters_set_three_states() { + let b = Builder::default().tcp_keepalive(Duration::from_secs(15)); + assert_eq!(b.tcp_keepalive, Some(Some(Duration::from_secs(15)))); + let b = Builder::default().tcp_keepalive(None); + assert_eq!(b.tcp_keepalive, Some(None)); + let mut b = Builder::default(); + b.set_tcp_keepalive(Some(None)); + assert_eq!(b.tcp_keepalive, Some(None)); + } + /// `tls_provider` transitions the type-state while preserving every /// configured field. Guards the hand-written field-by-field move in /// `tls_provider` against a dropped field on a future edit. @@ -470,6 +628,7 @@ mod tests { let b = Builder::default() .pool_idle_timeout(Duration::from_secs(30)) .tcp_nodelay(false) + .tcp_keepalive(Duration::from_secs(45)) .max_connections(100) .max_connections_per_host(10) .cross_partition_policy(CrossPartitionPolicy::PreferLocal) @@ -489,8 +648,9 @@ mod tests { let provider = tls::Provider::Rustls(tls::rustls_provider::CryptoMode::AwsLc); let b = b.tls_provider(provider); - assert_eq!(b.pool_idle_timeout, Some(Duration::from_secs(30))); + assert_eq!(b.pool_idle_timeout, Some(Some(Duration::from_secs(30)))); assert!(!b.tcp_nodelay); + assert_eq!(b.tcp_keepalive, Some(Some(Duration::from_secs(45)))); assert_eq!(b.max_connections, Some(100)); assert_eq!(b.max_connections_per_host, Some(10)); assert_eq!(b.cross_partition_policy, CrossPartitionPolicy::PreferLocal); From 174400987dccd7e137fefa96b1143d21c7ddfb78 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Tue, 16 Jun 2026 22:31:51 -0400 Subject: [PATCH 45/47] fix nic binding --- .../aws-smithy-http-client/src/client/pool.rs | 25 +++---- .../src/client/pool/builder.rs | 74 ++++++++++++++----- .../src/client/pool/partition.rs | 10 +-- 3 files changed, 73 insertions(+), 36 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index 58aeaf3edce..c90ca430a09 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -571,13 +571,14 @@ impl SharedPoolState { /// registering that cell's counters into the stats index and binding it to /// the shared budget. The connector is captured per partition so each can /// carry its own NIC binding. -pub(crate) fn build_pool( - connector: C, +pub(crate) fn build_pool( + connector_factory: F, config: PoolConfig, partitions: Vec, cross_partition_policy: partition::CrossPartitionPolicy, ) -> ConnectionPool where + F: Fn(&partition::Partition) -> C + Send + Sync + 'static, C: Service + Clone + Send + Sync + 'static, C::Error: Into + 'static, C::Future: Unpin + Send + 'static, @@ -592,13 +593,12 @@ where "pool: initialized" ); + let connector_factory = Arc::new(connector_factory); let make_stack_for = { - let connector = connector.clone(); - move |partition_id: partition::PartitionId, - spawner: &Arc| - -> MakeStack { - let connector = connector.clone(); - let spawner = spawner.clone(); + move |partition: &partition::Partition| -> MakeStack { + let connector = connector_factory(partition); + let partition_id = partition.id; + let spawner = partition.spawner.clone(); let cross_partition_policy = cross_partition_policy; Arc::new( move |uri: &http_1x::Uri, shared: &SharedPoolState| -> Box { @@ -1686,9 +1686,7 @@ mod tests { fn pool_with_config(config: PoolConfig) -> Arc { let shared = Arc::new(SharedPoolState::new(&config)); - let make_stack_for = |_partition_id: partition::PartitionId, - _spawner: &Arc| - -> MakeStack { + let make_stack_for = |_partition: &partition::Partition| -> MakeStack { Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box) }; let partitions = partition::normalize_partitions(Vec::new(), || { @@ -1709,9 +1707,8 @@ mod tests { /// A `make_stack_for` stub that produces `NullPoolEntry` stacks — enough /// to exercise registry indexing without standing up real connectors. - fn null_make_stack_for( - ) -> impl Fn(partition::PartitionId, &Arc) -> MakeStack { - |_id, _spawner| Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box) + fn null_make_stack_for() -> impl Fn(&partition::Partition) -> MakeStack { + |_partition| Arc::new(|_uri, _shared| Box::new(NullPoolEntry) as Box) } /// The registry's default partition is the first declared. diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index e2938e28308..0bdd5300983 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -348,6 +348,9 @@ impl Builder { /// that implements hyper's `Read`, `Write`, and `Connection` traits. /// The pool's Negotiate layer uses `Connection::connected().is_negotiated_h2()` /// to route connections to the H2 path. + /// + /// NIC binding is not applied to custom TCP connectors — the custom + /// connector owns its own socket configuration. #[cfg(all(feature = "test-util", aws_sdk_unstable))] #[doc(hidden)] pub fn build_http_with_tcp_connector(self, connector: C) -> super::SharedPool @@ -369,7 +372,8 @@ impl Builder { connection_event_listener: self.connection_event_listener.clone(), }; let policy = self.cross_partition_policy; - let pool = super::build_pool(connector, config, self.partitions, policy); + let connector_factory = move |_partition: &Partition| connector.clone(); + let pool = super::build_pool(connector_factory, config, self.partitions, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), @@ -408,10 +412,17 @@ where } if proxy_config.is_disabled() { - super::build_pool(tcp, config, partitions, policy) + let connector_factory = + move |partition: &Partition| bind_interface(&tcp, partition.nic.as_deref()); + super::build_pool(connector_factory, config, partitions, policy) } else { - let connector = crate::client::connect::HttpProxyConnector::new(tcp, proxy_config); - super::build_pool(connector, config, partitions, policy) + let connector_factory = move |partition: &Partition| { + crate::client::connect::HttpProxyConnector::new( + bind_interface(&tcp, partition.nic.as_deref()), + proxy_config.clone(), + ) + }; + super::build_pool(connector_factory, config, partitions, policy) } } @@ -480,13 +491,17 @@ impl Builder { feature = "rustls-ring" ))] tls::Provider::Rustls(crypto_mode) => { - let connector = tls::rustls_provider::build_connector::wrap_connector( - tcp, - crypto_mode.clone(), - &self.tls.context, - proxy_config, - ); - let pool = super::build_pool(connector, config, partitions, policy); + let crypto_mode = crypto_mode.clone(); + let tls_context = self.tls.context.clone(); + let connector_factory = move |partition: &Partition| { + tls::rustls_provider::build_connector::wrap_connector( + bind_interface(&tcp, partition.nic.as_deref()), + crypto_mode.clone(), + &tls_context, + proxy_config.clone(), + ) + }; + let pool = super::build_pool(connector_factory, config, partitions, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), @@ -496,12 +511,15 @@ impl Builder { } #[cfg(feature = "s2n-tls")] tls::Provider::S2nTls => { - let connector = tls::s2n_tls_provider::build_connector::wrap_connector( - tcp, - &self.tls.context, - proxy_config, - ); - let pool = super::build_pool(connector, config, partitions, policy); + let tls_context = self.tls.context.clone(); + let connector_factory = move |partition: &Partition| { + tls::s2n_tls_provider::build_connector::wrap_connector( + bind_interface(&tcp, partition.nic.as_deref()), + &tls_context, + proxy_config.clone(), + ) + }; + let pool = super::build_pool(connector_factory, config, partitions, policy); super::SharedPool { inner: Arc::new(super::SharedPoolInner { pool: Arc::new(pool), @@ -533,6 +551,28 @@ fn resolve_tcp_keepalive(configured: Option>) -> Option(base: &HyperHttpConnector, nic: Option<&str>) -> HyperHttpConnector +where + R: Clone, +{ + #[allow(unused_mut)] + let mut tcp = base.clone(); + #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] + if let Some(interface) = nic { + tcp.set_interface(interface); + } + let _ = nic; + tcp +} + /// Build the proxy URL matcher from a `ProxyConfig`, returning `None` when /// no proxy is configured. Wrapped in `Arc` for shared ownership: the pool /// stores it once and each `PooledConnector` clones the handle. diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index 36df5702f10..f1e4d6b58ab 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -242,9 +242,9 @@ pub enum CrossPartitionPolicy { /// for its lifetime. #[derive(Clone, Debug)] pub struct Partition { - id: PartitionId, - spawner: std::sync::Arc, - nic: Option, + pub(super) id: PartitionId, + pub(super) spawner: std::sync::Arc, + pub(super) nic: Option, } impl Partition { @@ -344,7 +344,7 @@ impl PartitionRegistry { /// no-topology case via [`normalize_partitions`] first). pub(crate) fn build( partitions: Vec, - make_stack_for: impl Fn(PartitionId, &std::sync::Arc) -> super::MakeStack, + make_stack_for: impl Fn(&Partition) -> super::MakeStack, ) -> Self { assert!( !partitions.is_empty(), @@ -357,7 +357,7 @@ impl PartitionRegistry { std::collections::HashMap::new(); for p in partitions { by_nic.entry(p.nic.clone()).or_default().push(p.id); - let make_stack = make_stack_for(p.id, &p.spawner); + let make_stack = make_stack_for(&p); let state = std::sync::Arc::new(PartitionState { id: p.id, spawner: p.spawner, From 998e5fb9254bf972aa9c6f6e82e521afd627613e Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 26 Jun 2026 10:03:40 -0400 Subject: [PATCH 46/47] Address pool review: H1 in-flight eviction, NIC setter gate, cap validation Fix three issues raised in PR review of the composable connection pool. H1 host entry evicted during an in-flight request A checked-out H1 connection is taken out of the cache's idle set, so the H1 retainer reports the host entry empty while the connection is still out. retain_idle then removes the entry and drops its cache; the in-flight checkout holds only a Weak, so on body drain the connection is dropped instead of returned. TypedPoolEntry::is_empty now reports the entry non-empty while active or establishing counters are non-zero, so an entry with in-flight checkouts or in-progress connects is never evicted. The counters field, previously write-only, gains this reader. Adds a characterization test (an in-flight request held across an eviction tick reuses its connection) and a no-eviction control. NIC interface setter available on platforms that cannot bind Partition::interface accepted an interface on every target, but the bind only applies on Android, Fuchsia, and Linux; elsewhere it was silently ignored. Gate the setter to those targets, matching the v1 connector's set_interface. The nic field stays cross-platform: it is also the cross-partition borrow-group key, which is platform-independent. Tests and the example that use interface() as a group label are gated to match. max_connections below max_connections_per_host A global cap below the per-host cap clamps every host to the global value, so the per-host limit can never be reached. build_pool logs a warning when both are set and the global is lower; the per-host setter documents the clamp. sdk-lints compatibility for vendored_cache.rs The copyright check scans the first ten lines; the vendored file's MIT attribution preamble pushed the Amazon header past that window, so it read as missing. Move the Amazon header above the preamble. The file's TODOs are upstream's and kept verbatim, so add it to the todos IGNORE_DIRS rather than reword them. --- .../examples/pool-partitioned.rs | 16 ++ .../aws-smithy-http-client/src/client/pool.rs | 34 ++++- .../src/client/pool/builder.rs | 4 +- .../src/client/pool/partition.rs | 14 ++ .../src/client/pool/vendored_cache.rs | 5 + .../tests/pool_behavior_test.rs | 142 ++++++++++++++++++ tools/ci-build/sdk-lints/src/todos.rs | 7 +- 7 files changed, 218 insertions(+), 4 deletions(-) diff --git a/rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs b/rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs index c67bab75ade..2401a878992 100644 --- a/rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs +++ b/rust-runtime/aws-smithy-http-client/examples/pool-partitioned.rs @@ -13,12 +13,19 @@ //! //! This example declares two interface-bound partitions, builds a `Client` //! for one, and reads per-partition connection state for an authority. +//! +//! Binding a partition to a network interface ([`Partition::interface`]) is +//! only available on Android, Fuchsia, and Linux, so this example builds its +//! interface-bound topology only on those targets. +#[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] use aws_smithy_http_client::pool::{ Authority, Client, CrossPartitionPolicy, Partition, PartitionId, SharedPool, TokioDriverSpawner, }; +#[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] use aws_smithy_http_client::tls::{self, rustls_provider::CryptoMode}; +#[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] #[tokio::main] async fn main() { // Declare one partition per interface. `PartitionId` is a caller-owned @@ -60,3 +67,12 @@ async fn main() { ); } } + +/// Interface binding is unavailable on this target; see the module docs. +#[cfg(not(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] +fn main() { + eprintln!( + "pool-partitioned: interface-bound partitions require Android, Fuchsia, \ + or Linux; nothing to demonstrate on this target" + ); +} diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool.rs b/rust-runtime/aws-smithy-http-client/src/client/pool.rs index c90ca430a09..ca8eba2a115 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool.rs @@ -584,6 +584,23 @@ where C::Future: Unpin + Send + 'static, IO: hyper::rt::Read + hyper::rt::Write + HyperConnection + Unpin + Send + 'static, { + // A global cap below the per-host cap clamps every host to the global + // value: the per-host limit can never be reached, so its configured + // value is silently ineffective. Benign (the pool still works, just + // more tightly bound than intended), so warn rather than reject. + if let (Some(global), Some(per_host)) = + (config.max_connections, config.max_connections_per_host) + { + if global < per_host { + tracing::warn!( + max_connections = global, + max_connections_per_host = per_host, + "pool: max_connections is below max_connections_per_host; the \ + global cap binds first, so the per-host limit cannot be reached" + ); + } + } + let shared = Arc::new(SharedPoolState::new(&config)); tracing::debug!( @@ -1452,8 +1469,9 @@ struct TypedPoolEntry { /// Empty until the first request against this host; safe to iterate /// at any point (no leg = no-op retain, trivially empty). retainers: RetainerSlot, - // Write handle for this cell's counters; the read path is `StatsIndex`. - #[allow(dead_code)] + // Write handle for this cell's counters; the `StatsIndex` is the other + // reader. `is_empty` consults `active`/`establishing` here to keep an + // entry with in-flight checkouts alive (see `is_empty`). counters: Arc, /// Cross-partition borrow handle. `Some` only under /// `CrossPartitionPolicy::PreferLocal`; `None` under `Never` (and for @@ -1622,6 +1640,18 @@ where } fn is_empty(&self) -> bool { + // An entry with in-flight checkouts (`active`) or in-progress + // connects (`establishing`) is not empty, even when the retainers + // report so. A checked-out H1 connection is `take()`n out of the + // cache's idle set, so the H1 retainer reports empty while the + // connection is still out; removing the entry would drop the cache + // the connection returns to, and the in-flight checkout (holding + // only a `Weak`) would lose it on body drain rather than reuse it. + if self.counters.active.load(Ordering::Relaxed) > 0 + || self.counters.establishing.load(Ordering::Relaxed) > 0 + { + return false; + } let retainers = self.retainers.lock().expect("retainer slot poisoned"); retainers.iter().all(|r| r.is_empty()) } diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs index 0bdd5300983..7b1ec987b5b 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/builder.rs @@ -167,7 +167,9 @@ impl Builder { /// /// Should be at least [`max_connections_per_host`](Self::max_connections_per_host) /// when both are set; the global limit applies on top of the per-host - /// limit. + /// limit. A global limit below the per-host limit clamps every host to + /// the global value (the per-host limit can never be reached) and logs a + /// warning at build time. pub fn max_connections(mut self, n: usize) -> Self { self.max_connections = Some(n); self diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs index f1e4d6b58ab..0b6e57baecf 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/partition.rs @@ -258,6 +258,20 @@ impl Partition { } /// Bind this partition's connections to a network interface. + /// + /// If a socket is bound to an interface, only packets received from that + /// particular interface are processed by the socket. Note that this only + /// works for some socket types (e.g. `AF_INET` sockets). + /// + /// On Linux it can be used to specify a [VRF], but the binary needs to + /// either have `CAP_NET_RAW` capability set or be run as root. + /// + /// This method is only available on Android, Fuchsia, and Linux; on other + /// targets the interface cannot be bound and the setter is not offered, + /// so an unsupported binding cannot be configured silently. + /// + /// [VRF]: https://www.kernel.org/doc/Documentation/networking/vrf.txt + #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] pub fn interface(mut self, nic: impl Into) -> Self { self.nic = Some(nic.into()); self diff --git a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs index 587240a2e04..9f5a05a7a46 100644 --- a/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs +++ b/rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs @@ -1,3 +1,8 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + /* * Portions of this file are derived from hyper-util * (https://github.com/hyperium/hyper-util), licensed under MIT: diff --git a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs index dc6eeb6bca9..b0acea4a391 100644 --- a/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs +++ b/rust-runtime/aws-smithy-http-client/tests/pool_behavior_test.rs @@ -1770,6 +1770,10 @@ async fn partition_spawners_are_isolated() { /// - P0's connection is freshly created (distinct `conn_id`), /// - the server accepted two connections (P1's reclaimed, P0's new), /// - P0 completed well under the idle timeout. +// `Partition::interface` (the NIC-group label these partitions share) is only +// available on Android/Fuchsia/Linux; the borrow/reclaim logic it exercises is +// platform-independent but cannot be configured elsewhere. +#[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] #[tokio::test] async fn v2_cross_partition_reclaim_frees_peer_idle() { use aws_smithy_http_client::pool::{ @@ -1934,6 +1938,10 @@ async fn v2_cross_partition_reclaim_frees_peer_idle() { /// reclaim: the connection stays the peer's (P0 runs on P1's exact /// connection), proven by `p0_conn_id == p1_conn_id` and a single TCP /// accept (P1's connection serves both requests via keep-alive). +// `Partition::interface` (the shared NIC-group label) is Android/Fuchsia/Linux +// only; the borrow path it exercises is platform-independent but unconfigurable +// elsewhere. +#[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] #[tokio::test] async fn v2_cross_partition_borrow_reuses_peer_connection() { use aws_smithy_http_client::pool::{ @@ -2029,6 +2037,10 @@ async fn v2_cross_partition_borrow_reuses_peer_connection() { /// not migrate it (connections never move). This is the wiring counterpart /// to the borrow test's `conn_id` equality: the id matches because the /// connection is P1's, and the residency confirms it stayed P1's. +// `Partition::interface` (the shared NIC-group label) is Android/Fuchsia/Linux +// only; the borrow residency it exercises is platform-independent but +// unconfigurable elsewhere. +#[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] #[tokio::test] async fn v2_borrowed_connection_stays_resident_in_peer() { use aws_smithy_http_client::pool::{ @@ -2119,6 +2131,10 @@ async fn v2_borrowed_connection_stays_resident_in_peer() { /// P0 then connects locally. Proven by `p0_conn_id != p1_conn_id` and two /// TCP accepts — the opposite of the same-NIC borrow case, which reuses /// P1's exact connection. +// `Partition::interface` (the distinct NIC-group labels that form the boundary +// under test) is Android/Fuchsia/Linux only; the boundary enforcement is +// platform-independent but unconfigurable elsewhere. +#[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] #[tokio::test] async fn v2_cross_partition_borrow_respects_nic_boundary() { use aws_smithy_http_client::pool::{ @@ -2629,3 +2645,129 @@ async fn v2_per_host_cap_isolates_hosts() { .expect("x2 task does not panic") .expect("x2 succeeds"); } + +/// An H1 host entry must survive eviction while a request is in flight: a +/// checked-out H1 connection is taken out of the cache's idle set, so the +/// retainers report the entry empty even though the connection is still +/// live and will return to the pool when its body drains. Hold an H1 +/// response (body undrained, connection checked out) across an eviction +/// tick, then drain it and issue a second request to the same host. If the +/// entry were removed mid-flight, the checked-out connection could not +/// return to its (dropped) cache and request 2 would have to reconnect +/// (`tcp_accepted == 2`). +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn h1_entry_survives_eviction_during_in_flight_request() { + use http_body_util::BodyExt; + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"one", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"two", + }, + ], + ) + .build() + .await; + + let idle_timeout = Duration::from_millis(100); + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .pool_idle_timeout(idle_timeout) + .build_http(); + + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + let client = PoolClient::new(&pool).into_shared(); + + // Request 1: hold the response without draining, connection checked out. + let resp1 = send_to(&client, &url).await.expect("req1 should succeed"); + assert_eq!(resp1.status().as_u16(), 200); + + // Let the eviction task tick (>= 2 ticks) while the request is in flight. + tokio::time::sleep(idle_timeout * 3).await; + + // Drain req1's body: the body guard drops, CachedConnection::Drop fires. + let _ = BodyExt::collect(resp1.into_body()) + .await + .expect("body1 readable") + .to_bytes(); + tokio::task::yield_now().await; + + // Request 2 to the same host must reuse the held-open connection. + let resp2 = send_to(&client, &url).await.expect("req2 should succeed"); + assert_eq!(resp2.status().as_u16(), 200); + let _ = BodyExt::collect(resp2.into_body()) + .await + .expect("body2 readable") + .to_bytes(); + + assert_eq!( + harness.tcp_accepted_count(), + 1, + "entry must not be evicted while a request is in flight; req2 reuses the connection" + ); +} + +/// Control for `h1_entry_survives_eviction_during_in_flight_request`: same +/// setup with eviction effectively disabled (60s idle timeout) reuses the +/// connection (one accept), confirming the repro's second request reuses +/// for the same reason in both cases — it isolates eviction as the only +/// variable. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn h1_in_flight_request_reuses_without_eviction() { + use http_body_util::BodyExt; + + let harness = ConnectionTestHarness::builder() + .endpoint( + IP1, + vec![ + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"one", + }, + ConnectionBehavior::RespondKeepAlive { + status: 200, + body: b"two", + }, + ], + ) + .build() + .await; + + let pool = SharedPool::builder() + .dns_resolver(harness.dns_resolver()) + .pool_idle_timeout(Duration::from_secs(60)) + .build_http(); + + let port = harness.endpoints[0].port(); + let url = format!("http://127.0.0.1:{port}/"); + let client = PoolClient::new(&pool).into_shared(); + + let resp1 = send_to(&client, &url).await.expect("req1 should succeed"); + assert_eq!(resp1.status().as_u16(), 200); + let _ = BodyExt::collect(resp1.into_body()) + .await + .expect("body1 readable") + .to_bytes(); + tokio::task::yield_now().await; + + let resp2 = send_to(&client, &url).await.expect("req2 should succeed"); + assert_eq!(resp2.status().as_u16(), 200); + let _ = BodyExt::collect(resp2.into_body()) + .await + .expect("body2 readable") + .to_bytes(); + + assert_eq!( + harness.tcp_accepted_count(), + 1, + "no eviction: req2 reuses the connection" + ); +} diff --git a/tools/ci-build/sdk-lints/src/todos.rs b/tools/ci-build/sdk-lints/src/todos.rs index d38046f3129..85c1b0212ad 100644 --- a/tools/ci-build/sdk-lints/src/todos.rs +++ b/tools/ci-build/sdk-lints/src/todos.rs @@ -11,7 +11,12 @@ use std::path::{Path, PathBuf}; // All "TODOs" must include (...) that gives them context pub(crate) struct TodosHaveContext; -const IGNORE_DIRS: &[&str] = &["tools/ci-build/sdk-lints/src/todos.rs"]; +const IGNORE_DIRS: &[&str] = &[ + "tools/ci-build/sdk-lints/src/todos.rs", + // Vendored from hyper-util; its TODOs are upstream's and kept verbatim so + // the file stays close to upstream for future syncs. + "rust-runtime/aws-smithy-http-client/src/client/pool/vendored_cache.rs", +]; impl Lint for TodosHaveContext { fn name(&self) -> &str { From e43e86d36a433b581ba45400e7604bce30cd8a31 Mon Sep 17 00:00:00 2001 From: Aaron J Todd Date: Fri, 26 Jun 2026 10:12:08 -0400 Subject: [PATCH 47/47] Update Cargo.lock files Pick up runtime crate version bumps in aws/rust-runtime (aws-credential-types, aws-runtime, aws-smithy-* and friends) and a transitive socket2 0.6.3 -> 0.6.4 bump in rust-runtime. Both lockfiles verified consistent with `cargo --locked`. --- aws/rust-runtime/Cargo.lock | 32 ++++++++++++++++---------------- rust-runtime/Cargo.lock | 2 +- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/aws/rust-runtime/Cargo.lock b/aws/rust-runtime/Cargo.lock index 89cd163f4c8..40b5e27156a 100644 --- a/aws/rust-runtime/Cargo.lock +++ b/aws/rust-runtime/Cargo.lock @@ -76,7 +76,7 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-credential-types" -version = "1.2.14" +version = "1.3.0" dependencies = [ "async-trait", "aws-smithy-async", @@ -123,7 +123,7 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.7.5" +version = "1.8.0" dependencies = [ "arbitrary", "aws-credential-types", @@ -161,11 +161,11 @@ dependencies = [ [[package]] name = "aws-runtime-api" -version = "1.1.12" +version = "1.2.0" [[package]] name = "aws-sigv4" -version = "1.4.5" +version = "1.5.0" dependencies = [ "aws-credential-types", "aws-smithy-eventstream", @@ -198,7 +198,7 @@ dependencies = [ [[package]] name = "aws-smithy-async" -version = "1.2.14" +version = "1.3.0" dependencies = [ "futures-util", "pin-project-lite", @@ -207,7 +207,7 @@ dependencies = [ [[package]] name = "aws-smithy-checksums" -version = "0.64.8" +version = "0.65.0" dependencies = [ "aws-smithy-http", "aws-smithy-types", @@ -226,7 +226,7 @@ dependencies = [ [[package]] name = "aws-smithy-eventstream" -version = "0.60.21" +version = "0.61.1" dependencies = [ "aws-smithy-types", "bytes", @@ -235,7 +235,7 @@ dependencies = [ [[package]] name = "aws-smithy-http" -version = "0.63.6" +version = "0.64.0" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -274,14 +274,14 @@ dependencies = [ [[package]] name = "aws-smithy-observability" -version = "0.2.6" +version = "0.3.0" dependencies = [ "aws-smithy-runtime-api", ] [[package]] name = "aws-smithy-protocol-test" -version = "0.63.14" +version = "0.64.0" dependencies = [ "assert-json-diff", "aws-smithy-runtime-api", @@ -298,7 +298,7 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.11.3" +version = "1.12.0" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -322,7 +322,7 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.12.3" +version = "1.13.0" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", @@ -338,7 +338,7 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api-macros" -version = "1.0.0" +version = "1.1.0" dependencies = [ "proc-macro2", "quote", @@ -347,7 +347,7 @@ dependencies = [ [[package]] name = "aws-smithy-schema" -version = "0.1.0" +version = "0.1.1" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -356,7 +356,7 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.5.1" +version = "1.6.0" dependencies = [ "base64-simd", "bytes", @@ -380,7 +380,7 @@ dependencies = [ [[package]] name = "aws-types" -version = "1.3.16" +version = "1.4.0" dependencies = [ "aws-credential-types", "aws-smithy-async", diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock index 77f9a408f7a..e5ab1c6eb17 100644 --- a/rust-runtime/Cargo.lock +++ b/rust-runtime/Cargo.lock @@ -532,7 +532,7 @@ dependencies = [ "serde", "serde_json", "serial_test", - "socket2 0.6.3", + "socket2 0.6.4", "tokio", "tokio-rustls 0.26.4", "tower 0.5.3",