Add stage support for init_component macro

kernel/libs/comp-sys/component-macro/src/lib.rs

@@ -17,13 +17,17 @@ pub(crate) const COMPONENT_FILE_NAME: &str = "Components.toml";
 
 /// Register a function to be called when the component system is initialized. The function should not public.
 ///
+/// You can specify the initialization stage by:
+/// - `#[init_component]` or `#[init_component(bootstrap)]` - the **Bootstrap** stage
+/// - `#[init_component(kthread)]` - the **Kthread** stage
+/// - `#[init_component(process)]` - the **Process** stage
+///
 /// Example:
 /// ```rust
 /// #[init_component]
 /// fn init() -> Result<(), component::ComponentInitError> {
 ///     Ok(())
 /// }
-///
 /// ```
 ///
 /// It will expand to
@@ -32,26 +36,28 @@ pub(crate) const COMPONENT_FILE_NAME: &str = "Components.toml";
 ///     Ok(())
 /// }
 ///
-/// const fn file() -> &'static str{
-///     file!()
-/// }
-///
-/// component::submit!(component::ComponentRegistry::new(&init,file()));
+/// component::submit!(component::ComponentRegistry::new(component::InitStage::Bootstrap, &init, file!()));
 /// ```
 /// The priority will calculate automatically
 ///
 #[proc_macro_attribute]
-pub fn init_component(_: TokenStream, input: TokenStream) -> proc_macro::TokenStream {
+pub fn init_component(args: TokenStream, input: TokenStream) -> proc_macro::TokenStream {
+    let stage = match args.to_string().as_str() {
+        "" | "bootstrap" => quote! { Bootstrap },
+        "kthread" => quote! { Kthread },
+        "process" => quote! { Process },
+        _ => panic!("Invalid argument for init_component"),
+    };
     let function = parse_macro_input!(input as ComponentInitFunction);
     let function_name = &function.function_name;
     quote! {
         #function
 
-        const fn file() -> &'static str{
-            file!()
-        }
-
-        component::submit!(component::ComponentRegistry::new(&#function_name,file()));
+        component::submit!(component::ComponentRegistry::new(
+            component::InitStage::#stage,
+            &#function_name,
+            file!())
+        );
     }
     .into()
 }
@@ -65,7 +71,7 @@ pub fn init_component(_: TokenStream, input: TokenStream) -> proc_macro::TokenSt
 /// Example:
 ///
 /// ```rust
-///     component::init_all(component::parse_metadata!());
+///     component::init_all(component::InitStage::Bootstrap, component::parse_metadata!());
 /// ```
 ///
 #[proc_macro]

kernel/libs/comp-sys/component/README.md

@@ -3,11 +3,24 @@
 ## Overview
 This crate is used for the initialization of the component system, which provides a priority initialization scheme based on the inventory crate.
 
+## Initialization Stages
+The component system supports three distinct initialization stages that execute in a specific order during system boot:
+1. **Bootstrap**: The earliest stage, called after OSTD initialization is complete but before kernel subsystem initialization begins. This stage runs on the BSP (Bootstrap Processor) only, before SMP (Symmetric Multi-Processing) is enabled. Components in this stage can initialize core kernel services that other components depend on.
+2. **Kthread**: The kernel thread stage, initialized after SMP is enabled and the first kernel thread is spawned. This stage runs in the context of the first kernel thread on the BSP.
+3. **Process**: The process stage, initialized after the first user process is created. This stage runs in the context of the first user process, and prepares the system for user-space execution.
+
 ## Usage
 
 ### Register component
 
-Registering a crate as component by marking a function in the lib.rs with `#[init_component]` macro. The specific definition of the function can refer to the comments in the macro.
+Registering a crate as component by marking a function in the lib.rs with `#[init_component]` macro. You can specify the initialization stage for your component:
+- `#[init_component]` or `#[init_component(bootstrap)]` - registers for the **Bootstrap** stage
+- `#[init_component(kthread)]` - registers for the **Kthread** stage
+- `#[init_component(process)]` - registers for the **Process** stage
+
+**Note**: Each crate can declare up to three initialization functions (one for each stage), but each function can only be associated with one specific stage. A given stage can only be declared once per crate.
+
+The specific definition of the function can refer to the comments in the macro.
 
 ### Component initialization
 
@@ -34,7 +47,7 @@ fn comp1_init() -> Result<(), component::ComponentInitError> {
 // src/main.rs
 use std::sync::atomic::Ordering::Relaxed;
 
-use component::init_component;
+use component::{init_component, InitStage};
 use comp1::INIT_COUNT;
 
 #[init_component]
@@ -45,7 +58,7 @@ fn init() -> Result<(), component::ComponentInitError> {
 }
 
 fn main(){
-    component::init_all(component::parse_metadata!()).unwrap();
+    component::init_all(InitStage::Bootstrap, component::parse_metadata!()).unwrap();
     assert_eq!(INIT_COUNT.load(Relaxed),2);
 }
 ```

kernel/libs/comp-sys/component/src/lib.rs

@@ -21,6 +21,26 @@ pub use component_macro::*;
 pub use inventory::submit;
 use log::{debug, error, info};
 
+/// The initialization stages of the component system.
+///
+/// - `Bootstrap`: The earliest stage, called after OSTD initialization is
+///   complete but before kernel subsystem initialization begins. This stage
+///   runs on the BSP (Bootstrap Processor) only, before SMP (Symmetric
+///   Multi-Processing) is enabled. Components in this stage can initialize
+///   core kernel services that other components depend on.
+/// - `Kthread`: The kernel thread stage, initialized after SMP is enabled
+///   and the first kernel thread is spawned. This stage runs in the context
+///   of the first kernel thread on the BSP.
+/// - `Process`: The process stage, initialized after the first user process
+///   is created. This stage runs in the context of the first user process,
+///   and prepares the system for user-space execution.
+#[derive(Debug, PartialEq, Eq)]
+pub enum InitStage {
+    Bootstrap,
+    Kthread,
+    Process,
+}
+
 #[derive(Debug)]
 pub enum ComponentInitError {
     UninitializedDependencies(String),
@@ -28,16 +48,22 @@ pub enum ComponentInitError {
 }
 
 pub struct ComponentRegistry {
+    stage: InitStage,
     function: &'static (dyn Fn() -> Result<(), ComponentInitError> + Sync),
     path: &'static str,
 }
 
 impl ComponentRegistry {
     pub const fn new(
+        stage: InitStage,
         function: &'static (dyn Fn() -> Result<(), ComponentInitError> + Sync),
         path: &'static str,
     ) -> Self {
-        Self { function, path }
+        Self {
+            stage,
+            function,
+            path,
+        }
     }
 }
 
@@ -46,6 +72,7 @@ inventory::collect!(ComponentRegistry);
 impl Debug for ComponentRegistry {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_struct("ComponentRegistry")
+            .field("stage", &self.stage)
             .field("path", &self.path)
             .finish()
     }
@@ -105,17 +132,24 @@ pub enum ComponentSystemInitError {
     NotIncludeAllComponent(String),
 }
 
-/// Component system initialization. It will collect invoke all functions that are marked by init_component based on dependencies between crates.
+/// Initializes the component system for a specific stage.
+///
+/// It collects all functions marked with the `init_component` macro, filters them
+/// according to the given stage, and invokes them in the correct order while honoring
+/// dependencies and priorities between crates.
 ///
 /// The collection of ComponentInfo usually generate by `parse_metadata` macro.
 ///
 /// ```rust
-///     component::init_all(component::parse_metadata!());
+///     component::init_all(component::InitStage::Bootstrap, component::parse_metadata!());
 /// ```
 ///
-pub fn init_all(components: Vec<ComponentInfo>) -> Result<(), ComponentSystemInitError> {
+pub fn init_all(
+    stage: InitStage,
+    components: Vec<ComponentInfo>,
+) -> Result<(), ComponentSystemInitError> {
     let components_info = parse_input(components);
-    match_and_call(components_info)?;
+    match_and_call(stage, components_info)?;
     Ok(())
 }
 
@@ -130,10 +164,15 @@ fn parse_input(components: Vec<ComponentInfo>) -> BTreeMap<String, ComponentInfo
 
 /// Match the ComponentInfo with ComponentRegistry. The key is the relative path of one component
 fn match_and_call(
+    stage: InitStage,
     mut components: BTreeMap<String, ComponentInfo>,
 ) -> Result<(), ComponentSystemInitError> {
     let mut infos = Vec::new();
     for registry in inventory::iter::<ComponentRegistry> {
+        if registry.stage != stage {
+            continue;
+        }
+
         // relative/path/to/comps/pci/src/lib.rs
         let mut str: String = registry.path.to_owned();
         str = str.replace('\\', "/");
@@ -169,7 +208,7 @@ fn match_and_call(
 
     infos.sort();
     debug!("component infos: {infos:?}");
-    info!("Components initializing...");
+    info!("Components initializing in {stage:?} stage...");
 
     for i in infos {
         info!("Component initializing:{:?}", i);
@@ -179,6 +218,6 @@ fn match_and_call(
             info!("Component initialize complete");
         }
     }
-    info!("All components initialization completed");
+    info!("All components initialization in {stage:?} stage completed");
     Ok(())
 }

kernel/libs/comp-sys/component/tests/init-order/second-init/tests/test.rs

@@ -6,6 +6,6 @@ use std::sync::atomic::Ordering::Relaxed;
 #[test]
 fn test() {
     simple_logger::init_with_level(log::Level::Debug).unwrap();
-    component::init_all(component::parse_metadata!()).unwrap();
+    component::init_all(component::InitStage::Bootstrap, component::parse_metadata!()).unwrap();
     assert_eq!(HAS_INIT.load(Relaxed), true);
 }

kernel/libs/comp-sys/component/tests/init-order/src/main.rs

@@ -17,7 +17,7 @@ fn kernel_init() -> Result<(), component::ComponentInitError> {
 
 fn main() {
     simple_logger::init_with_level(log::Level::Info).unwrap();
-    component::init_all(component::parse_metadata!()).unwrap();
+    component::init_all(component::InitStage::Bootstrap, component::parse_metadata!()).unwrap();
     assert_eq!(first_init::HAS_INIT.load(Relaxed), true);
     assert_eq!(second_init::HAS_INIT.load(Relaxed), true);
     assert_eq!(HAS_INIT.load(Relaxed), true);

kernel/libs/comp-sys/component/tests/init-order/tests/test.rs

@@ -12,6 +12,6 @@ fn kernel_init() -> Result<(), component::ComponentInitError> {
 #[test]
 fn test() {
     simple_logger::init_with_level(log::Level::Debug).unwrap();
-    component::init_all(component::parse_metadata!()).unwrap();
+    component::init_all(component::InitStage::Bootstrap, component::parse_metadata!()).unwrap();
     assert_eq!(HAS_INIT.load(Relaxed), true);
 }

kernel/libs/comp-sys/component/tests/init-stage/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "init-stage"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+component = { path = "../../../component" }
+foo = { path = "foo" }
+simple_logger = "4.0.0"
+log = "0.4"
+
+[workspace]
+members = [
+    "foo",
+]

kernel/libs/comp-sys/component/tests/init-stage/Components.toml

@@ -0,0 +1,4 @@
+# template
+[components]
+init-stage = { name = "init-stage" }
+foo = { name = "foo" }

kernel/libs/comp-sys/component/tests/init-stage/foo/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "foo"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+component = { path = "../../../../component" }

kernel/libs/comp-sys/component/tests/init-stage/foo/src/lib.rs

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
+
+use component::{init_component, InitStage};
+
+pub static INIT_BOOTSTRAP: AtomicBool = AtomicBool::new(false);
+pub static INIT_KTHREAD: AtomicBool = AtomicBool::new(false);
+pub static INIT_PROCESS: AtomicBool = AtomicBool::new(false);
+
+#[init_component]
+fn bootstrap() -> Result<(), component::ComponentInitError> {
+    assert_eq!(INIT_BOOTSTRAP.load(Relaxed), false);
+    INIT_BOOTSTRAP.store(true, Relaxed);
+    Ok(())
+}
+
+#[init_component(kthread)]
+fn kthread() -> Result<(), component::ComponentInitError> {
+    assert_eq!(INIT_KTHREAD.load(Relaxed), false);
+    INIT_KTHREAD.store(true, Relaxed);
+    Ok(())
+}
+
+#[init_component(process)]
+fn process() -> Result<(), component::ComponentInitError> {
+    assert_eq!(INIT_PROCESS.load(Relaxed), false);
+    INIT_PROCESS.store(true, Relaxed);
+    Ok(())
+}

kernel/libs/comp-sys/component/tests/init-stage/src/main.rs

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
+
+use component::{init_component, InitStage};
+use foo::{INIT_BOOTSTRAP, INIT_KTHREAD, INIT_PROCESS};
+
+fn main() {
+    simple_logger::init_with_level(log::Level::Debug).unwrap();
+    component::init_all(component::InitStage::Bootstrap, component::parse_metadata!()).unwrap();
+    assert_eq!(INIT_BOOTSTRAP.load(Relaxed), true);
+    component::init_all(component::InitStage::Kthread, component::parse_metadata!()).unwrap();
+    assert_eq!(INIT_KTHREAD.load(Relaxed), true);
+    component::init_all(component::InitStage::Process, component::parse_metadata!()).unwrap();
+    assert_eq!(INIT_PROCESS.load(Relaxed), true);
+}

kernel/src/lib.rs

@@ -29,6 +29,7 @@
 #![feature(associated_type_defaults)]
 #![register_tool(component_access_control)]
 
+use component::InitStage;
 use kcmdline::KCmdlineArg;
 use ostd::{
     arch::qemu::{exit_qemu, QemuExitCode},
@@ -83,7 +84,7 @@ mod vm;
 #[controlled]
 fn main() {
     ostd::early_println!("[kernel] OSTD initialized. Preparing components.");
-    component::init_all(component::parse_metadata!()).unwrap();
+    component::init_all(InitStage::Bootstrap, component::parse_metadata!()).unwrap();
     init();
 
     // Spawn all AP idle threads.
@@ -113,6 +114,7 @@ fn init_on_each_cpu() {
 }
 
 fn init_in_first_kthread(fs_resolver: &FsResolver) {
+    component::init_all(InitStage::Kthread, component::parse_metadata!()).unwrap();
     // Work queue should be initialized before interrupt is enabled,
     // in case any irq handler uses work queue as bottom half
     thread::work_queue::init_in_first_kthread();
@@ -124,6 +126,7 @@ fn init_in_first_kthread(fs_resolver: &FsResolver) {
 }
 
 fn init_in_first_process(ctx: &Context) {
+    component::init_all(InitStage::Process, component::parse_metadata!()).unwrap();
     device::init_in_first_process(ctx).unwrap();
     fs::init_in_first_process(ctx);
     process::init_in_first_process(ctx);