jonaskruckenberg.de
Next Generation WASM Microkernel Operating System
1// Copyright 2025 Jonas Kruckenberg
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use core::alloc::Allocator;
9use core::ops::DerefMut;
10use core::ptr::NonNull;
11use core::{cmp, mem};
12
13use anyhow::Context;
14use cranelift_entity::PrimaryMap;
15
16use crate::mem::Mmap;
17use crate::wasm::indices::{DefinedMemoryIndex, DefinedTableIndex};
18use crate::wasm::module::Module;
19use crate::wasm::translate::TranslatedModule;
20use crate::wasm::utils::round_usize_up_to_host_pages;
21use crate::wasm::vm::instance::Instance;
22use crate::wasm::vm::mmap_vec::MmapVec;
23use crate::wasm::vm::{InstanceHandle, TableElement, VMShape};
24use crate::wasm::{MEMORY_MAX, TABLE_MAX, translate, vm};
25
26/// A type that knows how to allocate backing memory for instance resources.
27pub trait InstanceAllocator {
28 unsafe fn allocate_instance_and_vmctx(
29 &self,
30 vmshape: &VMShape,
31 ) -> crate::Result<NonNull<Instance>>;
32 unsafe fn deallocate_instance_and_vmctx(&self, instance: NonNull<Instance>, vmshape: &VMShape);
33
34 /// Allocate a memory for an instance.
35 ///
36 /// # Errors
37 ///
38 /// Returns an error if any of the allocations fail.
39 ///
40 /// # Safety
41 ///
42 /// The safety of the entire VM depends on the correct implementation of this method.
43 unsafe fn allocate_memory(
44 &self,
45 memory: &translate::Memory,
46 memory_index: DefinedMemoryIndex,
47 ) -> crate::Result<vm::Memory>;
48
49 /// Deallocate an instance's previously allocated memory.
50 ///
51 /// # Safety
52 ///
53 /// The memory must have previously been allocated by
54 /// `Self::allocate_memory`, be at the given index, and must currently be
55 /// allocated. It must never be used again.
56 unsafe fn deallocate_memory(&self, memory_index: DefinedMemoryIndex, memory: vm::Memory);
57
58 /// Allocate a table for an instance.
59 ///
60 /// # Errors
61 ///
62 /// Returns an error if any of the allocations fail.
63 ///
64 /// # Safety
65 ///
66 /// The safety of the entire VM depends on the correct implementation of this method.
67 unsafe fn allocate_table(
68 &self,
69 table_desc: &translate::Table,
70 table_index: DefinedTableIndex,
71 ) -> crate::Result<vm::Table>;
72
73 /// Deallocate an instance's previously allocated table.
74 ///
75 /// # Safety
76 ///
77 /// The table must have previously been allocated by `Self::allocate_table`,
78 /// be at the given index, and must currently be allocated. It must never be
79 /// used again.
80 unsafe fn deallocate_table(&self, table_index: DefinedTableIndex, table: vm::Table);
81
82 /// Allocate multiple memories at once.
83 ///
84 /// By default, this will delegate the actual allocation to `Self::allocate_memory`.
85 ///
86 /// # Errors
87 ///
88 /// Returns an error if any of the allocations fail.
89 ///
90 /// # Safety
91 ///
92 /// The safety of the entire VM depends on the correct implementation of this method.
93 unsafe fn allocate_memories(
94 &self,
95 module: &TranslatedModule,
96 memories: &mut PrimaryMap<DefinedMemoryIndex, vm::Memory>,
97 ) -> crate::Result<()> {
98 for (index, plan) in &module.memories {
99 if let Some(def_index) = module.defined_memory_index(index) {
100 let new_def_index =
101 // Safety: caller has to ensure safety
102 memories.push(unsafe { self.allocate_memory(plan, def_index)? });
103 debug_assert_eq!(def_index, new_def_index);
104 }
105 }
106 Ok(())
107 }
108
109 /// Allocate multiple tables at once.
110 ///
111 /// By default, this will delegate the actual allocation to `Self::allocate_table`.
112 ///
113 /// # Errors
114 ///
115 /// Returns an error if any of the allocations fail.
116 ///
117 /// # Safety
118 ///
119 /// The safety of the entire VM depends on the correct implementation of this method.
120 unsafe fn allocate_tables(
121 &self,
122 module: &TranslatedModule,
123 tables: &mut PrimaryMap<DefinedTableIndex, vm::Table>,
124 ) -> crate::Result<()> {
125 for (index, plan) in &module.tables {
126 if let Some(def_index) = module.defined_table_index(index) {
127 let new_def_index =
128 // Safety: caller has to ensure safety
129 tables.push(unsafe { self.allocate_table(plan, def_index)? });
130 debug_assert_eq!(def_index, new_def_index);
131 }
132 }
133 Ok(())
134 }
135
136 /// Deallocates multiple memories at once.
137 ///
138 /// By default, this will delegate the actual deallocation to `Self::deallocate_memory`.
139 ///
140 /// # Safety
141 ///
142 /// Just like `Self::deallocate_memory` all memories must have been allocated by
143 /// `Self::allocate_memories`/`Self::allocate_memory` and must never be used again.
144 unsafe fn deallocate_memories(
145 &self,
146 memories: &mut PrimaryMap<DefinedMemoryIndex, vm::Memory>,
147 ) {
148 for (memory_index, memory) in mem::take(memories) {
149 // Because deallocating memory is infallible, we don't need to worry
150 // about leaking subsequent memories if the first memory failed to
151 // deallocate. If deallocating memory ever becomes fallible, we will
152 // need to be careful here!
153 // Safety: ensured by caller
154 unsafe {
155 self.deallocate_memory(memory_index, memory);
156 }
157 }
158 }
159
160 /// Deallocates multiple tables at once.
161 ///
162 /// By default, this will delegate the actual deallocation to `Self::deallocate_table`.
163 ///
164 /// # Safety
165 ///
166 /// Just like `Self::deallocate_table` all tables must have been allocated by
167 /// `Self::allocate_tables`/`Self::allocate_table` and must never be used again.
168 unsafe fn deallocate_tables(&self, tables: &mut PrimaryMap<DefinedTableIndex, vm::Table>) {
169 for (table_index, table) in mem::take(tables) {
170 // Safety: ensured by caller
171 unsafe {
172 self.deallocate_table(table_index, table);
173 }
174 }
175 }
176
177 /// Allocate all resources required to instantiate a module.
178 ///
179 /// By default, this will in-turn call `Self::allocate_vmctx`, `Self::allocate_tables` and
180 /// `Self::allocate_memories` as well as perform necessary clean up.
181 ///
182 /// # Errors
183 ///
184 /// Returns an error if any of the allocations fail. In this case, the resources are cleaned up
185 /// automatically.
186 fn allocate_module(&self, module: Module) -> crate::Result<InstanceHandle> {
187 let mut tables = PrimaryMap::with_capacity(
188 usize::try_from(module.translated().num_defined_tables()).unwrap(),
189 );
190 let mut memories = PrimaryMap::with_capacity(
191 usize::try_from(module.translated().num_defined_memories()).unwrap(),
192 );
193
194 // Safety: TODO
195 match (|| unsafe {
196 self.allocate_tables(module.translated(), &mut tables)?;
197 self.allocate_memories(module.translated(), &mut memories)?;
198 self.allocate_instance_and_vmctx(module.vmshape())
199 })() {
200 // Safety: we crated the instance handle and memories/tables from the same module description so should be fine
201 Ok(instance) => Ok(unsafe { Instance::from_parts(module, instance, tables, memories) }),
202 // Safety: memories and tables have just been allocated and will not be handed out
203 Err(e) => unsafe {
204 self.deallocate_memories(&mut memories);
205 self.deallocate_tables(&mut tables);
206 Err(e)
207 },
208 }
209 }
210
211 unsafe fn deallocate_module(&self, handle: &mut InstanceHandle) {
212 // Safety: ensured by caller
213 unsafe {
214 self.deallocate_memories(&mut handle.instance_mut().memories);
215 self.deallocate_tables(&mut handle.instance_mut().tables);
216 self.deallocate_instance_and_vmctx(handle.as_non_null(), handle.instance().vmshape());
217 }
218 }
219}
220
221pub struct PlaceholderAllocatorDontUse;
222impl InstanceAllocator for PlaceholderAllocatorDontUse {
223 unsafe fn allocate_instance_and_vmctx(
224 &self,
225 vmshape: &VMShape,
226 ) -> crate::Result<NonNull<Instance>> {
227 // FIXME this shouldn't allocate from the kernel heap
228 let ptr = alloc::alloc::Global.allocate(Instance::alloc_layout(vmshape))?;
229 Ok(ptr.cast())
230 }
231
232 unsafe fn deallocate_instance_and_vmctx(&self, instance: NonNull<Instance>, vmshape: &VMShape) {
233 // Safety: `NonNull<Instance>` is only ever created above using the same global allocator
234 unsafe {
235 // FIXME this shouldn't allocate from the kernel heap
236 let layout = Instance::alloc_layout(vmshape);
237 alloc::alloc::Global.deallocate(instance.cast(), layout);
238 }
239 }
240
241 unsafe fn allocate_memory(
242 &self,
243 memory: &translate::Memory,
244 _memory_index: DefinedMemoryIndex,
245 ) -> crate::Result<vm::Memory> {
246 // TODO we could call out to some resource management instance here to obtain
247 // dynamic "minimum" and "maximum" values that reflect the state of the real systems
248 // memory consumption
249
250 // If the minimum memory size overflows the size of our own address
251 // space, then we can't satisfy this request, but defer the error to
252 // later so the `store` can be informed that an effective oom is
253 // happening.
254 let minimum = memory
255 .minimum_byte_size()
256 .ok()
257 .and_then(|m| usize::try_from(m).ok())
258 .expect("memory minimum size exceeds memory limits");
259
260 // The plan stores the maximum size in units of wasm pages, but we
261 // use units of bytes. Unlike for the `minimum` size we silently clamp
262 // the effective maximum size to the limits of what we can track. If the
263 // maximum size exceeds `usize` or `u64` then there's no need to further
264 // keep track of it as some sort of runtime limit will kick in long
265 // before we reach the statically declared maximum size.
266 let maximum = memory
267 .maximum_byte_size()
268 .ok()
269 .and_then(|m| usize::try_from(m).ok());
270
271 let offset_guard_bytes = usize::try_from(memory.offset_guard_size).unwrap();
272 // Ensure that our guard regions are multiples of the host page size.
273 let offset_guard_bytes = round_usize_up_to_host_pages(offset_guard_bytes);
274
275 let bound_bytes = round_usize_up_to_host_pages(MEMORY_MAX);
276 let allocation_bytes = bound_bytes.min(maximum.unwrap_or(usize::MAX));
277 let request_bytes = allocation_bytes + offset_guard_bytes;
278
279 let mmap = crate::mem::with_kernel_aspace(|aspace| {
280 // attempt to use 2MiB alignment but if that's not available fallback to the largest
281 let align = cmp::min(2 * 1048576, aspace.lock().frame_alloc.max_alignment());
282
283 // TODO the align arg should be a named const not a weird number like this
284 Mmap::new_zeroed(aspace.clone(), request_bytes, align, None)
285 .context("Failed to mmap zeroed memory for Memory")
286 })?;
287
288 Ok(vm::Memory::from_parts(
289 mmap,
290 minimum,
291 maximum,
292 memory.page_size_log2,
293 offset_guard_bytes,
294 ))
295 }
296
297 unsafe fn deallocate_memory(&self, _memory_index: DefinedMemoryIndex, _memory: vm::Memory) {}
298
299 unsafe fn allocate_table(
300 &self,
301 table: &translate::Table,
302 _table_index: DefinedTableIndex,
303 ) -> crate::Result<vm::Table> {
304 // TODO we could call out to some resource management instance here to obtain
305 // dynamic "minimum" and "maximum" values that reflect the state of the real systems
306 // memory consumption
307 let maximum = table.limits.max.and_then(|m| usize::try_from(m).ok());
308 let reserve_size = TABLE_MAX.min(maximum.unwrap_or(usize::MAX));
309
310 let elements = if reserve_size == 0 {
311 MmapVec::new_empty()
312 } else {
313 crate::mem::with_kernel_aspace(|aspace| -> crate::Result<_> {
314 let mut elements = MmapVec::new_zeroed(aspace.clone(), reserve_size)?;
315 elements.extend_with(
316 aspace.lock().deref_mut(),
317 usize::try_from(table.limits.min).unwrap(),
318 TableElement::FuncRef(None),
319 );
320 Ok(elements)
321 })?
322 };
323
324 Ok(vm::Table::from_parts(elements, maximum))
325 }
326
327 unsafe fn deallocate_table(&self, _table_index: DefinedTableIndex, _table: vm::Table) {}
328}