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

Cargo.toml

@@ -303,3 +303,6 @@ inherits = "release"
 lto = "fat"
 # https://doc.rust-lang.org/rustc/codegen-options/index.html#codegen-units
 codegen-units = 1
+
+[patch.crates-io]
+parity-scale-codec = { git = "https://github.com/koute/parity-scale-codec.git", branch = "master_bytes_deserialization" }

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;
-
-		log::trace!(
-			target: LOG_TARGET,
-			"after allocation, total_size = {}, bumper = {}.",
-			self.total_size,
-			self.bumper,
-		);
+		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;
 
-		// 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]

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};

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.
 	///

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)
 	}