diff --git a/tokio/src/runtime/time_alt/entry.rs b/tokio/src/runtime/time_alt/entry.rs index f9a4746b8..40c1fc473 100644 --- a/tokio/src/runtime/time_alt/entry.rs +++ b/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; @@ -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) { diff --git a/tokio/src/runtime/time_alt/timer.rs b/tokio/src/runtime/time_alt/timer.rs index eb1505fd0..64ff94dc7 100644 --- a/tokio/src/runtime/time_alt/timer.rs +++ b/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), } }