From edb8f21a741358f7c80b744f008f1e5acc77b429 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Sat, 7 Dec 2024 14:45:16 +0100 Subject: [PATCH 1/7] Take critical section instead of unsafe --- embassy-executor/src/raw/mod.rs | 9 --------- embassy-executor/src/raw/run_queue_atomics.rs | 11 ++++++++++- .../src/raw/run_queue_critical_section.rs | 16 ++++++++++++---- 3 files changed, 22 insertions(+), 14 deletions(-) diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs index 6503b556f..c79fdae60 100644 --- a/embassy-executor/src/raw/mod.rs +++ b/embassy-executor/src/raw/mod.rs @@ -411,15 +411,6 @@ impl SyncExecutor { self.run_queue.dequeue_all(|p| { let task = p.header(); - if !task.state.run_dequeue() { - // If task is not running, ignore it. This can happen in the following scenario: - // - Task gets dequeued, poll starts - // - While task is being polled, it gets woken. It gets placed in the queue. - // - Task poll finishes, returning done=true - // - RUNNING bit is cleared, but the task is already in the queue. - return; - } - #[cfg(feature = "trace")] trace::task_exec_begin(self, &p); diff --git a/embassy-executor/src/raw/run_queue_atomics.rs b/embassy-executor/src/raw/run_queue_atomics.rs index efdafdff0..aad90d767 100644 --- a/embassy-executor/src/raw/run_queue_atomics.rs +++ b/embassy-executor/src/raw/run_queue_atomics.rs @@ -81,7 +81,16 @@ impl RunQueue { // safety: there are no concurrent accesses to `next` next = unsafe { task.header().run_queue_item.next.get() }; - on_task(task); + let run_task = task.header().state.run_dequeue(); + + if run_task { + // If task is not running, ignore it. This can happen in the following scenario: + // - Task gets dequeued, poll starts + // - While task is being polled, it gets woken. It gets placed in the queue. + // - Task poll finishes, returning done=true + // - RUNNING bit is cleared, but the task is already in the queue. + on_task(task); + } } } } diff --git a/embassy-executor/src/raw/run_queue_critical_section.rs b/embassy-executor/src/raw/run_queue_critical_section.rs index 90f09e8c8..4f1b2855a 100644 --- a/embassy-executor/src/raw/run_queue_critical_section.rs +++ b/embassy-executor/src/raw/run_queue_critical_section.rs @@ -63,11 +63,19 @@ impl RunQueue { // If the task re-enqueues itself, the `next` pointer will get overwritten. // Therefore, first read the next pointer, and only then process the task. - // safety: we know if the task is enqueued, no one else will touch the `next` pointer. - let cs = unsafe { CriticalSection::new() }; - next = task.header().run_queue_item.next.borrow(cs).get(); + let run_task = critical_section::with(|cs| { + next = task.header().run_queue_item.next.borrow(cs).get(); + task.header().state.run_dequeue(cs) + }); - on_task(task); + if run_task { + // If task is not running, ignore it. This can happen in the following scenario: + // - Task gets dequeued, poll starts + // - While task is being polled, it gets woken. It gets placed in the queue. + // - Task poll finishes, returning done=true + // - RUNNING bit is cleared, but the task is already in the queue. + on_task(task); + } } } } From 8fd08b1e97533c7526bb4937770060d18bb37410 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Tue, 17 Dec 2024 18:05:48 +0100 Subject: [PATCH 2/7] Swap poll_fn to allow polling exited tasks --- embassy-executor/src/raw/mod.rs | 17 +++++++++++++++-- embassy-executor/src/raw/run_queue_atomics.rs | 12 ++---------- .../src/raw/run_queue_critical_section.rs | 13 +++---------- embassy-executor/src/raw/state_atomics.rs | 5 ++--- embassy-executor/src/raw/state_atomics_arm.rs | 4 +--- .../src/raw/state_critical_section.rs | 8 ++------ 6 files changed, 25 insertions(+), 34 deletions(-) diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs index c79fdae60..242e9c365 100644 --- a/embassy-executor/src/raw/mod.rs +++ b/embassy-executor/src/raw/mod.rs @@ -202,16 +202,29 @@ impl TaskStorage { } } + unsafe fn poll_to_despawn(p: TaskRef) { + // The task's future has already been dropped, we just mark it as `!SPAWNED`. + let this = &*p.as_ptr().cast::>(); + this.raw.state.despawn(); + } + unsafe fn poll(p: TaskRef) { - let this = &*(p.as_ptr() as *const TaskStorage); + let this = &*p.as_ptr().cast::>(); let future = Pin::new_unchecked(this.future.as_mut()); let waker = waker::from_task(p); let mut cx = Context::from_waker(&waker); match future.poll(&mut cx) { Poll::Ready(_) => { + waker.wake_by_ref(); + + // As the future has finished and this function will not be called + // again, we can safely drop the future here. this.future.drop_in_place(); - this.raw.state.despawn(); + + // We replace the poll_fn with a despawn function, so that the task is cleaned up + // when the executor polls it next. + this.raw.poll_fn.set(Some(Self::poll_to_despawn)); } Poll::Pending => {} } diff --git a/embassy-executor/src/raw/run_queue_atomics.rs b/embassy-executor/src/raw/run_queue_atomics.rs index aad90d767..ce511d79a 100644 --- a/embassy-executor/src/raw/run_queue_atomics.rs +++ b/embassy-executor/src/raw/run_queue_atomics.rs @@ -81,16 +81,8 @@ impl RunQueue { // safety: there are no concurrent accesses to `next` next = unsafe { task.header().run_queue_item.next.get() }; - let run_task = task.header().state.run_dequeue(); - - if run_task { - // If task is not running, ignore it. This can happen in the following scenario: - // - Task gets dequeued, poll starts - // - While task is being polled, it gets woken. It gets placed in the queue. - // - Task poll finishes, returning done=true - // - RUNNING bit is cleared, but the task is already in the queue. - on_task(task); - } + task.header().state.run_dequeue(); + on_task(task); } } } diff --git a/embassy-executor/src/raw/run_queue_critical_section.rs b/embassy-executor/src/raw/run_queue_critical_section.rs index 4f1b2855a..86c4085ed 100644 --- a/embassy-executor/src/raw/run_queue_critical_section.rs +++ b/embassy-executor/src/raw/run_queue_critical_section.rs @@ -63,19 +63,12 @@ impl RunQueue { // If the task re-enqueues itself, the `next` pointer will get overwritten. // Therefore, first read the next pointer, and only then process the task. - let run_task = critical_section::with(|cs| { + critical_section::with(|cs| { next = task.header().run_queue_item.next.borrow(cs).get(); - task.header().state.run_dequeue(cs) + task.header().state.run_dequeue(cs); }); - if run_task { - // If task is not running, ignore it. This can happen in the following scenario: - // - Task gets dequeued, poll starts - // - While task is being polled, it gets woken. It gets placed in the queue. - // - Task poll finishes, returning done=true - // - RUNNING bit is cleared, but the task is already in the queue. - on_task(task); - } + on_task(task); } } } diff --git a/embassy-executor/src/raw/state_atomics.rs b/embassy-executor/src/raw/state_atomics.rs index bdd317b53..b6576bfc2 100644 --- a/embassy-executor/src/raw/state_atomics.rs +++ b/embassy-executor/src/raw/state_atomics.rs @@ -52,8 +52,7 @@ impl State { /// Unmark the task as run-queued. Return whether the task is spawned. #[inline(always)] - pub fn run_dequeue(&self) -> bool { - let state = self.state.fetch_and(!STATE_RUN_QUEUED, Ordering::AcqRel); - state & STATE_SPAWNED != 0 + pub fn run_dequeue(&self) { + self.state.fetch_and(!STATE_RUN_QUEUED, Ordering::AcqRel); } } diff --git a/embassy-executor/src/raw/state_atomics_arm.rs b/embassy-executor/src/raw/state_atomics_arm.rs index cbda0d89d..b743dcc2c 100644 --- a/embassy-executor/src/raw/state_atomics_arm.rs +++ b/embassy-executor/src/raw/state_atomics_arm.rs @@ -75,11 +75,9 @@ impl State { /// Unmark the task as run-queued. Return whether the task is spawned. #[inline(always)] - pub fn run_dequeue(&self) -> bool { + pub fn run_dequeue(&self) { compiler_fence(Ordering::Release); - let r = self.spawned.load(Ordering::Relaxed); self.run_queued.store(false, Ordering::Relaxed); - r } } diff --git a/embassy-executor/src/raw/state_critical_section.rs b/embassy-executor/src/raw/state_critical_section.rs index 4733af278..6b627ff79 100644 --- a/embassy-executor/src/raw/state_critical_section.rs +++ b/embassy-executor/src/raw/state_critical_section.rs @@ -67,11 +67,7 @@ impl State { /// Unmark the task as run-queued. Return whether the task is spawned. #[inline(always)] - pub fn run_dequeue(&self) -> bool { - self.update(|s| { - let ok = *s & STATE_SPAWNED != 0; - *s &= !STATE_RUN_QUEUED; - ok - }) + pub fn run_dequeue(&self, cs: CriticalSection<'_>) { + self.update_with_cs(cs, |s| *s &= !STATE_RUN_QUEUED) } } From b51bd9ad040754728142df9991763b4672c31ecd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Tue, 17 Dec 2024 18:09:51 +0100 Subject: [PATCH 3/7] Update tests --- embassy-executor/tests/test.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/embassy-executor/tests/test.rs b/embassy-executor/tests/test.rs index 78c49c071..d8c5a6ae3 100644 --- a/embassy-executor/tests/test.rs +++ b/embassy-executor/tests/test.rs @@ -69,6 +69,7 @@ fn executor_task() { &[ "pend", // spawning a task pends the executor "poll task1", // poll only once. + "pend", // task is done, wakes itself to exit ] ) } @@ -179,6 +180,7 @@ fn waking_after_completion_does_not_poll() { "pend", // manual wake, single pend for two wakes "pend", // respawning a task pends the executor "poll task1", // + "pend", // task is done, wakes itself to exit ] ) } @@ -266,6 +268,7 @@ fn waking_with_old_waker_after_respawn() { "yield_now", // "pend", // manual wake, gets cleared by poll "poll task1", // + "pend", // task is done, wakes itself to exit ] ); } From 7d5fbe26c955bae4bd394d1092702bd81f849c9b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Tue, 17 Dec 2024 18:17:36 +0100 Subject: [PATCH 4/7] Update state diagram --- embassy-executor/src/raw/mod.rs | 43 ++++++++++++++++----------------- 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs index 242e9c365..5df5ca9e1 100644 --- a/embassy-executor/src/raw/mod.rs +++ b/embassy-executor/src/raw/mod.rs @@ -50,33 +50,32 @@ use super::SpawnToken; /// A task's complete life cycle is as follows: /// /// ```text -/// ┌────────────┐ ┌────────────────────────┐ -/// ┌─►│Not spawned │◄─6┤Not spawned|Run enqueued│ -/// │ │ ├7─►│ │ -/// │ └─────┬──────┘ └──────▲─────────────────┘ -/// │ 1 │ -/// │ │ ┌────────────┘ -/// │ │ 5 -/// │ ┌─────▼────┴─────────┐ -/// │ │Spawned|Run enqueued│ -/// │ │ │ -/// │ └─────┬▲─────────────┘ -/// │ 2│ -/// │ │3 -/// │ ┌─────▼┴─────┐ -/// └─4┤ Spawned │ -/// │ │ -/// └────────────┘ +/// ┌────────────┐ ┌────────────────────────┐ +/// │Not spawned │◄─5┤Not spawned|Run enqueued│ +/// │ ├6─►│ │ +/// └─────┬──────┘ └──────▲─────────────────┘ +/// 1 │ +/// │ ┌────────────┘ +/// │ 4 +/// ┌─────▼────┴─────────┐ +/// │Spawned|Run enqueued│ +/// │ │ +/// └─────┬▲─────────────┘ +/// 2│ +/// │3 +/// ┌─────▼┴─────┐ +/// │ Spawned │ +/// │ │ +/// └────────────┘ /// ``` /// /// Transitions: /// - 1: Task is spawned - `AvailableTask::claim -> Executor::spawn` /// - 2: During poll - `RunQueue::dequeue_all -> State::run_dequeue` -/// - 3: Task wakes itself, waker wakes task - `Waker::wake -> wake_task -> State::run_enqueue` -/// - 4: Task exits - `TaskStorage::poll -> Poll::Ready` -/// - 5: A run-queued task exits - `TaskStorage::poll -> Poll::Ready` -/// - 6: Task is dequeued and then ignored via `State::run_dequeue` -/// - 7: A task is waken when it is not spawned - `wake_task -> State::run_enqueue` +/// - 3: Task wakes itself, waker wakes task, or task exits - `Waker::wake -> wake_task -> State::run_enqueue` +/// - 4: A run-queued task exits - `TaskStorage::poll -> Poll::Ready` +/// - 5: Task is dequeued. The task's future is not polled, because exiting the task replaces its `poll_fn`. +/// - 6: A task is waken when it is not spawned - `wake_task -> State::run_enqueue` pub(crate) struct TaskHeader { pub(crate) state: State, pub(crate) run_queue_item: RunQueueItem, From a011f487690465f8ae64fd74f4c51a8be3979890 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Tue, 17 Dec 2024 18:37:17 +0100 Subject: [PATCH 5/7] Make poll_to_despawn non-generic --- embassy-executor/src/raw/mod.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs index 5df5ca9e1..39d2d73ab 100644 --- a/embassy-executor/src/raw/mod.rs +++ b/embassy-executor/src/raw/mod.rs @@ -161,6 +161,12 @@ pub struct TaskStorage { future: UninitCell, // Valid if STATE_SPAWNED } +unsafe fn poll_to_despawn(p: TaskRef) { + // The task's future has already been dropped, we just mark it as `!SPAWNED`. + let this = p.header(); + this.state.despawn(); +} + impl TaskStorage { const NEW: Self = Self::new(); @@ -201,12 +207,6 @@ impl TaskStorage { } } - unsafe fn poll_to_despawn(p: TaskRef) { - // The task's future has already been dropped, we just mark it as `!SPAWNED`. - let this = &*p.as_ptr().cast::>(); - this.raw.state.despawn(); - } - unsafe fn poll(p: TaskRef) { let this = &*p.as_ptr().cast::>(); @@ -223,7 +223,7 @@ impl TaskStorage { // We replace the poll_fn with a despawn function, so that the task is cleaned up // when the executor polls it next. - this.raw.poll_fn.set(Some(Self::poll_to_despawn)); + this.raw.poll_fn.set(Some(poll_to_despawn)); } Poll::Pending => {} } From 2ca374fc9c0d0abe579716d1a7c2dc0724321ee7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Tue, 17 Dec 2024 18:46:32 +0100 Subject: [PATCH 6/7] Don't force a wake to despawn --- embassy-executor/src/raw/mod.rs | 6 ++++-- embassy-executor/tests/test.rs | 3 --- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs index 39d2d73ab..4a4ecf603 100644 --- a/embassy-executor/src/raw/mod.rs +++ b/embassy-executor/src/raw/mod.rs @@ -215,8 +215,6 @@ impl TaskStorage { let mut cx = Context::from_waker(&waker); match future.poll(&mut cx) { Poll::Ready(_) => { - waker.wake_by_ref(); - // As the future has finished and this function will not be called // again, we can safely drop the future here. this.future.drop_in_place(); @@ -224,6 +222,10 @@ impl TaskStorage { // We replace the poll_fn with a despawn function, so that the task is cleaned up // when the executor polls it next. this.raw.poll_fn.set(Some(poll_to_despawn)); + + // Make sure we despawn last, so that other threads can only spawn the task + // after we're done with it. + this.raw.state.despawn(); } Poll::Pending => {} } diff --git a/embassy-executor/tests/test.rs b/embassy-executor/tests/test.rs index d8c5a6ae3..78c49c071 100644 --- a/embassy-executor/tests/test.rs +++ b/embassy-executor/tests/test.rs @@ -69,7 +69,6 @@ fn executor_task() { &[ "pend", // spawning a task pends the executor "poll task1", // poll only once. - "pend", // task is done, wakes itself to exit ] ) } @@ -180,7 +179,6 @@ fn waking_after_completion_does_not_poll() { "pend", // manual wake, single pend for two wakes "pend", // respawning a task pends the executor "poll task1", // - "pend", // task is done, wakes itself to exit ] ) } @@ -268,7 +266,6 @@ fn waking_with_old_waker_after_respawn() { "yield_now", // "pend", // manual wake, gets cleared by poll "poll task1", // - "pend", // task is done, wakes itself to exit ] ); } From 76d8a896bbff612e3d2db27554891c71d28988af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Tue, 17 Dec 2024 18:51:22 +0100 Subject: [PATCH 7/7] Make poll_to_despawn a no_op --- embassy-executor/src/raw/mod.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs index 4a4ecf603..e38a2af66 100644 --- a/embassy-executor/src/raw/mod.rs +++ b/embassy-executor/src/raw/mod.rs @@ -161,10 +161,8 @@ pub struct TaskStorage { future: UninitCell, // Valid if STATE_SPAWNED } -unsafe fn poll_to_despawn(p: TaskRef) { - // The task's future has already been dropped, we just mark it as `!SPAWNED`. - let this = p.header(); - this.state.despawn(); +unsafe fn poll_exited(_p: TaskRef) { + // Nothing to do, the task is already !SPAWNED and dequeued. } impl TaskStorage { @@ -221,7 +219,7 @@ impl TaskStorage { // We replace the poll_fn with a despawn function, so that the task is cleaned up // when the executor polls it next. - this.raw.poll_fn.set(Some(poll_to_despawn)); + this.raw.poll_fn.set(Some(poll_exited)); // Make sure we despawn last, so that other threads can only spawn the task // after we're done with it.