Respect heap_pages parameter. (#3420)

substrate/core/executor/src/wasm_runtimes_cache.rs

@@ -86,7 +86,7 @@ struct StateSnapshot {
 	data_segments: Vec<(u32, Vec<u8>)>,
 	/// The list of all global mutable variables of the module in their sequential order.
 	global_mut_values: Vec<RuntimeValue>,
-	heap_pages: u32,
+	heap_pages: u64,
 }
 
 impl StateSnapshot {
@@ -94,7 +94,7 @@ impl StateSnapshot {
 	fn take(
 		module_instance: &WasmModuleInstanceRef,
 		data_segments: Vec<DataSegment>,
-		heap_pages: u32,
+		heap_pages: u64,
 	) -> Option<Self> {
 		let prepared_segments = data_segments
 			.into_iter()
@@ -250,6 +250,12 @@ impl RuntimesCache {
 			.original_storage_hash(well_known_keys::CODE)
 			.ok_or(Error::InvalidCode("`CODE` not found in storage.".into()))?;
 
+		let heap_pages = ext
+			.storage(well_known_keys::HEAP_PAGES)
+			.and_then(|pages| u64::decode(&mut &pages[..]).ok())
+			.or(default_heap_pages)
+			.unwrap_or(DEFAULT_HEAP_PAGES);
+
 		// This is direct result from fighting with borrowck.
 		let handle_result =
 			|cached_result: &Result<Rc<CachedRuntime>, CacheError>| match *cached_result {
@@ -258,10 +264,25 @@ impl RuntimesCache {
 			};
 
 		match self.instances.entry(code_hash.into()) {
-			Entry::Occupied(o) => handle_result(o.get()),
+			Entry::Occupied(mut o) => {
+				let result = o.get_mut();
+				if let Ok(ref cached_runtime) = result {
+					if cached_runtime.state_snapshot.heap_pages != heap_pages {
+						trace!(
+							target: "runtimes_cache",
+							"heap_pages were changed. Reinstantiating the instance"
+						);
+						*result = Self::create_wasm_instance(wasm_executor, ext, heap_pages);
+						if let Err(ref err) = result {
+							warn!(target: "runtimes_cache", "cannot create a runtime: {:?}", err);
+						}
+					}
+				}
+				handle_result(result)
+			},
 			Entry::Vacant(v) => {
 				trace!(target: "runtimes_cache", "no instance found in cache, creating now.");
-				let result = Self::create_wasm_instance(wasm_executor, ext, default_heap_pages);
+				let result = Self::create_wasm_instance(wasm_executor, ext, heap_pages);
 				if let Err(ref err) = result {
 					warn!(target: "runtimes_cache", "cannot create a runtime: {:?}", err);
 				}
@@ -273,7 +294,7 @@ impl RuntimesCache {
 	fn create_wasm_instance<E: Externalities<Blake2Hasher>>(
 		wasm_executor: &WasmExecutor,
 		ext: &mut E,
-		default_heap_pages: Option<u64>,
+		heap_pages: u64,
 	) -> Result<Rc<CachedRuntime>, CacheError> {
 		let code = ext
 			.original_storage(well_known_keys::CODE)
@@ -286,18 +307,12 @@ impl RuntimesCache {
 		// we just loaded and validated the `module` above.
 		let data_segments = extract_data_segments(&code).ok_or(CacheError::CantDeserializeWasm)?;
 
-		let heap_pages = ext
-			.storage(well_known_keys::HEAP_PAGES)
-			.and_then(|pages| u64::decode(&mut &pages[..]).ok())
-			.or(default_heap_pages)
-			.unwrap_or(DEFAULT_HEAP_PAGES);
-
 		// Instantiate this module.
 		let instance = WasmExecutor::instantiate_module::<E>(heap_pages as usize, &module)
 			.map_err(CacheError::Instantiation)?;
 
 		// Take state snapshot before executing anything.
-		let state_snapshot = StateSnapshot::take(&instance, data_segments, heap_pages as u32)
+		let state_snapshot = StateSnapshot::take(&instance, data_segments, heap_pages)
 			.expect(
 				"`take` returns `Err` if the module is not valid;
 				we already loaded module above, thus the `Module` is proven to be valid at this point;

substrate/core/test-runtime/src/lib.rs

@@ -261,6 +261,7 @@ cfg_if! {
 				fn benchmark_direct_call() -> u64;
 				fn returns_mutable_static() -> u64;
 				fn allocates_huge_stack_array(trap: bool) -> Vec<u8>;
+				fn vec_with_capacity(size: u32) -> Vec<u8>;
 				/// Returns the initialized block number.
 				fn get_block_number() -> u64;
 				/// Takes and returns the initialized block number.
@@ -302,6 +303,7 @@ cfg_if! {
 				fn benchmark_direct_call() -> u64;
 				fn returns_mutable_static() -> u64;
 				fn allocates_huge_stack_array(trap: bool) -> Vec<u8>;
+				fn vec_with_capacity(size: u32) -> Vec<u8>;
 				/// Returns the initialized block number.
 				fn get_block_number() -> u64;
 				/// Takes and returns the initialized block number.
@@ -561,6 +563,10 @@ cfg_if! {
 					unimplemented!("is not expected to be invoked from non-wasm builds");
 				}
 
+				fn vec_with_capacity(_size: u32) -> Vec<u8> {
+					unimplemented!("is not expected to be invoked from non-wasm builds");
+				}
+
 				fn get_block_number() -> u64 {
 					system::get_block_number().expect("Block number is initialized")
 				}
@@ -771,6 +777,10 @@ cfg_if! {
 					data.to_vec()
 				}
 
+				fn vec_with_capacity(size: u32) -> Vec<u8> {
+					Vec::with_capacity(size as usize)
+				}
+
 				fn get_block_number() -> u64 {
 					system::get_block_number().expect("Block number is initialized")
 				}
@@ -875,6 +885,7 @@ fn test_sr25519_crypto() -> (sr25519::AppSignature, sr25519::AppPublic) {
 mod tests {
 	use substrate_test_runtime_client::{
 		prelude::*,
+		consensus::BlockOrigin,
 		DefaultTestClientBuilderExt, TestClientBuilder,
 		runtime::TestAPI,
 	};
@@ -882,7 +893,9 @@ mod tests {
 		generic::BlockId,
 		traits::ProvideRuntimeApi,
 	};
+	use primitives::storage::well_known_keys::HEAP_PAGES;
 	use state_machine::ExecutionStrategy;
+	use codec::Encode;
 
 	#[test]
 	fn returns_mutable_static() {
@@ -930,4 +943,36 @@ mod tests {
 		assert!(ret.is_ok());
 	}
 
+	#[test]
+	fn heap_pages_is_respected() {
+		// This tests that the on-chain HEAP_PAGES parameter is respected.
+
+		// Create a client devoting only 8 pages of wasm memory. This gives us ~512k of heap memory.
+		let client = TestClientBuilder::new()
+			.set_execution_strategy(ExecutionStrategy::AlwaysWasm)
+			.set_heap_pages(8)
+			.build();
+		let runtime_api = client.runtime_api();
+		let block_id = BlockId::Number(client.info().chain.best_number);
+
+		// Try to allocate 1024k of memory on heap. This is going to fail since it is twice larger
+		// than the heap.
+		let ret = runtime_api.vec_with_capacity(&block_id, 1048576);
+		assert!(ret.is_err());
+
+		// Create a block that sets the `:heap_pages` to 32 pages of memory which corresponds to
+		// ~2048k of heap memory.
+		let new_block_id = {
+			let mut builder = client.new_block(Default::default()).unwrap();
+			builder.push_storage_change(HEAP_PAGES.to_vec(), Some(32u64.encode())).unwrap();
+			let block = builder.bake().unwrap();
+			let hash = block.header.hash();
+			client.import(BlockOrigin::Own, block).unwrap();
+			BlockId::Hash(hash)
+		};
+
+		// Allocation of 1024k while having ~2048k should succeed.
+		let ret = runtime_api.vec_with_capacity(&new_block_id, 1048576);
+		assert!(ret.is_ok());
+	}
 }