obkrnl/uma/
zone.rs

1use super::{Alloc, BucketHdr, FreeItem, Slab, StdFree, Uma, UmaBucket, UmaFlags, UmaKeg};
2use crate::context::{CpuLocal, config, current_thread};
3use crate::lock::Mutex;
4use crate::vm::Vm;
5use alloc::collections::VecDeque;
6use alloc::collections::linked_list::LinkedList;
7use alloc::string::String;
8use alloc::sync::Arc;
9use alloc::vec::Vec;
10use core::cell::RefCell;
11use core::cmp::min;
12use core::num::NonZero;
13use core::ops::DerefMut;
14use core::ptr::{NonNull, null_mut};
15use core::sync::atomic::{AtomicBool, Ordering};
16
17/// Implementation of `uma_zone` structure.
18pub struct UmaZone<T> {
19    bucket_enable: Arc<AtomicBool>,
20    bucket_keys: Arc<Vec<usize>>,
21    bucket_zones: Arc<Vec<UmaZone<StdFree>>>,
22    ty: ZoneType,
23    size: NonZero<usize>, // uz_size
24    slab: fn(&Self, &mut ZoneState<T>, Option<&Arc<UmaKeg<T>>>, Alloc) -> Option<NonNull<Slab<T>>>, // uz_slab
25    init: Option<fn(*mut u8, NonZero<usize>, Alloc) -> bool>, // uz_init
26    ctor: Option<fn(*mut u8, NonZero<usize>, Alloc) -> bool>, // uz_ctor
27    caches: CpuLocal<RefCell<UmaCache>>,                      // uz_cpu
28    flags: UmaFlags,                                          // uz_flags
29    state: Mutex<ZoneState<T>>,
30}
31
32impl<T: FreeItem> UmaZone<T> {
33    const ALIGN_CACHE: usize = 63; // uma_align_cache
34
35    /// See `zone_ctor` on Orbis for a reference.
36    ///
37    /// # Reference offsets
38    /// | Version | Offset |
39    /// |---------|--------|
40    /// |PS4 11.00|0x13D490|
41    #[allow(clippy::too_many_arguments)] // TODO: Find a better way.
42    pub(super) fn new(
43        vm: Arc<Vm>,
44        bucket_enable: Arc<AtomicBool>,
45        bucket_keys: Arc<Vec<usize>>,
46        bucket_zones: Arc<Vec<UmaZone<StdFree>>>,
47        name: impl Into<String>,
48        keg: Option<UmaKeg<T>>,
49        size: NonZero<usize>,
50        align: Option<usize>,
51        init: Option<fn()>,
52        flags: impl Into<UmaFlags>,
53    ) -> Self {
54        let name = name.into();
55        let flags = flags.into();
56        let (keg, mut flags) = if flags.has_any(UmaFlags::Secondary) {
57            todo!()
58        } else {
59            // We use a different approach here to make it idiomatic to Rust. On Orbis it will
60            // construct a keg here if it is passed from the caller. If not it will allocate a new
61            // keg from masterzone_k.
62            let keg = match keg {
63                Some(v) => v,
64                None => UmaKeg::new(vm, size, align.unwrap_or(Self::ALIGN_CACHE), init, flags),
65            };
66
67            (keg, UmaFlags::zeroed())
68        };
69
70        // Get type and uz_count.
71        let mut ty = ZoneType::Other;
72        let mut count = 0;
73
74        if !keg.flags().has_any(UmaFlags::Internal) {
75            count = if !keg.flags().has_any(UmaFlags::MaxBucket) {
76                min(keg.item_per_slab(), Uma::BUCKET_MAX)
77            } else {
78                Uma::BUCKET_MAX
79            };
80
81            match name.as_str() {
82                "mbuf_packet" => {
83                    ty = ZoneType::MbufPacket;
84                    count = 4;
85                }
86                "mbuf_cluster_pack" => {
87                    ty = ZoneType::MbufClusterPack;
88                    count = Uma::BUCKET_MAX;
89                }
90                "mbuf_jumbo_page" => {
91                    ty = ZoneType::MbufJumboPage;
92                    count = 1;
93                }
94                "mbuf" => {
95                    ty = ZoneType::Mbuf;
96                    count = 16;
97                }
98                "mbuf_cluster" => {
99                    ty = ZoneType::MbufCluster;
100                    count = 1;
101                }
102                _ => (),
103            }
104        }
105
106        // Construct uma_zone.
107        let inherit = UmaFlags::Offpage
108            | UmaFlags::Malloc
109            | UmaFlags::Hash
110            | UmaFlags::VToSlab
111            | UmaFlags::Bucket
112            | UmaFlags::Internal
113            | UmaFlags::CacheOnly;
114
115        flags |= keg.flags() & inherit;
116
117        Self {
118            bucket_enable,
119            bucket_keys,
120            bucket_zones,
121            ty,
122            size: keg.size(),
123            slab: Self::fetch_slab,
124            init: None,
125            ctor: None,
126            caches: CpuLocal::new(|_| RefCell::default()),
127            flags,
128            state: Mutex::new(ZoneState {
129                kegs: LinkedList::from([Arc::new(keg)]),
130                full_buckets: VecDeque::default(),
131                free_buckets: VecDeque::default(),
132                alloc_count: 0,
133                free_count: 0,
134                count,
135                fills: 0,
136            }),
137        }
138    }
139}
140
141impl<T> UmaZone<T> {
142    pub fn size(&self) -> NonZero<usize> {
143        self.size
144    }
145
146    /// See `uma_zalloc_arg` on the Orbis for a reference.
147    ///
148    /// # Reference offsets
149    /// | Version | Offset |
150    /// |---------|--------|
151    /// |PS4 11.00|0x13E750|
152    pub fn alloc(&self, flags: Alloc) -> *mut u8 {
153        if flags.has_any(Alloc::Wait) {
154            // TODO: The Orbis also modify td_pflags on a certain condition.
155            let td = current_thread();
156
157            if !td.can_sleep() {
158                panic!("attempt to do waitable heap allocation in a non-sleeping context");
159            }
160        }
161
162        loop {
163            // Try allocate from per-CPU cache first so we don't need to acquire a mutex lock.
164            let caches = self.caches.lock();
165            let mem = Self::alloc_from_cache(caches.borrow_mut().deref_mut());
166
167            if !mem.is_null() {
168                return mem;
169            }
170
171            drop(caches); // Exit from non-sleeping context before acquire the mutex.
172
173            // Cache not found, allocate from the zone. We need to re-check the cache again because
174            // we may on a different CPU since we drop the CPU pinning on the above.
175            let mut state = self.state.lock();
176            let caches = self.caches.lock();
177            let mut cache = caches.borrow_mut();
178            let mem = Self::alloc_from_cache(&mut cache);
179
180            if !mem.is_null() {
181                return mem;
182            }
183
184            // TODO: What actually we are doing here?
185            state.alloc_count += core::mem::take(&mut cache.allocs);
186            state.free_count += core::mem::take(&mut cache.frees);
187
188            if let Some(b) = cache.alloc.take() {
189                state.free_buckets.push_front(b);
190            }
191
192            if let Some(b) = state.full_buckets.pop_front() {
193                cache.alloc = Some(b);
194
195                // Seems like this should never fail.
196                let m = Self::alloc_from_cache(&mut cache);
197
198                assert!(!m.is_null());
199
200                return m;
201            }
202
203            drop(cache);
204            drop(caches);
205
206            // TODO: What is this?
207            if matches!(
208                self.ty,
209                ZoneType::MbufPacket
210                    | ZoneType::MbufJumboPage
211                    | ZoneType::Mbuf
212                    | ZoneType::MbufCluster
213            ) {
214                if flags.has_any(Alloc::Wait) {
215                    todo!()
216                }
217
218                todo!()
219            }
220
221            // TODO: What is this?
222            if !matches!(
223                self.ty,
224                ZoneType::MbufCluster
225                    | ZoneType::Mbuf
226                    | ZoneType::MbufJumboPage
227                    | ZoneType::MbufPacket
228                    | ZoneType::MbufClusterPack
229            ) && state.count < Uma::BUCKET_MAX
230            {
231                state.count += 1;
232            }
233
234            if self.alloc_bucket(&mut state, flags) {
235                return self.alloc_item(&mut state, flags);
236            }
237        }
238    }
239
240    fn alloc_from_cache(c: &mut UmaCache) -> *mut u8 {
241        while let Some(b) = c.alloc.map(|v| v.as_ptr()) {
242            if unsafe { (*b).hdr.len != 0 } {
243                todo!()
244            }
245
246            if c.free
247                .map(|v| v.as_ptr())
248                .is_some_and(|b| unsafe { (*b).hdr.len != 0 })
249            {
250                core::mem::swap(&mut c.alloc, &mut c.free);
251                continue;
252            }
253
254            break;
255        }
256
257        null_mut()
258    }
259
260    /// See `zone_alloc_bucket` on the Orbis for a reference.
261    ///
262    /// # Reference offsets
263    /// | Version | Offset |
264    /// |---------|--------|
265    /// |PS4 11.00|0x13EBA0|
266    fn alloc_bucket(&self, state: &mut ZoneState<T>, flags: Alloc) -> bool {
267        // Get bucket.
268        let b = match state.free_buckets.front() {
269            Some(_) => todo!(),
270            None => {
271                if self.bucket_enable.load(Ordering::Relaxed) {
272                    // Get allocation flags. On Orbis it will remove M_ZERO from the flags but we do
273                    // the opposite to eliminate the chance of dangling pointer in bucket items.
274                    let mut flags = flags | Alloc::Zero;
275
276                    if self.flags.has_any(UmaFlags::CacheOnly) {
277                        flags |= Alloc::NoVm;
278                    }
279
280                    // Alloc a bucket.
281                    let i = (state.count + 15) >> Uma::BUCKET_SHIFT;
282                    let k = self.bucket_keys[i];
283                    let b = &self.bucket_zones[k];
284                    let b = b.alloc_item(&mut b.state.lock(), flags);
285
286                    if b.is_null() {
287                        todo!()
288                    }
289
290                    // Initialize bucket.
291                    let h = BucketHdr { len: 0 };
292
293                    unsafe { core::ptr::write(b.cast(), h) };
294
295                    core::ptr::slice_from_raw_parts_mut(b, Uma::BUCKET_SIZES[k]) as *mut UmaBucket
296                } else {
297                    todo!()
298                }
299            }
300        };
301
302        // SAFETY: We have exclusive access to the bucket.
303        let b = unsafe { &mut *b };
304
305        if state.fills < config().cpu_count().get().into() {
306            let n = min(b.items.len(), state.count);
307            let mut k = None;
308            let mut f = flags;
309
310            state.fills += 1;
311
312            while b.hdr.len < n {
313                let s = match (self.slab)(self, state, k.as_ref(), f) {
314                    Some(v) => v.as_ptr(),
315                    None => todo!(),
316                };
317
318                while unsafe { (*s).hdr.free_count != 0 && b.hdr.len < n } {
319                    let i = unsafe { (*s).alloc_item() };
320
321                    b.items[b.hdr.len] = i;
322                    b.hdr.len += 1;
323                }
324
325                k = unsafe { Some((*s).hdr.keg.clone()) };
326                f |= Alloc::NoWait;
327            }
328
329            todo!()
330        }
331
332        true
333    }
334
335    /// See `zone_alloc_item` on the Orbis for a reference.
336    ///
337    /// # Reference offsets
338    /// | Version | Offset |
339    /// |---------|--------|
340    /// |PS4 11.00|0x13DD50|
341    fn alloc_item(&self, state: &mut ZoneState<T>, flags: Alloc) -> *mut u8 {
342        // Get a slab.
343        let slab = (self.slab)(self, state, None, flags);
344
345        if let Some(mut slab) = slab {
346            let item = unsafe { slab.as_mut().alloc_item() };
347
348            state.alloc_count += 1;
349
350            if self.init.is_none_or(|f| f(item, self.size, flags)) {
351                if self.ctor.is_none_or(|f| f(item, self.size, flags)) {
352                    if flags.has_any(Alloc::Zero) {
353                        unsafe { item.write_bytes(0, self.size.get()) };
354                    }
355
356                    return item;
357                } else {
358                    todo!()
359                }
360            } else {
361                todo!()
362            }
363        }
364
365        todo!()
366    }
367}
368
369impl<T: FreeItem> UmaZone<T> {
370    /// See `zone_fetch_slab` on the Orbis for a reference.
371    ///
372    /// # Reference offsets
373    /// | Version | Offset |
374    /// |---------|--------|
375    /// |PS4 11.00|0x141DB0|
376    fn fetch_slab(
377        &self,
378        state: &mut ZoneState<T>,
379        keg: Option<&Arc<UmaKeg<T>>>,
380        flags: Alloc,
381    ) -> Option<NonNull<Slab<T>>> {
382        let keg = keg.unwrap_or(state.kegs.front().unwrap());
383
384        if !keg.flags().has_any(UmaFlags::Bucket) || keg.recurse() == 0 {
385            loop {
386                if let Some(v) = keg.fetch_slab(flags) {
387                    return Some(v);
388                }
389
390                if flags.has_any(Alloc::NoWait | Alloc::NoVm) {
391                    break;
392                }
393            }
394        }
395
396        None
397    }
398}
399
400/// Contains mutable data for [UmaZone].
401struct ZoneState<T> {
402    kegs: LinkedList<Arc<UmaKeg<T>>>,           // uz_kegs + uz_klink
403    full_buckets: VecDeque<NonNull<UmaBucket>>, // uz_full_bucket
404    free_buckets: VecDeque<NonNull<UmaBucket>>, // uz_free_bucket
405    alloc_count: u64,                           // uz_allocs
406    free_count: u64,                            // uz_frees
407    count: usize,                               // uz_count
408    fills: u16,                                 // uz_fills
409}
410
411unsafe impl<T: Send> Send for ZoneState<T> {}
412
413/// Type of [UmaZone].
414#[derive(Clone, Copy)]
415enum ZoneType {
416    Other,
417    /// `zone_pack`.
418    MbufPacket,
419    /// `zone_jumbop`.
420    MbufJumboPage,
421    /// `zone_mbuf`.
422    Mbuf,
423    /// `zone_clust`.
424    MbufCluster,
425    /// `zone_clust_pack`.
426    MbufClusterPack,
427}
428
429/// Implementation of `uma_cache` structure.
430#[derive(Default)]
431struct UmaCache {
432    alloc: Option<NonNull<UmaBucket>>, // uc_allocbucket
433    free: Option<NonNull<UmaBucket>>,  // uc_freebucket
434    allocs: u64,                       // uc_allocs
435    frees: u64,                        // uc_frees
436}
437
438unsafe impl Send for UmaCache {}
obkrnl/uma/zone.rs

obkrnl/uma/
zone.rs
