From c335873d12574c450fdc523cc30a0d6a1e23d1cd Mon Sep 17 00:00:00 2001
From: Fabing Li <lifabing.lfb@antgroup.com>
Date: Thu, 18 Apr 2024 11:00:36 +0800
Subject: [PATCH] Enable kernel mode unit tests in kernel threads

---
 framework/aster-frame/src/boot/mod.rs         | 15 ++++++-
 framework/aster-frame/src/task/mod.rs         |  2 +-
 framework/aster-frame/src/task/scheduler.rs   | 39 +++++++++++++++++++
 framework/aster-frame/src/task/task.rs        | 23 +++++++++++
 .../aster-nix/src/sched/priority_scheduler.rs |  8 +---
 5 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/framework/aster-frame/src/boot/mod.rs b/framework/aster-frame/src/boot/mod.rs
index 353822117..2f1c5a01c 100644
--- a/framework/aster-frame/src/boot/mod.rs
+++ b/framework/aster-frame/src/boot/mod.rs
@@ -119,13 +119,26 @@ pub fn call_aster_main() -> ! {
     }
     #[cfg(ktest)]
     unsafe {
+        use alloc::boxed::Box;
+
+        use crate::task::{set_scheduler, FifoScheduler, Scheduler, TaskOptions};
+
         crate::init();
         // The whitelists that will be generated by OSDK runner as static consts.
         extern "Rust" {
             static KTEST_TEST_WHITELIST: Option<&'static [&'static str]>;
             static KTEST_CRATE_WHITELIST: Option<&'static [&'static str]>;
         }
-        run_ktests(KTEST_TEST_WHITELIST, KTEST_CRATE_WHITELIST);
+        // Set the global scheduler a FIFO scheduler.
+        let simple_scheduler = Box::new(FifoScheduler::new());
+        let static_scheduler: &'static dyn Scheduler = Box::leak(simple_scheduler);
+        set_scheduler(static_scheduler);
+
+        let test_task = move || {
+            run_ktests(KTEST_TEST_WHITELIST, KTEST_CRATE_WHITELIST);
+        };
+        let _ = TaskOptions::new(test_task).data(()).spawn();
+        unreachable!("The spawn method will NOT return in the boot context")
     }
 }
 
diff --git a/framework/aster-frame/src/task/mod.rs b/framework/aster-frame/src/task/mod.rs
index f986e2234..7ed565935 100644
--- a/framework/aster-frame/src/task/mod.rs
+++ b/framework/aster-frame/src/task/mod.rs
@@ -11,6 +11,6 @@ mod task;
 pub use self::{
     priority::Priority,
     processor::{current_task, disable_preempt, preempt, schedule, DisablePreemptGuard},
-    scheduler::{add_task, set_scheduler, Scheduler},
+    scheduler::{add_task, set_scheduler, FifoScheduler, Scheduler},
     task::{Task, TaskAdapter, TaskOptions, TaskStatus},
 };
diff --git a/framework/aster-frame/src/task/scheduler.rs b/framework/aster-frame/src/task/scheduler.rs
index c9518b13e..a7e515ce4 100644
--- a/framework/aster-frame/src/task/scheduler.rs
+++ b/framework/aster-frame/src/task/scheduler.rs
@@ -1,5 +1,7 @@
 // SPDX-License-Identifier: MPL-2.0
 
+use alloc::collections::VecDeque;
+
 use lazy_static::lazy_static;
 
 use crate::{prelude::*, sync::SpinLock, task::Task};
@@ -60,3 +62,40 @@ pub fn fetch_task() -> Option<Arc<Task>> {
 pub fn add_task(task: Arc<Task>) {
     GLOBAL_SCHEDULER.lock_irq_disabled().enqueue(task);
 }
+
+/// A simple FIFO (First-In-First-Out) task scheduler.
+pub struct FifoScheduler {
+    /// A thread-safe queue to hold tasks waiting to be executed.
+    task_queue: SpinLock<VecDeque<Arc<Task>>>,
+}
+
+impl FifoScheduler {
+    /// Creates a new instance of `FifoScheduler`.
+    pub const fn new() -> Self {
+        FifoScheduler {
+            task_queue: SpinLock::new(VecDeque::new()),
+        }
+    }
+}
+
+impl Default for FifoScheduler {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Scheduler for FifoScheduler {
+    /// Enqueues a task to the end of the queue.
+    fn enqueue(&self, task: Arc<Task>) {
+        self.task_queue.lock_irq_disabled().push_back(task);
+    }
+    /// Dequeues a task from the front of the queue, if any.
+    fn dequeue(&self) -> Option<Arc<Task>> {
+        self.task_queue.lock_irq_disabled().pop_front()
+    }
+    /// In this simple implementation, task preemption is not supported.
+    /// Once a task starts running, it will continue to run until completion.
+    fn should_preempt(&self, _task: &Arc<Task>) -> bool {
+        false
+    }
+}
diff --git a/framework/aster-frame/src/task/task.rs b/framework/aster-frame/src/task/task.rs
index 1171858ce..8e7891068 100644
--- a/framework/aster-frame/src/task/task.rs
+++ b/framework/aster-frame/src/task/task.rs
@@ -303,3 +303,26 @@ impl TaskOptions {
         Ok(task)
     }
 }
+
+#[cfg(ktest)]
+mod test {
+    #[ktest]
+    fn create_task() {
+        let task = || {
+            assert_eq!(1, 1);
+        };
+        let task_option = crate::task::TaskOptions::new(task)
+            .data(())
+            .build()
+            .unwrap();
+        task_option.run();
+    }
+
+    #[ktest]
+    fn spawn_task() {
+        let task = || {
+            assert_eq!(1, 1);
+        };
+        let _ = crate::task::TaskOptions::new(task).data(()).spawn();
+    }
+}
diff --git a/kernel/aster-nix/src/sched/priority_scheduler.rs b/kernel/aster-nix/src/sched/priority_scheduler.rs
index 2c4479cad..1ee3628da 100644
--- a/kernel/aster-nix/src/sched/priority_scheduler.rs
+++ b/kernel/aster-nix/src/sched/priority_scheduler.rs
@@ -36,13 +36,9 @@ impl PreemptScheduler {
 impl Scheduler for PreemptScheduler {
     fn enqueue(&self, task: Arc<Task>) {
         if task.is_real_time() {
-            self.real_time_tasks
-                .lock_irq_disabled()
-                .push_back(task.clone());
+            self.real_time_tasks.lock_irq_disabled().push_back(task);
         } else {
-            self.normal_tasks
-                .lock_irq_disabled()
-                .push_back(task.clone());
+            self.normal_tasks.lock_irq_disabled().push_back(task);
         }
     }