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>

substrate/client/executor/Cargo.toml

@@ -17,7 +17,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 lru = "0.8.1"
 parking_lot = "0.12.1"
 tracing = "0.1.29"
-wasmi = "0.13"
+wasmi = "0.13.2"
 
 codec = { package = "parity-scale-codec", version = "3.2.2" }
 sc-executor-common = { version = "0.10.0-dev", path = "common" }
@@ -43,6 +43,7 @@ sp-state-machine = { version = "0.13.0", path = "../../primitives/state-machine"
 sp-runtime = { version = "7.0.0", path = "../../primitives/runtime" }
 sp-maybe-compressed-blob = { version = "4.1.0-dev", path = "../../primitives/maybe-compressed-blob" }
 sc-tracing = { version = "4.0.0-dev", path = "../tracing" }
+sp-tracing = { version = "6.0.0", path = "../../primitives/tracing" }
 tracing-subscriber = "0.2.19"
 paste = "1.0"
 regex = "1.6.0"

substrate/client/executor/benches/bench.rs

@@ -21,7 +21,7 @@ use codec::Encode;
 
 use sc_executor_common::{
 	runtime_blob::RuntimeBlob,
-	wasm_runtime::{WasmInstance, WasmModule},
+	wasm_runtime::{HeapAllocStrategy, WasmInstance, WasmModule},
 };
 use sc_executor_wasmtime::InstantiationStrategy;
 use sc_runtime_test::wasm_binary_unwrap as test_runtime;
@@ -51,13 +51,13 @@ fn initialize(
 ) -> Arc<dyn WasmModule> {
 	let blob = RuntimeBlob::uncompress_if_needed(runtime).unwrap();
 	let host_functions = sp_io::SubstrateHostFunctions::host_functions();
-	let heap_pages = 2048;
+	let extra_pages = 2048;
 	let allow_missing_func_imports = true;
 
 	match method {
 		Method::Interpreted => sc_executor_wasmi::create_runtime(
 			blob,
-			heap_pages,
+			HeapAllocStrategy::Static { extra_pages },
 			host_functions,
 			allow_missing_func_imports,
 		)
@@ -67,12 +67,11 @@ fn initialize(
 				allow_missing_func_imports,
 				cache_path: None,
 				semantics: sc_executor_wasmtime::Semantics {
-					extra_heap_pages: heap_pages,
+					heap_alloc_strategy: HeapAllocStrategy::Static { extra_pages },
 					instantiation_strategy,
 					deterministic_stack_limit: None,
 					canonicalize_nans: false,
 					parallel_compilation: true,
-					max_memory_size: None,
 				},
 			};
 

substrate/client/executor/common/Cargo.toml

@@ -16,7 +16,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 [dependencies]
 thiserror = "1.0.30"
 wasm-instrument = "0.3"
-wasmi = "0.13"
+wasmi = "0.13.2"
 sc-allocator = { version = "4.1.0-dev", path = "../../allocator" }
 sp-maybe-compressed-blob = { version = "4.1.0-dev", path = "../../../primitives/maybe-compressed-blob" }
 sp-wasm-interface = { version = "7.0.0", path = "../../../primitives/wasm-interface" }

substrate/client/executor/common/src/runtime_blob/runtime_blob.rs

@@ -16,7 +16,7 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
-use crate::error::WasmError;
+use crate::{error::WasmError, wasm_runtime::HeapAllocStrategy};
 use wasm_instrument::{
 	export_mutable_globals,
 	parity_wasm::elements::{
@@ -157,18 +157,13 @@ impl RuntimeBlob {
 		Ok(())
 	}
 
-	/// Increases the number of memory pages requested by the WASM blob by
-	/// the given amount of `extra_heap_pages`.
+	/// Modifies the blob's memory section according to the given `heap_alloc_strategy`.
 	///
 	/// Will return an error in case there is no memory section present,
 	/// or if the memory section is empty.
-	///
-	/// Only modifies the initial size of the memory; the maximum is unmodified
-	/// unless it's smaller than the initial size, in which case it will be increased
-	/// so that it's at least as big as the initial size.
-	pub fn add_extra_heap_pages_to_memory_section(
+	pub fn setup_memory_according_to_heap_alloc_strategy(
 		&mut self,
-		extra_heap_pages: u32,
+		heap_alloc_strategy: HeapAllocStrategy,
 	) -> Result<(), WasmError> {
 		let memory_section = self
 			.raw_module
@@ -179,8 +174,17 @@ impl RuntimeBlob {
 			return Err(WasmError::Other("memory section is empty".into()))
 		}
 		for memory_ty in memory_section.entries_mut() {
-			let min = memory_ty.limits().initial().saturating_add(extra_heap_pages);
-			let max = memory_ty.limits().maximum().map(|max| std::cmp::max(min, max));
+			let initial = memory_ty.limits().initial();
+			let (min, max) = match heap_alloc_strategy {
+				HeapAllocStrategy::Dynamic { maximum_pages } => {
+					// Ensure `initial <= maximum_pages`
+					(maximum_pages.map(|m| m.min(initial)).unwrap_or(initial), maximum_pages)
+				},
+				HeapAllocStrategy::Static { extra_pages } => {
+					let pages = initial.saturating_add(extra_pages);
+					(pages, Some(pages))
+				},
+			};
 			*memory_ty = MemoryType::new(min, max);
 		}
 		Ok(())

substrate/client/executor/common/src/util.rs

@@ -26,11 +26,7 @@ use std::ops::Range;
 /// Returns None if the end of the range would exceed some maximum offset.
 pub fn checked_range(offset: usize, len: usize, max: usize) -> Option<Range<usize>> {
 	let end = offset.checked_add(len)?;
-	if end <= max {
-		Some(offset..end)
-	} else {
-		None
-	}
+	(end <= max).then(|| offset..end)
 }
 
 /// Provides safe memory access interface using an external buffer

substrate/client/executor/common/src/wasm_runtime.rs

@@ -119,3 +119,29 @@ pub trait WasmInstance: Send {
 		None
 	}
 }
+
+/// Defines the heap pages allocation strategy the wasm runtime should use.
+///
+/// A heap page is defined as 64KiB of memory.
+#[derive(Debug, Copy, Clone, PartialEq, Hash, Eq)]
+pub enum HeapAllocStrategy {
+	/// Allocate a static number of heap pages.
+	///
+	/// The total number of allocated heap pages is the initial number of heap pages requested by
+	/// the wasm file plus the `extra_pages`.
+	Static {
+		/// The number of pages that will be added on top of the initial heap pages requested by
+		/// the wasm file.
+		extra_pages: u32,
+	},
+	/// Allocate the initial heap pages as requested by the wasm file and then allow it to grow
+	/// dynamically.
+	Dynamic {
+		/// The absolute maximum size of the linear memory (in pages).
+		///
+		/// When `Some(_)` the linear memory will be allowed to grow up to this limit.
+		/// When `None` the linear memory will be allowed to grow up to the maximum limit supported
+		/// by WASM (4GB).
+		maximum_pages: Option<u32>,
+	},
+}

substrate/client/executor/runtime-test/src/lib.rs

@@ -328,6 +328,15 @@ sp_core::wasm_export_functions! {
 		assert_eq!(value, -66);
 	}
 
+	fn allocate_two_gigabyte() -> u32 {
+		let mut data = Vec::new();
+		for _ in 0..205 {
+			data.push(Vec::<u8>::with_capacity(10 * 1024 * 1024));
+		}
+
+		data.iter().map(|d| d.capacity() as u32).sum()
+	}
+
 	fn test_abort_on_panic() {
 		sp_io::panic_handler::abort_on_panic("test_abort_on_panic called");
 	}

substrate/client/executor/src/integration_tests/linux.rs

@@ -21,25 +21,60 @@
 use super::mk_test_runtime;
 use crate::WasmExecutionMethod;
 use codec::Encode as _;
+use sc_executor_common::wasm_runtime::HeapAllocStrategy;
 
 mod smaps;
 
 use self::smaps::Smaps;
 
+#[test]
+fn memory_consumption_interpreted() {
+	let _ = sp_tracing::try_init_simple();
+
+	if std::env::var("RUN_TEST").is_ok() {
+		memory_consumption(WasmExecutionMethod::Interpreted);
+	} else {
+		// We need to run the test in isolation, to not getting interfered by the other tests.
+		let executable = std::env::current_exe().unwrap();
+		let output = std::process::Command::new(executable)
+			.env("RUN_TEST", "1")
+			.args(&["--nocapture", "memory_consumption_interpreted"])
+			.output()
+			.unwrap();
+
+		assert!(output.status.success());
+	}
+}
+
 #[test]
 fn memory_consumption_compiled() {
+	let _ = sp_tracing::try_init_simple();
+
+	if std::env::var("RUN_TEST").is_ok() {
+		memory_consumption(WasmExecutionMethod::Compiled {
+			instantiation_strategy:
+				sc_executor_wasmtime::InstantiationStrategy::LegacyInstanceReuse,
+		});
+	} else {
+		// We need to run the test in isolation, to not getting interfered by the other tests.
+		let executable = std::env::current_exe().unwrap();
+		let status = std::process::Command::new(executable)
+			.env("RUN_TEST", "1")
+			.args(&["--nocapture", "memory_consumption_compiled"])
+			.status()
+			.unwrap();
+
+		assert!(status.success());
+	}
+}
+
+fn memory_consumption(wasm_method: WasmExecutionMethod) {
 	// This aims to see if linear memory stays backed by the physical memory after a runtime call.
 	//
 	// For that we make a series of runtime calls, probing the RSS for the VMA matching the linear
 	// memory. After the call we expect RSS to be equal to 0.
 
-	let runtime = mk_test_runtime(
-		WasmExecutionMethod::Compiled {
-			instantiation_strategy:
-				sc_executor_wasmtime::InstantiationStrategy::LegacyInstanceReuse,
-		},
-		1024,
-	);
+	let runtime = mk_test_runtime(wasm_method, HeapAllocStrategy::Static { extra_pages: 1024 });
 
 	let mut instance = runtime.new_instance().unwrap();
 	let heap_base = instance

substrate/client/executor/src/integration_tests/mod.rs

@@ -24,7 +24,7 @@ use codec::{Decode, Encode};
 use sc_executor_common::{
 	error::{Error, WasmError},
 	runtime_blob::RuntimeBlob,
-	wasm_runtime::WasmModule,
+	wasm_runtime::{HeapAllocStrategy, WasmModule},
 };
 use sc_runtime_test::wasm_binary_unwrap;
 use sp_core::{
@@ -52,11 +52,13 @@ macro_rules! test_wasm_execution {
 		paste::item! {
 			#[test]
 			fn [<$method_name _interpreted>]() {
+				let _ = sp_tracing::try_init_simple();
 				$method_name(WasmExecutionMethod::Interpreted);
 			}
 
 			#[test]
 			fn [<$method_name _compiled_recreate_instance_cow>]() {
+				let _ = sp_tracing::try_init_simple();
 				$method_name(WasmExecutionMethod::Compiled {
 					instantiation_strategy: sc_executor_wasmtime::InstantiationStrategy::RecreateInstanceCopyOnWrite
 				});
@@ -64,6 +66,7 @@ macro_rules! test_wasm_execution {
 
 			#[test]
 			fn [<$method_name _compiled_recreate_instance_vanilla>]() {
+				let _ = sp_tracing::try_init_simple();
 				$method_name(WasmExecutionMethod::Compiled {
 					instantiation_strategy: sc_executor_wasmtime::InstantiationStrategy::RecreateInstance
 				});
@@ -71,6 +74,7 @@ macro_rules! test_wasm_execution {
 
 			#[test]
 			fn [<$method_name _compiled_pooling_cow>]() {
+				let _ = sp_tracing::try_init_simple();
 				$method_name(WasmExecutionMethod::Compiled {
 					instantiation_strategy: sc_executor_wasmtime::InstantiationStrategy::PoolingCopyOnWrite
 				});
@@ -78,6 +82,7 @@ macro_rules! test_wasm_execution {
 
 			#[test]
 			fn [<$method_name _compiled_pooling_vanilla>]() {
+				let _ = sp_tracing::try_init_simple();
 				$method_name(WasmExecutionMethod::Compiled {
 					instantiation_strategy: sc_executor_wasmtime::InstantiationStrategy::Pooling
 				});
@@ -85,6 +90,7 @@ macro_rules! test_wasm_execution {
 
 			#[test]
 			fn [<$method_name _compiled_legacy_instance_reuse>]() {
+				let _ = sp_tracing::try_init_simple();
 				$method_name(WasmExecutionMethod::Compiled {
 					instantiation_strategy: sc_executor_wasmtime::InstantiationStrategy::LegacyInstanceReuse
 				});
@@ -474,7 +480,10 @@ fn should_trap_when_heap_exhausted(wasm_method: WasmExecutionMethod) {
 	}
 }
 
-fn mk_test_runtime(wasm_method: WasmExecutionMethod, pages: u64) -> Arc<dyn WasmModule> {
+fn mk_test_runtime(
+	wasm_method: WasmExecutionMethod,
+	pages: HeapAllocStrategy,
+) -> Arc<dyn WasmModule> {
 	let blob = RuntimeBlob::uncompress_if_needed(wasm_binary_unwrap())
 		.expect("failed to create a runtime blob out of test runtime");
 
@@ -490,7 +499,8 @@ fn mk_test_runtime(wasm_method: WasmExecutionMethod, pages: u64) -> Arc<dyn Wasm
 
 test_wasm_execution!(returns_mutable_static);
 fn returns_mutable_static(wasm_method: WasmExecutionMethod) {
-	let runtime = mk_test_runtime(wasm_method, 1024);
+	let runtime =
+		mk_test_runtime(wasm_method, HeapAllocStrategy::Dynamic { maximum_pages: Some(1024) });
 
 	let mut instance = runtime.new_instance().unwrap();
 	let res = instance.call_export("returns_mutable_static", &[0]).unwrap();
@@ -505,7 +515,8 @@ fn returns_mutable_static(wasm_method: WasmExecutionMethod) {
 
 test_wasm_execution!(returns_mutable_static_bss);
 fn returns_mutable_static_bss(wasm_method: WasmExecutionMethod) {
-	let runtime = mk_test_runtime(wasm_method, 1024);
+	let runtime =
+		mk_test_runtime(wasm_method, HeapAllocStrategy::Dynamic { maximum_pages: Some(1024) });
 
 	let mut instance = runtime.new_instance().unwrap();
 	let res = instance.call_export("returns_mutable_static_bss", &[0]).unwrap();
@@ -530,9 +541,12 @@ fn restoration_of_globals(wasm_method: WasmExecutionMethod) {
 	//
 	// The fixture performs 2 allocations of 768KB and this theoretically gives 1536KB, however, due
 	// to our allocator algorithm there are inefficiencies.
-	const REQUIRED_MEMORY_PAGES: u64 = 32;
+	const REQUIRED_MEMORY_PAGES: u32 = 32;
 
-	let runtime = mk_test_runtime(wasm_method, REQUIRED_MEMORY_PAGES);
+	let runtime = mk_test_runtime(
+		wasm_method,
+		HeapAllocStrategy::Static { extra_pages: REQUIRED_MEMORY_PAGES },
+	);
 	let mut instance = runtime.new_instance().unwrap();
 
 	// On the first invocation we allocate approx. 768KB (75%) of stack and then trap.
@@ -546,7 +560,7 @@ fn restoration_of_globals(wasm_method: WasmExecutionMethod) {
 
 test_wasm_execution!(interpreted_only heap_is_reset_between_calls);
 fn heap_is_reset_between_calls(wasm_method: WasmExecutionMethod) {
-	let runtime = mk_test_runtime(wasm_method, 1024);
+	let runtime = mk_test_runtime(wasm_method, HeapAllocStrategy::Static { extra_pages: 1024 });
 	let mut instance = runtime.new_instance().unwrap();
 
 	let heap_base = instance
@@ -651,6 +665,15 @@ fn wasm_tracing_should_work(wasm_method: WasmExecutionMethod) {
 	assert_eq!(len, 2);
 }
 
+test_wasm_execution!(allocate_two_gigabyte);
+fn allocate_two_gigabyte(wasm_method: WasmExecutionMethod) {
+	let runtime = mk_test_runtime(wasm_method, HeapAllocStrategy::Dynamic { maximum_pages: None });
+
+	let mut instance = runtime.new_instance().unwrap();
+	let res = instance.call_export("allocate_two_gigabyte", &[0]).unwrap();
+	assert_eq!(10 * 1024 * 1024 * 205, u32::decode(&mut &res[..]).unwrap());
+}
+
 test_wasm_execution!(memory_is_cleared_between_invocations);
 fn memory_is_cleared_between_invocations(wasm_method: WasmExecutionMethod) {
 	// This is based on the code generated by compiling a runtime *without*
@@ -713,7 +736,7 @@ fn memory_is_cleared_between_invocations(wasm_method: WasmExecutionMethod) {
 
 	let runtime = crate::wasm_runtime::create_wasm_runtime_with_code::<HostFunctions>(
 		wasm_method,
-		1024,
+		HeapAllocStrategy::Dynamic { maximum_pages: Some(1024) },
 		RuntimeBlob::uncompress_if_needed(&binary[..]).unwrap(),
 		true,
 		None,

substrate/client/executor/src/native_executor.rs

@@ -32,14 +32,15 @@ use std::{
 use codec::Encode;
 use sc_executor_common::{
 	runtime_blob::RuntimeBlob,
-	wasm_runtime::{AllocationStats, WasmInstance, WasmModule},
+	wasm_runtime::{AllocationStats, HeapAllocStrategy, WasmInstance, WasmModule},
 };
-use sp_core::traits::{CodeExecutor, Externalities, RuntimeCode};
+use sp_core::traits::{CallContext, CodeExecutor, Externalities, RuntimeCode};
 use sp_version::{GetNativeVersion, NativeVersion, RuntimeVersion};
 use sp_wasm_interface::{ExtendedHostFunctions, HostFunctions};
 
-/// Default num of pages for the heap
-const DEFAULT_HEAP_PAGES: u64 = 2048;
+/// Default heap allocation strategy.
+const DEFAULT_HEAP_ALLOC_STRATEGY: HeapAllocStrategy =
+	HeapAllocStrategy::Static { extra_pages: 2048 };
 
 /// Set up the externalities and safe calling environment to execute runtime calls.
 ///
@@ -79,13 +80,136 @@ pub trait NativeExecutionDispatch: Send + Sync {
 	fn native_version() -> NativeVersion;
 }
 
+fn unwrap_heap_pages(pages: Option<HeapAllocStrategy>) -> HeapAllocStrategy {
+	pages.unwrap_or_else(|| DEFAULT_HEAP_ALLOC_STRATEGY)
+}
+
+/// Builder for creating a [`WasmExecutor`] instance.
+pub struct WasmExecutorBuilder<H> {
+	_phantom: PhantomData<H>,
+	method: WasmExecutionMethod,
+	onchain_heap_alloc_strategy: Option<HeapAllocStrategy>,
+	offchain_heap_alloc_strategy: Option<HeapAllocStrategy>,
+	max_runtime_instances: usize,
+	cache_path: Option<PathBuf>,
+	allow_missing_host_functions: bool,
+	runtime_cache_size: u8,
+}
+
+impl<H> WasmExecutorBuilder<H> {
+	/// Create a new instance of `Self`
+	///
+	/// - `method`: The wasm execution method that should be used by the executor.
+	pub fn new(method: WasmExecutionMethod) -> Self {
+		Self {
+			_phantom: PhantomData,
+			method,
+			onchain_heap_alloc_strategy: None,
+			offchain_heap_alloc_strategy: None,
+			max_runtime_instances: 2,
+			runtime_cache_size: 4,
+			allow_missing_host_functions: false,
+			cache_path: None,
+		}
+	}
+
+	/// Create the wasm executor with the given number of `heap_alloc_strategy` for onchain runtime
+	/// calls.
+	pub fn with_onchain_heap_alloc_strategy(
+		mut self,
+		heap_alloc_strategy: HeapAllocStrategy,
+	) -> Self {
+		self.onchain_heap_alloc_strategy = Some(heap_alloc_strategy);
+		self
+	}
+
+	/// Create the wasm executor with the given number of `heap_alloc_strategy` for offchain runtime
+	/// calls.
+	pub fn with_offchain_heap_alloc_strategy(
+		mut self,
+		heap_alloc_strategy: HeapAllocStrategy,
+	) -> Self {
+		self.offchain_heap_alloc_strategy = Some(heap_alloc_strategy);
+		self
+	}
+
+	/// Create the wasm executor with the given maximum number of `instances`.
+	///
+	/// The number of `instances` defines how many different instances of a runtime the cache is
+	/// storing.
+	///
+	/// By default the maximum number of `instances` is `2`.
+	pub fn with_max_runtime_instances(mut self, instances: usize) -> Self {
+		self.max_runtime_instances = instances;
+		self
+	}
+
+	/// Create the wasm executor with the given `cache_path`.
+	///
+	/// The `cache_path` is A path to a directory where the executor can place its files for
+	/// purposes of caching. This may be important in cases when there are many different modules
+	/// with the compiled execution method is used.
+	///
+	/// By default there is no `cache_path` given.
+	pub fn with_cache_path(mut self, cache_path: impl Into<PathBuf>) -> Self {
+		self.cache_path = Some(cache_path.into());
+		self
+	}
+
+	/// Create the wasm executor and allow/forbid missing host functions.
+	///
+	/// If missing host functions are forbidden, the instantiation of a wasm blob will fail
+	/// for imported host functions that the executor is not aware of. If they are allowed,
+	/// a stub is generated that will return an error when being called while executing the wasm.
+	///
+	/// By default missing host functions are forbidden.
+	pub fn with_allow_missing_host_functions(mut self, allow: bool) -> Self {
+		self.allow_missing_host_functions = allow;
+		self
+	}
+
+	/// Create the wasm executor with the given `runtime_cache_size`.
+	///
+	/// Defines the number of different runtimes/instantiated wasm blobs the cache stores.
+	/// Runtimes/wasm blobs are differentiated based on the hash and the number of heap pages.
+	///
+	/// By default this value is set to `4`.
+	pub fn with_runtime_cache_size(mut self, runtime_cache_size: u8) -> Self {
+		self.runtime_cache_size = runtime_cache_size;
+		self
+	}
+
+	/// Build the configured [`WasmExecutor`].
+	pub fn build(self) -> WasmExecutor<H> {
+		WasmExecutor {
+			method: self.method,
+			default_offchain_heap_alloc_strategy: unwrap_heap_pages(
+				self.offchain_heap_alloc_strategy,
+			),
+			default_onchain_heap_alloc_strategy: unwrap_heap_pages(
+				self.onchain_heap_alloc_strategy,
+			),
+			cache: Arc::new(RuntimeCache::new(
+				self.max_runtime_instances,
+				self.cache_path.clone(),
+				self.runtime_cache_size,
+			)),
+			cache_path: self.cache_path,
+			allow_missing_host_functions: self.allow_missing_host_functions,
+			phantom: PhantomData,
+		}
+	}
+}
+
 /// An abstraction over Wasm code executor. Supports selecting execution backend and
 /// manages runtime cache.
 pub struct WasmExecutor<H> {
 	/// Method used to execute fallback Wasm code.
 	method: WasmExecutionMethod,
-	/// The number of 64KB pages to allocate for Wasm execution.
-	default_heap_pages: u64,
+	/// The heap allocation strategy for onchain Wasm calls.
+	default_onchain_heap_alloc_strategy: HeapAllocStrategy,
+	/// The heap allocation strategy for offchain Wasm calls.
+	default_offchain_heap_alloc_strategy: HeapAllocStrategy,
 	/// WASM runtime cache.
 	cache: Arc<RuntimeCache>,
 	/// The path to a directory which the executor can leverage for a file cache, e.g. put there
@@ -100,7 +224,8 @@ impl<H> Clone for WasmExecutor<H> {
 	fn clone(&self) -> Self {
 		Self {
 			method: self.method,
-			default_heap_pages: self.default_heap_pages,
+			default_onchain_heap_alloc_strategy: self.default_onchain_heap_alloc_strategy,
+			default_offchain_heap_alloc_strategy: self.default_offchain_heap_alloc_strategy,
 			cache: self.cache.clone(),
 			cache_path: self.cache_path.clone(),
 			allow_missing_host_functions: self.allow_missing_host_functions,
@@ -119,8 +244,10 @@ where
 	///
 	/// `method` - Method used to execute Wasm code.
 	///
-	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
-	///   Defaults to `DEFAULT_HEAP_PAGES` if `None` is provided.
+	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution. Internally this
+	/// will be mapped as [`HeapAllocStrategy::Static`] where `default_heap_pages` represent the
+	/// static number of heap pages to allocate. Defaults to `DEFAULT_HEAP_ALLOC_STRATEGY` if `None`
+	/// is provided.
 	///
 	/// `max_runtime_instances` - The number of runtime instances to keep in memory ready for reuse.
 	///
@@ -138,7 +265,12 @@ where
 	) -> Self {
 		WasmExecutor {
 			method,
-			default_heap_pages: default_heap_pages.unwrap_or(DEFAULT_HEAP_PAGES),
+			default_onchain_heap_alloc_strategy: unwrap_heap_pages(
+				default_heap_pages.map(|h| HeapAllocStrategy::Static { extra_pages: h as _ }),
+			),
+			default_offchain_heap_alloc_strategy: unwrap_heap_pages(
+				default_heap_pages.map(|h| HeapAllocStrategy::Static { extra_pages: h as _ }),
+			),
 			cache: Arc::new(RuntimeCache::new(
 				max_runtime_instances,
 				cache_path.clone(),
@@ -150,6 +282,11 @@ where
 		}
 	}
 
+	/// Instantiate a builder for creating an instance of `Self`.
+	pub fn builder(method: WasmExecutionMethod) -> WasmExecutorBuilder<H> {
+		WasmExecutorBuilder::new(method)
+	}
+
 	/// Ignore missing function imports if set true.
 	pub fn allow_missing_host_functions(&mut self, allow_missing_host_functions: bool) {
 		self.allow_missing_host_functions = allow_missing_host_functions
@@ -172,6 +309,7 @@ where
 		&self,
 		runtime_code: &RuntimeCode,
 		ext: &mut dyn Externalities,
+		heap_alloc_strategy: HeapAllocStrategy,
 		f: F,
 	) -> Result<R>
 	where
@@ -186,7 +324,7 @@ where
 			runtime_code,
 			ext,
 			self.method,
-			self.default_heap_pages,
+			heap_alloc_strategy,
 			self.allow_missing_host_functions,
 			|module, instance, version, ext| {
 				let module = AssertUnwindSafe(module);
@@ -259,7 +397,7 @@ where
 	) -> std::result::Result<Vec<u8>, Error> {
 		let module = crate::wasm_runtime::create_wasm_runtime_with_code::<H>(
 			self.method,
-			self.default_heap_pages,
+			self.default_onchain_heap_alloc_strategy,
 			runtime_blob,
 			allow_missing_host_functions,
 			self.cache_path.as_deref(),
@@ -334,6 +472,7 @@ where
 		method: &str,
 		data: &[u8],
 		_use_native: bool,
+		context: CallContext,
 	) -> (Result<Vec<u8>>, bool) {
 		tracing::trace!(
 			target: "executor",
@@ -341,10 +480,25 @@ where
 			"Executing function",
 		);
 
-		let result =
-			self.with_instance(runtime_code, ext, |_, mut instance, _onchain_version, mut ext| {
+		let on_chain_heap_alloc_strategy = runtime_code
+			.heap_pages
+			.map(|h| HeapAllocStrategy::Static { extra_pages: h as _ })
+			.unwrap_or_else(|| self.default_onchain_heap_alloc_strategy);
+
+		let heap_alloc_strategy = match context {
+			CallContext::Offchain => self.default_offchain_heap_alloc_strategy,
+			CallContext::Onchain => on_chain_heap_alloc_strategy,
+		};
+
+		let result = self.with_instance(
+			runtime_code,
+			ext,
+			heap_alloc_strategy,
+			|_, mut instance, _onchain_version, mut ext| {
 				with_externalities_safe(&mut **ext, move || instance.call_export(method, data))
-			});
+			},
+		);
+
 		(result, false)
 	}
 }
@@ -358,20 +512,25 @@ where
 		ext: &mut dyn Externalities,
 		runtime_code: &RuntimeCode,
 	) -> Result<RuntimeVersion> {
-		self.with_instance(runtime_code, ext, |_module, _instance, version, _ext| {
-			Ok(version.cloned().ok_or_else(|| Error::ApiError("Unknown version".into())))
-		})
+		let on_chain_heap_pages = runtime_code
+			.heap_pages
+			.map(|h| HeapAllocStrategy::Static { extra_pages: h as _ })
+			.unwrap_or_else(|| self.default_onchain_heap_alloc_strategy);
+
+		self.with_instance(
+			runtime_code,
+			ext,
+			on_chain_heap_pages,
+			|_module, _instance, version, _ext| {
+				Ok(version.cloned().ok_or_else(|| Error::ApiError("Unknown version".into())))
+			},
+		)
 	}
 }
 
 /// A generic `CodeExecutor` implementation that uses a delegate to determine wasm code equivalence
 /// and dispatch to native code when possible, falling back on `WasmExecutor` when not.
-pub struct NativeElseWasmExecutor<D>
-where
-	D: NativeExecutionDispatch,
-{
-	/// Dummy field to avoid the compiler complaining about us not using `D`.
-	_dummy: PhantomData<D>,
+pub struct NativeElseWasmExecutor<D: NativeExecutionDispatch> {
 	/// Native runtime version info.
 	native_version: NativeVersion,
 	/// Fallback wasm executor.
@@ -386,8 +545,10 @@ impl<D: NativeExecutionDispatch> NativeElseWasmExecutor<D> {
 	///
 	/// `fallback_method` - Method used to execute fallback Wasm code.
 	///
-	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
-	/// 	Defaults to `DEFAULT_HEAP_PAGES` if `None` is provided.
+	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution. Internally this
+	/// will be mapped as [`HeapAllocStrategy::Static`] where `default_heap_pages` represent the
+	/// static number of heap pages to allocate. Defaults to `DEFAULT_HEAP_ALLOC_STRATEGY` if `None`
+	/// is provided.
 	///
 	/// `max_runtime_instances` - The number of runtime instances to keep in memory ready for reuse.
 	///
@@ -406,11 +567,16 @@ impl<D: NativeExecutionDispatch> NativeElseWasmExecutor<D> {
 			runtime_cache_size,
 		);
 
-		NativeElseWasmExecutor {
-			_dummy: Default::default(),
-			native_version: D::native_version(),
-			wasm,
-		}
+		NativeElseWasmExecutor { native_version: D::native_version(), wasm }
+	}
+
+	/// Create a new instance using the given [`WasmExecutor`].
+	pub fn new_with_wasm_executor(
+		executor: WasmExecutor<
+			ExtendedHostFunctions<sp_io::SubstrateHostFunctions, D::ExtendHostFunctions>,
+		>,
+	) -> Self {
+		Self { native_version: D::native_version(), wasm: executor }
 	}
 
 	/// Ignore missing function imports if set true.
@@ -425,9 +591,7 @@ impl<D: NativeExecutionDispatch> RuntimeVersionOf for NativeElseWasmExecutor<D>
 		ext: &mut dyn Externalities,
 		runtime_code: &RuntimeCode,
 	) -> Result<RuntimeVersion> {
-		self.wasm.with_instance(runtime_code, ext, |_module, _instance, version, _ext| {
-			Ok(version.cloned().ok_or_else(|| Error::ApiError("Unknown version".into())))
-		})
+		self.wasm.runtime_version(ext, runtime_code)
 	}
 }
 
@@ -447,6 +611,7 @@ impl<D: NativeExecutionDispatch + 'static> CodeExecutor for NativeElseWasmExecut
 		method: &str,
 		data: &[u8],
 		use_native: bool,
+		context: CallContext,
 	) -> (Result<Vec<u8>>, bool) {
 		tracing::trace!(
 			target: "executor",
@@ -454,10 +619,21 @@ impl<D: NativeExecutionDispatch + 'static> CodeExecutor for NativeElseWasmExecut
 			"Executing function",
 		);
 
+		let on_chain_heap_alloc_strategy = runtime_code
+			.heap_pages
+			.map(|h| HeapAllocStrategy::Static { extra_pages: h as _ })
+			.unwrap_or_else(|| self.wasm.default_onchain_heap_alloc_strategy);
+
+		let heap_alloc_strategy = match context {
+			CallContext::Offchain => self.wasm.default_offchain_heap_alloc_strategy,
+			CallContext::Onchain => on_chain_heap_alloc_strategy,
+		};
+
 		let mut used_native = false;
 		let result = self.wasm.with_instance(
 			runtime_code,
 			ext,
+			heap_alloc_strategy,
 			|_, mut instance, onchain_version, mut ext| {
 				let onchain_version =
 					onchain_version.ok_or_else(|| Error::ApiError("Unknown version".into()))?;
@@ -496,11 +672,7 @@ impl<D: NativeExecutionDispatch + 'static> CodeExecutor for NativeElseWasmExecut
 
 impl<D: NativeExecutionDispatch> Clone for NativeElseWasmExecutor<D> {
 	fn clone(&self) -> Self {
-		NativeElseWasmExecutor {
-			_dummy: Default::default(),
-			native_version: D::native_version(),
-			wasm: self.wasm.clone(),
-		}
+		NativeElseWasmExecutor { native_version: D::native_version(), wasm: self.wasm.clone() }
 	}
 }
 

substrate/client/executor/src/wasm_runtime.rs

@@ -27,7 +27,7 @@ use lru::LruCache;
 use parking_lot::Mutex;
 use sc_executor_common::{
 	runtime_blob::RuntimeBlob,
-	wasm_runtime::{WasmInstance, WasmModule},
+	wasm_runtime::{HeapAllocStrategy, WasmInstance, WasmModule},
 };
 use sp_core::traits::{Externalities, FetchRuntimeCode, RuntimeCode};
 use sp_version::RuntimeVersion;
@@ -64,8 +64,8 @@ struct VersionedRuntimeId {
 	code_hash: Vec<u8>,
 	/// Wasm runtime type.
 	wasm_method: WasmExecutionMethod,
-	/// The number of WebAssembly heap pages this instance was created with.
-	heap_pages: u64,
+	/// The heap allocation strategy this runtime was created with.
+	heap_alloc_strategy: HeapAllocStrategy,
 }
 
 /// A Wasm runtime object along with its cached runtime version.
@@ -197,10 +197,12 @@ impl RuntimeCache {
 	///
 	/// `runtime_code` - The runtime wasm code used setup the runtime.
 	///
-	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
+	/// `ext` - The externalities to access the state.
 	///
 	/// `wasm_method` - Type of WASM backend to use.
 	///
+	/// `heap_alloc_strategy` - The heap allocation strategy to use.
+	///
 	/// `allow_missing_func_imports` - Ignore missing function imports.
 	///
 	/// `f` - Function to execute.
@@ -219,7 +221,7 @@ impl RuntimeCache {
 		runtime_code: &'c RuntimeCode<'c>,
 		ext: &mut dyn Externalities,
 		wasm_method: WasmExecutionMethod,
-		default_heap_pages: u64,
+		heap_alloc_strategy: HeapAllocStrategy,
 		allow_missing_func_imports: bool,
 		f: F,
 	) -> Result<Result<R, Error>, Error>
@@ -233,10 +235,9 @@ impl RuntimeCache {
 		) -> Result<R, Error>,
 	{
 		let code_hash = &runtime_code.hash;
-		let heap_pages = runtime_code.heap_pages.unwrap_or(default_heap_pages);
 
 		let versioned_runtime_id =
-			VersionedRuntimeId { code_hash: code_hash.clone(), heap_pages, wasm_method };
+			VersionedRuntimeId { code_hash: code_hash.clone(), heap_alloc_strategy, wasm_method };
 
 		let mut runtimes = self.runtimes.lock(); // this must be released prior to calling f
 		let versioned_runtime = if let Some(versioned_runtime) = runtimes.get(&versioned_runtime_id)
@@ -251,7 +252,7 @@ impl RuntimeCache {
 				&code,
 				ext,
 				wasm_method,
-				heap_pages,
+				heap_alloc_strategy,
 				allow_missing_func_imports,
 				self.max_runtime_instances,
 				self.cache_path.as_deref(),
@@ -289,7 +290,7 @@ impl RuntimeCache {
 /// Create a wasm runtime with the given `code`.
 pub fn create_wasm_runtime_with_code<H>(
 	wasm_method: WasmExecutionMethod,
-	heap_pages: u64,
+	heap_alloc_strategy: HeapAllocStrategy,
 	blob: RuntimeBlob,
 	allow_missing_func_imports: bool,
 	cache_path: Option<&Path>,
@@ -307,7 +308,7 @@ where
 
 			sc_executor_wasmi::create_runtime(
 				blob,
-				heap_pages,
+				heap_alloc_strategy,
 				H::host_functions(),
 				allow_missing_func_imports,
 			)
@@ -320,12 +321,11 @@ where
 					allow_missing_func_imports,
 					cache_path: cache_path.map(ToOwned::to_owned),
 					semantics: sc_executor_wasmtime::Semantics {
-						extra_heap_pages: heap_pages,
+						heap_alloc_strategy,
 						instantiation_strategy,
 						deterministic_stack_limit: None,
 						canonicalize_nans: false,
 						parallel_compilation: true,
-						max_memory_size: None,
 					},
 				},
 			)
@@ -393,7 +393,7 @@ fn create_versioned_wasm_runtime<H>(
 	code: &[u8],
 	ext: &mut dyn Externalities,
 	wasm_method: WasmExecutionMethod,
-	heap_pages: u64,
+	heap_alloc_strategy: HeapAllocStrategy,
 	allow_missing_func_imports: bool,
 	max_instances: usize,
 	cache_path: Option<&Path>,
@@ -408,11 +408,11 @@ where
 	// Use the runtime blob to scan if there is any metadata embedded into the wasm binary
 	// pertaining to runtime version. We do it before consuming the runtime blob for creating the
 	// runtime.
-	let mut version: Option<_> = read_embedded_version(&blob)?;
+	let mut version = read_embedded_version(&blob)?;
 
 	let runtime = create_wasm_runtime_with_code::<H>(
 		wasm_method,
-		heap_pages,
+		heap_alloc_strategy,
 		blob,
 		allow_missing_func_imports,
 		cache_path,

substrate/client/executor/wasmi/Cargo.toml

@@ -15,7 +15,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
 log = "0.4.17"
-wasmi = "0.13"
+wasmi = { version = "0.13.2", features = [ "virtual_memory" ] }
 sc-allocator = { version = "4.1.0-dev", path = "../../allocator" }
 sc-executor-common = { version = "0.10.0-dev", path = "../common" }
 sp-runtime-interface = { version = "7.0.0", path = "../../../primitives/runtime-interface" }

substrate/client/executor/wasmi/src/lib.rs

@@ -20,23 +20,58 @@
 
 use std::{cell::RefCell, str, sync::Arc};
 
-use log::{debug, error, trace};
+use log::{error, trace};
 use wasmi::{
 	memory_units::Pages,
-	FuncInstance, ImportsBuilder, MemoryInstance, MemoryRef, Module, ModuleInstance, ModuleRef,
+	FuncInstance, ImportsBuilder, MemoryRef, Module, ModuleInstance, ModuleRef,
 	RuntimeValue::{self, I32, I64},
 	TableRef,
 };
 
-use sc_allocator::AllocationStats;
+use sc_allocator::{AllocationStats, FreeingBumpHeapAllocator};
 use sc_executor_common::{
 	error::{Error, MessageWithBacktrace, WasmError},
 	runtime_blob::{DataSegmentsSnapshot, RuntimeBlob},
-	wasm_runtime::{InvokeMethod, WasmInstance, WasmModule},
+	wasm_runtime::{HeapAllocStrategy, InvokeMethod, WasmInstance, WasmModule},
 };
 use sp_runtime_interface::unpack_ptr_and_len;
 use sp_wasm_interface::{Function, FunctionContext, Pointer, Result as WResult, WordSize};
 
+/// Wrapper around [`MemorRef`] that implements [`sc_allocator::Memory`].
+struct MemoryWrapper<'a>(&'a MemoryRef);
+
+impl sc_allocator::Memory for MemoryWrapper<'_> {
+	fn with_access_mut<R>(&mut self, run: impl FnOnce(&mut [u8]) -> R) -> R {
+		self.0.with_direct_access_mut(run)
+	}
+
+	fn with_access<R>(&self, run: impl FnOnce(&[u8]) -> R) -> R {
+		self.0.with_direct_access(run)
+	}
+
+	fn pages(&self) -> u32 {
+		self.0.current_size().0 as _
+	}
+
+	fn max_pages(&self) -> Option<u32> {
+		self.0.maximum().map(|p| p.0 as _)
+	}
+
+	fn grow(&mut self, additional: u32) -> Result<(), ()> {
+		self.0
+			.grow(Pages(additional as _))
+			.map_err(|e| {
+				log::error!(
+					target: "wasm-executor",
+					"Failed to grow memory by {} pages: {}",
+					additional,
+					e,
+				)
+			})
+			.map(drop)
+	}
+}
+
 struct FunctionExecutor {
 	heap: RefCell<sc_allocator::FreeingBumpHeapAllocator>,
 	memory: MemoryRef,
@@ -55,7 +90,7 @@ impl FunctionExecutor {
 		missing_functions: Arc<Vec<String>>,
 	) -> Result<Self, Error> {
 		Ok(FunctionExecutor {
-			heap: RefCell::new(sc_allocator::FreeingBumpHeapAllocator::new(heap_base)),
+			heap: RefCell::new(FreeingBumpHeapAllocator::new(heap_base)),
 			memory: m,
 			host_functions,
 			allow_missing_func_imports,
@@ -75,15 +110,17 @@ impl FunctionContext for FunctionExecutor {
 	}
 
 	fn allocate_memory(&mut self, size: WordSize) -> WResult<Pointer<u8>> {
-		let heap = &mut self.heap.borrow_mut();
-		self.memory
-			.with_direct_access_mut(|mem| heap.allocate(mem, size).map_err(|e| e.to_string()))
+		self.heap
+			.borrow_mut()
+			.allocate(&mut MemoryWrapper(&self.memory), size)
+			.map_err(|e| e.to_string())
 	}
 
 	fn deallocate_memory(&mut self, ptr: Pointer<u8>) -> WResult<()> {
-		let heap = &mut self.heap.borrow_mut();
-		self.memory
-			.with_direct_access_mut(|mem| heap.deallocate(mem, ptr).map_err(|e| e.to_string()))
+		self.heap
+			.borrow_mut()
+			.deallocate(&mut MemoryWrapper(&self.memory), ptr)
+			.map_err(|e| e.to_string())
 	}
 
 	fn register_panic_error_message(&mut self, message: &str) {
@@ -101,26 +138,17 @@ struct Resolver<'a> {
 	allow_missing_func_imports: bool,
 	/// All the names of functions for that we did not provide a host function.
 	missing_functions: RefCell<Vec<String>>,
-	/// Will be used as initial and maximum size of the imported memory.
-	heap_pages: usize,
-	/// By default, runtimes should import memory and this is `Some(_)` after
-	/// resolving. However, to be backwards compatible, we also support memory
-	/// exported by the WASM blob (this will be `None` after resolving).
-	import_memory: RefCell<Option<MemoryRef>>,
 }
 
 impl<'a> Resolver<'a> {
 	fn new(
 		host_functions: &'a [&'static dyn Function],
 		allow_missing_func_imports: bool,
-		heap_pages: usize,
 	) -> Resolver<'a> {
 		Resolver {
 			host_functions,
 			allow_missing_func_imports,
 			missing_functions: RefCell::new(Vec::new()),
-			heap_pages,
-			import_memory: Default::default(),
 		}
 	}
 }
@@ -148,7 +176,11 @@ impl<'a> wasmi::ModuleImportResolver for Resolver<'a> {
 		}
 
 		if self.allow_missing_func_imports {
-			trace!(target: "wasm-executor", "Could not find function `{}`, a stub will be provided instead.", name);
+			trace!(
+				target: "wasm-executor",
+				"Could not find function `{}`, a stub will be provided instead.",
+				name,
+			);
 			let id = self.missing_functions.borrow().len() + self.host_functions.len();
 			self.missing_functions.borrow_mut().push(name.to_string());
 
@@ -160,44 +192,12 @@ impl<'a> wasmi::ModuleImportResolver for Resolver<'a> {
 
 	fn resolve_memory(
 		&self,
-		field_name: &str,
-		memory_type: &wasmi::MemoryDescriptor,
+		_: &str,
+		_: &wasmi::MemoryDescriptor,
 	) -> Result<MemoryRef, wasmi::Error> {
-		if field_name == "memory" {
-			match &mut *self.import_memory.borrow_mut() {
-				Some(_) =>
-					Err(wasmi::Error::Instantiation("Memory can not be imported twice!".into())),
-				memory_ref @ None => {
-					if memory_type
-						.maximum()
-						.map(|m| m.saturating_sub(memory_type.initial()))
-						.map(|m| self.heap_pages > m as usize)
-						.unwrap_or(false)
-					{
-						Err(wasmi::Error::Instantiation(format!(
-							"Heap pages ({}) is greater than imported memory maximum ({}).",
-							self.heap_pages,
-							memory_type
-								.maximum()
-								.map(|m| m.saturating_sub(memory_type.initial()))
-								.expect("Maximum is set, checked above; qed"),
-						)))
-					} else {
-						let memory = MemoryInstance::alloc(
-							Pages(memory_type.initial() as usize + self.heap_pages),
-							Some(Pages(memory_type.initial() as usize + self.heap_pages)),
-						)?;
-						*memory_ref = Some(memory.clone());
-						Ok(memory)
-					}
-				},
-			}
-		} else {
-			Err(wasmi::Error::Instantiation(format!(
-				"Unknown memory reference with name: {}",
-				field_name
-			)))
-		}
+		Err(wasmi::Error::Instantiation(
+			"Internal error, wasmi expects that the wasm blob exports memory.".into(),
+		))
 	}
 }
 
@@ -358,12 +358,11 @@ fn call_in_wasm_module(
 
 /// Prepare module instance
 fn instantiate_module(
-	heap_pages: usize,
 	module: &Module,
 	host_functions: &[&'static dyn Function],
 	allow_missing_func_imports: bool,
 ) -> Result<(ModuleRef, Vec<String>, MemoryRef), Error> {
-	let resolver = Resolver::new(host_functions, allow_missing_func_imports, heap_pages);
+	let resolver = Resolver::new(host_functions, allow_missing_func_imports);
 	// start module instantiation. Don't run 'start' function yet.
 	let intermediate_instance =
 		ModuleInstance::new(module, &ImportsBuilder::new().with_resolver("env", &resolver))?;
@@ -371,22 +370,10 @@ fn instantiate_module(
 	// Verify that the module has the heap base global variable.
 	let _ = get_heap_base(intermediate_instance.not_started_instance())?;
 
-	// Get the memory reference. Runtimes should import memory, but to be backwards
-	// compatible we also support exported memory.
-	let memory = match resolver.import_memory.into_inner() {
-		Some(memory) => memory,
-		None => {
-			debug!(
-				target: "wasm-executor",
-				"WASM blob does not imports memory, falling back to exported memory",
-			);
-
-			let memory = get_mem_instance(intermediate_instance.not_started_instance())?;
-			memory.grow(Pages(heap_pages)).map_err(|_| Error::Runtime)?;
-
-			memory
-		},
-	};
+	// The `module` should export the memory with the correct properties (min, max).
+	//
+	// This is ensured by modifying the `RuntimeBlob` before initializing the `Module`.
+	let memory = get_mem_instance(intermediate_instance.not_started_instance())?;
 
 	if intermediate_instance.has_start() {
 		// Runtime is not allowed to have the `start` function.
@@ -451,8 +438,6 @@ pub struct WasmiRuntime {
 	/// Enable stub generation for functions that are not available in `host_functions`.
 	/// These stubs will error when the wasm blob tries to call them.
 	allow_missing_func_imports: bool,
-	/// Numer of heap pages this runtime uses.
-	heap_pages: u64,
 
 	global_vals_snapshot: GlobalValsSnapshot,
 	data_segments_snapshot: DataSegmentsSnapshot,
@@ -461,13 +446,9 @@ pub struct WasmiRuntime {
 impl WasmModule for WasmiRuntime {
 	fn new_instance(&self) -> Result<Box<dyn WasmInstance>, Error> {
 		// Instantiate this module.
-		let (instance, missing_functions, memory) = instantiate_module(
-			self.heap_pages as usize,
-			&self.module,
-			&self.host_functions,
-			self.allow_missing_func_imports,
-		)
-		.map_err(|e| WasmError::Instantiation(e.to_string()))?;
+		let (instance, missing_functions, memory) =
+			instantiate_module(&self.module, &self.host_functions, self.allow_missing_func_imports)
+				.map_err(|e| WasmError::Instantiation(e.to_string()))?;
 
 		Ok(Box::new(WasmiInstance {
 			instance,
@@ -477,6 +458,7 @@ impl WasmModule for WasmiRuntime {
 			host_functions: self.host_functions.clone(),
 			allow_missing_func_imports: self.allow_missing_func_imports,
 			missing_functions: Arc::new(missing_functions),
+			memory_zeroed: true,
 		}))
 	}
 }
@@ -484,25 +466,26 @@ impl WasmModule for WasmiRuntime {
 /// Create a new `WasmiRuntime` given the code. This function loads the module and
 /// stores it in the instance.
 pub fn create_runtime(
-	blob: RuntimeBlob,
-	heap_pages: u64,
+	mut blob: RuntimeBlob,
+	heap_alloc_strategy: HeapAllocStrategy,
 	host_functions: Vec<&'static dyn Function>,
 	allow_missing_func_imports: bool,
 ) -> Result<WasmiRuntime, WasmError> {
 	let data_segments_snapshot =
 		DataSegmentsSnapshot::take(&blob).map_err(|e| WasmError::Other(e.to_string()))?;
 
+	// Make sure we only have exported memory to simplify the code of the wasmi executor.
+	blob.convert_memory_import_into_export()?;
+	// Ensure that the memory uses the correct heap pages.
+	blob.setup_memory_according_to_heap_alloc_strategy(heap_alloc_strategy)?;
+
 	let module =
 		Module::from_parity_wasm_module(blob.into_inner()).map_err(|_| WasmError::InvalidModule)?;
 
 	let global_vals_snapshot = {
-		let (instance, _, _) = instantiate_module(
-			heap_pages as usize,
-			&module,
-			&host_functions,
-			allow_missing_func_imports,
-		)
-		.map_err(|e| WasmError::Instantiation(e.to_string()))?;
+		let (instance, _, _) =
+			instantiate_module(&module, &host_functions, allow_missing_func_imports)
+				.map_err(|e| WasmError::Instantiation(e.to_string()))?;
 		GlobalValsSnapshot::take(&instance)
 	};
 
@@ -512,7 +495,6 @@ pub fn create_runtime(
 		global_vals_snapshot,
 		host_functions: Arc::new(host_functions),
 		allow_missing_func_imports,
-		heap_pages,
 	})
 }
 
@@ -522,6 +504,8 @@ pub struct WasmiInstance {
 	instance: ModuleRef,
 	/// The memory instance of used by the wasm module.
 	memory: MemoryRef,
+	/// Is the memory zeroed?
+	memory_zeroed: bool,
 	/// The snapshot of global variable values just after instantiation.
 	global_vals_snapshot: GlobalValsSnapshot,
 	/// The snapshot of data segments.
@@ -549,14 +533,16 @@ impl WasmiInstance {
 		// We reuse a single wasm instance for multiple calls and a previous call (if any)
 		// altered the state. Therefore, we need to restore the instance to original state.
 
-		// First, zero initialize the linear memory.
-		self.memory.erase().map_err(|e| {
-			// Snapshot restoration failed. This is pretty unexpected since this can happen
-			// if some invariant is broken or if the system is under extreme memory pressure
-			// (so erasing fails).
-			error!(target: "wasm-executor", "snapshot restoration failed: {}", e);
-			WasmError::ErasingFailed(e.to_string())
-		})?;
+		if !self.memory_zeroed {
+			// First, zero initialize the linear memory.
+			self.memory.erase().map_err(|e| {
+				// Snapshot restoration failed. This is pretty unexpected since this can happen
+				// if some invariant is broken or if the system is under extreme memory pressure
+				// (so erasing fails).
+				error!(target: "wasm-executor", "snapshot restoration failed: {}", e);
+				WasmError::ErasingFailed(e.to_string())
+			})?;
+		}
 
 		// Second, reapply data segments into the linear memory.
 		self.data_segments_snapshot
@@ -565,7 +551,7 @@ impl WasmiInstance {
 		// Third, restore the global variables to their initial values.
 		self.global_vals_snapshot.apply(&self.instance)?;
 
-		call_in_wasm_module(
+		let res = call_in_wasm_module(
 			&self.instance,
 			&self.memory,
 			method,
@@ -574,7 +560,12 @@ impl WasmiInstance {
 			self.allow_missing_func_imports,
 			self.missing_functions.clone(),
 			allocation_stats,
-		)
+		);
+
+		// If we couldn't unmap it, erase the memory.
+		self.memory_zeroed = self.memory.erase().is_ok();
+
+		res
 	}
 }
 
@@ -601,4 +592,8 @@ impl WasmInstance for WasmiInstance {
 			None => Ok(None),
 		}
 	}
+
+	fn linear_memory_base_ptr(&self) -> Option<*const u8> {
+		Some(self.memory.direct_access().as_ref().as_ptr())
+	}
 }

substrate/client/executor/wasmtime/Cargo.toml

@@ -13,9 +13,9 @@ readme = "README.md"
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
+log = "0.4.17"
 cfg-if = "1.0"
 libc = "0.2.121"
-log = "0.4.17"
 
 # When bumping wasmtime do not forget to also bump rustix
 # to exactly the same version as used by wasmtime!

substrate/client/executor/wasmtime/src/host.rs

@@ -24,20 +24,25 @@ use wasmtime::Caller;
 use sc_allocator::{AllocationStats, FreeingBumpHeapAllocator};
 use sp_wasm_interface::{Pointer, WordSize};
 
-use crate::{runtime::StoreData, util};
+use crate::{instance_wrapper::MemoryWrapper, runtime::StoreData, util};
 
 /// The state required to construct a HostContext context. The context only lasts for one host
 /// call, whereas the state is maintained for the duration of a Wasm runtime call, which may make
 /// many different host calls that must share state.
 pub struct HostState {
-	allocator: FreeingBumpHeapAllocator,
+	/// The allocator instance to keep track of allocated memory.
+	///
+	/// This is stored as an `Option` as we need to temporarly set this to `None` when we are
+	/// allocating/deallocating memory. The problem being that we can only mutable access `caller`
+	/// once.
+	allocator: Option<FreeingBumpHeapAllocator>,
 	panic_message: Option<String>,
 }
 
 impl HostState {
 	/// Constructs a new `HostState`.
 	pub fn new(allocator: FreeingBumpHeapAllocator) -> Self {
-		HostState { allocator, panic_message: None }
+		HostState { allocator: Some(allocator), panic_message: None }
 	}
 
 	/// Takes the error message out of the host state, leaving a `None` in its place.
@@ -46,7 +51,9 @@ impl HostState {
 	}
 
 	pub(crate) fn allocation_stats(&self) -> AllocationStats {
-		self.allocator.stats()
+		self.allocator.as_ref()
+			.expect("Allocator is always set and only unavailable when doing an allocation/deallocation; qed")
+			.stats()
 	}
 }
 
@@ -81,22 +88,38 @@ impl<'a> sp_wasm_interface::FunctionContext for HostContext<'a> {
 
 	fn allocate_memory(&mut self, size: WordSize) -> sp_wasm_interface::Result<Pointer<u8>> {
 		let memory = self.caller.data().memory();
-		let (memory, data) = memory.data_and_store_mut(&mut self.caller);
-		data.host_state_mut()
-			.expect("host state is not empty when calling a function in wasm; qed")
+		let mut allocator = self
+			.host_state_mut()
 			.allocator
-			.allocate(memory, size)
-			.map_err(|e| e.to_string())
+			.take()
+			.expect("allocator is not empty when calling a function in wasm; qed");
+
+		// We can not return on error early, as we need to store back allocator.
+		let res = allocator
+			.allocate(&mut MemoryWrapper(&memory, &mut self.caller), size)
+			.map_err(|e| e.to_string());
+
+		self.host_state_mut().allocator = Some(allocator);
+
+		res
 	}
 
 	fn deallocate_memory(&mut self, ptr: Pointer<u8>) -> sp_wasm_interface::Result<()> {
 		let memory = self.caller.data().memory();
-		let (memory, data) = memory.data_and_store_mut(&mut self.caller);
-		data.host_state_mut()
-			.expect("host state is not empty when calling a function in wasm; qed")
+		let mut allocator = self
+			.host_state_mut()
 			.allocator
-			.deallocate(memory, ptr)
-			.map_err(|e| e.to_string())
+			.take()
+			.expect("allocator is not empty when calling a function in wasm; qed");
+
+		// We can not return on error early, as we need to store back allocator.
+		let res = allocator
+			.deallocate(&mut MemoryWrapper(&memory, &mut self.caller), ptr)
+			.map_err(|e| e.to_string());
+
+		self.host_state_mut().allocator = Some(allocator);
+
+		res
 	}
 
 	fn register_panic_error_message(&mut self, message: &str) {

substrate/client/executor/wasmtime/src/instance_wrapper.rs

@@ -26,8 +26,7 @@ use sc_executor_common::{
 };
 use sp_wasm_interface::{Pointer, Value, WordSize};
 use wasmtime::{
-	AsContext, AsContextMut, Engine, Extern, Func, Global, Instance, InstancePre, Memory, Table,
-	Val,
+	AsContext, AsContextMut, Engine, Extern, Instance, InstancePre, Memory, Table, Val,
 };
 
 /// Invoked entrypoint format.
@@ -113,66 +112,58 @@ impl EntryPoint {
 	}
 }
 
+/// Wrapper around [`Memory`] that implements [`sc_allocator::Memory`].
+pub(crate) struct MemoryWrapper<'a, C>(pub &'a wasmtime::Memory, pub &'a mut C);
+
+impl<C: AsContextMut> sc_allocator::Memory for MemoryWrapper<'_, C> {
+	fn with_access<R>(&self, run: impl FnOnce(&[u8]) -> R) -> R {
+		run(self.0.data(&self.1))
+	}
+
+	fn with_access_mut<R>(&mut self, run: impl FnOnce(&mut [u8]) -> R) -> R {
+		run(self.0.data_mut(&mut self.1))
+	}
+
+	fn grow(&mut self, additional: u32) -> std::result::Result<(), ()> {
+		self.0
+			.grow(&mut self.1, additional as u64)
+			.map_err(|e| {
+				log::error!(
+					target: "wasm-executor",
+					"Failed to grow memory by {} pages: {}",
+					additional,
+					e,
+				)
+			})
+			.map(drop)
+	}
+
+	fn pages(&self) -> u32 {
+		self.0.size(&self.1) as u32
+	}
+
+	fn max_pages(&self) -> Option<u32> {
+		self.0.ty(&self.1).maximum().map(|p| p as _)
+	}
+}
+
 /// Wrap the given WebAssembly Instance of a wasm module with Substrate-runtime.
 ///
 /// This struct is a handy wrapper around a wasmtime `Instance` that provides substrate specific
 /// routines.
 pub struct InstanceWrapper {
 	instance: Instance,
+	/// The memory instance of the `instance`.
+	///
+	/// It is important to make sure that we don't make any copies of this to make it easier to
+	/// proof
 	memory: Memory,
 	store: Store,
 }
 
-fn extern_memory(extern_: &Extern) -> Option<&Memory> {
-	match extern_ {
-		Extern::Memory(mem) => Some(mem),
-		_ => None,
-	}
-}
-
-fn extern_global(extern_: &Extern) -> Option<&Global> {
-	match extern_ {
-		Extern::Global(glob) => Some(glob),
-		_ => None,
-	}
-}
-
-fn extern_table(extern_: &Extern) -> Option<&Table> {
-	match extern_ {
-		Extern::Table(table) => Some(table),
-		_ => None,
-	}
-}
-
-fn extern_func(extern_: &Extern) -> Option<&Func> {
-	match extern_ {
-		Extern::Func(func) => Some(func),
-		_ => None,
-	}
-}
-
-pub(crate) fn create_store(engine: &wasmtime::Engine, max_memory_size: Option<usize>) -> Store {
-	let limits = if let Some(max_memory_size) = max_memory_size {
-		wasmtime::StoreLimitsBuilder::new().memory_size(max_memory_size).build()
-	} else {
-		Default::default()
-	};
-
-	let mut store =
-		Store::new(engine, StoreData { limits, host_state: None, memory: None, table: None });
-	if max_memory_size.is_some() {
-		store.limiter(|s| &mut s.limits);
-	}
-	store
-}
-
 impl InstanceWrapper {
-	pub(crate) fn new(
-		engine: &Engine,
-		instance_pre: &InstancePre<StoreData>,
-		max_memory_size: Option<usize>,
-	) -> Result<Self> {
-		let mut store = create_store(engine, max_memory_size);
+	pub(crate) fn new(engine: &Engine, instance_pre: &InstancePre<StoreData>) -> Result<Self> {
+		let mut store = Store::new(engine, Default::default());
 		let instance = instance_pre.instantiate(&mut store).map_err(|error| {
 			WasmError::Other(format!(
 				"failed to instantiate a new WASM module instance: {:#}",
@@ -201,9 +192,10 @@ impl InstanceWrapper {
 					self.instance.get_export(&mut self.store, method).ok_or_else(|| {
 						Error::from(format!("Exported method {} is not found", method))
 					})?;
-				let func = extern_func(&export)
+				let func = export
+					.into_func()
 					.ok_or_else(|| Error::from(format!("Export {} is not a function", method)))?;
-				EntryPoint::direct(*func, &self.store).map_err(|_| {
+				EntryPoint::direct(func, &self.store).map_err(|_| {
 					Error::from(format!("Exported function '{}' has invalid signature.", method))
 				})?
 			},
@@ -259,7 +251,8 @@ impl InstanceWrapper {
 			.get_export(&mut self.store, "__heap_base")
 			.ok_or_else(|| Error::from("__heap_base is not found"))?;
 
-		let heap_base_global = extern_global(&heap_base_export)
+		let heap_base_global = heap_base_export
+			.into_global()
 			.ok_or_else(|| Error::from("__heap_base is not a global"))?;
 
 		let heap_base = heap_base_global
@@ -277,7 +270,7 @@ impl InstanceWrapper {
 			None => return Ok(None),
 		};
 
-		let global = extern_global(&global).ok_or_else(|| format!("`{}` is not a global", name))?;
+		let global = global.into_global().ok_or_else(|| format!("`{}` is not a global", name))?;
 
 		match global.get(&mut self.store) {
 			Val::I32(val) => Ok(Some(Value::I32(val))),
@@ -300,7 +293,8 @@ fn get_linear_memory(instance: &Instance, ctx: impl AsContextMut) -> Result<Memo
 		.get_export(ctx, "memory")
 		.ok_or_else(|| Error::from("memory is not exported under `memory` name"))?;
 
-	let memory = *extern_memory(&memory_export)
+	let memory = memory_export
+		.into_memory()
 		.ok_or_else(|| Error::from("the `memory` export should have memory type"))?;
 
 	Ok(memory)
@@ -311,8 +305,8 @@ fn get_table(instance: &Instance, ctx: &mut Store) -> Option<Table> {
 	instance
 		.get_export(ctx, "__indirect_function_table")
 		.as_ref()
-		.and_then(extern_table)
 		.cloned()
+		.and_then(Extern::into_table)
 }
 
 /// Functions related to memory.
@@ -403,7 +397,7 @@ fn decommit_works() {
 	let module = wasmtime::Module::new(&engine, code).unwrap();
 	let linker = wasmtime::Linker::new(&engine);
 	let instance_pre = linker.instantiate_pre(&module).unwrap();
-	let mut wrapper = InstanceWrapper::new(&engine, &instance_pre, None).unwrap();
+	let mut wrapper = InstanceWrapper::new(&engine, &instance_pre).unwrap();
 	unsafe { *wrapper.memory.data_ptr(&wrapper.store) = 42 };
 	assert_eq!(unsafe { *wrapper.memory.data_ptr(&wrapper.store) }, 42);
 	wrapper.decommit();

substrate/client/executor/wasmtime/src/runtime.rs

@@ -20,7 +20,7 @@
 
 use crate::{
 	host::HostState,
-	instance_wrapper::{EntryPoint, InstanceWrapper},
+	instance_wrapper::{EntryPoint, InstanceWrapper, MemoryWrapper},
 	util::{self, replace_strategy_if_broken},
 };
 
@@ -31,7 +31,7 @@ use sc_executor_common::{
 		self, DataSegmentsSnapshot, ExposedMutableGlobalsSet, GlobalsSnapshot, RuntimeBlob,
 	},
 	util::checked_range,
-	wasm_runtime::{InvokeMethod, WasmInstance, WasmModule},
+	wasm_runtime::{HeapAllocStrategy, InvokeMethod, WasmInstance, WasmModule},
 };
 use sp_runtime_interface::unpack_ptr_and_len;
 use sp_wasm_interface::{HostFunctions, Pointer, Value, WordSize};
@@ -42,12 +42,10 @@ use std::{
 		Arc,
 	},
 };
-use wasmtime::{AsContext, Engine, Memory, StoreLimits, Table};
+use wasmtime::{AsContext, Engine, Memory, Table};
 
+#[derive(Default)]
 pub(crate) struct StoreData {
-	/// The limits we apply to the store. We need to store it here to return a reference to this
-	/// object when we have the limits enabled.
-	pub(crate) limits: StoreLimits,
 	/// This will only be set when we call into the runtime.
 	pub(crate) host_state: Option<HostState>,
 	/// This will be always set once the store is initialized.
@@ -83,12 +81,11 @@ enum Strategy {
 struct InstanceCreator {
 	engine: wasmtime::Engine,
 	instance_pre: Arc<wasmtime::InstancePre<StoreData>>,
-	max_memory_size: Option<usize>,
 }
 
 impl InstanceCreator {
 	fn instantiate(&mut self) -> Result<InstanceWrapper> {
-		InstanceWrapper::new(&self.engine, &self.instance_pre, self.max_memory_size)
+		InstanceWrapper::new(&self.engine, &self.instance_pre)
 	}
 }
 
@@ -128,18 +125,13 @@ pub struct WasmtimeRuntime {
 	engine: wasmtime::Engine,
 	instance_pre: Arc<wasmtime::InstancePre<StoreData>>,
 	instantiation_strategy: InternalInstantiationStrategy,
-	config: Config,
 }
 
 impl WasmModule for WasmtimeRuntime {
 	fn new_instance(&self) -> Result<Box<dyn WasmInstance>> {
 		let strategy = match self.instantiation_strategy {
 			InternalInstantiationStrategy::LegacyInstanceReuse(ref snapshot_data) => {
-				let mut instance_wrapper = InstanceWrapper::new(
-					&self.engine,
-					&self.instance_pre,
-					self.config.semantics.max_memory_size,
-				)?;
+				let mut instance_wrapper = InstanceWrapper::new(&self.engine, &self.instance_pre)?;
 				let heap_base = instance_wrapper.extract_heap_base()?;
 
 				// This function panics if the instance was created from a runtime blob different
@@ -161,7 +153,6 @@ impl WasmModule for WasmtimeRuntime {
 			InternalInstantiationStrategy::Builtin => Strategy::RecreateInstance(InstanceCreator {
 				engine: self.engine.clone(),
 				instance_pre: self.instance_pre.clone(),
-				max_memory_size: self.config.semantics.max_memory_size,
 			}),
 		};
 
@@ -350,25 +341,22 @@ fn common_config(semantics: &Semantics) -> std::result::Result<wasmtime::Config,
 		InstantiationStrategy::LegacyInstanceReuse => (false, false),
 	};
 
+	const WASM_PAGE_SIZE: u64 = 65536;
+
 	config.memory_init_cow(use_cow);
-	config.memory_guaranteed_dense_image_size(
-		semantics.max_memory_size.map(|max| max as u64).unwrap_or(u64::MAX),
-	);
+	config.memory_guaranteed_dense_image_size(match semantics.heap_alloc_strategy {
+		HeapAllocStrategy::Dynamic { maximum_pages } =>
+			maximum_pages.map(|p| p as u64 * WASM_PAGE_SIZE).unwrap_or(u64::MAX),
+		HeapAllocStrategy::Static { .. } => u64::MAX,
+	});
 
 	if use_pooling {
-		const WASM_PAGE_SIZE: u64 = 65536;
 		const MAX_WASM_PAGES: u64 = 0x10000;
 
-		let memory_pages = if let Some(max_memory_size) = semantics.max_memory_size {
-			let max_memory_size = max_memory_size as u64;
-			let mut pages = max_memory_size / WASM_PAGE_SIZE;
-			if max_memory_size % WASM_PAGE_SIZE != 0 {
-				pages += 1;
-			}
-
-			std::cmp::min(MAX_WASM_PAGES, pages)
-		} else {
-			MAX_WASM_PAGES
+		let memory_pages = match semantics.heap_alloc_strategy {
+			HeapAllocStrategy::Dynamic { maximum_pages } =>
+				maximum_pages.map(|p| p as u64).unwrap_or(MAX_WASM_PAGES),
+			HeapAllocStrategy::Static { .. } => MAX_WASM_PAGES,
 		};
 
 		let mut pooling_config = wasmtime::PoolingAllocationConfig::default();
@@ -514,25 +502,8 @@ pub struct Semantics {
 	/// Configures wasmtime to use multiple threads for compiling.
 	pub parallel_compilation: bool,
 
-	/// The number of extra WASM pages which will be allocated
-	/// on top of what is requested by the WASM blob itself.
-	pub extra_heap_pages: u64,
-
-	/// The total amount of memory in bytes an instance can request.
-	///
-	/// If specified, the runtime will be able to allocate only that much of wasm memory.
-	/// This is the total number and therefore the [`Semantics::extra_heap_pages`] is accounted
-	/// for.
-	///
-	/// That means that the initial number of pages of a linear memory plus the
-	/// [`Semantics::extra_heap_pages`] multiplied by the wasm page size (64KiB) should be less
-	/// than or equal to `max_memory_size`, otherwise the instance won't be created.
-	///
-	/// Moreover, `memory.grow` will fail (return -1) if the sum of sizes of currently mounted
-	/// and additional pages exceeds `max_memory_size`.
-	///
-	/// The default is `None`.
-	pub max_memory_size: Option<usize>,
+	/// The heap allocation strategy to use.
+	pub heap_alloc_strategy: HeapAllocStrategy,
 }
 
 #[derive(Clone)]
@@ -689,12 +660,7 @@ where
 		.instantiate_pre(&module)
 		.map_err(|e| WasmError::Other(format!("cannot preinstantiate module: {:#}", e)))?;
 
-	Ok(WasmtimeRuntime {
-		engine,
-		instance_pre: Arc::new(instance_pre),
-		instantiation_strategy,
-		config,
-	})
+	Ok(WasmtimeRuntime { engine, instance_pre: Arc::new(instance_pre), instantiation_strategy })
 }
 
 fn prepare_blob_for_compilation(
@@ -717,12 +683,7 @@ fn prepare_blob_for_compilation(
 	// now automatically take care of creating the memory for us, and it is also necessary
 	// to enable `wasmtime`'s instance pooling. (Imported memories are ineligible for pooling.)
 	blob.convert_memory_import_into_export()?;
-	blob.add_extra_heap_pages_to_memory_section(
-		semantics
-			.extra_heap_pages
-			.try_into()
-			.map_err(|e| WasmError::Other(format!("invalid `extra_heap_pages`: {}", e)))?,
-	)?;
+	blob.setup_memory_according_to_heap_alloc_strategy(semantics.heap_alloc_strategy)?;
 
 	Ok(blob)
 }
@@ -783,9 +744,8 @@ fn inject_input_data(
 ) -> Result<(Pointer<u8>, WordSize)> {
 	let mut ctx = instance.store_mut();
 	let memory = ctx.data().memory();
-	let memory = memory.data_mut(&mut ctx);
 	let data_len = data.len() as WordSize;
-	let data_ptr = allocator.allocate(memory, data_len)?;
+	let data_ptr = allocator.allocate(&mut MemoryWrapper(&memory, &mut ctx), data_len)?;
 	util::write_memory_from(instance.store_mut(), data_ptr, data)?;
 	Ok((data_ptr, data_len))
 }

substrate/client/executor/wasmtime/src/tests.rs

@@ -17,7 +17,11 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 use codec::{Decode as _, Encode as _};
-use sc_executor_common::{error::Error, runtime_blob::RuntimeBlob, wasm_runtime::WasmModule};
+use sc_executor_common::{
+	error::Error,
+	runtime_blob::RuntimeBlob,
+	wasm_runtime::{HeapAllocStrategy, WasmModule},
+};
 use sc_runtime_test::wasm_binary_unwrap;
 
 use crate::InstantiationStrategy;
@@ -77,8 +81,7 @@ struct RuntimeBuilder {
 	instantiation_strategy: InstantiationStrategy,
 	canonicalize_nans: bool,
 	deterministic_stack: bool,
-	extra_heap_pages: u64,
-	max_memory_size: Option<usize>,
+	heap_pages: HeapAllocStrategy,
 	precompile_runtime: bool,
 	tmpdir: Option<tempfile::TempDir>,
 }
@@ -90,8 +93,7 @@ impl RuntimeBuilder {
 			instantiation_strategy,
 			canonicalize_nans: false,
 			deterministic_stack: false,
-			extra_heap_pages: 1024,
-			max_memory_size: None,
+			heap_pages: HeapAllocStrategy::Static { extra_pages: 1024 },
 			precompile_runtime: false,
 			tmpdir: None,
 		}
@@ -117,8 +119,8 @@ impl RuntimeBuilder {
 		self
 	}
 
-	fn max_memory_size(mut self, max_memory_size: Option<usize>) -> Self {
-		self.max_memory_size = max_memory_size;
+	fn heap_alloc_strategy(mut self, heap_pages: HeapAllocStrategy) -> Self {
+		self.heap_pages = heap_pages;
 		self
 	}
 
@@ -152,8 +154,7 @@ impl RuntimeBuilder {
 				},
 				canonicalize_nans: self.canonicalize_nans,
 				parallel_compilation: true,
-				extra_heap_pages: self.extra_heap_pages,
-				max_memory_size: self.max_memory_size,
+				heap_alloc_strategy: self.heap_pages,
 			},
 		};
 
@@ -227,7 +228,7 @@ fn deep_call_stack_wat(depth: usize) -> String {
 // We need two limits here since depending on whether the code is compiled in debug
 // or in release mode the maximum call depth is slightly different.
 const CALL_DEPTH_LOWER_LIMIT: usize = 65455;
-const CALL_DEPTH_UPPER_LIMIT: usize = 65503;
+const CALL_DEPTH_UPPER_LIMIT: usize = 65509;
 
 test_wasm_execution!(test_consume_under_1mb_of_stack_does_not_trap);
 fn test_consume_under_1mb_of_stack_does_not_trap(instantiation_strategy: InstantiationStrategy) {
@@ -344,29 +345,25 @@ fn test_max_memory_pages(
 	import_memory: bool,
 	precompile_runtime: bool,
 ) {
-	fn try_instantiate(
-		max_memory_size: Option<usize>,
+	fn call(
+		heap_alloc_strategy: HeapAllocStrategy,
 		wat: String,
 		instantiation_strategy: InstantiationStrategy,
 		precompile_runtime: bool,
 	) -> Result<(), Box<dyn std::error::Error>> {
 		let mut builder = RuntimeBuilder::new(instantiation_strategy)
 			.use_wat(wat)
-			.max_memory_size(max_memory_size)
+			.heap_alloc_strategy(heap_alloc_strategy)
 			.precompile_runtime(precompile_runtime);
 
 		let runtime = builder.build();
-		let mut instance = runtime.new_instance()?;
+		let mut instance = runtime.new_instance().unwrap();
 		let _ = instance.call_export("main", &[])?;
 		Ok(())
 	}
 
-	fn memory(initial: u32, maximum: Option<u32>, import: bool) -> String {
-		let memory = if let Some(maximum) = maximum {
-			format!("(memory $0 {} {})", initial, maximum)
-		} else {
-			format!("(memory $0 {})", initial)
-		};
+	fn memory(initial: u32, maximum: u32, import: bool) -> String {
+		let memory = format!("(memory $0 {} {})", initial, maximum);
 
 		if import {
 			format!("(import \"env\" \"memory\" {})", memory)
@@ -375,152 +372,90 @@ fn test_max_memory_pages(
 		}
 	}
 
-	const WASM_PAGE_SIZE: usize = 65536;
+	let assert_grow_ok = |alloc_strategy: HeapAllocStrategy, initial_pages: u32, max_pages: u32| {
+		eprintln!("assert_grow_ok({alloc_strategy:?}, {initial_pages}, {max_pages})");
 
-	// check the old behavior if preserved. That is, if no limit is set we allow 4 GiB of memory.
-	try_instantiate(
-		None,
-		format!(
-			r#"
-			(module
-				{}
-				(global (export "__heap_base") i32 (i32.const 0))
-				(func (export "main")
-					(param i32 i32) (result i64)
-					(i64.const 0)
-				)
-			)
-			"#,
-			/*
-				We want to allocate the maximum number of pages supported in wasm for this test.
-				However, due to a bug in wasmtime (I think wasmi is also affected) it is only possible
-				to allocate 65536 - 1 pages.
+		call(
+			alloc_strategy,
+			format!(
+				r#"
+					(module
+						{}
+						(global (export "__heap_base") i32 (i32.const 0))
+						(func (export "main")
+							(param i32 i32) (result i64)
 
-				Then, during creation of the Substrate Runtime instance, 1024 (heap_pages) pages are
-				mounted.
-
-				Thus 65535 = 64511 + 1024
-			*/
-			memory(64511, None, import_memory)
-		),
-		instantiation_strategy,
-		precompile_runtime,
-	)
-	.unwrap();
-
-	// max is not specified, therefore it's implied to be 65536 pages (4 GiB).
-	//
-	// max_memory_size = (1 (initial) + 1024 (heap_pages)) * WASM_PAGE_SIZE
-	try_instantiate(
-		Some((1 + 1024) * WASM_PAGE_SIZE),
-		format!(
-			r#"
-			(module
-				{}
-				(global (export "__heap_base") i32 (i32.const 0))
-				(func (export "main")
-					(param i32 i32) (result i64)
-					(i64.const 0)
-				)
-			)
-			"#,
-			// 1 initial, max is not specified.
-			memory(1, None, import_memory)
-		),
-		instantiation_strategy,
-		precompile_runtime,
-	)
-	.unwrap();
-
-	// max is specified explicitly to 2048 pages.
-	try_instantiate(
-		Some((1 + 1024) * WASM_PAGE_SIZE),
-		format!(
-			r#"
-			(module
-				{}
-				(global (export "__heap_base") i32 (i32.const 0))
-				(func (export "main")
-					(param i32 i32) (result i64)
-					(i64.const 0)
-				)
-			)
-			"#,
-			// Max is 2048.
-			memory(1, Some(2048), import_memory)
-		),
-		instantiation_strategy,
-		precompile_runtime,
-	)
-	.unwrap();
-
-	// memory grow should work as long as it doesn't exceed 1025 pages in total.
-	try_instantiate(
-		Some((0 + 1024 + 25) * WASM_PAGE_SIZE),
-		format!(
-			r#"
-			(module
-				{}
-				(global (export "__heap_base") i32 (i32.const 0))
-				(func (export "main")
-					(param i32 i32) (result i64)
-
-					;; assert(memory.grow returns != -1)
-					(if
-						(i32.eq
-							(memory.grow
-								(i32.const 25)
+							;; assert(memory.grow returns != -1)
+							(if
+								(i32.eq
+									(memory.grow
+										(i32.const 1)
+									)
+									(i32.const -1)
+								)
+								(unreachable)
 							)
-							(i32.const -1)
+
+							(i64.const 0)
 						)
-						(unreachable)
 					)
-
-					(i64.const 0)
-				)
-			)
 			"#,
-			// Zero starting pages.
-			memory(0, None, import_memory)
-		),
-		instantiation_strategy,
-		precompile_runtime,
-	)
-	.unwrap();
+				memory(initial_pages, max_pages, import_memory)
+			),
+			instantiation_strategy,
+			precompile_runtime,
+		)
+		.unwrap()
+	};
 
-	// We start with 1025 pages and try to grow at least one.
-	try_instantiate(
-		Some((1 + 1024) * WASM_PAGE_SIZE),
-		format!(
-			r#"
-			(module
-				{}
-				(global (export "__heap_base") i32 (i32.const 0))
-				(func (export "main")
-					(param i32 i32) (result i64)
+	let assert_grow_fail =
+		|alloc_strategy: HeapAllocStrategy, initial_pages: u32, max_pages: u32| {
+			eprintln!("assert_grow_fail({alloc_strategy:?}, {initial_pages}, {max_pages})");
 
-					;; assert(memory.grow returns == -1)
-					(if
-						(i32.ne
-							(memory.grow
-								(i32.const 1)
+			call(
+				alloc_strategy,
+				format!(
+					r#"
+						(module
+							{}
+							(global (export "__heap_base") i32 (i32.const 0))
+							(func (export "main")
+								(param i32 i32) (result i64)
+
+								;; assert(memory.grow returns == -1)
+								(if
+									(i32.ne
+										(memory.grow
+											(i32.const 1)
+										)
+										(i32.const -1)
+									)
+									(unreachable)
+								)
+
+								(i64.const 0)
 							)
-							(i32.const -1)
 						)
-						(unreachable)
-					)
-
-					(i64.const 0)
-				)
+					"#,
+					memory(initial_pages, max_pages, import_memory)
+				),
+				instantiation_strategy,
+				precompile_runtime,
 			)
-			"#,
-			// Initial=1, meaning after heap pages mount the total will be already 1025.
-			memory(1, None, import_memory)
-		),
-		instantiation_strategy,
-		precompile_runtime,
-	)
-	.unwrap();
+			.unwrap()
+		};
+
+	assert_grow_ok(HeapAllocStrategy::Dynamic { maximum_pages: Some(10) }, 1, 10);
+	assert_grow_ok(HeapAllocStrategy::Dynamic { maximum_pages: Some(10) }, 9, 10);
+	assert_grow_fail(HeapAllocStrategy::Dynamic { maximum_pages: Some(10) }, 10, 10);
+
+	assert_grow_ok(HeapAllocStrategy::Dynamic { maximum_pages: None }, 1, 10);
+	assert_grow_ok(HeapAllocStrategy::Dynamic { maximum_pages: None }, 9, 10);
+	assert_grow_ok(HeapAllocStrategy::Dynamic { maximum_pages: None }, 10, 10);
+
+	assert_grow_fail(HeapAllocStrategy::Static { extra_pages: 10 }, 1, 10);
+	assert_grow_fail(HeapAllocStrategy::Static { extra_pages: 10 }, 9, 10);
+	assert_grow_fail(HeapAllocStrategy::Static { extra_pages: 10 }, 10, 10);
 }
 
 // This test takes quite a while to execute in a debug build (over 6 minutes on a TR 3970x)
@@ -538,8 +473,7 @@ fn test_instances_without_reuse_are_not_leaked() {
 				deterministic_stack_limit: None,
 				canonicalize_nans: false,
 				parallel_compilation: true,
-				extra_heap_pages: 2048,
-				max_memory_size: None,
+				heap_alloc_strategy: HeapAllocStrategy::Static { extra_pages: 2048 },
 			},
 		},
 	)
@@ -583,6 +517,10 @@ fn test_rustix_version_matches_with_wasmtime() {
 		.unwrap();
 
 	if wasmtime_rustix.req != our_rustix.req {
-		panic!("our version of rustix ({0}) doesn't match wasmtime's ({1}); bump the version in `sc-executor-wasmtime`'s `Cargo.toml' to '{1}' and try again", our_rustix.req, wasmtime_rustix.req);
+		panic!(
+			"our version of rustix ({0}) doesn't match wasmtime's ({1}); \
+				bump the version in `sc-executor-wasmtime`'s `Cargo.toml' to '{1}' and try again",
+			our_rustix.req, wasmtime_rustix.req,
+		);
 	}
 }