time: Fix race condition in timer drop (#3229)

Dropping a timer on the millisecond that it was scheduled for, when it was on
the pending list, could result in a panic previously, as we did not record the
pending-list state in cached_when.

Hopefully fixes: ZcashFoundation/zebra#1452

2 files changed, 19 insertions, 9 deletions

diff --git a/tokio/src/time/driver/entry.rs b/tokio/src/time/driver/entry.rs
index e0926797..87ba0c17 100644
--- a/tokio/src/time/driver/entry.rs
+++ b/tokio/src/time/driver/entry.rs
@@ -437,6 +437,17 @@ impl TimerShared {
         true_when
     }
 
+    /// Sets the cached time-of-expiration value.
+    ///
+    /// SAFETY: Must be called with the driver lock held, and when this entry is
+    /// not in any timer wheel lists.
+    unsafe fn set_cached_when(&self, when: u64) {
+        self.driver_state
+            .0
+            .cached_when
+            .store(when, Ordering::Relaxed);
+    }
+
     /// Returns the true time-of-expiration value, with relaxed memory ordering.
     pub(super) fn true_when(&self) -> u64 {
         self.state.when().expect("Timer already fired")
@@ -643,14 +654,13 @@ impl TimerHandle {
     /// After returning Ok, the entry must be added to the pending list.
     pub(super) unsafe fn mark_pending(&self, not_after: u64) -> Result<(), u64> {
         match self.inner.as_ref().state.mark_pending(not_after) {
-            Ok(()) => Ok(()),
+            Ok(()) => {
+                // mark this as being on the pending queue in cached_when
+                self.inner.as_ref().set_cached_when(u64::max_value());
+                Ok(())
+            }
             Err(tick) => {
-                self.inner
-                    .as_ref()
-                    .driver_state
-                    .0
-                    .cached_when
-                    .store(tick, Ordering::Relaxed);
+                self.inner.as_ref().set_cached_when(tick);
                 Err(tick)
             }
         }
diff --git a/tokio/src/time/driver/wheel/mod.rs b/tokio/src/time/driver/wheel/mod.rs
index e9df87af..164cac46 100644
--- a/tokio/src/time/driver/wheel/mod.rs
+++ b/tokio/src/time/driver/wheel/mod.rs
@@ -118,10 +118,10 @@ impl Wheel {
     /// Remove `item` from the timing wheel.
     pub(crate) unsafe fn remove(&mut self, item: NonNull<TimerShared>) {
         unsafe {
-            if !item.as_ref().might_be_registered() {
+            let when = item.as_ref().cached_when();
+            if when == u64::max_value() {
                 self.pending.remove(item);
             } else {
-                let when = item.as_ref().cached_when();
                 let level = self.level_for(when);
 
                 self.levels[level].remove_entry(item);