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