Prevent double allocation of the payload when calling sp_io::storage::get (#11523)

* Expose allocation stats in `FreeingBumpHeapAllocator`

* Return allocation stats when calling into the runtime

* Bump `parity-scale-codec` to 3.1.3 (fork)

* Prevent double allocation of the payload when calling `sp_io::storage::get`

* Fix tests

* Remove unnecessary `mut`

* Enable the `bytes` feature for `parity-scale-codec` in `sp-runtime-interface`

* Update client/allocator/src/freeing_bump.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Bump `parity-scale-codec` to 3.1.3

* Fix some of the UI tests

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

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

@@ -21,6 +21,8 @@
 use crate::error::Error;
 use sp_wasm_interface::Value;
 
+pub use sc_allocator::AllocationStats;
+
 /// A method to be used to find the entrypoint when calling into the runtime
 ///
 /// Contains variants on how to resolve wasm function that will be invoked.
@@ -78,7 +80,20 @@ pub trait WasmInstance: Send {
 	/// Before execution, instance is reset.
 	///
 	/// Returns the encoded result on success.
-	fn call(&mut self, method: InvokeMethod, data: &[u8]) -> Result<Vec<u8>, Error>;
+	fn call(&mut self, method: InvokeMethod, data: &[u8]) -> Result<Vec<u8>, Error> {
+		self.call_with_allocation_stats(method, data).0
+	}
+
+	/// Call a method on this WASM instance.
+	///
+	/// Before execution, instance is reset.
+	///
+	/// Returns the encoded result on success.
+	fn call_with_allocation_stats(
+		&mut self,
+		method: InvokeMethod,
+		data: &[u8],
+	) -> (Result<Vec<u8>, Error>, Option<AllocationStats>);
 
 	/// Call an exported method on this WASM instance.
 	///

substrate/client/executor/src/native_executor.rs

@@ -37,7 +37,7 @@ use std::{
 use codec::{Decode, Encode};
 use sc_executor_common::{
 	runtime_blob::RuntimeBlob,
-	wasm_runtime::{InvokeMethod, WasmInstance, WasmModule},
+	wasm_runtime::{AllocationStats, InvokeMethod, WasmInstance, WasmModule},
 };
 use sp_core::{
 	traits::{CodeExecutor, Externalities, RuntimeCode, RuntimeSpawn, RuntimeSpawnExt},
@@ -219,6 +219,47 @@ where
 		allow_missing_host_functions: bool,
 		export_name: &str,
 		call_data: &[u8],
+	) -> std::result::Result<Vec<u8>, Error> {
+		self.uncached_call_impl(
+			runtime_blob,
+			ext,
+			allow_missing_host_functions,
+			export_name,
+			call_data,
+			&mut None,
+		)
+	}
+
+	/// Same as `uncached_call`, except it also returns allocation statistics.
+	#[doc(hidden)] // We use this function in tests.
+	pub fn uncached_call_with_allocation_stats(
+		&self,
+		runtime_blob: RuntimeBlob,
+		ext: &mut dyn Externalities,
+		allow_missing_host_functions: bool,
+		export_name: &str,
+		call_data: &[u8],
+	) -> (std::result::Result<Vec<u8>, Error>, Option<AllocationStats>) {
+		let mut allocation_stats = None;
+		let result = self.uncached_call_impl(
+			runtime_blob,
+			ext,
+			allow_missing_host_functions,
+			export_name,
+			call_data,
+			&mut allocation_stats,
+		);
+		(result, allocation_stats)
+	}
+
+	fn uncached_call_impl(
+		&self,
+		runtime_blob: RuntimeBlob,
+		ext: &mut dyn Externalities,
+		allow_missing_host_functions: bool,
+		export_name: &str,
+		call_data: &[u8],
+		allocation_stats_out: &mut Option<AllocationStats>,
 	) -> std::result::Result<Vec<u8>, Error> {
 		let module = crate::wasm_runtime::create_wasm_runtime_with_code::<H>(
 			self.method,
@@ -235,10 +276,14 @@ where
 		let mut instance = AssertUnwindSafe(instance);
 		let mut ext = AssertUnwindSafe(ext);
 		let module = AssertUnwindSafe(module);
+		let mut allocation_stats_out = AssertUnwindSafe(allocation_stats_out);
 
 		with_externalities_safe(&mut **ext, move || {
 			preregister_builtin_ext(module.clone());
-			instance.call_export(export_name, call_data)
+			let (result, allocation_stats) =
+				instance.call_with_allocation_stats(export_name.into(), call_data);
+			**allocation_stats_out = allocation_stats;
+			result
 		})
 		.and_then(|r| r)
 	}

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

@@ -29,6 +29,7 @@ use wasmi::{
 };
 
 use codec::{Decode, Encode};
+use sc_allocator::AllocationStats;
 use sc_executor_common::{
 	error::{Error, MessageWithBacktrace, WasmError},
 	runtime_blob::{DataSegmentsSnapshot, RuntimeBlob},
@@ -489,6 +490,7 @@ fn call_in_wasm_module(
 	host_functions: Arc<Vec<&'static dyn Function>>,
 	allow_missing_func_imports: bool,
 	missing_functions: Arc<Vec<String>>,
+	allocation_stats: &mut Option<AllocationStats>,
 ) -> Result<Vec<u8>, Error> {
 	// Initialize FunctionExecutor.
 	let table: Option<TableRef> = module_instance
@@ -561,6 +563,8 @@ fn call_in_wasm_module(
 		},
 	};
 
+	*allocation_stats = Some(function_executor.heap.borrow().stats());
+
 	match result {
 		Ok(Some(I64(r))) => {
 			let (ptr, length) = unpack_ptr_and_len(r as u64);
@@ -761,8 +765,13 @@ pub struct WasmiInstance {
 // `self.instance`
 unsafe impl Send for WasmiInstance {}
 
-impl WasmInstance for WasmiInstance {
-	fn call(&mut self, method: InvokeMethod, data: &[u8]) -> Result<Vec<u8>, Error> {
+impl WasmiInstance {
+	fn call_impl(
+		&mut self,
+		method: InvokeMethod,
+		data: &[u8],
+		allocation_stats: &mut Option<AllocationStats>,
+	) -> Result<Vec<u8>, Error> {
 		// 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.
 
@@ -790,8 +799,21 @@ impl WasmInstance for WasmiInstance {
 			self.host_functions.clone(),
 			self.allow_missing_func_imports,
 			self.missing_functions.clone(),
+			allocation_stats,
 		)
 	}
+}
+
+impl WasmInstance for WasmiInstance {
+	fn call_with_allocation_stats(
+		&mut self,
+		method: InvokeMethod,
+		data: &[u8],
+	) -> (Result<Vec<u8>, Error>, Option<AllocationStats>) {
+		let mut allocation_stats = None;
+		let result = self.call_impl(method, data, &mut allocation_stats);
+		(result, allocation_stats)
+	}
 
 	fn get_global_const(&mut self, name: &str) -> Result<Option<sp_wasm_interface::Value>, Error> {
 		match self.instance.export_by_name(name) {

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

@@ -23,7 +23,7 @@ use log::trace;
 use wasmtime::{Caller, Func, Val};
 
 use codec::{Decode, Encode};
-use sc_allocator::FreeingBumpHeapAllocator;
+use sc_allocator::{AllocationStats, FreeingBumpHeapAllocator};
 use sc_executor_common::{
 	error::Result,
 	sandbox::{self, SupervisorFuncIndex},
@@ -66,6 +66,10 @@ impl HostState {
 	pub fn take_panic_message(&mut self) -> Option<String> {
 		self.panic_message.take()
 	}
+
+	pub(crate) fn allocation_stats(&self) -> AllocationStats {
+		self.allocator.stats()
+	}
 }
 
 /// A `HostContext` implements `FunctionContext` for making host calls from a Wasmtime

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

@@ -24,7 +24,7 @@ use crate::{
 	util::{self, replace_strategy_if_broken},
 };
 
-use sc_allocator::FreeingBumpHeapAllocator;
+use sc_allocator::{AllocationStats, FreeingBumpHeapAllocator};
 use sc_executor_common::{
 	error::{Result, WasmError},
 	runtime_blob::{
@@ -184,8 +184,13 @@ pub struct WasmtimeInstance {
 	strategy: Strategy,
 }
 
-impl WasmInstance for WasmtimeInstance {
-	fn call(&mut self, method: InvokeMethod, data: &[u8]) -> Result<Vec<u8>> {
+impl WasmtimeInstance {
+	fn call_impl(
+		&mut self,
+		method: InvokeMethod,
+		data: &[u8],
+		allocation_stats: &mut Option<AllocationStats>,
+	) -> Result<Vec<u8>> {
 		match &mut self.strategy {
 			Strategy::LegacyInstanceReuse {
 				ref mut instance_wrapper,
@@ -205,7 +210,8 @@ impl WasmInstance for WasmtimeInstance {
 				globals_snapshot.apply(&mut InstanceGlobals { instance: instance_wrapper });
 				let allocator = FreeingBumpHeapAllocator::new(*heap_base);
 
-				let result = perform_call(data, instance_wrapper, entrypoint, allocator);
+				let result =
+					perform_call(data, instance_wrapper, entrypoint, allocator, allocation_stats);
 
 				// Signal to the OS that we are done with the linear memory and that it can be
 				// reclaimed.
@@ -219,10 +225,22 @@ impl WasmInstance for WasmtimeInstance {
 				let entrypoint = instance_wrapper.resolve_entrypoint(method)?;
 
 				let allocator = FreeingBumpHeapAllocator::new(heap_base);
-				perform_call(data, &mut instance_wrapper, entrypoint, allocator)
+				perform_call(data, &mut instance_wrapper, entrypoint, allocator, allocation_stats)
 			},
 		}
 	}
+}
+
+impl WasmInstance for WasmtimeInstance {
+	fn call_with_allocation_stats(
+		&mut self,
+		method: InvokeMethod,
+		data: &[u8],
+	) -> (Result<Vec<u8>>, Option<AllocationStats>) {
+		let mut allocation_stats = None;
+		let result = self.call_impl(method, data, &mut allocation_stats);
+		(result, allocation_stats)
+	}
 
 	fn get_global_const(&mut self, name: &str) -> Result<Option<Value>> {
 		match &mut self.strategy {
@@ -741,6 +759,7 @@ fn perform_call(
 	instance_wrapper: &mut InstanceWrapper,
 	entrypoint: EntryPoint,
 	mut allocator: FreeingBumpHeapAllocator,
+	allocation_stats: &mut Option<AllocationStats>,
 ) -> Result<Vec<u8>> {
 	let (data_ptr, data_len) = inject_input_data(instance_wrapper, &mut allocator, data)?;
 
@@ -754,7 +773,10 @@ fn perform_call(
 		.map(unpack_ptr_and_len);
 
 	// Reset the host state
-	instance_wrapper.store_mut().data_mut().host_state = None;
+	let host_state = instance_wrapper.store_mut().data_mut().host_state.take().expect(
+		"the host state is always set before calling into WASM so it can't be None here; qed",
+	);
+	*allocation_stats = Some(host_state.allocation_stats());
 
 	let (output_ptr, output_len) = ret?;
 	let output = extract_output_data(instance_wrapper, output_ptr, output_len)?;