Fair reusing of wasm runtime instances (#3011)

* Add test from original bug report

Original is from @pepyakin in 3d7b27f3421818e8d6de568e02fbc2947a06246b.
I adapted it to work with the latest master.

* No longer cleanup module instance

* Replace runtime cache with synchronous clone

* Fix test

* Preserve initial runtime memory and restore it on fetch

* Remove leftover comment

* Fix style

* Improve variable naming

* Replace get_into() with get()

* Handle missing memory export better

* Return earlier when creating runtime first time

* Improve comments

* fmt

* Fix #2967.

* Eradicate `code` from `Error::InvalidCode`

* tidy

* A state snapshot doc.

* Store multiple runtimes by hash.

* Get rid of deref.

* Docs

* Use Self for instantiate_module

* REVERT ME

* Should be ok

* Commit

* Remove dbg

* Use fast-memory's erase

* Clean and undo hacks.

* Introduce a dedicated error for heap_base

* Ban the start function.

* Clean, docs and refactor

* Add rustflags.

* Update Cargo.lock

* Apply Basti's suggestions

Co-Authored-By: Bastian Köcher <bkchr@users.noreply.github.com>

* Rename allocates_huge_stack_array

* Extend TestClientBuilder with set_heap_pages

* Update the test.

* Update core/executor/src/wasm_executor.rs

Co-Authored-By: Bastian Köcher <bkchr@users.noreply.github.com>

* Update core/executor/src/wasm_runtimes_cache.rs

Co-Authored-By: Bastian Köcher <bkchr@users.noreply.github.com>

* Update core/executor/src/error.rs

Co-Authored-By: Bastian Köcher <bkchr@users.noreply.github.com>

* Update core/executor/src/error.rs

Co-Authored-By: Bastian Köcher <bkchr@users.noreply.github.com>

* Fix tests.

* Update cargo-lock

* Use wasmi master

* Use master wasmi

* Move tests.

* Use wasmi crates.io

* Update Cargo.lock

* Fix build.rs

* Bump runtime version

* Revert initial_heap_pages renaming

* Bump wasmi up to 0.5.0

* Bump runtime version

* Don't restore an instance every now and then

* Update core/executor/src/wasm_runtimes_cache.rs

Co-Authored-By: DemiMarie-parity <48690212+DemiMarie-parity@users.noreply.github.com>

* Propagate error in CacheError

* Clarify the get_heap_base call in instantiation

* Supply --export=__heap_base

See https://reviews.llvm.org/D62744

Co-authored-by: Jim Posen <jim.posen@gmail.com>

* Bump version.

* Use combinators for segments.

* Fix build.rs

* Fix build.rs for runtime-test

substrate/core/executor/src/allocator.rs

@@ -53,22 +53,13 @@ impl FreeingBumpHeapAllocator {
 	///
 	/// # Arguments
 	///
-	/// * `ptr_offset` - The pointers returned by `allocate()` start from this
-	///   offset on. The pointer offset needs to be aligned to a multiple of 8,
-	///   hence a padding might be added to align `ptr_offset` properly.
-	///
-	/// * `heap_size` - The size available to this heap instance (in bytes) for
-	///   allocating memory.
-	///
-	/// * `heap` - A `MemoryRef` to the available `MemoryInstance` which is
-	///   used as the heap.
-	///
-	pub fn new(mem: MemoryRef) -> Self {
+	/// - `mem` - reference to the linear memory instance on which this allocator operates.
+	/// - `heap_base` - the offset from the beginning of the linear memory where the heap starts.
+	pub fn new(mem: MemoryRef, heap_base: u32) -> Self {
 		let current_size: Bytes = mem.current_size().into();
 		let current_size = current_size.0 as u32;
-		let used_size = mem.used_size().0 as u32;
 
-		let mut ptr_offset = used_size;
+		let mut ptr_offset = heap_base;
 		let padding = ptr_offset % ALIGNMENT;
 		if padding != 0 {
 			ptr_offset += ALIGNMENT - padding;
@@ -195,16 +186,11 @@ mod tests {
 
 	const PAGE_SIZE: u32 = 65536;
 
-	fn set_offset(mem: MemoryRef, offset: usize) {
-		let offset: Vec<u8> = vec![255; offset];
-		mem.set(0, &offset).unwrap();
-	}
-
 	#[test]
 	fn should_allocate_properly() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 0);
 
 		// when
 		let ptr = heap.allocate(1).unwrap();
@@ -217,8 +203,7 @@ mod tests {
 	fn should_always_align_pointers_to_multiples_of_8() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		set_offset(mem.clone(), 13);
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 13);
 
 		// when
 		let ptr = heap.allocate(1).unwrap();
@@ -233,7 +218,7 @@ mod tests {
 	fn should_increment_pointers_properly() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 0);
 
 		// when
 		let ptr1 = heap.allocate(1).unwrap();
@@ -256,7 +241,7 @@ mod tests {
 	fn should_free_properly() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 0);
 		let ptr1 = heap.allocate(1).unwrap();
 		// the prefix of 8 bytes is prepended to the pointer
 		assert_eq!(ptr1, 8);
@@ -278,9 +263,8 @@ mod tests {
 	fn should_deallocate_and_reallocate_properly() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		set_offset(mem.clone(), 13);
 		let padded_offset = 16;
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 13);
 
 		let ptr1 = heap.allocate(1).unwrap();
 		// the prefix of 8 bytes is prepended to the pointer
@@ -306,7 +290,7 @@ mod tests {
 	fn should_build_linked_list_of_free_areas_properly() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 0);
 
 		let ptr1 = heap.allocate(8).unwrap();
 		let ptr2 = heap.allocate(8).unwrap();
@@ -333,8 +317,7 @@ mod tests {
 	fn should_not_allocate_if_too_large() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), Some(Pages(1))).unwrap();
-		set_offset(mem.clone(), 13);
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 13);
 
 		// when
 		let ptr = heap.allocate(PAGE_SIZE - 13);
@@ -353,7 +336,7 @@ mod tests {
 	fn should_not_allocate_if_full() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), Some(Pages(1))).unwrap();
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 0);
 		let ptr1 = heap.allocate((PAGE_SIZE / 2) - 8).unwrap();
 		assert_eq!(ptr1, 8);
 
@@ -376,7 +359,7 @@ mod tests {
 		// given
 		let pages_needed = (MAX_POSSIBLE_ALLOCATION as usize / PAGE_SIZE as usize) + 1;
 		let mem = MemoryInstance::alloc(Pages(pages_needed), Some(Pages(pages_needed))).unwrap();
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 0);
 
 		// when
 		let ptr = heap.allocate(MAX_POSSIBLE_ALLOCATION).unwrap();
@@ -389,7 +372,7 @@ mod tests {
 	fn should_not_allocate_if_requested_size_too_large() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 0);
 
 		// when
 		let ptr = heap.allocate(MAX_POSSIBLE_ALLOCATION + 1);
@@ -408,8 +391,7 @@ mod tests {
 	fn should_include_prefixes_in_total_heap_size() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		set_offset(mem.clone(), 1);
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 1);
 
 		// when
 		// an item size of 16 must be used then
@@ -423,8 +405,7 @@ mod tests {
 	fn should_calculate_total_heap_size_to_zero() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		set_offset(mem.clone(), 13);
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 13);
 
 		// when
 		let ptr = heap.allocate(42).unwrap();
@@ -439,8 +420,7 @@ mod tests {
 	fn should_calculate_total_size_of_zero() {
 		// given
 		let mem = MemoryInstance::alloc(Pages(1), None).unwrap();
-		set_offset(mem.clone(), 19);
-		let mut heap = FreeingBumpHeapAllocator::new(mem);
+		let mut heap = FreeingBumpHeapAllocator::new(mem, 19);
 
 		// when
 		for _ in 1..10 {

substrate/core/executor/src/error.rs

@@ -38,8 +38,8 @@ pub enum Error {
 	#[display(fmt="Method not found: '{}'", _0)]
 	MethodNotFound(String),
 	/// Code is invalid (expected single byte)
-	#[display(fmt="Invalid Code: {:?}", _0)]
-	InvalidCode(Vec<u8>),
+	#[display(fmt="Invalid Code")]
+	InvalidCode,
 	/// Could not get runtime version.
 	#[display(fmt="On-chain runtime does not specify version")]
 	VersionInvalid,
@@ -58,6 +58,13 @@ pub enum Error {
 	/// Invalid memory reference.
 	#[display(fmt="Invalid memory reference")]
 	InvalidMemoryReference,
+	/// The runtime must provide a global named `__heap_base` of type i32 for specifying where the
+	/// allocator is allowed to place its data.
+	#[display(fmt="The runtime doesn't provide a global named `__heap_base`")]
+	HeapBaseNotFoundOrInvalid,
+	/// The runtime WebAssembly module is not allowed to have the `start` function.
+	#[display(fmt="The runtime has the `start` function")]
+	RuntimeHasStartFn,
 	/// Some other error occurred
 	Other(&'static str),
 	/// Some error occurred in the allocator

substrate/core/executor/src/lib.rs

@@ -35,11 +35,13 @@ mod wasm_executor;
 mod native_executor;
 mod sandbox;
 mod allocator;
+mod wasm_runtimes_cache;
 
 pub mod error;
 pub use wasmi;
 pub use wasm_executor::WasmExecutor;
 pub use native_executor::{with_native_environment, NativeExecutor, NativeExecutionDispatch};
+pub use wasm_runtimes_cache::RuntimesCache;
 pub use state_machine::Externalities;
 pub use runtime_version::{RuntimeVersion, NativeVersion};
 pub use parity_codec::Codec;

substrate/core/executor/src/native_executor.rs

@@ -14,86 +14,20 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
-use std::{borrow::BorrowMut, result, cell::{RefMut, RefCell}};
+use std::{result, cell::RefCell, panic::UnwindSafe};
 use crate::error::{Error, Result};
 use state_machine::{CodeExecutor, Externalities};
 use crate::wasm_executor::WasmExecutor;
-use wasmi::{Module as WasmModule, ModuleRef as WasmModuleInstanceRef};
 use runtime_version::{NativeVersion, RuntimeVersion};
-use std::{collections::HashMap, panic::UnwindSafe};
 use parity_codec::{Decode, Encode};
 use crate::RuntimeInfo;
 use primitives::{Blake2Hasher, NativeOrEncoded};
-use primitives::storage::well_known_keys;
 use log::trace;
 
-/// Default num of pages for the heap
-const DEFAULT_HEAP_PAGES: u64 = 1024;
-
-// For the internal Runtime Cache:
-// Is it compatible enough to run this natively or do we need to fall back on the WasmModule
-
-enum RuntimePreproc {
-	InvalidCode,
-	ValidCode(WasmModuleInstanceRef, Option<RuntimeVersion>),
-}
-
-type CacheType = HashMap<[u8; 32], RuntimePreproc>;
+use crate::RuntimesCache;
 
 thread_local! {
-	static RUNTIMES_CACHE: RefCell<CacheType> = RefCell::new(HashMap::new());
-}
-
-/// fetch a runtime version from the cache or if there is no cached version yet, create
-/// the runtime version entry for `code`, determines whether `Compatibility::IsCompatible`
-/// can be used by comparing returned RuntimeVersion to `ref_version`
-fn fetch_cached_runtime_version<'a, E: Externalities<Blake2Hasher>>(
-	wasm_executor: &WasmExecutor,
-	cache: &'a mut RefMut<CacheType>,
-	ext: &mut E,
-	default_heap_pages: Option<u64>,
-) -> Result<(&'a WasmModuleInstanceRef, &'a Option<RuntimeVersion>)> {
-	let code_hash = match ext.original_storage_hash(well_known_keys::CODE) {
-		Some(code_hash) => code_hash,
-		None => return Err(Error::InvalidCode(vec![])),
-	};
-
-	let maybe_runtime_preproc = cache.borrow_mut().entry(code_hash.into())
-		.or_insert_with(|| {
-			let code = match ext.original_storage(well_known_keys::CODE) {
-				Some(code) => code,
-				None => return RuntimePreproc::InvalidCode,
-			};
-			let heap_pages = ext.storage(well_known_keys::HEAP_PAGES)
-				.and_then(|pages| u64::decode(&mut &pages[..]))
-				.or(default_heap_pages)
-				.unwrap_or(DEFAULT_HEAP_PAGES);
-			match WasmModule::from_buffer(code)
-				.map_err(|_| Error::InvalidCode(vec![]))
-				.and_then(|module| wasm_executor.prepare_module(ext, heap_pages as usize, &module))
-			{
-				Ok(module) => {
-					let version = wasm_executor.call_in_wasm_module(ext, &module, "Core_version", &[])
-						.ok()
-						.and_then(|v| RuntimeVersion::decode(&mut v.as_slice()));
-					RuntimePreproc::ValidCode(module, version)
-				}
-				Err(e) => {
-					trace!(target: "executor", "Invalid code presented to executor ({:?})", e);
-					RuntimePreproc::InvalidCode
-				}
-			}
-		});
-
-	match maybe_runtime_preproc {
-		RuntimePreproc::InvalidCode => {
-			let code = ext.original_storage(well_known_keys::CODE).unwrap_or(vec![]);
-			Err(Error::InvalidCode(code))
-		},
-		RuntimePreproc::ValidCode(m, v) => {
-			Ok((m, v))
-		}
-	}
+	static RUNTIMES_CACHE: RefCell<RuntimesCache> = RefCell::new(RuntimesCache::new());
 }
 
 fn safe_call<F, U>(f: F) -> Result<U>
@@ -140,7 +74,7 @@ pub struct NativeExecutor<D> {
 	fallback: WasmExecutor,
 	/// Native runtime version info.
 	native_version: NativeVersion,
-	/// The default number of 64KB pages to allocate for Wasm execution.
+	/// The number of 64KB pages to allocate for Wasm execution.
 	default_heap_pages: Option<u64>,
 }
 
@@ -151,7 +85,7 @@ impl<D: NativeExecutionDispatch> NativeExecutor<D> {
 			_dummy: Default::default(),
 			fallback: WasmExecutor::new(),
 			native_version: D::native_version(),
-			default_heap_pages,
+			default_heap_pages: default_heap_pages,
 		}
 	}
 }
@@ -176,10 +110,11 @@ impl<D: NativeExecutionDispatch> RuntimeInfo for NativeExecutor<D> {
 		&self,
 		ext: &mut E,
 	) -> Option<RuntimeVersion> {
-		RUNTIMES_CACHE.with(|c|
-			fetch_cached_runtime_version(&self.fallback, &mut c.borrow_mut(), ext, self.default_heap_pages)
-				.ok()?.1.clone()
-		)
+		RUNTIMES_CACHE.with(|cache| {
+			let cache = &mut cache.borrow_mut();
+			cache.fetch_runtime(&self.fallback, ext, self.default_heap_pages)
+				.ok()?.version().clone()
+		})
 	}
 }
 
@@ -198,14 +133,16 @@ impl<D: NativeExecutionDispatch> CodeExecutor<Blake2Hasher> for NativeExecutor<D
 		data: &[u8],
 		use_native: bool,
 		native_call: Option<NC>,
-	) -> (Result<NativeOrEncoded<R>>, bool) {
-		RUNTIMES_CACHE.with(|c| {
-			let mut c = c.borrow_mut();
-			let (module, onchain_version) = match fetch_cached_runtime_version(
-				&self.fallback, &mut c, ext, self.default_heap_pages) {
-					Ok((module, onchain_version)) => (module, onchain_version),
-					Err(e) => return (Err(e), false),
+	) -> (Result<NativeOrEncoded<R>>, bool){
+		RUNTIMES_CACHE.with(|cache| {
+			let cache = &mut cache.borrow_mut();
+			let cached_runtime = match cache.fetch_runtime(
+				&self.fallback, ext, self.default_heap_pages,
+			) {
+				Ok(cached_runtime) => cached_runtime,
+				Err(e) => return (Err(e), false),
 			};
+			let onchain_version = cached_runtime.version();
 			match (
 				use_native,
 				onchain_version
@@ -223,17 +160,21 @@ impl<D: NativeExecutionDispatch> CodeExecutor<Blake2Hasher> for NativeExecutor<D
 							.map_or_else(||"<None>".into(), |v| format!("{}", v))
 					);
 					(
-						self.fallback
-							.call_in_wasm_module(ext, module, method, data)
-							.map(NativeOrEncoded::Encoded),
+						cached_runtime.with(|module|
+							self.fallback
+								.call_in_wasm_module(ext, module, method, data)
+								.map(NativeOrEncoded::Encoded)
+						),
 						false
 					)
 				}
 				(false, _, _) => {
 					(
-						self.fallback
-							.call_in_wasm_module(ext, module, method, data)
-							.map(NativeOrEncoded::Encoded),
+						cached_runtime.with(|module|
+							self.fallback
+								.call_in_wasm_module(ext, module, method, data)
+								.map(NativeOrEncoded::Encoded)
+						),
 						false
 					)
 				}

substrate/core/executor/src/wasm_executor.rs

@@ -56,10 +56,10 @@ struct FunctionExecutor<'e, E: Externalities<Blake2Hasher> + 'e> {
 }
 
 impl<'e, E: Externalities<Blake2Hasher>> FunctionExecutor<'e, E> {
-	fn new(m: MemoryRef, t: Option<TableRef>, e: &'e mut E) -> Result<Self> {
+	fn new(m: MemoryRef, heap_base: u32, t: Option<TableRef>, e: &'e mut E) -> Result<Self> {
 		Ok(FunctionExecutor {
 			sandbox_store: sandbox::Store::new(),
-			heap: allocator::FreeingBumpHeapAllocator::new(m.clone()),
+			heap: allocator::FreeingBumpHeapAllocator::new(m.clone(), heap_base),
 			memory: m,
 			table: t,
 			ext: e,
@@ -1270,7 +1270,7 @@ impl WasmExecutor {
 		data: &[u8],
 	) -> Result<Vec<u8>> {
 		let module = ::wasmi::Module::from_buffer(code)?;
-		let module = self.prepare_module(ext, heap_pages, &module)?;
+		let module = Self::instantiate_module::<E>(heap_pages, &module)?;
 		self.call_in_wasm_module(ext, &module, method, data)
 	}
 
@@ -1292,7 +1292,7 @@ impl WasmExecutor {
 		filter_result: FR,
 	) -> Result<R> {
 		let module = wasmi::Module::from_buffer(code)?;
-		let module = self.prepare_module(ext, heap_pages, &module)?;
+		let module = Self::instantiate_module::<E>(heap_pages, &module)?;
 		self.call_in_wasm_module_with_custom_signature(
 			ext,
 			&module,
@@ -1311,6 +1311,22 @@ impl WasmExecutor {
 			.clone())
 	}
 
+	/// Find the global named `__heap_base` in the given wasm module instance and
+	/// tries to get its value.
+	fn get_heap_base(module: &ModuleRef) -> Result<u32> {
+		let heap_base_val = module
+			.export_by_name("__heap_base")
+			.ok_or_else(|| Error::HeapBaseNotFoundOrInvalid)?
+			.as_global()
+			.ok_or_else(|| Error::HeapBaseNotFoundOrInvalid)?
+			.get();
+
+		Ok(match heap_base_val {
+			wasmi::RuntimeValue::I32(v) => v as u32,
+			_ => return Err(Error::HeapBaseNotFoundOrInvalid),
+		})
+	}
+
 	/// Call a given method in the given wasm-module runtime.
 	pub fn call_in_wasm_module<E: Externalities<Blake2Hasher>>(
 		&self,
@@ -1359,10 +1375,9 @@ impl WasmExecutor {
 		let table: Option<TableRef> = module_instance
 			.export_by_name("__indirect_function_table")
 			.and_then(|e| e.as_table().cloned());
+		let heap_base = Self::get_heap_base(module_instance)?;
 
-		let low = memory.lowest_used();
-		let used_mem = memory.used_size();
-		let mut fec = FunctionExecutor::new(memory.clone(), table, ext)?;
+		let mut fec = FunctionExecutor::new(memory.clone(), heap_base, table, ext)?;
 		let parameters = create_parameters(&mut |data: &[u8]| {
 			let offset = fec.heap.allocate(data.len() as u32)?;
 			memory.set(offset, &data)?;
@@ -1385,24 +1400,14 @@ impl WasmExecutor {
 			},
 		};
 
-		// cleanup module instance for next use
-		let new_low = memory.lowest_used();
-		if new_low < low {
-			memory.zero(new_low as usize, (low - new_low) as usize)?;
-			memory.reset_lowest_used(low);
-		}
-		memory.with_direct_access_mut(|buf| buf.resize(used_mem.0, 0));
 		result
 	}
 
 	/// Prepare module instance
-	pub fn prepare_module<E: Externalities<Blake2Hasher>>(
-		&self,
-		ext: &mut E,
+	pub fn instantiate_module<E: Externalities<Blake2Hasher>>(
 		heap_pages: usize,
 		module: &Module,
-		) -> Result<ModuleRef>
-	{
+	) -> Result<ModuleRef> {
 		// start module instantiation. Don't run 'start' function yet.
 		let intermediate_instance = ModuleInstance::new(
 			module,
@@ -1410,18 +1415,19 @@ impl WasmExecutor {
 			.with_resolver("env", FunctionExecutor::<E>::resolver())
 		)?;
 
-		// extract a reference to a linear memory, optional reference to a table
-		// and then initialize FunctionExecutor.
+		// Verify that the module has the heap base global variable.
+		let _ = Self::get_heap_base(intermediate_instance.not_started_instance())?;
+
+		// Extract a reference to a linear memory.
 		let memory = Self::get_mem_instance(intermediate_instance.not_started_instance())?;
 		memory.grow(Pages(heap_pages)).map_err(|_| Error::Runtime)?;
-		let table: Option<TableRef> = intermediate_instance
-			.not_started_instance()
-			.export_by_name("__indirect_function_table")
-			.and_then(|e| e.as_table().cloned());
-		let mut fec = FunctionExecutor::new(memory.clone(), table, ext)?;
 
-		// finish instantiation by running 'start' function (if any).
-		Ok(intermediate_instance.run_start(&mut fec)?)
+		if intermediate_instance.has_start() {
+			// Runtime is not allowed to have the `start` function.
+			Err(Error::RuntimeHasStartFn)
+		} else {
+			Ok(intermediate_instance.assert_no_start())
+		}
 	}
 }
 

substrate/core/executor/src/wasm_runtimes_cache.rs

@@ -0,0 +1,331 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Implements a cache for pre-created Wasm runtime module instances.
+
+use crate::error::Error;
+use crate::wasm_executor::WasmExecutor;
+use log::{trace, warn};
+use parity_codec::Decode;
+use parity_wasm::elements::{deserialize_buffer, DataSegment, Instruction, Module as RawModule};
+use primitives::storage::well_known_keys;
+use primitives::Blake2Hasher;
+use runtime_version::RuntimeVersion;
+use state_machine::Externalities;
+use std::collections::hash_map::{Entry, HashMap};
+use std::mem;
+use std::rc::Rc;
+use wasmi::{Module as WasmModule, ModuleRef as WasmModuleInstanceRef, RuntimeValue};
+
+#[derive(Debug)]
+enum CacheError {
+	CodeNotFound,
+	ApplySnapshotFailed,
+	InvalidModule,
+	CantDeserializeWasm,
+	Instantiation(Error),
+}
+
+/// A runtime along with its version and initial state snapshot.
+#[derive(Clone)]
+pub struct CachedRuntime {
+	/// A wasm module instance.
+	instance: WasmModuleInstanceRef,
+	/// Runtime version according to `Core_version`.
+	///
+	/// Can be `None` if the runtime doesn't expose this function.
+	version: Option<RuntimeVersion>,
+	/// The snapshot of the instance's state taken just after the instantiation.
+	state_snapshot: StateSnapshot,
+}
+
+impl CachedRuntime {
+	/// Perform an operation with the clean version of the runtime wasm instance.
+	pub fn with<R, F>(&self, f: F) -> R
+	where
+		F: FnOnce(&WasmModuleInstanceRef) -> R,
+	{
+		self.state_snapshot.apply(&self.instance).expect(
+			"applying the snapshot can only fail if the passed instance is different
+			from the one that was used for creation of the snapshot;
+			we use the snapshot that is directly associated with the instance;
+			thus the snapshot was created using the instance;
+			qed",
+		);
+		f(&self.instance)
+	}
+
+	/// Returns the version of this cached runtime.
+	///
+	/// Returns `None` if the runtime doesn't provide the information or there was an error
+	/// while fetching it.
+	pub fn version(&self) -> Option<RuntimeVersion> {
+		self.version.clone()
+	}
+}
+
+/// A state snapshot of an instance taken just after instantiation.
+///
+/// It is used for restoring the state of the module after execution.
+#[derive(Clone)]
+struct StateSnapshot {
+	/// The offset and the content of the memory segments that should be used to restore the snapshot
+	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,
+}
+
+impl StateSnapshot {
+	// Returns `None` if instance is not valid.
+	fn take(
+		module_instance: &WasmModuleInstanceRef,
+		data_segments: Vec<DataSegment>,
+		heap_pages: u32,
+	) -> Option<Self> {
+		let prepared_segments = data_segments
+			.into_iter()
+			.map(|mut segment| {
+				// Just replace contents of the segment since the segments will be discarded later
+				// anyway.
+				let contents = mem::replace(segment.value_mut(), vec![]);
+
+				let init_expr = segment.offset().code();
+				// [op, End]
+				if init_expr.len() != 2 {
+					return None;
+				}
+				let offset = match init_expr[0] {
+					Instruction::I32Const(v) => v as u32,
+					Instruction::GetGlobal(idx) => {
+						let global_val = module_instance.globals().get(idx as usize)?.get();
+						match global_val {
+							RuntimeValue::I32(v) => v as u32,
+							_ => return None,
+						}
+					}
+					_ => return None,
+				};
+
+				Some((offset, contents))
+			})
+			.collect::<Option<Vec<_>>>()?;
+
+		// Collect all values of mutable globals.
+		let global_mut_values = module_instance
+			.globals()
+			.iter()
+			.filter(|g| g.is_mutable())
+			.map(|g| g.get())
+			.collect();
+
+		Some(Self {
+			data_segments: prepared_segments,
+			global_mut_values,
+			heap_pages,
+		})
+	}
+
+	/// Reset the runtime instance to the initial version by restoring
+	/// the preserved memory and globals.
+	///
+	/// Returns `Err` if applying the snapshot is failed.
+	fn apply(&self, instance: &WasmModuleInstanceRef) -> Result<(), CacheError> {
+		let memory = instance
+			.export_by_name("memory")
+			.ok_or(CacheError::ApplySnapshotFailed)?
+			.as_memory()
+			.cloned()
+			.ok_or(CacheError::ApplySnapshotFailed)?;
+
+		// First, erase the memory and copy the data segments into it.
+		memory
+			.erase()
+			.map_err(|_| CacheError::ApplySnapshotFailed)?;
+		for (offset, contents) in &self.data_segments {
+			memory
+				.set(*offset, contents)
+				.map_err(|_| CacheError::ApplySnapshotFailed)?;
+		}
+
+		// Second, restore the values of mutable globals.
+		for (global_ref, global_val) in instance
+			.globals()
+			.iter()
+			.filter(|g| g.is_mutable())
+			.zip(self.global_mut_values.iter())
+		{
+			// the instance should be the same as used for preserving and
+			// we iterate the same way it as we do it for preserving values that means that the
+			// types should be the same and all the values are mutable. So no error is expected/
+			global_ref
+				.set(*global_val)
+				.map_err(|_| CacheError::ApplySnapshotFailed)?;
+		}
+		Ok(())
+	}
+}
+
+/// Default num of pages for the heap
+const DEFAULT_HEAP_PAGES: u64 = 1024;
+
+/// Cache for the runtimes.
+///
+/// When an instance is requested for the first time it is added to this
+/// cache. Furthermore its initial memory and values of mutable globals are preserved here. Follow-up
+/// requests to fetch a runtime return this one instance with the memory
+/// reset to the initial memory. So, one runtime instance is reused for
+/// every fetch request.
+///
+/// For now the cache grows indefinitely, but that should be fine for now since runtimes can only be
+/// upgraded rarely and there are no other ways to make the node to execute some other runtime.
+pub struct RuntimesCache {
+	/// A cache of runtime instances along with metadata, ready to be reused.
+	///
+	/// Instances are keyed by the hash of their code.
+	instances: HashMap<[u8; 32], Result<Rc<CachedRuntime>, CacheError>>,
+}
+
+impl RuntimesCache {
+	/// Creates a new instance of a runtimes cache.
+	pub fn new() -> RuntimesCache {
+		RuntimesCache {
+			instances: HashMap::new(),
+		}
+	}
+
+	/// Fetches an instance of the runtime.
+	///
+	/// On first use we create a new runtime instance, save it to the cache
+	/// and persist its initial memory.
+	///
+	/// Each subsequent request will return this instance, with its memory restored
+	/// to the persisted initial memory. Thus, we reuse one single runtime instance
+	/// for every `fetch_runtime` invocation.
+	///
+	/// # Parameters
+	///
+	/// `wasm_executor`- Rust wasm executor. Executes the provided code in a
+	/// sandboxed Wasm runtime.
+	///
+	/// `ext` - Externalities to use for the runtime. This is used for setting
+	/// up an initial runtime instance. The parameter is only needed for calling
+	/// into the Wasm module to find out the `Core_version`.
+	///
+	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
+	/// Defaults to `DEFAULT_HEAP_PAGES` if `None` is provided.
+	///
+	/// # Return value
+	///
+	/// If no error occurred a tuple `(wasmi::ModuleRef, Option<RuntimeVersion>)` is
+	/// returned. `RuntimeVersion` is contained if the call to `Core_version` returned
+	/// a version.
+	///
+	/// In case of failure one of two errors can be returned:
+	///
+	/// `Err::InvalidCode` is returned for runtime code issues.
+	///
+	/// `Error::InvalidMemoryReference` is returned if no memory export with the
+	/// identifier `memory` can be found in the runtime.
+	pub fn fetch_runtime<E: Externalities<Blake2Hasher>>(
+		&mut self,
+		wasm_executor: &WasmExecutor,
+		ext: &mut E,
+		default_heap_pages: Option<u64>,
+	) -> Result<Rc<CachedRuntime>, Error> {
+		let code_hash = ext
+			.original_storage_hash(well_known_keys::CODE)
+			.ok_or(Error::InvalidCode)?;
+
+		// This is direct result from fighting with borrowck.
+		let handle_result =
+			|cached_result: &Result<Rc<CachedRuntime>, CacheError>| match *cached_result {
+				Err(_) => Err(Error::InvalidCode),
+				Ok(ref cached_runtime) => Ok(Rc::clone(cached_runtime)),
+			};
+
+		match self.instances.entry(code_hash.into()) {
+			Entry::Occupied(o) => handle_result(o.get()),
+			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);
+				if let Err(ref err) = result {
+					warn!(target: "runtimes_cache", "cannot create a runtime: {:?}", err);
+				}
+				handle_result(v.insert(result))
+			}
+		}
+	}
+
+	fn create_wasm_instance<E: Externalities<Blake2Hasher>>(
+		wasm_executor: &WasmExecutor,
+		ext: &mut E,
+		default_heap_pages: Option<u64>,
+	) -> Result<Rc<CachedRuntime>, CacheError> {
+		let code = ext
+			.original_storage(well_known_keys::CODE)
+			.ok_or(CacheError::CodeNotFound)?;
+		let module = WasmModule::from_buffer(&code).map_err(|_| CacheError::InvalidModule)?;
+
+		// Extract the data segments from the wasm code.
+		//
+		// A return of this error actually indicates that there is a problem in logic, since
+		// 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[..]))
+			.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)
+			.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;
+				qed
+				",
+			);
+
+		let version = wasm_executor
+			.call_in_wasm_module(ext, &instance, "Core_version", &[])
+			.ok()
+			.and_then(|v| RuntimeVersion::decode(&mut v.as_slice()));
+		Ok(Rc::new(CachedRuntime {
+			instance,
+			version,
+			state_snapshot,
+		}))
+	}
+}
+
+/// Extract the data segments from the given wasm code.
+///
+/// Returns `Err` if the given wasm code cannot be deserialized.
+fn extract_data_segments(wasm_code: &[u8]) -> Option<Vec<DataSegment>> {
+	let raw_module: RawModule = deserialize_buffer(wasm_code).ok()?;
+	let segments = raw_module
+		.data_section()
+		.map(|ds| ds.entries())
+		.unwrap_or(&[])
+		.to_vec();
+	Some(segments)
+}