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/allocator/src/freeing_bump.rs

@@ -314,27 +314,50 @@ impl IndexMut<Order> for FreeLists {
 	}
 }
 
+/// Memory allocation stats gathered during the lifetime of the allocator.
+#[derive(Clone, Debug, Default)]
+#[non_exhaustive]
+pub struct AllocationStats {
+	/// The current number of bytes allocated.
+	///
+	/// This represents how many bytes are allocated *right now*.
+	pub bytes_allocated: u32,
+
+	/// The peak number of bytes ever allocated.
+	///
+	/// This is the maximum the `bytes_allocated` ever reached.
+	pub bytes_allocated_peak: u32,
+
+	/// The sum of every allocation ever made.
+	///
+	/// This increases every time a new allocation is made.
+	pub bytes_allocated_sum: u32,
+
+	/// The amount of address space (in bytes) used by the allocator.
+	///
+	/// This is calculated as the difference between the allocator's bumper
+	/// and the heap base.
+	///
+	/// Currently the bumper's only ever incremented, so this is simultaneously
+	/// the current value as well as the peak value.
+	pub address_space_used: u32,
+}
+
 /// An implementation of freeing bump allocator.
 ///
 /// Refer to the module-level documentation for further details.
 pub struct FreeingBumpHeapAllocator {
+	original_heap_base: u32,
 	bumper: u32,
 	free_lists: FreeLists,
-	total_size: u32,
 	poisoned: bool,
-	max_total_size: u32,
-	max_bumper: u32,
 	last_observed_memory_size: u32,
+	stats: AllocationStats,
 }
 
 impl Drop for FreeingBumpHeapAllocator {
 	fn drop(&mut self) {
-		log::debug!(
-			target: LOG_TARGET,
-			"allocator being destroyed, max_total_size {}, max_bumper {}",
-			self.max_total_size,
-			self.max_bumper,
-		)
+		log::debug!(target: LOG_TARGET, "allocator dropped: {:?}", self.stats)
 	}
 }
 
@@ -348,13 +371,12 @@ impl FreeingBumpHeapAllocator {
 		let aligned_heap_base = (heap_base + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT;
 
 		FreeingBumpHeapAllocator {
+			original_heap_base: aligned_heap_base,
 			bumper: aligned_heap_base,
 			free_lists: FreeLists::new(),
-			total_size: 0,
 			poisoned: false,
-			max_total_size: 0,
-			max_bumper: aligned_heap_base,
 			last_observed_memory_size: 0,
+			stats: AllocationStats::default(),
 		}
 	}
 
@@ -412,22 +434,13 @@ impl FreeingBumpHeapAllocator {
 		// Write the order in the occupied header.
 		Header::Occupied(order).write_into(mem, header_ptr)?;
 
-		self.total_size += order.size() + HEADER_SIZE;
+		self.stats.bytes_allocated += order.size() + HEADER_SIZE;
+		self.stats.bytes_allocated_sum += order.size() + HEADER_SIZE;
+		self.stats.bytes_allocated_peak =
+			std::cmp::max(self.stats.bytes_allocated_peak, self.stats.bytes_allocated);
+		self.stats.address_space_used = self.bumper - self.original_heap_base;
 
-		log::trace!(
-			target: LOG_TARGET,
-			"after allocation, total_size = {}, bumper = {}.",
-			self.total_size,
-			self.bumper,
-		);
-
-		// update trackers if needed.
-		if self.total_size > self.max_total_size {
-			self.max_total_size = self.total_size;
-		}
-		if self.bumper > self.max_bumper {
-			self.max_bumper = self.bumper;
-		}
+		log::trace!(target: LOG_TARGET, "after allocation: {:?}", self.stats);
 
 		bomb.disarm();
 		Ok(Pointer::new(header_ptr + HEADER_SIZE))
@@ -469,21 +482,23 @@ impl FreeingBumpHeapAllocator {
 		let prev_head = self.free_lists.replace(order, Link::Ptr(header_ptr));
 		Header::Free(prev_head).write_into(mem, header_ptr)?;
 
-		// Do the total_size book keeping.
-		self.total_size = self
-			.total_size
+		self.stats.bytes_allocated = self
+			.stats
+			.bytes_allocated
 			.checked_sub(order.size() + HEADER_SIZE)
-			.ok_or_else(|| error("Unable to subtract from total heap size without overflow"))?;
-		log::trace!(
-			"after deallocation, total_size = {}, bumper = {}.",
-			self.total_size,
-			self.bumper,
-		);
+			.ok_or_else(|| error("underflow of the currently allocated bytes count"))?;
+
+		log::trace!("after deallocation: {:?}", self.stats);
 
 		bomb.disarm();
 		Ok(())
 	}
 
+	/// Returns the allocation stats for this allocator.
+	pub fn stats(&self) -> AllocationStats {
+		self.stats.clone()
+	}
+
 	/// Increases the `bumper` by `size`.
 	///
 	/// Returns the `bumper` from before the increase. Returns an `Error::AllocatorOutOfSpace` if
@@ -791,13 +806,13 @@ mod tests {
 		let ptr1 = heap.allocate(&mut mem[..], 32).unwrap();
 		assert_eq!(ptr1, to_pointer(HEADER_SIZE));
 		heap.deallocate(&mut mem[..], ptr1).expect("failed freeing ptr1");
-		assert_eq!(heap.total_size, 0);
+		assert_eq!(heap.stats.bytes_allocated, 0);
 		assert_eq!(heap.bumper, 40);
 
 		let ptr2 = heap.allocate(&mut mem[..], 16).unwrap();
 		assert_eq!(ptr2, to_pointer(48));
 		heap.deallocate(&mut mem[..], ptr2).expect("failed freeing ptr2");
-		assert_eq!(heap.total_size, 0);
+		assert_eq!(heap.stats.bytes_allocated, 0);
 		assert_eq!(heap.bumper, 64);
 
 		// when
@@ -825,7 +840,7 @@ mod tests {
 		heap.allocate(&mut mem[..], 9).unwrap();
 
 		// then
-		assert_eq!(heap.total_size, HEADER_SIZE + 16);
+		assert_eq!(heap.stats.bytes_allocated, HEADER_SIZE + 16);
 	}
 
 	#[test]
@@ -840,7 +855,7 @@ mod tests {
 		heap.deallocate(&mut mem[..], ptr).unwrap();
 
 		// then
-		assert_eq!(heap.total_size, 0);
+		assert_eq!(heap.stats.bytes_allocated, 0);
 	}
 
 	#[test]
@@ -856,7 +871,7 @@ mod tests {
 		}
 
 		// then
-		assert_eq!(heap.total_size, 0);
+		assert_eq!(heap.stats.bytes_allocated, 0);
 	}
 
 	#[test]

substrate/client/allocator/src/lib.rs

@@ -26,4 +26,4 @@ mod error;
 mod freeing_bump;
 
 pub use error::Error;
-pub use freeing_bump::FreeingBumpHeapAllocator;
+pub use freeing_bump::{AllocationStats, FreeingBumpHeapAllocator};

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