time: fix alt timer cancellation and insertion race (#8252)

Co-authored-by: ADD-SP <[email protected]>

tokio/src/runtime/time_alt/entry.rs

@@ -210,12 +210,16 @@ impl Handle {
         }
     }
 
-    pub(crate) fn register_cancel_tx(&self, cancel_tx: Sender) {
+    /// Returns `false` if the `self` has already been cancelled or woken up.
+    pub(crate) fn register_cancel_tx(&self, cancel_tx: Sender) -> bool {
         let mut lock = self.entry.state.lock();
         if !lock.cancelled && !lock.woken_up {
             let old_tx = lock.cancel_tx.replace(cancel_tx);
             // don't unlock — poisoning the `Mutex` stops others from using the bad state.
             assert!(old_tx.is_none(), "cancel_tx is already registered");
+            true
+        } else {
+            false
         }
     }
 

tokio/src/runtime/time_alt/tests.rs

@@ -12,8 +12,12 @@ const NUM_ITEMS: usize = 16;
 const NUM_ITEMS: usize = 64;
 
 fn new_handle() -> (EntryHandle, AwokenCount) {
+    new_handle_with_deadline(0)
+}
+
+fn new_handle_with_deadline(deadline: u64) -> (EntryHandle, AwokenCount) {
     let (waker, count) = new_count_waker();
-    let entry = EntryHandle::new(0);
+    let entry = EntryHandle::new(deadline);
     _ = entry.poll(&mut Context::from_waker(&waker));
     (entry, count)
 }
@@ -66,7 +70,7 @@ fn cancel_in_the_same_thread() {
 
         for _ in 0..NUM_ITEMS {
             let (hdl, count) = new_handle();
-            hdl.register_cancel_tx(cancel_tx.clone());
+            assert!(hdl.register_cancel_tx(cancel_tx.clone()));
             counts.push(count);
             unsafe {
                 reg_queue.push_front(hdl.clone());
@@ -91,6 +95,97 @@ fn cancel_in_the_same_thread() {
     });
 }
 
+#[test]
+fn insert_of_already_cancelled_entry_does_not_enter_wheel() {
+    let (cancel_tx, mut cancel_rx) = cancellation_queue::new();
+    let mut wheel = Wheel::new();
+
+    // do not expire during the test
+    let far_future = 10_000_000;
+    let (hdl, awoken_count) = new_handle_with_deadline(far_future);
+
+    // cancel the timer before inserting it into the wheel
+    hdl.cancel();
+    assert!(hdl.is_cancelled());
+
+    // try to insert the cancelled entry into the wheel
+    unsafe {
+        wheel.insert(hdl, cancel_tx);
+    }
+
+    // a cancelled entry should not be inserted into the wheel
+    assert!(
+        wheel.next_expiration_time().is_none(),
+        "an already-cancelled entry leaked into the wheel on insert"
+    );
+
+    // It also must not have been queued for cancellation removal, since
+    // it was never actually placed in the wheel for `remove` to find.
+    assert_eq!(cancel_rx.recv_all().count(), 0);
+    assert_eq!(awoken_count.get(), 0);
+
+    // drain the wheel unconditionally, otherwise loom will complain
+    // about the leaked entry, which confuses developers in case
+    // this test fails for some other reason.
+    let mut wake_queue = WakeQueue::new();
+    wheel.take_expired(u64::MAX, &mut wake_queue);
+    wake_queue.wake_all();
+}
+
+#[test]
+fn cancel_races_with_insert() {
+    model(|| {
+        let (cancel_tx, mut cancel_rx) = cancellation_queue::new();
+        let mut wheel = Wheel::new();
+
+        // do not expire during the test
+        let far_future = 10_000_000;
+        let (hdl, count) = new_handle_with_deadline(far_future);
+
+        let hdl2 = hdl.clone();
+        let jh = thread::spawn(move || {
+            // cancel the timer concurrently with insertion into the wheel
+            hdl2.cancel();
+        });
+
+        // try to insert the entry into the wheel concurrently with cancellation
+        unsafe {
+            wheel.insert(hdl, cancel_tx);
+        }
+
+        // ensure the cancellation thread has exited
+        jh.join().unwrap();
+
+        // Whichever way the race went, the entry must end up in exactly one
+        // of two consistent end states -- never "in the wheel with no way
+        // to ever remove it again":
+        //
+        // (a) `cancel()` won the race and ran first: `insert` sees it's
+        //     already cancelled and refuses to add it to the wheel.
+        // (b) `insert` won the race and registered `cancel_tx` first:
+        //     `cancel()` then finds `cancel_tx` and pushes the entry into
+        //     the cancellation queue for later removal from the wheel.
+        //
+        // Either way, the entry is not left stuck in the wheel with no
+        // corresponding entry in the cancellation queue.
+        let cancelled_via_queue = cancel_rx.recv_all().count();
+        let leaked_in_wheel = wheel.next_expiration_time().is_some();
+
+        assert!(
+            !leaked_in_wheel || cancelled_via_queue == 1,
+            "entry is stuck in the wheel with no way to ever be removed"
+        );
+        assert_eq!(count.get(), 0, "a cancelled entry must never be woken");
+
+        // drain the wheel unconditionally, otherwise loom will complain
+        // about the leaked entry, which confuses developers in case
+        // this test fails for some other reason.
+        let mut wake_queue = WakeQueue::new();
+        wheel.take_expired(u64::MAX, &mut wake_queue);
+        wake_queue.wake_all();
+    });
+}
+
 #[test]
 fn wake_up_in_the_different_thread() {
     model(|| {
@@ -137,7 +232,7 @@ fn cancel_in_the_different_thread() {
 
         for _ in 0..NUM_ITEMS {
             let (hdl, count) = new_handle();
-            hdl.register_cancel_tx(cancel_tx.clone());
+            assert!(hdl.register_cancel_tx(cancel_tx.clone()));
             counts.push(count);
             hdls.push(hdl.clone());
             unsafe {

tokio/src/runtime/time_alt/wheel/mod.rs

@@ -54,7 +54,7 @@ impl Wheel {
         self.elapsed
     }
 
-    /// Inserts an entry into the timing wheel.
+    /// Inserts an entry into the timing wheel
     ///
     /// # Arguments
     ///
@@ -70,7 +70,10 @@ impl Wheel {
 
         assert!(deadline > self.elapsed);
 
-        hdl.register_cancel_tx(cancel_tx);
+        if !hdl.register_cancel_tx(cancel_tx) {
+            // `hdl` has been cancelled or woken up concurrently
+            return;
+        }
 
         // Get the level at which the entry should be stored
         let level = self.level_for(deadline);