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