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