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/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)?;