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