From ebf1a61fd5665211eaa32c54e73c064c063e356b Mon Sep 17 00:00:00 2001 From: Putta Khunchalee Date: Sun, 8 Sep 2024 18:19:23 +0700 Subject: [PATCH] Initializes malloc stats --- src/obconf/src/lib.rs | 10 +++++- src/obkrnl/src/config/mod.rs | 25 ++++++++------- src/obkrnl/src/context/aarch64.rs | 4 +++ src/obkrnl/src/context/mod.rs | 51 ++++++++++++++++++++++++++----- src/obkrnl/src/context/x86_64.rs | 26 ++++++++++++++-- src/obkrnl/src/main.rs | 13 +++++--- src/obkrnl/src/malloc/stage2.rs | 40 +++++++++++++++++++++--- src/obkrnl/src/proc/thread.rs | 12 +++++++- src/obkrnl/src/uma/mod.rs | 13 ++++++-- 9 files changed, 160 insertions(+), 34 deletions(-) diff --git a/src/obconf/src/lib.rs b/src/obconf/src/lib.rs index 75574c17..c67374db 100644 --- a/src/obconf/src/lib.rs +++ b/src/obconf/src/lib.rs @@ -1,11 +1,19 @@ #![no_std] +use core::num::NonZero; + pub use self::env::*; mod env; -/// Information about the boot environment. +/// Contains information about the boot environment. #[repr(C)] pub enum BootEnv { Vm(Vm), } + +/// Runtime configurations for the kernel. +#[repr(C)] +pub struct Config { + pub max_cpu: NonZero, +} diff --git a/src/obkrnl/src/config/mod.rs b/src/obkrnl/src/config/mod.rs index 87d27437..ff075b14 100644 --- a/src/obkrnl/src/config/mod.rs +++ b/src/obkrnl/src/config/mod.rs @@ -1,30 +1,33 @@ use core::ptr::null; use macros::elf_note; -use obconf::BootEnv; +use obconf::{BootEnv, Config}; -#[cfg(target_arch = "aarch64")] -pub use self::aarch64::*; -#[cfg(target_arch = "x86_64")] -pub use self::x86_64::*; +pub use self::arch::*; -#[cfg(target_arch = "aarch64")] -mod aarch64; -#[cfg(target_arch = "x86_64")] -mod x86_64; +#[cfg_attr(target_arch = "aarch64", path = "aarch64.rs")] +#[cfg_attr(target_arch = "x86_64", path = "x86_64.rs")] +mod arch; pub fn boot_env() -> &'static BootEnv { - // SAFETY: This is safe because the set_boot_env() requirements. + // SAFETY: This is safe because the setup() requirements. unsafe { &*BOOT_ENV } } +pub fn config() -> &'static Config { + // SAFETY: This is safe because the setup() requirements. + unsafe { &*CONFIG } +} + /// # Safety /// This function must be called immediately in the kernel entry point. After that it must never /// be called again. -pub unsafe fn set_boot_env(env: &'static BootEnv) { +pub unsafe fn setup(env: &'static BootEnv, conf: &'static Config) { BOOT_ENV = env; + CONFIG = conf; } static mut BOOT_ENV: *const BootEnv = null(); +static mut CONFIG: *const Config = null(); #[elf_note(section = ".note.obkrnl.page-size", name = "obkrnl", ty = 0)] static NOTE_PAGE_SIZE: [u8; size_of::()] = PAGE_SIZE.to_ne_bytes(); diff --git a/src/obkrnl/src/context/aarch64.rs b/src/obkrnl/src/context/aarch64.rs index 24c7598c..153ff13b 100644 --- a/src/obkrnl/src/context/aarch64.rs +++ b/src/obkrnl/src/context/aarch64.rs @@ -8,3 +8,7 @@ pub unsafe fn activate(_: *mut Context) { pub unsafe fn thread() -> *const Thread { todo!(); } + +pub unsafe fn current() -> *const Context { + todo!(); +} diff --git a/src/obkrnl/src/context/mod.rs b/src/obkrnl/src/context/mod.rs index 8669f9a1..d4a967ec 100644 --- a/src/obkrnl/src/context/mod.rs +++ b/src/obkrnl/src/context/mod.rs @@ -12,20 +12,16 @@ mod arch; /// not safe to have a temporary a pointer or reference to this struct or its field because the CPU /// might get interupted, which mean it is possible for the next instruction to get executed on /// a different CPU if the interupt cause the CPU to switch the task. -/// -/// We don't support `pc_cpuid` field here because it value is 100% unpredictable due to the above -/// reason. Once we have loaded `pc_cpuid` the next instruction might get executed on a different -/// CPU, which render the loaded value incorrect. The only way to prevent this issue is to disable -/// interupt before reading `pc_cpuid`, which can make the CPU missed some events from the other -/// hardwares. pub struct Context { + cpu: usize, // pc_cpuid thread: AtomicPtr, // pc_curthread } impl Context { /// See `pcpu_init` on the PS4 for a reference. - pub fn new(td: Arc) -> Self { + pub fn new(cpu: usize, td: Arc) -> Self { Self { + cpu, thread: AtomicPtr::new(Arc::into_raw(td).cast_mut()), } } @@ -35,10 +31,28 @@ impl Context { // it is going to be the same one since it represent the current thread. let td = unsafe { self::arch::thread() }; + // We cannot return a reference here because it requires 'static lifetime, which allow the + // caller to store it at a global level. Once the thread is destroyed that reference will be + // invalid. unsafe { Arc::increment_strong_count(td) }; unsafe { Arc::from_raw(td) } } + /// See `critical_enter` and `critical_exit` on the PS4 for a reference. + pub fn pin() -> PinnedContext { + // TODO: Verify if memory ordering here is correct. We need a call to self::arch::current() + // to execute after the thread is in a critical section. The CPU must not reorder this. Our + // current implementation follow how Drop on Arc is implemented. + let td = unsafe { self::arch::thread() }; + + unsafe { (*td).critical_sections().fetch_add(1, Ordering::Release) }; + core::sync::atomic::fence(Ordering::Acquire); + + // Once the thread is in a critical section it will never be switch a CPU so it is safe to + // keep a pointer to a context here. + PinnedContext(unsafe { self::arch::current() }) + } + /// # Safety /// The only place this method is safe to call is in the CPU entry point. Once this method /// return this instance must outlive the CPU lifetime and it must never be accessed via this @@ -54,3 +68,26 @@ impl Drop for Context { unsafe { drop(Arc::from_raw(self.thread.load(Ordering::Relaxed))) }; } } + +/// RAII struct to pin the current thread to current CPU. +/// +/// This struct must not implement [`Send`] and [`Sync`]. Currently it stored a pointer, which will +/// make it `!Send` and `!Sync`. +pub struct PinnedContext(*const Context); + +impl PinnedContext { + pub fn cpu(&self) -> usize { + unsafe { (*self.0).cpu } + } +} + +impl Drop for PinnedContext { + fn drop(&mut self) { + // TODO: Verify if memory ordering here is correct. + let td = unsafe { (*self.0).thread.load(Ordering::Relaxed) }; + + unsafe { (*td).critical_sections().fetch_sub(1, Ordering::Release) }; + + // TODO: Implement td_owepreempt. + } +} diff --git a/src/obkrnl/src/context/x86_64.rs b/src/obkrnl/src/context/x86_64.rs index 28dbd29b..d3d59532 100644 --- a/src/obkrnl/src/context/x86_64.rs +++ b/src/obkrnl/src/context/x86_64.rs @@ -15,7 +15,7 @@ pub unsafe fn activate(cx: *mut Context) { in("ecx") 0xc0000101u32, in("edx") cx >> 32, in("eax") cx, - options(preserves_flags, nostack) + options(nomem, preserves_flags, nostack) ); // Clear FS and GS for user mode. @@ -24,7 +24,7 @@ pub unsafe fn activate(cx: *mut Context) { in("ecx") 0xc0000100u32, in("edx") 0, in("eax") 0, - options(preserves_flags, nostack) + options(nomem, preserves_flags, nostack) ); asm!( @@ -32,7 +32,7 @@ pub unsafe fn activate(cx: *mut Context) { in("ecx") 0xc0000102u32, in("edx") 0, in("eax") 0, - options(preserves_flags, nostack) + options(nomem, preserves_flags, nostack) ); } @@ -50,3 +50,23 @@ pub unsafe fn thread() -> *const Thread { td } + +pub unsafe fn current() -> *const Context { + // Load current GS. + let mut edx: u32; + let mut eax: u32; + + asm!( + "rdmsr", + in("ecx") 0xc0000101u32, + out("edx") edx, + out("eax") eax, + options(pure, nomem, preserves_flags, nostack) + ); + + // Combine EDX and EAX. + let edx = edx as usize; + let eax = eax as usize; + + ((edx << 32) | eax) as *const Context +} diff --git a/src/obkrnl/src/main.rs b/src/obkrnl/src/main.rs index 41391937..7a7e0a57 100644 --- a/src/obkrnl/src/main.rs +++ b/src/obkrnl/src/main.rs @@ -1,15 +1,15 @@ #![no_std] #![cfg_attr(not(test), no_main)] -use crate::config::set_boot_env; use crate::context::Context; use crate::malloc::KernelHeap; use crate::proc::Thread; use alloc::sync::Arc; use core::arch::asm; use core::mem::zeroed; +use core::num::NonZero; use core::panic::PanicInfo; -use obconf::BootEnv; +use obconf::{BootEnv, Config}; mod config; mod console; @@ -37,8 +37,13 @@ extern crate alloc; #[allow(dead_code)] #[cfg_attr(target_os = "none", no_mangle)] extern "C" fn _start(env: &'static BootEnv) -> ! { + // TODO: Accept config from bootloader/hypervisor. + static CONFIG: Config = Config { + max_cpu: unsafe { NonZero::new_unchecked(1) }, + }; + // SAFETY: This is safe because we called it as the first thing here. - unsafe { set_boot_env(env) }; + unsafe { crate::config::setup(env, &CONFIG) }; info!("Starting Obliteration Kernel."); @@ -48,7 +53,7 @@ extern "C" fn _start(env: &'static BootEnv) -> ! { // Setup CPU context. We use a different mechanism here. The PS4 put all of pcpu at a global // level but we put it on each CPU stack instead. let thread0 = Arc::new(thread0); - let mut cx = Context::new(thread0); + let mut cx = Context::new(0, thread0); // SAFETY: We are in the main CPU entry point and we move all the remaining code after this into // a dedicated no-return function. diff --git a/src/obkrnl/src/malloc/stage2.rs b/src/obkrnl/src/malloc/stage2.rs index 22ff5aa6..33125f21 100644 --- a/src/obkrnl/src/malloc/stage2.rs +++ b/src/obkrnl/src/malloc/stage2.rs @@ -1,10 +1,11 @@ -use crate::config::PAGE_SIZE; +use crate::config::{config, PAGE_SIZE}; use crate::context::Context; use crate::uma::UmaZone; use alloc::string::ToString; use alloc::sync::Arc; use alloc::vec::Vec; use core::alloc::Layout; +use core::sync::atomic::{AtomicU64, Ordering}; /// Stage 2 kernel heap. /// @@ -12,6 +13,7 @@ use core::alloc::Layout; /// `malloc_type` and `malloc_type_internal` structure. pub struct Stage2 { zones: [Vec>; (usize::BITS - 1) as usize], // kmemsize + kmemzones + stats: Vec, // mti_stats } impl Stage2 { @@ -54,7 +56,14 @@ impl Stage2 { zones }); - Self { zones } + // TODO: Is there a better way than this? + let mut stats = Vec::with_capacity(config().max_cpu.get()); + + for _ in 0..config().max_cpu.get() { + stats.push(Stats::default()); + } + + Self { zones, stats } } /// See `malloc` on the PS4 for a reference. @@ -82,8 +91,24 @@ impl Stage2 { size }; - // TODO: There are more logic after this on the PS4. - self.zones[align][size >> Self::KMEM_ZSHIFT].alloc() + // Allocate a memory from UMA zone. + let zone = &self.zones[align][size >> Self::KMEM_ZSHIFT]; + let mem = zone.alloc(); + + // Update stats. + let cx = Context::pin(); + let stats = &self.stats[cx.cpu()]; + let size = if mem.is_null() { 0 } else { zone.size() }; + + if size != 0 { + stats + .alloc_bytes + .fetch_add(size.try_into().unwrap(), Ordering::Relaxed); + stats.alloc_count.fetch_add(1, Ordering::Relaxed); + } + + // TODO: How to update mts_size here since our zone table also indexed by alignment? + mem } else { todo!() } @@ -96,3 +121,10 @@ impl Stage2 { todo!() } } + +/// Implementation of `malloc_type_stats` structure. +#[derive(Default)] +struct Stats { + alloc_bytes: AtomicU64, // mts_memalloced + alloc_count: AtomicU64, // mts_numallocs +} diff --git a/src/obkrnl/src/proc/thread.rs b/src/obkrnl/src/proc/thread.rs index 9fbffab4..e1d4cdbc 100644 --- a/src/obkrnl/src/proc/thread.rs +++ b/src/obkrnl/src/proc/thread.rs @@ -1,9 +1,12 @@ +use core::sync::atomic::AtomicU32; + /// Implementation of `thread` structure. /// /// All thread **must** run to completion once execution has been started otherwise resource will be /// leak if the thread is dropped while its execution currently in the kernel space. pub struct Thread { - active_interrupts: usize, // td_intr_nesting_level + critical_sections: AtomicU32, // td_critnest + active_interrupts: usize, // td_intr_nesting_level } impl Thread { @@ -12,11 +15,18 @@ impl Thread { /// responsibility to configure the thread after this so it have a proper states and trigger /// necessary events. pub unsafe fn new_bare() -> Self { + // td_critnest on the PS4 started with 1 but this does not work in our case because we use + // RAII to increase and decrease it. Self { + critical_sections: AtomicU32::new(0), active_interrupts: 0, } } + pub fn critical_sections(&self) -> &AtomicU32 { + &self.critical_sections + } + pub fn active_interrupts(&self) -> usize { self.active_interrupts } diff --git a/src/obkrnl/src/uma/mod.rs b/src/obkrnl/src/uma/mod.rs index 3f7f3645..90eae17c 100644 --- a/src/obkrnl/src/uma/mod.rs +++ b/src/obkrnl/src/uma/mod.rs @@ -1,12 +1,19 @@ use alloc::borrow::Cow; /// Implementation of `uma_zone` structure. -pub struct UmaZone {} +pub struct UmaZone { + size: usize, // uz_size +} impl UmaZone { /// See `uma_zcreate` on the PS4 for a reference. - pub fn new(_: Cow<'static, str>, _: usize, _: usize) -> Self { - Self {} + pub fn new(_: Cow<'static, str>, size: usize, _: usize) -> Self { + // TODO: Check if size is allowed to be zero. If not, change it to NonZero. + Self { size } + } + + pub fn size(&self) -> usize { + self.size } /// See `uma_zalloc_arg` on the PS4 for a reference.