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