Skip to main content

ax_alloc/
lib.rs

1//! [ArceOS](https://github.com/arceos-org/arceos) global memory allocator.
2//!
3//! It provides [`GlobalAllocator`], which implements the trait
4//! [`core::alloc::GlobalAlloc`]. A static global variable of type
5//! [`GlobalAllocator`] is defined with the `#[global_allocator]` attribute, to
6//! be registered as the standard library's default allocator.
7
8#![no_std]
9
10#[allow(unused_imports)]
11#[macro_use]
12extern crate log;
13extern crate alloc;
14
15use core::{alloc::Layout, fmt, ptr::NonNull};
16
17use ax_errno::AxError;
18use strum::{IntoStaticStr, VariantArray};
19
20const PAGE_SIZE: usize = 0x1000;
21
22/// A function that tries to reclaim physical pages (e.g. by evicting
23/// clean file-backed page cache pages). Returns the number of pages freed.
24pub type PageReclaimFn = fn(num_pages: usize) -> usize;
25
26static PAGE_RECLAIM_FN: ax_kspin::SpinNoIrq<Option<PageReclaimFn>> = ax_kspin::SpinNoIrq::new(None);
27
28/// Register a callback that the allocator will invoke when a page allocation
29/// cannot be satisfied.
30pub fn register_page_reclaim_fn(f: PageReclaimFn) {
31    *PAGE_RECLAIM_FN.lock() = Some(f);
32}
33
34/// Try to reclaim physical pages by invoking the registered callback.
35/// Returns the number of pages actually freed.
36///
37/// The `SpinNoIrq` guard is released before calling into the reclaim
38/// function so that the reclaim path (and any evict listeners it
39/// triggers) runs with interrupts enabled.
40pub fn try_page_reclaim(num_pages: usize) -> usize {
41    let reclaim_fn = { *PAGE_RECLAIM_FN.lock() };
42    reclaim_fn.map_or(0, |f| f(num_pages))
43}
44
45mod page;
46pub use page::GlobalPage;
47
48/// Tracking of memory usage, enabled with the `tracking` feature.
49#[cfg(feature = "tracking")]
50pub mod tracking;
51
52/// Kinds of memory usage for tracking.
53#[repr(u8)]
54#[derive(Debug, Clone, Copy, PartialEq, Eq, VariantArray, IntoStaticStr)]
55pub enum UsageKind {
56    /// Heap allocations made by kernel Rust code.
57    RustHeap,
58    /// Virtual memory, usually used for user space.
59    VirtMem,
60    /// Page cache for file systems.
61    PageCache,
62    /// Page tables.
63    PageTable,
64    /// DMA memory.
65    Dma,
66    /// Memory used by [`GlobalPage`].
67    Global,
68}
69
70/// Statistics of memory usages.
71#[derive(Clone, Copy)]
72pub struct Usages([usize; UsageKind::VARIANTS.len()]);
73
74impl Usages {
75    const fn new() -> Self {
76        Self([0; UsageKind::VARIANTS.len()])
77    }
78
79    #[allow(dead_code)]
80    fn alloc(&mut self, kind: UsageKind, size: usize) {
81        self.0[kind as usize] += size;
82    }
83
84    #[allow(dead_code)]
85    fn dealloc(&mut self, kind: UsageKind, size: usize) {
86        self.0[kind as usize] -= size;
87    }
88
89    /// Get the memory usage for a specific kind.
90    pub fn get(&self, kind: UsageKind) -> usize {
91        self.0[kind as usize]
92    }
93}
94
95impl fmt::Debug for Usages {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        let mut d = f.debug_struct("UsageStats");
98        for &kind in UsageKind::VARIANTS {
99            d.field(kind.into(), &self.0[kind as usize]);
100        }
101        d.finish()
102    }
103}
104
105/// The error type used for allocation operations in `ax-alloc`.
106#[derive(Debug, Clone, Copy, PartialEq, Eq)]
107pub enum AllocError {
108    /// Invalid size, alignment, or other input parameter.
109    InvalidParam,
110    /// The allocator has already been initialized.
111    AlreadyInitialized,
112    /// A region overlaps with an existing managed region.
113    MemoryOverlap,
114    /// Not enough memory is available to satisfy the request.
115    NoMemory,
116    /// Attempted to deallocate memory that was not allocated.
117    NotAllocated,
118    /// The allocator has not been initialized.
119    NotInitialized,
120    /// The requested address or entity was not found.
121    NotFound,
122}
123
124/// A [`Result`] alias with [`AllocError`] as the error type.
125pub type AllocResult<T = ()> = Result<T, AllocError>;
126
127impl From<AllocError> for AxError {
128    fn from(value: AllocError) -> Self {
129        match value {
130            AllocError::NoMemory => AxError::NoMemory,
131            AllocError::NotFound => AxError::NotFound,
132            AllocError::NotInitialized | AllocError::AlreadyInitialized => AxError::BadState,
133            AllocError::MemoryOverlap => AxError::AlreadyExists,
134            AllocError::InvalidParam | AllocError::NotAllocated => AxError::InvalidInput,
135        }
136    }
137}
138
139/// Unified allocator operations provided by all `ax-alloc` backends.
140pub trait AllocatorOps {
141    /// Returns the allocator name.
142    fn name(&self) -> &'static str;
143
144    /// Initializes the allocator with the given region.
145    fn init(&self, start_vaddr: usize, size: usize) -> AllocResult;
146
147    /// Adds an extra memory region to the allocator.
148    fn add_memory(&self, start_vaddr: usize, size: usize) -> AllocResult;
149
150    /// Allocates arbitrary bytes.
151    fn alloc(&self, layout: Layout) -> AllocResult<NonNull<u8>>;
152
153    /// Deallocates a prior byte allocation.
154    fn dealloc(&self, pos: NonNull<u8>, layout: Layout);
155
156    /// Allocates contiguous pages.
157    ///
158    /// `align` is the requested byte alignment, not a log2/exponent.
159    /// It must be a power-of-two byte alignment accepted by the backend page allocator.
160    fn alloc_pages(&self, num_pages: usize, align: usize, kind: UsageKind) -> AllocResult<usize>;
161
162    /// Allocates contiguous DMA32 pages.
163    ///
164    /// `align` is the requested byte alignment, not a log2/exponent.
165    /// It must be a power-of-two byte alignment accepted by the backend page allocator.
166    fn alloc_dma32_pages(
167        &self,
168        num_pages: usize,
169        align: usize,
170        kind: UsageKind,
171    ) -> AllocResult<usize>;
172
173    /// Allocates contiguous pages starting from the given address.
174    ///
175    /// `align` is the requested byte alignment, not a log2/exponent.
176    /// It must be a power-of-two byte alignment accepted by the backend page allocator.
177    fn alloc_pages_at(
178        &self,
179        start: usize,
180        num_pages: usize,
181        align: usize,
182        kind: UsageKind,
183    ) -> AllocResult<usize>;
184
185    /// Deallocates a prior page allocation.
186    fn dealloc_pages(&self, pos: usize, num_pages: usize, kind: UsageKind);
187
188    /// Returns used byte count.
189    fn used_bytes(&self) -> usize;
190
191    /// Returns available byte count.
192    fn available_bytes(&self) -> usize;
193
194    /// Returns used page count.
195    fn used_pages(&self) -> usize;
196
197    /// Returns available page count.
198    fn available_pages(&self) -> usize;
199
200    /// Returns usage statistics.
201    fn usages(&self) -> Usages;
202}
203
204// Select implementation based on build.rs-generated cfg flags.
205#[cfg(buddy_slab)]
206mod buddy_slab;
207#[cfg(not(any(tlsf, buddy_slab)))]
208mod stub_impl;
209#[cfg(tlsf)]
210mod tlsf_impl;
211
212#[cfg(buddy_slab)]
213use buddy_slab as imp;
214pub use imp::{
215    DefaultByteAllocator, GlobalAllocator, global_add_memory, global_init, init_percpu_slab,
216};
217#[cfg(not(any(tlsf, buddy_slab)))]
218use stub_impl as imp;
219#[cfg(tlsf)]
220use tlsf_impl as imp;
221
222/// Returns the reference to the global allocator.
223pub fn global_allocator() -> &'static GlobalAllocator {
224    imp::global_allocator()
225}