percpu: add `percpu_shared`

arch/percpu-macros/src/arch.rs

@@ -5,16 +5,19 @@ use syn::{Ident, Type};
 /// Get the base address for percpu variables of the current thread.
 pub fn get_percpu_pointer(percpu: &Ident, ty: &Type) -> TokenStream {
     quote! {
-        #[cfg(target_arch = "x86_64")] {
-            let base: *mut #ty;
-            ::core::arch::asm!(
-                "mov %gs:0, {address}",
-                "add ${percpu_pointer}, {address}",
-                percpu_pointer = sym #percpu,
-                address = out(reg) base,
-                options(att_syntax)
-            );
-            base
+        {
+            #[cfg(target_arch = "x86_64")]
+            {
+                let base: *mut #ty;
+                ::core::arch::asm!(
+                    "mov %gs:0, {address}",
+                    "add ${percpu_pointer}, {address}",
+                    percpu_pointer = sym #percpu,
+                    address = out(reg) base,
+                    options(att_syntax)
+                );
+                base
+            }
         }
     }
     .into()

arch/percpu-macros/src/lib.rs

@@ -115,3 +115,67 @@ pub fn define_percpu(attrs: TokenStream, item: TokenStream) -> TokenStream {
     }
     .into()
 }
+
+#[proc_macro_attribute]
+pub fn define_percpu_shared(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    if !attrs.is_empty() {
+        panic!("`define_percpu_shared` attribute does not take any arguments");
+    }
+
+    let item = parse_macro_input!(item as ItemStatic);
+    let vis = &item.vis;
+    let ident = &item.ident;
+    let ty = &item.ty;
+    let expr = &item.expr;
+
+    let inner_ident = format_ident!("_percpu_shared_inner_{}", ident);
+    let access_ident = format_ident!("_access_shared_{}", ident);
+
+    let as_ptr = arch::get_percpu_pointer(&inner_ident, &ty);
+
+    quote! {
+        #[link_section = ".percpu"]
+        #[allow(non_upper_case_globals)]
+        static #inner_ident: ::eonix_sync::Spin< #ty > = ::eonix_sync::Spin::new( #expr );
+        #[allow(non_camel_case_types)]
+        #vis struct #access_ident;
+        #vis static #ident: #access_ident = #access_ident;
+
+        impl #access_ident {
+            fn as_ptr(&self) -> *const ::eonix_sync::Spin< #ty > {
+                unsafe { ( #as_ptr ) as *const _ }
+            }
+
+            pub fn get_ref(&self) -> &::eonix_sync::Spin< #ty > {
+                // SAFETY: This is safe because `as_ptr()` is guaranteed to be valid.
+                unsafe { self.as_ptr().as_ref().unwrap() }
+            }
+
+            pub fn get_for_cpu(&self, cpuid: usize) -> Option<&::eonix_sync::Spin< #ty >> {
+                let offset = & #inner_ident as *const _ as usize;
+                let base = ::arch::PercpuArea::get_for(cpuid);
+                base.map(|base| unsafe { base.byte_add(offset).cast().as_ref() })
+            }
+        }
+
+        impl ::core::ops::Deref for #access_ident {
+            type Target = ::eonix_sync::Spin< #ty >;
+
+            fn deref(&self) -> &Self::Target {
+                self.get_ref()
+            }
+        }
+
+        impl<T> ::core::convert::AsRef<T> for #access_ident
+        where
+            <Self as ::core::ops::Deref>::Target: ::core::convert::AsRef<T>,
+        {
+            fn as_ref(&self) -> &T {
+                use ::core::ops::Deref;
+
+                self.deref().as_ref()
+            }
+        }
+    }
+    .into()
+}

arch/src/x86_64/init.rs

@@ -1,10 +1,5 @@
-use core::{
-    alloc::Layout,
-    pin::Pin,
-    ptr::{addr_of, NonNull},
-};
-
-use super::{enable_sse, percpu::init_percpu_area_thiscpu, GDTEntry, InterruptControl, GDT};
+use super::{enable_sse, GDTEntry, InterruptControl, GDT};
+use core::{pin::Pin, ptr::addr_of};
 
 #[repr(C)]
 #[derive(Debug, Clone, Copy)]
@@ -15,7 +10,7 @@ struct TSS_SP {
 }
 
 #[repr(C)]
-pub struct TSS {
+pub(crate) struct TSS {
     _reserved1: u32,
     rsp: [TSS_SP; 3],
     _reserved2: u32,
@@ -48,46 +43,22 @@ impl TSS {
     }
 }
 
-/// Architecture-specific per-cpu status.
-#[allow(dead_code)]
-pub struct CPUStatus {
-    id: usize,
+/// Architecture-specific cpu status data.
+pub struct CPU {
+    cpuid: usize,
     gdt: GDT,
     tss: TSS,
-
-    percpu_area: NonNull<u8>,
     pub interrupt: InterruptControl,
 }
 
-impl CPUStatus {
-    pub unsafe fn new_thiscpu<F>(allocate: F) -> Self
-    where
-        F: FnOnce(Layout) -> NonNull<u8>,
-    {
-        const PAGE_SIZE: usize = 0x1000;
-        extern "C" {
-            static PERCPU_PAGES: usize;
-            fn _PERCPU_DATA_START();
-        }
-
-        let percpu_area = allocate(Layout::from_size_align_unchecked(
-            PERCPU_PAGES * PAGE_SIZE,
-            PAGE_SIZE,
-        ));
-
-        percpu_area.copy_from_nonoverlapping(
-            NonNull::new(_PERCPU_DATA_START as *mut u8).unwrap(),
-            PERCPU_PAGES * PAGE_SIZE,
-        );
-
+impl CPU {
+    pub fn new() -> Self {
         let (interrupt_control, cpuid) = InterruptControl::new();
 
-        init_percpu_area_thiscpu(percpu_area);
         Self {
-            id: cpuid,
+            cpuid,
             gdt: GDT::new(),
             tss: TSS::new(),
-            percpu_area,
             interrupt: interrupt_control,
         }
     }
@@ -126,7 +97,7 @@ impl CPUStatus {
     }
 
     pub fn cpuid(&self) -> usize {
-        self.id
+        self.cpuid
     }
 }
 

arch/src/x86_64/interrupt.rs

@@ -426,7 +426,7 @@ impl APICRegs {
 impl InterruptControl {
     /// # Return
     /// Returns a tuple of InterruptControl and the cpu id of the current cpu.
-    pub unsafe fn new() -> (Self, usize) {
+    pub(crate) fn new() -> (Self, usize) {
         extern "C" {
             static ISR_START_ADDR: usize;
         }

arch/src/x86_64/mm.rs

@@ -0,0 +1 @@
+pub const PAGE_SIZE: usize = 0x1000;

arch/src/x86_64/mod.rs

@@ -3,17 +3,19 @@ mod gdt;
 mod init;
 mod interrupt;
 mod io;
+mod mm;
+mod percpu;
 mod user;
 
-pub(self) mod percpu;
-
 pub use self::context::*;
 pub use self::gdt::*;
 pub use self::init::*;
 pub use self::interrupt::*;
 pub use self::io::*;
 pub use self::user::*;
-pub use percpu_macros::define_percpu;
+pub use mm::*;
+pub use percpu::*;
+pub use percpu_macros::{define_percpu, define_percpu_shared};
 
 use core::arch::asm;
 

arch/src/x86_64/percpu.rs

@@ -1,16 +1,86 @@
 use super::wrmsr;
-use core::{arch::asm, ptr::NonNull};
+use crate::x86_64::mm::PAGE_SIZE;
+use core::{
+    alloc::Layout,
+    arch::asm,
+    cell::UnsafeCell,
+    ptr::{null_mut, NonNull},
+    sync::atomic::{AtomicPtr, Ordering},
+};
 
-fn save_percpu_pointer(percpu_area_base: NonNull<u8>) {
-    wrmsr(0xC0000101, percpu_area_base.as_ptr() as u64);
+pub const MAX_CPUS: usize = 256;
+
+#[repr(align(4096))]
+struct PercpuData(UnsafeCell<()>); // Not `Sync`.
+
+pub struct PercpuArea {
+    data: NonNull<PercpuData>,
 }
 
-pub unsafe fn init_percpu_area_thiscpu(percpu_area_base: NonNull<u8>) {
-    save_percpu_pointer(percpu_area_base);
+static PERCPU_POINTERS: [AtomicPtr<PercpuData>; MAX_CPUS] =
+    [const { AtomicPtr::new(null_mut()) }; MAX_CPUS];
+
+impl PercpuArea {
+    fn page_count() -> usize {
+        extern "C" {
+            static PERCPU_PAGES: usize;
+        }
+        // SAFETY: `PERCPU_PAGES` is defined in linker script and never change.
+        let page_count = unsafe { PERCPU_PAGES };
+        assert_ne!(page_count, 0);
+        page_count
+    }
+
+    fn data_start() -> NonNull<u8> {
+        extern "C" {
+            fn _PERCPU_DATA_START();
+        }
+
+        NonNull::new(_PERCPU_DATA_START as usize as *mut _)
+            .expect("Percpu data should not be null.")
+    }
+
+    fn layout() -> Layout {
+        Layout::from_size_align(Self::page_count() * PAGE_SIZE, PAGE_SIZE).expect("Invalid layout.")
+    }
+
+    pub fn new<F>(allocate: F) -> Self
+    where
+        F: FnOnce(Layout) -> NonNull<u8>,
+    {
+        let data_pointer = allocate(Self::layout());
+
+        unsafe {
+            // SAFETY: The `data_pointer` is of valid length and properly aligned.
+            data_pointer
+                .copy_from_nonoverlapping(Self::data_start(), Self::page_count() * PAGE_SIZE);
+        }
+
+        Self {
+            data: data_pointer.cast(),
+        }
+    }
+
+    /// Set up the percpu area for the current CPU.
+    pub fn setup(&self) {
+        wrmsr(0xC0000101, self.data.as_ptr() as u64);
+
+        unsafe {
+            // SAFETY: %gs:0 points to the start of the percpu area.
+            asm!(
+                "movq {}, %gs:0",
+                in(reg) self.data.as_ptr(),
+                options(nostack, preserves_flags, att_syntax)
+            );
+        }
+    }
+
+    pub fn register(self: Self, cpuid: usize) {
+        PERCPU_POINTERS[cpuid].store(self.data.as_ptr(), Ordering::Release);
+    }
 
-    asm!(
-        "movq {}, %gs:0",
-        in(reg) percpu_area_base.as_ptr(),
-        options(att_syntax)
-    );
+    pub fn get_for(cpuid: usize) -> Option<NonNull<()>> {
+        let pointer = PERCPU_POINTERS[cpuid].load(Ordering::Acquire);
+        NonNull::new(pointer.cast())
+    }
 }

arch/src/x86_64/user.rs

@@ -1,6 +1,6 @@
 use core::pin::Pin;
 
-use super::{CPUStatus, GDTEntry};
+use super::{CPU, GDTEntry};
 
 #[derive(Debug, Clone)]
 pub enum UserTLS {
@@ -28,7 +28,7 @@ impl UserTLS {
         )
     }
 
-    pub fn load(&self, cpu_status: Pin<&mut CPUStatus>) {
+    pub fn load(&self, cpu_status: Pin<&mut CPU>) {
         match self {
             Self::TLS64(base) => {
                 const IA32_KERNEL_GS_BASE: u32 = 0xc0000102;
@@ -48,7 +48,7 @@ impl UserTLS {
     }
 }
 
-pub unsafe fn load_interrupt_stack(cpu_status: Pin<&mut CPUStatus>, stack: u64) {
+pub unsafe fn load_interrupt_stack(cpu_status: Pin<&mut CPU>, stack: u64) {
     // SAFETY: We don't move the CPUStatus object.
     cpu_status.get_unchecked_mut().set_rsp0(stack);
 }

crates/eonix_runtime/src/ready_queue.rs

@@ -1,10 +1,9 @@
 use crate::task::Task;
 use alloc::{collections::VecDeque, sync::Arc};
-use eonix_sync::{LazyLock, Spin};
+use eonix_sync::Spin;
 
-#[arch::define_percpu]
-static READYQUEUE: LazyLock<Spin<FifoReadyQueue>> =
-    LazyLock::new(|| Spin::new(FifoReadyQueue::new()));
+#[arch::define_percpu_shared]
+static READYQUEUE: FifoReadyQueue = FifoReadyQueue::new();
 
 pub trait ReadyQueue {
     fn get(&mut self) -> Option<Arc<Task>>;
@@ -34,6 +33,5 @@ impl ReadyQueue for FifoReadyQueue {
 }
 
 pub fn local_rq() -> &'static Spin<dyn ReadyQueue> {
-    // SAFETY: The inner rq is protected by `Spin`.
-    unsafe { &**READYQUEUE.as_ref() }
+    &*READYQUEUE
 }

src/kernel/cpu.rs

@@ -1,32 +1,42 @@
-use core::{pin::Pin, ptr::NonNull};
-
-use arch::CPUStatus;
-
 use super::mem::{paging::Page, phys::PhysPtr as _};
+use arch::{PercpuArea, CPU};
+use core::{alloc::Layout, mem::ManuallyDrop, pin::Pin, ptr::NonNull};
+use eonix_sync::LazyLock;
 
 #[arch::define_percpu]
-static CPU_STATUS: Option<CPUStatus> = None;
+static CPU: LazyLock<CPU> = LazyLock::new(CPU::new);
 
 /// # Safety
 /// This function is unsafe because it needs preemption to be disabled.
-pub unsafe fn current_cpu() -> Pin<&'static mut CPUStatus> {
+pub unsafe fn local_cpu() -> Pin<&'static mut CPU> {
     // SAFETY: `CPU_STATUS` is global static and initialized only once.
-    unsafe { Pin::new_unchecked(CPU_STATUS.as_mut().as_mut().unwrap()) }
+    unsafe { Pin::new_unchecked(CPU.as_mut().get_mut()) }
 }
 
-pub unsafe fn init_thiscpu() {
-    let status = arch::CPUStatus::new_thiscpu(|layout| {
-        // TODO: Use page size defined in `arch`.
-        let page_count = (layout.size() + 0x1000 - 1) / 0x1000;
-        let page = Page::early_alloc_ceil(page_count);
-        let pointer = page.as_cached().as_ptr();
-        core::mem::forget(page);
+pub fn percpu_allocate(layout: Layout) -> NonNull<u8> {
+    // TODO: Use page size defined in `arch`.
+    let page_count = layout.size().div_ceil(arch::PAGE_SIZE);
+    let page = ManuallyDrop::new(Page::early_alloc_ceil(page_count));
+    let pointer = page.as_cached().as_ptr();
 
-        NonNull::new(pointer).expect("Allocated page pfn should be non-null")
-    });
+    NonNull::new(pointer).expect("Allocated page pfn should be non-null.")
+}
 
-    CPU_STATUS.set(Some(status));
+pub fn init_localcpu() {
+    let percpu_area = PercpuArea::new(percpu_allocate);
 
-    // SAFETY: `CPU_STATUS` is global static and initialized only once.
-    current_cpu().init();
+    // Preemption count is percpu. So we need to initialize percpu area first.
+    percpu_area.setup();
+
+    eonix_preempt::disable();
+
+    // SAFETY: Preemption is disabled.
+    let mut cpu = unsafe { local_cpu() };
+
+    unsafe {
+        cpu.as_mut().init();
+    }
+    percpu_area.register(cpu.cpuid());
+
+    eonix_preempt::enable();
 }

src/kernel/interrupt.rs

@@ -1,4 +1,4 @@
-use super::cpu::current_cpu;
+use super::cpu::local_cpu;
 use super::mem::handle_page_fault;
 use super::syscall::handle_syscall32;
 use super::task::{ProcessList, Signal};
@@ -95,5 +95,5 @@ pub fn init() -> KResult<()> {
 
 pub fn end_of_interrupt() {
     // SAFETY: We only use this function in irq context, where preemption is disabled.
-    unsafe { current_cpu() }.interrupt.end_of_interrupt();
+    unsafe { local_cpu() }.interrupt.end_of_interrupt();
 }

src/kernel/smp.rs

@@ -1,7 +1,7 @@
-use super::cpu::init_thiscpu;
+use super::cpu::init_localcpu;
 use crate::{
     kernel::{
-        cpu::current_cpu,
+        cpu::local_cpu,
         mem::{paging::Page, phys::PhysPtr as _},
         task::KernelStack,
     },
@@ -18,9 +18,9 @@ define_smp_bootstrap!(4, ap_entry, {
 });
 
 unsafe extern "C" fn ap_entry() -> ! {
-    init_thiscpu();
+    init_localcpu();
     Scheduler::init_local_scheduler::<KernelStack>();
-    println_debug!("AP{} started", current_cpu().cpuid());
+    println_debug!("AP{} started", local_cpu().cpuid());
 
     eonix_preempt::disable();
     arch::enable_irqs();
@@ -32,7 +32,12 @@ unsafe extern "C" fn ap_entry() -> ! {
     }
 }
 
-pub unsafe fn bootstrap_smp() {
-    current_cpu().bootstrap_cpus();
-    wait_cpus_online();
+pub fn bootstrap_smp() {
+    eonix_preempt::disable();
+    unsafe {
+        // SAFETY: Preemption is disabled.
+        local_cpu().bootstrap_cpus();
+        wait_cpus_online();
+    }
+    eonix_preempt::enable();
 }

src/kernel/task/thread.rs

@@ -4,7 +4,7 @@ use super::{
 };
 use crate::{
     kernel::{
-        cpu::current_cpu,
+        cpu::local_cpu,
         mem::VAddr,
         user::dataflow::CheckedUserPointer,
         vfs::{filearray::FileArray, FsContext},
@@ -243,7 +243,7 @@ impl Thread {
     pub unsafe fn load_thread_area32(&self) {
         if let Some(tls) = self.inner.lock().tls.as_ref() {
             // SAFETY: Preemption is disabled.
-            tls.load(current_cpu());
+            tls.load(local_cpu());
         }
     }
 
@@ -331,7 +331,7 @@ impl Contexted for ThreadRunnable {
             0 => {}
             sp => unsafe {
                 // SAFETY: Preemption is disabled.
-                arch::load_interrupt_stack(current_cpu(), sp as u64);
+                arch::load_interrupt_stack(local_cpu(), sp as u64);
             },
         }
 
@@ -391,7 +391,7 @@ impl PinRun for ThreadRunnable {
 
         unsafe {
             // SAFETY: Preemption is disabled.
-            arch::load_interrupt_stack(current_cpu(), sp as u64);
+            arch::load_interrupt_stack(local_cpu(), sp as u64);
         }
 
         eonix_preempt::enable();

src/lib.rs

@@ -30,7 +30,7 @@ use core::alloc::{GlobalAlloc, Layout};
 use elf::ParsedElf32;
 use eonix_runtime::{run::FutureRun, scheduler::Scheduler};
 use kernel::{
-    cpu::init_thiscpu,
+    cpu::init_localcpu,
     mem::Page,
     task::{KernelStack, ProcessBuilder, ProcessList, ThreadBuilder, ThreadRunnable},
     vfs::{
@@ -98,7 +98,7 @@ extern "C" {
 pub extern "C" fn rust_kinit(early_kstack_pfn: usize) -> ! {
     // We don't call global constructors.
     // Rust doesn't need that, and we're not going to use global variables in C++.
-    unsafe { init_thiscpu() };
+    init_localcpu();
 
     unsafe { init_allocator() };
 
@@ -138,7 +138,7 @@ async fn init_process(early_kstack_pfn: usize) {
     fs::procfs::init();
     fs::fat32::init();
 
-    unsafe { kernel::smp::bootstrap_smp() };
+    kernel::smp::bootstrap_smp();
 
     let (ip, sp, mm_list) = {
         // mount fat32 /mnt directory