bitreq: Fix connection persistence, pool refresh

bitreq/src/client.rs

@@ -66,10 +66,13 @@ impl Client {
 
         // Try to get cached connection
         let conn_opt = {
-            let state = self.r#async.lock().unwrap();
+            let mut state = self.r#async.lock().unwrap();
 
             if let Some(conn) = state.connections.get(&owned_key) {
-                Some(Arc::clone(conn))
+                let conn = Arc::clone(conn);
+                state.lru_order.retain(|key| key != &owned_key);
+                state.lru_order.push_back(owned_key.clone());
+                Some(conn)
             } else {
                 None
             }

bitreq/src/connection.rs

@@ -24,12 +24,19 @@ use tokio::sync::Mutex as AsyncMutex;
 
 use crate::request::{ConnectionParams, OwnedConnectionParams, ParsedRequest};
 #[cfg(feature = "async")]
+use crate::response::HttpVersion;
+#[cfg(feature = "async")]
 use crate::Response;
 use crate::{Error, Method, ResponseLazy};
 
 #[cfg(feature = "async")]
 const BACKING_READ_BUFFER_LENGTH: usize = 16 * 1024;
 
+#[cfg(feature = "async")]
+fn has_connection_option(value: &str, option: &str) -> bool {
+    value.split(',').any(|value| value.trim().eq_ignore_ascii_case(option))
+}
+
 type UnsecuredStream = TcpStream;
 
 #[cfg(feature = "rustls")]
@@ -407,6 +414,7 @@ impl AsyncConnection {
     ) -> Pin<Box<dyn Future<Output = Result<Response, Error>> + Send + 'a>> {
         Box::pin(async move {
             let conn = Arc::clone(&*self.0.lock().unwrap());
+            let sent_close = request.has_connection_option("close");
             #[cfg(debug_assertions)]
             {
                 let next_read = conn.readable_request_id.load(Ordering::Acquire);
@@ -495,6 +503,9 @@ impl AsyncConnection {
                     // need to resend the request on a new connection.
                     retry_new_connection!(CONNECTION_STILL_READABLE, write);
                 }
+                if sent_close {
+                    conn.permits.store(0, Ordering::Release);
+                }
                 request_id = conn.next_request_id.fetch_add(1, Ordering::Relaxed);
                 #[cfg(feature = "log")]
                 log::trace!(
@@ -566,7 +577,7 @@ impl AsyncConnection {
                     request.connection_params(),
                 );
 
-                let response = Response::create_async(
+                let (response, http_version) = Response::create_async(
                     &mut *read,
                     request.config.method == Method::Head,
                     request.config.max_headers_size,
@@ -575,13 +586,19 @@ impl AsyncConnection {
                 )
                 .await?;
 
-                let mut found_keep_alive = false;
-                if let Some(header) = response.headers.get("connection") {
-                    if header.eq_ignore_ascii_case("keep-alive") {
-                        found_keep_alive = true;
-                    }
-                }
-                if !found_keep_alive {
+                let connection = response.headers.get("connection");
+                let received_close =
+                    connection.is_some_and(|value| has_connection_option(value, "close"));
+                let received_keep_alive =
+                    connection.is_some_and(|value| has_connection_option(value, "keep-alive"));
+                let persistent = !sent_close
+                    && !received_close
+                    && match http_version {
+                        HttpVersion::Http11 => true,
+                        HttpVersion::Http10 => received_keep_alive,
+                        HttpVersion::Other => false,
+                    };
+                if !persistent {
                     conn.permits.store(0, Ordering::Release);
                     conn.readable_request_id.store(usize::MAX, Ordering::Release);
                 } else {

bitreq/src/request.rs

@@ -527,6 +527,14 @@ impl ParsedRequest {
     pub(crate) fn connection_params(&self) -> ConnectionParams<'_> {
         ConnectionParams::from_request(self)
     }
+
+    #[cfg(feature = "async")]
+    pub(crate) fn has_connection_option(&self, option: &str) -> bool {
+        self.config.headers.iter().any(|(name, value)| {
+            name.eq_ignore_ascii_case("connection")
+                && value.split(',').any(|value| value.trim().eq_ignore_ascii_case(option))
+        })
+    }
 }
 
 /// A key which determines whether an existing connection can be reused

bitreq/src/response.rs

@@ -17,6 +17,14 @@ const BACKING_READ_BUFFER_LENGTH: usize = 16 * 1024;
 #[cfg(feature = "std")]
 const MAX_CONTENT_LENGTH: usize = 16 * 1024;
 
+#[cfg(feature = "async")]
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(crate) enum HttpVersion {
+    Http10,
+    Http11,
+    Other,
+}
+
 /// An HTTP response.
 ///
 /// Returned by [`Request::send`](struct.Request.html#method.send).
@@ -87,10 +95,11 @@ impl Response {
         max_headers_size: Option<usize>,
         max_status_line_len: Option<usize>,
         max_body_size: Option<usize>,
-    ) -> Result<Response, Error> {
+    ) -> Result<(Response, HttpVersion), Error> {
         use HttpStreamState::*;
 
         let ResponseMetadata {
+            http_version,
             status_code,
             reason_phrase,
             mut headers,
@@ -148,7 +157,10 @@ impl Response {
             }
         }
 
-        Ok(Response { status_code, reason_phrase, headers, url: String::new(), body })
+        Ok((
+            Response { status_code, reason_phrase, headers, url: String::new(), body },
+            http_version,
+        ))
     }
 
     /// Returns the body as an `&str`.
@@ -331,6 +343,7 @@ impl ResponseLazy {
             headers,
             state,
             max_trailing_headers_size,
+            ..
         } = read_metadata(&mut stream, max_headers_size, max_status_line_len)?;
 
         Ok(ResponseLazy {
@@ -444,6 +457,8 @@ enum HttpStreamState {
 // response.meta.status_code or similar.)
 #[cfg(feature = "std")]
 struct ResponseMetadata {
+    #[cfg(feature = "async")]
+    http_version: HttpVersion,
     status_code: i32,
     reason_phrase: String,
     headers: BTreeMap<String, String>,
@@ -612,6 +627,8 @@ macro_rules! define_read_methods {
             max_status_line_len: Option<usize>,
         ) -> Result<ResponseMetadata, Error> {
             let line = maybe_await!($read_line(stream, max_status_line_len, Error::StatusLineOverflow), $($await)?)?;
+            #[cfg(feature = "async")]
+            let http_version = parse_http_version(&line);
             let (status_code, reason_phrase) = parse_status_line(&line);
 
             let mut headers = BTreeMap::new();
@@ -657,6 +674,8 @@ macro_rules! define_read_methods {
             };
 
             Ok(ResponseMetadata {
+                #[cfg(feature = "async")]
+                http_version,
                 status_code,
                 reason_phrase,
                 headers,
@@ -702,6 +721,15 @@ define_read_methods!((read_until_closed, read_with_content_length, read_trailers
 #[cfg(feature = "async")]
 define_read_methods!((read_until_closed_async, read_with_content_length_async, read_trailers_async, read_chunked_async, read_metadata_async, read_line_async)<R: AsyncRead | Unpin>, R, async, await);
 
+#[cfg(feature = "async")]
+fn parse_http_version(line: &str) -> HttpVersion {
+    match line.split_once(' ').map(|(version, _)| version) {
+        Some("HTTP/1.0") => HttpVersion::Http10,
+        Some("HTTP/1.1") => HttpVersion::Http11,
+        _ => HttpVersion::Other,
+    }
+}
+
 #[cfg(feature = "std")]
 fn parse_status_line(line: &str) -> (i32, String) {
     // sample status line format

bitreq/tests/async_pool.rs

@@ -0,0 +1,68 @@
+#![cfg(feature = "async")]
+
+use std::time::Duration;
+
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+
+const EMPTY_RESPONSE: &[u8] = b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n";
+
+async fn read_request(stream: &mut TcpStream) -> Vec<u8> {
+    let mut request = Vec::new();
+    let mut buf = [0; 1024];
+    loop {
+        let read = stream.read(&mut buf).await.unwrap();
+        if read == 0 {
+            return request;
+        }
+        request.extend_from_slice(&buf[..read]);
+        if request.windows(4).any(|window| window == b"\r\n\r\n") {
+            return request;
+        }
+    }
+}
+
+async fn spawn_counting_server(
+    expected_requests: usize,
+) -> (String, tokio::task::JoinHandle<usize>) {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let server = tokio::spawn(async move {
+        let mut accepted = 0;
+        let mut served = 0;
+        while served < expected_requests {
+            let (mut stream, _) = listener.accept().await.unwrap();
+            accepted += 1;
+            loop {
+                if read_request(&mut stream).await.is_empty() {
+                    break;
+                }
+                served += 1;
+                stream.write_all(EMPTY_RESPONSE).await.unwrap();
+                if served == expected_requests {
+                    return accepted;
+                }
+            }
+        }
+        accepted
+    });
+    (format!("http://{addr}/"), server)
+}
+
+#[tokio::test]
+async fn pool_hits_refresh_lru_order() {
+    let (url_a, server_a) = spawn_counting_server(3).await;
+    let (url_b, server_b) = spawn_counting_server(1).await;
+    let (url_c, server_c) = spawn_counting_server(1).await;
+    let client = bitreq::Client::new(2);
+
+    for url in [&url_a, &url_b, &url_a, &url_c, &url_a] {
+        assert_eq!(client.send_async(bitreq::get(url)).await.unwrap().status_code, 200);
+    }
+
+    let accepted_a = tokio::time::timeout(Duration::from_secs(5), server_a).await.unwrap().unwrap();
+    let accepted_b = server_b.await.unwrap();
+    let accepted_c = server_c.await.unwrap();
+    assert_eq!(accepted_a, 1, "the recently used connection was evicted");
+    assert_eq!((accepted_b, accepted_c), (1, 1));
+}