wasm-executor: Support growing the memory (#12520)

* As always, start with something :P * Add support for max_heap_pages * Add support for wasmtime * Make it compile * Fix compilation * Copy wrongly merged code * Fix compilation * Some fixes * Fix * Get stuff working * More work * More fixes * ... * More * FIXEs * Switch wasmi to use `RuntimeBlob` like wasmtime * Removed unused stuff * Cleanup * More cleanups * Introduce `CallContext` * Fixes * More fixes * Add builder for creating the `WasmExecutor` * Adds some docs * FMT * First round of feedback. * Review feedback round 2 * More fixes * Fix try-runtime * Update client/executor/wasmtime/src/instance_wrapper.rs Co-authored-by: Koute <koute@users.noreply.github.com> * Update client/executor/common/src/wasm_runtime.rs Co-authored-by: Koute <koute@users.noreply.github.com> * Update client/executor/common/src/runtime_blob/runtime_blob.rs Co-authored-by: Koute <koute@users.noreply.github.com> * Update client/executor/common/src/wasm_runtime.rs Co-authored-by: Koute <koute@users.noreply.github.com> * Update client/allocator/src/freeing_bump.rs Co-authored-by: Koute <koute@users.noreply.github.com> * Update client/allocator/src/freeing_bump.rs Co-authored-by: Koute <koute@users.noreply.github.com> * Feedback round 3 * FMT * Review comments --------- Co-authored-by: Koute <koute@users.noreply.github.com>
2026-06-12 21:41:12 +00:00 · 2023-02-24 12:43:01 +01:00
parent c848d40775
commit 941288c6d0
37 changed files with 1092 additions and 667 deletions
@@ -16,7 +16,7 @@
 // limitations under the License.

 /// The error type used by the allocators.
-#[derive(thiserror::Error, Debug)]
+#[derive(thiserror::Error, Debug, PartialEq)]
 pub enum Error {
 	/// Someone tried to allocate more memory than the allowed maximum per allocation.
 	#[error("Requested allocation size is too large")]
@@ -67,10 +67,11 @@
 //!   wasted. This is more pronounced (in terms of absolute heap amounts) with larger allocation
 //!   sizes.

-use crate::Error;
+use crate::{Error, Memory, MAX_WASM_PAGES, PAGE_SIZE};
 pub use sp_core::MAX_POSSIBLE_ALLOCATION;
 use sp_wasm_interface::{Pointer, WordSize};
 use std::{
+	cmp::{max, min},
 	mem,
 	ops::{Index, IndexMut, Range},
 };
@@ -237,7 +238,7 @@ impl Header {
 	///
 	/// Returns an error if the `header_ptr` is out of bounds of the linear memory or if the read
 	/// header is corrupted (e.g. the order is incorrect).
-	fn read_from<M: Memory + ?Sized>(memory: &M, header_ptr: u32) -> Result<Self, Error> {
+	fn read_from(memory: &impl Memory, header_ptr: u32) -> Result<Self, Error> {
 		let raw_header = memory.read_le_u64(header_ptr)?;

 		// Check if the header represents an occupied or free allocation and extract the header data
@@ -255,7 +256,7 @@ impl Header {
 	/// Write out this header to memory.
 	///
 	/// Returns an error if the `header_ptr` is out of bounds of the linear memory.
-	fn write_into<M: Memory + ?Sized>(&self, memory: &mut M, header_ptr: u32) -> Result<(), Error> {
+	fn write_into(&self, memory: &mut impl Memory, header_ptr: u32) -> Result<(), Error> {
 		let (header_data, occupied_mask) = match *self {
 			Self::Occupied(order) => (order.into_raw(), 0x00000001_00000000),
 			Self::Free(link) => (link.into_raw(), 0x00000000_00000000),
@@ -343,6 +344,15 @@ pub struct AllocationStats {
 	pub address_space_used: u32,
 }

+/// Convert the given `size` in bytes into the number of pages.
+///
+/// The returned number of pages is ensured to be big enough to hold memory with the given `size`.
+///
+/// Returns `None` if the number of pages to not fit into `u32`.
+fn pages_from_size(size: u64) -> Option<u32> {
+	u32::try_from((size + PAGE_SIZE as u64 - 1) / PAGE_SIZE as u64).ok()
+}
+
 /// An implementation of freeing bump allocator.
 ///
 /// Refer to the module-level documentation for further details.
@@ -351,7 +361,7 @@ pub struct FreeingBumpHeapAllocator {
 	bumper: u32,
 	free_lists: FreeLists,
 	poisoned: bool,
-	last_observed_memory_size: u32,
+	last_observed_memory_size: u64,
 	stats: AllocationStats,
 }

@@ -395,9 +405,9 @@ impl FreeingBumpHeapAllocator {
 	///
 	/// - `mem` - a slice representing the linear memory on which this allocator operates.
 	/// - `size` - size in bytes of the allocation request
-	pub fn allocate<M: Memory + ?Sized>(
+	pub fn allocate(
 		&mut self,
-		mem: &mut M,
+		mem: &mut impl Memory,
 		size: WordSize,
 	) -> Result<Pointer<u8>, Error> {
 		if self.poisoned {
@@ -412,7 +422,7 @@ impl FreeingBumpHeapAllocator {
 		let header_ptr: u32 = match self.free_lists[order] {
 			Link::Ptr(header_ptr) => {
 				assert!(
-					header_ptr + order.size() + HEADER_SIZE <= mem.size(),
+					u64::from(header_ptr + order.size() + HEADER_SIZE) <= mem.size(),
 					"Pointer is looked up in list of free entries, into which
 					only valid values are inserted; qed"
 				);
@@ -427,7 +437,7 @@ impl FreeingBumpHeapAllocator {
 			},
 			Link::Nil => {
 				// Corresponding free list is empty. Allocate a new item.
-				Self::bump(&mut self.bumper, order.size() + HEADER_SIZE, mem.size())?
+				Self::bump(&mut self.bumper, order.size() + HEADER_SIZE, mem)?
 			},
 		};

@@ -437,7 +447,7 @@ impl FreeingBumpHeapAllocator {
 		self.stats.bytes_allocated += order.size() + HEADER_SIZE;
 		self.stats.bytes_allocated_sum += u128::from(order.size() + HEADER_SIZE);
 		self.stats.bytes_allocated_peak =
-			std::cmp::max(self.stats.bytes_allocated_peak, self.stats.bytes_allocated);
+			max(self.stats.bytes_allocated_peak, self.stats.bytes_allocated);
 		self.stats.address_space_used = self.bumper - self.original_heap_base;

 		log::trace!(target: LOG_TARGET, "after allocation: {:?}", self.stats);
@@ -457,11 +467,7 @@ impl FreeingBumpHeapAllocator {
 	///
 	/// - `mem` - a slice representing the linear memory on which this allocator operates.
 	/// - `ptr` - pointer to the allocated chunk
-	pub fn deallocate<M: Memory + ?Sized>(
-		&mut self,
-		mem: &mut M,
-		ptr: Pointer<u8>,
-	) -> Result<(), Error> {
+	pub fn deallocate(&mut self, mem: &mut impl Memory, ptr: Pointer<u8>) -> Result<(), Error> {
 		if self.poisoned {
 			return Err(error("the allocator has been poisoned"))
 		}
@@ -503,15 +509,52 @@ impl FreeingBumpHeapAllocator {
 	///
 	/// Returns the `bumper` from before the increase. Returns an `Error::AllocatorOutOfSpace` if
 	/// the operation would exhaust the heap.
-	fn bump(bumper: &mut u32, size: u32, heap_end: u32) -> Result<u32, Error> {
-		if *bumper + size > heap_end {
-			log::error!(
-				target: LOG_TARGET,
-				"running out of space with current bumper {}, mem size {}",
-				bumper,
-				heap_end
+	fn bump(bumper: &mut u32, size: u32, memory: &mut impl Memory) -> Result<u32, Error> {
+		let required_size = u64::from(*bumper) + u64::from(size);
+
+		if required_size > memory.size() {
+			let required_pages =
+				pages_from_size(required_size).ok_or_else(|| Error::AllocatorOutOfSpace)?;
+
+			let current_pages = memory.pages();
+			let max_pages = memory.max_pages().unwrap_or(MAX_WASM_PAGES);
+			debug_assert!(
+				current_pages < required_pages,
+				"current pages {current_pages} < required pages {required_pages}"
 			);
-			return Err(Error::AllocatorOutOfSpace)
+
+			if current_pages >= max_pages {
+				log::debug!(
+					target: LOG_TARGET,
+					"Wasm pages ({current_pages}) are already at the maximum.",
+				);
+
+				return Err(Error::AllocatorOutOfSpace)
+			} else if required_pages > max_pages {
+				log::debug!(
+					target: LOG_TARGET,
+					"Failed to grow memory from {current_pages} pages to at least {required_pages}\
+						 pages due to the maximum limit of {max_pages} pages",
+				);
+				return Err(Error::AllocatorOutOfSpace)
+			}
+
+			// Ideally we want to double our current number of pages,
+			// as long as it's less than the absolute maximum we can have.
+			let next_pages = min(current_pages * 2, max_pages);
+			// ...but if even more pages are required then try to allocate that many.
+			let next_pages = max(next_pages, required_pages);
+
+			if memory.grow(next_pages - current_pages).is_err() {
+				log::error!(
+					target: LOG_TARGET,
+					"Failed to grow memory from {current_pages} pages to {next_pages} pages",
+				);
+
+				return Err(Error::AllocatorOutOfSpace)
+			}
+
+			debug_assert_eq!(memory.pages(), next_pages, "Number of pages should have increased!");
 		}

 		let res = *bumper;
@@ -519,9 +562,9 @@ impl FreeingBumpHeapAllocator {
 		Ok(res)
 	}

-	fn observe_memory_size<M: Memory + ?Sized>(
-		last_observed_memory_size: &mut u32,
-		mem: &mut M,
+	fn observe_memory_size(
+		last_observed_memory_size: &mut u64,
+		mem: &mut impl Memory,
 	) -> Result<(), Error> {
 		if mem.size() < *last_observed_memory_size {
 			return Err(Error::MemoryShrinked)
@@ -538,37 +581,41 @@ impl FreeingBumpHeapAllocator {
 /// accessible up to the reported size.
 ///
 /// The linear memory can grow in size with the wasm page granularity (64KiB), but it cannot shrink.
-pub trait Memory {
+trait MemoryExt: Memory {
 	/// Read a u64 from the heap in LE form. Returns an error if any of the bytes read are out of
 	/// bounds.
-	fn read_le_u64(&self, ptr: u32) -> Result<u64, Error>;
+	fn read_le_u64(&self, ptr: u32) -> Result<u64, Error> {
+		self.with_access(|memory| {
+			let range =
+				heap_range(ptr, 8, memory.len()).ok_or_else(|| error("read out of heap bounds"))?;
+			let bytes = memory[range]
+				.try_into()
+				.expect("[u8] slice of length 8 must be convertible to [u8; 8]");
+			Ok(u64::from_le_bytes(bytes))
+		})
+	}
+
 	/// Write a u64 to the heap in LE form. Returns an error if any of the bytes written are out of
 	/// bounds.
-	fn write_le_u64(&mut self, ptr: u32, val: u64) -> Result<(), Error>;
+	fn write_le_u64(&mut self, ptr: u32, val: u64) -> Result<(), Error> {
+		self.with_access_mut(|memory| {
+			let range = heap_range(ptr, 8, memory.len())
+				.ok_or_else(|| error("write out of heap bounds"))?;
+			let bytes = val.to_le_bytes();
+			memory[range].copy_from_slice(&bytes[..]);
+			Ok(())
+		})
+	}
+
 	/// Returns the full size of the memory in bytes.
-	fn size(&self) -> u32;
+	fn size(&self) -> u64 {
+		debug_assert!(self.pages() <= MAX_WASM_PAGES);
+
+		self.pages() as u64 * PAGE_SIZE as u64
+	}
 }

-impl Memory for [u8] {
-	fn read_le_u64(&self, ptr: u32) -> Result<u64, Error> {
-		let range =
-			heap_range(ptr, 8, self.len()).ok_or_else(|| error("read out of heap bounds"))?;
-		let bytes = self[range]
-			.try_into()
-			.expect("[u8] slice of length 8 must be convertible to [u8; 8]");
-		Ok(u64::from_le_bytes(bytes))
-	}
-	fn write_le_u64(&mut self, ptr: u32, val: u64) -> Result<(), Error> {
-		let range =
-			heap_range(ptr, 8, self.len()).ok_or_else(|| error("write out of heap bounds"))?;
-		let bytes = val.to_le_bytes();
-		self[range].copy_from_slice(&bytes[..]);
-		Ok(())
-	}
-	fn size(&self) -> u32 {
-		u32::try_from(self.len()).expect("size of Wasm linear memory is <2^32; qed")
-	}
-}
+impl<T: Memory> MemoryExt for T {}

 fn heap_range(offset: u32, length: u32, heap_len: usize) -> Option<Range<usize>> {
 	let start = offset as usize;
@@ -601,21 +648,72 @@ impl<'a> Drop for PoisonBomb<'a> {
 mod tests {
 	use super::*;

-	const PAGE_SIZE: u32 = 65536;
-
 	/// Makes a pointer out of the given address.
 	fn to_pointer(address: u32) -> Pointer<u8> {
 		Pointer::new(address)
 	}

+	#[derive(Debug)]
+	struct MemoryInstance {
+		data: Vec<u8>,
+		max_wasm_pages: u32,
+	}
+
+	impl MemoryInstance {
+		fn with_pages(pages: u32) -> Self {
+			Self { data: vec![0; (pages * PAGE_SIZE) as usize], max_wasm_pages: MAX_WASM_PAGES }
+		}
+
+		fn set_max_wasm_pages(&mut self, max_pages: u32) {
+			self.max_wasm_pages = max_pages;
+		}
+	}
+
+	impl Memory for MemoryInstance {
+		fn with_access<R>(&self, run: impl FnOnce(&[u8]) -> R) -> R {
+			run(&self.data)
+		}
+
+		fn with_access_mut<R>(&mut self, run: impl FnOnce(&mut [u8]) -> R) -> R {
+			run(&mut self.data)
+		}
+
+		fn pages(&self) -> u32 {
+			pages_from_size(self.data.len() as u64).unwrap()
+		}
+
+		fn max_pages(&self) -> Option<u32> {
+			Some(self.max_wasm_pages)
+		}
+
+		fn grow(&mut self, pages: u32) -> Result<(), ()> {
+			if self.pages() + pages > self.max_wasm_pages {
+				Err(())
+			} else {
+				self.data.resize(((self.pages() + pages) * PAGE_SIZE) as usize, 0);
+				Ok(())
+			}
+		}
+	}
+
+	#[test]
+	fn test_pages_from_size() {
+		assert_eq!(pages_from_size(0).unwrap(), 0);
+		assert_eq!(pages_from_size(1).unwrap(), 1);
+		assert_eq!(pages_from_size(65536).unwrap(), 1);
+		assert_eq!(pages_from_size(65536 + 1).unwrap(), 2);
+		assert_eq!(pages_from_size(2 * 65536).unwrap(), 2);
+		assert_eq!(pages_from_size(2 * 65536 + 1).unwrap(), 3);
+	}
+
 	#[test]
 	fn should_allocate_properly() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

 		// when
-		let ptr = heap.allocate(&mut mem[..], 1).unwrap();
+		let ptr = heap.allocate(&mut mem, 1).unwrap();

 		// then
 		// returned pointer must start right after `HEADER_SIZE`
@@ -625,11 +723,11 @@ mod tests {
 	#[test]
 	fn should_always_align_pointers_to_multiples_of_8() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(13);

 		// when
-		let ptr = heap.allocate(&mut mem[..], 1).unwrap();
+		let ptr = heap.allocate(&mut mem, 1).unwrap();

 		// then
 		// the pointer must start at the next multiple of 8 from 13
@@ -640,13 +738,13 @@ mod tests {
 	#[test]
 	fn should_increment_pointers_properly() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

 		// when
-		let ptr1 = heap.allocate(&mut mem[..], 1).unwrap();
-		let ptr2 = heap.allocate(&mut mem[..], 9).unwrap();
-		let ptr3 = heap.allocate(&mut mem[..], 1).unwrap();
+		let ptr1 = heap.allocate(&mut mem, 1).unwrap();
+		let ptr2 = heap.allocate(&mut mem, 9).unwrap();
+		let ptr3 = heap.allocate(&mut mem, 1).unwrap();

 		// then
 		// a prefix of 8 bytes is prepended to each pointer
@@ -663,18 +761,18 @@ mod tests {
 	#[test]
 	fn should_free_properly() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);
-		let ptr1 = heap.allocate(&mut mem[..], 1).unwrap();
+		let ptr1 = heap.allocate(&mut mem, 1).unwrap();
 		// the prefix of 8 bytes is prepended to the pointer
 		assert_eq!(ptr1, to_pointer(HEADER_SIZE));

-		let ptr2 = heap.allocate(&mut mem[..], 1).unwrap();
+		let ptr2 = heap.allocate(&mut mem, 1).unwrap();
 		// the prefix of 8 bytes + the content of ptr 1 is prepended to the pointer
 		assert_eq!(ptr2, to_pointer(24));

 		// when
-		heap.deallocate(&mut mem[..], ptr2).unwrap();
+		heap.deallocate(&mut mem, ptr2).unwrap();

 		// then
 		// then the heads table should contain a pointer to the
@@ -685,23 +783,23 @@ mod tests {
 	#[test]
 	fn should_deallocate_and_reallocate_properly() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let padded_offset = 16;
 		let mut heap = FreeingBumpHeapAllocator::new(13);

-		let ptr1 = heap.allocate(&mut mem[..], 1).unwrap();
+		let ptr1 = heap.allocate(&mut mem, 1).unwrap();
 		// the prefix of 8 bytes is prepended to the pointer
 		assert_eq!(ptr1, to_pointer(padded_offset + HEADER_SIZE));

-		let ptr2 = heap.allocate(&mut mem[..], 9).unwrap();
+		let ptr2 = heap.allocate(&mut mem, 9).unwrap();
 		// the padded_offset + the previously allocated ptr (8 bytes prefix +
 		// 8 bytes content) + the prefix of 8 bytes which is prepended to the
 		// current pointer
 		assert_eq!(ptr2, to_pointer(padded_offset + 16 + HEADER_SIZE));

 		// when
-		heap.deallocate(&mut mem[..], ptr2).unwrap();
-		let ptr3 = heap.allocate(&mut mem[..], 9).unwrap();
+		heap.deallocate(&mut mem, ptr2).unwrap();
+		let ptr3 = heap.allocate(&mut mem, 9).unwrap();

 		// then
 		// should have re-allocated
@@ -712,22 +810,22 @@ mod tests {
 	#[test]
 	fn should_build_linked_list_of_free_areas_properly() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

-		let ptr1 = heap.allocate(&mut mem[..], 8).unwrap();
-		let ptr2 = heap.allocate(&mut mem[..], 8).unwrap();
-		let ptr3 = heap.allocate(&mut mem[..], 8).unwrap();
+		let ptr1 = heap.allocate(&mut mem, 8).unwrap();
+		let ptr2 = heap.allocate(&mut mem, 8).unwrap();
+		let ptr3 = heap.allocate(&mut mem, 8).unwrap();

 		// when
-		heap.deallocate(&mut mem[..], ptr1).unwrap();
-		heap.deallocate(&mut mem[..], ptr2).unwrap();
-		heap.deallocate(&mut mem[..], ptr3).unwrap();
+		heap.deallocate(&mut mem, ptr1).unwrap();
+		heap.deallocate(&mut mem, ptr2).unwrap();
+		heap.deallocate(&mut mem, ptr3).unwrap();

 		// then
 		assert_eq!(heap.free_lists.heads[0], Link::Ptr(u32::from(ptr3) - HEADER_SIZE));

-		let ptr4 = heap.allocate(&mut mem[..], 8).unwrap();
+		let ptr4 = heap.allocate(&mut mem, 8).unwrap();
 		assert_eq!(ptr4, ptr3);

 		assert_eq!(heap.free_lists.heads[0], Link::Ptr(u32::from(ptr2) - HEADER_SIZE));
@@ -736,29 +834,28 @@ mod tests {
 	#[test]
 	fn should_not_allocate_if_too_large() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
+		mem.set_max_wasm_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(13);

 		// when
-		let ptr = heap.allocate(&mut mem[..], PAGE_SIZE - 13);
+		let ptr = heap.allocate(&mut mem, PAGE_SIZE - 13);

 		// then
-		match ptr.unwrap_err() {
-			Error::AllocatorOutOfSpace => {},
-			e => panic!("Expected allocator out of space error, got: {:?}", e),
-		}
+		assert_eq!(Error::AllocatorOutOfSpace, ptr.unwrap_err());
 	}

 	#[test]
 	fn should_not_allocate_if_full() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
+		mem.set_max_wasm_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);
-		let ptr1 = heap.allocate(&mut mem[..], (PAGE_SIZE / 2) - HEADER_SIZE).unwrap();
+		let ptr1 = heap.allocate(&mut mem, (PAGE_SIZE / 2) - HEADER_SIZE).unwrap();
 		assert_eq!(ptr1, to_pointer(HEADER_SIZE));

 		// when
-		let ptr2 = heap.allocate(&mut mem[..], PAGE_SIZE / 2);
+		let ptr2 = heap.allocate(&mut mem, PAGE_SIZE / 2);

 		// then
 		// there is no room for another half page incl. its 8 byte prefix
@@ -771,11 +868,11 @@ mod tests {
 	#[test]
 	fn should_allocate_max_possible_allocation_size() {
 		// given
-		let mut mem = vec![0u8; (MAX_POSSIBLE_ALLOCATION + PAGE_SIZE) as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

 		// when
-		let ptr = heap.allocate(&mut mem[..], MAX_POSSIBLE_ALLOCATION).unwrap();
+		let ptr = heap.allocate(&mut mem, MAX_POSSIBLE_ALLOCATION).unwrap();

 		// then
 		assert_eq!(ptr, to_pointer(HEADER_SIZE));
@@ -784,60 +881,62 @@ mod tests {
 	#[test]
 	fn should_not_allocate_if_requested_size_too_large() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

 		// when
-		let ptr = heap.allocate(&mut mem[..], MAX_POSSIBLE_ALLOCATION + 1);
+		let ptr = heap.allocate(&mut mem, MAX_POSSIBLE_ALLOCATION + 1);

 		// then
-		match ptr.unwrap_err() {
-			Error::RequestedAllocationTooLarge => {},
-			e => panic!("Expected allocation size too large error, got: {:?}", e),
-		}
+		assert_eq!(Error::RequestedAllocationTooLarge, ptr.unwrap_err());
 	}

 	#[test]
 	fn should_return_error_when_bumper_greater_than_heap_size() {
 		// given
-		let mut mem = [0u8; 64];
+		let mut mem = MemoryInstance::with_pages(1);
+		mem.set_max_wasm_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

-		let ptr1 = heap.allocate(&mut mem[..], 32).unwrap();
-		assert_eq!(ptr1, to_pointer(HEADER_SIZE));
-		heap.deallocate(&mut mem[..], ptr1).expect("failed freeing ptr1");
-		assert_eq!(heap.stats.bytes_allocated, 0);
-		assert_eq!(heap.bumper, 40);
+		let mut ptrs = Vec::new();
+		for _ in 0..(PAGE_SIZE as usize / 40) {
+			ptrs.push(heap.allocate(&mut mem, 32).expect("Allocate 32 byte"));
+		}
+
+		assert_eq!(heap.stats.bytes_allocated, PAGE_SIZE - 16);
+		assert_eq!(heap.bumper, PAGE_SIZE - 16);
+
+		ptrs.into_iter()
+			.for_each(|ptr| heap.deallocate(&mut mem, ptr).expect("Deallocate 32 byte"));

-		let ptr2 = heap.allocate(&mut mem[..], 16).unwrap();
-		assert_eq!(ptr2, to_pointer(48));
-		heap.deallocate(&mut mem[..], ptr2).expect("failed freeing ptr2");
 		assert_eq!(heap.stats.bytes_allocated, 0);
-		assert_eq!(heap.bumper, 64);
+		assert_eq!(heap.stats.bytes_allocated_peak, PAGE_SIZE - 16);
+		assert_eq!(heap.bumper, PAGE_SIZE - 16);
+
+		// Allocate another 8 byte to use the full heap.
+		heap.allocate(&mut mem, 8).expect("Allocate 8 byte");

 		// when
 		// the `bumper` value is equal to `size` here and any
 		// further allocation which would increment the bumper must fail.
 		// we try to allocate 8 bytes here, which will increment the
-		// bumper since no 8 byte item has been allocated+freed before.
-		let ptr = heap.allocate(&mut mem[..], 8);
+		// bumper since no 8 byte item has been freed before.
+		assert_eq!(heap.bumper as u64, mem.size());
+		let ptr = heap.allocate(&mut mem, 8);

 		// then
-		match ptr.unwrap_err() {
-			Error::AllocatorOutOfSpace => {},
-			e => panic!("Expected allocator out of space error, got: {:?}", e),
-		}
+		assert_eq!(Error::AllocatorOutOfSpace, ptr.unwrap_err());
 	}

 	#[test]
 	fn should_include_prefixes_in_total_heap_size() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(1);

 		// when
 		// an item size of 16 must be used then
-		heap.allocate(&mut mem[..], 9).unwrap();
+		heap.allocate(&mut mem, 9).unwrap();

 		// then
 		assert_eq!(heap.stats.bytes_allocated, HEADER_SIZE + 16);
@@ -846,13 +945,13 @@ mod tests {
 	#[test]
 	fn should_calculate_total_heap_size_to_zero() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(13);

 		// when
-		let ptr = heap.allocate(&mut mem[..], 42).unwrap();
+		let ptr = heap.allocate(&mut mem, 42).unwrap();
 		assert_eq!(ptr, to_pointer(16 + HEADER_SIZE));
-		heap.deallocate(&mut mem[..], ptr).unwrap();
+		heap.deallocate(&mut mem, ptr).unwrap();

 		// then
 		assert_eq!(heap.stats.bytes_allocated, 0);
@@ -861,13 +960,13 @@ mod tests {
 	#[test]
 	fn should_calculate_total_size_of_zero() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(19);

 		// when
 		for _ in 1..10 {
-			let ptr = heap.allocate(&mut mem[..], 42).unwrap();
-			heap.deallocate(&mut mem[..], ptr).unwrap();
+			let ptr = heap.allocate(&mut mem, 42).unwrap();
+			heap.deallocate(&mut mem, ptr).unwrap();
 		}

 		// then
@@ -877,13 +976,13 @@ mod tests {
 	#[test]
 	fn should_read_and_write_u64_correctly() {
 		// given
-		let mut mem = [0u8; PAGE_SIZE as usize];
+		let mut mem = MemoryInstance::with_pages(1);

 		// when
-		Memory::write_le_u64(mem.as_mut(), 40, 4480113).unwrap();
+		mem.write_le_u64(40, 4480113).unwrap();

 		// then
-		let value = Memory::read_le_u64(mem.as_mut(), 40).unwrap();
+		let value = MemoryExt::read_le_u64(&mut mem, 40).unwrap();
 		assert_eq!(value, 4480113);
 	}

@@ -913,24 +1012,25 @@ mod tests {

 	#[test]
 	fn deallocate_needs_to_maintain_linked_list() {
-		let mut mem = [0u8; 8 * 2 * 4 + ALIGNMENT as usize];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

 		// Allocate and free some pointers
-		let ptrs = (0..4).map(|_| heap.allocate(&mut mem[..], 8).unwrap()).collect::<Vec<_>>();
-		ptrs.into_iter().for_each(|ptr| heap.deallocate(&mut mem[..], ptr).unwrap());
+		let ptrs = (0..4).map(|_| heap.allocate(&mut mem, 8).unwrap()).collect::<Vec<_>>();
+		ptrs.iter().rev().for_each(|ptr| heap.deallocate(&mut mem, *ptr).unwrap());

-		// Second time we should be able to allocate all of them again.
-		let _ = (0..4).map(|_| heap.allocate(&mut mem[..], 8).unwrap()).collect::<Vec<_>>();
+		// Second time we should be able to allocate all of them again and get the same pointers!
+		let new_ptrs = (0..4).map(|_| heap.allocate(&mut mem, 8).unwrap()).collect::<Vec<_>>();
+		assert_eq!(ptrs, new_ptrs);
 	}

 	#[test]
 	fn header_read_write() {
 		let roundtrip = |header: Header| {
-			let mut memory = [0u8; 32];
-			header.write_into(memory.as_mut(), 0).unwrap();
+			let mut memory = MemoryInstance::with_pages(1);
+			header.write_into(&mut memory, 0).unwrap();

-			let read_header = Header::read_from(memory.as_mut(), 0).unwrap();
+			let read_header = Header::read_from(&memory, 0).unwrap();
 			assert_eq!(header, read_header);
 		};

@@ -944,18 +1044,18 @@ mod tests {
 	#[test]
 	fn poison_oom() {
 		// given
-		// a heap of 32 bytes. Should be enough for two allocations.
-		let mut mem = [0u8; 32];
+		let mut mem = MemoryInstance::with_pages(1);
+		mem.set_max_wasm_pages(1);
+
 		let mut heap = FreeingBumpHeapAllocator::new(0);

 		// when
-		assert!(heap.allocate(mem.as_mut(), 8).is_ok());
-		let alloc_ptr = heap.allocate(mem.as_mut(), 8).unwrap();
-		assert!(heap.allocate(mem.as_mut(), 8).is_err());
+		let alloc_ptr = heap.allocate(&mut mem, PAGE_SIZE / 2).unwrap();
+		assert_eq!(Error::AllocatorOutOfSpace, heap.allocate(&mut mem, PAGE_SIZE).unwrap_err());

 		// then
 		assert!(heap.poisoned);
-		assert!(heap.deallocate(mem.as_mut(), alloc_ptr).is_err());
+		assert!(heap.deallocate(&mut mem, alloc_ptr).is_err());
 	}

 	#[test]
@@ -969,36 +1069,41 @@ mod tests {

 	#[test]
 	fn accepts_growing_memory() {
-		const ITEM_SIZE: u32 = 16;
-		const ITEM_ON_HEAP_SIZE: usize = 16 + HEADER_SIZE as usize;
-
-		let mut mem = vec![0u8; ITEM_ON_HEAP_SIZE * 2];
+		let mut mem = MemoryInstance::with_pages(1);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

-		let _ = heap.allocate(&mut mem[..], ITEM_SIZE).unwrap();
-		let _ = heap.allocate(&mut mem[..], ITEM_SIZE).unwrap();
+		heap.allocate(&mut mem, PAGE_SIZE / 2).unwrap();
+		heap.allocate(&mut mem, PAGE_SIZE / 2).unwrap();

-		mem.extend_from_slice(&[0u8; ITEM_ON_HEAP_SIZE]);
+		mem.grow(1).unwrap();

-		let _ = heap.allocate(&mut mem[..], ITEM_SIZE).unwrap();
+		heap.allocate(&mut mem, PAGE_SIZE / 2).unwrap();
 	}

 	#[test]
 	fn doesnt_accept_shrinking_memory() {
-		const ITEM_SIZE: u32 = 16;
-		const ITEM_ON_HEAP_SIZE: usize = 16 + HEADER_SIZE as usize;
-
-		let initial_size = ITEM_ON_HEAP_SIZE * 3;
-		let mut mem = vec![0u8; initial_size];
+		let mut mem = MemoryInstance::with_pages(2);
 		let mut heap = FreeingBumpHeapAllocator::new(0);

-		let _ = heap.allocate(&mut mem[..], ITEM_SIZE).unwrap();
+		heap.allocate(&mut mem, PAGE_SIZE / 2).unwrap();

-		mem.truncate(initial_size - 1);
+		mem.data.truncate(PAGE_SIZE as usize);

-		match heap.allocate(&mut mem[..], ITEM_SIZE).unwrap_err() {
+		match heap.allocate(&mut mem, PAGE_SIZE / 2).unwrap_err() {
 			Error::MemoryShrinked => (),
 			_ => panic!(),
 		}
 	}
+
+	#[test]
+	fn should_grow_memory_when_running_out_of_memory() {
+		let mut mem = MemoryInstance::with_pages(1);
+		let mut heap = FreeingBumpHeapAllocator::new(0);
+
+		assert_eq!(1, mem.pages());
+
+		heap.allocate(&mut mem, PAGE_SIZE * 2).unwrap();
+
+		assert_eq!(3, mem.pages());
+	}
 }
@@ -27,3 +27,35 @@ mod freeing_bump;

 pub use error::Error;
 pub use freeing_bump::{AllocationStats, FreeingBumpHeapAllocator};
+
+/// The size of one wasm page in bytes.
+///
+/// The wasm memory is divided into pages, meaning the minimum size of a memory is one page.
+const PAGE_SIZE: u32 = 65536;
+
+/// The maximum number of wasm pages that can be allocated.
+///
+/// 4GiB / [`PAGE_SIZE`].
+const MAX_WASM_PAGES: u32 = (4u64 * 1024 * 1024 * 1024 / PAGE_SIZE as u64) as u32;
+
+/// Grants access to the memory for the allocator.
+///
+/// Memory of wasm is allocated in pages. A page has a constant size of 64KiB. The maximum allowed
+/// memory size as defined in the wasm specification is 4GiB (65536 pages).
+pub trait Memory {
+	/// Run the given closure `run` and grant it write access to the raw memory.
+	fn with_access_mut<R>(&mut self, run: impl FnOnce(&mut [u8]) -> R) -> R;
+	/// Run the given closure `run` and grant it read access to the raw memory.
+	fn with_access<R>(&self, run: impl FnOnce(&[u8]) -> R) -> R;
+	/// Grow the memory by `additional` pages.
+	fn grow(&mut self, additional: u32) -> Result<(), ()>;
+	/// Returns the current number of pages this memory has allocated.
+	fn pages(&self) -> u32;
+	/// Returns the maximum number of pages this memory is allowed to allocate.
+	///
+	/// The returned number needs to be smaller or equal to `65536`. The returned number needs to be
+	/// bigger or equal to [`Self::pages`].
+	///
+	/// If `None` is returned, there is no maximum (besides the maximum defined in the wasm spec).
+	fn max_pages(&self) -> Option<u32>;
+}