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/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);
+}