sync: make AtomicWaker panic safe (#3689)

tokio/src/sync/task/atomic_waker.rs

@@ -4,6 +4,7 @@ use crate::loom::cell::UnsafeCell;
 use crate::loom::sync::atomic::{self, AtomicUsize};
 
 use std::fmt;
+use std::panic::{resume_unwind, AssertUnwindSafe, RefUnwindSafe, UnwindSafe};
 use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
 use std::task::Waker;
 
@@ -27,6 +28,9 @@ pub(crate) struct AtomicWaker {
     waker: UnsafeCell<Option<Waker>>,
 }
 
+impl RefUnwindSafe for AtomicWaker {}
+impl UnwindSafe for AtomicWaker {}
+
 // `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell
 // stores a `Waker` value produced by calls to `register` and many threads can
 // race to take the waker by calling `wake`.
@@ -84,7 +88,7 @@ pub(crate) struct AtomicWaker {
 // back to `WAITING`. This transition must succeed as, at this point, the state
 // cannot be transitioned by another thread.
 //
-// If the thread is unable to obtain the lock, the `WAKING` bit is still.
+// If the thread is unable to obtain the lock, the `WAKING` bit is still set.
 // This is because it has either been set by the current thread but the previous
 // value included the `REGISTERING` bit **or** a concurrent thread is in the
 // `WAKING` critical section. Either way, no action must be taken.
@@ -171,6 +175,10 @@ impl AtomicWaker {
     where
         W: WakerRef,
     {
+        fn catch_unwind<F: FnOnce() -> R, R>(f: F) -> std::thread::Result<R> {
+            std::panic::catch_unwind(AssertUnwindSafe(f))
+        }
+
         match self
             .state
             .compare_exchange(WAITING, REGISTERING, Acquire, Acquire)
@@ -178,8 +186,24 @@ impl AtomicWaker {
         {
             WAITING => {
                 unsafe {
-                    // Locked acquired, update the waker cell
-                    self.waker.with_mut(|t| *t = Some(waker.into_waker()));
+                    // If `into_waker` panics (because it's code outside of
+                    // AtomicWaker) we need to prime a guard that is called on
+                    // unwind to restore the waker to a WAITING state. Otherwise
+                    // any future calls to register will incorrectly be stuck
+                    // believing it's being updated by someone else.
+                    let new_waker_or_panic = catch_unwind(move || waker.into_waker());
+
+                    // Set the field to contain the new waker, or if
+                    // `into_waker` panicked, leave the old value.
+                    let mut maybe_panic = None;
+                    let mut old_waker = None;
+                    match new_waker_or_panic {
+                        Ok(new_waker) => {
+                            old_waker = self.waker.with_mut(|t| (*t).take());
+                            self.waker.with_mut(|t| *t = Some(new_waker));
+                        }
+                        Err(panic) => maybe_panic = Some(panic),
+                    }
 
                     // Release the lock. If the state transitioned to include
                     // the `WAKING` bit, this means that a wake has been
@@ -193,33 +217,67 @@ impl AtomicWaker {
                         .compare_exchange(REGISTERING, WAITING, AcqRel, Acquire);
 
                     match res {
-                        Ok(_) => {}
+                        Ok(_) => {
+                            // We don't want to give the caller the panic if it
+                            // was someone else who put in that waker.
+                            let _ = catch_unwind(move || {
+                                drop(old_waker);
+                            });
+                        }
                         Err(actual) => {
                             // This branch can only be reached if a
                             // concurrent thread called `wake`. In this
                             // case, `actual` **must** be `REGISTERING |
-                            // `WAKING`.
+                            // WAKING`.
                             debug_assert_eq!(actual, REGISTERING | WAKING);
 
                             // Take the waker to wake once the atomic operation has
                             // completed.
-                            let waker = self.waker.with_mut(|t| (*t).take()).unwrap();
+                            let mut waker = self.waker.with_mut(|t| (*t).take());
 
                             // Just swap, because no one could change state
                             // while state == `Registering | `Waking`
                             self.state.swap(WAITING, AcqRel);
 
-                            // The atomic swap was complete, now
-                            // wake the waker and return.
-                            waker.wake();
+                            // If `into_waker` panicked, then the waker in the
+                            // waker slot is actually the old waker.
+                            if maybe_panic.is_some() {
+                                old_waker = waker.take();
+                            }
+
+                            // We don't want to give the caller the panic if it
+                            // was someone else who put in that waker.
+                            if let Some(old_waker) = old_waker {
+                                let _ = catch_unwind(move || {
+                                    old_waker.wake();
+                                });
+                            }
+
+                            // The atomic swap was complete, now wake the waker
+                            // and return.
+                            //
+                            // If this panics, we end up in a consumed state and
+                            // return the panic to the caller.
+                            if let Some(waker) = waker {
+                                debug_assert!(maybe_panic.is_none());
+                                waker.wake();
+                            }
                         }
                     }
+
+                    if let Some(panic) = maybe_panic {
+                        // If `into_waker` panicked, return the panic to the caller.
+                        resume_unwind(panic);
+                    }
                 }
             }
             WAKING => {
                 // Currently in the process of waking the task, i.e.,
                 // `wake` is currently being called on the old waker.
                 // So, we call wake on the new waker.
+                //
+                // If this panics, someone else is responsible for restoring the
+                // state of the waker.
                 waker.wake();
 
                 // This is equivalent to a spin lock, so use a spin hint.
@@ -245,6 +303,8 @@ impl AtomicWaker {
     /// If `register` has not been called yet, then this does nothing.
     pub(crate) fn wake(&self) {
         if let Some(waker) = self.take_waker() {
+            // If wake panics, we've consumed the waker which is a legitimate
+            // outcome.
             waker.wake();
         }
     }

tokio/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());
+}

tokio/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
+        }));
+    });
+}