Body dropped too late when an h1 handshake hits an error #4122

use bytes::Bytes;
use http_body::Body;
use hyper::rt::{Read, ReadBufCursor, Write};
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier};
use std::task::{Context, Poll, Waker};

fn main() {
    // start up some background work to stress scheduling a bit more.
    for _ in 0..16 {
        std::thread::spawn(|| loop {
            std::hint::black_box(std::time::Instant::now());
        });
    }

    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let _guard = rt.enter();

    for _ in 0..2_000_000 {
        // Handshake against the mock IO. handshake() itself does no I/O.
        let (mut sender, mut conn) = match rt
            .block_on(hyper::client::conn::http1::handshake::<_, TrackedBody>(
                MockIo::default(),
            )) {
            Ok(p) => p,
            Err(_) => continue,
        };

        let dropped = Arc::new(AtomicBool::new(false));
        let req = hyper::Request::connect("http://x/")
            .body(TrackedBody {
                dropped: dropped.clone(),
            })
            .unwrap();

        let barrier = Arc::new(Barrier::new(2));

        // Thread B: dispatch the request the instant the barrier releases.
        let b2 = barrier.clone();
        let sender_thread = std::thread::spawn(move || {
            b2.wait();
            let fut = sender.send_request(req);
            (fut, sender)
        });

        barrier.wait();
        let waker = Waker::noop();
        let mut cx = Context::from_waker(&waker);
        for _ in 0..4 {
            if Pin::new(&mut conn).poll(&mut cx).is_ready() {
                break;
            }
        }
        drop(conn);

        let (fut, sender2) = sender_thread.join().unwrap();
        drop(fut);

        // At this point the future from `send_request` is gone, the receiver
        // through `conn` is gone, so it should be the case that our body in the
        // request that was sent was dropped. However this isn't always true so
        // sometimes this will print out.
        if !dropped.load(Ordering::SeqCst) {
            println!("body wasn't dropped");
        }
        drop(sender2);
        assert!(
            dropped.load(Ordering::SeqCst),
            "request must be reclaimed once the SendRequest drops"
        );
    }
}

#[derive(Default)]
struct MockIo {
    gave: bool,
}

impl Read for MockIo {
    fn poll_read(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        mut buf: ReadBufCursor<'_>,
    ) -> Poll<io::Result<()>> {
        // read enough bytes that the connection should terminate with an error.
        if !self.gave {
            self.gave = true;
            buf.put_slice(&[0u8; 9]);
        }
        Poll::Ready(Ok(()))
    }
}

impl Write for MockIo {
    fn poll_write(
        self: Pin<&mut Self>,
        _: &mut Context<'_>,
        _buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        // shouldn't be called for h1
        unreachable!()
    }
    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
}

/// An empty body that sets `dropped` to true when dropped.
struct TrackedBody {
    dropped: Arc<AtomicBool>,
}

impl Body for TrackedBody {
    type Data = Bytes;
    type Error = std::convert::Infallible;
    fn poll_frame(
        self: Pin<&mut Self>,
        _: &mut Context<'_>,
    ) -> Poll<Option<Result<http_body::Frame<Bytes>, Self::Error>>> {
        Poll::Ready(None)
    }
}

impl Drop for TrackedBody {
    fn drop(&mut self) {
        self.dropped.store(true, Ordering::SeqCst);
    }
}