Initializes zone allocation for uma_zalloc_arg (#1097)

.vscode/settings.json

@@ -20,5 +20,6 @@
   ],
   "rust-analyzer.cargo.features": "all",
   "rust-analyzer.imports.granularity.group": "module",
-  "rust-analyzer.imports.group.enable": false
+  "rust-analyzer.imports.group.enable": false,
+  "rust-analyzer.imports.prefix": "self"
 }

kernel/src/config/mod.rs

@@ -17,11 +17,12 @@ pub fn boot_env() -> &'static BootEnv {
     unsafe { &*BOOT_ENV }
 }
 
-/// # Interupt safety
-/// This function is interupt safe.
+/// # Context safety
+/// This function does not require a CPU context.
+///
+/// # Interrupt safety
+/// This function can be called from interrupt handler.
 pub fn config() -> &'static Config {
-    // This function is not allowed to access the CPU context due to it can be called before the
-    // context has been activated.
     // SAFETY: This is safe because the setup() requirements.
     unsafe { &*CONFIG }
 }

kernel/src/context/local.rs

@@ -11,6 +11,8 @@ use core::ops::Deref;
 pub struct CpuLocal<T>(Vec<T>);
 
 impl<T> CpuLocal<T> {
+    /// # Context safety
+    /// This function does not require a CPU context on **stage 1** heap as long as `f` does not.
     pub fn new(mut f: impl FnMut(usize) -> T) -> Self {
         let len = config().max_cpu.get();
         let mut vec = Vec::with_capacity(len);

kernel/src/lock/gutex/mod.rs

@@ -88,7 +88,7 @@ pub struct GutexGroup {
 
 impl GutexGroup {
     /// # Context safety
-    /// This function does not require a CPU context.
+    /// This function does not require a CPU context on **stage 1** heap.
     pub fn new() -> Arc<Self> {
         Arc::new(Self {
             owning: AtomicUsize::new(MTX_UNOWNED),
@@ -108,21 +108,25 @@ impl GutexGroup {
 
     #[inline(never)]
     fn lock(&self) -> GroupGuard {
-        // Check if the calling thread already own the lock.
+        // Acquire the lock.
         let td = current_thread();
         let id = BorrowedArc::as_ptr(&td) as usize;
 
-        if id == self.owning.load(Ordering::Relaxed) {
-            // SAFETY: This is safe because the current thread own the lock.
-            return unsafe { GroupGuard::new(self) };
-        }
+        loop {
+            let owning = match self.owning.compare_exchange(
+                MTX_UNOWNED,
+                id,
+                Ordering::Acquire,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => break,
+                Err(v) => v,
+            };
+
+            if owning == id {
+                break;
+            }
 
-        // Acquire the lock.
-        while self
-            .owning
-            .compare_exchange(MTX_UNOWNED, id, Ordering::Acquire, Ordering::Relaxed)
-            .is_err()
-        {
             todo!()
         }
 

kernel/src/uma/mod.rs

@@ -1,26 +1,37 @@
-use self::cache::UmaCache;
+use self::bucket::UmaBucket;
 use crate::context::{current_thread, CpuLocal};
+use crate::lock::{Gutex, GutexGroup};
 use alloc::borrow::Cow;
 use core::cell::RefCell;
 use core::num::NonZero;
+use core::ops::DerefMut;
+use core::ptr::null_mut;
 
 mod bucket;
-mod cache;
 
 /// Implementation of `uma_zone` structure.
 pub struct UmaZone {
     size: NonZero<usize>,                // uz_size
     caches: CpuLocal<RefCell<UmaCache>>, // uz_cpu
+    allocs: Gutex<u64>,                  // uz_allocs
+    frees: Gutex<u64>,                   // uz_frees
 }
 
 impl UmaZone {
     /// See `uma_zcreate` on the PS4 for a reference.
+    ///
+    /// # Context safety
+    /// This function does not require a CPU context on **stage 1** heap.
     pub fn new(_: Cow<'static, str>, size: NonZero<usize>, _: usize) -> Self {
         // Ths PS4 allocate a new uma_zone from masterzone_z but we don't have that. This method
         // basically an implementation of zone_ctor.
+        let gg = GutexGroup::new();
+
         Self {
             size,
             caches: CpuLocal::new(|_| RefCell::default()),
+            allocs: gg.clone().spawn(0),
+            frees: gg.spawn(0),
         }
     }
 
@@ -37,22 +48,57 @@ impl UmaZone {
             panic!("heap allocation in a non-sleeping context is not supported");
         }
 
-        // Try to allocate from per-CPU cache.
-        let pin = self.caches.lock();
-        let mut cache = pin.borrow_mut();
-        let bucket = cache.alloc_mut();
+        // Try allocate from per-CPU cache first so we don't need to acquire a mutex lock.
+        let caches = self.caches.lock();
+        let mem = Self::alloc_from_cache(caches.borrow_mut().deref_mut());
 
-        while let Some(bucket) = bucket {
-            if bucket.len() != 0 {
-                todo!()
-            }
+        if !mem.is_null() {
+            return mem;
+        }
+
+        drop(caches); // Exit from non-sleeping context before acquire the mutex.
 
-            todo!()
+        // Cache not found, allocate from the zone. We need to re-check the cache again because we
+        // may on a different CPU since we drop the CPU pinning on the above.
+        let mut allocs = self.allocs.write();
+        let mut frees = self.frees.write();
+        let caches = self.caches.lock();
+        let mut cache = caches.borrow_mut();
+        let mem = Self::alloc_from_cache(&mut cache);
+
+        if !mem.is_null() {
+            return mem;
         }
 
-        drop(cache);
-        drop(pin);
+        *allocs += core::mem::take(&mut cache.allocs);
+        *frees += core::mem::take(&mut cache.frees);
 
         todo!()
     }
+
+    fn alloc_from_cache(c: &mut UmaCache) -> *mut u8 {
+        while let Some(b) = &mut c.alloc {
+            if b.len() != 0 {
+                todo!()
+            }
+
+            if c.free.as_ref().is_some_and(|b| b.len() != 0) {
+                core::mem::swap(&mut c.alloc, &mut c.free);
+                continue;
+            }
+
+            break;
+        }
+
+        null_mut()
+    }
+}
+
+/// Implementation of `uma_cache` structure.
+#[derive(Default)]
+struct UmaCache {
+    alloc: Option<UmaBucket>, // uc_allocbucket
+    free: Option<UmaBucket>,  // uc_freebucket
+    allocs: u64,              // uc_allocs
+    frees: u64,               // uc_frees
 }