From a274d68cc11e8bdcf85fac36be46fdb4c13c658c Mon Sep 17 00:00:00 2001
From: Weijia Jiang <wjjiang19@163.com>
Date: Sun, 10 Sep 2023 09:47:25 +0800
Subject: [PATCH] fix unsafe get_scheduler

---
 tokio/src/runtime/scheduler/mod.rs | 10 ++++++----
 tokio/src/runtime/task/core.rs     |  4 ++--
 tokio/src/runtime/task/raw.rs      |  2 +-
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs
index 6548564e214..c7fa17260c6 100644
--- 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),
@@ -168,7 +170,7 @@ cfg_rt! {
                 Handle::MultiThreadAlt(h) => h.yield_now(task),
             }
         }
-
+            
         pub(crate) fn shutdown(&self) {
             match *self {
                 Handle::CurrentThread(_) => {},
diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs
index de0c47bdb4f..45315c07926 100644
--- 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
index 76ee1f323ea..f62f4abcf91 100644
--- 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),