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

@@ -6516,6 +6516,7 @@ dependencies = [
  "arrayvec 0.7.2",
  "bitvec",
  "byte-slice-cast",
+ "bytes",
  "impl-trait-for-tuples",
  "parity-scale-codec-derive",
  "serde",
@@ -9731,6 +9732,7 @@ dependencies = [
 name = "sp-io"
 version = "6.0.0"
 dependencies = [
+ "bytes",
  "futures",
  "hash-db",
  "libsecp256k1",
@@ -9889,6 +9891,7 @@ dependencies = [
 name = "sp-runtime-interface"
 version = "6.0.0"
 dependencies = [
+ "bytes",
  "impl-trait-for-tuples",
  "parity-scale-codec",
  "primitive-types",
@@ -9938,6 +9941,7 @@ dependencies = [
 name = "sp-runtime-interface-test-wasm"
 version = "2.0.0"
 dependencies = [
+ "bytes",
  "sp-core",
  "sp-io",
  "sp-runtime-interface",

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

substrate/frame/support/src/storage/unhashed.rs

@@ -155,7 +155,7 @@ pub fn clear_prefix(
 
 /// Get a Vec of bytes from storage.
 pub fn get_raw(key: &[u8]) -> Option<Vec<u8>> {
-	sp_io::storage::get(key)
+	sp_io::storage::get(key).map(|value| value.to_vec())
 }
 
 /// Put a raw byte slice into storage.

substrate/frame/support/test/tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen.stderr

@@ -28,7 +28,7 @@ error[E0277]: the trait bound `Bar: EncodeLike` is not satisfied
              <&[(T,)] as EncodeLike<BinaryHeap<LikeT>>>
              <&[(T,)] as EncodeLike<LinkedList<LikeT>>>
              <&[T] as EncodeLike<Vec<U>>>
-           and 273 others
+           and 278 others
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`
    = note: required because of the requirements on the impl of `PartialStorageInfoTrait` for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>`
@@ -47,8 +47,8 @@ error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
              Cow<'a, T>
              Rc<T>
              Vec<T>
-             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
-           and 2 others
+             bytes::bytes::Bytes
+           and 3 others
    = note: required because of the requirements on the impl of `Encode` for `Bar`
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`
@@ -103,7 +103,7 @@ error[E0277]: the trait bound `Bar: EncodeLike` is not satisfied
              <&[(T,)] as EncodeLike<BinaryHeap<LikeT>>>
              <&[(T,)] as EncodeLike<LinkedList<LikeT>>>
              <&[T] as EncodeLike<Vec<U>>>
-           and 273 others
+           and 278 others
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`
    = note: required because of the requirements on the impl of `StorageEntryMetadataBuilder` for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>`
@@ -122,8 +122,8 @@ error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
              Cow<'a, T>
              Rc<T>
              Vec<T>
-             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
-           and 2 others
+             bytes::bytes::Bytes
+           and 3 others
    = note: required because of the requirements on the impl of `Encode` for `Bar`
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`

substrate/frame/support/test/tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen_unnamed.stderr

@@ -28,7 +28,7 @@ error[E0277]: the trait bound `Bar: EncodeLike` is not satisfied
              <&[(T,)] as EncodeLike<BinaryHeap<LikeT>>>
              <&[(T,)] as EncodeLike<LinkedList<LikeT>>>
              <&[T] as EncodeLike<Vec<U>>>
-           and 273 others
+           and 278 others
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`
    = note: required because of the requirements on the impl of `PartialStorageInfoTrait` for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>`
@@ -47,8 +47,8 @@ error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
              Cow<'a, T>
              Rc<T>
              Vec<T>
-             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
-           and 2 others
+             bytes::bytes::Bytes
+           and 3 others
    = note: required because of the requirements on the impl of `Encode` for `Bar`
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`
@@ -103,7 +103,7 @@ error[E0277]: the trait bound `Bar: EncodeLike` is not satisfied
              <&[(T,)] as EncodeLike<BinaryHeap<LikeT>>>
              <&[(T,)] as EncodeLike<LinkedList<LikeT>>>
              <&[T] as EncodeLike<Vec<U>>>
-           and 273 others
+           and 278 others
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`
    = note: required because of the requirements on the impl of `StorageEntryMetadataBuilder` for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>`
@@ -122,8 +122,8 @@ error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
              Cow<'a, T>
              Rc<T>
              Vec<T>
-             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
-           and 2 others
+             bytes::bytes::Bytes
+           and 3 others
    = note: required because of the requirements on the impl of `Encode` for `Bar`
    = note: required because of the requirements on the impl of `FullEncode` for `Bar`
    = note: required because of the requirements on the impl of `FullCodec` for `Bar`

substrate/primitives/io/Cargo.toml

@@ -15,7 +15,8 @@ targets = ["x86_64-unknown-linux-gnu"]
 
 
 [dependencies]
-codec = { package = "parity-scale-codec", version = "3.0.0", default-features = false }
+bytes = { version = "1.1.0", default-features = false }
+codec = { package = "parity-scale-codec", version = "3.1.3", default-features = false, features = ["bytes"] }
 hash-db = { version = "0.15.2", default-features = false }
 sp-core = { version = "6.0.0", default-features = false, path = "../core" }
 sp-keystore = { version = "0.12.0", default-features = false, optional = true, path = "../keystore" }

substrate/primitives/io/src/lib.rs

@@ -125,8 +125,8 @@ impl From<MultiRemovalResults> for KillStorageResult {
 #[runtime_interface]
 pub trait Storage {
 	/// Returns the data for `key` in the storage or `None` if the key can not be found.
-	fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
-		self.storage(key).map(|s| s.to_vec())
+	fn get(&self, key: &[u8]) -> Option<bytes::Bytes> {
+		self.storage(key).map(|s| bytes::Bytes::from(s.to_vec()))
 	}
 
 	/// Get `key` from storage, placing the value into `value_out` and return the number of
@@ -1787,7 +1787,7 @@ mod tests {
 		t.execute_with(|| {
 			assert_eq!(storage::get(b"hello"), None);
 			storage::set(b"hello", b"world");
-			assert_eq!(storage::get(b"hello"), Some(b"world".to_vec()));
+			assert_eq!(storage::get(b"hello"), Some(b"world".to_vec().into()));
 			assert_eq!(storage::get(b"foo"), None);
 			storage::set(b"foo", &[1, 2, 3][..]);
 		});
@@ -1799,7 +1799,7 @@ mod tests {
 
 		t.execute_with(|| {
 			assert_eq!(storage::get(b"hello"), None);
-			assert_eq!(storage::get(b"foo"), Some(b"bar".to_vec()));
+			assert_eq!(storage::get(b"foo"), Some(b"bar".to_vec().into()));
 		});
 
 		let value = vec![7u8; 35];
@@ -1809,7 +1809,7 @@ mod tests {
 
 		t.execute_with(|| {
 			assert_eq!(storage::get(b"hello"), None);
-			assert_eq!(storage::get(b"foo00"), Some(value.clone()));
+			assert_eq!(storage::get(b"foo00"), Some(value.clone().into()));
 		});
 	}
 

substrate/primitives/runtime-interface/Cargo.toml

@@ -14,12 +14,13 @@ readme = "README.md"
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
+bytes = { version = "1.1.0", default-features = false }
 sp-wasm-interface = { version = "6.0.0", path = "../wasm-interface", default-features = false }
 sp-std = { version = "4.0.0", default-features = false, path = "../std" }
 sp-tracing = { version = "5.0.0", default-features = false, path = "../tracing" }
 sp-runtime-interface-proc-macro = { version = "5.0.0", path = "proc-macro" }
 sp-externalities = { version = "0.12.0", default-features = false, path = "../externalities" }
-codec = { package = "parity-scale-codec", version = "3.0.0", default-features = false }
+codec = { package = "parity-scale-codec", version = "3.0.0", default-features = false, features = ["bytes"] }
 static_assertions = "1.0.0"
 primitive-types = { version = "0.11.1", default-features = false }
 sp-storage = { version = "6.0.0", default-features = false, path = "../storage" }

substrate/primitives/runtime-interface/src/pass_by.rs

@@ -248,12 +248,12 @@ impl<T: codec::Codec> PassByImpl<T> for Codec<T> {
 		let len = len as usize;
 
 		let encoded = if len == 0 {
-			Vec::new()
+			bytes::Bytes::new()
 		} else {
-			unsafe { Vec::from_raw_parts(ptr as *mut u8, len, len) }
+			bytes::Bytes::from(unsafe { Vec::from_raw_parts(ptr as *mut u8, len, len) })
 		};
 
-		T::decode(&mut &encoded[..]).expect("Host to wasm values are encoded correctly; qed")
+		codec::decode_from_bytes(encoded).expect("Host to wasm values are encoded correctly; qed")
 	}
 }
 

substrate/primitives/runtime-interface/test-wasm/Cargo.toml

@@ -13,6 +13,7 @@ publish = false
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
+bytes = { version = "1.1.0", default-features = false }
 sp-core = { version = "6.0.0", default-features = false, path = "../../core" }
 sp-io = { version = "6.0.0", default-features = false, path = "../../io" }
 sp-runtime-interface = { version = "6.0.0", default-features = false, path = "../" }

substrate/primitives/runtime-interface/test-wasm/src/lib.rs

@@ -60,6 +60,18 @@ pub trait TestApi {
 		vec![0; 4 * 1024]
 	}
 
+	fn return_option_vec() -> Option<Vec<u8>> {
+		let mut vec = Vec::new();
+		vec.resize(16 * 1024, 0xAA);
+		Some(vec)
+	}
+
+	fn return_option_bytes() -> Option<bytes::Bytes> {
+		let mut vec = Vec::new();
+		vec.resize(16 * 1024, 0xAA);
+		Some(vec.into())
+	}
+
 	/// Set the storage at key with value.
 	fn set_storage(&mut self, key: &[u8], data: &[u8]) {
 		self.place_storage(key.to_vec(), Some(data.to_vec()));
@@ -300,4 +312,12 @@ wasm_export_functions! {
 		assert_eq!(c, res.2);
 		assert_eq!(d, res.3);
 	}
+
+	fn test_return_option_vec() {
+		test_api::return_option_vec();
+	}
+
+	fn test_return_option_bytes() {
+		test_api::return_option_bytes();
+	}
 }

substrate/primitives/runtime-interface/test/src/lib.rs

@@ -23,7 +23,7 @@ use sp_runtime_interface::*;
 use sp_runtime_interface_test_wasm::{test_api::HostFunctions, wasm_binary_unwrap};
 use sp_runtime_interface_test_wasm_deprecated::wasm_binary_unwrap as wasm_binary_deprecated_unwrap;
 
-use sc_executor_common::runtime_blob::RuntimeBlob;
+use sc_executor_common::{runtime_blob::RuntimeBlob, wasm_runtime::AllocationStats};
 use sp_wasm_interface::{ExtendedHostFunctions, HostFunctions as HostFunctionsT};
 
 use std::{
@@ -36,7 +36,7 @@ type TestExternalities = sp_state_machine::TestExternalities<sp_runtime::traits:
 fn call_wasm_method_with_result<HF: HostFunctionsT>(
 	binary: &[u8],
 	method: &str,
-) -> Result<TestExternalities, String> {
+) -> (Result<TestExternalities, String>, Option<AllocationStats>) {
 	let mut ext = TestExternalities::default();
 	let mut ext_ext = ext.ext();
 
@@ -44,20 +44,21 @@ fn call_wasm_method_with_result<HF: HostFunctionsT>(
 		ExtendedHostFunctions<sp_io::SubstrateHostFunctions, HF>,
 	>::new(sc_executor::WasmExecutionMethod::Interpreted, Some(8), 8, None, 2);
 
-	executor
-		.uncached_call(
-			RuntimeBlob::uncompress_if_needed(binary).expect("Failed to parse binary"),
-			&mut ext_ext,
-			false,
-			method,
-			&[],
-		)
-		.map_err(|e| format!("Failed to execute `{}`: {}", method, e))?;
-	Ok(ext)
+	let (result, allocation_stats) = executor.uncached_call_with_allocation_stats(
+		RuntimeBlob::uncompress_if_needed(binary).expect("Failed to parse binary"),
+		&mut ext_ext,
+		false,
+		method,
+		&[],
+	);
+	let result = result
+		.map_err(|e| format!("Failed to execute `{}`: {}", method, e))
+		.map(|_| ext);
+	(result, allocation_stats)
 }
 
 fn call_wasm_method<HF: HostFunctionsT>(binary: &[u8], method: &str) -> TestExternalities {
-	call_wasm_method_with_result::<HF>(binary, method).unwrap()
+	call_wasm_method_with_result::<HF>(binary, method).0.unwrap()
 }
 
 #[test]
@@ -103,8 +104,9 @@ fn test_return_input_public_key() {
 
 #[test]
 fn host_function_not_found() {
-	let err =
-		call_wasm_method_with_result::<()>(wasm_binary_unwrap(), "test_return_data").unwrap_err();
+	let err = call_wasm_method_with_result::<()>(wasm_binary_unwrap(), "test_return_data")
+		.0
+		.unwrap_err();
 
 	assert!(err.contains("Instantiation: Export "));
 	assert!(err.contains(" not found"));
@@ -236,3 +238,43 @@ fn test_tracing() {
 fn test_return_input_as_tuple() {
 	call_wasm_method::<HostFunctions>(wasm_binary_unwrap(), "test_return_input_as_tuple");
 }
+
+#[test]
+fn test_returning_option_bytes_from_a_host_function_is_efficient() {
+	let (result, stats_vec) = call_wasm_method_with_result::<HostFunctions>(
+		wasm_binary_unwrap(),
+		"test_return_option_vec",
+	);
+	result.unwrap();
+	let (result, stats_bytes) = call_wasm_method_with_result::<HostFunctions>(
+		wasm_binary_unwrap(),
+		"test_return_option_bytes",
+	);
+	result.unwrap();
+
+	let stats_vec = stats_vec.unwrap();
+	let stats_bytes = stats_bytes.unwrap();
+
+	// The way we currently implement marshaling of `Option<Vec<u8>>` through
+	// the WASM FFI boundary from the host to the runtime requires that it is
+	// marshaled through SCALE. This is quite inefficient as it requires two
+	// memory allocations inside of the runtime:
+	//
+	//   1) the first allocation to copy the SCALE-encoded blob into the runtime;
+	//   2) and another allocation for the resulting `Vec<u8>` when decoding that blob.
+	//
+	// Both of these allocations are are as big as the `Vec<u8>` which is being
+	// passed to the runtime. This is especially bad when fetching big values
+	// from storage, as it can lead to an out-of-memory situation.
+	//
+	// Our `Option<Bytes>` marshaling is better; it still must go through SCALE,
+	// and it still requires two allocations, however since `Bytes` is zero-copy
+	// only the first allocation is `Vec<u8>`-sized, and the second allocation
+	// which creates the deserialized `Bytes` is tiny, and is only necessary because
+	// the underlying `Bytes` buffer from which we're deserializing gets automatically
+	// turned into an `Arc`.
+	//
+	// So this assertion tests that deserializing `Option<Bytes>` allocates less than
+	// deserializing `Option<Vec<u8>>`.
+	assert_eq!(stats_bytes.bytes_allocated_sum + 16 * 1024 + 8, stats_vec.bytes_allocated_sum);
+}