From 2e46edb68de2de5db1d5a0ec097978ba8afe76ef Mon Sep 17 00:00:00 2001
From: Zejun Zhao <jelly.zhao.42@gmail.com>
Date: Thu, 18 Sep 2025 13:38:30 +0800
Subject: [PATCH] Register some timer callbacks on all CPUs

---
 kernel/src/lib.rs                           |  2 ++
 kernel/src/process/mod.rs                   |  5 ++++-
 kernel/src/process/process/mod.rs           |  4 ++--
 kernel/src/process/process/timer_manager.rs |  2 +-
 kernel/src/sched/mod.rs                     |  4 +++-
 kernel/src/sched/sched_class/mod.rs         |  8 ++++++--
 ostd/src/task/scheduler/mod.rs              | 13 +++++++++++++
 7 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/kernel/src/lib.rs b/kernel/src/lib.rs
index 4bb3b7335..54433f85e 100644
--- a/kernel/src/lib.rs
+++ b/kernel/src/lib.rs
@@ -110,6 +110,8 @@ fn init() {
 }
 
 fn init_on_each_cpu() {
+    sched::init_on_each_cpu();
+    process::init_on_each_cpu();
     fs::init_on_each_cpu();
 }
 
diff --git a/kernel/src/process/mod.rs b/kernel/src/process/mod.rs
index 72cbc935c..54f872cd1 100644
--- a/kernel/src/process/mod.rs
+++ b/kernel/src/process/mod.rs
@@ -44,10 +44,13 @@ pub use wait::{do_wait, WaitOptions, WaitStatus};
 use crate::context::Context;
 
 pub(super) fn init() {
-    process::init();
     posix_thread::futex::init();
 }
 
+pub(super) fn init_on_each_cpu() {
+    process::init_on_each_cpu();
+}
+
 pub(super) fn init_in_first_process(ctx: &Context) {
     process::init_in_first_process(ctx);
 }
diff --git a/kernel/src/process/process/mod.rs b/kernel/src/process/process/mod.rs
index 99b1b96e1..9ebe70923 100644
--- a/kernel/src/process/process/mod.rs
+++ b/kernel/src/process/process/mod.rs
@@ -55,8 +55,8 @@ pub type Sid = u32;
 
 pub type ExitCode = u32;
 
-pub(super) fn init() {
-    timer_manager::init();
+pub(super) fn init_on_each_cpu() {
+    timer_manager::init_on_each_cpu();
 }
 
 pub(super) fn init_in_first_process(ctx: &Context) {
diff --git a/kernel/src/process/process/timer_manager.rs b/kernel/src/process/process/timer_manager.rs
index 17475988e..440f69941 100644
--- a/kernel/src/process/process/timer_manager.rs
+++ b/kernel/src/process/process/timer_manager.rs
@@ -66,7 +66,7 @@ fn update_cpu_time() {
 
 /// Registers a function to update the CPU clock in processes and
 /// threads during the system timer interrupt.
-pub(super) fn init() {
+pub(super) fn init_on_each_cpu() {
     timer::register_callback(update_cpu_time);
 }
 
diff --git a/kernel/src/sched/mod.rs b/kernel/src/sched/mod.rs
index 3f8842d04..db5019387 100644
--- a/kernel/src/sched/mod.rs
+++ b/kernel/src/sched/mod.rs
@@ -6,6 +6,8 @@ mod stats;
 
 pub use self::{
     nice::{AtomicNice, Nice},
-    sched_class::{init, RealTimePolicy, RealTimePriority, SchedAttr, SchedPolicy},
+    sched_class::{
+        init, init_on_each_cpu, RealTimePolicy, RealTimePriority, SchedAttr, SchedPolicy,
+    },
     stats::{loadavg, nr_queued_and_running},
 };
diff --git a/kernel/src/sched/sched_class/mod.rs b/kernel/src/sched/sched_class/mod.rs
index 98af15c54..f2375941d 100644
--- a/kernel/src/sched/sched_class/mod.rs
+++ b/kernel/src/sched/sched_class/mod.rs
@@ -13,8 +13,8 @@ use ostd::{
     sync::SpinLock,
     task::{
         scheduler::{
-            info::CommonSchedInfo, inject_scheduler, EnqueueFlags, LocalRunQueue, Scheduler,
-            UpdateFlags,
+            enable_preemption, info::CommonSchedInfo, inject_scheduler, EnqueueFlags,
+            LocalRunQueue, Scheduler, UpdateFlags,
         },
         AtomicCpuId, Task,
     },
@@ -55,6 +55,10 @@ pub fn init() {
     set_stats_from_scheduler(scheduler);
 }
 
+pub fn init_on_each_cpu() {
+    enable_preemption();
+}
+
 /// Represents the middle layer between scheduling classes and generic scheduler
 /// traits. It consists of all the sets of run queues for CPU cores. Other global
 /// information may also be stored here.
diff --git a/ostd/src/task/scheduler/mod.rs b/ostd/src/task/scheduler/mod.rs
index a734a6e54..fb6965f00 100644
--- a/ostd/src/task/scheduler/mod.rs
+++ b/ostd/src/task/scheduler/mod.rs
@@ -84,7 +84,20 @@ use crate::{
 /// before any [`Task`]-related APIs are invoked.
 pub fn inject_scheduler(scheduler: &'static dyn Scheduler<Task>) {
     SCHEDULER.call_once(|| scheduler);
+}
 
+/// Enables preemptive scheduling.
+///
+/// After calling this function on a CPU,
+/// a task that is executing in the user mode may get preempted
+/// if another runnable task on the same CPU is deemed more urgent by the scheduler.
+///
+/// OSTD achieves task preemption by registering a per-CPU timer callback
+/// to invoke the scheduler periodically.
+/// Thus, this function should be called _once_ on every CPU
+/// by an OSTD-based kernel during its initialization phase,
+/// after it has injected its scheduler via [`inject_scheduler`].
+pub fn enable_preemption() {
     timer::register_callback(|| {
         SCHEDULER.get().unwrap().mut_local_rq_with(&mut |local_rq| {
             let should_pick_next = local_rq.update_current(UpdateFlags::Tick);