Implements bucket_zones (#1272)

kernel/src/lock/gutex/mod.rs

@@ -112,7 +112,6 @@ impl GutexGroup {
         }
     }
 
-    #[inline(never)]
     fn lock(&self) -> GroupGuard {
         // Acquire the lock.
         let td = current_thread();
@@ -133,13 +132,18 @@ impl GutexGroup {
                 break;
             }
 
-            todo!()
+            self.wait();
         }
 
         // SAFETY: This is safe because the current thread acquire the lock successfully by the
         // above compare_exchange().
         unsafe { GroupGuard::new(self) }
     }
+
+    #[inline(never)]
+    fn wait(&self) {
+        todo!()
+    }
 }
 
 unsafe impl Send for GutexGroup {}
@@ -164,12 +168,17 @@ impl<'a> GroupGuard<'a> {
             phantom: PhantomData,
         }
     }
+
+    #[inline(never)]
+    fn release(&mut self) {
+        self.group.owning.store(MTX_UNOWNED, Ordering::Release);
+
+        todo!("wakeup waiting thread");
+    }
 }
 
 impl Drop for GroupGuard<'_> {
-    #[inline(never)]
     fn drop(&mut self) {
-        // Decrease the active lock.
         unsafe {
             let active = self.group.active.get();
 
@@ -180,9 +189,6 @@ impl Drop for GroupGuard<'_> {
             }
         }
 
-        // Release the lock.
-        self.group.owning.store(MTX_UNOWNED, Ordering::Release);
-
-        todo!("wakeup waiting thread");
+        self.release();
     }
 }

kernel/src/malloc/stage2.rs

@@ -1,6 +1,6 @@
 use crate::config::PAGE_SIZE;
 use crate::context::{current_thread, current_uma, CpuLocal};
-use crate::uma::{UmaFlags, UmaZone};
+use crate::uma::{Alloc, UmaFlags, UmaZone};
 use alloc::string::ToString;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
@@ -106,7 +106,7 @@ impl VmHeap {
 
             // Allocate a memory from UMA zone.
             let zone = &self.zones[align][size >> Self::KMEM_ZSHIFT];
-            let mem = zone.alloc();
+            let mem = zone.alloc(Alloc::Wait | Alloc::Zero);
 
             // Update stats.
             let stats = self.stats.lock();

kernel/src/uma/boxed.rs

@@ -0,0 +1,15 @@
+use core::ops::Deref;
+
+/// Encapsulates an object allocated from a UMA zone.
+pub struct UmaBox<T: ?Sized>(*mut T);
+
+impl<T: ?Sized> Deref for UmaBox<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { &*self.0 }
+    }
+}
+
+unsafe impl<T: Send + ?Sized> Send for UmaBox<T> {}
+unsafe impl<T: Sync + ?Sized> Sync for UmaBox<T> {}

kernel/src/uma/bucket.rs

@@ -1,10 +1,15 @@
 /// Implementation of `uma_bucket` structure.
-pub struct UmaBucket {
+#[repr(C)]
+pub struct UmaBucket<I: ?Sized> {
     len: usize, // ub_cnt
+    items: I,   // ub_bucket
 }
 
-impl UmaBucket {
+impl<I: ?Sized> UmaBucket<I> {
     pub fn len(&self) -> usize {
         self.len
     }
 }
+
+/// Each item in the [`UmaBucket::items`].
+pub struct BucketItem {}

kernel/src/uma/mod.rs

@@ -1,12 +1,18 @@
+pub use self::boxed::*;
 pub use self::zone::*;
 
+use self::bucket::{BucketItem, UmaBucket};
 use crate::config::PAGE_SIZE;
+use alloc::format;
 use alloc::string::String;
 use alloc::sync::Arc;
+use alloc::vec::Vec;
+use core::alloc::Layout;
 use core::num::NonZero;
 use core::sync::atomic::AtomicBool;
 use macros::bitflag;
 
+mod boxed;
 mod bucket;
 mod keg;
 mod slab;
@@ -15,6 +21,8 @@ mod zone;
 /// Implementation of UMA system.
 pub struct Uma {
     bucket_enable: Arc<AtomicBool>,
+    bucket_keys: Arc<Vec<usize>>,    // bucket_size
+    bucket_zones: Arc<Vec<UmaZone>>, // bucket_zones
 }
 
 impl Uma {
@@ -25,7 +33,6 @@ impl Uma {
     const MAX_WASTE: NonZero<usize> = NonZero::new(PAGE_SIZE.get() / 10).unwrap();
     const BUCKET_MAX: usize = 128;
     const BUCKET_SHIFT: usize = 4;
-    const BUCKET_ZONES: usize = ((Self::BUCKET_MAX >> Self::BUCKET_SHIFT) + 1);
 
     /// `bucket_zones`.
     const BUCKET_SIZES: [usize; 4] = [16, 32, 64, 128];
@@ -38,18 +45,41 @@ impl Uma {
     /// |PS4 11.00|0x13CA70|
     pub fn new() -> Arc<Self> {
         let bucket_enable = Arc::new(AtomicBool::new(true)); // TODO: Use a proper value.
-        let mut bucket_keys = [0; Self::BUCKET_ZONES];
+        let mut bucket_keys = Vec::new();
+        let mut bucket_zones = Vec::with_capacity(Self::BUCKET_SIZES.len());
         let mut ki = 0;
 
         // Create bucket zones.
         for (si, size) in Self::BUCKET_SIZES.into_iter().enumerate() {
+            let items = Layout::array::<BucketItem>(size).unwrap();
+            let layout = Layout::new::<UmaBucket<()>>()
+                .extend(items)
+                .unwrap()
+                .0
+                .pad_to_align();
+
+            bucket_zones.push(UmaZone::new(
+                bucket_enable.clone(),
+                Arc::default(),
+                Arc::default(),
+                format!("{size} Bucket"),
+                None,
+                layout.size().try_into().unwrap(),
+                Some(layout.align() - 1),
+                UmaFlags::Bucket | UmaFlags::Internal,
+            ));
+
             while ki <= size {
-                bucket_keys[ki >> Self::BUCKET_SHIFT] = si;
+                bucket_keys.push(si);
                 ki += 1 << Self::BUCKET_SHIFT;
             }
         }
 
-        Arc::new(Self { bucket_enable })
+        Arc::new(Self {
+            bucket_enable,
+            bucket_keys: Arc::new(bucket_keys),
+            bucket_zones: Arc::new(bucket_zones),
+        })
     }
 
     /// See `uma_zcreate` on the Orbis for a reference.
@@ -67,7 +97,16 @@ impl Uma {
     ) -> UmaZone {
         // The Orbis will allocate a new zone from masterzone_z. We choose to remove this since it
         // does not idomatic to Rust, which mean our uma_zone itself can live on the stack.
-        UmaZone::new(self.bucket_enable.clone(), name, None, size, align, flags)
+        UmaZone::new(
+            self.bucket_enable.clone(),
+            self.bucket_keys.clone(),
+            self.bucket_zones.clone(),
+            name,
+            None,
+            size,
+            align,
+            flags,
+        )
     }
 }
 
@@ -96,8 +135,21 @@ pub enum UmaFlags {
     CacheSpread = 0x1000,
     /// `UMA_ZONE_VTOSLAB`.
     VToSlab = 0x2000,
+    /// `UMA_ZFLAG_BUCKET`.
+    Bucket = 0x2000000,
     /// `UMA_ZFLAG_INTERNAL`.
     Internal = 0x20000000,
     /// `UMA_ZFLAG_CACHEONLY`.
     CacheOnly = 0x80000000,
 }
+
+/// Implementation of `malloc` flags.
+#[bitflag(u32)]
+pub enum Alloc {
+    /// `M_WAITOK`.
+    Wait = 0x2,
+    /// `M_ZERO`.
+    Zero = 0x100,
+    /// `M_NOVM`.
+    NoVm = 0x200,
+}

kernel/src/uma/zone.rs

@@ -1,11 +1,12 @@
-use super::bucket::UmaBucket;
+use super::bucket::{BucketItem, UmaBucket};
 use super::keg::UmaKeg;
-use super::{Uma, UmaFlags};
+use super::{Alloc, Uma, UmaBox, UmaFlags};
 use crate::context::{current_thread, CpuLocal};
 use crate::lock::{Gutex, GutexGroup, GutexWrite};
 use alloc::collections::VecDeque;
 use alloc::string::String;
 use alloc::sync::Arc;
+use alloc::vec::Vec;
 use core::cell::RefCell;
 use core::cmp::min;
 use core::num::NonZero;
@@ -16,15 +17,17 @@ use core::sync::atomic::{AtomicBool, Ordering};
 /// Implementation of `uma_zone` structure.
 pub struct UmaZone {
     bucket_enable: Arc<AtomicBool>,
+    bucket_keys: Arc<Vec<usize>>,
+    bucket_zones: Arc<Vec<UmaZone>>,
     ty: ZoneType,
-    size: NonZero<usize>,                     // uz_size
-    caches: CpuLocal<RefCell<UmaCache>>,      // uz_cpu
-    full_buckets: Gutex<VecDeque<UmaBucket>>, // uz_full_bucket
-    free_buckets: Gutex<VecDeque<UmaBucket>>, // uz_free_bucket
-    alloc_count: Gutex<u64>,                  // uz_allocs
-    free_count: Gutex<u64>,                   // uz_frees
-    count: Gutex<usize>,                      // uz_count
-    flags: UmaFlags,                          // uz_flags
+    size: NonZero<usize>,                                           // uz_size
+    caches: CpuLocal<RefCell<UmaCache>>,                            // uz_cpu
+    full_buckets: Gutex<VecDeque<UmaBox<UmaBucket<[BucketItem]>>>>, // uz_full_bucket
+    free_buckets: Gutex<VecDeque<UmaBox<UmaBucket<[BucketItem]>>>>, // uz_free_bucket
+    alloc_count: Gutex<u64>,                                        // uz_allocs
+    free_count: Gutex<u64>,                                         // uz_frees
+    count: Gutex<usize>,                                            // uz_count
+    flags: UmaFlags,                                                // uz_flags
 }
 
 impl UmaZone {
@@ -36,8 +39,11 @@ impl UmaZone {
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13D490|
+    #[allow(clippy::too_many_arguments)] // TODO: Find a better way.
     pub(super) fn new(
         bucket_enable: Arc<AtomicBool>,
+        bucket_keys: Arc<Vec<usize>>,
+        bucket_zones: Arc<Vec<UmaZone>>,
         name: impl Into<String>,
         keg: Option<UmaKeg>,
         size: NonZero<usize>,
@@ -100,6 +106,8 @@ impl UmaZone {
 
         Self {
             bucket_enable,
+            bucket_keys,
+            bucket_zones,
             ty,
             size: keg.size(),
             caches: CpuLocal::new(|_| RefCell::default()),
@@ -122,13 +130,14 @@ impl UmaZone {
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13E750|
-    pub fn alloc(&self) -> *mut u8 {
-        // Our implementation imply M_WAITOK and M_ZERO. Beware that we can't call into global
-        // allocator here otherwise it will cause a recursive call, which will end up panic.
-        let td = current_thread();
+    pub fn alloc(&self, flags: Alloc) -> *mut u8 {
+        if flags.has(Alloc::Wait) {
+            // TODO: The Orbis also modify td_pflags on a certain condition.
+            let td = current_thread();
 
-        if !td.can_sleep() {
-            panic!("heap allocation in a non-sleeping context is not supported");
+            if !td.can_sleep() {
+                panic!("attempt to do waitable heap allocation in a non-sleeping context");
+            }
         }
 
         loop {
@@ -184,6 +193,10 @@ impl UmaZone {
                     | ZoneType::Mbuf
                     | ZoneType::MbufCluster
             ) {
+                if flags.has(Alloc::Wait) {
+                    todo!()
+                }
+
                 todo!()
             }
 
@@ -200,8 +213,8 @@ impl UmaZone {
                 *count += 1;
             }
 
-            if self.alloc_bucket(frees) {
-                return self.alloc_item();
+            if self.alloc_bucket(frees, count, flags) {
+                return self.alloc_item(flags);
             }
         }
     }
@@ -229,16 +242,30 @@ impl UmaZone {
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13EBA0|
-    fn alloc_bucket(&self, frees: GutexWrite<VecDeque<UmaBucket>>) -> bool {
+    fn alloc_bucket(
+        &self,
+        frees: GutexWrite<VecDeque<UmaBox<UmaBucket<[BucketItem]>>>>,
+        count: GutexWrite<usize>,
+        flags: Alloc,
+    ) -> bool {
         match frees.front() {
             Some(_) => todo!(),
             None => {
                 if self.bucket_enable.load(Ordering::Relaxed) {
+                    // Get allocation flags.
+                    let mut flags = flags & !Alloc::Zero;
+
                     if self.flags.has(UmaFlags::CacheOnly) {
-                        todo!()
-                    } else {
-                        todo!()
+                        flags |= Alloc::NoVm;
                     }
+
+                    // Alloc a bucket.
+                    let i = (*count + 15) >> Uma::BUCKET_SHIFT;
+                    let k = self.bucket_keys[i];
+
+                    self.bucket_zones[k].alloc_item(flags);
+
+                    todo!()
                 }
             }
         }
@@ -252,7 +279,7 @@ impl UmaZone {
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13DD50|
-    fn alloc_item(&self) -> *mut u8 {
+    fn alloc_item(&self, _: Alloc) -> *mut u8 {
         todo!()
     }
 }
@@ -276,8 +303,8 @@ enum ZoneType {
 /// 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
+    alloc: Option<UmaBox<UmaBucket<[BucketItem]>>>, // uc_allocbucket
+    free: Option<UmaBox<UmaBucket<[BucketItem]>>>,  // uc_freebucket
+    allocs: u64,                                    // uc_allocs
+    frees: u64,                                     // uc_frees
 }