5 files changed, 13 insertions, 57 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 5a476213b..997db6756 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -195,25 +195,7 @@ impl<F: Future + 'static> TaskStorage<F> {
         match future.poll(&mut cx) {
             Poll::Ready(_) => {
                 this.future.drop_in_place();
-
-                // Mark this task to be timer queued.
-                // We're splitting the enqueue in two parts, so that we can change task state
-                // to something that prevent re-queueing.
-                let op = this.raw.state.timer_enqueue();
-
-                // Now mark the task as not spawned, so that
-                // - it can be spawned again once it has been removed from the timer queue
-                // - it can not be timer-queued again
-                // We must do this before scheduling the wake, to prevent the task from being
-                // dequeued by the time driver while it's still SPAWNED.
                 this.raw.state.despawn();
-
-                // Now let's finish enqueueing. While we shouldn't get an `Ignore` here, it's
-                // better to be safe.
-                if op == timer_queue::TimerEnqueueOperation::Enqueue {
-                    // Schedule the task in the past, so it gets dequeued ASAP.
-                    unsafe { _embassy_time_schedule_wake(0, &waker) }
-                }
             }
             Poll::Pending => {}
         }
@@ -232,10 +214,6 @@ impl<F: Future + 'static> TaskStorage<F> {
     }
 }
 
-extern "Rust" {
-    fn _embassy_time_schedule_wake(at: u64, waker: &core::task::Waker);
-}
-
 /// An uninitialized [`TaskStorage`].
 pub struct AvailableTask<F: Future + 'static> {
     task: &'static TaskStorage<F>,
diff --git a/embassy-time-driver/src/lib.rs b/embassy-time-driver/src/lib.rs
index 57a9f7587..c776fbdf8 100644
--- a/embassy-time-driver/src/lib.rs
+++ b/embassy-time-driver/src/lib.rs
@@ -131,6 +131,7 @@ pub trait Driver: Send + Sync + 'static {
 
 extern "Rust" {
     fn _embassy_time_now() -> u64;
+    fn _embassy_time_schedule_wake(at: u64, waker: &Waker);
 }
 
 /// See [`Driver::now`]
@@ -138,6 +139,11 @@ pub fn now() -> u64 {
     unsafe { _embassy_time_now() }
 }
 
+/// Schedule the given waker to be woken at `at`.
+pub fn schedule_wake(at: u64, waker: &Waker) {
+    unsafe { _embassy_time_schedule_wake(at, waker) }
+}
+
 /// Set the time Driver implementation.
 ///
 /// See the module documentation for an example.
diff --git a/embassy-time-queue-driver/src/lib.rs b/embassy-time-queue-driver/src/lib.rs
index 97c81a124..333b6124d 100644
--- a/embassy-time-queue-driver/src/lib.rs
+++ b/embassy-time-queue-driver/src/lib.rs
@@ -10,8 +10,6 @@
 //! As a HAL implementer, you need to depend on this crate if you want to implement a time driver,
 //! but how you should do so is documented in `embassy-time-driver`.
 
-use core::task::Waker;
-
 #[cfg(feature = "_generic-queue")]
 pub mod queue_generic;
 #[cfg(not(feature = "_generic-queue"))]
@@ -21,29 +19,3 @@ pub mod queue_integrated;
 pub use queue_generic::Queue;
 #[cfg(not(feature = "_generic-queue"))]
 pub use queue_integrated::Queue;
-
-extern "Rust" {
-    fn _embassy_time_schedule_wake(at: u64, waker: &Waker);
-}
-
-/// Schedule the given waker to be woken at `at`.
-pub fn schedule_wake(at: u64, waker: &Waker) {
-    // This function is not implemented in embassy-time-driver because it needs access to executor
-    // internals. The function updates task state, then delegates to the implementation provided
-    // by the time driver.
-    #[cfg(not(feature = "_generic-queue"))]
-    {
-        use embassy_executor::raw::task_from_waker;
-        use embassy_executor::raw::timer_queue::TimerEnqueueOperation;
-        // The very first thing we must do, before we even access the timer queue, is to
-        // mark the task a TIMER_QUEUED. This ensures that the task that is being scheduled
-        // can not be respawn while we are accessing the timer queue.
-        let task = task_from_waker(waker);
-        if unsafe { task.timer_enqueue() } == TimerEnqueueOperation::Ignore {
-            // We are not allowed to enqueue the task in the timer queue. This is because the
-            // task is not spawned, and so it makes no sense to schedule it.
-            return;
-        }
-    }
-    unsafe { _embassy_time_schedule_wake(at, waker) }
-}
diff --git a/embassy-time/Cargo.toml b/embassy-time/Cargo.toml
index e3074119f..4f4ea0b14 100644
--- a/embassy-time/Cargo.toml
+++ b/embassy-time/Cargo.toml
@@ -24,8 +24,8 @@ target = "x86_64-unknown-linux-gnu"
 features = ["defmt", "std"]
 
 [features]
-std = ["tick-hz-1_000_000", "critical-section/std"]
-wasm = ["dep:wasm-bindgen", "dep:js-sys", "dep:wasm-timer", "tick-hz-1_000_000"]
+std = ["tick-hz-1_000_000", "critical-section/std", "dep:embassy-time-queue-driver"]
+wasm = ["dep:wasm-bindgen", "dep:js-sys", "dep:wasm-timer", "tick-hz-1_000_000", "dep:embassy-time-queue-driver"]
 
 ## Display the time since startup next to defmt log messages.
 ## At most 1 `defmt-timestamp-uptime-*` feature can be used.
@@ -40,7 +40,7 @@ defmt-timestamp-uptime-tms = ["defmt"]
 defmt-timestamp-uptime-tus = ["defmt"]
 
 ## Create a `MockDriver` that can be manually advanced for testing purposes.
-mock-driver = ["tick-hz-1_000_000"]
+mock-driver = ["tick-hz-1_000_000", "dep:embassy-time-queue-driver"]
 
 #! ### Tick Rate
 #!
@@ -384,7 +384,7 @@ tick-hz-5_242_880_000 = ["embassy-time-driver/tick-hz-5_242_880_000"]
 
 [dependencies]
 embassy-time-driver = { version = "0.1.0", path = "../embassy-time-driver" }
-embassy-time-queue-driver = { version = "0.1.0", path = "../embassy-time-queue-driver" }
+embassy-time-queue-driver = { version = "0.1.0", path = "../embassy-time-queue-driver", optional = true}
 
 defmt = { version = "0.3", optional = true }
 log = { version = "0.4.14", optional = true }
diff --git a/embassy-time/src/timer.rs b/embassy-time/src/timer.rs
index 4d7194b20..295ddbd9b 100644
--- a/embassy-time/src/timer.rs
+++ b/embassy-time/src/timer.rs
@@ -157,7 +157,7 @@ impl Future for Timer {
         if self.yielded_once && self.expires_at <= Instant::now() {
             Poll::Ready(())
         } else {
-            embassy_time_queue_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
+            embassy_time_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
             self.yielded_once = true;
             Poll::Pending
         }
@@ -238,7 +238,7 @@ impl Ticker {
                 self.expires_at += dur;
                 Poll::Ready(())
             } else {
-                embassy_time_queue_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
+                embassy_time_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
                 Poll::Pending
             }
         })
@@ -255,7 +255,7 @@ impl Stream for Ticker {
             self.expires_at += dur;
             Poll::Ready(Some(()))
         } else {
-            embassy_time_queue_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
+            embassy_time_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
             Poll::Pending
         }
     }