From 3dbdcfa220df1406b3627772194eff859104e17a Mon Sep 17 00:00:00 2001 From: Andrew Brown Date: Fri, 29 Apr 2022 08:12:38 -0700 Subject: [PATCH] runtime: refactor `Memory` to always use `Box` (#4086) While working with the runtime `Memory` object, it became clear that some refactoring was needed. In order to implement shared memory from the threads proposal, we must be able to atomically change the memory size. Previously, the split into variants, `Memory::Static` and `Memory::Dynamic`, made any attempt to lock forced us to duplicate logic in various places. This change moves `enum Memory { Static..., Dynamic... }` to simply `struct Memory(Box)`. A new type, `ExternalMemory`, takes the place of `Memory::Static` and also implements the `RuntimeLinearMemory` trait, allowing `Memory` to contain the same two options as before: `MmapMemory` for `Memory::Dynamic` and `ExternalMemory` for `Memory::Static`. To interface with the `PoolingAllocator`, this change also required the ability to downcast to the internal representation. --- .../runtime/src/instance/allocator/pooling.rs | 41 +-- crates/runtime/src/memory.rs | 296 +++++++++--------- crates/wasmtime/src/trampoline/memory.rs | 7 +- 3 files changed, 176 insertions(+), 168 deletions(-) diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index bb7ab1938d..136a17da99 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -445,34 +445,25 @@ impl InstancePool { instance_index: usize, memories: &mut PrimaryMap, ) { - // Decommit any linear memories that were used - for ((def_mem_idx, memory), base) in - memories.iter_mut().zip(self.memories.get(instance_index)) + // Decommit any linear memories that were used. + let memories = mem::take(memories); + for ((def_mem_idx, mut memory), base) in + memories.into_iter().zip(self.memories.get(instance_index)) { - let memory = mem::take(memory); assert!(memory.is_static()); - - match memory { - Memory::Static { - memory_image: Some(mut image), - .. - } => { - // If there was any error clearing the image, just - // drop it here, and let the drop handler for the - // slot unmap in a way that retains the - // address space reservation. - if image.clear_and_remain_ready().is_ok() { - self.memories - .return_memory_image_slot(instance_index, def_mem_idx, image); - } - } - - _ => { - let size = memory.byte_size(); - drop(memory); - decommit_memory_pages(base, size) - .expect("failed to decommit linear memory pages"); + let size = memory.byte_size(); + if let Some(mut image) = memory.unwrap_static_image() { + // Reset the image slot. If there is any error clearing the + // image, just drop it here, and let the drop handler for the + // slot unmap in a way that retains the address space + // reservation. + if image.clear_and_remain_ready().is_ok() { + self.memories + .return_memory_image_slot(instance_index, def_mem_idx, image); } + } else { + // Otherwise, decommit the memory pages. + decommit_memory_pages(base, size).expect("failed to decommit linear memory pages"); } } } diff --git a/crates/runtime/src/memory.rs b/crates/runtime/src/memory.rs index 4295830b5d..0d402dad4e 100644 --- a/crates/runtime/src/memory.rs +++ b/crates/runtime/src/memory.rs @@ -68,12 +68,17 @@ pub trait RuntimeLinearMemory: Send + Sync { /// Return a `VMMemoryDefinition` for exposing the memory to compiled wasm /// code. - fn vmmemory(&self) -> VMMemoryDefinition; + fn vmmemory(&mut self) -> VMMemoryDefinition; /// Does this memory need initialization? It may not if it already /// has initial contents courtesy of the `MemoryImage` passed to /// `RuntimeMemoryCreator::new_memory()`. fn needs_init(&self) -> bool; + + /// For the pooling allocator, we must be able to downcast this trait to its + /// underlying structure. + #[cfg(feature = "pooling-allocator")] + fn as_any_mut(&mut self) -> &mut dyn std::any::Any; } /// A linear memory instance. @@ -242,7 +247,7 @@ impl RuntimeLinearMemory for MmapMemory { Ok(()) } - fn vmmemory(&self) -> VMMemoryDefinition { + fn vmmemory(&mut self) -> VMMemoryDefinition { VMMemoryDefinition { base: unsafe { self.mmap.as_mut_ptr().add(self.pre_guard_size) }, current_length: self.accessible, @@ -254,35 +259,130 @@ impl RuntimeLinearMemory for MmapMemory { // is needed. self.memory_image.is_none() } + + #[cfg(feature = "pooling-allocator")] + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } } -/// Representation of a runtime wasm linear memory. -pub enum Memory { - /// A "static" memory where the lifetime of the backing memory is managed - /// elsewhere. Currently used with the pooling allocator. - Static { - /// The memory in the host for this wasm memory. The length of this - /// slice is the maximum size of the memory that can be grown to. - base: &'static mut [u8], +/// A "static" memory where the lifetime of the backing memory is managed +/// elsewhere. Currently used with the pooling allocator. +struct ExternalMemory { + /// The memory in the host for this wasm memory. The length of this + /// slice is the maximum size of the memory that can be grown to. + base: &'static mut [u8], - /// The current size, in bytes, of this memory. - size: usize, + /// The current size, in bytes, of this memory. + size: usize, - /// A callback which makes portions of `base` accessible for when memory - /// is grown. Otherwise it's expected that accesses to `base` will - /// fault. - make_accessible: Option Result<()>>, + /// A callback which makes portions of `base` accessible for when memory + /// is grown. Otherwise it's expected that accesses to `base` will + /// fault. + make_accessible: Option Result<()>>, + + /// The image management, if any, for this memory. Owned here and + /// returned to the pooling allocator when termination occurs. + memory_image: Option, +} - /// The image management, if any, for this memory. Owned here and - /// returned to the pooling allocator when termination occurs. +impl ExternalMemory { + fn new( + base: &'static mut [u8], + initial_size: usize, + maximum_size: Option, + make_accessible: Option Result<()>>, memory_image: Option, - }, + ) -> Result { + if base.len() < initial_size { + bail!( + "initial memory size of {} exceeds the pooling allocator's \ + configured maximum memory size of {} bytes", + initial_size, + base.len(), + ); + } + + // Only use the part of the slice that is necessary. + let base = match maximum_size { + Some(max) if max < base.len() => &mut base[..max], + _ => base, + }; + + if let Some(make_accessible) = make_accessible { + if initial_size > 0 { + make_accessible(base.as_mut_ptr(), initial_size)?; + } + } - /// A "dynamic" memory whose data is managed at runtime and lifetime is tied - /// to this instance. - Dynamic(Box), + Ok(Self { + base, + size: initial_size, + make_accessible, + memory_image, + }) + } } +impl RuntimeLinearMemory for ExternalMemory { + fn byte_size(&self) -> usize { + self.size + } + + fn maximum_byte_size(&self) -> Option { + Some(self.base.len()) + } + + fn grow_to(&mut self, new_byte_size: usize) -> Result<()> { + // Never exceed the static memory size; this check should have been made + // prior to arriving here. + assert!(new_byte_size <= self.base.len()); + + // Actually grow the memory. + if let Some(image) = &mut self.memory_image { + image.set_heap_limit(new_byte_size)?; + } else { + let make_accessible = self + .make_accessible + .expect("make_accessible must be Some if this is not a CoW memory"); + + // Operating system can fail to make memory accessible. + let old_byte_size = self.byte_size(); + make_accessible( + unsafe { self.base.as_mut_ptr().add(old_byte_size) }, + new_byte_size - old_byte_size, + )?; + } + + // Update our accounting of the available size. + self.size = new_byte_size; + Ok(()) + } + + fn vmmemory(&mut self) -> VMMemoryDefinition { + VMMemoryDefinition { + base: self.base.as_mut_ptr().cast(), + current_length: self.size, + } + } + + fn needs_init(&self) -> bool { + if let Some(slot) = &self.memory_image { + !slot.has_image() + } else { + true + } + } + + #[cfg(feature = "pooling-allocator")] + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } +} + +/// Representation of a runtime wasm linear memory. +pub struct Memory(Box); + impl Memory { /// Create a new dynamic (movable) memory instance for the specified plan. pub fn new_dynamic( @@ -292,7 +392,7 @@ impl Memory { memory_image: Option<&Arc>, ) -> Result { let (minimum, maximum) = Self::limit_new(plan, store)?; - Ok(Memory::Dynamic(creator.new_memory( + Ok(Memory(creator.new_memory( plan, minimum, maximum, @@ -309,33 +409,9 @@ impl Memory { store: &mut dyn Store, ) -> Result { let (minimum, maximum) = Self::limit_new(plan, store)?; - - if base.len() < minimum { - bail!( - "initial memory size of {} exceeds the pooling allocator's \ - configured maximum memory size of {} bytes", - minimum, - base.len(), - ); - } - - let base = match maximum { - Some(max) if max < base.len() => &mut base[..max], - _ => base, - }; - - if let Some(make_accessible) = make_accessible { - if minimum > 0 { - make_accessible(base.as_mut_ptr(), minimum)?; - } - } - - Ok(Memory::Static { - base, - size: minimum, - make_accessible, - memory_image, - }) + let pooled_memory = + ExternalMemory::new(base, minimum, maximum, make_accessible, memory_image)?; + Ok(Memory(Box::new(pooled_memory))) } /// Calls the `store`'s limiter to optionally prevent a memory from being allocated. @@ -423,10 +499,7 @@ impl Memory { /// Returns the number of allocated wasm pages. pub fn byte_size(&self) -> usize { - match self { - Memory::Static { size, .. } => *size, - Memory::Dynamic(mem) => mem.byte_size(), - } + self.0.byte_size() } /// Returns the maximum number of pages the memory can grow to at runtime. @@ -436,34 +509,14 @@ impl Memory { /// The runtime maximum may not be equal to the maximum from the linear memory's /// Wasm type when it is being constrained by an instance allocator. pub fn maximum_byte_size(&self) -> Option { - match self { - Memory::Static { base, .. } => Some(base.len()), - Memory::Dynamic(mem) => mem.maximum_byte_size(), - } - } - - /// Returns whether or not the underlying storage of the memory is "static". - #[cfg(feature = "pooling-allocator")] - pub(crate) fn is_static(&self) -> bool { - if let Memory::Static { .. } = self { - true - } else { - false - } + self.0.maximum_byte_size() } /// Returns whether or not this memory needs initialization. It /// may not if it already has initial content thanks to a CoW /// mechanism. pub(crate) fn needs_init(&self) -> bool { - match self { - Memory::Static { - memory_image: Some(slot), - .. - } => !slot.has_image(), - Memory::Dynamic(mem) => mem.needs_init(), - _ => true, - } + self.0.needs_init() } /// Grow memory by the specified amount of wasm pages. @@ -489,6 +542,7 @@ impl Memory { store: &mut dyn Store, ) -> Result, Error> { let old_byte_size = self.byte_size(); + // Wasm spec: when growing by 0 pages, always return the current size. if delta_pages == 0 { return Ok(Some(old_byte_size)); @@ -524,80 +578,38 @@ impl Memory { } } - match self { - Memory::Static { - base, - size, - memory_image: Some(image), - .. - } => { - // Never exceed static memory size - if new_byte_size > base.len() { - store.memory_grow_failed(&format_err!("static memory size exceeded")); - return Ok(None); - } - - if let Err(e) = image.set_heap_limit(new_byte_size) { - store.memory_grow_failed(&e); - return Ok(None); - } - *size = new_byte_size; - } - Memory::Static { - base, - size, - make_accessible, - .. - } => { - let make_accessible = make_accessible - .expect("make_accessible must be Some if this is not a CoW memory"); - - // Never exceed static memory size - if new_byte_size > base.len() { - store.memory_grow_failed(&format_err!("static memory size exceeded")); - return Ok(None); - } - - // Operating system can fail to make memory accessible - if let Err(e) = make_accessible( - base.as_mut_ptr().add(old_byte_size), - new_byte_size - old_byte_size, - ) { - store.memory_grow_failed(&e); - return Ok(None); - } - *size = new_byte_size; - } - Memory::Dynamic(mem) => { - if let Err(e) = mem.grow_to(new_byte_size) { - store.memory_grow_failed(&e); - return Ok(None); - } + match self.0.grow_to(new_byte_size) { + Ok(_) => Ok(Some(old_byte_size)), + Err(e) => { + store.memory_grow_failed(&e); + Ok(None) } } - Ok(Some(old_byte_size)) } /// Return a `VMMemoryDefinition` for exposing the memory to compiled wasm code. pub fn vmmemory(&mut self) -> VMMemoryDefinition { - match self { - Memory::Static { base, size, .. } => VMMemoryDefinition { - base: base.as_mut_ptr().cast(), - current_length: *size, - }, - Memory::Dynamic(mem) => mem.vmmemory(), - } + self.0.vmmemory() + } + + /// Check if the inner implementation of [`Memory`] is a memory created with + /// [`Memory::new_static()`]. + #[cfg(feature = "pooling-allocator")] + pub fn is_static(&mut self) -> bool { + let as_any = self.0.as_any_mut(); + as_any.downcast_ref::().is_some() } -} -// The default memory representation is an empty memory that cannot grow. -impl Default for Memory { - fn default() -> Self { - Memory::Static { - base: &mut [], - size: 0, - make_accessible: Some(|_, _| unreachable!()), - memory_image: None, + /// Consume the memory, returning its [`MemoryImageSlot`] if any is present. + /// The image should only be present for a subset of memories created with + /// [`Memory::new_static()`]. + #[cfg(feature = "pooling-allocator")] + pub fn unwrap_static_image(mut self) -> Option { + let as_any = self.0.as_any_mut(); + if let Some(m) = as_any.downcast_mut::() { + std::mem::take(&mut m.memory_image) + } else { + None } } } diff --git a/crates/wasmtime/src/trampoline/memory.rs b/crates/wasmtime/src/trampoline/memory.rs index 2bb91d644e..1a5858182f 100644 --- a/crates/wasmtime/src/trampoline/memory.rs +++ b/crates/wasmtime/src/trampoline/memory.rs @@ -42,7 +42,7 @@ impl RuntimeLinearMemory for LinearMemoryProxy { self.mem.grow_to(new_size) } - fn vmmemory(&self) -> VMMemoryDefinition { + fn vmmemory(&mut self) -> VMMemoryDefinition { VMMemoryDefinition { base: self.mem.as_ptr(), current_length: self.mem.byte_size(), @@ -52,6 +52,11 @@ impl RuntimeLinearMemory for LinearMemoryProxy { fn needs_init(&self) -> bool { true } + + #[cfg(feature = "pooling-allocator")] + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } } #[derive(Clone)]