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