Fix the preempt info on APs

ostd/src/cpu/local/mod.rs

@@ -64,7 +64,9 @@ extern "C" {
 ///
 /// It should be called early to let [`crate::task::disable_preempt`] work,
 /// which needs to update a CPU-local preemption info. Otherwise it may
-/// panic when calling [`crate::task::disable_preempt`].
+/// panic when calling [`crate::task::disable_preempt`]. It is needed since
+/// heap allocations need to disable preemption, which would happen in the
+/// very early stage of the kernel.
 ///
 /// # Safety
 ///
@@ -140,16 +142,12 @@ pub unsafe fn init_on_ap(cpu_id: u32) {
 
     let ap_pages_ptr = paddr_to_vaddr(ap_pages.start_paddr()) as *mut u32;
 
-    debug_assert_eq!(
-        cpu_id,
-        // SAFETY: the CPU ID is stored at the beginning of the CPU local area.
-        unsafe { ap_pages_ptr.read() }
-    );
-
     // SAFETY: the memory will be dedicated to the AP. And we are on the AP.
     unsafe {
         arch::cpu::local::set_base(ap_pages_ptr as u64);
     }
+
+    crate::task::reset_preempt_info();
 }
 
 mod has_init {

ostd/src/task/mod.rs

@@ -9,6 +9,8 @@ pub mod scheduler;
 #[allow(clippy::module_inception)]
 mod task;
 
+pub(crate) use preempt::cpu_local::reset_preempt_info;
+
 pub use self::{
     preempt::{disable_preempt, DisabledPreemptGuard},
     task::{AtomicCpuId, Priority, Task, TaskAdapter, TaskContextApi, TaskOptions},

ostd/src/task/preempt/cpu_local.rs

@@ -57,5 +57,19 @@ cpu_local_cell! {
     static PREEMPT_INFO: u32 = NEED_PREEMPT_MASK;
 }
 
+/// Resets the preempt info to the initial state.
+///
+/// # Safety
+///
+/// This function is only useful for the initialization of application
+/// processors' CPU-local storage. Because that the BSP should access the CPU-
+/// local storage (`PREEMPT_INFO`) (when doing heap allocation) before we can
+/// initialize the CPU-local storage for APs, the value of the AP's
+/// `PREEMPT_INFO` would be that of the BSP's. Therefore, we need to reset the
+/// `PREEMPT_INFO` to the initial state on APs' initialization.
+pub(crate) unsafe fn reset_preempt_info() {
+    PREEMPT_INFO.store(NEED_PREEMPT_MASK);
+}
+
 const NEED_PREEMPT_MASK: u32 = 1 << 31;
 const GUARD_COUNT_MASK: u32 = (1 << 31) - 1;