time: consolidate mutex locks on spurious poll (#8124)

tokio/src/runtime/time_alt/entry.rs

@@ -4,7 +4,7 @@ use crate::util::linked_list;
 
 use std::marker::PhantomPinned;
 use std::ptr::NonNull;
-use std::task::Waker;
+use std::task::{Context, Poll, Waker};
 
 pub(super) type EntryList = linked_list::LinkedList<Entry, Entry>;
 
@@ -228,19 +228,24 @@ impl Handle {
         }
     }
 
-    pub(crate) fn register_waker(&self, waker: &Waker) {
+    pub(crate) fn poll(&self, cx: &mut Context<'_>) -> Poll<()> {
         let mut lock = self.entry.state.lock();
-        if !lock.cancelled && !lock.woken_up {
-            // PANIC: no intermediary state is possible should the user-controllable `Waker`
-            // panic on `Clone` or `Drop`.
-            let maybe_old_waker = match &lock.waker {
-                Some(current_waker) if current_waker.will_wake(waker) => None,
-                _ => lock.waker.replace(waker.clone()),
-            };
-            // unlock before dropping waker
-            drop(lock);
-            drop(maybe_old_waker);
+        if lock.woken_up {
+            return Poll::Ready(());
+        } else if lock.cancelled {
+            return Poll::Pending;
         }
+        // PANIC: no intermediary state is possible should the user-controllable `Waker`
+        // panic on `Clone` or `Drop`.
+        let maybe_old_waker = match &lock.waker {
+            Some(current_waker) if current_waker.will_wake(cx.waker()) => None,
+            _ => lock.waker.replace(cx.waker().clone()),
+        };
+        // unlock before dropping waker
+        drop(lock);
+        drop(maybe_old_waker);
+
+        Poll::Pending
     }
 
     pub(crate) fn cancel(&self) {

tokio/src/runtime/time_alt/timer.rs

@@ -95,11 +95,7 @@ impl Timer {
 
     pub(crate) fn poll_elapsed(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
         match self.entry.as_ref() {
-            Some(entry) if entry.is_woken_up() => Poll::Ready(()),
-            Some(entry) => {
-                entry.register_waker(cx.waker());
-                Poll::Pending
-            }
+            Some(entry) => entry.poll(cx),
             None => self.register(cx),
         }
     }