Reuse wasmtime instances, the PR (#5567)

* Preserve a single wasmtime instance.

* Sketch of wasm instance reusing.

* Refactor and docs.

* Rename state_snapshot to util module.

* Renaming.

* Comments.

* Error handling

* More fixes.

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

@@ -20,11 +20,55 @@
 use crate::util;
 use crate::imports::Imports;
 
-use sc_executor_common::error::{Error, Result};
+use std::{slice, marker};
+use sc_executor_common::{
+	error::{Error, Result},
+	util::{WasmModuleInfo, DataSegmentsSnapshot},
+};
 use sp_wasm_interface::{Pointer, WordSize, Value};
-use std::slice;
-use std::marker;
-use wasmtime::{Instance, Module, Memory, Table, Val};
+use wasmtime::{Store, Instance, Module, Memory, Table, Val};
+
+mod globals_snapshot;
+
+pub use globals_snapshot::GlobalsSnapshot;
+
+pub struct ModuleWrapper {
+	imported_globals_count: u32,
+	globals_count: u32,
+	module: Module,
+	data_segments_snapshot: DataSegmentsSnapshot,
+}
+
+impl ModuleWrapper {
+	pub fn new(store: &Store, code: &[u8]) -> Result<Self> {
+		let module = Module::new(&store, code)
+			.map_err(|e| Error::from(format!("cannot create module: {}", e)))?;
+
+		let module_info = WasmModuleInfo::new(code)
+			.ok_or_else(|| Error::from("cannot deserialize module".to_string()))?;
+		let declared_globals_count = module_info.declared_globals_count();
+		let imported_globals_count = module_info.imported_globals_count();
+		let globals_count = imported_globals_count + declared_globals_count;
+
+		let data_segments_snapshot = DataSegmentsSnapshot::take(&module_info)
+			.map_err(|e| Error::from(format!("cannot take data segments snapshot: {}", e)))?;
+
+		Ok(Self {
+			module,
+			imported_globals_count,
+			globals_count,
+			data_segments_snapshot,
+		})
+	}
+
+	pub fn module(&self) -> &Module {
+		&self.module
+	}
+
+	pub fn data_segments_snapshot(&self) -> &DataSegmentsSnapshot {
+		&self.data_segments_snapshot
+	}
+}
 
 /// Wrap the given WebAssembly Instance of a wasm module with Substrate-runtime.
 ///
@@ -32,6 +76,8 @@ use wasmtime::{Instance, Module, Memory, Table, Val};
 /// routines.
 pub struct InstanceWrapper {
 	instance: Instance,
+	globals_count: u32,
+	imported_globals_count: u32,
 	// The memory instance of the `instance`.
 	//
 	// It is important to make sure that we don't make any copies of this to make it easier to proof
@@ -44,8 +90,8 @@ pub struct InstanceWrapper {
 
 impl InstanceWrapper {
 	/// Create a new instance wrapper from the given wasm module.
-	pub fn new(module: &Module, imports: &Imports, heap_pages: u32) -> Result<Self> {
-		let instance = Instance::new(module, &imports.externs)
+	pub fn new(module_wrapper: &ModuleWrapper, imports: &Imports, heap_pages: u32) -> Result<Self> {
+		let instance = Instance::new(&module_wrapper.module, &imports.externs)
 			.map_err(|e| Error::from(format!("cannot instantiate: {}", e)))?;
 
 		let memory = match imports.memory_import_index {
@@ -66,8 +112,10 @@ impl InstanceWrapper {
 
 		Ok(Self {
 			table: get_table(&instance),
-			memory,
 			instance,
+			globals_count: module_wrapper.globals_count,
+			imported_globals_count: module_wrapper.imported_globals_count,
+			memory,
 			_not_send_nor_sync: marker::PhantomData,
 		})
 	}

substrate/client/executor/wasmtime/src/instance_wrapper/globals_snapshot.rs

@@ -0,0 +1,130 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+use super::InstanceWrapper;
+use sc_executor_common::{
+	error::{Error, Result},
+};
+use sp_wasm_interface::Value;
+use cranelift_codegen::ir;
+use cranelift_wasm::GlobalIndex;
+
+/// A snapshot of a global variables values. This snapshot can be used later for restoring the
+/// values to the preserved state.
+///
+/// Technically, a snapshot stores only values of mutable global variables. This is because
+/// immutable global variables always have the same values.
+pub struct GlobalsSnapshot {
+	handle: wasmtime_runtime::InstanceHandle,
+	preserved_mut_globals: Vec<(*mut wasmtime_runtime::VMGlobalDefinition, Value)>,
+}
+
+impl GlobalsSnapshot {
+	/// Take a snapshot of global variables for a given instance.
+	pub fn take(instance_wrapper: &InstanceWrapper) -> Result<Self> {
+		// EVIL:
+		// Usage of an undocumented function.
+		let handle = instance_wrapper.instance.handle().clone();
+
+		let mut preserved_mut_globals = vec![];
+
+		for global_idx in instance_wrapper.imported_globals_count..instance_wrapper.globals_count {
+			let (def, global) = match handle.lookup_by_declaration(
+				&wasmtime_environ::Export::Global(GlobalIndex::from_u32(global_idx)),
+			) {
+				wasmtime_runtime::Export::Global {
+					definition, global, ..
+				} => (definition, global),
+				_ => unreachable!("only globals can be returned for a global request"),
+			};
+
+			// skip immutable globals.
+			if !global.mutability {
+				continue;
+			}
+
+			let value = unsafe {
+				// Safety of this function solely depends on the correctness of the reference and
+				// the type information of the global.
+				read_global(def, global.ty)?
+			};
+			preserved_mut_globals.push((def, value));
+		}
+
+		Ok(Self {
+			preserved_mut_globals,
+			handle,
+		})
+	}
+
+	/// Apply the snapshot to the given instance.
+	///
+	/// This instance must be the same that was used for creation of this snapshot.
+	pub fn apply(&self, instance_wrapper: &InstanceWrapper) -> Result<()> {
+		if instance_wrapper.instance.handle() != &self.handle {
+			return Err(Error::from("unexpected instance handle".to_string()));
+		}
+
+		for (def, value) in &self.preserved_mut_globals {
+			unsafe {
+				// The following writes are safe if the precondition that this is the same instance
+				// this snapshot was created with:
+				//
+				// 1. These pointers must be still not-NULL and allocated.
+				// 2. The set of global variables is fixed for the lifetime of the same instance.
+				// 3. We obviously assume that the wasmtime references are correct in the first place.
+				// 4. We write the data with the same type it was read in the first place.
+				write_global(*def, *value)?;
+			}
+		}
+		Ok(())
+	}
+}
+
+unsafe fn read_global(
+	def: *const wasmtime_runtime::VMGlobalDefinition,
+	ty: ir::Type,
+) -> Result<Value> {
+	let def = def
+		.as_ref()
+		.ok_or_else(|| Error::from("wasmtime global reference is null during read".to_string()))?;
+	let val = match ty {
+		ir::types::I32 => Value::I32(*def.as_i32()),
+		ir::types::I64 => Value::I64(*def.as_i64()),
+		ir::types::F32 => Value::F32(*def.as_u32()),
+		ir::types::F64 => Value::F64(*def.as_u64()),
+		_ => {
+			return Err(Error::from(format!(
+				"unsupported global variable type: {}",
+				ty
+			)))
+		}
+	};
+	Ok(val)
+}
+
+unsafe fn write_global(def: *mut wasmtime_runtime::VMGlobalDefinition, value: Value) -> Result<()> {
+	let def = def
+		.as_mut()
+		.ok_or_else(|| Error::from("wasmtime global reference is null during write".to_string()))?;
+	match value {
+		Value::I32(v) => *def.as_i32_mut() = v,
+		Value::I64(v) => *def.as_i64_mut() = v,
+		Value::F32(v) => *def.as_u32_mut() = v,
+		Value::F64(v) => *def.as_u64_mut() = v,
+	}
+	Ok(())
+}

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

@@ -15,14 +15,14 @@
 // 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::{Imports, resolve_imports};
-use crate::instance_wrapper::InstanceWrapper;
+use crate::instance_wrapper::{ModuleWrapper, InstanceWrapper, GlobalsSnapshot};
 use crate::state_holder;
 
+use std::rc::Rc;
+use std::sync::Arc;
 use sc_executor_common::{
 	error::{Error, Result, WasmError},
 	wasm_runtime::{WasmModule, WasmInstance},
@@ -30,12 +30,12 @@ use sc_executor_common::{
 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};
+use wasmtime::{Config, Engine, Store};
 
 /// A `WasmModule` implementation using wasmtime to compile the runtime module to machine code
 /// and execute the compiled code.
 pub struct WasmtimeRuntime {
-	module: Arc<Module>,
+	module_wrapper: Arc<ModuleWrapper>,
 	heap_pages: u32,
 	allow_missing_func_imports: bool,
 	host_functions: Vec<&'static dyn Function>,
@@ -46,16 +46,24 @@ impl WasmModule for WasmtimeRuntime {
 		// Scan all imports, find the matching host functions, and create stubs that adapt arguments
 		// and results.
 		let imports = resolve_imports(
-			&self.module,
+			self.module_wrapper.module(),
 			&self.host_functions,
 			self.heap_pages,
 			self.allow_missing_func_imports,
 		)?;
 
+		let instance_wrapper =
+			InstanceWrapper::new(&self.module_wrapper, &imports, self.heap_pages)?;
+		let heap_base = instance_wrapper.extract_heap_base()?;
+		let globals_snapshot = GlobalsSnapshot::take(&instance_wrapper)?;
+
 		Ok(Box::new(WasmtimeInstance {
-			module: self.module.clone(),
+			instance_wrapper: Rc::new(instance_wrapper),
+			module_wrapper: Arc::clone(&self.module_wrapper),
 			imports,
+			globals_snapshot,
 			heap_pages: self.heap_pages,
+			heap_base,
 		}))
 	}
 }
@@ -63,9 +71,12 @@ impl WasmModule for WasmtimeRuntime {
 /// A `WasmInstance` implementation that reuses compiled module and spawns instances
 /// to execute the compiled code.
 pub struct WasmtimeInstance {
-	module: Arc<Module>,
+	module_wrapper: Arc<ModuleWrapper>,
+	instance_wrapper: Rc<InstanceWrapper>,
+	globals_snapshot: GlobalsSnapshot,
 	imports: Imports,
 	heap_pages: u32,
+	heap_base: u32,
 }
 
 // This is safe because `WasmtimeInstance` does not leak reference to `self.imports`
@@ -74,23 +85,32 @@ 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,
+		let entrypoint = self.instance_wrapper.resolve_entrypoint(method)?;
+		let allocator = FreeingBumpHeapAllocator::new(self.heap_base);
+
+		self.module_wrapper
+			.data_segments_snapshot()
+			.apply(|offset, contents| {
+				self.instance_wrapper
+					.write_memory_from(Pointer::new(offset), contents)
+			})?;
+
+		self.globals_snapshot.apply(&*self.instance_wrapper)?;
+
+		perform_call(
 			data,
+			Rc::clone(&self.instance_wrapper),
+			entrypoint,
+			allocator,
 		)
 	}
 
 	fn get_global_const(&self, name: &str) -> Result<Option<Value>> {
-		let instance = InstanceWrapper::new(&self.module, &self.imports, self.heap_pages)?;
+		let instance = InstanceWrapper::new(&self.module_wrapper, &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_runtime(
@@ -105,30 +125,18 @@ pub fn create_runtime(
 
 	let engine = Engine::new(&config);
 	let store = Store::new(&engine);
-	let module = Module::new(&store, code)
+
+	let module_wrapper = ModuleWrapper::new(&store, code)
 		.map_err(|e| WasmError::Other(format!("cannot create module: {}", e)))?;
 
 	Ok(WasmtimeRuntime {
-		module: Arc::new(module),
+		module_wrapper: Arc::new(module_wrapper),
 		heap_pages: heap_pages as u32,
 		allow_missing_func_imports,
 		host_functions,
 	})
 }
 
-/// Call a function inside a precompiled Wasm module.
-fn call_method(
-	instance_wrapper: Rc<InstanceWrapper>,
-	method: &str,
-	data: &[u8],
-) -> Result<Vec<u8>> {
-	let entrypoint = instance_wrapper.resolve_entrypoint(method)?;
-	let heap_base = instance_wrapper.extract_heap_base()?;
-	let allocator = FreeingBumpHeapAllocator::new(heap_base);
-
-	perform_call(data, instance_wrapper, entrypoint, allocator)
-}
-
 fn perform_call(
 	data: &[u8],
 	instance_wrapper: Rc<InstanceWrapper>,