diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 11dfa274e..c619f1d47 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -280,10 +280,43 @@ pin_project! { tokio_thread_local!(static CURRENT: LocalData = const { LocalData { ctx: RcCell::new(), + wake_on_schedule: Cell::new(false), } }); struct LocalData { ctx: RcCell, + wake_on_schedule: Cell, +} + +impl LocalData { + /// Should be called except when we call `LocalSet::enter`. + /// Especially when we poll a LocalSet. + #[must_use = "dropping this guard will reset the entered state"] + fn enter(&self, ctx: Rc) -> LocalDataEnterGuard<'_> { + let ctx = self.ctx.replace(Some(ctx)); + let wake_on_schedule = self.wake_on_schedule.replace(false); + LocalDataEnterGuard { + local_data_ref: self, + ctx, + wake_on_schedule, + } + } +} + +/// A guard for `LocalData::enter()` +struct LocalDataEnterGuard<'a> { + local_data_ref: &'a LocalData, + ctx: Option>, + wake_on_schedule: bool, +} + +impl<'a> Drop for LocalDataEnterGuard<'a> { + fn drop(&mut self) { + self.local_data_ref.ctx.set(self.ctx.take()); + self.local_data_ref + .wake_on_schedule + .set(self.wake_on_schedule) + } } cfg_rt! { @@ -360,13 +393,26 @@ const MAX_TASKS_PER_TICK: usize = 61; const REMOTE_FIRST_INTERVAL: u8 = 31; /// Context guard for LocalSet -pub struct LocalEnterGuard(Option>); +pub struct LocalEnterGuard { + ctx: Option>, + + /// Distinguishes whether the context was entered or being polled. + /// When we enter it, the value `wake_on_schedule` is set. In this case + /// `spawn_local` refers the context, whereas it is not being polled now. + wake_on_schedule: bool, +} impl Drop for LocalEnterGuard { fn drop(&mut self) { - CURRENT.with(|LocalData { ctx, .. }| { - ctx.set(self.0.take()); - }) + CURRENT.with( + |LocalData { + ctx, + wake_on_schedule, + }| { + ctx.set(self.ctx.take()); + wake_on_schedule.set(self.wake_on_schedule); + }, + ) } } @@ -408,10 +454,20 @@ impl LocalSet { /// /// [`spawn_local`]: fn@crate::task::spawn_local pub fn enter(&self) -> LocalEnterGuard { - CURRENT.with(|LocalData { ctx, .. }| { - let old = ctx.replace(Some(self.context.clone())); - LocalEnterGuard(old) - }) + CURRENT.with( + |LocalData { + ctx, + wake_on_schedule, + .. + }| { + let ctx = ctx.replace(Some(self.context.clone())); + let wake_on_schedule = wake_on_schedule.replace(true); + LocalEnterGuard { + ctx, + wake_on_schedule, + } + }, + ) } /// Spawns a `!Send` task onto the local task set. @@ -667,23 +723,8 @@ impl LocalSet { } fn with(&self, f: impl FnOnce() -> T) -> T { - CURRENT.with(|LocalData { ctx, .. }| { - struct Reset<'a> { - ctx_ref: &'a RcCell, - val: Option>, - } - impl<'a> Drop for Reset<'a> { - fn drop(&mut self) { - self.ctx_ref.set(self.val.take()); - } - } - let old = ctx.replace(Some(self.context.clone())); - - let _reset = Reset { - ctx_ref: ctx, - val: old, - }; - + CURRENT.with(|local_data| { + let _guard = local_data.enter(self.context.clone()); f() }) } @@ -693,23 +734,8 @@ impl LocalSet { fn with_if_possible(&self, f: impl FnOnce() -> T) -> T { let mut f = Some(f); - let res = CURRENT.try_with(|LocalData { ctx, .. }| { - struct Reset<'a> { - ctx_ref: &'a RcCell, - val: Option>, - } - impl<'a> Drop for Reset<'a> { - fn drop(&mut self) { - self.ctx_ref.replace(self.val.take()); - } - } - let old = ctx.replace(Some(self.context.clone())); - - let _reset = Reset { - ctx_ref: ctx, - val: old, - }; - + let res = CURRENT.try_with(|local_data| { + let _guard = local_data.enter(self.context.clone()); (f.take().unwrap())() }); @@ -967,7 +993,10 @@ impl Shared { fn schedule(&self, task: task::Notified>) { CURRENT.with(|localdata| { match localdata.ctx.get() { - Some(cx) if cx.shared.ptr_eq(self) => unsafe { + // If the current `LocalSet` is being polled, we don't need to wake it. + // When we `enter` it, then the value `wake_on_schedule` is set to be true. + // In this case it is not being polled, so we need to wake it. + Some(cx) if cx.shared.ptr_eq(self) && !localdata.wake_on_schedule.get() => unsafe { // Safety: if the current `LocalSet` context points to this // `LocalSet`, then we are on the thread that owns it. cx.shared.local_state.task_push_back(task); diff --git a/tokio/tests/task_local_set.rs b/tokio/tests/task_local_set.rs index baca58b77..168a05808 100644 --- a/tokio/tests/task_local_set.rs +++ b/tokio/tests/task_local_set.rs @@ -573,6 +573,27 @@ async fn spawn_wakes_localset() { } } +/// Checks that the task wakes up with `enter`. +/// Reproduces . +#[tokio::test] +async fn sleep_with_local_enter_guard() { + let local = LocalSet::new(); + let _guard = local.enter(); + + let (tx, rx) = oneshot::channel(); + + local + .run_until(async move { + tokio::task::spawn_local(async move { + time::sleep(Duration::ZERO).await; + + tx.send(()).expect("failed to send"); + }); + assert_eq!(rx.await, Ok(())); + }) + .await; +} + #[test] fn store_local_set_in_thread_local_with_runtime() { use tokio::runtime::Runtime;