Register some timer callbacks on all CPUs

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();
 }
 

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);
 }

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) {

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);
 }
 

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},
 };

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.

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);