5 files changed, 168 insertions, 0 deletions
diff --git a/src/sync/tests/atomic_waker.rs b/src/sync/tests/atomic_waker.rs
index c832d62..b167a5d 100644
--- a/src/sync/tests/atomic_waker.rs
+++ b/src/sync/tests/atomic_waker.rs
@@ -32,3 +32,42 @@ fn wake_without_register() {
 
     assert!(!waker.is_woken());
 }
+
+#[test]
+fn atomic_waker_panic_safe() {
+    use std::panic;
+    use std::ptr;
+    use std::task::{RawWaker, RawWakerVTable, Waker};
+
+    static PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new(
+        |_| panic!("clone"),
+        |_| unimplemented!("wake"),
+        |_| unimplemented!("wake_by_ref"),
+        |_| (),
+    );
+
+    static NONPANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new(
+        |_| RawWaker::new(ptr::null(), &NONPANICKING_VTABLE),
+        |_| unimplemented!("wake"),
+        |_| unimplemented!("wake_by_ref"),
+        |_| (),
+    );
+
+    let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) };
+    let nonpanicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &NONPANICKING_VTABLE)) };
+
+    let atomic_waker = AtomicWaker::new();
+
+    let panicking = panic::AssertUnwindSafe(&panicking);
+
+    let result = panic::catch_unwind(|| {
+        let panic::AssertUnwindSafe(panicking) = panicking;
+        atomic_waker.register_by_ref(panicking);
+    });
+
+    assert!(result.is_err());
+    assert!(atomic_waker.take_waker().is_none());
+
+    atomic_waker.register_by_ref(&nonpanicking);
+    assert!(atomic_waker.take_waker().is_some());
+}
diff --git a/src/sync/tests/loom_atomic_waker.rs b/src/sync/tests/loom_atomic_waker.rs
index c148bcb..f8bae65 100644
--- a/src/sync/tests/loom_atomic_waker.rs
+++ b/src/sync/tests/loom_atomic_waker.rs
@@ -43,3 +43,58 @@ fn basic_notification() {
         }));
     });
 }
+
+#[test]
+fn test_panicky_waker() {
+    use std::panic;
+    use std::ptr;
+    use std::task::{RawWaker, RawWakerVTable, Waker};
+
+    static PANICKING_VTABLE: RawWakerVTable =
+        RawWakerVTable::new(|_| panic!("clone"), |_| (), |_| (), |_| ());
+
+    let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) };
+
+    // If you're working with this test (and I sure hope you never have to!),
+    // uncomment the following section because there will be a lot of panics
+    // which would otherwise log.
+    //
+    // We can't however leaved it uncommented, because it's global.
+    // panic::set_hook(Box::new(|_| ()));
+
+    const NUM_NOTIFY: usize = 2;
+
+    loom::model(move || {
+        let chan = Arc::new(Chan {
+            num: AtomicUsize::new(0),
+            task: AtomicWaker::new(),
+        });
+
+        for _ in 0..NUM_NOTIFY {
+            let chan = chan.clone();
+
+            thread::spawn(move || {
+                chan.num.fetch_add(1, Relaxed);
+                chan.task.wake();
+            });
+        }
+
+        // Note: this panic should have no effect on the overall state of the
+        // waker and it should proceed as normal.
+        //
+        // A thread above might race to flag a wakeup, and a WAKING state will
+        // be preserved if this expected panic races with that so the below
+        // procedure should be allowed to continue uninterrupted.
+        let _ = panic::catch_unwind(|| chan.task.register_by_ref(&panicking));
+
+        block_on(poll_fn(move |cx| {
+            chan.task.register_by_ref(cx.waker());
+
+            if NUM_NOTIFY == chan.num.load(Relaxed) {
+                return Ready(());
+            }
+
+            Pending
+        }));
+    });
+}
diff --git a/src/sync/tests/loom_oneshot.rs b/src/sync/tests/loom_oneshot.rs
index 9729cfb..c5f7972 100644
--- a/src/sync/tests/loom_oneshot.rs
+++ b/src/sync/tests/loom_oneshot.rs
@@ -55,6 +55,35 @@ fn changing_rx_task() {
     });
 }
 
+#[test]
+fn try_recv_close() {
+    // reproduces https://github.com/tokio-rs/tokio/issues/4225
+    loom::model(|| {
+        let (tx, mut rx) = oneshot::channel();
+        thread::spawn(move || {
+            let _ = tx.send(());
+        });
+
+        rx.close();
+        let _ = rx.try_recv();
+    })
+}
+
+#[test]
+fn recv_closed() {
+    // reproduces https://github.com/tokio-rs/tokio/issues/4225
+    loom::model(|| {
+        let (tx, mut rx) = oneshot::channel();
+
+        thread::spawn(move || {
+            let _ = tx.send(1);
+        });
+
+        rx.close();
+        let _ = block_on(rx);
+    });
+}
+
 // TODO: Move this into `oneshot` proper.
 
 use std::future::Future;
diff --git a/src/sync/tests/mod.rs b/src/sync/tests/mod.rs
index c5d5601..ee76418 100644
--- a/src/sync/tests/mod.rs
+++ b/src/sync/tests/mod.rs
@@ -1,5 +1,6 @@
 cfg_not_loom! {
     mod atomic_waker;
+    mod notify;
     mod semaphore_batch;
 }
 
diff --git a/src/sync/tests/notify.rs b/src/sync/tests/notify.rs
new file mode 100644
index 0000000..8c9a573
--- /dev/null
+++ b/src/sync/tests/notify.rs
@@ -0,0 +1,44 @@
+use crate::sync::Notify;
+use std::future::Future;
+use std::mem::ManuallyDrop;
+use std::sync::Arc;
+use std::task::{Context, RawWaker, RawWakerVTable, Waker};
+
+#[test]
+fn notify_clones_waker_before_lock() {
+    const VTABLE: &RawWakerVTable = &RawWakerVTable::new(clone_w, wake, wake_by_ref, drop_w);
+
+    unsafe fn clone_w(data: *const ()) -> RawWaker {
+        let arc = ManuallyDrop::new(Arc::<Notify>::from_raw(data as *const Notify));
+        // Or some other arbitrary code that shouldn't be executed while the
+        // Notify wait list is locked.
+        arc.notify_one();
+        let _arc_clone: ManuallyDrop<_> = arc.clone();
+        RawWaker::new(data, VTABLE)
+    }
+
+    unsafe fn drop_w(data: *const ()) {
+        let _ = Arc::<Notify>::from_raw(data as *const Notify);
+    }
+
+    unsafe fn wake(_data: *const ()) {
+        unreachable!()
+    }
+
+    unsafe fn wake_by_ref(_data: *const ()) {
+        unreachable!()
+    }
+
+    let notify = Arc::new(Notify::new());
+    let notify2 = notify.clone();
+
+    let waker =
+        unsafe { Waker::from_raw(RawWaker::new(Arc::into_raw(notify2) as *const _, VTABLE)) };
+    let mut cx = Context::from_waker(&waker);
+
+    let future = notify.notified();
+    pin!(future);
+
+    // The result doesn't matter, we're just testing that we don't deadlock.
+    let _ = future.poll(&mut cx);
+}