From 78a695f21b8a34be8d1ed9adc667862a53bb85cf Mon Sep 17 00:00:00 2001
From: Zhang Junyang <junyang@stu.pku.edu.cn>
Date: Wed, 13 Nov 2024 21:28:58 +0800
Subject: [PATCH] Use `RcuOption` for lazy-initialized soft-IRQ callbacks

---
 kernel/src/time/clocks/system_wide.rs | 21 +++++++----
 kernel/src/time/softirq.rs            | 51 +++++++++++++++++++--------
 2 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/kernel/src/time/clocks/system_wide.rs b/kernel/src/time/clocks/system_wide.rs
index a757030b..fea81550 100644
--- a/kernel/src/time/clocks/system_wide.rs
+++ b/kernel/src/time/clocks/system_wide.rs
@@ -225,14 +225,18 @@ macro_rules! define_timer_managers {
         fn _init_system_wide_timer_managers() {
             $(
                 let clock = paste! {[<$clock_id _INSTANCE>].get().unwrap().clone()};
-                let clock_manager = TimerManager::new(clock);
+                let timer_manager = TimerManager::new(clock);
                 for cpu in ostd::cpu::all_cpus() {
                     paste! {
-                        [<$clock_id _MANAGER>].get_on_cpu(cpu).call_once(|| clock_manager.clone());
+                        [<$clock_id _MANAGER>].get_on_cpu(cpu).call_once(|| timer_manager.clone());
                     }
                 }
-                let callback = move || {
-                    clock_manager.process_expired_timers();
+                let callback = || {
+                    let preempt_guard = ostd::task::disable_preempt();
+                    let cpu = preempt_guard.current_cpu();
+                    paste! {
+                        [<$clock_id _MANAGER>].get_on_cpu(cpu).get().unwrap().process_expired_timers();
+                    }
                 };
                 time::softirq::register_callback(callback);
             )*
@@ -267,10 +271,13 @@ pub static JIFFIES_TIMER_MANAGER: Once<Arc<TimerManager>> = Once::new();
 fn init_jiffies_clock_manager() {
     let jiffies_clock = JiffiesClock { _private: () };
     let jiffies_timer_manager = TimerManager::new(Arc::new(jiffies_clock));
-    JIFFIES_TIMER_MANAGER.call_once(|| jiffies_timer_manager.clone());
+    JIFFIES_TIMER_MANAGER.call_once(|| jiffies_timer_manager);
 
-    let callback = move || {
-        jiffies_timer_manager.process_expired_timers();
+    let callback = || {
+        JIFFIES_TIMER_MANAGER
+            .get()
+            .unwrap()
+            .process_expired_timers();
     };
     time::softirq::register_callback(callback);
 }
diff --git a/kernel/src/time/softirq.rs b/kernel/src/time/softirq.rs
index a2c4bdfe..d75952b4 100644
--- a/kernel/src/time/softirq.rs
+++ b/kernel/src/time/softirq.rs
@@ -1,15 +1,13 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use alloc::{boxed::Box, vec::Vec};
+use alloc::{boxed::Box, vec, vec::Vec};
 
 use aster_softirq::{softirq_id::TIMER_SOFTIRQ_ID, SoftIrqLine};
-use ostd::{
-    sync::{LocalIrqDisabled, RwLock},
-    timer,
-};
+use ostd::{sync::RcuOption, timer};
 
-static TIMER_SOFTIRQ_CALLBACKS: RwLock<Vec<Box<dyn Fn() + Sync + Send>>, LocalIrqDisabled> =
-    RwLock::new(Vec::new());
+#[allow(clippy::type_complexity)]
+#[allow(clippy::box_collection)]
+static TIMER_SOFTIRQ_CALLBACKS: RcuOption<Box<Vec<fn()>>> = RcuOption::new_none();
 
 pub(super) fn init() {
     SoftIrqLine::get(TIMER_SOFTIRQ_ID).enable(timer_softirq_handler);
@@ -20,16 +18,41 @@ pub(super) fn init() {
 }
 
 /// Registers a function that will be executed during timer softirq.
-pub(super) fn register_callback<F>(func: F)
-where
-    F: Fn() + Sync + Send + 'static,
-{
-    TIMER_SOFTIRQ_CALLBACKS.write().push(Box::new(func));
+pub(super) fn register_callback(func: fn()) {
+    loop {
+        let callbacks = TIMER_SOFTIRQ_CALLBACKS.read();
+        match callbacks.get() {
+            // Initialized, copy the vector, push the function and update.
+            Some(callbacks_vec) => {
+                let mut callbacks_cloned = callbacks_vec.clone();
+                callbacks_cloned.push(func);
+                if callbacks
+                    .compare_exchange(Some(Box::new(callbacks_cloned)))
+                    .is_ok()
+                {
+                    break;
+                }
+            }
+            // Uninitialized, initialize it.
+            None => {
+                if callbacks
+                    .compare_exchange(Some(Box::new(vec![func])))
+                    .is_ok()
+                {
+                    break;
+                }
+            }
+        }
+        // Contention on initialization or pushing, retry.
+        core::hint::spin_loop();
+    }
 }
 
 fn timer_softirq_handler() {
     let callbacks = TIMER_SOFTIRQ_CALLBACKS.read();
-    for callback in callbacks.iter() {
-        (callback)();
+    if let Some(callbacks) = callbacks.get() {
+        for callback in callbacks.iter() {
+            (callback)();
+        }
     }
 }