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