Better wasm instance cache (#5109)

* Wasm instance cache

* Reduce slot locking

* Fixed test

* Dispose of instance in case of error

* Fixed benches

* Style, comments, some renames

* Replaced upgradable lock with mutex

* Bump dependencies

* Re-export CallInWasm

* Update client/executor/src/wasm_runtime.rs

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

* Update client/executor/src/native_executor.rs

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

* Update client/executor/src/native_executor.rs

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

* Update client/executor/src/wasm_runtime.rs

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

* Update client/executor/wasmtime/src/runtime.rs

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

* Update client/executor/src/wasm_runtime.rs

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

* Update client/executor/src/wasm_runtime.rs

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

* Update client/executor/src/wasm_runtime.rs

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

* Indents

* Whitespace

* Formatting

* Added issue link

Co-authored-by: Benjamin Kampmann <ben.kampmann@googlemail.com>
Co-authored-by: Gavin Wood <github@gavwood.com>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/client/executor/src/integration_tests/mod.rs

@@ -21,7 +21,7 @@ use hex_literal::hex;
 use sp_core::{
 	blake2_128, blake2_256, ed25519, sr25519, map, Pair,
 	offchain::{OffchainExt, testing},
-	traits::Externalities,
+	traits::{Externalities, CallInWasm},
 };
 use sc_runtime_test::WASM_BINARY;
 use sp_state_machine::TestExternalities as CoreTestExternalities;
@@ -40,15 +40,18 @@ fn call_in_wasm<E: Externalities>(
 	call_data: &[u8],
 	execution_method: WasmExecutionMethod,
 	ext: &mut E,
-) -> crate::error::Result<Vec<u8>> {
-	crate::call_in_wasm::<HostFunctions>(
+) -> Result<Vec<u8>, String> {
+	let executor = crate::WasmExecutor::new(
+		execution_method,
+		Some(1024),
+		HostFunctions::host_functions(),
+		true,
+	);
+	executor.call_in_wasm(
+		&WASM_BINARY[..],
 		function,
 		call_data,
-		execution_method,
 		ext,
-		&WASM_BINARY[..],
-		1024,
-		true,
 	)
 }
 
@@ -84,12 +87,12 @@ fn call_not_existing_function(wasm_method: WasmExecutionMethod) {
 			match wasm_method {
 				WasmExecutionMethod::Interpreted => assert_eq!(
 					&format!("{:?}", e),
-					"Wasmi(Trap(Trap { kind: Host(Other(\"Function `missing_external` is only a stub. Calling a stub is not allowed.\")) }))"
+					"\"Trap: Trap { kind: Host(Other(\\\"Function `missing_external` is only a stub. Calling a stub is not allowed.\\\")) }\""
 				),
 				#[cfg(feature = "wasmtime")]
 				WasmExecutionMethod::Compiled => assert_eq!(
 					&format!("{:?}", e),
-					"Other(\"Wasm execution trapped: call to a missing function env:missing_external\")"
+					"\"Wasm execution trapped: call to a missing function env:missing_external\""
 				),
 			}
 		}
@@ -113,12 +116,12 @@ fn call_yet_another_not_existing_function(wasm_method: WasmExecutionMethod) {
 			match wasm_method {
 				WasmExecutionMethod::Interpreted => assert_eq!(
 					&format!("{:?}", e),
-					"Wasmi(Trap(Trap { kind: Host(Other(\"Function `yet_another_missing_external` is only a stub. Calling a stub is not allowed.\")) }))"
+					"\"Trap: Trap { kind: Host(Other(\\\"Function `yet_another_missing_external` is only a stub. Calling a stub is not allowed.\\\")) }\""
 				),
 				#[cfg(feature = "wasmtime")]
 				WasmExecutionMethod::Compiled => assert_eq!(
 					&format!("{:?}", e),
-					"Other(\"Wasm execution trapped: call to a missing function env:yet_another_missing_external\")"
+					"\"Wasm execution trapped: call to a missing function env:yet_another_missing_external\""
 				),
 			}
 		}
@@ -502,29 +505,32 @@ fn offchain_http_should_work(wasm_method: WasmExecutionMethod) {
 fn should_trap_when_heap_exhausted(wasm_method: WasmExecutionMethod) {
 	let mut ext = TestExternalities::default();
 
-	crate::call_in_wasm::<HostFunctions>(
+	let executor = crate::WasmExecutor::new(
+		wasm_method,
+		Some(17),  // `17` is the initial number of pages compiled into the binary.
+		HostFunctions::host_functions(),
+		true,
+	);
+	executor.call_in_wasm(
+		&WASM_BINARY[..],
 		"test_exhaust_heap",
 		&[0],
-		wasm_method,
 		&mut ext.ext(),
-		&WASM_BINARY[..],
-		// `17` is the initial number of pages compiled into the binary.
-		17,
-		true,
 	).unwrap();
 }
 
 #[test_case(WasmExecutionMethod::Interpreted)]
 #[cfg_attr(feature = "wasmtime", test_case(WasmExecutionMethod::Compiled))]
 fn returns_mutable_static(wasm_method: WasmExecutionMethod) {
-	let mut instance = crate::wasm_runtime::create_wasm_runtime_with_code(
+	let runtime = crate::wasm_runtime::create_wasm_runtime_with_code(
 		wasm_method,
 		1024,
 		&WASM_BINARY[..],
 		HostFunctions::host_functions(),
 		true,
-	).expect("Creates instance");
+	).expect("Creates runtime");
 
+	let instance = runtime.new_instance().unwrap();
 	let res = instance.call("returns_mutable_static", &[0]).unwrap();
 	assert_eq!(33, u64::decode(&mut &res[..]).unwrap());
 
@@ -550,13 +556,14 @@ fn restoration_of_globals(wasm_method: WasmExecutionMethod) {
 	// to our allocator algorithm there are inefficiencies.
 	const REQUIRED_MEMORY_PAGES: u64 = 32;
 
-	let mut instance = crate::wasm_runtime::create_wasm_runtime_with_code(
+	let runtime = crate::wasm_runtime::create_wasm_runtime_with_code(
 		wasm_method,
 		REQUIRED_MEMORY_PAGES,
 		&WASM_BINARY[..],
 		HostFunctions::host_functions(),
 		true,
-	).expect("Creates instance");
+	).expect("Creates runtime");
+	let instance = runtime.new_instance().unwrap();
 
 	// On the first invocation we allocate approx. 768KB (75%) of stack and then trap.
 	let res = instance.call("allocates_huge_stack_array", &true.encode());
@@ -569,15 +576,16 @@ fn restoration_of_globals(wasm_method: WasmExecutionMethod) {
 
 #[test_case(WasmExecutionMethod::Interpreted)]
 fn heap_is_reset_between_calls(wasm_method: WasmExecutionMethod) {
-	let mut instance = crate::wasm_runtime::create_wasm_runtime_with_code(
+	let runtime = crate::wasm_runtime::create_wasm_runtime_with_code(
 		wasm_method,
 		1024,
 		&WASM_BINARY[..],
 		HostFunctions::host_functions(),
 		true,
-	).expect("Creates instance");
+	).expect("Creates runtime");
+	let instance = runtime.new_instance().unwrap();
 
-	let heap_base = instance.get_global_val("__heap_base")
+	let heap_base = instance.get_global_const("__heap_base")
 		.expect("`__heap_base` is valid")
 		.expect("`__heap_base` exists")
 		.as_i32()

substrate/client/executor/src/lib.rs

@@ -36,82 +36,24 @@ mod wasm_runtime;
 mod integration_tests;
 
 pub use wasmi;
-pub use native_executor::{with_externalities_safe, NativeExecutor, NativeExecutionDispatch};
+pub use native_executor::{with_externalities_safe, NativeExecutor, WasmExecutor, NativeExecutionDispatch};
 pub use sp_version::{RuntimeVersion, NativeVersion};
 pub use codec::Codec;
 #[doc(hidden)]
-pub use sp_core::traits::Externalities;
+pub use sp_core::traits::{Externalities, CallInWasm};
 #[doc(hidden)]
 pub use sp_wasm_interface;
 pub use wasm_runtime::WasmExecutionMethod;
 
 pub use sc_executor_common::{error, sandbox};
 
-/// Call the given `function` in the given wasm `code`.
-///
-/// The signature of `function` needs to follow the default Substrate function signature.
-///
-/// - `call_data`: Will be given as input parameters to `function`
-/// - `execution_method`: The execution method to use.
-/// - `ext`: The externalities that should be set while executing the wasm function.
-///          If `None` is given, no externalities will be set.
-/// - `heap_pages`: The number of heap pages to allocate.
-///
-/// Returns the `Vec<u8>` that contains the return value of the function.
-pub fn call_in_wasm<HF: sp_wasm_interface::HostFunctions>(
-	function: &str,
-	call_data: &[u8],
-	execution_method: WasmExecutionMethod,
-	ext: &mut dyn Externalities,
-	code: &[u8],
-	heap_pages: u64,
-	allow_missing_func_imports: bool,
-) -> error::Result<Vec<u8>> {
-	call_in_wasm_with_host_functions(
-		function,
-		call_data,
-		execution_method,
-		ext,
-		code,
-		heap_pages,
-		HF::host_functions(),
-		allow_missing_func_imports,
-	)
-}
-
-/// Non-generic version of [`call_in_wasm`] that takes the `host_functions` as parameter.
-/// For more information please see [`call_in_wasm`].
-pub fn call_in_wasm_with_host_functions(
-	function: &str,
-	call_data: &[u8],
-	execution_method: WasmExecutionMethod,
-	ext: &mut dyn Externalities,
-	code: &[u8],
-	heap_pages: u64,
-	host_functions: Vec<&'static dyn sp_wasm_interface::Function>,
-	allow_missing_func_imports: bool,
-) -> error::Result<Vec<u8>> {
-	let instance = wasm_runtime::create_wasm_runtime_with_code(
-		execution_method,
-		heap_pages,
-		code,
-		host_functions,
-		allow_missing_func_imports,
-	)?;
-
-	// It is safe, as we delete the instance afterwards.
-	let mut instance = std::panic::AssertUnwindSafe(instance);
-
-	with_externalities_safe(ext, move || instance.call(function, call_data)).and_then(|r| r)
-}
-
 /// Provides runtime information.
 pub trait RuntimeInfo {
 	/// Native runtime information.
 	fn native_version(&self) -> &NativeVersion;
 
 	/// Extract RuntimeVersion of given :code block
-	fn runtime_version<E: Externalities> (&self, ext: &mut E) -> error::Result<RuntimeVersion>;
+	fn runtime_version(&self, ext: &mut dyn Externalities) -> error::Result<RuntimeVersion>;
 }
 
 #[cfg(test)]
@@ -119,19 +61,25 @@ mod tests {
 	use super::*;
 	use sc_runtime_test::WASM_BINARY;
 	use sp_io::TestExternalities;
+	use sp_wasm_interface::HostFunctions;
+	use sp_core::traits::CallInWasm;
 
 	#[test]
 	fn call_in_interpreted_wasm_works() {
 		let mut ext = TestExternalities::default();
 		let mut ext = ext.ext();
-		let res = call_in_wasm::<sp_io::SubstrateHostFunctions>(
+
+		let executor = WasmExecutor::new(
+			WasmExecutionMethod::Interpreted,
+			Some(8),
+			sp_io::SubstrateHostFunctions::host_functions(),
+			true,
+		);
+		let res = executor.call_in_wasm(
+			&WASM_BINARY[..],
 			"test_empty_return",
 			&[],
-			WasmExecutionMethod::Interpreted,
 			&mut ext,
-			&WASM_BINARY,
-			8,
-			true,
 		).unwrap();
 		assert_eq!(res, vec![0u8; 0]);
 	}

substrate/client/executor/src/native_executor.rs

@@ -16,19 +16,15 @@
 
 use crate::{
 	RuntimeInfo, error::{Error, Result},
-	wasm_runtime::{RuntimesCache, WasmExecutionMethod},
+	wasm_runtime::{RuntimeCache, WasmExecutionMethod, CodeSource},
 };
 use sp_version::{NativeVersion, RuntimeVersion};
 use codec::{Decode, Encode};
 use sp_core::{NativeOrEncoded, traits::{CodeExecutor, Externalities}};
 use log::trace;
-use std::{result, cell::RefCell, panic::{UnwindSafe, AssertUnwindSafe}, sync::Arc};
+use std::{result, panic::{UnwindSafe, AssertUnwindSafe}, sync::Arc};
 use sp_wasm_interface::{HostFunctions, Function};
-use sc_executor_common::wasm_runtime::WasmRuntime;
-
-thread_local! {
-	static RUNTIMES_CACHE: RefCell<RuntimesCache> = RefCell::new(RuntimesCache::new());
-}
+use sc_executor_common::wasm_runtime::WasmInstance;
 
 /// Default num of pages for the heap
 const DEFAULT_HEAP_PAGES: u64 = 1024;
@@ -75,42 +71,43 @@ pub trait NativeExecutionDispatch: Send + Sync {
 	fn native_version() -> NativeVersion;
 }
 
-/// A generic `CodeExecutor` implementation that uses a delegate to determine wasm code equivalence
-/// and dispatch to native code when possible, falling back on `WasmExecutor` when not.
-pub struct NativeExecutor<D> {
-	/// Dummy field to avoid the compiler complaining about us not using `D`.
-	_dummy: std::marker::PhantomData<D>,
+/// An abstraction over Wasm code executor. Supports selecting execution backend and
+/// manages runtime cache.
+#[derive(Clone)]
+pub struct WasmExecutor {
 	/// Method used to execute fallback Wasm code.
-	fallback_method: WasmExecutionMethod,
-	/// Native runtime version info.
-	native_version: NativeVersion,
+	method: WasmExecutionMethod,
 	/// The number of 64KB pages to allocate for Wasm execution.
 	default_heap_pages: u64,
 	/// The host functions registered with this instance.
 	host_functions: Arc<Vec<&'static dyn Function>>,
+	/// WASM runtime cache.
+	cache: Arc<RuntimeCache>,
+	/// Allow missing function imports.
+	allow_missing_func_imports: bool,
 }
 
-impl<D: NativeExecutionDispatch> NativeExecutor<D> {
+impl WasmExecutor {
 	/// Create new instance.
 	///
 	/// # Parameters
 	///
-	/// `fallback_method` - Method used to execute fallback Wasm code.
+	/// `method` - Method used to execute Wasm code.
 	///
 	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
 	/// 	Defaults to `DEFAULT_HEAP_PAGES` if `None` is provided.
-	pub fn new(fallback_method: WasmExecutionMethod, default_heap_pages: Option<u64>) -> Self {
-		let mut host_functions = sp_io::SubstrateHostFunctions::host_functions();
-
-		// Add the custom host functions provided by the user.
-		host_functions.extend(D::ExtendHostFunctions::host_functions());
-
-		NativeExecutor {
-			_dummy: Default::default(),
-			fallback_method,
-			native_version: D::native_version(),
+	pub fn new(
+		method: WasmExecutionMethod,
+		default_heap_pages: Option<u64>,
+		host_functions: Vec<&'static dyn Function>,
+		allow_missing_func_imports: bool,
+	) -> Self {
+		WasmExecutor {
+			method,
 			default_heap_pages: default_heap_pages.unwrap_or(DEFAULT_HEAP_PAGES),
 			host_functions: Arc::new(host_functions),
+			cache: Arc::new(RuntimeCache::new()),
+			allow_missing_func_imports,
 		}
 	}
 
@@ -127,46 +124,90 @@ impl<D: NativeExecutionDispatch> NativeExecutor<D> {
 	/// runtime is invalidated on any `panic!` to prevent a poisoned state. `ext` is already
 	/// implicitly handled as unwind safe, as we store it in a global variable while executing the
 	/// native runtime.
-	fn with_runtime<E, R>(
+	fn with_instance<'c, R, F>(
 		&self,
-		ext: &mut E,
-		f: impl for<'a> FnOnce(
-			AssertUnwindSafe<&'a mut (dyn WasmRuntime + 'static)>,
-			&'a RuntimeVersion,
-			AssertUnwindSafe<&'a mut E>,
+		code: CodeSource<'c>,
+		ext: &mut dyn Externalities,
+		f: F,
+	) -> Result<R>
+		where F: FnOnce(
+			AssertUnwindSafe<&dyn WasmInstance>,
+			Option<&RuntimeVersion>,
+			AssertUnwindSafe<&mut dyn Externalities>,
 		) -> Result<Result<R>>,
-	) -> Result<R> where E: Externalities {
-		RUNTIMES_CACHE.with(|cache| {
-			let mut cache = cache.borrow_mut();
-			let (runtime, version, code_hash) = cache.fetch_runtime(
-				ext,
-				self.fallback_method,
-				self.default_heap_pages,
-				&*self.host_functions,
-			)?;
-
-			let runtime = AssertUnwindSafe(runtime);
-			let ext = AssertUnwindSafe(ext);
-
-			match f(runtime, version, ext) {
-				Ok(res) => res,
-				Err(e) => {
-					cache.invalidate_runtime(self.fallback_method, code_hash);
-					Err(e)
-				}
+	{
+		match self.cache.with_instance(
+			code,
+			ext,
+			self.method,
+			self.default_heap_pages,
+			&*self.host_functions,
+			self.allow_missing_func_imports,
+			|instance, version, ext| {
+				let instance = AssertUnwindSafe(instance);
+				let ext = AssertUnwindSafe(ext);
+				f(instance, version, ext)
 			}
-		})
+		)? {
+			Ok(r) => r,
+			Err(e) => Err(e),
+		}
 	}
 }
 
-impl<D: NativeExecutionDispatch> Clone for NativeExecutor<D> {
-	fn clone(&self) -> Self {
+impl sp_core::traits::CallInWasm for WasmExecutor {
+	fn call_in_wasm(
+		&self,
+		wasm_blob: &[u8],
+		method: &str,
+		call_data: &[u8],
+		ext: &mut dyn Externalities,
+	) -> std::result::Result<Vec<u8>, String> {
+		self.with_instance(CodeSource::Custom(wasm_blob), ext, |instance, _, mut ext| {
+			with_externalities_safe(
+				&mut **ext,
+				move || instance.call(method, call_data),
+			)
+		}).map_err(|e| e.to_string())
+	}
+}
+
+/// A generic `CodeExecutor` implementation that uses a delegate to determine wasm code equivalence
+/// and dispatch to native code when possible, falling back on `WasmExecutor` when not.
+pub struct NativeExecutor<D> {
+	/// Dummy field to avoid the compiler complaining about us not using `D`.
+	_dummy: std::marker::PhantomData<D>,
+	/// Native runtime version info.
+	native_version: NativeVersion,
+	/// Fallback wasm executor.
+	wasm: WasmExecutor,
+}
+
+impl<D: NativeExecutionDispatch> NativeExecutor<D> {
+	/// Create new instance.
+	///
+	/// # Parameters
+	///
+	/// `fallback_method` - Method used to execute fallback Wasm code.
+	///
+	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
+	/// 	Defaults to `DEFAULT_HEAP_PAGES` if `None` is provided.
+	pub fn new(fallback_method: WasmExecutionMethod, default_heap_pages: Option<u64>) -> Self {
+		let mut host_functions = sp_io::SubstrateHostFunctions::host_functions();
+
+		// Add the custom host functions provided by the user.
+		host_functions.extend(D::ExtendHostFunctions::host_functions());
+		let wasm_executor = WasmExecutor::new(
+			fallback_method,
+			default_heap_pages,
+			host_functions,
+			false,
+		);
+
 		NativeExecutor {
 			_dummy: Default::default(),
-			fallback_method: self.fallback_method,
 			native_version: D::native_version(),
-			default_heap_pages: self.default_heap_pages,
-			host_functions: self.host_functions.clone(),
+			wasm: wasm_executor,
 		}
 	}
 }
@@ -176,90 +217,109 @@ impl<D: NativeExecutionDispatch> RuntimeInfo for NativeExecutor<D> {
 		&self.native_version
 	}
 
-	fn runtime_version<E: Externalities>(
+	fn runtime_version(
 		&self,
-		ext: &mut E,
+		ext: &mut dyn Externalities,
 	) -> Result<RuntimeVersion> {
-		self.with_runtime(ext, |_runtime, version, _ext| Ok(Ok(version.clone())))
+		self.wasm.with_instance(CodeSource::Externalities, ext,
+			|_instance, version, _ext|
+				Ok(version.cloned().ok_or_else(|| Error::ApiError("Unknown version".into())))
+		)
 	}
 }
 
 impl<D: NativeExecutionDispatch + 'static> CodeExecutor for NativeExecutor<D> {
 	type Error = Error;
 
-	fn call
-	<
-		E: Externalities,
+	fn call<
 		R: Decode + Encode + PartialEq,
 		NC: FnOnce() -> result::Result<R, String> + UnwindSafe,
 	>(
 		&self,
-		ext: &mut E,
+		ext: &mut dyn Externalities,
 		method: &str,
 		data: &[u8],
 		use_native: bool,
 		native_call: Option<NC>,
-	) -> (Result<NativeOrEncoded<R>>, bool){
+	) -> (Result<NativeOrEncoded<R>>, bool) {
 		let mut used_native = false;
-		let result = self.with_runtime(ext, |mut runtime, onchain_version, mut ext| {
-			match (
-				use_native,
-				onchain_version.can_call_with(&self.native_version.runtime_version),
-				native_call,
-			) {
-				(_, false, _) => {
-					trace!(
-						target: "executor",
-						"Request for native execution failed (native: {}, chain: {})",
-						self.native_version.runtime_version,
-						onchain_version,
-					);
-
-					with_externalities_safe(
-						&mut **ext,
-						move || runtime.call(method, data).map(NativeOrEncoded::Encoded)
-					)
-				}
-				(false, _, _) => {
-					with_externalities_safe(
-						&mut **ext,
-						move || runtime.call(method, data).map(NativeOrEncoded::Encoded)
-					)
-				},
-				(true, true, Some(call)) => {
-					trace!(
-						target: "executor",
-						"Request for native execution with native call succeeded (native: {}, chain: {}).",
-						self.native_version.runtime_version,
-						onchain_version,
-					);
-
-					used_native = true;
-					let res = with_externalities_safe(&mut **ext, move || (call)())
-						.and_then(|r| r
-							.map(NativeOrEncoded::Native)
-							.map_err(|s| Error::ApiError(s.to_string()))
+		let result = self.wasm.with_instance(
+			CodeSource::Externalities,
+			ext,
+			|instance, onchain_version, mut ext| {
+				let onchain_version = onchain_version.ok_or_else(
+					|| Error::ApiError("Unknown version".into())
+				)?;
+				match (
+					use_native,
+					onchain_version.can_call_with(&self.native_version.runtime_version),
+					native_call,
+				) {
+					(_, false, _) => {
+						trace!(
+							target: "executor",
+							"Request for native execution failed (native: {}, chain: {})",
+							self.native_version.runtime_version,
+							onchain_version,
 						);
 
-					Ok(res)
-				}
-				_ => {
-					trace!(
-						target: "executor",
-						"Request for native execution succeeded (native: {}, chain: {})",
-						self.native_version.runtime_version,
-						onchain_version
-					);
+						with_externalities_safe(
+							&mut **ext,
+							move || instance.call(method, data).map(NativeOrEncoded::Encoded)
+						)
+					}
+					(false, _, _) => {
+						with_externalities_safe(
+							&mut **ext,
+							move || instance.call(method, data).map(NativeOrEncoded::Encoded)
+						)
+					},
+					(true, true, Some(call)) => {
+						trace!(
+							target: "executor",
+							"Request for native execution with native call succeeded \
+							(native: {}, chain: {}).",
+							self.native_version.runtime_version,
+							onchain_version,
+						);
 
-					used_native = true;
-					Ok(D::dispatch(&mut **ext, method, data).map(NativeOrEncoded::Encoded))
+						used_native = true;
+						let res = with_externalities_safe(&mut **ext, move || (call)())
+							.and_then(|r| r
+								.map(NativeOrEncoded::Native)
+								.map_err(|s| Error::ApiError(s.to_string()))
+							);
+
+						Ok(res)
+					}
+					_ => {
+						trace!(
+							target: "executor",
+							"Request for native execution succeeded (native: {}, chain: {})",
+							self.native_version.runtime_version,
+							onchain_version
+						);
+
+						used_native = true;
+						Ok(D::dispatch(&mut **ext, method, data).map(NativeOrEncoded::Encoded))
+					}
 				}
 			}
-		});
+		);
 		(result, used_native)
 	}
 }
 
+impl<D: NativeExecutionDispatch> Clone for NativeExecutor<D> {
+	fn clone(&self) -> Self {
+		NativeExecutor {
+			_dummy: Default::default(),
+			native_version: D::native_version(),
+			wasm: self.wasm.clone(),
+		}
+	}
+}
+
 impl<D: NativeExecutionDispatch> sp_core::traits::CallInWasm for NativeExecutor<D> {
 	fn call_in_wasm(
 		&self,
@@ -268,16 +328,7 @@ impl<D: NativeExecutionDispatch> sp_core::traits::CallInWasm for NativeExecutor<
 		call_data: &[u8],
 		ext: &mut dyn Externalities,
 	) -> std::result::Result<Vec<u8>, String> {
-		crate::call_in_wasm_with_host_functions(
-			method,
-			call_data,
-			self.fallback_method,
-			ext,
-			wasm_blob,
-			self.default_heap_pages,
-			(*self.host_functions).clone(),
-			false,
-		).map_err(|e| e.to_string())
+		sp_core::traits::CallInWasm::call_in_wasm(&self.wasm, wasm_blob, method, call_data, ext)
 	}
 }
 
@@ -379,7 +430,7 @@ mod tests {
 		let executor = NativeExecutor::<MyExecutor>::new(WasmExecutionMethod::Interpreted, None);
 		my_interface::HostFunctions::host_functions().iter().for_each(|function| {
 			assert_eq!(
-				executor.host_functions.iter().filter(|f| f == &function).count(),
+				executor.wasm.host_functions.iter().filter(|f| f == &function).count(),
 				2,
 			);
 		});

substrate/client/executor/src/wasm_runtime.rs

@@ -19,13 +19,15 @@
 //! The primary means of accessing the runtimes is through a cache which saves the reusable
 //! components of the runtime that are expensive to initialize.
 
+use std::sync::Arc;
+use std::borrow::Cow;
 use crate::error::{Error, WasmError};
-use log::{trace, warn};
+use parking_lot::{Mutex, RwLock};
 use codec::Decode;
 use sp_core::{storage::well_known_keys, traits::Externalities};
 use sp_version::RuntimeVersion;
-use std::{collections::hash_map::{Entry, HashMap}, panic::AssertUnwindSafe};
-use sc_executor_common::wasm_runtime::WasmRuntime;
+use std::panic::AssertUnwindSafe;
+use sc_executor_common::wasm_runtime::{WasmModule, WasmInstance};
 
 use sp_wasm_interface::Function;
 
@@ -39,15 +41,33 @@ pub enum WasmExecutionMethod {
 	Compiled,
 }
 
+/// Executoed code origin.
+pub enum CodeSource<'a> {
+	/// Take code from storage,
+	Externalities,
+	/// Use provided code,
+	Custom(&'a [u8]),
+}
+
 /// A Wasm runtime object along with its cached runtime version.
 struct VersionedRuntime {
-	runtime: Box<dyn WasmRuntime>,
+	/// Runtime code hash.
+	code_hash: Vec<u8>,
+	/// Wasm runtime type.
+	wasm_method: WasmExecutionMethod,
+	/// Shared runtime that can spawn instances.
+	module: Box<dyn WasmModule>,
 	/// The number of WebAssembly heap pages this instance was created with.
 	heap_pages: u64,
-	/// Runtime version according to `Core_version`.
-	version: RuntimeVersion,
+	/// Runtime version according to `Core_version` if any.
+	version: Option<RuntimeVersion>,
+	/// Cached instance pool.
+	instances: RwLock<[Option<Arc<Mutex<Box<dyn WasmInstance>>>>; MAX_INSTANCES]>,
 }
 
+const MAX_RUNTIMES: usize = 2;
+const MAX_INSTANCES: usize = 8;
+
 /// Cache for the runtimes.
 ///
 /// When an instance is requested for the first time it is added to this cache. Metadata is kept
@@ -60,130 +80,184 @@ struct VersionedRuntime {
 ///
 /// 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.
+pub struct RuntimeCache {
+	/// A cache of runtimes along with metadata.
 	///
-	/// Instances are keyed by the Wasm execution method and the hash of their code.
-	instances: HashMap<(WasmExecutionMethod, Vec<u8>), Result<VersionedRuntime, WasmError>>,
+	/// Runtimes sorted by recent usage. The most recently used is at the front.
+	runtimes: Mutex<[Option<Arc<VersionedRuntime>>; MAX_RUNTIMES]>,
 }
 
-impl RuntimesCache {
+impl RuntimeCache {
 	/// Creates a new instance of a runtimes cache.
-	pub fn new() -> RuntimesCache {
-		RuntimesCache {
-			instances: HashMap::new(),
+	pub fn new() -> RuntimeCache {
+		RuntimeCache {
+			runtimes: Default::default(),
 		}
 	}
 
-	/// 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.
+	/// Prepares a WASM module instance and executes given function for it.
 	///
+	/// This uses internal cache to find avaiable instance or create a new one.
 	/// # Parameters
 	///
+	/// `code` - Provides external code or tells the executor to fetch it from storage.
+	///
 	/// `ext` - Externalities to use for the runtime. This is used for setting
 	/// up an initial runtime instance.
 	///
 	/// `default_heap_pages` - Number of 64KB pages to allocate for Wasm execution.
 	///
+	/// `wasm_method` - Type of WASM backend to use.
+	///
 	/// `host_functions` - The host functions that should be registered for the Wasm runtime.
 	///
-	/// # Return value
+	/// `allow_missing_func_imports` - Ignore missing function imports.
 	///
-	/// If no error occurred a tuple `(&mut WasmRuntime, H256)` is
-	/// returned. `H256` is the hash of the runtime code.
+	/// `f` - Function to execute.
 	///
+	/// # Returns result of `f` wrapped in an additonal result.
 	/// 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>(
-		&mut self,
-		ext: &mut E,
+	pub fn with_instance<'c, R, F>(
+		&self,
+		code: CodeSource<'c>,
+		ext: &mut dyn Externalities,
 		wasm_method: WasmExecutionMethod,
 		default_heap_pages: u64,
 		host_functions: &[&'static dyn Function],
-	) -> Result<(&mut (dyn WasmRuntime + 'static), &RuntimeVersion, Vec<u8>), Error> {
-		let code_hash = ext
-			.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())
-			.unwrap_or(default_heap_pages);
-
-		let result = match self.instances.entry((wasm_method, code_hash.clone())) {
-			Entry::Occupied(o) => {
-				let result = o.into_mut();
-				if let Ok(ref mut cached_runtime) = result {
-					let heap_pages_changed = cached_runtime.heap_pages != heap_pages;
-					let host_functions_changed = cached_runtime.runtime.host_functions()
-						!= host_functions;
-					if heap_pages_changed || host_functions_changed {
-						let changed = if heap_pages_changed {
-							"heap_pages"
-						} else {
-							"host functions"
-						};
-
-						trace!(
-							target: "runtimes_cache",
-							"{} were changed. Reinstantiating the instance",
-							changed,
-						);
-						*result = create_versioned_wasm_runtime(
-							ext,
-							wasm_method,
-							heap_pages,
-							host_functions.into(),
-						);
-						if let Err(ref err) = result {
-							warn!(target: "runtimes_cache", "cannot create a runtime: {:?}", err);
-						}
-					}
-				}
-				result
+		allow_missing_func_imports: bool,
+		f: F,
+	) -> Result<Result<R, Error>, Error>
+		where F: FnOnce(
+			&dyn WasmInstance,
+			Option<&RuntimeVersion>,
+			&mut dyn Externalities)
+		-> Result<R, Error>,
+	{
+		let (code_hash, heap_pages) = match &code {
+			CodeSource::Externalities => {
+				(
+					ext
+						.original_storage_hash(well_known_keys::CODE)
+						.ok_or(Error::InvalidCode("`CODE` not found in storage.".into()))?,
+					ext
+						.storage(well_known_keys::HEAP_PAGES)
+						.and_then(|pages| u64::decode(&mut &pages[..]).ok())
+						.unwrap_or(default_heap_pages),
+				)
 			},
-			Entry::Vacant(v) => {
-				trace!(target: "runtimes_cache", "no instance found in cache, creating now.");
+			CodeSource::Custom(code) => {
+				(sp_core::blake2_256(code).to_vec(), default_heap_pages)
+			}
+		};
+
+		let mut runtimes = self.runtimes.lock(); // this must be released prior to calling f
+		let pos = runtimes.iter().position(|r| r.as_ref().map_or(
+			false,
+			|r| r.wasm_method == wasm_method &&
+				r.code_hash == code_hash &&
+				r.heap_pages == heap_pages
+		));
+
+		let runtime = match pos {
+			Some(n) => runtimes[n]
+				.clone()
+				.expect("`position` only returns `Some` for entries that are `Some`"),
+			None =>  {
+				let code = match code {
+					CodeSource::Externalities => {
+						Cow::Owned(ext.original_storage(well_known_keys::CODE)
+							.ok_or(WasmError::CodeNotFound)?)
+					}
+					CodeSource::Custom(code) => {
+						Cow::Borrowed(code)
+					}
+				};
+
 				let result = create_versioned_wasm_runtime(
+					&code,
+					code_hash,
 					ext,
 					wasm_method,
 					heap_pages,
 					host_functions.into(),
+					allow_missing_func_imports,
 				);
 				if let Err(ref err) = result {
-					warn!(target: "runtimes_cache", "cannot create a runtime: {:?}", err);
+					log::warn!(target: "wasm-runtime", "Cannot create a runtime: {:?}", err);
 				}
-				v.insert(result)
+				Arc::new(result?)
 			}
 		};
 
-		result.as_mut()
-			.map(|entry| (entry.runtime.as_mut(), &entry.version, code_hash))
-			.map_err(|ref e| Error::InvalidCode(format!("{:?}", e)))
-	}
+		// Rearrange runtimes by last recently used.
+		match pos {
+			Some(0) => {},
+			Some(n) => {
+				for i in (1 .. n + 1).rev() {
+					runtimes.swap(i, i - 1);
+				}
+			}
+			None => {
+				runtimes[MAX_RUNTIMES-1] = Some(runtime.clone());
+				for i in (1 .. MAX_RUNTIMES).rev() {
+					runtimes.swap(i, i - 1);
+				}
+			}
+		}
+		drop(runtimes);
 
-	/// Invalidate the runtime for the given `wasm_method` and `code_hash`.
-	///
-	/// Invalidation of a runtime is useful when there was a `panic!` in native while executing it.
-	/// The `panic!` maybe have brought the runtime into a poisoned state and so, it is better to
-	/// invalidate this runtime instance.
-	pub fn invalidate_runtime(
-		&mut self,
-		wasm_method: WasmExecutionMethod,
-		code_hash: Vec<u8>,
-	) {
-		// Just remove the instance, it will be re-created the next time it is requested.
-		self.instances.remove(&(wasm_method, code_hash));
+		let result = {
+			// Find a free instance
+			let instance_pool = runtime.instances.read().clone();
+			let instance = instance_pool
+				.iter()
+				.find_map(|i| i.as_ref().and_then(|i| i.try_lock()));
+			if let Some(mut locked) = instance {
+				let result = f(&**locked, runtime.version.as_ref(), ext);
+				if let Err(e) = &result {
+					log::warn!(target: "wasm-runtime", "Evicting failed runtime instance: {:?}", e);
+					*locked = runtime.module.new_instance()?;
+				}
+				result
+			} else {
+				// Allocate a new instance
+				let instance = runtime.module.new_instance()?;
+
+				let result = f(&*instance, runtime.version.as_ref(), ext);
+				match &result {
+					Ok(_) => {
+						let mut instance_pool = runtime.instances.write();
+						if let Some(ref mut slot) = instance_pool.iter_mut().find(|s| s.is_none()) {
+							**slot = Some(Arc::new(Mutex::new(instance)));
+							log::debug!(
+								target: "wasm-runtime",
+								"Allocated WASM instance {}/{}",
+								instance_pool.len(),
+								MAX_INSTANCES,
+							);
+						} else {
+							log::warn!(target: "wasm-runtime", "Ran out of free WASM instances");
+						}
+					}
+					Err(e) => {
+						log::warn!(
+							target:
+							"wasm-runtime",
+							"Fresh runtime instance failed with {:?}",
+							e,
+						);
+					}
+				}
+				result
+			}
+		};
+
+		Ok(result)
 	}
 }
 
@@ -194,28 +268,43 @@ pub fn create_wasm_runtime_with_code(
 	code: &[u8],
 	host_functions: Vec<&'static dyn Function>,
 	allow_missing_func_imports: bool,
-) -> Result<Box<dyn WasmRuntime>, WasmError> {
+) -> Result<Box<dyn WasmModule>, WasmError> {
 	match wasm_method {
 		WasmExecutionMethod::Interpreted =>
-			sc_executor_wasmi::create_instance(code, heap_pages, host_functions, allow_missing_func_imports)
-				.map(|runtime| -> Box<dyn WasmRuntime> { Box::new(runtime) }),
+			sc_executor_wasmi::create_runtime(
+				code,
+				heap_pages,
+				host_functions,
+				allow_missing_func_imports
+			).map(|runtime| -> Box<dyn WasmModule> { Box::new(runtime) }),
 		#[cfg(feature = "wasmtime")]
 		WasmExecutionMethod::Compiled =>
-			sc_executor_wasmtime::create_instance(code, heap_pages, host_functions, allow_missing_func_imports)
-				.map(|runtime| -> Box<dyn WasmRuntime> { Box::new(runtime) }),
+			sc_executor_wasmtime::create_runtime(
+				code,
+				heap_pages,
+				host_functions,
+				allow_missing_func_imports
+			).map(|runtime| -> Box<dyn WasmModule> { Box::new(runtime) }),
 	}
 }
 
-fn create_versioned_wasm_runtime<E: Externalities>(
-	ext: &mut E,
+fn create_versioned_wasm_runtime(
+	code: &[u8],
+	code_hash: Vec<u8>,
+	ext: &mut dyn Externalities,
 	wasm_method: WasmExecutionMethod,
 	heap_pages: u64,
 	host_functions: Vec<&'static dyn Function>,
+	allow_missing_func_imports: bool,
 ) -> Result<VersionedRuntime, WasmError> {
-	let code = ext
-		.original_storage(well_known_keys::CODE)
-		.ok_or(WasmError::CodeNotFound)?;
-	let mut runtime = create_wasm_runtime_with_code(wasm_method, heap_pages, &code, host_functions, false)?;
+	let time = std::time::Instant::now();
+	let mut runtime = create_wasm_runtime_with_code(
+		wasm_method,
+		heap_pages,
+		&code,
+		host_functions,
+		allow_missing_func_imports,
+	)?;
 
 	// Call to determine runtime version.
 	let version_result = {
@@ -224,21 +313,33 @@ fn create_versioned_wasm_runtime<E: Externalities>(
 
 		// The following unwind safety assertion is OK because if the method call panics, the
 		// runtime will be dropped.
-		let mut runtime = AssertUnwindSafe(runtime.as_mut());
+		let runtime = AssertUnwindSafe(runtime.as_mut());
 		crate::native_executor::with_externalities_safe(
 			&mut **ext,
-			move || runtime.call("Core_version", &[])
+			move || runtime.new_instance()?.call("Core_version", &[])
 		).map_err(|_| WasmError::Instantiation("panic in call to get runtime version".into()))?
 	};
-	let encoded_version = version_result
-		.map_err(|e| WasmError::Instantiation(format!("failed to call \"Core_version\": {}", e)))?;
-	let version = RuntimeVersion::decode(&mut encoded_version.as_slice())
-		.map_err(|_| WasmError::Instantiation("failed to decode \"Core_version\" result".into()))?;
+	let version = match version_result {
+		Ok(version) => Some(RuntimeVersion::decode(&mut version.as_slice())
+			.map_err(|_|
+				WasmError::Instantiation("failed to decode \"Core_version\" result".into())
+			)?),
+		Err(_) => None,
+	};
+	log::debug!(
+		target: "wasm-runtime",
+		"Prepared new runtime version {:?} in {} ms.",
+		version,
+		time.elapsed().as_millis(),
+	);
 
 	Ok(VersionedRuntime {
-		runtime,
+		code_hash,
+		module: runtime,
 		version,
 		heap_pages,
+		wasm_method,
+		instances: Default::default(),
 	})
 }