1 use super::{
2     InstanceAllocationRequest, InstanceAllocator, MemoryAllocationIndex, TableAllocationIndex,
3 };
4 use crate::prelude::*;
5 use crate::runtime::vm::CompiledModuleId;
6 use crate::runtime::vm::instance::RuntimeMemoryCreator;
7 use crate::runtime::vm::memory::{DefaultMemoryCreator, Memory};
8 use crate::runtime::vm::mpk::ProtectionKey;
9 use crate::runtime::vm::table::Table;
10 use alloc::sync::Arc;
11 use wasmtime_environ::{
12     DefinedMemoryIndex, DefinedTableIndex, HostPtr, Module, Tunables, VMOffsets,
13 };
14 
15 #[cfg(feature = "gc")]
16 use crate::runtime::vm::{GcHeap, GcHeapAllocationIndex, GcRuntime};
17 
18 #[cfg(feature = "async")]
19 use wasmtime_fiber::RuntimeFiberStackCreator;
20 
21 #[cfg(feature = "component-model")]
22 use wasmtime_environ::{
23     StaticModuleIndex,
24     component::{Component, VMComponentOffsets},
25 };
26 
27 /// Represents the on-demand instance allocator.
28 #[derive(Clone)]
29 pub struct OnDemandInstanceAllocator {
30     mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>,
31     #[cfg(feature = "async")]
32     stack_creator: Option<Arc<dyn RuntimeFiberStackCreator>>,
33     #[cfg(feature = "async")]
34     stack_size: usize,
35     #[cfg(feature = "async")]
36     stack_zeroing: bool,
37 }
38 
39 impl OnDemandInstanceAllocator {
40     /// Creates a new on-demand instance allocator.
41     pub fn new(
42         mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>,
43         stack_size: usize,
44         stack_zeroing: bool,
45     ) -> Self {
46         let _ = (stack_size, stack_zeroing); // suppress warnings when async feature is disabled.
47         Self {
48             mem_creator,
49             #[cfg(feature = "async")]
50             stack_creator: None,
51             #[cfg(feature = "async")]
52             stack_size,
53             #[cfg(feature = "async")]
54             stack_zeroing,
55         }
56     }
57 
58     /// Set the stack creator.
59     #[cfg(feature = "async")]
60     pub fn set_stack_creator(&mut self, stack_creator: Arc<dyn RuntimeFiberStackCreator>) {
61         self.stack_creator = Some(stack_creator);
62     }
63 }
64 
65 impl Default for OnDemandInstanceAllocator {
66     fn default() -> Self {
67         Self {
68             mem_creator: None,
69             #[cfg(feature = "async")]
70             stack_creator: None,
71             #[cfg(feature = "async")]
72             stack_size: 0,
73             #[cfg(feature = "async")]
74             stack_zeroing: false,
75         }
76     }
77 }
78 
79 unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
80     #[cfg(feature = "component-model")]
81     fn validate_component<'a>(
82         &self,
83         _component: &Component,
84         _offsets: &VMComponentOffsets<HostPtr>,
85         _get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module,
86     ) -> Result<()> {
87         Ok(())
88     }
89 
90     fn validate_module(&self, _module: &Module, _offsets: &VMOffsets<HostPtr>) -> Result<()> {
91         Ok(())
92     }
93 
94     #[cfg(feature = "gc")]
95     fn validate_memory(&self, _memory: &wasmtime_environ::Memory) -> Result<()> {
96         Ok(())
97     }
98 
99     #[cfg(feature = "component-model")]
100     fn increment_component_instance_count(&self) -> Result<()> {
101         Ok(())
102     }
103 
104     #[cfg(feature = "component-model")]
105     fn decrement_component_instance_count(&self) {}
106 
107     fn increment_core_instance_count(&self) -> Result<()> {
108         Ok(())
109     }
110 
111     fn decrement_core_instance_count(&self) {}
112 
113     fn allocate_memory(
114         &self,
115         request: &mut InstanceAllocationRequest,
116         ty: &wasmtime_environ::Memory,
117         tunables: &Tunables,
118         memory_index: Option<DefinedMemoryIndex>,
119     ) -> Result<(MemoryAllocationIndex, Memory)> {
120         let creator = self
121             .mem_creator
122             .as_deref()
123             .unwrap_or_else(|| &DefaultMemoryCreator);
124 
125         let image = if let Some(memory_index) = memory_index {
126             request.runtime_info.memory_image(memory_index)?
127         } else {
128             None
129         };
130 
131         let allocation_index = MemoryAllocationIndex::default();
132         let memory = Memory::new_dynamic(
133             ty,
134             tunables,
135             creator,
136             unsafe {
137                 request
138                     .store
139                     .get()
140                     .expect("if module has memory plans, store is not empty")
141             },
142             image,
143         )?;
144         Ok((allocation_index, memory))
145     }
146 
147     unsafe fn deallocate_memory(
148         &self,
149         _memory_index: Option<DefinedMemoryIndex>,
150         allocation_index: MemoryAllocationIndex,
151         _memory: Memory,
152     ) {
153         debug_assert_eq!(allocation_index, MemoryAllocationIndex::default());
154         // Normal destructors do all the necessary clean up.
155     }
156 
157     fn allocate_table(
158         &self,
159         request: &mut InstanceAllocationRequest,
160         ty: &wasmtime_environ::Table,
161         tunables: &Tunables,
162         _table_index: DefinedTableIndex,
163     ) -> Result<(TableAllocationIndex, Table)> {
164         let allocation_index = TableAllocationIndex::default();
165         let table = Table::new_dynamic(ty, tunables, unsafe {
166             request
167                 .store
168                 .get()
169                 .expect("if module has table plans, store is not empty")
170         })?;
171         Ok((allocation_index, table))
172     }
173 
174     unsafe fn deallocate_table(
175         &self,
176         _table_index: DefinedTableIndex,
177         allocation_index: TableAllocationIndex,
178         _table: Table,
179     ) {
180         debug_assert_eq!(allocation_index, TableAllocationIndex::default());
181         // Normal destructors do all the necessary clean up.
182     }
183 
184     #[cfg(feature = "async")]
185     fn allocate_fiber_stack(&self) -> Result<wasmtime_fiber::FiberStack> {
186         if self.stack_size == 0 {
187             anyhow::bail!("fiber stacks are not supported by the allocator")
188         }
189         let stack = match &self.stack_creator {
190             Some(stack_creator) => {
191                 let stack = stack_creator.new_stack(self.stack_size, self.stack_zeroing)?;
192                 wasmtime_fiber::FiberStack::from_custom(stack)
193             }
194             None => wasmtime_fiber::FiberStack::new(self.stack_size, self.stack_zeroing),
195         }?;
196         Ok(stack)
197     }
198 
199     #[cfg(feature = "async")]
200     unsafe fn deallocate_fiber_stack(&self, stack: wasmtime_fiber::FiberStack) {
201         // The on-demand allocator has no further bookkeeping for fiber stacks
202         // beyond dropping them.
203         let _ = stack;
204     }
205 
206     fn purge_module(&self, _: CompiledModuleId) {}
207 
208     fn next_available_pkey(&self) -> Option<ProtectionKey> {
209         // The on-demand allocator cannot use protection keys--it requires
210         // back-to-back allocation of memory slots that this allocator cannot
211         // guarantee.
212         None
213     }
214 
215     fn restrict_to_pkey(&self, _: ProtectionKey) {
216         // The on-demand allocator cannot use protection keys; an on-demand
217         // allocator will never hand out protection keys to the stores its
218         // engine creates.
219         unreachable!()
220     }
221 
222     fn allow_all_pkeys(&self) {
223         // The on-demand allocator cannot use protection keys; an on-demand
224         // allocator will never hand out protection keys to the stores its
225         // engine creates.
226         unreachable!()
227     }
228 
229     #[cfg(feature = "gc")]
230     fn allocate_gc_heap(
231         &self,
232         engine: &crate::Engine,
233         gc_runtime: &dyn GcRuntime,
234         memory_alloc_index: MemoryAllocationIndex,
235         memory: Memory,
236     ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)> {
237         debug_assert_eq!(memory_alloc_index, MemoryAllocationIndex::default());
238         let mut heap = gc_runtime.new_gc_heap(engine)?;
239         heap.attach(memory);
240         Ok((GcHeapAllocationIndex::default(), heap))
241     }
242 
243     #[cfg(feature = "gc")]
244     fn deallocate_gc_heap(
245         &self,
246         allocation_index: GcHeapAllocationIndex,
247         mut gc_heap: Box<dyn crate::runtime::vm::GcHeap>,
248     ) -> (MemoryAllocationIndex, Memory) {
249         debug_assert_eq!(allocation_index, GcHeapAllocationIndex::default());
250         (MemoryAllocationIndex::default(), gc_heap.detach())
251     }
252 }
253