Update wasmtime to 0.29.0 (#9552)

* Start

* Move to ctx

* Make it compile for now

* More work

* Get rid off state-holder

* Use less Refcells

* 🤦

* Don't use RefCell

* Use names for parameters

* Fixes after merge

* Fixes after merge

* Review feedback

* FMT

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

@@ -22,13 +22,15 @@ use crate::{
 	host::HostState,
 	imports::{resolve_imports, Imports},
 	instance_wrapper::{EntryPoint, InstanceWrapper},
-	state_holder,
+	util,
 };
 
 use sc_allocator::FreeingBumpHeapAllocator;
 use sc_executor_common::{
 	error::{Result, WasmError},
-	runtime_blob::{DataSegmentsSnapshot, ExposedMutableGlobalsSet, GlobalsSnapshot, RuntimeBlob},
+	runtime_blob::{
+		self, DataSegmentsSnapshot, ExposedMutableGlobalsSet, GlobalsSnapshot, RuntimeBlob,
+	},
 	wasm_runtime::{InvokeMethod, WasmInstance, WasmModule},
 };
 use sp_runtime_interface::unpack_ptr_and_len;
@@ -38,7 +40,24 @@ use std::{
 	rc::Rc,
 	sync::Arc,
 };
-use wasmtime::{Engine, Store};
+use wasmtime::{AsContext, AsContextMut, Engine, StoreLimits};
+
+pub(crate) struct StoreData {
+	/// The limits we aply to the store. We need to store it here to return a reference to this
+	/// object when we have the limits enabled.
+	limits: StoreLimits,
+	/// This will only be set when we call into the runtime.
+	host_state: Option<Rc<HostState>>,
+}
+
+impl StoreData {
+	/// Returns a reference to the host state.
+	pub fn host_state(&self) -> Option<&Rc<HostState>> {
+		self.host_state.as_ref()
+	}
+}
+
+pub(crate) type Store = wasmtime::Store<StoreData>;
 
 enum Strategy {
 	FastInstanceReuse {
@@ -46,6 +65,7 @@ enum Strategy {
 		globals_snapshot: GlobalsSnapshot<wasmtime::Global>,
 		data_segments_snapshot: Arc<DataSegmentsSnapshot>,
 		heap_base: u32,
+		store: Store,
 	},
 	RecreateInstance(InstanceCreator),
 }
@@ -58,8 +78,33 @@ struct InstanceCreator {
 }
 
 impl InstanceCreator {
-	fn instantiate(&self) -> Result<InstanceWrapper> {
-		InstanceWrapper::new(&self.store, &*self.module, &*self.imports, self.heap_pages)
+	fn instantiate(&mut self) -> Result<InstanceWrapper> {
+		InstanceWrapper::new(&*self.module, &*self.imports, self.heap_pages, &mut self.store)
+	}
+}
+
+struct InstanceGlobals<'a, C> {
+	ctx: &'a mut C,
+	instance: &'a InstanceWrapper,
+}
+
+impl<'a, C: AsContextMut> runtime_blob::InstanceGlobals for InstanceGlobals<'a, C> {
+	type Global = wasmtime::Global;
+
+	fn get_global(&mut self, export_name: &str) -> Self::Global {
+		self.instance
+			.get_global(&mut self.ctx, export_name)
+			.expect("get_global is guaranteed to be called with an export name of a global; qed")
+	}
+
+	fn get_global_value(&mut self, global: &Self::Global) -> Value {
+		util::from_wasmtime_val(global.get(&mut self.ctx))
+	}
+
+	fn set_global_value(&mut self, global: &Self::Global, value: Value) {
+		global.set(&mut self.ctx, util::into_wasmtime_val(value)).expect(
+			"the value is guaranteed to be of the same value; the global is guaranteed to be mutable; qed",
+		);
 	}
 }
 
@@ -82,19 +127,25 @@ pub struct WasmtimeRuntime {
 impl WasmtimeRuntime {
 	/// Creates the store respecting the set limits.
 	fn new_store(&self) -> Store {
-		match self.config.max_memory_pages {
-			Some(max_memory_pages) => Store::new_with_limits(
-				&self.engine,
-				wasmtime::StoreLimitsBuilder::new().memory_pages(max_memory_pages).build(),
-			),
-			None => Store::new(&self.engine),
+		let limits = if let Some(max_memory_pages) = self.config.max_memory_pages {
+			wasmtime::StoreLimitsBuilder::new().memory_pages(max_memory_pages).build()
+		} else {
+			Default::default()
+		};
+
+		let mut store = Store::new(&self.engine, StoreData { limits, host_state: None });
+
+		if self.config.max_memory_pages.is_some() {
+			store.limiter(|s| &mut s.limits);
 		}
+
+		store
 	}
 }
 
 impl WasmModule for WasmtimeRuntime {
 	fn new_instance(&self) -> Result<Box<dyn WasmInstance>> {
-		let store = self.new_store();
+		let mut store = self.new_store();
 
 		// Scan all imports, find the matching host functions, and create stubs that adapt arguments
 		// and results.
@@ -103,7 +154,7 @@ impl WasmModule for WasmtimeRuntime {
 		//       However, I am not sure if that's a good idea since it would be pushing our luck
 		// further       by assuming that `Store` not only `Send` but also `Sync`.
 		let imports = resolve_imports(
-			&store,
+			&mut store,
 			&self.module,
 			&self.host_functions,
 			self.config.heap_pages,
@@ -112,21 +163,24 @@ impl WasmModule for WasmtimeRuntime {
 
 		let strategy = if let Some(ref snapshot_data) = self.snapshot_data {
 			let instance_wrapper =
-				InstanceWrapper::new(&store, &self.module, &imports, self.config.heap_pages)?;
-			let heap_base = instance_wrapper.extract_heap_base()?;
+				InstanceWrapper::new(&self.module, &imports, self.config.heap_pages, &mut store)?;
+			let heap_base = instance_wrapper.extract_heap_base(&mut store)?;
 
 			// This function panics if the instance was created from a runtime blob different from
 			// which the mutable globals were collected. Here, it is easy to see that there is only
 			// a single runtime blob and thus it's the same that was used for both creating the
 			// instance and collecting the mutable globals.
-			let globals_snapshot =
-				GlobalsSnapshot::take(&snapshot_data.mutable_globals, &instance_wrapper);
+			let globals_snapshot = GlobalsSnapshot::take(
+				&snapshot_data.mutable_globals,
+				&mut InstanceGlobals { ctx: &mut store, instance: &instance_wrapper },
+			);
 
 			Strategy::FastInstanceReuse {
 				instance_wrapper: Rc::new(instance_wrapper),
 				globals_snapshot,
 				data_segments_snapshot: snapshot_data.data_segments_snapshot.clone(),
 				heap_base,
+				store,
 			}
 		} else {
 			Strategy::RecreateInstance(InstanceCreator {
@@ -152,48 +206,63 @@ pub struct WasmtimeInstance {
 unsafe impl Send for WasmtimeInstance {}
 
 impl WasmInstance for WasmtimeInstance {
-	fn call(&self, method: InvokeMethod, data: &[u8]) -> Result<Vec<u8>> {
-		match &self.strategy {
+	fn call(&mut self, method: InvokeMethod, data: &[u8]) -> Result<Vec<u8>> {
+		match &mut self.strategy {
 			Strategy::FastInstanceReuse {
 				instance_wrapper,
 				globals_snapshot,
 				data_segments_snapshot,
 				heap_base,
+				ref mut store,
 			} => {
-				let entrypoint = instance_wrapper.resolve_entrypoint(method)?;
+				let entrypoint = instance_wrapper.resolve_entrypoint(method, &mut *store)?;
 
 				data_segments_snapshot.apply(|offset, contents| {
-					instance_wrapper.write_memory_from(Pointer::new(offset), contents)
+					instance_wrapper.write_memory_from(&mut *store, Pointer::new(offset), contents)
 				})?;
-				globals_snapshot.apply(&**instance_wrapper);
+				globals_snapshot
+					.apply(&mut InstanceGlobals { ctx: &mut *store, instance: &*instance_wrapper });
 				let allocator = FreeingBumpHeapAllocator::new(*heap_base);
 
-				let result =
-					perform_call(data, Rc::clone(&instance_wrapper), entrypoint, allocator);
+				let result = perform_call(
+					&mut *store,
+					data,
+					instance_wrapper.clone(),
+					entrypoint,
+					allocator,
+				);
 
 				// Signal to the OS that we are done with the linear memory and that it can be
 				// reclaimed.
-				instance_wrapper.decommit();
+				instance_wrapper.decommit(&store);
 
 				result
 			},
-			Strategy::RecreateInstance(instance_creator) => {
+			Strategy::RecreateInstance(ref mut instance_creator) => {
 				let instance_wrapper = instance_creator.instantiate()?;
-				let heap_base = instance_wrapper.extract_heap_base()?;
-				let entrypoint = instance_wrapper.resolve_entrypoint(method)?;
+				let heap_base = instance_wrapper.extract_heap_base(&mut instance_creator.store)?;
+				let entrypoint =
+					instance_wrapper.resolve_entrypoint(method, &mut instance_creator.store)?;
 
 				let allocator = FreeingBumpHeapAllocator::new(heap_base);
-				perform_call(data, Rc::new(instance_wrapper), entrypoint, allocator)
+				perform_call(
+					&mut instance_creator.store,
+					data,
+					Rc::new(instance_wrapper),
+					entrypoint,
+					allocator,
+				)
 			},
 		}
 	}
 
-	fn get_global_const(&self, name: &str) -> Result<Option<Value>> {
-		match &self.strategy {
-			Strategy::FastInstanceReuse { instance_wrapper, .. } =>
-				instance_wrapper.get_global_val(name),
-			Strategy::RecreateInstance(instance_creator) =>
-				instance_creator.instantiate()?.get_global_val(name),
+	fn get_global_const(&mut self, name: &str) -> Result<Option<Value>> {
+		match &mut self.strategy {
+			Strategy::FastInstanceReuse { instance_wrapper, ref mut store, .. } =>
+				instance_wrapper.get_global_val(&mut *store, name),
+			Strategy::RecreateInstance(ref mut instance_creator) => instance_creator
+				.instantiate()?
+				.get_global_val(&mut instance_creator.store, name),
 		}
 	}
 
@@ -204,8 +273,8 @@ impl WasmInstance for WasmtimeInstance {
 				// associated with it.
 				None
 			},
-			Strategy::FastInstanceReuse { instance_wrapper, .. } =>
-				Some(instance_wrapper.base_ptr()),
+			Strategy::FastInstanceReuse { instance_wrapper, store, .. } =>
+				Some(instance_wrapper.base_ptr(&store)),
 		}
 	}
 }
@@ -536,40 +605,50 @@ pub fn prepare_runtime_artifact(
 }
 
 fn perform_call(
+	mut ctx: impl AsContextMut<Data = StoreData>,
 	data: &[u8],
 	instance_wrapper: Rc<InstanceWrapper>,
 	entrypoint: EntryPoint,
 	mut allocator: FreeingBumpHeapAllocator,
 ) -> Result<Vec<u8>> {
-	let (data_ptr, data_len) = inject_input_data(&instance_wrapper, &mut allocator, data)?;
+	let (data_ptr, data_len) =
+		inject_input_data(&mut ctx, &instance_wrapper, &mut allocator, data)?;
 
 	let host_state = HostState::new(allocator, instance_wrapper.clone());
-	let ret = state_holder::with_initialized_state(&host_state, || -> Result<_> {
-		Ok(unpack_ptr_and_len(entrypoint.call(data_ptr, data_len)?))
-	});
+
+	// Set the host state before calling into wasm.
+	ctx.as_context_mut().data_mut().host_state = Some(Rc::new(host_state));
+
+	let ret = entrypoint.call(&mut ctx, data_ptr, data_len).map(unpack_ptr_and_len);
+
+	// Reset the host state
+	ctx.as_context_mut().data_mut().host_state = None;
+
 	let (output_ptr, output_len) = ret?;
-	let output = extract_output_data(&instance_wrapper, output_ptr, output_len)?;
+	let output = extract_output_data(ctx, &instance_wrapper, output_ptr, output_len)?;
 
 	Ok(output)
 }
 
 fn inject_input_data(
+	mut ctx: impl AsContextMut,
 	instance: &InstanceWrapper,
 	allocator: &mut FreeingBumpHeapAllocator,
 	data: &[u8],
 ) -> Result<(Pointer<u8>, WordSize)> {
 	let data_len = data.len() as WordSize;
-	let data_ptr = instance.allocate(allocator, data_len)?;
-	instance.write_memory_from(data_ptr, data)?;
+	let data_ptr = instance.allocate(&mut ctx, allocator, data_len)?;
+	instance.write_memory_from(ctx, data_ptr, data)?;
 	Ok((data_ptr, data_len))
 }
 
 fn extract_output_data(
+	ctx: impl AsContext,
 	instance: &InstanceWrapper,
 	output_ptr: u32,
 	output_len: u32,
 ) -> Result<Vec<u8>> {
 	let mut output = vec![0; output_len as usize];
-	instance.read_memory_into(Pointer::new(output_ptr), &mut output)?;
+	instance.read_memory_into(ctx, Pointer::new(output_ptr), &mut output)?;
 	Ok(output)
 }