hyper-util legacy client: each dropped pool `Checkout` scans every parked waiter under the pool lock, O(n²) at cold start

[package]
name = "pool-waiters-repro"
version = "0.0.0"
edition = "2021"
publish = false

[dependencies]
hyper-util = { version = "=0.1.20", features = ["client-legacy", "http1", "tokio"] }
hyper = { version = "1", features = ["client", "server", "http1"] }
tokio = { version = "1", features = ["full"] }
http-body-util = "0.1"
tower-service = "0.3"

//! Fire N concurrent GETs at one authority through hyper-util's legacy client,
//! starting from a cold pool. The connector hands out in-memory duplex pipes
//! served by hyper's HTTP/1 server, so there are no sockets or fd limits.
//!
//! Usage: cargo run --release -- n=65536 connect_ms=500 resp_ms=1000 cancel_pct=0 cancel_ms=10 threads=0 deadline_s=120

use std::convert::Infallible;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

use http_body_util::{BodyExt, Empty, Full};
use hyper::body::{Bytes, Incoming};
use hyper::service::service_fn;
use hyper::{Request, Response, Uri};
use hyper_util::client::legacy::connect::{Connected, Connection};
use hyper_util::client::legacy::Client;
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
use tokio::io::DuplexStream;

// Implements hyper's IO traits directly, since newer hyper-util gates
// `Connection for TokioIo<T>` behind the HTTP connector feature.
struct Pipe(TokioIo<DuplexStream>);

impl Connection for Pipe {
    fn connected(&self) -> Connected {
        Connected::new()
    }
}

impl hyper::rt::Read for Pipe {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: hyper::rt::ReadBufCursor<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.0).poll_read(cx, buf)
    }
}

impl hyper::rt::Write for Pipe {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        b: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.0).poll_write(cx, b)
    }
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.0).poll_flush(cx)
    }
    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.0).poll_shutdown(cx)
    }
}

#[derive(Clone)]
struct DuplexConnector {
    connect_delay: Duration,
    resp_delay: Duration,
    connects: Arc<AtomicU64>,
}

impl tower_service::Service<Uri> for DuplexConnector {
    type Response = Pipe;
    type Error = Infallible;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Infallible>> + Send>>;

    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, _: Uri) -> Self::Future {
        let this = self.clone();
        Box::pin(async move {
            if !this.connect_delay.is_zero() {
                tokio::time::sleep(this.connect_delay).await;
            }
            this.connects.fetch_add(1, Relaxed);
            let (client, server) = tokio::io::duplex(16 * 1024);
            tokio::spawn(serve(server, this.resp_delay));
            Ok(Pipe(TokioIo::new(client)))
        })
    }
}

async fn serve(io: DuplexStream, resp_delay: Duration) {
    let svc = service_fn(move |_req: Request<Incoming>| async move {
        if !resp_delay.is_zero() {
            tokio::time::sleep(resp_delay).await;
        }
        Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok"))))
    });
    let _ = hyper::server::conn::http1::Builder::new()
        .serve_connection(TokioIo::new(io), svc)
        .await;
}

struct Args {
    n: usize,
    connect_ms: u64,
    resp_ms: u64,
    cancel_pct: u64,
    cancel_ms: u64,
    threads: usize,
    deadline_s: u64,
}

fn parse_args() -> Args {
    let mut a = Args {
        n: 65536,
        connect_ms: 500,
        resp_ms: 1000,
        cancel_pct: 0,
        cancel_ms: 10,
        threads: 0,
        deadline_s: 120,
    };
    for kv in std::env::args().skip(1) {
        let (k, v) = kv.split_once('=').expect("args are key=value");
        let v: u64 = v.parse().expect("numeric value");
        match k {
            "n" => a.n = v as usize,
            "connect_ms" => a.connect_ms = v,
            "resp_ms" => a.resp_ms = v,
            "cancel_pct" => a.cancel_pct = v,
            "cancel_ms" => a.cancel_ms = v,
            "threads" => a.threads = v as usize,
            "deadline_s" => a.deadline_s = v,
            _ => panic!("unknown arg {k}"),
        }
    }
    a
}

fn main() {
    let args = parse_args();
    let mut rt = tokio::runtime::Builder::new_multi_thread();
    if args.threads > 0 {
        rt.worker_threads(args.threads);
    }
    let rt = rt.enable_all().build().unwrap();
    let threads = rt.metrics().num_workers();
    rt.block_on(run(args, threads));
    // Skip dropping tens of thousands of pooled connections.
    std::process::exit(0);
}

async fn run(args: Args, threads: usize) {
    let connects = Arc::new(AtomicU64::new(0));
    let connector = DuplexConnector {
        connect_delay: Duration::from_millis(args.connect_ms),
        resp_delay: Duration::from_millis(args.resp_ms),
        connects: connects.clone(),
    };
    // Mirrors reqwest's defaults: pool timer on, 90 s idle timeout, unbounded idle per host.
    let client: Client<_, Empty<Bytes>> = Client::builder(TokioExecutor::new())
        .pool_timer(TokioTimer::new())
        .pool_idle_timeout(Duration::from_secs(90))
        .build(connector);
    let uri: Uri = "http://bench.invalid/obj".parse().unwrap();

    let ok = Arc::new(AtomicU64::new(0));
    let canceled = Arc::new(AtomicU64::new(0));
    let first_ns = Arc::new(AtomicU64::new(u64::MAX));
    let t0 = Instant::now();

    let mut set = tokio::task::JoinSet::new();
    for i in 0..args.n {
        let client = client.clone();
        let uri = uri.clone();
        let (ok, canceled, first_ns) = (ok.clone(), canceled.clone(), first_ns.clone());
        let cancel = (i as u64 % 100) < args.cancel_pct;
        let cancel_after = Duration::from_millis(args.cancel_ms);
        set.spawn(async move {
            let req = async {
                let resp = client.get(uri).await.expect("request");
                resp.into_body().collect().await.expect("body");
            };
            let done = if cancel {
                tokio::time::timeout(cancel_after, req).await.is_ok()
            } else {
                req.await;
                true
            };
            if done {
                first_ns.fetch_min(t0.elapsed().as_nanos() as u64, Relaxed);
                ok.fetch_add(1, Relaxed);
            } else {
                canceled.fetch_add(1, Relaxed);
            }
        });
    }
    let spawn_ms = t0.elapsed().as_secs_f64() * 1e3;

    let deadline = Duration::from_secs(args.deadline_s);
    let all = async { while set.join_next().await.is_some() {} };
    let finished = tokio::time::timeout(deadline, all).await.is_ok();
    let wall = t0.elapsed().as_secs_f64();

    let first = first_ns.load(Relaxed);
    let first_ms = if first == u64::MAX {
        f64::NAN
    } else {
        first as f64 / 1e6
    };
    print!(
        "crate={} threads={} n={} connect_ms={} resp_ms={} cancel_pct={} finished={} wall_s={:.3} spawn_ms={:.0} first_ms={:.1} ok={} canceled={} connects={}",
        env!("CARGO_PKG_NAME"),
        threads,
        args.n,
        args.connect_ms,
        args.resp_ms,
        args.cancel_pct,
        finished,
        wall,
        spawn_ms,
        first_ms,
        ok.load(Relaxed),
        canceled.load(Relaxed),
        connects.load(Relaxed),
    );
    println!();
}

--- a/src/client/legacy/pool.rs
+++ b/src/client/legacy/pool.rs
@@ -92,7 +92,7 @@
     // this list is checked for any parked Checkouts, and tries to notify
     // them that the Conn could be used instead of waiting for a brand new
     // connection.
-    waiters: HashMap<K, VecDeque<oneshot::Sender<T>>>,
+    waiters: HashMap<K, Waiters<T>>,
     // A oneshot channel is used to allow the interval to be notified when
     // the Pool completely drops. That way, the interval can cancel immediately.
     idle_interval_ref: Option<oneshot::Sender<Infallible>>,
@@ -101,6 +101,23 @@
     timeout: Option<Duration>,
 }
 
+/// The parked `Checkout`s for one key.
+struct Waiters<T> {
+    queue: VecDeque<oneshot::Sender<T>>,
+    // Checkouts dropped since the queue was last purged. Purging only once
+    // this reaches half the queue keeps each drop amortized O(1).
+    canceled: usize,
+}
+
+impl<T> Waiters<T> {
+    fn new() -> Self {
+        Waiters {
+            queue: VecDeque::new(),
+            canceled: 0,
+        }
+    }
+}
+
 // This is because `Weak::new()` *allocates* space for `T`, even if it
 // doesn't need it!
 struct WeakOpt<T>(Option<Weak<T>>);
@@ -354,7 +371,7 @@
         let mut remove_waiters = false;
         let mut value = Some(value);
         if let Some(waiters) = self.waiters.get_mut(&key) {
-            while let Some(tx) = waiters.pop_front() {
+            while let Some(tx) = waiters.queue.pop_front() {
                 if !tx.is_canceled() {
                     let reserved = value.take().expect("value already sent");
                     let reserved = match reserved.reserve() {
@@ -381,7 +398,7 @@
 
                 trace!("put; removing canceled waiter for {:?}", key);
             }
-            remove_waiters = waiters.is_empty();
+            remove_waiters = waiters.queue.is_empty();
         }
         if remove_waiters {
             self.waiters.remove(&key);
@@ -466,11 +483,20 @@
     /// and possibly inserted into the pool that it is waiting for an idle
     /// connection. If a user ever dropped that future, we need to clean out
     /// those parked senders.
+    ///
+    /// Scanning the whole queue on every drop is O(n^2) under the pool lock
+    /// when many checkouts are parked, so the scan waits until at least half
+    /// the queue may be canceled. That bounds the queue to about twice the
+    /// live waiters.
     fn clean_waiters(&mut self, key: &K) {
         let mut remove_waiters = false;
         if let Some(waiters) = self.waiters.get_mut(key) {
-            waiters.retain(|tx| !tx.is_canceled());
-            remove_waiters = waiters.is_empty();
+            waiters.canceled += 1;
+            if waiters.canceled * 2 >= waiters.queue.len() {
+                waiters.queue.retain(|tx| !tx.is_canceled());
+                waiters.canceled = 0;
+                remove_waiters = waiters.queue.is_empty();
+            }
         }
         if remove_waiters {
             self.waiters.remove(key);
@@ -687,7 +713,8 @@
                 inner
                     .waiters
                     .entry(self.key.clone())
-                    .or_insert_with(VecDeque::new)
+                    .or_insert_with(Waiters::new)
+                    .queue
                     .push_back(tx);
 
                 // register the waker with this oneshot
@@ -1075,18 +1102,45 @@
 
         // first poll needed to get into Pool's parked
         poll_once1.await;
-        assert_eq!(pool.locked().waiters.get(&key).unwrap().len(), 1);
+        assert_eq!(pool.locked().waiters.get(&key).unwrap().queue.len(), 1);
         poll_once2.await;
-        assert_eq!(pool.locked().waiters.get(&key).unwrap().len(), 2);
+        assert_eq!(pool.locked().waiters.get(&key).unwrap().queue.len(), 2);
 
         // on drop, clean up Pool
         drop(checkout1);
-        assert_eq!(pool.locked().waiters.get(&key).unwrap().len(), 1);
+        assert_eq!(pool.locked().waiters.get(&key).unwrap().queue.len(), 1);
 
         drop(checkout2);
         assert!(!pool.locked().waiters.contains_key(&key));
     }
 
+    #[tokio::test]
+    async fn test_pool_checkout_drop_cleans_up_waiters_amortized() {
+        let pool = pool_no_timer::<Uniq<i32>, KeyImpl>();
+        let key = host_key("foo");
+
+        let mut checkouts: Vec<_> = (0..64).map(|_| pool.checkout(key.clone())).collect();
+        for checkout in &mut checkouts {
+            PollOnce(checkout).await;
+        }
+        assert_eq!(pool.locked().waiters.get(&key).unwrap().queue.len(), 64);
+
+        // A single drop does not scan the queue.
+        checkouts.pop();
+        assert_eq!(pool.locked().waiters.get(&key).unwrap().queue.len(), 64);
+
+        // Canceled senders never outnumber the live ones.
+        while checkouts.pop().is_some() {
+            let live = checkouts.len();
+            let queued = pool.locked().waiters.get(&key).map_or(0, |w| w.queue.len());
+            assert!(
+                queued <= 2 * live,
+                "queued {queued} for {live} live waiters"
+            );
+        }
+        assert!(!pool.locked().waiters.contains_key(&key));
+    }
+
     #[derive(Debug)]
     struct CanClose {
         #[allow(unused)]