From 923e72345c77aef4d11875d874ba67421b0412d3 Mon Sep 17 00:00:00 2001 From: Tim Vilgot Mikael Fredenberg <26655508+vilgotf@users.noreply.github.com> Date: Mon, 25 May 2026 02:11:09 +0200 Subject: [PATCH] time: move lazy-registration state into `Sleep` (#8132) --- tokio/src/runtime/mod.rs | 94 ++++------- tokio/src/runtime/scheduler/mod.rs | 23 --- tokio/src/runtime/time/entry.rs | 112 +++---------- tokio/src/runtime/time/tests/mod.rs | 44 ++--- .../time_alt/cancellation_queue/tests.rs | 4 +- tokio/src/runtime/time_alt/context.rs | 2 - tokio/src/runtime/time_alt/entry.rs | 15 +- .../time_alt/registration_queue/tests.rs | 4 +- tokio/src/runtime/time_alt/tests.rs | 6 +- tokio/src/runtime/time_alt/timer.rs | 131 ++++----------- .../src/runtime/time_alt/wake_queue/tests.rs | 6 +- tokio/src/time/interval.rs | 10 +- tokio/src/time/sleep.rs | 158 +++++++++--------- tokio/tests/time_rt.rs | 30 ++++ tokio/tests/time_sleep.rs | 16 +- tokio/tests/tracing_time.rs | 26 ++- 16 files changed, 260 insertions(+), 421 deletions(-) diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 713f60de0..2b3f0ad9b 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -416,7 +416,7 @@ cfg_process_driver! { mod process; } -#[cfg_attr(not(feature = "time"), allow(dead_code))] +#[allow(dead_code)] #[derive(Debug, Copy, Clone, PartialEq)] pub(crate) enum TimerFlavor { Traditional, @@ -430,6 +430,8 @@ cfg_time! { #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] pub(crate) mod time_alt; + use crate::time::Instant; + use std::task::{Context, Poll}; use std::pin::Pin; @@ -442,27 +444,30 @@ cfg_time! { } impl Timer { + #[cfg_attr(not(all(tokio_unstable, feature = "rt-multi-thread")), allow(unused_variables))] #[track_caller] - pub(crate) fn new( - handle: crate::runtime::scheduler::Handle, - deadline: crate::time::Instant, - ) -> Self { + pub(crate) fn new(handle: scheduler::Handle, deadline: Instant) -> Self { match handle.timer_flavor() { - crate::runtime::TimerFlavor::Traditional => { - Timer::Traditional(time::TimerEntry::new(handle, deadline)) + TimerFlavor::Traditional => { + Timer::Traditional(time::TimerEntry::new(handle)) } #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - crate::runtime::TimerFlavor::Alternative => { + TimerFlavor::Alternative => { Timer::Alternative(time_alt::Timer::new(handle, deadline)) } } } - pub(crate) fn deadline(&self) -> crate::time::Instant { - match self { - Timer::Traditional(entry) => entry.deadline(), + pub(crate) fn init(self: Pin<&mut Self>, deadline: Instant) { + // Safety: we never move the inner entries. + let this = unsafe { self.get_unchecked_mut() }; + match this { + // Safety: we never move the inner entries. + Timer::Traditional(entry) => unsafe { + Pin::new_unchecked(entry).init(deadline) + } #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - Timer::Alternative(entry) => entry.deadline(), + Timer::Alternative(_) => {}, } } @@ -474,28 +479,20 @@ cfg_time! { } } - pub(crate) fn flavor(self: Pin<&Self>) -> TimerFlavor { - match self.get_ref() { - Timer::Traditional(_) => TimerFlavor::Traditional, - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - Timer::Alternative(_) => TimerFlavor::Alternative, - } - } - - pub(crate) fn reset( - self: Pin<&mut Self>, - new_time: crate::time::Instant, - reregister: bool - ) { + #[cfg_attr(not(all(tokio_unstable, feature = "rt-multi-thread")), allow(unused_variables))] + pub(crate) fn reset(self: Pin<&mut Self>, handle: scheduler::Handle, deadline: Instant) { // Safety: we never move the inner entries. let this = unsafe { self.get_unchecked_mut() }; match this { - Timer::Traditional(entry) => { - // Safety: we never move the inner entries. - unsafe { Pin::new_unchecked(entry).reset(new_time, reregister); } + // Safety: we never move the inner entries. + Timer::Traditional(entry) => unsafe { + Pin::new_unchecked(entry).reset(deadline) } + // Safety: we never move the inner entries. #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - Timer::Alternative(_) => panic!("not implemented yet"), + Timer::Alternative(entry) => unsafe { + Pin::new_unchecked(entry).set(time_alt::Timer::new(handle, deadline)) + }, } } @@ -506,44 +503,17 @@ cfg_time! { // Safety: we never move the inner entries. let this = unsafe { self.get_unchecked_mut() }; match this { - Timer::Traditional(entry) => { - // Safety: we never move the inner entries. - unsafe { Pin::new_unchecked(entry).poll_elapsed(cx) } + // Safety: we never move the inner entries. + Timer::Traditional(entry) => unsafe { + Pin::new_unchecked(entry).poll_elapsed(cx) } + // Safety: we never move the inner entries. #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - Timer::Alternative(entry) => { - // Safety: we never move the inner entries. - unsafe { Pin::new_unchecked(entry).poll_elapsed(cx).map(Ok) } + Timer::Alternative(entry) => unsafe { + Pin::new_unchecked(entry).poll_elapsed(cx).map(Ok) } } } - - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - pub(crate) fn scheduler_handle(&self) -> &crate::runtime::scheduler::Handle { - match self { - Timer::Traditional(_) => unreachable!("we should not call this on Traditional Timer"), - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - Timer::Alternative(entry) => entry.scheduler_handle(), - } - } - - #[cfg(all(tokio_unstable, feature = "tracing"))] - pub(crate) fn driver(self: Pin<&Self>) -> &crate::runtime::time::Handle { - match self.get_ref() { - Timer::Traditional(entry) => entry.driver(), - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - Timer::Alternative(entry) => entry.driver(), - } - } - - #[cfg(all(tokio_unstable, feature = "tracing"))] - pub(crate) fn clock(self: Pin<&Self>) -> &crate::time::Clock { - match self.get_ref() { - Timer::Traditional(entry) => entry.clock(), - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - Timer::Alternative(entry) => entry.clock(), - } - } } } diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs index f991e8abd..8bbd110cb 100644 --- a/tokio/src/runtime/scheduler/mod.rs +++ b/tokio/src/runtime/scheduler/mod.rs @@ -119,18 +119,6 @@ cfg_rt! { } } - #[cfg(all(tokio_unstable, feature = "rt-multi-thread", feature = "time"))] - /// Returns true if both handles belong to the same runtime instance. - pub(crate) fn is_same_runtime(&self, other: &Handle) -> bool { - match (self, other) { - (Handle::CurrentThread(a), Handle::CurrentThread(b)) => Arc::ptr_eq(a, b), - #[cfg(feature = "rt-multi-thread")] - (Handle::MultiThread(a), Handle::MultiThread(b)) => Arc::ptr_eq(a, b), - #[cfg(feature = "rt-multi-thread")] - _ => false, // different runtime types - } - } - #[cfg(all(tokio_unstable, feature = "rt-multi-thread", feature = "time"))] /// Returns true if the runtime is shutting down. pub(crate) fn is_shutdown(&self) -> bool { @@ -302,17 +290,6 @@ cfg_rt! { } } - #[cfg(all(tokio_unstable, feature = "time", feature = "rt-multi-thread"))] - pub(crate) fn with_time_temp_local_context(&self, f: F) -> R - where - F: FnOnce(Option>) -> R, - { - match self { - Context::CurrentThread(_) => panic!("the alternative timer implementation is not supported on CurrentThread runtime"), - Context::MultiThread(context) => context.with_time_temp_local_context(f), - } - } - cfg_rt_multi_thread! { #[track_caller] pub(crate) fn expect_multi_thread(&self) -> &multi_thread::Context { diff --git a/tokio/src/runtime/time/entry.rs b/tokio/src/runtime/time/entry.rs index 736105d5a..bfb465eee 100644 --- a/tokio/src/runtime/time/entry.rs +++ b/tokio/src/runtime/time/entry.rs @@ -294,12 +294,7 @@ pin_project! { // // This is manipulated only under the inner mutex. #[pin] - inner: Option, - // Deadline for the timer. This is used to register on the first - // poll, as we can't register prior to being pinned. - deadline: Instant, - // Whether the deadline has been registered. - registered: bool, + inner: TimerShared, } impl PinnedDrop for TimerEntry { @@ -478,71 +473,29 @@ unsafe impl linked_list::Link for TimerShared { // ===== impl Entry ===== impl TimerEntry { - #[track_caller] - pub(crate) fn new(handle: scheduler::Handle, deadline: Instant) -> Self { - // Panic if the time driver is not enabled - let _ = handle.driver().time(); - + pub(crate) fn new(handle: scheduler::Handle) -> Self { Self { driver: handle, - inner: None, - deadline, - registered: false, + inner: TimerShared::new(), } } - fn inner(&self) -> Option<&TimerShared> { - self.inner.as_ref() - } + pub(crate) fn init(self: Pin<&mut Self>, deadline: Instant) { + let tick = self.driver().time_source().deadline_to_tick(deadline); - fn init_inner(self: Pin<&mut Self>) { - match self.inner { - Some(_) => {} - None => self.project().inner.set(Some(TimerShared::new())), + unsafe { + self.driver() + .reregister(&self.driver.driver().io, tick, (&self.inner).into()); } } - pub(crate) fn deadline(&self) -> Instant { - self.deadline - } - pub(crate) fn is_elapsed(&self) -> bool { - let Some(inner) = self.inner() else { - return false; - }; - // Is this timer still in the timer wheel? - let deregistered = !inner.might_be_registered(); - - // Once the timer has expired, - // it will be taken out of the wheel and be fired. - // - // So if we have already registered the timer into the wheel, - // but now it is not in the wheel, it means that it has been - // fired. - // - // +--------------+-----------------+----------+ - // | deregistered | self.registered | output | - // +--------------+-----------------+----------+ - // | true | false | false | <- never been registered - // +--------------+-----------------+----------+ - // | false | false | false | <- never been registered - // +--------------+-----------------+----------+ - // | true | true | true | <- registered into the wheel, - // | | | | and then taken out of the wheel. - // +--------------+-----------------+----------+ - // | false | true | false | <- still registered in the wheel - // +--------------+-----------------+----------+ - deregistered && self.registered + !self.inner.might_be_registered() } /// Cancels and deregisters the timer. This operation is irreversible. pub(crate) fn cancel(self: Pin<&mut Self>) { - // Avoid calling the `clear_entry` method, because it has not been initialized yet. - let Some(inner) = self.inner() else { - return; - }; - // We need to perform an acq/rel fence with the driver thread, and the // simplest way to do so is to grab the driver lock. // @@ -565,38 +518,24 @@ impl TimerEntry { // driver did so far and happens-before everything the driver does in // the future. While we have the lock held, we also go ahead and // deregister the entry if necessary. - unsafe { self.driver().clear_entry(NonNull::from(inner)) }; + unsafe { self.driver().clear_entry(NonNull::from(&self.inner)) }; } - pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant, reregister: bool) { - let this = self.as_mut().project(); - *this.deadline = new_time; - *this.registered = reregister; + pub(crate) fn reset(self: Pin<&mut Self>, deadline: Instant) { + let tick = self.driver().time_source().deadline_to_tick(deadline); - let tick = self.driver().time_source().deadline_to_tick(new_time); - let inner = match self.inner() { - Some(inner) => inner, - None => { - self.as_mut().init_inner(); - self.inner() - .expect("inner should already be initialized by `this.init_inner()`") - } - }; - - if inner.extend_expiration(tick).is_ok() { + if self.inner.extend_expiration(tick).is_ok() { return; } - if reregister { - unsafe { - self.driver() - .reregister(&self.driver.driver().io, tick, inner.into()); - } + unsafe { + self.driver() + .reregister(&self.driver.driver().io, tick, (&self.inner).into()); } } pub(crate) fn poll_elapsed( - mut self: Pin<&mut Self>, + self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { assert!( @@ -605,25 +544,12 @@ impl TimerEntry { crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR ); - if !self.registered { - let deadline = self.deadline; - self.as_mut().reset(deadline, true); - } - - let inner = self - .inner() - .expect("inner should already be initialized by `self.reset()`"); - inner.state.poll(cx.waker()) + self.inner.state.poll(cx.waker()) } - pub(crate) fn driver(&self) -> &super::Handle { + fn driver(&self) -> &super::Handle { self.driver.driver().time() } - - #[cfg(all(tokio_unstable, feature = "tracing"))] - pub(crate) fn clock(&self) -> &super::Clock { - self.driver.driver().clock() - } } impl TimerHandle { diff --git a/tokio/src/runtime/time/tests/mod.rs b/tokio/src/runtime/time/tests/mod.rs index 33c4a5366..84c765af6 100644 --- a/tokio/src/runtime/time/tests/mod.rs +++ b/tokio/src/runtime/time/tests/mod.rs @@ -48,11 +48,11 @@ fn single_timer() { let handle_ = handle.clone(); let jh = thread::spawn(move || { - let entry = TimerEntry::new( - handle_.inner.clone(), - handle_.inner.driver().clock().now() + Duration::from_secs(1), - ); + let entry = TimerEntry::new(handle_.inner.clone()); pin!(entry); + entry + .as_mut() + .init(handle_.inner.driver().clock().now() + Duration::from_secs(1)); block_on(std::future::poll_fn(|cx| entry.as_mut().poll_elapsed(cx))).unwrap(); }); @@ -77,11 +77,11 @@ fn drop_timer() { let handle_ = handle.clone(); let jh = thread::spawn(move || { - let entry = TimerEntry::new( - handle_.inner.clone(), - handle_.inner.driver().clock().now() + Duration::from_secs(1), - ); + let entry = TimerEntry::new(handle_.inner.clone()); pin!(entry); + entry + .as_mut() + .init(handle_.inner.driver().clock().now() + Duration::from_secs(1)); let _ = entry .as_mut() @@ -111,11 +111,11 @@ fn change_waker() { let handle_ = handle.clone(); let jh = thread::spawn(move || { - let entry = TimerEntry::new( - handle_.inner.clone(), - handle_.inner.driver().clock().now() + Duration::from_secs(1), - ); + let entry = TimerEntry::new(handle_.inner.clone()); pin!(entry); + entry + .as_mut() + .init(handle_.inner.driver().clock().now() + Duration::from_secs(1)); let _ = entry .as_mut() @@ -149,14 +149,15 @@ fn reset_future() { let start = handle.inner.driver().clock().now(); let jh = thread::spawn(move || { - let entry = TimerEntry::new(handle_.inner.clone(), start + Duration::from_secs(1)); + let entry = TimerEntry::new(handle_.inner.clone()); pin!(entry); + entry.as_mut().init(start + Duration::from_secs(1)); let _ = entry .as_mut() .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref())); - entry.as_mut().reset(start + Duration::from_secs(2), true); + entry.as_mut().reset(start + Duration::from_secs(2)); // shouldn't complete before 2s block_on(std::future::poll_fn(|cx| entry.as_mut().poll_elapsed(cx))).unwrap(); @@ -206,10 +207,10 @@ fn poll_process_levels() { let mut entries = vec![]; for i in 0..normal_or_miri(1024, 64) { - let mut entry = Box::pin(TimerEntry::new( - handle.inner.clone(), - handle.inner.driver().clock().now() + Duration::from_millis(i), - )); + let mut entry = Box::pin(TimerEntry::new(handle.inner.clone())); + entry + .as_mut() + .init(handle.inner.driver().clock().now() + Duration::from_millis(i)); let _ = entry .as_mut() @@ -240,11 +241,10 @@ fn poll_process_levels_targeted() { let rt = rt(true); let handle = rt.handle(); - let e1 = TimerEntry::new( - handle.inner.clone(), - handle.inner.driver().clock().now() + Duration::from_millis(193), - ); + let e1 = TimerEntry::new(handle.inner.clone()); pin!(e1); + e1.as_mut() + .init(handle.inner.driver().clock().now() + Duration::from_millis(193)); let handle = handle.inner.driver().time(); diff --git a/tokio/src/runtime/time_alt/cancellation_queue/tests.rs b/tokio/src/runtime/time_alt/cancellation_queue/tests.rs index b20e316ac..bd6eb58c8 100644 --- a/tokio/src/runtime/time_alt/cancellation_queue/tests.rs +++ b/tokio/src/runtime/time_alt/cancellation_queue/tests.rs @@ -1,7 +1,5 @@ use super::*; -use futures::task::noop_waker; - #[cfg(loom)] const NUM_ITEMS: usize = 16; @@ -9,7 +7,7 @@ const NUM_ITEMS: usize = 16; const NUM_ITEMS: usize = 64; fn new_handle() -> EntryHandle { - EntryHandle::new(0, noop_waker()) + EntryHandle::new(0) } fn model(f: F) { diff --git a/tokio/src/runtime/time_alt/context.rs b/tokio/src/runtime/time_alt/context.rs index 76035634a..38c7cbf01 100644 --- a/tokio/src/runtime/time_alt/context.rs +++ b/tokio/src/runtime/time_alt/context.rs @@ -25,7 +25,6 @@ pub(crate) enum TempLocalContext<'a> { /// The runtime is running, we can access it. Running { registration_queue: &'a mut RegistrationQueue, - elapsed: u64, }, #[cfg(feature = "rt-multi-thread")] /// The runtime is shutting down, no timers can be registered. @@ -36,7 +35,6 @@ impl<'a> TempLocalContext<'a> { pub(crate) fn new_running(cx: &'a mut LocalContext) -> Self { TempLocalContext::Running { registration_queue: &mut cx.registration_queue, - elapsed: cx.wheel.elapsed(), } } diff --git a/tokio/src/runtime/time_alt/entry.rs b/tokio/src/runtime/time_alt/entry.rs index 40c1fc473..ecf8cd8c5 100644 --- a/tokio/src/runtime/time_alt/entry.rs +++ b/tokio/src/runtime/time_alt/entry.rs @@ -8,7 +8,7 @@ use std::task::{Context, Poll, Waker}; pub(super) type EntryList = linked_list::LinkedList; -#[derive(Debug)] +#[derive(Debug, Default)] struct State { cancelled: bool, woken_up: bool, @@ -34,7 +34,7 @@ pub(crate) struct Entry { /// /// And then, before parking the resource driver, /// the scheduler removes the entry from the [`RegistrationQueue`] - /// [`RegistrationQueue`] and insert it into the [`Wheel`]. + /// and insert it into the [`Wheel`]. /// /// Finally, after parking the resource driver, the scheduler removes /// the entry from the [`Wheel`] and insert it into the [`WakeQueue`]. @@ -186,19 +186,12 @@ impl From<&Handle> for NonNull { } impl Handle { - pub(crate) fn new(deadline: u64, waker: Waker) -> Self { - let state = State { - cancelled: false, - woken_up: false, - waker: Some(waker), - cancel_tx: None, - }; - + pub(crate) fn new(deadline: u64) -> Self { let entry = Arc::new(Entry { cancel_pointers: linked_list::Pointers::new(), extra_pointers: linked_list::Pointers::new(), deadline, - state: Mutex::new(state), + state: Mutex::new(State::default()), _pin: PhantomPinned, }); diff --git a/tokio/src/runtime/time_alt/registration_queue/tests.rs b/tokio/src/runtime/time_alt/registration_queue/tests.rs index b6b3699fa..76e2523b1 100644 --- a/tokio/src/runtime/time_alt/registration_queue/tests.rs +++ b/tokio/src/runtime/time_alt/registration_queue/tests.rs @@ -1,7 +1,5 @@ use super::*; -use futures::task::noop_waker; - #[cfg(loom)] const NUM_ITEMS: usize = 16; @@ -9,7 +7,7 @@ const NUM_ITEMS: usize = 16; const NUM_ITEMS: usize = 64; fn new_handle() -> EntryHandle { - EntryHandle::new(0, noop_waker()) + EntryHandle::new(0) } fn model(f: F) { diff --git a/tokio/src/runtime/time_alt/tests.rs b/tokio/src/runtime/time_alt/tests.rs index 29015f3bd..e1a7196fa 100644 --- a/tokio/src/runtime/time_alt/tests.rs +++ b/tokio/src/runtime/time_alt/tests.rs @@ -1,6 +1,8 @@ use super::*; use crate::loom::thread; +use std::task::Context; + use futures_test::task::{new_count_waker, AwokenCount}; #[cfg(loom)] @@ -11,7 +13,9 @@ const NUM_ITEMS: usize = 64; fn new_handle() -> (EntryHandle, AwokenCount) { let (waker, count) = new_count_waker(); - (EntryHandle::new(0, waker), count) + let entry = EntryHandle::new(0); + _ = entry.poll(&mut Context::from_waker(&waker)); + (entry, count) } fn model(f: F) { diff --git a/tokio/src/runtime/time_alt/timer.rs b/tokio/src/runtime/time_alt/timer.rs index 64ff94dc7..bd87d05b6 100644 --- a/tokio/src/runtime/time_alt/timer.rs +++ b/tokio/src/runtime/time_alt/timer.rs @@ -1,5 +1,5 @@ use super::{EntryHandle, TempLocalContext}; -use crate::runtime::scheduler::Handle as SchedulerHandle; +use crate::runtime::scheduler; use crate::time::Instant; use std::pin::Pin; @@ -9,119 +9,61 @@ use std::task::{Context, Poll}; use crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR; pub(crate) struct Timer { - sched_handle: SchedulerHandle, - /// The entry in the timing wheel. - /// - /// - `Some` if the timer is registered / pending / woken up / cancelling. - /// - `None` if the timer is unregistered. - entry: Option, - - /// The deadline for the timer. - deadline: Instant, + entry: EntryHandle, } impl std::fmt::Debug for Timer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Timer") - .field("deadline", &self.deadline) - .finish() + f.debug_struct("Timer").finish() } } impl Drop for Timer { fn drop(&mut self) { - if let Some(entry) = self.entry.take() { - entry.cancel(); - } + self.entry.cancel(); } } impl Timer { #[track_caller] - pub(crate) fn new(sched_hdl: SchedulerHandle, deadline: Instant) -> Self { - // Panic if the time driver is not enabled - let _ = sched_hdl.driver().time(); - Timer { - sched_handle: sched_hdl, - entry: None, - deadline, - } - } + pub(crate) fn new(handle: scheduler::Handle, deadline: Instant) -> Self { + let tick = deadline_to_tick(&handle, deadline); + let entry = with_current_temp_local_context(|ctx| match ctx { + Some(TempLocalContext::Running { registration_queue }) => { + let entry = EntryHandle::new(tick); + unsafe { registration_queue.push_front(entry.clone()) } + entry + } + #[cfg(feature = "rt-multi-thread")] + Some(TempLocalContext::Shutdown) => panic!("{RUNTIME_SHUTTING_DOWN_ERROR}"), - pub(crate) fn deadline(&self) -> Instant { - self.deadline + _ => { + let entry = EntryHandle::new(tick); + push_from_remote(&handle, entry.clone()); + entry + } + }); + + Timer { entry } } pub(crate) fn is_elapsed(&self) -> bool { - self.entry.as_ref().is_some_and(|entry| entry.is_woken_up()) - } - - fn register(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { - let this = self.get_mut(); - - with_current_temp_local_context(&this.sched_handle, |maybe_time_cx| { - let deadline = deadline_to_tick(&this.sched_handle, this.deadline); - - match maybe_time_cx { - Some(TempLocalContext::Running { - registration_queue: _, - elapsed, - }) if deadline <= elapsed => Poll::Ready(()), - - Some(TempLocalContext::Running { - registration_queue, - elapsed: _, - }) => { - let hdl = EntryHandle::new(deadline, cx.waker().clone()); - this.entry = Some(hdl.clone()); - unsafe { - registration_queue.push_front(hdl); - } - Poll::Pending - } - #[cfg(feature = "rt-multi-thread")] - Some(TempLocalContext::Shutdown) => panic!("{RUNTIME_SHUTTING_DOWN_ERROR}"), - - _ => { - let hdl = EntryHandle::new(deadline, cx.waker().clone()); - this.entry = Some(hdl.clone()); - push_from_remote(&this.sched_handle, hdl); - Poll::Pending - } - } - }) + self.entry.is_woken_up() } pub(crate) fn poll_elapsed(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { - match self.entry.as_ref() { - Some(entry) => entry.poll(cx), - None => self.register(cx), - } - } - - pub(crate) fn scheduler_handle(&self) -> &SchedulerHandle { - &self.sched_handle - } - - #[cfg(all(tokio_unstable, feature = "tracing"))] - pub(crate) fn driver(&self) -> &crate::runtime::time::Handle { - self.sched_handle.driver().time() - } - - #[cfg(all(tokio_unstable, feature = "tracing"))] - pub(crate) fn clock(&self) -> &crate::time::Clock { - self.sched_handle.driver().clock() + self.entry.poll(cx) } } -fn with_current_temp_local_context(hdl: &SchedulerHandle, f: F) -> R +fn with_current_temp_local_context(f: F) -> R where F: FnOnce(Option>) -> R, { #[cfg(not(feature = "rt"))] { - let (_, _) = (hdl, f); + let _ = f; panic!("Tokio runtime is not enabled, cannot access the current wheel"); } @@ -129,23 +71,14 @@ where { use crate::runtime::context; - let is_same_rt = - context::with_current(|cur_hdl| cur_hdl.is_same_runtime(hdl)).unwrap_or_default(); - - if !is_same_rt { - // We don't want to create the timer in one runtime, - // but register it in a different runtime's timer wheel. - f(None) - } else { - context::with_scheduler(|maybe_cx| match maybe_cx { - Some(cx) => cx.with_time_temp_local_context(f), - None => f(None), - }) - } + context::with_scheduler(|maybe_cx| match maybe_cx { + Some(cx) => cx.expect_multi_thread().with_time_temp_local_context(f), + None => f(None), + }) } } -fn push_from_remote(sched_hdl: &SchedulerHandle, entry_hdl: EntryHandle) { +fn push_from_remote(sched_hdl: &scheduler::Handle, entry_hdl: EntryHandle) { #[cfg(not(feature = "rt"))] { let (_, _) = (sched_hdl, entry_hdl); @@ -159,7 +92,7 @@ fn push_from_remote(sched_hdl: &SchedulerHandle, entry_hdl: EntryHandle) { } } -fn deadline_to_tick(sched_hdl: &SchedulerHandle, deadline: Instant) -> u64 { +fn deadline_to_tick(sched_hdl: &scheduler::Handle, deadline: Instant) -> u64 { let time_hdl = sched_hdl.driver().time(); time_hdl.time_source().deadline_to_tick(deadline) } diff --git a/tokio/src/runtime/time_alt/wake_queue/tests.rs b/tokio/src/runtime/time_alt/wake_queue/tests.rs index f0449ee91..bc6fb0051 100644 --- a/tokio/src/runtime/time_alt/wake_queue/tests.rs +++ b/tokio/src/runtime/time_alt/wake_queue/tests.rs @@ -1,5 +1,7 @@ use super::*; +use std::task::Context; + use futures_test::task::{new_count_waker, AwokenCount}; #[cfg(loom)] @@ -10,7 +12,9 @@ const NUM_ITEMS: usize = 64; fn new_handle() -> (EntryHandle, AwokenCount) { let (waker, count) = new_count_waker(); - (EntryHandle::new(0, waker), count) + let entry = EntryHandle::new(0); + _ = entry.poll(&mut Context::from_waker(&waker)); + (entry, count) } fn model(f: F) { diff --git a/tokio/src/time/interval.rs b/tokio/src/time/interval.rs index 02cecc6ec..fb8baa983 100644 --- a/tokio/src/time/interval.rs +++ b/tokio/src/time/interval.rs @@ -131,14 +131,8 @@ fn internal_interval_at( ) }; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let delay = resource_span.in_scope(|| Box::pin(sleep_until(start))); - - #[cfg(not(all(tokio_unstable, feature = "tracing")))] - let delay = Box::pin(sleep_until(start)); - Interval { - delay, + delay: Box::pin(sleep_until(start)), period, missed_tick_behavior: MissedTickBehavior::default(), #[cfg(all(tokio_unstable, feature = "tracing"))] @@ -487,7 +481,7 @@ impl Interval { // When we arrive here, the internal delay returned `Poll::Ready`. // Reset the delay but do not register it. It should be registered with // the next call to [`poll_tick`]. - self.delay.as_mut().reset_without_reregister(next); + self.delay.as_mut().reset_without_timer(next); // Return the time when we were scheduled to tick Poll::Ready(timeout) diff --git a/tokio/src/time/sleep.rs b/tokio/src/time/sleep.rs index 2aa19b666..e4535e4d6 100644 --- a/tokio/src/time/sleep.rs +++ b/tokio/src/time/sleep.rs @@ -1,4 +1,4 @@ -use crate::runtime::Timer; +use crate::runtime::{scheduler, Timer}; use crate::time::{error::Error, Duration, Instant}; use crate::util::trace; @@ -223,11 +223,11 @@ pin_project! { #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Sleep { + deadline: Instant, + driver: scheduler::Handle, inner: Inner, - - // The link between the `Sleep` instance and the timer that drives it. #[pin] - entry: Timer, + timer: Option, } } @@ -251,18 +251,11 @@ impl Sleep { deadline: Instant, location: Option<&'static Location<'static>>, ) -> Sleep { - use crate::runtime::scheduler; let handle = scheduler::Handle::current(); - let entry = Timer::new(handle, deadline); + // Panic if the time driver is not enabled (backwards compat) + _ = handle.driver().time(); #[cfg(all(tokio_unstable, feature = "tracing"))] let inner = { - let handle = scheduler::Handle::current(); - let clock = handle.driver().clock(); - let handle = &handle.driver().time(); - let time_source = handle.time_source(); - let deadline_tick = time_source.deadline_to_tick(deadline); - let duration = deadline_tick.saturating_sub(time_source.now(clock)); - let location = location.expect("should have location if tracing"); let resource_span = tracing::trace_span!( parent: None, @@ -274,19 +267,14 @@ impl Sleep { loc.col = location.column(), ); - let async_op_span = resource_span.in_scope(|| { - tracing::trace!( - target: "runtime::resource::state_update", - duration = duration, - duration.unit = "ms", - duration.op = "override", - ); - - tracing::trace_span!("runtime.resource.async_op", source = "Sleep::new_timeout") - }); + let async_op_span = tracing::trace_span!( + parent: &resource_span, + "runtime.resource.async_op", + source = "Sleep::new_timeout", + ); let async_op_poll_span = - async_op_span.in_scope(|| tracing::trace_span!("runtime.resource.async_op.poll")); + tracing::trace_span!(parent: &async_op_span, "runtime.resource.async_op.poll"); let ctx = trace::AsyncOpTracingCtx { async_op_span, @@ -300,7 +288,12 @@ impl Sleep { #[cfg(not(all(tokio_unstable, feature = "tracing")))] let inner = Inner {}; - Sleep { inner, entry } + Sleep { + deadline, + driver: handle, + inner, + timer: None, + } } pub(crate) fn far_future(location: Option<&'static Location<'static>>) -> Sleep { @@ -309,14 +302,14 @@ impl Sleep { /// Returns the instant at which the future will complete. pub fn deadline(&self) -> Instant { - self.entry.deadline() + self.deadline } /// Returns `true` if `Sleep` has elapsed. /// /// A `Sleep` instance is elapsed when the requested duration has elapsed. pub fn is_elapsed(&self) -> bool { - self.entry.is_elapsed() + self.timer.as_ref().is_some_and(Timer::is_elapsed) } /// Resets the `Sleep` instance to a new deadline. @@ -349,74 +342,53 @@ impl Sleep { /// /// [`Pin::as_mut`]: fn@std::pin::Pin::as_mut pub fn reset(self: Pin<&mut Self>, deadline: Instant) { - self.reset_inner(deadline); - } + let mut this = self.project(); + *this.deadline = deadline; - /// Resets the `Sleep` instance to a new deadline without reregistering it - /// to be woken up. - /// - /// Calling this function allows changing the instant at which the `Sleep` - /// future completes without having to create new associated state and - /// without having it registered. This is required in e.g. the - /// [`crate::time::Interval`] where we want to reset the internal [Sleep] - /// without having it wake up the last task that polled it. - pub(crate) fn reset_without_reregister(self: Pin<&mut Self>, deadline: Instant) { - let mut me = self.project(); - match me.entry.as_ref().flavor() { - crate::runtime::TimerFlavor::Traditional => { - me.entry.as_mut().reset(deadline, false); - } - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - crate::runtime::TimerFlavor::Alternative => { - let handle = me.entry.as_ref().scheduler_handle().clone(); - me.entry.set(Timer::new(handle, deadline)); - } - } - } - - fn reset_inner(self: Pin<&mut Self>, deadline: Instant) { - let mut me = self.project(); - match me.entry.as_ref().flavor() { - crate::runtime::TimerFlavor::Traditional => { - me.entry.as_mut().reset(deadline, true); - } - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - crate::runtime::TimerFlavor::Alternative => { - let handle = me.entry.as_ref().scheduler_handle().clone(); - me.entry.set(Timer::new(handle, deadline)); - } - } + let handle = this.driver; #[cfg(all(tokio_unstable, feature = "tracing"))] { - let _resource_enter = me.inner.ctx.resource_span.enter(); - me.inner.ctx.async_op_span = + let _resource_enter = this.inner.ctx.resource_span.enter(); + this.inner.ctx.async_op_span = tracing::trace_span!("runtime.resource.async_op", source = "Sleep::reset"); - let _async_op_enter = me.inner.ctx.async_op_span.enter(); + let _async_op_enter = this.inner.ctx.async_op_span.enter(); - me.inner.ctx.async_op_poll_span = + this.inner.ctx.async_op_poll_span = tracing::trace_span!("runtime.resource.async_op.poll"); - let duration = { - let clock = me.entry.as_ref().clock(); - let time_source = me.entry.as_ref().driver().time_source(); - let now = time_source.now(clock); - let deadline_tick = time_source.deadline_to_tick(deadline); - deadline_tick.saturating_sub(now) - }; - + let clock = handle.driver().clock(); + let time_source = handle.driver().time().time_source(); + let now = time_source.now(clock); + let tick = time_source.deadline_to_tick(deadline); tracing::trace!( target: "runtime::resource::state_update", - duration = duration, + duration = tick.saturating_sub(now), duration.unit = "ms", duration.op = "override", ); } + + match this.timer.as_mut().as_pin_mut() { + Some(timer) => timer.reset(handle.clone(), deadline), + None => { + let timer = Timer::new(handle.clone(), deadline); + this.timer.set(Some(timer)); + this.timer.as_pin_mut().unwrap().init(deadline); + } + } + } + + /// Resets the `Sleep` instance to a new deadline. + /// + /// Unlike [`reset`][Self::reset], this __removes__ the internal timer. + pub(super) fn reset_without_timer(self: Pin<&mut Self>, deadline: Instant) { + let mut this = self.project(); + *this.deadline = deadline; + this.timer.set(None); } fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll> { - let me = self.project(); - ready!(crate::trace::trace_leaf()); // Keep track of task budget @@ -429,7 +401,35 @@ impl Sleep { #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] let coop = ready!(crate::task::coop::poll_proceed(cx)); - let result = me.entry.poll_elapsed(cx).map(move |r| { + let mut this = self.project(); + let timer = match this.timer.as_mut().as_pin_mut() { + Some(timer) => timer, + None => { + let handle = this.driver; + + #[cfg(all(tokio_unstable, feature = "tracing"))] + { + let clock = handle.driver().clock(); + let time_source = handle.driver().time().time_source(); + let now = time_source.now(clock); + let tick = time_source.deadline_to_tick(*this.deadline); + tracing::trace!( + target: "runtime::resource::state_update", + duration = tick.saturating_sub(now), + duration.unit = "ms", + duration.op = "override", + ); + } + + let timer = Timer::new(handle.clone(), *this.deadline); + this.timer.set(Some(timer)); + let mut timer = this.timer.as_pin_mut().unwrap(); + timer.as_mut().init(*this.deadline); + timer + } + }; + + let result = timer.poll_elapsed(cx).map(move |r| { coop.made_progress(); r }); diff --git a/tokio/tests/time_rt.rs b/tokio/tests/time_rt.rs index 283967798..d86866ac6 100644 --- a/tokio/tests/time_rt.rs +++ b/tokio/tests/time_rt.rs @@ -1,8 +1,10 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "full")] +use futures_test::task::noop_context; use tokio::runtime::Runtime; use tokio::time::*; +use tokio_test::assert_pending; use std::sync::mpsc; @@ -164,3 +166,31 @@ fn timeout_value() { }); } } + +#[test] +fn tickspace() { + use std::future::Future as _; + use std::thread; + let rt = || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .start_paused(true) + .build() + .unwrap() + }; + + let rt_past = rt(); + thread::sleep(Duration::from_millis(1)); + let rt = rt(); + + let _guard = rt_past.enter(); + let mut sleep = std::pin::pin!(sleep(Duration::from_millis(1))); + assert_pending!(sleep.as_mut().poll(&mut noop_context())); + + let deadline = sleep.deadline(); + rt.block_on(async { sleep.as_mut().reset(deadline + Duration::from_millis(1)) }); + + let now = Instant::now(); + rt_past.block_on(sleep); + assert_eq!(now.elapsed(), Duration::from_millis(2)); +} diff --git a/tokio/tests/time_sleep.rs b/tokio/tests/time_sleep.rs index 857b3822d..b82b1cc6a 100644 --- a/tokio/tests/time_sleep.rs +++ b/tokio/tests/time_sleep.rs @@ -10,15 +10,19 @@ use futures::task::noop_waker_ref; use tokio::time::{self, Duration, Instant}; use tokio_test::{assert_elapsed, assert_pending, assert_ready, task}; -#[tokio::test] +#[tokio::test(start_paused = true)] async fn immediate_sleep() { - time::pause(); - let now = Instant::now(); - // Ready! - time::sleep_until(now).await; - assert_elapsed!(now, ms(1)); + let sleep = time::sleep_until(now); + + tokio::pin!(sleep); + + assert!(!sleep.is_elapsed()); + + sleep.as_mut().await; + assert_elapsed!(now, ms(0)); + assert!(sleep.is_elapsed()); } #[tokio::test] diff --git a/tokio/tests/tracing_time.rs b/tokio/tests/tracing_time.rs index f251cc780..261237b7c 100644 --- a/tokio/tests/tracing_time.rs +++ b/tokio/tests/tracing_time.rs @@ -17,6 +17,12 @@ async fn test_sleep_creates_span() { .named("runtime.resource") .with_target("tokio::time::sleep"); + let poll_op = || { + expect::event() + .with_target("runtime::resource::poll_op") + .with_fields(expect::field("op_name").with_value(&"poll_elapsed")) + }; + let state_update = expect::event() .with_target("runtime::resource::state_update") .with_fields( @@ -40,31 +46,35 @@ async fn test_sleep_creates_span() { let (subscriber, handle) = subscriber::mock() .new_span(sleep_span.clone().with_ancestry(expect::is_explicit_root())) - .enter(sleep_span.clone()) - .event(state_update) .new_span( async_op_span .clone() - .with_ancestry(expect::has_contextual_parent(&sleep_span_id)) + .with_ancestry(expect::has_explicit_parent(&sleep_span_id)) .with_fields(expect::field("source").with_value(&"Sleep::new_timeout")), ) - .exit(sleep_span.clone()) - .enter(async_op_span.clone()) .new_span( async_op_poll_span .clone() - .with_ancestry(expect::has_contextual_parent(&async_op_span_id)), + .with_ancestry(expect::has_explicit_parent(&async_op_span_id)), ) + .enter(sleep_span.clone()) + .enter(async_op_span.clone()) + .enter(async_op_poll_span.clone()) + .event(poll_op()) + .event(state_update) + .event(poll_op()) + .exit(async_op_poll_span.clone()) + .drop_span(async_op_poll_span) .exit(async_op_span.clone()) .drop_span(async_op_span) - .drop_span(async_op_poll_span) + .exit(sleep_span.clone()) .drop_span(sleep_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); - _ = tokio::time::sleep(Duration::from_millis(7)); + tokio::time::sleep(Duration::from_millis(7)).await; } handle.assert_finished();