Incoming 与 body 通道

pub struct Incoming {
    kind: Kind,
}

enum Kind {
    Empty,
    #[cfg(all(feature = "http1", any(feature = "client", feature = "server")))]
    Chan {
        content_length: DecodedLength,
        rx: chan::Receiver,
    },
    #[cfg(all(feature = "http2", any(feature = "client", feature = "server")))]
    H2 {
        content_length: DecodedLength,
        data_done: bool,
        ping: ping::Recorder,
        recv: h2::RecvStream,
    },
    #[cfg(feature = "ffi")]
    Ffi(crate::ffi::UserBody),
}

            Kind::Chan {
                content_length: len,
                rx,
            } => {
                if let Some(chunk) = ready!(rx.poll_next(cx)?) {
                    len.sub_if(chunk.len() as u64);
                    return Poll::Ready(Some(Ok(Frame::data(chunk))));
                }
                Poll::Ready(rx.take_trailers().map(Frame::trailers).map(Ok))
            }

#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) struct DecodedLength(u64);

// ...
const MAX_LEN: u64 = u64::MAX - 2;

impl DecodedLength {
    pub(crate) const CLOSE_DELIMITED: DecodedLength = DecodedLength(u64::MAX);
    pub(crate) const CHUNKED: DecodedLength = DecodedLength(u64::MAX - 1);
    pub(crate) const ZERO: DecodedLength = DecodedLength(0);

struct Shared {
    state: Mutex<State>,
    sender_waker: AtomicWaker,
    receiver_waker: AtomicWaker,
}

struct State {
    item: Option<Result<Bytes, crate::Error>>,
    // An error must not displace data which was already accepted. The old
    // mpsc channel achieved this by sending the error from a cloned sender.
    pending_error: Option<crate::Error>,
    trailers: Option<HeaderMap>,
    sender_open: bool,
    receiver_open: bool,
    want: bool,
}

    pub(crate) fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
        self.shared.sender_waker.register(cx.waker());
        let state = self.shared.state.lock().panic_if_poisoned();
        if !state.receiver_open {
            Poll::Ready(Err(crate::Error::new_closed()))
        } else if state.want && state.item.is_none() && state.pending_error.is_none() {
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

        self.shared.receiver_waker.register(cx.waker());
        let mut state = self.shared.state.lock().panic_if_poisoned();
        let wake_sender = if !state.want {
            state.want = true;
            true
        } else {
            false
        };
        if let Some(item) = state.item.take() {
            drop(state);
            self.shared.sender_waker.wake();
            return Poll::Ready(Some(item));
        }
        if let Some(err) = state.pending_error.take() {
            // ...
            return Poll::Ready(Some(Err(err)));
        }
        let sender_open = state.sender_open;
        drop(state);
        if wake_sender {
            self.shared.sender_waker.wake();
        }
        if sender_open {
            Poll::Pending
        } else {
            self.terminated = true;
            Poll::Ready(None)
        }

            } else if let Some(mut body) = self.body_tx.take() {
                if self.conn.can_read_body() {
                    match body.poll_ready(cx) {
                        Poll::Ready(Ok(())) => (),
                        Poll::Pending => {
                            self.body_tx.set(body);
                            return Poll::Pending;
                        }
                        Poll::Ready(Err(_canceled)) => {
                            // user doesn't care about the body
                            // so we should stop reading
                            trace!("body receiver dropped before eof, draining or closing");
                            self.conn.poll_drain_or_close_read(cx);
                            continue;
                        }
                    }
                    match self.conn.poll_read_body(cx) {
                        Poll::Ready(Some(Ok(frame))) => {
                            if frame.is_data() {
                                let chunk = frame.into_data().unwrap_or_else(|_| unreachable!());
                                match body.try_send_data(chunk) {
                                    Ok(()) => {
                                        self.body_tx.set(body);
                                    }
                                    // ...
                                }
                            // ... trailers：try_send_trailers
                            }
                        }
                        Poll::Ready(None) => {
                            // just drop, the body will close automatically
                        }
                        Poll::Pending => {
                            self.body_tx.set(body);
                            return Poll::Pending;
                        }
                        Poll::Ready(Some(Err(e))) => {
                            body.send_error(crate::Error::new_body(e));
                        }
                    }
                } else {
                    // just drop, the body will close automatically
                }

    /// If the read side can be cheaply drained, do so. Otherwise, close.
    pub(super) fn poll_drain_or_close_read(&mut self, cx: &mut Context<'_>) {
        if let Reading::Continue(decoder) = &mut self.state.reading {
            // skip sending the 100-continue
            // just move forward to a read, in case a tiny body was included
            self.state.reading = Reading::Body(decoder.clone());
        }

        let _ = self.poll_read_body(cx);

        // If still in Reading::Body, just give up
        match self.state.reading {
            Reading::Init | Reading::KeepAlive => {
                trace!("body drained")
            }
            _ => self.close_read(),
        }
    }

hyper = { version = "=1.12.0", features = ["server", "http1"] }
hyper-util = { version = "=0.1.21", features = ["tokio"] }
http-body-util = "=0.1.5"
tokio = { version = "=1.54.1", features = ["full"] }

use std::{
  convert::Infallible,
  io,
  pin::Pin,
  sync::{
    Arc,
    atomic::{AtomicUsize, Ordering::SeqCst},
  },
  task::{Context, Poll},
  time::Duration,
};
use http_body_util::Full;
use hyper::{
  Request, Response,
  body::{Bytes, Incoming},
  server::conn::http1,
  service::service_fn,
};
use hyper_util::rt::TokioIo;
use tokio::{
  io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf},
  net::{TcpListener, TcpStream},
  time::timeout,
};

// Counts the bytes the server pulls out of the socket.
struct Counting {
  inner: TcpStream,
  read: Arc<AtomicUsize>,
}
impl AsyncRead for Counting {
  fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
    let before = buf.filled().len();
    let r = Pin::new(&mut self.inner).poll_read(cx, buf);
    self.read.fetch_add(buf.filled().len() - before, SeqCst);
    r
  }
}
impl AsyncWrite for Counting {
  fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, b: &[u8]) -> Poll<io::Result<usize>> {
    Pin::new(&mut self.inner).poll_write(cx, b)
  }
  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 handler(req: Request<Incoming>) -> Result<Response<Full<Bytes>>, Infallible> {
  if req.uri().path() == "/drop-then-wait" {
    drop(req); // body receiver dropped before the response exists
    tokio::time::sleep(Duration::from_millis(50)).await;
  }
  // "/ignore": the body is dropped together with `req` when this future completes
  Ok(Response::new(Full::new(Bytes::from_static(b"ok"))))
}

async fn read_response(rd: &mut tokio::net::tcp::OwnedReadHalf) -> Result<String, String> {
  let mut got = Vec::new();
  let mut buf = [0u8; 4096];
  loop {
    if let Some(i) = got.windows(4).position(|w| w == b"\r\n\r\n") {
      if got.len() >= i + 4 + 2 {
        let head = String::from_utf8_lossy(&got[..i]).to_lowercase();
        let close = if head.contains("connection: close") { ", connection: close" } else { "" };
        return Ok(format!("{}{close}", head.lines().next().unwrap().to_uppercase()));
      }
    }
    match timeout(Duration::from_millis(500), rd.read(&mut buf)).await {
      Ok(Ok(0)) => return Err("EOF".into()),
      Ok(Ok(n)) => got.extend_from_slice(&buf[..n]),
      Ok(Err(e)) => return Err(format!("{:?}", e.kind())),
      Err(_) => return Err("no response in 500ms".into()),
    }
  }
}

#[tokio::main]
async fn main() {
  let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move {
    loop {
      let (stream, _) = listener.accept().await.unwrap();
      let read = Arc::new(AtomicUsize::new(0));
      let io = Counting { inner: stream, read: read.clone() };
      tokio::spawn(async move {
        let res = http1::Builder::new().serve_connection(TokioIo::new(io), service_fn(handler)).await;
        println!("  server: connection ended {res:?}, read {} bytes in total", read.load(SeqCst));
      });
    }
  });

  for path in ["/ignore", "/drop-then-wait"] {
    for size in [10usize, 12 * 1024, 1024 * 1024] {
      println!("== POST {path}, body {size} bytes");
      let tcp = TcpStream::connect(addr).await.unwrap();
      let (mut rd, mut wr) = tcp.into_split();
      let mut req = format!("POST {path} HTTP/1.1\r\nHost: x\r\nContent-Length: {size}\r\n\r\n").into_bytes();
      req.extend(std::iter::repeat_n(b'a', size));
      let total = req.len();
      let writer = tokio::spawn(async move {
        let r = wr.write_all(&req).await;
        (wr, r)
      });
      match read_response(&mut rd).await {
        Ok(status) => println!("  client: {status}"),
        Err(e) => println!("  client: response failed: {e}"),
      }
      let (mut wr, wrote) = match timeout(Duration::from_secs(2), writer).await {
        Ok(j) => j.unwrap(),
        Err(_) => panic!("writer stuck"),
      };
      match wrote {
        Ok(()) => println!("  client: wrote all {total} bytes of the request"),
        Err(e) => println!("  client: writing the request failed: {:?}", e.kind()),
      }
      let _ = wr.write_all(b"GET /next HTTP/1.1\r\nHost: x\r\n\r\n").await;
      match read_response(&mut rd).await {
        Ok(status) => println!("  client: second request on same connection -> {status}"),
        Err(e) => println!("  client: second request on same connection -> {e}"),
      }
      drop((rd, wr));
      tokio::time::sleep(Duration::from_millis(100)).await;
    }
  }
}

== POST /ignore, body 10 bytes
  client: HTTP/1.1 200 OK
  client: wrote all 64 bytes of the request
  client: second request on same connection -> HTTP/1.1 200 OK
  server: connection ended Ok(()), read 95 bytes in total
== POST /ignore, body 12288 bytes
  client: HTTP/1.1 200 OK
  client: wrote all 12345 bytes of the request
  client: second request on same connection -> HTTP/1.1 200 OK
  server: connection ended Ok(()), read 12376 bytes in total
== POST /ignore, body 1048576 bytes
  server: connection ended Ok(()), read 24576 bytes in total
  client: HTTP/1.1 200 OK
  client: writing the request failed: BrokenPipe
  client: second request on same connection -> ConnectionReset
== POST /drop-then-wait, body 10 bytes
  client: HTTP/1.1 200 OK
  client: wrote all 72 bytes of the request
  client: second request on same connection -> HTTP/1.1 200 OK
  server: connection ended Ok(()), read 103 bytes in total
== POST /drop-then-wait, body 12288 bytes
  client: HTTP/1.1 200 OK
  client: wrote all 12353 bytes of the request
  client: second request on same connection -> HTTP/1.1 200 OK
  server: connection ended Ok(()), read 12384 bytes in total
== POST /drop-then-wait, body 1048576 bytes
  server: connection ended Ok(()), read 24576 bytes in total
  client: HTTP/1.1 200 OK, connection: close
  client: wrote all 1048643 bytes of the request
  client: second request on same connection -> ConnectionReset

client: wrote 262196 bytes
before first frame: socket read   8192, size_hint Some(262144)
frame   8140 B: socket read before   8192, after   8192; consumed   8140, size_hint Some(254004)
frame  16384 B: socket read before  24576, after  24576; consumed  24524, size_hint Some(237620)
frame  32768 B: socket read before  57344, after  57344; consumed  57292, size_hint Some(204852)
frame  65536 B: socket read before 122880, after 122880; consumed 122828, size_hint Some(139316)
frame 131072 B: socket read before 253952, after 253952; consumed 253900, size_hint Some(8244)
frame   8244 B: socket read before 262196, after 262196; consumed 262144, size_hint Some(0)
client: got "HTTP/1.1 200 OK"