From 0f890b944b4efadafab1484c3ae063872baac016 Mon Sep 17 00:00:00 2001
From: Alex Rudy <github@alexrudy.net>
Date: Thu, 31 Oct 2024 18:38:31 +0000
Subject: [PATCH] feature: client pool can use anything from request Parts (not
 just URI) for pool key

---
 src/client/pool/key.rs     | 24 +++++++++++++++++++-----
 src/client/pool/mod.rs     | 13 +++++++++++--
 src/client/pool/service.rs | 26 ++++++++++++++++----------
 3 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/src/client/pool/key.rs b/src/client/pool/key.rs
index 7efc892..b8e4d5a 100644
--- a/src/client/pool/key.rs
+++ b/src/client/pool/key.rs
@@ -89,14 +89,28 @@ impl From<(http::uri::Scheme, http::uri::Authority)> for UriKey {
 impl TryFrom<http::Uri> for UriKey {
     type Error = UriError;
 
-    fn try_from(value: http::Uri) -> Result<Self, Self::Error> {
-        let parts = value.clone().into_parts();
+    fn try_from(uri: http::Uri) -> Result<Self, Self::Error> {
+        let parts = uri.into_parts();
+        let authority = parts.authority.clone();
+        let scheme = parts
+            .scheme
+            .clone()
+            .ok_or_else(|| UriError::MissingScheme(http::Uri::from_parts(parts).unwrap()))?;
+        Ok::<_, UriError>(Self(scheme, authority))
+    }
+}
+
+impl TryFrom<&http::request::Parts> for UriKey {
+    type Error = UriError;
 
+    fn try_from(parts: &http::request::Parts) -> Result<Self, Self::Error> {
         Ok::<_, UriError>(Self(
             parts
-                .scheme
-                .ok_or_else(|| UriError::MissingScheme(value.clone()))?,
-            parts.authority,
+                .uri
+                .scheme()
+                .ok_or_else(|| UriError::MissingScheme(parts.uri.clone()))?
+                .clone(),
+            parts.uri.authority().cloned(),
         ))
     }
 }
diff --git a/src/client/pool/mod.rs b/src/client/pool/mod.rs
index 2441cb6..1244835 100644
--- a/src/client/pool/mod.rs
+++ b/src/client/pool/mod.rs
@@ -47,9 +47,18 @@ use super::conn::Protocol;
 use super::conn::Transport;
 
 /// Key which links a URI to a connection.
-pub trait Key: Eq + std::hash::Hash + fmt::Debug + TryFrom<http::Uri, Error = UriError> {}
+pub trait Key:
+    Eq + std::hash::Hash + fmt::Debug + for<'a> TryFrom<&'a http::request::Parts, Error = UriError>
+{
+}
 
-impl<K> Key for K where K: Eq + std::hash::Hash + fmt::Debug + TryFrom<http::Uri, Error = UriError> {}
+impl<K> Key for K where
+    K: Eq
+        + std::hash::Hash
+        + fmt::Debug
+        + for<'a> TryFrom<&'a http::request::Parts, Error = UriError>
+{
+}
 
 /// The URI used for connecting to a server is invalid.
 ///
diff --git a/src/client/pool/service.rs b/src/client/pool/service.rs
index 78908bb..a7d4cc5 100644
--- a/src/client/pool/service.rs
+++ b/src/client/pool/service.rs
@@ -10,7 +10,6 @@ use tower::ServiceExt;
 use crate::client::conn::connection::ConnectionError;
 use crate::client::conn::connection::HttpConnection;
 use crate::client::conn::protocol::auto::HttpConnectionBuilder;
-use crate::client::conn::protocol::HttpProtocol;
 use crate::client::conn::transport::tcp::TcpTransport;
 use crate::client::conn::Connection;
 use crate::client::conn::Protocol;
@@ -225,14 +224,19 @@ where
     #[allow(clippy::type_complexity)]
     fn connect_to(
         &self,
-        uri: http::Uri,
-        http_protocol: HttpProtocol,
+        request_parts: &http::request::Parts,
     ) -> Result<Checkout<T, P, BIn>, ConnectionError> {
-        let key: K = K::try_from(uri.clone())?;
+        let key: K = K::try_from(request_parts)?;
         let protocol = self.protocol.clone();
         let transport = self.transport.clone();
+        let http_protocol = request_parts.version.into();
 
-        let connector = Connector::new(transport, protocol, uri, http_protocol);
+        let connector = Connector::new(
+            transport,
+            protocol,
+            request_parts.uri.clone(),
+            http_protocol,
+        );
 
         if let Some(pool) = self.pool.as_ref() {
             tracing::trace!(?key, "checking out connection");
@@ -276,12 +280,14 @@ where
     }
 
     fn call(&mut self, request: http::Request<BIn>) -> Self::Future {
-        let uri = request.uri().clone();
+        let (parts, body) = request.into_parts();
 
-        let protocol: HttpProtocol = request.version().into();
-
-        match self.connect_to(uri, protocol) {
-            Ok(checkout) => ResponseFuture::new(checkout, request, self.service.clone()),
+        match self.connect_to(&parts) {
+            Ok(checkout) => ResponseFuture::new(
+                checkout,
+                http::Request::from_parts(parts, body),
+                self.service.clone(),
+            ),
             Err(error) => ResponseFuture::error(error),
         }
     }