1 //! An allocator definition for this embedding.
2 //!
3 //! The Rust standard library and Wasmtime require a memory allocator to be
4 //! configured. For custom embeddings of Wasmtime this might likely already be
5 //! defined elsewhere in the system in which case that should be used. This file
6 //! contains an example implementation using the Rust `dlmalloc` crate using
7 //! memory created by `wasmtime_*` platform symbols. This provides a file that
8 //! manages memory without any extra runtime dependencies, but this is just an
9 //! example.
10 //!
11 //! Allocators in Rust are configured with the `#[global_allocator]` attribute
12 //! and the `GlobalAlloc for T` trait impl. This should be used when hooking
13 //! up to an allocator elsewhere in the system.
14 
15 use alloc::alloc::{GlobalAlloc, Layout};
16 use core::cell::UnsafeCell;
17 use core::ops::{Deref, DerefMut};
18 use core::ptr;
19 use core::sync::atomic::{
20     AtomicBool,
21     Ordering::{Acquire, Release},
22 };
23 use dlmalloc::Dlmalloc;
24 
25 #[global_allocator]
26 static MALLOC: MyGlobalDmalloc = MyGlobalDmalloc {
27     dlmalloc: Mutex::new(Dlmalloc::new_with_allocator(MyAllocator)),
28 };
29 
30 struct MyGlobalDmalloc {
31     dlmalloc: Mutex<Dlmalloc<MyAllocator>>,
32 }
33 
34 struct MyAllocator;
35 
36 unsafe impl GlobalAlloc for MyGlobalDmalloc {
37     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
38         self.dlmalloc
39             .try_lock()
40             .unwrap()
41             .malloc(layout.size(), layout.align())
42     }
43 
44     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
45         self.dlmalloc
46             .try_lock()
47             .unwrap()
48             .calloc(layout.size(), layout.align())
49     }
50 
51     unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
52         self.dlmalloc
53             .try_lock()
54             .unwrap()
55             .realloc(ptr, layout.size(), layout.align(), new_size)
56     }
57 
58     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
59         self.dlmalloc
60             .try_lock()
61             .unwrap()
62             .free(ptr, layout.size(), layout.align())
63     }
64 }
65 
66 #[cfg(not(feature = "wasi"))]
67 const INITIAL_HEAP_SIZE: usize = 64 * 1024;
68 // The wasi component requires a larger heap than the module tests
69 #[cfg(feature = "wasi")]
70 const INITIAL_HEAP_SIZE: usize = 4 * 1024 * 1024;
71 
72 static mut INITIAL_HEAP: [u8; INITIAL_HEAP_SIZE] = [0; INITIAL_HEAP_SIZE];
73 static mut INITIAL_HEAP_ALLOCATED: bool = false;
74 
75 unsafe impl dlmalloc::Allocator for MyAllocator {
76     fn alloc(&self, _size: usize) -> (*mut u8, usize, u32) {
77         unsafe {
78             if INITIAL_HEAP_ALLOCATED {
79                 (ptr::null_mut(), 0, 0)
80             } else {
81                 INITIAL_HEAP_ALLOCATED = true;
82                 ((&raw mut INITIAL_HEAP).cast(), INITIAL_HEAP_SIZE, 0)
83             }
84         }
85     }
86 
87     fn remap(&self, _ptr: *mut u8, _old: usize, _new: usize, _can_move: bool) -> *mut u8 {
88         core::ptr::null_mut()
89     }
90 
91     fn free_part(&self, _ptr: *mut u8, _old: usize, _new: usize) -> bool {
92         false
93     }
94 
95     fn free(&self, _ptr: *mut u8, _size: usize) -> bool {
96         false
97     }
98 
99     fn can_release_part(&self, _flags: u32) -> bool {
100         false
101     }
102 
103     fn allocates_zeros(&self) -> bool {
104         true
105     }
106 
107     fn page_size(&self) -> usize {
108         4096
109     }
110 }
111 
112 // Simple mutex which only supports `try_lock` at this time. This would probably
113 // be replaced with a "real" mutex in a "real" embedding.
114 struct Mutex<T> {
115     data: UnsafeCell<T>,
116     locked: AtomicBool,
117 }
118 
119 unsafe impl<T: Send> Send for Mutex<T> {}
120 unsafe impl<T: Send> Sync for Mutex<T> {}
121 
122 impl<T> Mutex<T> {
123     const fn new(val: T) -> Mutex<T> {
124         Mutex {
125             data: UnsafeCell::new(val),
126             locked: AtomicBool::new(false),
127         }
128     }
129 
130     fn try_lock(&self) -> Option<impl DerefMut<Target = T> + '_> {
131         if self.locked.swap(true, Acquire) {
132             None
133         } else {
134             Some(MutexGuard { lock: self })
135         }
136     }
137 }
138 
139 struct MutexGuard<'a, T> {
140     lock: &'a Mutex<T>,
141 }
142 
143 impl<T> Deref for MutexGuard<'_, T> {
144     type Target = T;
145 
146     fn deref(&self) -> &T {
147         unsafe { &*self.lock.data.get() }
148     }
149 }
150 
151 impl<T> DerefMut for MutexGuard<'_, T> {
152     fn deref_mut(&mut self) -> &mut T {
153         unsafe { &mut *self.lock.data.get() }
154     }
155 }
156 
157 impl<T> Drop for MutexGuard<'_, T> {
158     fn drop(&mut self) {
159         self.lock.locked.store(false, Release);
160     }
161 }
162