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-04-30 11:57:56 +00:00 · 2023-02-24 12:43:01 +01:00
parent c848d40775
commit 941288c6d0
37 changed files with 1092 additions and 667 deletions
@@ -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
@@ -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,
@@ -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() }
 	}
 }

@@ -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,