fix(http1): synchronize dispatch shutdown with sends

src/client/dispatch.rs

@@ -1,3 +1,5 @@
+#[cfg(feature = "http1")]
+use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
 #[cfg(feature = "http2")]
 use std::{future::Future, pin::Pin};
@@ -10,6 +12,8 @@ use http_body::Body;
 use pin_project_lite::pin_project;
 use tokio::sync::{mpsc, oneshot};
 
+#[cfg(feature = "http1")]
+use crate::common::lock::LockResultExt;
 #[cfg(feature = "http2")]
 use crate::{body::Incoming, proto::h2::client::ResponseFutMap};
 
@@ -31,13 +35,22 @@ pub struct TrySendError<T> {
 pub(crate) fn channel<T, U>() -> (Sender<T, U>, Receiver<T, U>) {
     let (tx, rx) = mpsc::unbounded_channel();
     let (giver, taker) = want::new();
+    #[cfg(feature = "http1")]
+    let send_lock = Arc::new(Mutex::new(()));
     let tx = Sender {
         #[cfg(feature = "http1")]
         buffered_once: false,
+        #[cfg(feature = "http1")]
+        send_lock: send_lock.clone(),
         giver,
         inner: tx,
     };
-    let rx = Receiver { inner: rx, taker };
+    let rx = Receiver {
+        inner: rx,
+        taker,
+        #[cfg(feature = "http1")]
+        send_lock,
+    };
     (tx, rx)
 }
 
@@ -51,6 +64,10 @@ pub(crate) struct Sender<T, U> {
     /// without notice.
     #[cfg(feature = "http1")]
     buffered_once: bool,
+    /// Synchronizes HTTP/1 sends with receiver shutdown so an envelope cannot
+    /// be published after the receiver has finished draining the channel.
+    #[cfg(feature = "http1")]
+    send_lock: Arc<Mutex<()>>,
     /// The Giver helps watch that the Receiver side has been polled
     /// when the queue is empty. This helps us know when a request and
     /// response have been fully processed, and a connection is ready
@@ -109,6 +126,7 @@ impl<T, U> Sender<T, U> {
             return Err(val);
         }
         let (tx, rx) = oneshot::channel();
+        let _guard = self.send_lock.lock().panic_if_poisoned();
         self.inner
             .send(Envelope(Some((val, Callback::Retry(Some(tx))))))
             .map(move |_| rx)
@@ -121,6 +139,7 @@ impl<T, U> Sender<T, U> {
             return Err(val);
         }
         let (tx, rx) = oneshot::channel();
+        let _guard = self.send_lock.lock().panic_if_poisoned();
         self.inner
             .send(Envelope(Some((val, Callback::NoRetry(Some(tx))))))
             .map(move |_| rx)
@@ -176,6 +195,8 @@ impl<T, U> Clone for UnboundedSender<T, U> {
 pub(crate) struct Receiver<T, U> {
     inner: mpsc::UnboundedReceiver<Envelope<T, U>>,
     taker: want::Taker,
+    #[cfg(feature = "http1")]
+    send_lock: Arc<Mutex<()>>,
 }
 
 impl<T, U> Receiver<T, U> {
@@ -193,6 +214,7 @@ impl<T, U> Receiver<T, U> {
 
     #[cfg(feature = "http1")]
     pub(crate) fn close(&mut self) {
+        let _guard = self.send_lock.lock().panic_if_poisoned();
         self.taker.cancel();
         self.inner.close();
     }
@@ -208,9 +230,13 @@ impl<T, U> Receiver<T, U> {
 
 impl<T, U> Drop for Receiver<T, U> {
     fn drop(&mut self) {
+        #[cfg(feature = "http1")]
+        let _guard = self.send_lock.lock().panic_if_poisoned();
         // Notify the giver about the closure first, before dropping
         // the mpsc::Receiver.
         self.taker.cancel();
+        #[cfg(feature = "http1")]
+        self.inner.close();
     }
 }
 
@@ -394,11 +420,65 @@ mod tests {
     use std::pin::Pin;
     use std::task::{Context, Poll};
 
+    #[cfg(feature = "http1")]
+    use std::sync::atomic::{AtomicBool, Ordering};
+    #[cfg(feature = "http1")]
+    use std::sync::{mpsc as std_mpsc, Arc, Barrier};
+
     use super::{channel, Callback, Receiver};
 
     #[derive(Debug)]
     struct Custom(#[allow(dead_code)] i32);
 
+    #[cfg(feature = "http1")]
+    struct TrackDrop(Arc<AtomicBool>);
+
+    #[cfg(feature = "http1")]
+    impl Drop for TrackDrop {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+
+    #[cfg(feature = "http1")]
+    #[test]
+    fn receiver_shutdown_reclaims_concurrent_send() {
+        let barrier = Arc::new(Barrier::new(2));
+        let worker_barrier = barrier.clone();
+        let (work_tx, work_rx) = std_mpsc::channel::<(super::Sender<TrackDrop, ()>, TrackDrop)>();
+        let (done_tx, done_rx) = std_mpsc::channel();
+        let worker = std::thread::spawn(move || {
+            while let Ok((mut tx, val)) = work_rx.recv() {
+                worker_barrier.wait();
+                drop(tx.send(val));
+                worker_barrier.wait();
+                done_tx.send(tx).unwrap();
+            }
+        });
+
+        // Tokio's unbounded channel reserves capacity before publishing the
+        // envelope. Repeat enough times to exercise shutdown in that window.
+        for _ in 0..10_000 {
+            let (tx, mut rx) = channel::<TrackDrop, ()>();
+            let dropped = Arc::new(AtomicBool::new(false));
+            work_tx.send((tx, TrackDrop(dropped.clone()))).unwrap();
+            barrier.wait();
+            rx.close();
+            drop(rx.try_recv());
+            drop(rx);
+            barrier.wait();
+            let tx = done_rx.recv().unwrap();
+            assert!(
+                dropped.load(Ordering::SeqCst),
+                "value remained queued after the receiver was dropped"
+            );
+            drop(tx);
+        }
+
+        drop(work_tx);
+        worker.join().unwrap();
+    }
+
     impl<T, U> Future for Receiver<T, U> {
         type Output = Option<(T, Callback<T, U>)>;