HTTP/2 client sends HEADERS and an immediately ready Full body in separate socket writes

[package]
name = "unary-h2-writes"
version = "0.1.0"
edition = "2021"

[dependencies]
bytes = "1"
h2 = "=0.4.20"
http = "1"
http-body-util = "0.1"
hyper = { version = "=1.12.0", features = ["client", "http2"] }
hyper-util = { version = "=0.1.21", features = ["tokio"] }
tokio = { version = "=1.53.2", features = ["macros", "rt", "net", "io-util", "time"] }
tonic = "=0.14.6"
tonic-prost = "=0.14.6"
prost = "0.14"
tower = { version = "0.5", features = ["util"] }
parking_lot = "0.12"

use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper_util::rt::{TokioExecutor, TokioIo};
use std::io::{self, IoSlice};
use std::pin::Pin;
use std::sync::Arc;
use parking_lot::Mutex;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::{TcpListener, TcpStream};

type Trace = Arc<Mutex<Vec<Vec<u8>>>>;

struct Counted {
    inner: TcpStream,
    trace: Trace,
}

impl AsyncRead for Counted {
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_read(cx, buf)
    }
}

impl AsyncWrite for Counted {
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        match Pin::new(&mut self.inner).poll_write(cx, buf) {
            Poll::Ready(Ok(n)) => {
                self.trace.lock().push(buf[..n].to_vec());
                Poll::Ready(Ok(n))
            }
            result => result,
        }
    }

    fn poll_write_vectored(mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>]) -> Poll<io::Result<usize>> {
        match Pin::new(&mut self.inner).poll_write_vectored(cx, bufs) {
            Poll::Ready(Ok(n)) => {
                let mut bytes = Vec::with_capacity(n);
                let mut remaining = n;
                for buf in bufs {
                    let used = remaining.min(buf.len());
                    bytes.extend_from_slice(&buf[..used]);
                    remaining -= used;
                    if remaining == 0 { break; }
                }
                self.trace.lock().push(bytes);
                Poll::Ready(Ok(n))
            }
            result => result,
        }
    }

    fn is_write_vectored(&self) -> bool { true }
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }
    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

async fn server() -> std::net::SocketAddr {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        while let Ok((stream, _)) = listener.accept().await {
            stream.set_nodelay(true).unwrap();
            tokio::spawn(async move {
                let mut connection = h2::server::handshake(stream).await.unwrap();
                while let Some(Ok((request, mut respond))) = connection.accept().await {
                    tokio::spawn(async move {
                        let mut body = request.into_body();
                        let mut bytes = Vec::new();
                        while let Some(chunk) = body.data().await {
                            let chunk = chunk.unwrap();
                            body.flow_control().release_capacity(chunk.len()).unwrap();
                            bytes.extend_from_slice(&chunk);
                        }
                        let response = http::Response::builder().header("content-type", "application/grpc").body(()).unwrap();
                        let mut stream = respond.send_response(response, false).unwrap();
                        stream.send_data(Bytes::from(bytes), false).unwrap();
                        let mut trailers = http::HeaderMap::new();
                        trailers.insert("grpc-status", http::HeaderValue::from_static("0"));
                        stream.send_trailers(trailers).unwrap();
                    });
                }
            });
        }
    });
    addr
}

fn print_trace(name: &str, trace: &Trace) {
    let writes = trace.lock();
    println!("{name}: {} socket writes", writes.len());
    for (i, write) in writes.iter().enumerate() {
        print!("  write {}:", i + 1);
        let mut bytes = write.as_slice();
        while bytes.len() >= 9 {
            let len = ((bytes[0] as usize) << 16) | ((bytes[1] as usize) << 8) | bytes[2] as usize;
            let kind = match bytes[3] { 0 => "DATA", 1 => "HEADERS", 4 => "SETTINGS", 8 => "WINDOW_UPDATE", _ => "OTHER" };
            let stream = u32::from_be_bytes(bytes[5..9].try_into().unwrap()) & 0x7fff_ffff;
            print!(" {kind}(stream={stream},len={len},END_STREAM={})", bytes[4] & 1 != 0);
            if bytes.len() < 9 + len { break; }
            bytes = &bytes[9 + len..];
        }
        println!();
    }
}

#[derive(Clone, PartialEq, prost::Message)]
struct Message {
    #[prost(string, tag = "1")]
    value: String,
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let addr = server().await;
    let trace: Trace = Arc::default();
    let stream = TcpStream::connect(addr).await?;
    stream.set_nodelay(true)?;
    let (mut sender, connection) = hyper::client::conn::http2::Builder::new(TokioExecutor::new())
        .handshake(TokioIo::new(Counted { inner: stream, trace: trace.clone() })).await?;
    tokio::spawn(async move { let _ = connection.await; });
    for round in 0..2 {
        if round == 1 { trace.lock().clear(); }
        let request = http::Request::post(format!("http://{addr}/echo"))
            .body(Full::new(Bytes::from_static(b"hello")))?;
        let response = sender.send_request(request).await?;
        assert_eq!(response.into_body().collect().await?.to_bytes(), "hello");
    }
    print_trace("hyper Full body", &trace);

    let trace: Trace = Arc::default();
    let stream = TcpStream::connect(addr).await?;
    stream.set_nodelay(true)?;
    let (mut sender, connection) = h2::client::handshake(Counted { inner: stream, trace: trace.clone() }).await?;
    tokio::spawn(async move { let _ = connection.await; });
    for round in 0..2 {
        if round == 1 { trace.lock().clear(); }
        sender = sender.ready().await?;
        let request = http::Request::post(format!("http://{addr}/echo")).body(())?;
        let (response, mut send) = sender.send_request(request, false)?;
        send.send_data(Bytes::from_static(b"hello"), true)?;
        let mut body = response.await?.into_body();
        let mut bytes = Vec::new();
        while let Some(chunk) = body.data().await {
            let chunk = chunk?;
            body.flow_control().release_capacity(chunk.len())?;
            bytes.extend_from_slice(&chunk);
        }
        assert_eq!(bytes, b"hello");
    }
    print_trace("h2 immediate body", &trace);

    let trace: Trace = Arc::default();
    let captured = trace.clone();
    let channel = tonic::transport::Endpoint::from_shared(format!("http://{addr}"))?
        .connect_with_connector(tower::service_fn(move |_uri| {
            let trace = captured.clone();
            async move {
                let inner = TcpStream::connect(addr).await?;
                inner.set_nodelay(true)?;
                Ok::<_, io::Error>(TokioIo::new(Counted { inner, trace }))
            }
        })).await?;
    let mut grpc = tonic::client::Grpc::new(channel);
    for round in 0..2 {
        if round == 1 { trace.lock().clear(); }
        grpc.ready().await?;
        let response: tonic::Response<Message> = grpc.unary(
            tonic::Request::new(Message { value: "hello".into() }),
            http::uri::PathAndQuery::from_static("/echo.Echo/Unary"),
            tonic_prost::ProstCodec::default(),
        ).await?;
        assert_eq!(response.into_inner().value, "hello");
    }
    print_trace("tonic unary", &trace);
    Ok(())
}