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