feat: add new allocator design for runtime

Cargo.toml

@@ -412,6 +412,7 @@ members = [
 	"substrate/frame/utility",
 	"substrate/frame/vesting",
 	"substrate/frame/whitelist",
+	"substrate/primitives/allocator-v1",
 	"substrate/primitives/api",
 	"substrate/primitives/api/proc-macro",
 	"substrate/primitives/api/test",

substrate/bin/node-template/runtime/Cargo.toml

@@ -29,6 +29,7 @@ frame-try-runtime = { path = "../../../frame/try-runtime", default-features = fa
 pallet-timestamp = { path = "../../../frame/timestamp", default-features = false }
 pallet-transaction-payment = { path = "../../../frame/transaction-payment", default-features = false }
 frame-executive = { path = "../../../frame/executive", default-features = false }
+sp-allocator-v1 = { default-features = false, path = "../../../primitives/allocator-v1" }
 sp-api = { path = "../../../primitives/api", default-features = false }
 sp-block-builder = { path = "../../../primitives/block-builder", default-features = false }
 sp-consensus-aura = { path = "../../../primitives/consensus/aura", default-features = false, features = ["serde"] }
@@ -60,7 +61,10 @@ pallet-template = { path = "../pallets/template", default-features = false }
 substrate-wasm-builder = { path = "../../../utils/wasm-builder", optional = true }
 
 [features]
-default = ["std"]
+default = [ "std" ]
+allocator-v1 = [
+	"sp-allocator-v1/allocator-v1",
+]
 std = [
 	"codec/std",
 	"frame-benchmarking?/std",
@@ -79,6 +83,7 @@ std = [
 	"pallet-transaction-payment-rpc-runtime-api/std",
 	"pallet-transaction-payment/std",
 	"scale-info/std",
+	"sp-allocator-v1/std",
 	"serde_json/std",
 	"sp-api/std",
 	"sp-block-builder/std",

substrate/bin/node-template/runtime/build.rs

@@ -5,6 +5,7 @@ fn main() {
 			.with_current_project()
 			.export_heap_base()
 			.import_memory()
+			.enable_feature("allocator-v1")
 			.build();
 	}
 }

substrate/bin/node-template/runtime/src/lib.rs

@@ -6,6 +6,8 @@
 #[cfg(feature = "std")]
 include!(concat!(env!("OUT_DIR"), "/wasm_binary.rs"));
 
+use sp_allocator_v1 as _;
+
 use pallet_grandpa::AuthorityId as GrandpaId;
 use sp_api::impl_runtime_apis;
 use sp_consensus_aura::sr25519::AuthorityId as AuraId;

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

@@ -52,6 +52,14 @@ impl RuntimeBlob {
 		Ok(Self { raw_module })
 	}
 
+	/// Return true if wasm contains a export function.
+	pub fn contain_export_func(&self, func: &str) -> bool {
+		self.raw_module
+			.export_section()
+			.map(|section| section.entries().iter().any(|entry| entry.field() == func))
+			.unwrap_or(false)
+	}
+
 	/// The number of globals defined in locally in this module.
 	pub fn declared_globals_count(&self) -> u32 {
 		self.raw_module
@@ -160,6 +168,8 @@ impl RuntimeBlob {
 		&mut self,
 		heap_alloc_strategy: HeapAllocStrategy,
 	) -> Result<(), WasmError> {
+		let is_allocator_v1 = self.contain_export_func("v1");
+
 		let memory_section = self
 			.raw_module
 			.memory_section_mut()
@@ -168,6 +178,7 @@ impl RuntimeBlob {
 		if memory_section.entries().is_empty() {
 			return Err(WasmError::Other("memory section is empty".into()))
 		}
+
 		for memory_ty in memory_section.entries_mut() {
 			let initial = memory_ty.limits().initial();
 			let (min, max) = match heap_alloc_strategy {
@@ -177,7 +188,14 @@ impl RuntimeBlob {
 				},
 				HeapAllocStrategy::Static { extra_pages } => {
 					let pages = initial.saturating_add(extra_pages);
-					(pages, Some(pages))
+					// The runtime-customized allocator may rely on this init memory page,
+					// so this value cannot be modified at will, otherwise it will cause errors in
+					// the allocator logic.
+					if is_allocator_v1 {
+						(initial, Some(pages))
+					} else {
+						(pages, Some(pages))
+					}
 				},
 			};
 			*memory_ty = MemoryType::new(min, max);

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

@@ -19,7 +19,7 @@
 //! This module defines `HostState` and `HostContext` structs which provide logic and state
 //! required for execution of host.
 
-use wasmtime::Caller;
+use wasmtime::{AsContextMut, Caller, Val};
 
 use sc_allocator::{AllocationStats, FreeingBumpHeapAllocator};
 use sp_wasm_interface::{Pointer, WordSize};
@@ -41,8 +41,8 @@ pub struct HostState {
 
 impl HostState {
 	/// Constructs a new `HostState`.
-	pub fn new(allocator: FreeingBumpHeapAllocator) -> Self {
-		HostState { allocator: Some(allocator), panic_message: None }
+	pub fn new(allocator: Option<FreeingBumpHeapAllocator>) -> Self {
+		HostState { allocator, panic_message: None }
 	}
 
 	/// Takes the error message out of the host state, leaving a `None` in its place.
@@ -88,38 +88,56 @@ 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 mut allocator = self
-			.host_state_mut()
-			.allocator
-			.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
+		if let Some(alloc) = self.caller.data().alloc {
+			let params = [Val::I32(size as i32)];
+			let mut results = [Val::I32(0)];
+
+			alloc
+				.call(self.caller.as_context_mut(), &params, &mut results)
+				.expect("alloc must success; qed");
+			let data_ptr = results[0].i32().unwrap();
+			let data_ptr = Pointer::new(data_ptr as u32);
+
+			Ok(data_ptr)
+		} else {
+			let mut allocator = self
+				.host_state_mut()
+				.allocator
+				.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 mut allocator = self
-			.host_state_mut()
-			.allocator
-			.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
+		if let Some(_alloc) = self.caller.data().alloc {
+			// TODO: maybe we need to support it.
+			unreachable!("The `deallocate_memory` never be used in allocator v1");
+		} else {
+			let mut allocator = self
+				.host_state_mut()
+				.allocator
+				.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

@@ -28,7 +28,7 @@ use sc_executor_common::{
 };
 use sp_wasm_interface::{Pointer, Value, WordSize};
 use wasmtime::{
-	AsContext, AsContextMut, Engine, Extern, Instance, InstancePre, Memory, Table, Val,
+	AsContext, AsContextMut, Engine, Extern, Func, Instance, InstancePre, Memory, Table, Val,
 };
 
 /// Invoked entrypoint format.
@@ -179,9 +179,11 @@ impl InstanceWrapper {
 
 		let memory = get_linear_memory(&instance, &mut store)?;
 		let table = get_table(&instance, &mut store);
+		let alloc = get_export_alloc(&instance, &mut store);
 
 		store.data_mut().memory = Some(memory);
 		store.data_mut().table = table;
+		store.data_mut().alloc = alloc;
 
 		Ok(InstanceWrapper { instance, store, _release_instance_handle })
 	}
@@ -310,6 +312,11 @@ fn get_table(instance: &Instance, ctx: &mut Store) -> Option<Table> {
 		.and_then(Extern::into_table)
 }
 
+/// Extract export `alloc` func from the given instance.
+fn get_export_alloc(instance: &Instance, ctx: impl AsContextMut) -> Option<Func> {
+	instance.get_export(ctx, "alloc").and_then(Extern::into_func)
+}
+
 /// Functions related to memory.
 impl InstanceWrapper {
 	pub(crate) fn store(&self) -> &Store {

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

@@ -41,7 +41,7 @@ use std::{
 		Arc,
 	},
 };
-use wasmtime::{AsContext, Engine, Memory, Table};
+use wasmtime::{AsContext, Engine, Func, Memory, Table, Val};
 
 const MAX_INSTANCE_COUNT: u32 = 64;
 
@@ -53,6 +53,8 @@ pub(crate) struct StoreData {
 	pub(crate) memory: Option<Memory>,
 	/// This will be set only if the runtime actually contains a table.
 	pub(crate) table: Option<Table>,
+	/// This will be set only if the runtime actually contains a table.
+	pub(crate) alloc: Option<Func>,
 }
 
 impl StoreData {
@@ -171,11 +173,21 @@ impl WasmtimeInstance {
 		match &mut self.strategy {
 			Strategy::RecreateInstance(ref mut instance_creator) => {
 				let mut instance_wrapper = instance_creator.instantiate()?;
-				let heap_base = instance_wrapper.extract_heap_base()?;
 				let entrypoint = instance_wrapper.resolve_entrypoint(method)?;
-				let allocator = FreeingBumpHeapAllocator::new(heap_base);
 
-				perform_call(data, &mut instance_wrapper, entrypoint, allocator, allocation_stats)
+				if let Some(alloc) = instance_wrapper.store().data().alloc {
+					perform_call_v1(data, &mut instance_wrapper, entrypoint, alloc)
+				} else {
+					let heap_base = instance_wrapper.extract_heap_base()?;
+					let allocator = FreeingBumpHeapAllocator::new(heap_base);
+					perform_call(
+						data,
+						&mut instance_wrapper,
+						entrypoint,
+						allocator,
+						allocation_stats,
+					)
+				}
 			},
 		}
 	}
@@ -694,7 +706,7 @@ fn perform_call(
 ) -> Result<Vec<u8>> {
 	let (data_ptr, data_len) = inject_input_data(instance_wrapper, &mut allocator, data)?;
 
-	let host_state = HostState::new(allocator);
+	let host_state = HostState::new(Some(allocator));
 
 	// Set the host state before calling into wasm.
 	instance_wrapper.store_mut().data_mut().host_state = Some(host_state);
@@ -715,6 +727,29 @@ fn perform_call(
 	Ok(output)
 }
 
+fn perform_call_v1(
+	data: &[u8],
+	instance_wrapper: &mut InstanceWrapper,
+	entrypoint: EntryPoint,
+	alloc: Func,
+) -> Result<Vec<u8>> {
+	let (data_ptr, data_len) = inject_input_data_v1(instance_wrapper, alloc, data)?;
+
+	let host_state = HostState::new(None);
+
+	// Set the host state before calling into wasm.
+	instance_wrapper.store_mut().data_mut().host_state = Some(host_state);
+
+	let ret = entrypoint
+		.call(instance_wrapper.store_mut(), data_ptr, data_len)
+		.map(unpack_ptr_and_len);
+
+	let (output_ptr, output_len) = ret?;
+	let output = extract_output_data(instance_wrapper, output_ptr, output_len)?;
+
+	Ok(output)
+}
+
 fn inject_input_data(
 	instance: &mut InstanceWrapper,
 	allocator: &mut FreeingBumpHeapAllocator,
@@ -728,6 +763,25 @@ fn inject_input_data(
 	Ok((data_ptr, data_len))
 }
 
+fn inject_input_data_v1(
+	instance: &mut InstanceWrapper,
+	alloc: Func,
+	data: &[u8],
+) -> Result<(Pointer<u8>, WordSize)> {
+	let data_len = data.len() as WordSize;
+	let params = [Val::I32(data_len as _)];
+	let mut results = [Val::I32(0)];
+
+	alloc
+		.call(instance.store_mut(), &params, &mut results)
+		.expect("alloc must success; qed");
+	let data_ptr = results[0].i32().unwrap();
+	let data_ptr = Pointer::new(data_ptr as u32);
+	util::write_memory_from(instance.store_mut(), data_ptr, data)?;
+
+	Ok((data_ptr, data_len))
+}
+
 fn extract_output_data(
 	instance: &InstanceWrapper,
 	output_ptr: u32,