diff --git a/benches/main.rs b/benches/main.rs index b1e64edf4..55fcff85f 100644 --- a/benches/main.rs +++ b/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> { @@ -111,8 +162,174 @@ async fn send_requests(addr: &str) -> Result<(), Box> { Ok(()) } +async fn write_contention_benchmark() -> Result<(), Box> { + 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> { + 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, + mut respond: SendResponse, + config: WriteContentionConfig, +) -> Result<(), Box> { + 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, + 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, +} + +impl Future for Capacity<'_> { + type Output = Option>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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(); + } } diff --git a/src/codec/framed_write.rs b/src/codec/framed_write.rs index 17e557623..547270460 100644 --- a/src/codec/framed_write.rs +++ b/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 diff --git a/src/codec/mod.rs b/src/codec/mod.rs index 6cbdc1e18..d7aecee91 100644 --- a/src/codec/mod.rs +++ b/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 diff --git a/src/proto/streams/prioritize.rs b/src/proto/streams/prioritize.rs index 1b0b82fda..df39cf3b7 100644 --- a/src/proto/streams/prioritize.rs +++ b/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( + pub fn buffer_pending( &mut self, - cx: &mut Context, buffer: &mut Buffer>, store: &mut Store, counts: &mut Counts, dst: &mut Codec>, - ) -> Poll> + ) -> io::Result 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( + &mut self, + buffer: &mut Buffer>, + store: &mut Store, + dst: &mut Codec>, + ) -> 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. diff --git a/src/proto/streams/recv.rs b/src/proto/streams/recv.rs index a21b74cf4..481d1c55a 100644 --- a/src/proto/streams/recv.rs +++ b/src/proto/streams/recv.rs @@ -1005,15 +1005,16 @@ impl Recv { /// Send any pending refusals. pub fn send_pending_refusal( &mut self, - cx: &mut Context, dst: &mut Codec>, - ) -> Poll> + ) -> io::Result where T: AsyncWrite + Unpin, B: Buf, { if let Some(stream_id) = self.refused { - ready!(dst.poll_ready(cx))?; + if !dst.has_send_capacity() { + return Ok(false); + } // Create the RST_STREAM frame let frame = frame::Reset::new(stream_id, Reason::REFUSED_STREAM); @@ -1024,7 +1025,7 @@ impl Recv { self.refused = None; - Poll::Ready(Ok(())) + Ok(true) } pub fn clear_expired_reset_streams(&mut self, store: &mut Store, counts: &mut Counts) { @@ -1078,32 +1079,34 @@ impl Recv { } } - pub fn poll_complete( + pub fn buffer_pending( &mut self, - cx: &mut Context, store: &mut Store, counts: &mut Counts, dst: &mut Codec>, - ) -> Poll> + ) -> io::Result where T: AsyncWrite + Unpin, B: Buf, { // Send any pending connection level window updates - ready!(self.send_connection_window_update(cx, dst))?; + if !self.send_connection_window_update(dst)? { + return Ok(false); + } // Send any pending stream level window updates - ready!(self.send_stream_window_updates(cx, store, counts, dst))?; + if !self.send_stream_window_updates(store, counts, dst)? { + return Ok(false); + } - Poll::Ready(Ok(())) + Ok(true) } /// Send connection level window update fn send_connection_window_update( &mut self, - cx: &mut Context, dst: &mut Codec>, - ) -> Poll> + ) -> io::Result where T: AsyncWrite + Unpin, B: Buf, @@ -1112,7 +1115,9 @@ impl Recv { let frame = frame::WindowUpdate::new(StreamId::zero(), incr); // Ensure the codec has capacity - ready!(dst.poll_ready(cx))?; + if !dst.has_send_capacity() { + return Ok(false); + } // Buffer the WINDOW_UPDATE frame dst.buffer(frame.into()) @@ -1124,29 +1129,30 @@ impl Recv { .expect("unexpected flow control state"); } - Poll::Ready(Ok(())) + Ok(true) } /// Send stream level window update pub fn send_stream_window_updates( &mut self, - cx: &mut Context, store: &mut Store, counts: &mut Counts, dst: &mut Codec>, - ) -> Poll> + ) -> io::Result where T: AsyncWrite + Unpin, B: Buf, { loop { // Ensure the codec has capacity - ready!(dst.poll_ready(cx))?; + if !dst.has_send_capacity() { + return Ok(false); + } // Get the next stream let stream = match self.pending_window_updates.pop(store) { Some(stream) => stream, - None => return Poll::Ready(Ok(())), + None => return Ok(true), }; counts.transition(stream, |_, stream| { diff --git a/src/proto/streams/send.rs b/src/proto/streams/send.rs index 4dd114950..72cb65aaa 100644 --- a/src/proto/streams/send.rs +++ b/src/proto/streams/send.rs @@ -334,20 +334,30 @@ impl Send { Ok(()) } - pub fn poll_complete( + pub fn buffer_pending( &mut self, - cx: &mut Context, buffer: &mut Buffer>, store: &mut Store, counts: &mut Counts, dst: &mut Codec>, - ) -> Poll> + ) -> io::Result where T: AsyncWrite + Unpin, B: Buf, { - self.prioritize - .poll_complete(cx, buffer, store, counts, dst) + self.prioritize.buffer_pending(buffer, store, counts, dst) + } + + pub fn reclaim_written_frame( + &mut self, + buffer: &mut Buffer>, + store: &mut Store, + dst: &mut Codec>, + ) -> bool + where + B: Buf, + { + self.prioritize.reclaim_written_frame(buffer, store, dst) } /// Request capacity to send data diff --git a/src/proto/streams/streams.rs b/src/proto/streams/streams.rs index 33bcfd67d..9fb5ed47c 100644 --- a/src/proto/streams/streams.rs +++ b/src/proto/streams/streams.rs @@ -161,9 +161,15 @@ where where T: AsyncWrite + Unpin, { + ready!(dst.poll_ready(cx))?; + let mut me = self.inner.lock().unwrap(); let me = &mut *me; - me.actions.recv.send_pending_refusal(cx, dst) + if me.actions.recv.send_pending_refusal(dst)? { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } } pub fn clear_expired_reset_streams(&mut self) { @@ -182,8 +188,36 @@ where where T: AsyncWrite + Unpin, { - let mut me = self.inner.lock().unwrap(); - me.poll_complete(&self.send_buffer, cx, dst) + loop { + // Make any required socket progress before taking stream locks. + ready!(dst.poll_ready(cx))?; + + let drained = { + let mut me = self.inner.lock().unwrap(); + me.buffer_pending(&self.send_buffer, dst)? + }; + + if !drained { + continue; + } + + // Flush any frames staged by `buffer_pending` without holding the + // stream-state or send-buffer mutexes. + ready!(dst.flush(cx))?; + + let reclaimed = { + let mut me = self.inner.lock().unwrap(); + let reclaimed = me.reclaim_written_frame(&self.send_buffer, dst); + if !reclaimed { + me.actions.task = Some(cx.waker().clone()); + } + reclaimed + }; + + if !reclaimed { + return Poll::Ready(Ok(())); + } + } } pub fn apply_remote_settings( @@ -902,12 +936,11 @@ impl Inner { Ok(()) } - fn poll_complete( + fn buffer_pending( &mut self, send_buffer: &SendBuffer, - cx: &mut Context, dst: &mut Codec>, - ) -> Poll> + ) -> io::Result where T: AsyncWrite + Unpin, B: Buf, @@ -919,24 +952,40 @@ impl Inner { // // TODO: It would probably be better to interleave updates w/ data // frames. - ready!(self + if !self .actions .recv - .poll_complete(cx, &mut self.store, &mut self.counts, dst))?; + .buffer_pending(&mut self.store, &mut self.counts, dst)? + { + return Ok(false); + } // Send any other pending frames - ready!(self.actions.send.poll_complete( - cx, - send_buffer, - &mut self.store, - &mut self.counts, - dst - ))?; + if !self + .actions + .send + .buffer_pending(send_buffer, &mut self.store, &mut self.counts, dst)? + { + return Ok(false); + } - // Nothing else to do, track the task - self.actions.task = Some(cx.waker().clone()); + Ok(true) + } - Poll::Ready(Ok(())) + fn reclaim_written_frame( + &mut self, + send_buffer: &SendBuffer, + dst: &mut Codec>, + ) -> bool + where + B: Buf, + { + let mut send_buffer = send_buffer.inner.lock().unwrap(); + let send_buffer = &mut *send_buffer; + + self.actions + .send + .reclaim_written_frame(send_buffer, &mut self.store, dst) } fn send_reset( diff --git a/tests/h2-tests/tests/flow_control.rs b/tests/h2-tests/tests/flow_control.rs index 2f355fd50..06dbeb756 100644 --- a/tests/h2-tests/tests/flow_control.rs +++ b/tests/h2-tests/tests/flow_control.rs @@ -836,6 +836,43 @@ async fn recv_window_update_on_stream_closed_by_data_frame() { join(srv, h2).await; } +#[tokio::test] +async fn small_data_frames_reclaimed_before_buffering_next() { + h2_support::trace_init!(); + let (io, mut srv) = mock::new(); + + let h2 = async move { + let (mut client, mut h2) = client::handshake(io).await.unwrap(); + let request = Request::builder() + .method(Method::POST) + .uri("https://http2.akamai.com/") + .body(()) + .unwrap(); + + let (response, mut stream) = client.send_request(request, false).unwrap(); + + stream.send_data("hello".into(), false).unwrap(); + stream.send_data("world".into(), true).unwrap(); + + let response = h2.drive(response).await.unwrap(); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + h2.await.unwrap(); + }; + + let srv = async move { + let settings = srv.assert_client_handshake().await; + assert_default_settings!(settings); + srv.recv_frame(frames::headers(1).request("POST", "https://http2.akamai.com/")) + .await; + srv.recv_frame(frames::data(1, "hello")).await; + srv.recv_frame(frames::data(1, "world").eos()).await; + srv.send_frame(frames::headers(1).response(204).eos()).await; + }; + + join(srv, h2).await; +} + #[tokio::test] async fn reserved_capacity_assigned_in_multi_window_updates() { h2_support::trace_init!();