fix unsafe get_scheduler

tokio-rs · Sep 10, 2023 · 4326c7a · 4326c7a
1 parent 0126a7c
commit 4326c7a
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs
@@ -27,7 +27,8 @@ cfg_rt_multi_thread! {
 
 use crate::runtime::driver;
 
-use super::task::Schedule;
+#[cfg(feature = "rt")]
+use crate::runtime::task::Schedule;
 
 #[derive(Debug, Clone)]
 pub(crate) enum Handle {
@@ -130,7 +131,7 @@ cfg_rt! {
                 Handle::MultiThreadAlt(h) => multi_thread_alt::Handle::spawn(h, future, id),
             }
         }
-
+        #[cfg(feature = "rt")]
         pub(crate) fn release(&self, task: &Task) -> Option<Task>
         {
             match self {
@@ -144,6 +145,7 @@ cfg_rt! {
             }
         }
 
+        #[cfg(feature = "rt")]
         pub(crate) fn schedule(&self, task: Notified)
         {
             match self {
@@ -156,7 +158,7 @@ cfg_rt! {
                 Handle::MultiThreadAlt(h) => h.schedule(task),
             }
         }
-
+        #[cfg(feature = "rt")]
         pub(crate) fn yield_now(&self, task: Notified) {
             match self {
                 Handle::CurrentThread(h) => h.yield_now(task),

diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs
@@ -431,9 +431,9 @@ impl Header {
     ///
     /// The generic type S must be set to the correct scheduler type for this
     /// task.
-    pub(super) unsafe fn get_scheduler<S>(me: NonNull<Header>) -> NonNull<S> {
+    pub(super) unsafe fn get_scheduler<S>(me: NonNull<Header>) -> NonNull<Option<S>> {
         let offset = me.as_ref().vtable.scheduler_offset;
-        let scheduler = me.as_ptr().cast::<u8>().add(offset).cast::<S>();
+        let scheduler = me.as_ptr().cast::<u8>().add(offset).cast::<Option<S>>();
         NonNull::new_unchecked(scheduler)
     }
 

diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs
@@ -279,7 +279,7 @@ unsafe fn poll<T: Future, S: Schedule>(ptr: NonNull<Header>) {
 unsafe fn schedule<S: Schedule>(ptr: NonNull<Header>) {
     use crate::runtime::task::{Notified, Task};
 
-    let scheduler = Header::get_scheduler::<Option<S>>(ptr).as_ref();
+    let scheduler = Header::get_scheduler::<S>(ptr).as_ref();
     let notify = Notified(Task::from_raw(ptr.cast()));
     match scheduler {
         Some(scheduler) => scheduler.schedule(notify),