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