Avoid flushing h2 socket while holding stream locks

benches/main.rs

@@ -8,12 +8,63 @@ use http::Request;
 
 use std::{
     error::Error,
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
     time::{Duration, Instant},
 };
 
 use tokio::net::{TcpListener, TcpStream};
 
 const NUM_REQUESTS_TO_SEND: usize = 100_000;
+const WRITE_CONTENTION_STREAMS: usize = 512;
+const WRITE_CONTENTION_CHUNKS_PER_STREAM: usize = 128;
+const WRITE_CONTENTION_CHUNK_SIZE: usize = 16 * 1024;
+const WRITE_CONTENTION_WINDOW: u32 = 16 * 1024 * 1024;
+const WRITE_CONTENTION_MAX_FRAME_SIZE: u32 = 64 * 1024;
+const WRITE_CONTENTION_MAX_SEND_BUFFER: usize = 8 * 1024 * 1024;
+const WRITE_CONTENTION_WORKER_THREADS: usize = 4;
+
+#[derive(Clone, Copy)]
+struct WriteContentionConfig {
+    streams: usize,
+    chunks_per_stream: usize,
+    chunk_size: usize,
+}
+
+impl WriteContentionConfig {
+    fn from_env() -> Self {
+        Self {
+            streams: env_usize("H2_WRITE_CONTENTION_STREAMS", WRITE_CONTENTION_STREAMS),
+            chunks_per_stream: env_usize(
+                "H2_WRITE_CONTENTION_CHUNKS_PER_STREAM",
+                WRITE_CONTENTION_CHUNKS_PER_STREAM,
+            ),
+            chunk_size: env_usize(
+                "H2_WRITE_CONTENTION_CHUNK_SIZE",
+                WRITE_CONTENTION_CHUNK_SIZE,
+            ),
+        }
+    }
+
+    fn bytes(&self) -> usize {
+        self.streams * self.chunks_per_stream * self.chunk_size
+    }
+}
+
+fn write_contention_worker_threads() -> usize {
+    env_usize(
+        "H2_WRITE_CONTENTION_WORKER_THREADS",
+        WRITE_CONTENTION_WORKER_THREADS,
+    )
+}
+
+fn env_usize(name: &str, default: usize) -> usize {
+    std::env::var(name)
+        .ok()
+        .and_then(|value| value.parse().ok())
+        .unwrap_or(default)
+}
 
 // The actual server.
 async fn server(addr: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
@@ -111,8 +162,174 @@ async fn send_requests(addr: &str) -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
+async fn write_contention_benchmark() -> Result<(), Box<dyn Error>> {
+    let config = WriteContentionConfig::from_env();
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let addr = listener.local_addr()?;
+
+    println!(
+        "H2 write contention: {} streams x {} chunks x {}B = {:.1} MiB",
+        config.streams,
+        config.chunks_per_stream,
+        config.chunk_size,
+        config.bytes() as f64 / (1024.0 * 1024.0),
+    );
+
+    tokio::spawn(async move {
+        let (socket, _peer_addr) = listener.accept().await.unwrap();
+        if let Err(e) = serve_write_contention(socket, config).await {
+            println!("write contention server error: {e:?}");
+        }
+    });
+
+    let tcp = TcpStream::connect(addr).await?;
+    let mut builder = client::Builder::new();
+    builder
+        .initial_window_size(WRITE_CONTENTION_WINDOW)
+        .initial_connection_window_size(WRITE_CONTENTION_WINDOW)
+        .max_frame_size(WRITE_CONTENTION_MAX_FRAME_SIZE)
+        .max_concurrent_streams(config.streams as u32)
+        .max_send_buffer_size(WRITE_CONTENTION_MAX_SEND_BUFFER);
+
+    let (client, h2) = builder.handshake::<_, Bytes>(tcp).await?;
+    tokio::spawn(async move {
+        if let Err(e) = h2.await {
+            println!("write contention client connection error: {e:?}");
+        }
+    });
+
+    let mut handles = Vec::with_capacity(config.streams);
+    let started = Instant::now();
+    for _ in 0..config.streams {
+        let client = client.clone();
+        let expected = config.chunks_per_stream * config.chunk_size;
+        handles.push(tokio::spawn(async move {
+            let request = Request::builder().body(()).unwrap();
+            let mut client = client.ready().await.unwrap();
+            let (response, _) = client.send_request(request, true).unwrap();
+            let response = response.await.unwrap();
+            let mut body = response.into_body();
+            let mut received = 0;
+
+            while let Some(chunk) = body.data().await {
+                let chunk = chunk.unwrap();
+                received += chunk.len();
+                let _ = body.flow_control().release_capacity(chunk.len());
+            }
+
+            assert_eq!(received, expected);
+        }));
+    }
+
+    for handle in handles {
+        handle.await.unwrap();
+    }
+
+    let elapsed = started.elapsed();
+    let mib = config.bytes() as f64 / (1024.0 * 1024.0);
+    println!("Overall: {}ms.", elapsed.as_millis());
+    println!("Throughput: {:.1} MiB/s", mib / elapsed.as_secs_f64());
+
+    Ok(())
+}
+
+async fn serve_write_contention(
+    socket: TcpStream,
+    config: WriteContentionConfig,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
+    let mut builder = server::Builder::new();
+    builder
+        .initial_window_size(WRITE_CONTENTION_WINDOW)
+        .initial_connection_window_size(WRITE_CONTENTION_WINDOW)
+        .max_frame_size(WRITE_CONTENTION_MAX_FRAME_SIZE)
+        .max_concurrent_streams(config.streams as u32)
+        .max_send_buffer_size(WRITE_CONTENTION_MAX_SEND_BUFFER);
+
+    let mut connection = builder.handshake(socket).await?;
+    while let Some(result) = connection.accept().await {
+        let (request, respond) = result?;
+        tokio::spawn(async move {
+            if let Err(e) = handle_write_contention_request(request, respond, config).await {
+                println!("write contention request error: {e}");
+            }
+        });
+    }
+
+    Ok(())
+}
+
+async fn handle_write_contention_request(
+    mut request: Request<RecvStream>,
+    mut respond: SendResponse<Bytes>,
+    config: WriteContentionConfig,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
+    let body = request.body_mut();
+    while let Some(data) = body.data().await {
+        let data = data?;
+        let _ = body.flow_control().release_capacity(data.len());
+    }
+
+    let response = http::Response::new(());
+    let mut send = respond.send_response(response, false)?;
+    let chunk = Bytes::from(vec![b'x'; config.chunk_size]);
+
+    for idx in 0..config.chunks_per_stream {
+        let end_of_stream = idx + 1 == config.chunks_per_stream;
+        send_chunk(&mut send, chunk.clone(), end_of_stream).await?;
+    }
+
+    Ok(())
+}
+
+async fn send_chunk(
+    send: &mut h2::SendStream<Bytes>,
+    chunk: Bytes,
+    end_of_stream: bool,
+) -> Result<(), h2::Error> {
+    let len = chunk.len();
+    send.reserve_capacity(len);
+    loop {
+        if send.capacity() >= len {
+            send.send_data(chunk, end_of_stream)?;
+            return Ok(());
+        }
+
+        match (Capacity { send }).await {
+            Some(Ok(_)) => {}
+            Some(Err(err)) => return Err(err),
+            None => return Err(h2::Reason::INTERNAL_ERROR.into()),
+        }
+    }
+}
+
+struct Capacity<'a> {
+    send: &'a mut h2::SendStream<Bytes>,
+}
+
+impl Future for Capacity<'_> {
+    type Output = Option<Result<usize, h2::Error>>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.send.poll_capacity(cx)
+    }
+}
+
 fn main() {
     let _ = env_logger::try_init();
+    let bench = std::env::var("H2_BENCH").unwrap_or_else(|_| "all".to_string());
+
+    if bench == "write-contention" {
+        let worker_threads = write_contention_worker_threads();
+        println!("H2 write contention worker threads: {worker_threads}");
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(worker_threads)
+            .enable_all()
+            .build()
+            .unwrap();
+        rt.block_on(write_contention_benchmark()).unwrap();
+        return;
+    }
+
     let addr = "127.0.0.1:5928";
     println!("H2 running in current-thread runtime at {addr}:");
     std::thread::spawn(|| {
@@ -145,4 +362,15 @@ fn main() {
         .build()
         .unwrap();
     rt.block_on(send_requests(addr)).unwrap();
+
+    if bench == "all" {
+        let worker_threads = write_contention_worker_threads();
+        println!("H2 write contention worker threads: {worker_threads}");
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(worker_threads)
+            .enable_all()
+            .build()
+            .unwrap();
+        rt.block_on(write_contention_benchmark()).unwrap();
+    }
 }

src/codec/framed_write.rs

@@ -120,6 +120,12 @@ where
         Poll::Ready(Ok(()))
     }
 
+    /// Returns whether a frame can be buffered without first flushing the
+    /// underlying I/O object.
+    pub(crate) fn has_capacity(&self) -> bool {
+        self.encoder.has_capacity()
+    }
+
     /// Buffer a frame.
     ///
     /// `poll_ready` must be called first to ensure that a frame may be

src/codec/mod.rs

@@ -136,6 +136,12 @@ where
         self.framed_write().poll_ready(cx)
     }
 
+    /// Returns whether the codec can buffer a frame without flushing the
+    /// underlying I/O object.
+    pub(crate) fn has_send_capacity(&mut self) -> bool {
+        self.framed_write().has_capacity()
+    }
+
     /// Buffer a frame.
     ///
     /// `poll_ready` must be called first to ensure that a frame may be

src/proto/streams/prioritize.rs

@@ -10,7 +10,7 @@ use bytes::buf::Take;
 use std::{
     cmp::{self, Ordering},
     fmt, io, mem,
-    task::{Context, Poll, Waker},
+    task::Waker,
 };
 
 /// # Warning
@@ -505,30 +505,30 @@ impl Prioritize {
         }
     }
 
-    pub fn poll_complete<T, B>(
+    pub fn buffer_pending<T, B>(
         &mut self,
-        cx: &mut Context,
         buffer: &mut Buffer<Frame<B>>,
         store: &mut Store,
         counts: &mut Counts,
         dst: &mut Codec<T, Prioritized<B>>,
-    ) -> Poll<io::Result<()>>
+    ) -> io::Result<bool>
     where
         T: AsyncWrite + Unpin,
         B: Buf,
     {
-        // Ensure codec is ready
-        ready!(dst.poll_ready(cx))?;
-
         // Reclaim any frame that has previously been written
         self.reclaim_frame(buffer, store, dst);
 
         // The max frame length
         let max_frame_len = dst.max_send_frame_size();
 
-        tracing::trace!("poll_complete");
+        tracing::trace!("buffer_pending");
 
         loop {
+            if !dst.has_send_capacity() {
+                return Ok(false);
+            }
+
             if let Some(mut stream) = self.pop_pending_open(store, counts) {
                 self.pending_send.push_front(&mut stream);
                 self.try_assign_capacity(&mut stream);
@@ -544,28 +544,31 @@ impl Prioritize {
                     }
                     dst.buffer(frame).expect("invalid frame");
 
-                    // Ensure the codec is ready to try the loop again.
-                    ready!(dst.poll_ready(cx))?;
-
-                    // Because, always try to reclaim...
+                    // Small DATA frames can be fully encoded by `buffer`,
+                    // which records completion in a single codec slot. Reclaim
+                    // before accepting another frame so that slot is not
+                    // overwritten.
                     self.reclaim_frame(buffer, store, dst);
                 }
                 None => {
-                    // Try to flush the codec.
-                    ready!(dst.flush(cx))?;
-
-                    // This might release a data frame...
-                    if !self.reclaim_frame(buffer, store, dst) {
-                        return Poll::Ready(Ok(()));
-                    }
-
-                    // No need to poll ready as poll_complete() does this for
-                    // us...
+                    return Ok(true);
                 }
             }
         }
     }
 
+    pub fn reclaim_written_frame<T, B>(
+        &mut self,
+        buffer: &mut Buffer<Frame<B>>,
+        store: &mut Store,
+        dst: &mut Codec<T, Prioritized<B>>,
+    ) -> bool
+    where
+        B: Buf,
+    {
+        self.reclaim_frame(buffer, store, dst)
+    }
+
     /// Tries to reclaim a pending data frame from the codec.
     ///
     /// Returns true if a frame was reclaimed.