Skip to main content

wasmtime/runtime/vm/instance/allocator/
pooling.rs

1//! Implements the pooling instance allocator.
2//!
3//! The pooling instance allocator maps memory in advance and allocates
4//! instances, memories, tables, and stacks from a pool of available resources.
5//! Using the pooling instance allocator can speed up module instantiation when
6//! modules can be constrained based on configurable limits
7//! ([`InstanceLimits`]). Each new instance is stored in a "slot"; as instances
8//! are allocated and freed, these slots are either filled or emptied:
9//!
10//! ```text
11//! ┌──────┬──────┬──────┬──────┬──────┐
12//! │Slot 0│Slot 1│Slot 2│Slot 3│......│
13//! └──────┴──────┴──────┴──────┴──────┘
14//! ```
15//!
16//! Each slot has a "slot ID"--an index into the pool. Slot IDs are handed out
17//! by the [`index_allocator`] module. Note that each kind of pool-allocated
18//! item is stored in its own separate pool: [`memory_pool`], [`table_pool`],
19//! [`stack_pool`]. See those modules for more details.
20
21mod decommit_queue;
22mod index_allocator;
23mod memory_pool;
24mod metrics;
25mod table_pool;
26
27#[cfg(feature = "gc")]
28mod gc_heap_pool;
29
30#[cfg(all(feature = "async"))]
31mod generic_stack_pool;
32#[cfg(all(feature = "async", unix, not(miri)))]
33mod unix_stack_pool;
34
35#[cfg(all(feature = "async"))]
36cfg_select! {
37    all(unix, not(miri), not(asan)) => {
38        use unix_stack_pool as stack_pool;
39    }
40    _ => {
41        use generic_stack_pool as stack_pool;
42    }
43}
44
45use self::decommit_queue::DecommitQueue;
46use self::memory_pool::MemoryPool;
47pub use self::metrics::PoolingAllocatorMetrics;
48use self::table_pool::TablePool;
49use super::{
50    InstanceAllocationRequest, InstanceAllocator, MemoryAllocationIndex, TableAllocationIndex,
51};
52use crate::Enabled;
53use crate::config::PoolingAllocationConfig;
54use crate::prelude::*;
55use crate::runtime::vm::{
56    CompiledModuleId, Memory, Table,
57    instance::Instance,
58    mpk::{self, ProtectionKey, ProtectionMask},
59    sys::vm::PageMap,
60};
61use core::future::Future;
62use core::pin::Pin;
63use core::sync::atomic::AtomicUsize;
64use std::borrow::Cow;
65use std::fmt::Display;
66use std::sync::{Mutex, MutexGuard};
67use std::{
68    mem,
69    sync::atomic::{AtomicU64, Ordering},
70};
71use wasmtime_environ::{
72    DefinedMemoryIndex, DefinedTableIndex, HostPtr, MemoryKind, Module, Tunables, VMOffsets,
73};
74
75#[cfg(feature = "gc")]
76use super::GcHeapAllocationIndex;
77#[cfg(feature = "gc")]
78use crate::runtime::vm::{GcHeap, GcRuntime};
79#[cfg(feature = "gc")]
80use gc_heap_pool::GcHeapPool;
81
82/// Pad a value out to a full cache line (or two, on aarch64 prefetch
83/// granularity) so neighboring shards don't false-share.
84#[repr(align(128))]
85#[derive(Debug)]
86struct CachePadded<T>(T);
87
88/// Identifier of one shard of the pooling allocator's sharded data
89/// structures (the decommit queues and each pool's index allocator).
90#[derive(Copy, Clone, Debug, PartialEq, Eq)]
91pub(crate) struct ShardId(u32);
92
93impl ShardId {
94    pub(crate) fn from_index(index: usize) -> ShardId {
95        ShardId(u32::try_from(index).unwrap())
96    }
97
98    pub(crate) fn index(self) -> usize {
99        usize::try_from(self.0).unwrap()
100    }
101}
102
103/// The number of shards used for the pooling allocator's sharded data
104/// structures: one per available CPU, capped to 16.
105///
106/// The cap bounds worst-case probing when pools run near-full, the
107/// dilution of per-shard warm-slot budgets, and per-shard memory
108/// overhead, while still being enough shards to make lock collisions
109/// rare given the very short critical sections involved.
110pub(crate) fn default_shard_count() -> u32 {
111    let n = std::thread::available_parallelism()
112        .map(|n| n.get())
113        .unwrap_or(1)
114        .min(16);
115    u32::try_from(n).unwrap()
116}
117
118/// Pick this thread's shard (used for both the sharded decommit queue and
119/// the sharded index allocators): assigned round-robin at first use per
120/// thread, cached in a thread-local.
121pub(crate) fn thread_shard(nshards: usize) -> ShardId {
122    static NEXT_SHARD: AtomicUsize = AtomicUsize::new(0);
123    std::thread_local! {
124        static SHARD: usize = NEXT_SHARD.fetch_add(1, Ordering::Relaxed);
125    }
126    ShardId::from_index(SHARD.with(|s| *s) % nshards)
127}
128
129/// Enumerate all shard ids for a sharded structure with `nshards` shards,
130/// starting with the current thread's home shard and wrapping around.
131pub(crate) fn shard_ids_from_home(nshards: usize) -> impl Iterator<Item = ShardId> {
132    let home = thread_shard(nshards).index();
133    (0..nshards).map(move |i| ShardId::from_index((home + i) % nshards))
134}
135
136#[cfg(feature = "async")]
137use stack_pool::StackPool;
138
139#[cfg(feature = "component-model")]
140use wasmtime_environ::{
141    StaticModuleIndex,
142    component::{Component, VMComponentOffsets},
143};
144
145impl PoolingAllocationConfig {
146    /// Tests whether [`Self::pagemap_scan`] is available or not on the host
147    /// system.
148    pub fn is_pagemap_scan_available() -> bool {
149        PageMap::new().is_some()
150    }
151}
152
153/// An error returned when the pooling allocator cannot allocate a table,
154/// memory, etc... because the maximum number of concurrent allocations for that
155/// entity has been reached.
156#[derive(Debug)]
157pub struct PoolConcurrencyLimitError {
158    limit: usize,
159    kind: Cow<'static, str>,
160}
161
162impl core::error::Error for PoolConcurrencyLimitError {}
163
164impl Display for PoolConcurrencyLimitError {
165    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
166        let limit = self.limit;
167        let kind = &self.kind;
168        write!(f, "maximum concurrent limit of {limit} for {kind} reached")
169    }
170}
171
172impl PoolConcurrencyLimitError {
173    fn new(limit: usize, kind: impl Into<Cow<'static, str>>) -> Self {
174        Self {
175            limit,
176            kind: kind.into(),
177        }
178    }
179}
180
181/// Implements the pooling instance allocator.
182///
183/// This allocator internally maintains pools of instances, memories, tables,
184/// and stacks.
185///
186/// Note: the resource pools are manually dropped so that the fault handler
187/// terminates correctly.
188#[derive(Debug)]
189pub struct PoolingInstanceAllocator {
190    // The number of live core module and component instances at any given
191    // time. Note that this can temporarily go over the configured limit. This
192    // doesn't mean we have actually overshot, but that we attempted to allocate
193    // a new instance and incremented the counter, we've seen (or are about to
194    // see) that the counter is beyond the configured threshold, and are going
195    // to decrement the counter and return an error but haven't done so yet. See
196    // the increment trait methods for more details.
197    live_core_instances: AtomicU64,
198    live_component_instances: AtomicU64,
199
200    /// Sharded to avoid a single global mutex on every deallocation when
201    /// decommit batching is enabled: each thread appends to its own shard
202    /// (assigned round-robin at first use) and flushes that shard when it
203    /// reaches the configured batch size. Slot-exhaustion paths flush all
204    /// shards.
205    decommit_queues: Box<[CachePadded<Mutex<DecommitQueue>>]>,
206
207    memories: MemoryPool,
208    live_memories: AtomicUsize,
209
210    tables: TablePool,
211    live_tables: AtomicUsize,
212
213    #[cfg(feature = "gc")]
214    gc_heaps: Option<GcHeapPool>,
215    #[cfg(feature = "gc")]
216    live_gc_heaps: AtomicUsize,
217
218    #[cfg(feature = "async")]
219    stacks: StackPool,
220    #[cfg(feature = "async")]
221    live_stacks: AtomicUsize,
222
223    pagemap: Option<PageMap>,
224    config: PoolingAllocationConfig,
225}
226
227impl Drop for PoolingInstanceAllocator {
228    fn drop(&mut self) {
229        if !cfg!(debug_assertions) {
230            return;
231        }
232
233        // NB: when cfg(not(debug_assertions)) it is okay that we don't flush
234        // the queue, as the sub-pools will unmap those ranges anyways, so
235        // there's no point in decommitting them. But we do need to flush the
236        // queue when debug assertions are enabled to make sure that all
237        // entities get returned to their associated sub-pools and we can
238        // differentiate between a leaking slot and an enqueued-for-decommit
239        // slot.
240        self.flush_all_decommit_queues();
241
242        debug_assert_eq!(self.live_component_instances.load(Ordering::Acquire), 0);
243        debug_assert_eq!(self.live_core_instances.load(Ordering::Acquire), 0);
244        debug_assert_eq!(self.live_memories.load(Ordering::Acquire), 0);
245        debug_assert_eq!(self.live_tables.load(Ordering::Acquire), 0);
246
247        debug_assert!(self.memories.is_empty());
248        debug_assert!(self.tables.is_empty());
249
250        #[cfg(feature = "gc")]
251        if let Some(gc_heaps) = &self.gc_heaps {
252            debug_assert!(gc_heaps.is_empty());
253            debug_assert_eq!(self.live_gc_heaps.load(Ordering::Acquire), 0);
254        }
255
256        #[cfg(feature = "async")]
257        {
258            debug_assert!(self.stacks.is_empty());
259            debug_assert_eq!(self.live_stacks.load(Ordering::Acquire), 0);
260        }
261    }
262}
263
264impl PoolingInstanceAllocator {
265    /// Releases memory kept resident for unused slots, returning the number
266    /// of bytes released. See
267    /// [`MemoryPool::release_resident_unused_memory`].
268    pub fn release_resident_unused_memory(&self) -> usize {
269        self.memories.release_resident_unused_memory()
270    }
271
272    /// Creates a new pooling instance allocator with the given strategy and limits.
273    pub fn new(config: &PoolingAllocationConfig, tunables: &Tunables) -> Result<Self> {
274        Ok(Self {
275            live_component_instances: AtomicU64::new(0),
276            live_core_instances: AtomicU64::new(0),
277            decommit_queues: (0..default_shard_count())
278                .map(|_| CachePadded(Mutex::new(DecommitQueue::default())))
279                .try_collect::<Box<[_]>, OutOfMemory>()?,
280            memories: MemoryPool::new(config, tunables)?,
281            live_memories: AtomicUsize::new(0),
282            tables: TablePool::new(config)?,
283            live_tables: AtomicUsize::new(0),
284            #[cfg(feature = "gc")]
285            gc_heaps: if tunables.collector.is_some() {
286                Some(GcHeapPool::new(config, tunables)?)
287            } else {
288                None
289            },
290            #[cfg(feature = "gc")]
291            live_gc_heaps: AtomicUsize::new(0),
292            #[cfg(feature = "async")]
293            stacks: StackPool::new(config)?,
294            #[cfg(feature = "async")]
295            live_stacks: AtomicUsize::new(0),
296            pagemap: match config.pagemap_scan {
297                Enabled::Auto => PageMap::new(),
298                Enabled::Yes => Some(PageMap::new().ok_or_else(|| {
299                    format_err!(
300                        "required to enable PAGEMAP_SCAN but this system \
301                         does not support it"
302                    )
303                })?),
304                Enabled::No => None,
305            },
306            config: config.clone(),
307        })
308    }
309
310    fn max_core_instance_size(&self) -> Option<usize> {
311        self.config
312            .limits
313            .core_instance_size
314            .checked_next_multiple_of(mem::align_of::<Instance>())
315    }
316
317    fn validate_table_plans(&self, module: &Module) -> Result<()> {
318        self.tables.validate(module)
319    }
320
321    fn validate_memory_plans(&self, module: &Module) -> Result<()> {
322        self.memories.validate_memories(module)
323    }
324
325    fn validate_core_instance_size(&self, offsets: &VMOffsets<HostPtr>) -> Result<()> {
326        let layout = Instance::alloc_layout(offsets);
327        let Some(max) = self.max_core_instance_size() else {
328            return Ok(());
329        };
330
331        if layout.size() <= max {
332            return Ok(());
333        }
334
335        // If this `module` exceeds the allocation size allotted to it then an
336        // error will be reported here. The error of "required N bytes but
337        // cannot allocate that" is pretty opaque, however, because it's not
338        // clear what the breakdown of the N bytes are and what to optimize
339        // next. To help provide a better error message here some fancy-ish
340        // logic is done here to report the breakdown of the byte request into
341        // the largest portions and where it's coming from.
342        let mut message = format!(
343            "instance allocation for this module \
344             requires {} bytes which exceeds the configured maximum \
345             of {max} bytes; breakdown of allocation requirement:\n\n",
346            layout.size(),
347        );
348
349        let mut remaining = layout.size();
350        let mut push = |name: &str, bytes: usize| {
351            assert!(remaining >= bytes);
352            remaining -= bytes;
353
354            // If the `name` region is more than 5% of the allocation request
355            // then report it here, otherwise ignore it. We have less than 20
356            // fields so we're guaranteed that something should be reported, and
357            // otherwise it's not particularly interesting to learn about 5
358            // different fields that are all 8 or 0 bytes. Only try to report
359            // the "major" sources of bytes here.
360            if bytes > layout.size() / 20 {
361                message.push_str(&format!(
362                    " * {:.02}% - {} bytes - {}\n",
363                    ((bytes as f32) / (layout.size() as f32)) * 100.0,
364                    bytes,
365                    name,
366                ));
367            }
368        };
369
370        // The `Instance` itself requires some size allocated to it.
371        push("instance state management", mem::size_of::<Instance>());
372
373        // Afterwards the `VMContext`'s regions are why we're requesting bytes,
374        // so ask it for descriptions on each region's byte size.
375        for (desc, size) in offsets.region_sizes() {
376            push(desc, size as usize);
377        }
378
379        // double-check we accounted for all the bytes
380        assert_eq!(remaining, 0);
381
382        bail!("{message}")
383    }
384
385    #[cfg(feature = "component-model")]
386    fn validate_component_instance_size(
387        &self,
388        offsets: &VMComponentOffsets<HostPtr>,
389        core_instances_aggregate_size: usize,
390    ) -> Result<()> {
391        let vmcomponentctx_size = usize::try_from(offsets.size_of_vmctx()).unwrap();
392        let total_instance_size = core_instances_aggregate_size.saturating_add(vmcomponentctx_size);
393        if total_instance_size <= self.config.limits.component_instance_size {
394            return Ok(());
395        }
396
397        // TODO: Add context with detailed accounting of what makes up all the
398        // `VMComponentContext`'s space like we do for module instances.
399        bail!(
400            "instance allocation for this component requires {total_instance_size} bytes of `VMComponentContext` \
401             and aggregated core instance runtime space which exceeds the configured maximum of {} bytes. \
402             `VMComponentContext` used {vmcomponentctx_size} bytes, `core module instances` used \
403             {core_instances_aggregate_size} bytes.",
404            self.config.limits.component_instance_size
405        )
406    }
407
408    /// Returns the decommit-queue shard for `shard`.
409    fn decommit_queue(&self, shard: ShardId) -> &Mutex<DecommitQueue> {
410        &self.decommit_queues[shard.index()].0
411    }
412
413    /// Enumerate all decommit-queue shard ids, starting with the current
414    /// thread's home shard.
415    fn decommit_shard_ids(&self) -> impl Iterator<Item = ShardId> {
416        shard_ids_from_home(self.decommit_queues.len())
417    }
418
419    fn flush_decommit_queue(&self, mut locked_queue: MutexGuard<'_, DecommitQueue>) -> bool {
420        // Take the queue out of the mutex and drop the lock, to minimize
421        // contention.
422        let queue = mem::take(&mut *locked_queue);
423        drop(locked_queue);
424        queue.flush(self)
425    }
426
427    /// Flush every shard of the decommit queue, e.g. on allocator drop.
428    /// Returns whether any slot was returned to any pool.
429    fn flush_all_decommit_queues(&self) -> bool {
430        let mut any = false;
431        for shard in self.decommit_shard_ids() {
432            let queue = self.decommit_queue(shard).lock().unwrap();
433            any |= self.flush_decommit_queue(queue);
434        }
435        any
436    }
437
438    /// Execute `f` and if it returns `Err(PoolConcurrencyLimitError)`, then try
439    /// flushing the decommit queue. If flushing the queue freed up slots, then
440    /// try running `f` again.
441    ///
442    /// Queue shards are flushed one at a time, retrying `f` after each flush
443    /// that returned slots to a pool, rather than eagerly flushing all
444    /// shards: one flushed shard is often enough to satisfy the allocation,
445    /// and this avoids acquiring every shard's lock (at the cost of raising
446    /// the chances that another thread steals the freshly-flushed slots
447    /// before we get a chance to grab one, in which case we keep flushing).
448    ///
449    /// Note that [`Self::flush_decommit_queue`] takes the shard's queue out
450    /// of its mutex and drops the lock immediately, so no queue lock is held
451    /// while decommitting or while `f` runs.
452    #[cfg(feature = "async")]
453    fn with_flush_and_retry<T>(&self, mut f: impl FnMut() -> Result<T>) -> Result<T> {
454        let mut result = f();
455        for shard in self.decommit_shard_ids() {
456            match &result {
457                Err(e) if e.is::<PoolConcurrencyLimitError>() => {}
458                _ => break,
459            }
460            let queue = self.decommit_queue(shard).lock().unwrap();
461            if self.flush_decommit_queue(queue) {
462                result = f();
463            }
464        }
465        result
466    }
467
468    fn merge_or_flush(&self, mut local_queue: DecommitQueue) {
469        match local_queue.raw_len() {
470            // If we didn't enqueue any regions for decommit, then we must have
471            // either memset the whole entity or eagerly remapped it to zero
472            // because we don't have linux's `madvise(DONTNEED)` semantics. In
473            // either case, the entity slot is ready for reuse immediately.
474            0 => {
475                local_queue.flush(self);
476            }
477
478            // We enqueued at least our batch size of regions for decommit, so
479            // flush the local queue immediately. Don't bother inspecting (or
480            // locking!) the shared queue.
481            n if n >= self.config.decommit_batch_size => {
482                local_queue.flush(self);
483            }
484
485            // If we enqueued some regions for decommit, but did not reach our
486            // batch size, so we don't want to flush it yet, then merge the
487            // local queue into this thread's shard of the shared queue.
488            n => {
489                debug_assert!(n < self.config.decommit_batch_size);
490                let shard = thread_shard(self.decommit_queues.len());
491                let mut shared_queue = self.decommit_queue(shard).lock().unwrap();
492                shared_queue.append(&mut local_queue);
493                // And if this shard now has at least as many regions enqueued
494                // for decommit as our batch size, then we can flush it.
495                if shared_queue.raw_len() >= self.config.decommit_batch_size {
496                    self.flush_decommit_queue(shared_queue);
497                }
498            }
499        }
500    }
501
502    pub fn config(&self) -> &PoolingAllocationConfig {
503        &self.config
504    }
505}
506
507unsafe impl InstanceAllocator for PoolingInstanceAllocator {
508    #[cfg(feature = "component-model")]
509    fn validate_component<'a>(
510        &self,
511        component: &Component,
512        offsets: &VMComponentOffsets<HostPtr>,
513        get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module,
514    ) -> Result<()> {
515        let mut num_core_instances = 0;
516        let mut num_memories = 0;
517        let mut num_tables = 0;
518        let mut core_instances_aggregate_size: usize = 0;
519        for init in &component.initializers {
520            use wasmtime_environ::component::GlobalInitializer::*;
521            use wasmtime_environ::component::InstantiateModule;
522            match init {
523                InstantiateModule(InstantiateModule::Import(_, _), _) => {
524                    num_core_instances += 1;
525                    // Can't statically account for the total vmctx size, number
526                    // of memories, and number of tables in this component.
527                }
528                InstantiateModule(InstantiateModule::Static(static_module_index, _), _) => {
529                    let module = get_module(*static_module_index);
530                    let offsets = VMOffsets::new(HostPtr, &module);
531                    let layout = Instance::alloc_layout(&offsets);
532                    self.validate_module(module, &offsets)?;
533                    num_core_instances += 1;
534                    num_memories += module.num_defined_memories();
535                    num_tables += module.num_defined_tables();
536                    core_instances_aggregate_size += layout.size();
537                }
538                LowerImport { .. }
539                | ExtractMemory(_)
540                | ExtractTable(_)
541                | ExtractRealloc(_)
542                | ExtractCallback(_)
543                | ExtractPostReturn(_)
544                | Resource(_) => {}
545            }
546        }
547
548        if num_core_instances
549            > usize::try_from(self.config.limits.max_core_instances_per_component).unwrap()
550        {
551            bail!(
552                "The component transitively contains {num_core_instances} core module instances, \
553                 which exceeds the configured maximum of {} in the pooling allocator",
554                self.config.limits.max_core_instances_per_component
555            );
556        }
557
558        if num_memories > usize::try_from(self.config.limits.max_memories_per_component).unwrap() {
559            bail!(
560                "The component transitively contains {num_memories} Wasm linear memories, which \
561                 exceeds the configured maximum of {} in the pooling allocator",
562                self.config.limits.max_memories_per_component
563            );
564        }
565
566        if num_tables > usize::try_from(self.config.limits.max_tables_per_component).unwrap() {
567            bail!(
568                "The component transitively contains {num_tables} tables, which exceeds the \
569                 configured maximum of {} in the pooling allocator",
570                self.config.limits.max_tables_per_component
571            );
572        }
573
574        self.validate_component_instance_size(offsets, core_instances_aggregate_size)
575            .context("component instance size does not fit in pooling allocator requirements")?;
576
577        Ok(())
578    }
579
580    fn validate_module(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()> {
581        self.validate_memory_plans(module)
582            .context("module memory does not fit in pooling allocator requirements")?;
583        self.validate_table_plans(module)
584            .context("module table does not fit in pooling allocator requirements")?;
585        self.validate_core_instance_size(offsets)
586            .context("module instance size does not fit in pooling allocator requirements")?;
587        Ok(())
588    }
589
590    #[cfg(feature = "gc")]
591    fn validate_memory(&self, memory: &wasmtime_environ::Memory) -> Result<()> {
592        self.memories.validate_memory(memory)
593    }
594
595    #[cfg(feature = "component-model")]
596    fn increment_component_instance_count(&self) -> Result<()> {
597        let old_count = self.live_component_instances.fetch_add(1, Ordering::AcqRel);
598        if old_count >= u64::from(self.config.limits.total_component_instances) {
599            self.decrement_component_instance_count();
600            return Err(PoolConcurrencyLimitError::new(
601                usize::try_from(self.config.limits.total_component_instances).unwrap(),
602                "component instances",
603            )
604            .into());
605        }
606        Ok(())
607    }
608
609    #[cfg(feature = "component-model")]
610    fn decrement_component_instance_count(&self) {
611        self.live_component_instances.fetch_sub(1, Ordering::AcqRel);
612    }
613
614    fn increment_core_instance_count(&self) -> Result<()> {
615        let old_count = self.live_core_instances.fetch_add(1, Ordering::AcqRel);
616        if old_count >= u64::from(self.config.limits.total_core_instances) {
617            self.decrement_core_instance_count();
618            return Err(PoolConcurrencyLimitError::new(
619                usize::try_from(self.config.limits.total_core_instances).unwrap(),
620                "core instances",
621            )
622            .into());
623        }
624        Ok(())
625    }
626
627    fn decrement_core_instance_count(&self) {
628        self.live_core_instances.fetch_sub(1, Ordering::AcqRel);
629    }
630
631    fn allocate_memory<'a, 'b: 'a, 'c: 'a>(
632        &'a self,
633        request: &'a mut InstanceAllocationRequest<'b, 'c>,
634        ty: &'a wasmtime_environ::Memory,
635        memory_index: Option<DefinedMemoryIndex>,
636        _memory_kind: MemoryKind,
637    ) -> Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>> {
638        crate::runtime::box_future(async move {
639            async {
640                // FIXME(rust-lang/rust#145127) this should ideally use a version of
641                // `with_flush_and_retry` but adapted for async closures instead of only
642                // sync closures. Right now that won't compile though so this is the
643                // manually expanded version of the method.
644                let mut e = match self.memories.allocate(request, ty, memory_index).await {
645                    Ok(result) => return Ok(result),
646                    Err(e) => e,
647                };
648
649                for shard in self.decommit_shard_ids() {
650                    if !e.is::<PoolConcurrencyLimitError>() {
651                        break;
652                    }
653                    let queue = self.decommit_queue(shard).lock().unwrap();
654                    if self.flush_decommit_queue(queue) {
655                        match self.memories.allocate(request, ty, memory_index).await {
656                            Ok(result) => return Ok(result),
657                            Err(err) => e = err,
658                        }
659                    }
660                }
661
662                Err(e)
663            }
664            .await
665            .inspect(|_| {
666                self.live_memories.fetch_add(1, Ordering::Relaxed);
667            })
668        })
669    }
670
671    unsafe fn deallocate_memory(
672        &self,
673        _memory_index: Option<DefinedMemoryIndex>,
674        allocation_index: MemoryAllocationIndex,
675        memory: Memory,
676    ) {
677        let prev = self.live_memories.fetch_sub(1, Ordering::Relaxed);
678        debug_assert!(prev > 0);
679
680        // Reset the image slot. Depending on whether this is successful or not
681        // the `image` is preserved for future use. On success it's queued up to
682        // get deallocated later, and on failure the slot is deallocated
683        // immediately without preserving the image.
684        let mut image = memory.unwrap_static_image();
685        let mut queue = DecommitQueue::default();
686        let bytes_resident = image.clear_and_remain_ready(
687            self.pagemap.as_ref(),
688            self.memories.keep_resident,
689            |ptr, len| {
690                // SAFETY: the memory in `image` won't be used until this
691                // decommit queue is flushed, and by definition the memory is
692                // not in use when calling this function.
693                unsafe {
694                    queue.push_raw(ptr, len);
695                }
696            },
697        );
698
699        match bytes_resident {
700            Ok(bytes_resident) => {
701                // SAFETY: this image is not in use and its memory regions were enqueued
702                // with `push_raw` above.
703                unsafe {
704                    queue.push_memory(allocation_index, image, bytes_resident);
705                }
706                self.merge_or_flush(queue);
707            }
708            Err(e) => {
709                log::warn!("ignoring clear_and_remain_ready error {e}");
710                // SAFETY: `allocation_index` comes from this pool, as an unsafe
711                // contract of this function itself, and it's guaranteed to be no
712                // longer in use so safe to deallocate. The slot couldn't be
713                // preserved so it's dropped here.
714                //
715                // Note that at this point it's not clear how many bytes are
716                // resident in memory, so it's inevitably going to leave statistics
717                // a little off. Also note though that non-Linux platforms don't
718                // keep track of resident bytes anyway, and this path is only
719                // reachable on non-Linux platforms because Linux can't return an
720                // error.
721                unsafe {
722                    self.memories.deallocate(allocation_index, None, 0);
723                }
724            }
725        }
726    }
727
728    fn allocate_table<'a, 'b: 'a, 'c: 'a>(
729        &'a self,
730        request: &'a mut InstanceAllocationRequest<'b, 'c>,
731        ty: &'a wasmtime_environ::Table,
732        _table_index: DefinedTableIndex,
733    ) -> Pin<Box<dyn Future<Output = Result<(super::TableAllocationIndex, Table)>> + Send + 'a>>
734    {
735        crate::runtime::box_future(async move {
736            async {
737                // FIXME: see `allocate_memory` above for comments about duplication
738                // with `with_flush_and_retry`.
739                let mut e = match self.tables.allocate(request, ty).await {
740                    Ok(result) => return Ok(result),
741                    Err(e) => e,
742                };
743
744                for shard in self.decommit_shard_ids() {
745                    if !e.is::<PoolConcurrencyLimitError>() {
746                        break;
747                    }
748                    let queue = self.decommit_queue(shard).lock().unwrap();
749                    if self.flush_decommit_queue(queue) {
750                        match self.tables.allocate(request, ty).await {
751                            Ok(result) => return Ok(result),
752                            Err(err) => e = err,
753                        }
754                    }
755                }
756
757                Err(e)
758            }
759            .await
760            .inspect(|_| {
761                self.live_tables.fetch_add(1, Ordering::Relaxed);
762            })
763        })
764    }
765
766    unsafe fn deallocate_table(
767        &self,
768        _table_index: DefinedTableIndex,
769        allocation_index: TableAllocationIndex,
770        mut table: Table,
771    ) {
772        let prev = self.live_tables.fetch_sub(1, Ordering::Relaxed);
773        debug_assert!(prev > 0);
774
775        let mut queue = DecommitQueue::default();
776        // SAFETY: This table is no longer in use by the allocator when this
777        // method is called and additionally all image ranges are pushed with
778        // the understanding that the memory won't get used until the whole
779        // queue is flushed.
780        let bytes_resident = unsafe {
781            self.tables.reset_table_pages_to_zero(
782                self.pagemap.as_ref(),
783                allocation_index,
784                &mut table,
785                |ptr, len| {
786                    queue.push_raw(ptr, len);
787                },
788            )
789        };
790
791        // SAFETY: the table has had all its memory regions enqueued above.
792        unsafe {
793            queue.push_table(allocation_index, table, bytes_resident);
794        }
795        self.merge_or_flush(queue);
796    }
797
798    #[cfg(feature = "async")]
799    fn allocate_fiber_stack(&self) -> Result<wasmtime_fiber::FiberStack> {
800        let ret = self.with_flush_and_retry(|| self.stacks.allocate())?;
801        self.live_stacks.fetch_add(1, Ordering::Relaxed);
802        Ok(ret)
803    }
804
805    #[cfg(feature = "async")]
806    unsafe fn deallocate_fiber_stack(&self, mut stack: wasmtime_fiber::FiberStack) {
807        self.live_stacks.fetch_sub(1, Ordering::Relaxed);
808        let mut queue = DecommitQueue::default();
809        // SAFETY: the stack is no longer in use by definition when this
810        // function is called and memory ranges pushed here are otherwise no
811        // longer in use.
812        let bytes_resident = unsafe {
813            self.stacks
814                .zero_stack(&mut stack, |ptr, len| queue.push_raw(ptr, len))
815        };
816        // SAFETY: this stack's memory regions were enqueued above.
817        unsafe {
818            queue.push_stack(stack, bytes_resident);
819        }
820        self.merge_or_flush(queue);
821    }
822
823    fn purge_module(&self, module: CompiledModuleId) {
824        self.memories.purge_module(module);
825    }
826
827    fn next_available_pkey(&self) -> Option<ProtectionKey> {
828        self.memories.next_available_pkey()
829    }
830
831    fn restrict_to_pkey(&self, pkey: ProtectionKey) {
832        mpk::allow(ProtectionMask::zero().or(pkey));
833    }
834
835    fn allow_all_pkeys(&self) {
836        mpk::allow(ProtectionMask::all());
837    }
838
839    #[cfg(feature = "gc")]
840    fn allocate_gc_heap(
841        &self,
842        engine: &crate::Engine,
843        gc_runtime: &dyn GcRuntime,
844        memory_alloc_index: MemoryAllocationIndex,
845    ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)> {
846        let ret =
847            self.gc_heaps
848                .as_ref()
849                .unwrap()
850                .allocate(engine, gc_runtime, memory_alloc_index)?;
851        self.live_gc_heaps.fetch_add(1, Ordering::Relaxed);
852        Ok(ret)
853    }
854
855    #[cfg(feature = "gc")]
856    fn deallocate_gc_heap(
857        &self,
858        allocation_index: GcHeapAllocationIndex,
859        gc_heap: Box<dyn GcHeap>,
860    ) -> MemoryAllocationIndex {
861        let gc_heaps = self.gc_heaps.as_ref().unwrap();
862        self.live_gc_heaps.fetch_sub(1, Ordering::Relaxed);
863        gc_heaps.deallocate(allocation_index, gc_heap)
864    }
865
866    fn as_pooling(&self) -> Option<&PoolingInstanceAllocator> {
867        Some(self)
868    }
869}
870
871#[cfg(test)]
872#[cfg(target_pointer_width = "64")]
873mod test {
874    use super::*;
875    use crate::config::InstanceLimits;
876
877    #[test]
878    fn test_pooling_allocator_with_memory_pages_exceeded() {
879        let config = PoolingAllocationConfig {
880            limits: InstanceLimits {
881                total_memories: 1,
882                max_memory_size: 0x100010000,
883                ..Default::default()
884            },
885            ..PoolingAllocationConfig::default()
886        };
887        assert_eq!(
888            PoolingInstanceAllocator::new(
889                &config,
890                &Tunables {
891                    memory_reservation: 0x10000,
892                    ..Tunables::default_host()
893                },
894            )
895            .map_err(|e| e.to_string())
896            .expect_err("expected a failure constructing instance allocator"),
897            "maximum memory size of 0x100010000 bytes exceeds the configured \
898             memory reservation of 0x10000 bytes"
899        );
900    }
901
902    #[cfg(all(
903        unix,
904        target_pointer_width = "64",
905        feature = "async",
906        not(miri),
907        not(asan)
908    ))]
909    #[test]
910    fn test_stack_zeroed() -> Result<()> {
911        let config = PoolingAllocationConfig {
912            max_unused_warm_slots: 0,
913            limits: InstanceLimits {
914                total_stacks: 1,
915                total_memories: 0,
916                total_tables: 0,
917                ..Default::default()
918            },
919            stack_size: 128,
920            async_stack_zeroing: true,
921            ..PoolingAllocationConfig::default()
922        };
923        let allocator = PoolingInstanceAllocator::new(&config, &Tunables::default_host())?;
924
925        unsafe {
926            for _ in 0..255 {
927                let stack = allocator.allocate_fiber_stack()?;
928
929                // The stack pointer is at the top, so decrement it first
930                let addr = stack.top().unwrap().sub(1);
931
932                assert_eq!(*addr, 0);
933                *addr = 1;
934
935                allocator.deallocate_fiber_stack(stack);
936            }
937        }
938
939        Ok(())
940    }
941
942    #[cfg(all(
943        unix,
944        target_pointer_width = "64",
945        feature = "async",
946        not(miri),
947        not(asan)
948    ))]
949    #[test]
950    fn test_stack_unzeroed() -> Result<()> {
951        let config = PoolingAllocationConfig {
952            max_unused_warm_slots: 0,
953            limits: InstanceLimits {
954                total_stacks: 1,
955                total_memories: 0,
956                total_tables: 0,
957                ..Default::default()
958            },
959            stack_size: 128,
960            async_stack_zeroing: false,
961            ..PoolingAllocationConfig::default()
962        };
963        let allocator = PoolingInstanceAllocator::new(&config, &Tunables::default_host())?;
964
965        unsafe {
966            for i in 0..255 {
967                let stack = allocator.allocate_fiber_stack()?;
968
969                // The stack pointer is at the top, so decrement it first
970                let addr = stack.top().unwrap().sub(1);
971
972                assert_eq!(*addr, i);
973                *addr = i + 1;
974
975                allocator.deallocate_fiber_stack(stack);
976            }
977        }
978
979        Ok(())
980    }
981}