From 91d3b4c0bccf2234fc3ed19e605e2cd402f19437 Mon Sep 17 00:00:00 2001 From: Alexander Kireyev Date: Mon, 20 Jul 2026 17:32:46 +0700 Subject: [PATCH] time: fix alt timer cancellation and insertion race (#8252) Co-authored-by: ADD-SP --- tokio/src/runtime/time_alt/entry.rs | 6 +- tokio/src/runtime/time_alt/tests.rs | 101 +++++++++++++++++++++++- tokio/src/runtime/time_alt/wheel/mod.rs | 7 +- 3 files changed, 108 insertions(+), 6 deletions(-) diff --git a/tokio/src/runtime/time_alt/entry.rs b/tokio/src/runtime/time_alt/entry.rs index 9ebadd26d..2b204a130 100644 --- a/tokio/src/runtime/time_alt/entry.rs +++ b/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 } } diff --git a/tokio/src/runtime/time_alt/tests.rs b/tokio/src/runtime/time_alt/tests.rs index e1a7196fa..c713a60bc 100644 --- a/tokio/src/runtime/time_alt/tests.rs +++ b/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 { diff --git a/tokio/src/runtime/time_alt/wheel/mod.rs b/tokio/src/runtime/time_alt/wheel/mod.rs index e234f71cb..071ccbac6 100644 --- a/tokio/src/runtime/time_alt/wheel/mod.rs +++ b/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);