From 714e5b571f2f6b329dfd830b9d3986fac7e7a0c8 Mon Sep 17 00:00:00 2001
From: Qi <add_sp@outlook.com>
Date: Mon, 9 Jun 2025 15:32:59 +0800
Subject: [PATCH] runtime: move ` impl Schedule for Arc<Handle>` (#7398)

---
 .../runtime/scheduler/multi_thread/handle.rs  | 21 +++++++++++++++
 .../runtime/scheduler/multi_thread/worker.rs  | 26 +------------------
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs
index 4075713c9..47299659f 100644
--- a/tokio/src/runtime/scheduler/multi_thread/handle.rs
+++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs
@@ -1,6 +1,7 @@
 use crate::future::Future;
 use crate::loom::sync::Arc;
 use crate::runtime::scheduler::multi_thread::worker;
+use crate::runtime::task::{Notified, Task, TaskHarnessScheduleHooks};
 use crate::runtime::{
     blocking, driver,
     task::{self, JoinHandle},
@@ -66,6 +67,26 @@ impl Handle {
     }
 }
 
+impl task::Schedule for Arc<Handle> {
+    fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
+        self.shared.owned.remove(task)
+    }
+
+    fn schedule(&self, task: Notified<Self>) {
+        self.schedule_task(task, false);
+    }
+
+    fn hooks(&self) -> TaskHarnessScheduleHooks {
+        TaskHarnessScheduleHooks {
+            task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
+        }
+    }
+
+    fn yield_now(&self, task: Notified<Self>) {
+        self.schedule_task(task, true);
+    }
+}
+
 cfg_unstable! {
     use std::num::NonZeroU64;
 
diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs
index 73f033720..838c15840 100644
--- a/tokio/src/runtime/scheduler/multi_thread/worker.rs
+++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs
@@ -62,7 +62,7 @@ use crate::runtime::scheduler::multi_thread::{
     idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker,
 };
 use crate::runtime::scheduler::{inject, Defer, Lock};
-use crate::runtime::task::{OwnedTasks, TaskHarnessScheduleHooks};
+use crate::runtime::task::OwnedTasks;
 use crate::runtime::{blocking, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics};
 use crate::runtime::{context, TaskHooks};
 use crate::task::coop;
@@ -225,9 +225,6 @@ pub(crate) struct Launch(Vec<Arc<Worker>>);
 /// to stop processing.
 type RunResult = Result<Box<Core>, ()>;
 
-/// A task handle
-type Task = task::Task<Arc<Handle>>;
-
 /// A notified task handle
 type Notified = task::Notified<Arc<Handle>>;
 
@@ -1048,27 +1045,6 @@ impl Worker {
     }
 }
 
-// TODO: Move `Handle` impls into handle.rs
-impl task::Schedule for Arc<Handle> {
-    fn release(&self, task: &Task) -> Option<Task> {
-        self.shared.owned.remove(task)
-    }
-
-    fn schedule(&self, task: Notified) {
-        self.schedule_task(task, false);
-    }
-
-    fn hooks(&self) -> TaskHarnessScheduleHooks {
-        TaskHarnessScheduleHooks {
-            task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
-        }
-    }
-
-    fn yield_now(&self, task: Notified) {
-        self.schedule_task(task, true);
-    }
-}
-
 impl Handle {
     pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) {
         with_current(|maybe_cx| {