From 27148d6110aa36de2b225e2a9f10ccd6e743cc0f Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Wed, 13 Mar 2019 10:38:14 -0700
Subject: [PATCH] sync: free chan Blocks when Chan is dropped (#978)

---
 tokio-sync/src/mpsc/chan.rs |  3 +++
 tokio-sync/src/mpsc/list.rs | 37 +++++++++++++++++++++++++++++++------
 tokio-sync/tests/mpsc.rs    |  2 +-
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/tokio-sync/src/mpsc/chan.rs b/tokio-sync/src/mpsc/chan.rs
index 05a5634de..894da468f 100644
--- a/tokio-sync/src/mpsc/chan.rs
+++ b/tokio-sync/src/mpsc/chan.rs
@@ -326,10 +326,13 @@ impl<T, S> Drop for Chan<T, S> {
     fn drop(&mut self) {
         use super::block::Read::Value;
 
+        // Safety: the only owner of the rx fields is Chan, and eing
+        // inside its own Drop means we're the last ones to touch it.
         self.rx_fields.with_mut(|rx_fields_ptr| {
             let rx_fields = unsafe { &mut *rx_fields_ptr };
 
             while let Some(Value(_)) = rx_fields.list.pop(&self.tx) {}
+            unsafe { rx_fields.list.free_blocks() };
         });
     }
 }
diff --git a/tokio-sync/src/mpsc/list.rs b/tokio-sync/src/mpsc/list.rs
index 165f1adc5..d7441c9a3 100644
--- a/tokio-sync/src/mpsc/list.rs
+++ b/tokio-sync/src/mpsc/list.rs
@@ -58,7 +58,7 @@ impl<T> Tx<T> {
     /// Push a value into the list.
     pub(crate) fn push(&self, value: T) {
         // First, claim a slot for the value. `Acquire` is used here to
-        // synchronize with the `fetch_add` in `free_blocks`.
+        // synchronize with the `fetch_add` in `reclaim_blocks`.
         let slot_index = self.tail_position.fetch_add(1, Acquire);
 
         // Load the current block and write the value
@@ -170,6 +170,7 @@ impl<T> Tx<T> {
     }
 
     pub(crate) unsafe fn reclaim_block(&self, mut block: NonNull<Block<T>>) {
+        debug!("+ reclaim_block({:p})", block);
         // The block has been removed from the linked list and ownership
         // is reclaimed.
         //
@@ -206,6 +207,7 @@ impl<T> Tx<T> {
         }
 
         if !reused {
+            debug!(" + block freed {:p}", block);
             let _ = Box::from_raw(block.as_ptr());
         }
     }
@@ -231,7 +233,7 @@ impl<T> Rx<T> {
             return None;
         }
 
-        self.free_blocks(tx);
+        self.reclaim_blocks(tx);
 
         unsafe {
             let block = self.head.as_ref();
@@ -276,8 +278,8 @@ impl<T> Rx<T> {
         }
     }
 
-    fn free_blocks(&mut self, tx: &Tx<T>) {
-        debug!("+ free_blocks()");
+    fn reclaim_blocks(&mut self, tx: &Tx<T>) {
+        debug!("+ reclaim_blocks()");
 
         while self.free_head != self.head {
             unsafe {
@@ -297,8 +299,8 @@ impl<T> Rx<T> {
                 }
 
                 // We may read the next pointer with `Relaxed` ordering as it is
-                // guaranteed that the `free_blocks` routine trails the `recv`
-                // routine. Any memory accessed by `free_blocks` has already
+                // guaranteed that the `reclaim_blocks` routine trails the `recv`
+                // routine. Any memory accessed by `reclaim_blocks` has already
                 // been acquired by `recv`.
                 let next_block = block.as_ref().load_next(Relaxed);
 
@@ -313,6 +315,29 @@ impl<T> Rx<T> {
             loom::yield_now();
         }
     }
+
+    /// Effectively `Drop` all the blocks. Should only be called once, when
+    /// the list is dropping.
+    pub(super) unsafe fn free_blocks(&mut self) {
+        debug!("+ free_blocks()");
+        debug_assert_ne!(self.free_head, NonNull::dangling());
+
+        let mut cur = Some(self.free_head);
+
+        #[cfg(debug_assertions)]
+        {
+            // to trigger the debug assert above so as to catch that we
+            // don't call `free_blocks` more than once.
+            self.free_head = NonNull::dangling();
+            self.head = NonNull::dangling();
+        }
+
+        while let Some(block) = cur {
+            cur = block.as_ref().load_next(Relaxed);
+            debug!(" + free: block = {:p}", block);
+            drop(Box::from_raw(block.as_ptr()));
+        }
+    }
 }
 
 impl<T> fmt::Debug for Rx<T> {
diff --git a/tokio-sync/tests/mpsc.rs b/tokio-sync/tests/mpsc.rs
index 69386ef61..3572b94b9 100644
--- a/tokio-sync/tests/mpsc.rs
+++ b/tokio-sync/tests/mpsc.rs
@@ -356,7 +356,7 @@ fn dropping_rx_closes_channel_for_try() {
 }
 
 #[test]
-fn unconsumed_messagers_are_dropped() {
+fn unconsumed_messages_are_dropped() {
     let msg = Arc::new(());
 
     let (mut tx, rx) = mpsc::channel(100);