1 use super::{
2     InstanceAllocationRequest, InstanceAllocatorImpl, MemoryAllocationIndex, TableAllocationIndex,
3 };
4 use crate::prelude::*;
5 use crate::runtime::vm::instance::RuntimeMemoryCreator;
6 use crate::runtime::vm::memory::{DefaultMemoryCreator, Memory};
7 use crate::runtime::vm::mpk::ProtectionKey;
8 use crate::runtime::vm::table::Table;
9 use crate::runtime::vm::CompiledModuleId;
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     component::{Component, VMComponentOffsets},
24     StaticModuleIndex,
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 InstanceAllocatorImpl for OnDemandInstanceAllocator {
80     #[cfg(feature = "component-model")]
81     fn validate_component_impl<'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_impl(&self, _module: &Module, _offsets: &VMOffsets<HostPtr>) -> Result<()> {
91         Ok(())
92     }
93 
94     #[cfg(feature = "gc")]
95     fn validate_memory_impl(&self, _memory: &wasmtime_environ::Memory) -> Result<()> {
96         Ok(())
97     }
98 
99     fn increment_component_instance_count(&self) -> Result<()> {
100         Ok(())
101     }
102 
103     fn decrement_component_instance_count(&self) {}
104 
105     fn increment_core_instance_count(&self) -> Result<()> {
106         Ok(())
107     }
108 
109     fn decrement_core_instance_count(&self) {}
110 
111     unsafe fn allocate_memory(
112         &self,
113         request: &mut InstanceAllocationRequest,
114         ty: &wasmtime_environ::Memory,
115         tunables: &Tunables,
116         memory_index: Option<DefinedMemoryIndex>,
117     ) -> Result<(MemoryAllocationIndex, Memory)> {
118         let creator = self
119             .mem_creator
120             .as_deref()
121             .unwrap_or_else(|| &DefaultMemoryCreator);
122 
123         let image = if let Some(memory_index) = memory_index {
124             request.runtime_info.memory_image(memory_index)?
125         } else {
126             None
127         };
128 
129         let allocation_index = MemoryAllocationIndex::default();
130         let memory = Memory::new_dynamic(
131             ty,
132             tunables,
133             creator,
134             request
135                 .store
136                 .get()
137                 .expect("if module has memory plans, store is not empty"),
138             image,
139         )?;
140         Ok((allocation_index, memory))
141     }
142 
143     unsafe fn deallocate_memory(
144         &self,
145         _memory_index: Option<DefinedMemoryIndex>,
146         allocation_index: MemoryAllocationIndex,
147         _memory: Memory,
148     ) {
149         debug_assert_eq!(allocation_index, MemoryAllocationIndex::default());
150         // Normal destructors do all the necessary clean up.
151     }
152 
153     unsafe fn allocate_table(
154         &self,
155         request: &mut InstanceAllocationRequest,
156         ty: &wasmtime_environ::Table,
157         tunables: &Tunables,
158         _table_index: DefinedTableIndex,
159     ) -> Result<(TableAllocationIndex, Table)> {
160         let allocation_index = TableAllocationIndex::default();
161         let table = Table::new_dynamic(
162             ty,
163             tunables,
164             request
165                 .store
166                 .get()
167                 .expect("if module has table plans, store is not empty"),
168         )?;
169         Ok((allocation_index, table))
170     }
171 
172     unsafe fn deallocate_table(
173         &self,
174         _table_index: DefinedTableIndex,
175         allocation_index: TableAllocationIndex,
176         _table: Table,
177     ) {
178         debug_assert_eq!(allocation_index, TableAllocationIndex::default());
179         // Normal destructors do all the necessary clean up.
180     }
181 
182     #[cfg(feature = "async")]
183     fn allocate_fiber_stack(&self) -> Result<wasmtime_fiber::FiberStack> {
184         if self.stack_size == 0 {
185             anyhow::bail!("fiber stacks are not supported by the allocator")
186         }
187         let stack = match &self.stack_creator {
188             Some(stack_creator) => {
189                 let stack = stack_creator.new_stack(self.stack_size, self.stack_zeroing)?;
190                 wasmtime_fiber::FiberStack::from_custom(stack)
191             }
192             None => wasmtime_fiber::FiberStack::new(self.stack_size, self.stack_zeroing),
193         }?;
194         Ok(stack)
195     }
196 
197     #[cfg(feature = "async")]
198     unsafe fn deallocate_fiber_stack(&self, stack: wasmtime_fiber::FiberStack) {
199         // The on-demand allocator has no further bookkeeping for fiber stacks
200         // beyond dropping them.
201         let _ = stack;
202     }
203 
204     fn purge_module(&self, _: CompiledModuleId) {}
205 
206     fn next_available_pkey(&self) -> Option<ProtectionKey> {
207         // The on-demand allocator cannot use protection keys--it requires
208         // back-to-back allocation of memory slots that this allocator cannot
209         // guarantee.
210         None
211     }
212 
213     fn restrict_to_pkey(&self, _: ProtectionKey) {
214         // The on-demand allocator cannot use protection keys; an on-demand
215         // allocator will never hand out protection keys to the stores its
216         // engine creates.
217         unreachable!()
218     }
219 
220     fn allow_all_pkeys(&self) {
221         // The on-demand allocator cannot use protection keys; an on-demand
222         // allocator will never hand out protection keys to the stores its
223         // engine creates.
224         unreachable!()
225     }
226 
227     #[cfg(feature = "gc")]
228     fn allocate_gc_heap(
229         &self,
230         engine: &crate::Engine,
231         gc_runtime: &dyn GcRuntime,
232         memory_alloc_index: MemoryAllocationIndex,
233         memory: Memory,
234     ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)> {
235         debug_assert_eq!(memory_alloc_index, MemoryAllocationIndex::default());
236         let mut heap = gc_runtime.new_gc_heap(engine)?;
237         heap.attach(memory);
238         Ok((GcHeapAllocationIndex::default(), heap))
239     }
240 
241     #[cfg(feature = "gc")]
242     fn deallocate_gc_heap(
243         &self,
244         allocation_index: GcHeapAllocationIndex,
245         mut gc_heap: Box<dyn crate::runtime::vm::GcHeap>,
246     ) -> (MemoryAllocationIndex, Memory) {
247         debug_assert_eq!(allocation_index, GcHeapAllocationIndex::default());
248         (MemoryAllocationIndex::default(), gc_heap.detach())
249     }
250 }
251