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/common/src/wasm_runtime.rs

@@ -17,18 +17,25 @@
 //! Definitions for a wasm runtime.
 
 use crate::error::Error;
-use sp_wasm_interface::{Function, Value};
+use sp_wasm_interface::Value;
 
-/// A trait that defines an abstract wasm runtime.
+/// A trait that defines an abstract WASM runtime module.
 ///
 /// This can be implemented by an execution engine.
-pub trait WasmRuntime {
-	/// Return the host functions that are registered for this Wasm runtime.
-	fn host_functions(&self) -> &[&'static dyn Function];
+pub trait WasmModule: Sync + Send {
+	/// Create a new instance.
+	fn new_instance(&self) -> Result<Box<dyn WasmInstance>, Error>;
+}
 
-	/// Call a method in the Substrate runtime by name. Returns the encoded result on success.
-	fn call(&mut self, method: &str, data: &[u8]) -> Result<Vec<u8>, Error>;
+/// A trait that defines an abstract wasm module instance.
+///
+/// This can be implemented by an execution engine.
+pub trait WasmInstance: Send {
+	/// Call a method on this WASM instance and reset it afterwards.
+	/// Returns the encoded result on success.
+	fn call(&self, method: &str, data: &[u8]) -> Result<Vec<u8>, Error>;
 
 	/// Get the value from a global with the given `name`.
-	fn get_global_val(&self, name: &str) -> Result<Option<Value>, Error>;
+	/// This method is only suitable for getting immutable globals.
+	fn get_global_const(&self, name: &str) -> Result<Option<Value>, Error>;
 }

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(),
 	})
 }
 

substrate/client/executor/wasmi/src/lib.rs

@@ -14,10 +14,10 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
-//! This crate provides an implementation of `WasmRuntime` that is baked by wasmi.
+//! This crate provides an implementation of `WasmModule` that is baked by wasmi.
 
 use sc_executor_common::{error::{Error, WasmError}, sandbox};
-use std::{str, mem, cell::RefCell};
+use std::{str, mem, cell::RefCell, sync::Arc};
 use wasmi::{
 	Module, ModuleInstance, MemoryInstance, MemoryRef, TableRef, ImportsBuilder, ModuleRef,
 	memory_units::Pages, RuntimeValue::{I32, I64, self},
@@ -30,7 +30,7 @@ use sp_wasm_interface::{
 	FunctionContext, Pointer, WordSize, Sandbox, MemoryId, Result as WResult, Function,
 };
 use sp_runtime_interface::unpack_ptr_and_len;
-use sc_executor_common::wasm_runtime::WasmRuntime;
+use sc_executor_common::wasm_runtime::{WasmModule, WasmInstance};
 
 struct FunctionExecutor<'a> {
 	sandbox_store: sandbox::Store<wasmi::FuncRef>,
@@ -623,9 +623,77 @@ impl StateSnapshot {
 	}
 }
 
-/// A runtime along with its initial state snapshot.
-#[derive(Clone)]
+/// A runtime along with initial copy of data segments.
 pub struct WasmiRuntime {
+	/// A wasm module.
+	module: Module,
+	/// The host functions registered for this instance.
+	host_functions: Arc<Vec<&'static dyn Function>>,
+	/// Enable stub generation for functions that are not available in `host_functions`.
+	/// These stubs will error when the wasm blob tries to call them.
+	allow_missing_func_imports: bool,
+	/// Numer of heap pages this runtime uses.
+	heap_pages: u64,
+	/// Data segments created for each new instance.
+	data_segments: Vec<DataSegment>,
+}
+
+impl WasmModule for WasmiRuntime {
+	fn new_instance(&self) -> Result<Box<dyn WasmInstance>, Error> {
+		// Instantiate this module.
+		let (instance, missing_functions, memory) = instantiate_module(
+			self.heap_pages as usize,
+			&self.module,
+			&self.host_functions,
+			self.allow_missing_func_imports,
+		).map_err(|e| WasmError::Instantiation(e.to_string()))?;
+
+		// Take state snapshot before executing anything.
+		let state_snapshot = StateSnapshot::take(&instance, self.data_segments.clone())
+			.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
+					",
+			);
+
+		Ok(Box::new(WasmiInstance {
+			instance,
+			memory,
+			state_snapshot,
+			host_functions: self.host_functions.clone(),
+			allow_missing_func_imports: self.allow_missing_func_imports,
+			missing_functions,
+		}))
+	}
+}
+
+/// Create a new `WasmiRuntime` given the code. This function loads the module and
+/// stores it in the instance.
+pub fn create_runtime(
+	code: &[u8],
+	heap_pages: u64,
+	host_functions: Vec<&'static dyn Function>,
+	allow_missing_func_imports: bool,
+) -> Result<WasmiRuntime, WasmError> {
+	let module = Module::from_buffer(&code).map_err(|_| WasmError::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(WasmiRuntime {
+		module,
+		data_segments,
+		host_functions: Arc::new(host_functions),
+		allow_missing_func_imports,
+		heap_pages,
+	})
+}
+
+/// Wasmi instance wrapper along with the state snapshot.
+pub struct WasmiInstance {
 	/// A wasm module instance.
 	instance: ModuleRef,
 	/// The memory instance of used by the wasm module.
@@ -633,7 +701,7 @@ pub struct WasmiRuntime {
 	/// The snapshot of the instance's state taken just after the instantiation.
 	state_snapshot: StateSnapshot,
 	/// The host functions registered for this instance.
-	host_functions: Vec<&'static dyn Function>,
+	host_functions: Arc<Vec<&'static dyn Function>>,
 	/// Enable stub generation for functions that are not available in `host_functions`.
 	/// These stubs will error when the wasm blob tries to call them.
 	allow_missing_func_imports: bool,
@@ -641,13 +709,12 @@ pub struct WasmiRuntime {
 	missing_functions: Vec<String>,
 }
 
-impl WasmRuntime for WasmiRuntime {
-	fn host_functions(&self) -> &[&'static dyn Function] {
-		&self.host_functions
-	}
+// This is safe because `WasmiInstance` does not leak any references to `self.memory` and `self.instance`
+unsafe impl Send for WasmiInstance {}
 
+impl WasmInstance for WasmiInstance {
 	fn call(
-		&mut self,
+		&self,
 		method: &str,
 		data: &[u8],
 	) -> Result<Vec<u8>, Error> {
@@ -664,67 +731,26 @@ impl WasmRuntime for WasmiRuntime {
 			&self.memory,
 			method,
 			data,
-			&self.host_functions,
+			self.host_functions.as_ref(),
 			self.allow_missing_func_imports,
-			&self.missing_functions,
+			self.missing_functions.as_ref(),
 		)
 	}
 
-	fn get_global_val(&self, name: &str) -> Result<Option<sp_wasm_interface::Value>, Error> {
+	fn get_global_const(&self, name: &str) -> Result<Option<sp_wasm_interface::Value>, Error> {
 		match self.instance.export_by_name(name) {
 			Some(global) => Ok(Some(
 				global
-					 .as_global()
-					 .ok_or_else(|| format!("`{}` is not a global", name))?
-					 .get()
-					 .into()
+					.as_global()
+					.ok_or_else(|| format!("`{}` is not a global", name))?
+					.get()
+					.into()
 			)),
 			None => Ok(None),
 		}
 	}
 }
 
-pub fn create_instance(
-	code: &[u8],
-	heap_pages: u64,
-	host_functions: Vec<&'static dyn Function>,
-	allow_missing_func_imports: bool,
-) -> Result<WasmiRuntime, WasmError> {
-	let module = Module::from_buffer(&code).map_err(|_| WasmError::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)?;
-
-	// Instantiate this module.
-	let (instance, missing_functions, memory) = instantiate_module(
-		heap_pages as usize,
-		&module,
-		&host_functions,
-		allow_missing_func_imports,
-	).map_err(|e| WasmError::Instantiation(e.to_string()))?;
-
-	// Take state snapshot before executing anything.
-	let state_snapshot = StateSnapshot::take(&instance, data_segments)
-		.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
-				",
-		);
-
-	Ok(WasmiRuntime {
-		instance,
-		memory,
-		state_snapshot,
-		host_functions,
-		allow_missing_func_imports,
-		missing_functions,
-	})
-}
-
 /// Extract the data segments from the given wasm code.
 ///
 /// Returns `Err` if the given wasm code cannot be deserialized.

substrate/client/executor/wasmtime/Cargo.toml

@@ -10,7 +10,7 @@ description = "Defines a `WasmRuntime` that uses the Wasmtime JIT to execute."
 
 [dependencies]
 log = "0.4.8"
-wasmi = "0.6.2"
+scoped-tls = "1.0"
 parity-wasm = "0.41.0"
 codec = { package = "parity-scale-codec", version = "1.2.0" }
 sc-executor-common = { version = "0.8.0-alpha.2", path = "../common" }
@@ -18,8 +18,7 @@ sp-wasm-interface = { version = "2.0.0-alpha.2", path = "../../../primitives/was
 sp-runtime-interface = { version = "2.0.0-alpha.2", path = "../../../primitives/runtime-interface" }
 sp-core = { version = "2.0.0-alpha.2", path = "../../../primitives/core" }
 sp-allocator = { version = "2.0.0-alpha.2", path = "../../../primitives/allocator" }
-
-wasmtime = "0.11"
+wasmtime = { git = "https://github.com/paritytech/wasmtime", branch = "a-thread-safe-api" }
 
 [dev-dependencies]
 assert_matches = "1.3.0"

substrate/client/executor/wasmtime/src/host.rs

@@ -19,7 +19,7 @@
 
 use crate::instance_wrapper::InstanceWrapper;
 use crate::util;
-use std::cell::RefCell;
+use std::{cell::RefCell, rc::Rc};
 use log::trace;
 use codec::{Encode, Decode};
 use sp_allocator::FreeingBumpHeapAllocator;
@@ -51,12 +51,12 @@ pub struct HostState {
 	// borrow after performing necessary queries/changes.
 	sandbox_store: RefCell<sandbox::Store<SupervisorFuncRef>>,
 	allocator: RefCell<FreeingBumpHeapAllocator>,
-	instance: InstanceWrapper,
+	instance: Rc<InstanceWrapper>,
 }
 
 impl HostState {
 	/// Constructs a new `HostState`.
-	pub fn new(allocator: FreeingBumpHeapAllocator, instance: InstanceWrapper) -> Self {
+	pub fn new(allocator: FreeingBumpHeapAllocator, instance: Rc<InstanceWrapper>) -> Self {
 		HostState {
 			sandbox_store: RefCell::new(sandbox::Store::new()),
 			allocator: RefCell::new(allocator),
@@ -64,11 +64,6 @@ impl HostState {
 		}
 	}
 
-	/// Destruct the host state and extract the `InstanceWrapper` passed at the creation.
-	pub fn into_instance(self) -> InstanceWrapper {
-		self.instance
-	}
-
 	/// Materialize `HostContext` that can be used to invoke a substrate host `dyn Function`.
 	pub fn materialize<'a>(&'a self) -> HostContext<'a> {
 		HostContext(self)

substrate/client/executor/wasmtime/src/imports.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
-use crate::state_holder::StateHolder;
+use crate::state_holder;
 use sc_executor_common::error::WasmError;
 use sp_wasm_interface::{Function, Value, ValueType};
 use std::any::Any;
@@ -34,7 +34,6 @@ pub struct Imports {
 /// Goes over all imports of a module and prepares a vector of `Extern`s that can be used for
 /// instantiation of the module. Returns an error if there are imports that cannot be satisfied.
 pub fn resolve_imports(
-	state_holder: &StateHolder,
 	module: &Module,
 	host_functions: &[&'static dyn Function],
 	heap_pages: u32,
@@ -58,7 +57,6 @@ pub fn resolve_imports(
 			}
 			_ => resolve_func_import(
 				module,
-				state_holder,
 				import_ty,
 				host_functions,
 				allow_missing_func_imports,
@@ -112,7 +110,6 @@ fn resolve_memory_import(
 
 fn resolve_func_import(
 	module: &Module,
-	state_holder: &StateHolder,
 	import_ty: &ImportType,
 	host_functions: &[&'static dyn Function],
 	allow_missing_func_imports: bool,
@@ -152,7 +149,7 @@ fn resolve_func_import(
 		)));
 	}
 
-	Ok(HostFuncHandler::new(&state_holder, *host_func).into_extern(module))
+	Ok(HostFuncHandler::new(*host_func).into_extern(module))
 }
 
 /// Returns `true` if `lhs` and `rhs` represent the same signature.
@@ -163,14 +160,12 @@ fn signature_matches(lhs: &wasmtime::FuncType, rhs: &wasmtime::FuncType) -> bool
 /// This structure implements `Callable` and acts as a bridge between wasmtime and
 /// substrate host functions.
 struct HostFuncHandler {
-	state_holder: StateHolder,
 	host_func: &'static dyn Function,
 }
 
 impl HostFuncHandler {
-	fn new(state_holder: &StateHolder, host_func: &'static dyn Function) -> Self {
+	fn new(host_func: &'static dyn Function) -> Self {
 		Self {
-			state_holder: state_holder.clone(),
 			host_func,
 		}
 	}
@@ -188,7 +183,7 @@ impl Callable for HostFuncHandler {
 		wasmtime_params: &[Val],
 		wasmtime_results: &mut [Val],
 	) -> Result<(), wasmtime::Trap> {
-		let unwind_result = self.state_holder.with_context(|host_ctx| {
+		let unwind_result = state_holder::with_context(|host_ctx| {
 			let mut host_ctx = host_ctx.expect(
 				"host functions can be called only from wasm instance;
 				wasm instance is always called initializing context;

substrate/client/executor/wasmtime/src/lib.rs

@@ -23,4 +23,4 @@ mod imports;
 mod instance_wrapper;
 mod util;
 
-pub use runtime::create_instance;
+pub use runtime::create_runtime;

substrate/client/executor/wasmtime/src/runtime.rs

@@ -15,57 +15,85 @@
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
 //! Defines the compiled Wasm runtime that uses Wasmtime internally.
+use std::rc::Rc;
+use std::sync::Arc;
 
 use crate::host::HostState;
-use crate::imports::{resolve_imports, Imports};
+use crate::imports::{Imports, resolve_imports};
 use crate::instance_wrapper::InstanceWrapper;
-use crate::state_holder::StateHolder;
+use crate::state_holder;
 
 use sc_executor_common::{
 	error::{Error, Result, WasmError},
-	wasm_runtime::WasmRuntime,
+	wasm_runtime::{WasmModule, WasmInstance},
 };
 use sp_allocator::FreeingBumpHeapAllocator;
 use sp_runtime_interface::unpack_ptr_and_len;
 use sp_wasm_interface::{Function, Pointer, WordSize, Value};
 use wasmtime::{Config, Engine, Module, Store};
 
-/// A `WasmRuntime` implementation using wasmtime to compile the runtime module to machine code
+/// A `WasmModule` implementation using wasmtime to compile the runtime module to machine code
 /// and execute the compiled code.
 pub struct WasmtimeRuntime {
-	module: Module,
-	imports: Imports,
-	state_holder: StateHolder,
+	module: Arc<Module>,
 	heap_pages: u32,
+	allow_missing_func_imports: bool,
 	host_functions: Vec<&'static dyn Function>,
 }
 
-impl WasmRuntime for WasmtimeRuntime {
-	fn host_functions(&self) -> &[&'static dyn Function] {
-		&self.host_functions
-	}
-
-	fn call(&mut self, method: &str, data: &[u8]) -> Result<Vec<u8>> {
-		call_method(
+impl WasmModule for WasmtimeRuntime {
+	fn new_instance(&self) -> Result<Box<dyn WasmInstance>> {
+		// Scan all imports, find the matching host functions, and create stubs that adapt arguments
+		// and results.
+		let imports = resolve_imports(
 			&self.module,
-			&mut self.imports,
-			&self.state_holder,
-			method,
-			data,
+			&self.host_functions,
 			self.heap_pages,
-		)
-	}
+			self.allow_missing_func_imports,
+		)?;
 
-	fn get_global_val(&self, name: &str) -> Result<Option<Value>> {
-		// Yeah, there is no better way currently :(
-		InstanceWrapper::new(&self.module, &self.imports, self.heap_pages)?
-			.get_global_val(name)
+		Ok(Box::new(WasmtimeInstance {
+			module: self.module.clone(),
+			imports,
+			heap_pages: self.heap_pages,
+		}))
 	}
 }
 
+/// A `WasmInstance` implementation that reuses compiled module and spawns instances
+/// to execute the compiled code.
+pub struct WasmtimeInstance {
+	module: Arc<Module>,
+	imports: Imports,
+	heap_pages: u32,
+}
+
+// This is safe because `WasmtimeInstance` does not leak reference to `self.imports`
+// and all imports don't reference any anything, other than host functions and memory
+unsafe impl Send for WasmtimeInstance {}
+
+impl WasmInstance for WasmtimeInstance {
+	fn call(&self, method: &str, data: &[u8]) -> Result<Vec<u8>> {
+		// TODO: reuse the instance and reset globals after call
+		// https://github.com/paritytech/substrate/issues/5141
+		let instance = Rc::new(InstanceWrapper::new(&self.module, &self.imports, self.heap_pages)?);
+		call_method(
+			instance,
+			method,
+			data,
+		)
+	}
+
+	fn get_global_const(&self, name: &str) -> Result<Option<Value>> {
+		let instance = InstanceWrapper::new(&self.module, &self.imports, self.heap_pages)?;
+		instance.get_global_val(name)
+	}
+}
+
+
 /// Create a new `WasmtimeRuntime` given the code. This function performs translation from Wasm to
 /// machine code, which can be computationally heavy.
-pub fn create_instance(
+pub fn create_runtime(
 	code: &[u8],
 	heap_pages: u64,
 	host_functions: Vec<&'static dyn Function>,
@@ -80,55 +108,37 @@ pub fn create_instance(
 	let module = Module::new(&store, code)
 		.map_err(|e| WasmError::Other(format!("cannot create module: {}", e)))?;
 
-	let state_holder = StateHolder::empty();
-
-	// Scan all imports, find the matching host functions, and create stubs that adapt arguments
-	// and results.
-	let imports = resolve_imports(
-		&state_holder,
-		&module,
-		&host_functions,
-		heap_pages as u32,
-		allow_missing_func_imports,
-	)?;
-
 	Ok(WasmtimeRuntime {
-		module,
-		imports,
-		state_holder,
+		module: Arc::new(module),
 		heap_pages: heap_pages as u32,
+		allow_missing_func_imports,
 		host_functions,
 	})
 }
 
 /// Call a function inside a precompiled Wasm module.
 fn call_method(
-	module: &Module,
-	imports: &mut Imports,
-	state_holder: &StateHolder,
+	instance_wrapper: Rc<InstanceWrapper>,
 	method: &str,
 	data: &[u8],
-	heap_pages: u32,
 ) -> Result<Vec<u8>> {
-	let instance_wrapper = InstanceWrapper::new(module, imports, heap_pages)?;
 	let entrypoint = instance_wrapper.resolve_entrypoint(method)?;
 	let heap_base = instance_wrapper.extract_heap_base()?;
 	let allocator = FreeingBumpHeapAllocator::new(heap_base);
 
-	perform_call(data, state_holder, instance_wrapper, entrypoint, allocator)
+	perform_call(data, instance_wrapper, entrypoint, allocator)
 }
 
 fn perform_call(
 	data: &[u8],
-	state_holder: &StateHolder,
-	instance_wrapper: InstanceWrapper,
+	instance_wrapper: Rc<InstanceWrapper>,
 	entrypoint: wasmtime::Func,
 	mut allocator: FreeingBumpHeapAllocator,
 ) -> Result<Vec<u8>> {
 	let (data_ptr, data_len) = inject_input_data(&instance_wrapper, &mut allocator, data)?;
 
-	let host_state = HostState::new(allocator, instance_wrapper);
-	let (ret, host_state) = state_holder.with_initialized_state(host_state, || {
+	let host_state = HostState::new(allocator, instance_wrapper.clone());
+	let ret = state_holder::with_initialized_state(&host_state, || {
 		match entrypoint.call(&[
 			wasmtime::Val::I32(u32::from(data_ptr) as i32),
 			wasmtime::Val::I32(u32::from(data_len) as i32),
@@ -146,9 +156,7 @@ fn perform_call(
 		}
 	});
 	let (output_ptr, output_len) = ret?;
-
-	let instance = host_state.into_instance();
-	let output = extract_output_data(&instance, output_ptr, output_len)?;
+	let output = extract_output_data(&instance_wrapper, output_ptr, output_len)?;
 
 	Ok(output)
 }

substrate/client/executor/wasmtime/src/state_holder.rs

@@ -15,63 +15,29 @@
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
 use crate::host::{HostContext, HostState};
-use std::cell::RefCell;
-use std::rc::Rc;
 
-/// A common place to store a reference to the `HostState`.
+scoped_tls::scoped_thread_local!(static HOST_STATE: HostState);
+
+/// Provide `HostState` for the runtime method call and execute the given function `f`.
 ///
-/// This structure is passed into each host function handler and retained in the implementation of
-/// `WasmRuntime`. Whenever a call into a runtime method is initiated, the host state is populated
-/// with the state for that runtime method call.
-///
-/// During the execution of the runtime method call, wasm can call imported host functions. When
-/// that happens the host function handler gets a `HostContext` (obtainable through having a
-/// `HostState` reference).
-#[derive(Clone)]
-pub struct StateHolder {
-	// This is `Some` only during a call.
-	state: Rc<RefCell<Option<HostState>>>,
+/// During the execution of the provided function `with_context` will be callable.
+pub fn with_initialized_state<R, F>(s: &HostState, f: F) -> R
+where
+	F: FnOnce() -> R,
+{
+	HOST_STATE.set(s, f)
 }
 
-impl StateHolder {
-	/// Create a placeholder `StateHolder`.
-	pub fn empty() -> StateHolder {
-		StateHolder {
-			state: Rc::new(RefCell::new(None)),
-		}
-	}
-
-	/// Provide `HostState` for the runtime method call and execute the given function `f`.
-	///
-	/// During the execution of the provided function `with_context` will be callable.
-	pub fn with_initialized_state<R, F>(&self, state: HostState, f: F) -> (R, HostState)
-	where
-		F: FnOnce() -> R,
-	{
-		*self.state.borrow_mut() = Some(state);
-
-		let ret = f();
-		let state = self
-			.state
-			.borrow_mut()
-			.take()
-			.expect("cannot be None since was just assigned; qed");
-
-		(ret, state)
-	}
-
-	/// Create a `HostContext` from the contained `HostState` and execute the given function `f`.
-	///
-	/// This function is only callable within closure passed to `init_state`. Otherwise, the passed
-	/// context will be `None`.
-	pub fn with_context<R, F>(&self, f: F) -> R
-	where
-		F: FnOnce(Option<HostContext>) -> R,
-	{
-		let state = self.state.borrow();
-		match *state {
-			Some(ref state) => f(Some(state.materialize())),
-			None => f(None),
-		}
+/// Create a `HostContext` from the contained `HostState` and execute the given function `f`.
+///
+/// This function is only callable within closure passed to `init_state`. Otherwise, the passed
+/// context will be `None`.
+pub fn with_context<R, F>(f: F) -> R
+where
+	F: FnOnce(Option<HostContext>) -> R,
+{
+	if !HOST_STATE.is_set() {
+		return f(None)
 	}
+	HOST_STATE.with(|state| f(Some(state.materialize())))
 }