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 core::future::Future;
12 use core::pin::Pin;
13 use wasmtime_environ::{DefinedMemoryIndex, DefinedTableIndex, HostPtr, Module, VMOffsets};
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<'a, 'b: 'a, 'c: 'a>(
114         &'a self,
115         request: &'a mut InstanceAllocationRequest<'b, 'c>,
116         ty: &'a wasmtime_environ::Memory,
117         memory_index: Option<DefinedMemoryIndex>,
118     ) -> Result<
119         Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>>,
120         OutOfMemory,
121     > {
122         let creator = self
123             .mem_creator
124             .as_deref()
125             .unwrap_or_else(|| &DefaultMemoryCreator);
126 
127         Ok(Box::into_pin(try_new::<Box<_>>(async move {
128             let image = if let Some(memory_index) = memory_index {
129                 request.runtime_info.memory_image(memory_index)?
130             } else {
131                 None
132             };
133 
134             let allocation_index = MemoryAllocationIndex::default();
135             let memory = Memory::new_dynamic(
136                 ty,
137                 request.store.engine(),
138                 creator,
139                 image,
140                 request.limiter.as_deref_mut(),
141             )
142             .await?;
143             Ok((allocation_index, memory))
144         })?))
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<'a, 'b: 'a, 'c: 'a>(
158         &'a self,
159         request: &'a mut InstanceAllocationRequest<'b, 'c>,
160         ty: &'a wasmtime_environ::Table,
161         _table_index: DefinedTableIndex,
162     ) -> Result<
163         Pin<Box<dyn Future<Output = Result<(TableAllocationIndex, Table)>> + Send + 'a>>,
164         OutOfMemory,
165     > {
166         Ok(Box::into_pin(try_new::<Box<_>>(async move {
167             let allocation_index = TableAllocationIndex::default();
168             let table = Table::new_dynamic(
169                 ty,
170                 request.store.engine().tunables(),
171                 request.limiter.as_deref_mut(),
172             )
173             .await?;
174             Ok((allocation_index, table))
175         })?))
176     }
177 
178     unsafe fn deallocate_table(
179         &self,
180         _table_index: DefinedTableIndex,
181         allocation_index: TableAllocationIndex,
182         _table: Table,
183     ) {
184         debug_assert_eq!(allocation_index, TableAllocationIndex::default());
185         // Normal destructors do all the necessary clean up.
186     }
187 
188     #[cfg(feature = "async")]
189     fn allocate_fiber_stack(&self) -> Result<wasmtime_fiber::FiberStack> {
190         if self.stack_size == 0 {
191             crate::bail!("fiber stacks are not supported by the allocator")
192         }
193         let stack = match &self.stack_creator {
194             Some(stack_creator) => {
195                 let stack = stack_creator.new_stack(self.stack_size, self.stack_zeroing)?;
196                 wasmtime_fiber::FiberStack::from_custom(stack)
197             }
198             None => wasmtime_fiber::FiberStack::new(self.stack_size, self.stack_zeroing),
199         }?;
200         Ok(stack)
201     }
202 
203     #[cfg(feature = "async")]
204     unsafe fn deallocate_fiber_stack(&self, stack: wasmtime_fiber::FiberStack) {
205         // The on-demand allocator has no further bookkeeping for fiber stacks
206         // beyond dropping them.
207         let _ = stack;
208     }
209 
210     fn purge_module(&self, _: CompiledModuleId) {}
211 
212     fn next_available_pkey(&self) -> Option<ProtectionKey> {
213         // The on-demand allocator cannot use protection keys--it requires
214         // back-to-back allocation of memory slots that this allocator cannot
215         // guarantee.
216         None
217     }
218 
219     fn restrict_to_pkey(&self, _: ProtectionKey) {
220         // The on-demand allocator cannot use protection keys; an on-demand
221         // allocator will never hand out protection keys to the stores its
222         // engine creates.
223         unreachable!()
224     }
225 
226     fn allow_all_pkeys(&self) {
227         // The on-demand allocator cannot use protection keys; an on-demand
228         // allocator will never hand out protection keys to the stores its
229         // engine creates.
230         unreachable!()
231     }
232 
233     #[cfg(feature = "gc")]
234     fn allocate_gc_heap(
235         &self,
236         engine: &crate::Engine,
237         gc_runtime: &dyn GcRuntime,
238         memory_alloc_index: MemoryAllocationIndex,
239         memory: Memory,
240     ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)> {
241         debug_assert_eq!(memory_alloc_index, MemoryAllocationIndex::default());
242         let mut heap = gc_runtime.new_gc_heap(engine)?;
243         heap.attach(memory);
244         Ok((GcHeapAllocationIndex::default(), heap))
245     }
246 
247     #[cfg(feature = "gc")]
248     fn deallocate_gc_heap(
249         &self,
250         allocation_index: GcHeapAllocationIndex,
251         mut gc_heap: Box<dyn crate::runtime::vm::GcHeap>,
252     ) -> (MemoryAllocationIndex, Memory) {
253         debug_assert_eq!(allocation_index, GcHeapAllocationIndex::default());
254         (MemoryAllocationIndex::default(), gc_heap.detach())
255     }
256 }
257