xref: /DragonOS/kernel/crates/rust-slabmalloc/src/zone.rs (revision c635d8a9cfe25bc11779f323ef0c7d7a0f597d4a)
1 //! A ZoneAllocator to allocate arbitrary object sizes (up to `ZoneAllocator::MAX_ALLOC_SIZE`)
2 //!
3 //! The ZoneAllocator achieves this by having many `SCAllocator`
4 
5 use crate::*;
6 use core::sync::atomic::Ordering;
7 
8 /// Creates an instance of a zone, we do this in a macro because we
9 /// re-use the code in const and non-const functions
10 ///
11 /// We can get rid of this once the const fn feature is fully stabilized.
12 macro_rules! new_zone {
13     () => {
14         ZoneAllocator {
15             // TODO(perf): We should probably pick better classes
16             // rather than powers-of-two (see SuperMalloc etc.)
17             small_slabs: [
18                 SCAllocator::new(1 << 3),  // 8
19                 SCAllocator::new(1 << 4),  // 16
20                 SCAllocator::new(1 << 5),  // 32
21                 SCAllocator::new(1 << 6),  // 64
22                 SCAllocator::new(1 << 7),  // 128
23                 SCAllocator::new(1 << 8),  // 256
24                 SCAllocator::new(1 << 9),  // 512
25                 SCAllocator::new(1 << 10), // 1024
26                 SCAllocator::new(1 << 11), // 2048            ],
27             ],
28             total: 0,
29         }
30     };
31 }
32 
33 /// A zone allocator for arbitrary sized allocations.
34 ///
35 /// Has a bunch of `SCAllocator` and through that can serve allocation
36 /// requests for many different object sizes up to (MAX_SIZE_CLASSES) by selecting
37 /// the right `SCAllocator` for allocation and deallocation.
38 ///
39 /// The allocator provides to refill functions `refill` and `refill_large`
40 /// to provide the underlying `SCAllocator` with more memory in case it runs out.
41 pub struct ZoneAllocator<'a> {
42     small_slabs: [SCAllocator<'a, ObjectPage<'a>>; ZoneAllocator::MAX_BASE_SIZE_CLASSES],
43     total: u64,
44 }
45 
46 impl<'a> Default for ZoneAllocator<'a> {
47     fn default() -> ZoneAllocator<'a> {
48         new_zone!()
49     }
50 }
51 
52 enum Slab {
53     Base(usize),
54     Unsupported,
55 }
56 
57 impl<'a> ZoneAllocator<'a> {
58     /// Maximum size which is allocated with ObjectPages (4 KiB pages).
59     ///
60     /// e.g. this is 4 KiB - 80 bytes of meta-data.
61     pub const MAX_BASE_ALLOC_SIZE: usize = 1 << 11;
62 
63     /// How many allocators of type SCAllocator<ObjectPage> we have.
64     pub const MAX_BASE_SIZE_CLASSES: usize = 9;
65 
66     #[cfg(feature = "unstable")]
67     pub const fn new() -> ZoneAllocator<'a> {
68         new_zone!()
69     }
70 
71     #[cfg(not(feature = "unstable"))]
72     pub fn new() -> ZoneAllocator<'a> {
73         new_zone!()
74     }
75 
76     /// Return maximum size an object of size `current_size` can use.
77     ///
78     /// Used to optimize `realloc`.
79     pub fn get_max_size(current_size: usize) -> Option<usize> {
80         match current_size {
81             0..=8 => Some(8),
82             9..=16 => Some(16),
83             17..=32 => Some(32),
84             33..=64 => Some(64),
85             65..=128 => Some(128),
86             129..=256 => Some(256),
87             257..=512 => Some(512),
88             513..=1024 => Some(1024),
89             1025..=2048 => Some(2048),
90             _ => None,
91         }
92     }
93 
94     /// Figure out index into zone array to get the correct slab allocator for that size.
95     fn get_slab(requested_size: usize) -> Slab {
96         match requested_size {
97             0..=8 => Slab::Base(0),
98             9..=16 => Slab::Base(1),
99             17..=32 => Slab::Base(2),
100             33..=64 => Slab::Base(3),
101             65..=128 => Slab::Base(4),
102             129..=256 => Slab::Base(5),
103             257..=512 => Slab::Base(6),
104             513..=1024 => Slab::Base(7),
105             1025..=2048 => Slab::Base(8),
106             _ => Slab::Unsupported,
107         }
108     }
109 
110     /// Reclaims empty pages by calling `dealloc` on it and removing it from the
111     /// empty lists in the [`SCAllocator`].
112     ///
113     /// The `dealloc` function is called at most `reclaim_base_max` times for
114     /// base pages, and at most `reclaim_large_max` for large pages.
115     pub fn try_reclaim_base_pages<F>(&mut self, mut to_reclaim: usize, mut dealloc: F)
116     where
117         F: Fn(*mut ObjectPage),
118     {
119         for i in 0..ZoneAllocator::MAX_BASE_SIZE_CLASSES {
120             let slab = &mut self.small_slabs[i];
121             // reclaim的page数
122             let just_reclaimed = slab.try_reclaim_pages(to_reclaim, &mut dealloc);
123             self.total -= (just_reclaimed * OBJECT_PAGE_SIZE) as u64;
124             to_reclaim = to_reclaim.saturating_sub(just_reclaimed);
125             if to_reclaim == 0 {
126                 break;
127             }
128         }
129     }
130 
131     /// 获取scallocator中的还未被分配的空间
132     pub fn free_space(&mut self) -> u64 {
133         // 记录空闲空间
134         let mut free = 0;
135         // 遍历所有scallocator
136         for count in 0..ZoneAllocator::MAX_BASE_SIZE_CLASSES {
137             // 获取scallocator
138             let scallocator = &mut self.small_slabs[count];
139 
140             // 遍历scallocator中的部分分配的page(partial_page)
141             for slab_page in scallocator.slabs.iter_mut() {
142                 // 统计page中还可以分配多少个object
143                 let mut free_obj_count = 0;
144                 // 遍历page中的bitfield(用来统计内存分配情况的u64数组)
145                 for b in slab_page.bitfield().iter() {
146                     let bitval = b.load(Ordering::Relaxed);
147                     let free_count = bitval.count_zeros() as usize;
148                     free_obj_count += free_count;
149                 }
150                 // 剩余可分配object数乘上page中规定的每个object的大小,即空闲空间
151                 free += free_obj_count * scallocator.size();
152             }
153             // 遍历scallocator中的empty_page,把空页空间也加上去
154             free +=
155                 scallocator.empty_slabs.elements * (scallocator.obj_per_page * scallocator.size());
156         }
157         free as u64
158     }
159 
160     pub fn usage(&mut self) -> SlabUsage {
161         let free_num = self.free_space();
162         SlabUsage::new(self.total, free_num)
163     }
164 }
165 
166 unsafe impl<'a> crate::Allocator<'a> for ZoneAllocator<'a> {
167     /// Allocate a pointer to a block of memory described by `layout`.
168     fn allocate(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocationError> {
169         match ZoneAllocator::get_slab(layout.size()) {
170             Slab::Base(idx) => self.small_slabs[idx].allocate(layout),
171             Slab::Unsupported => Err(AllocationError::InvalidLayout),
172         }
173     }
174 
175     /// Deallocates a pointer to a block of memory, which was
176     /// previously allocated by `allocate`.
177     ///
178     /// # Arguments
179     ///  * `ptr` - Address of the memory location to free.
180     ///  * `layout` - Memory layout of the block pointed to by `ptr`.
181     fn deallocate(&mut self, ptr: NonNull<u8>, layout: Layout) -> Result<(), AllocationError> {
182         match ZoneAllocator::get_slab(layout.size()) {
183             Slab::Base(idx) => self.small_slabs[idx].deallocate(ptr, layout),
184             Slab::Unsupported => Err(AllocationError::InvalidLayout),
185         }
186     }
187 
188     /// Refills the SCAllocator for a given Layout with an ObjectPage.
189     ///
190     /// # Safety
191     /// ObjectPage needs to be emtpy etc.
192     unsafe fn refill(
193         &mut self,
194         layout: Layout,
195         new_page: &'a mut ObjectPage<'a>,
196     ) -> Result<(), AllocationError> {
197         match ZoneAllocator::get_slab(layout.size()) {
198             Slab::Base(idx) => {
199                 self.small_slabs[idx].refill(new_page);
200                 // 每refill一个page就为slab的总空间统计加上4KB
201                 self.total += OBJECT_PAGE_SIZE as u64;
202                 Ok(())
203             }
204             Slab::Unsupported => Err(AllocationError::InvalidLayout),
205         }
206     }
207 }
208 
209 /// Slab内存空间使用情况
210 pub struct SlabUsage {
211     // slab总共使用的内存空间
212     total: u64,
213     // slab的空闲空间
214     free: u64,
215 }
216 
217 impl SlabUsage {
218     /// 初始化SlabUsage
219     pub fn new(total: u64, free: u64) -> Self {
220         Self { total, free }
221     }
222 
223     pub fn total(&self) -> u64 {
224         self.total
225     }
226 
227     pub fn used(&self) -> u64 {
228         self.total - self.free
229     }
230 
231     pub fn free(&self) -> u64 {
232         self.free
233     }
234 }
235