Add extra WASM heap pages when precompiling the runtime blob (parityt…

…ech#11107)

* Add extra WASM heap pages when precompiling the runtime blob

* Fix compilation

* Fix rustdoc

* Fix rustdoc for real this time

* Fix benches compilation

* Improve the builder in `sc-executor-wasmtime`'s tests

client/executor/benches/bench.rs

@@ -55,11 +55,11 @@ fn initialize(runtime: &[u8], method: Method) -> Arc<dyn WasmModule> {
 			sc_executor_wasmtime::create_runtime::<sp_io::SubstrateHostFunctions>(
 				blob,
 				sc_executor_wasmtime::Config {
-					heap_pages,
 					max_memory_size: None,
 					allow_missing_func_imports,
 					cache_path: None,
 					semantics: sc_executor_wasmtime::Semantics {
+						extra_heap_pages: heap_pages,
 						fast_instance_reuse,
 						deterministic_stack_limit: None,
 						canonicalize_nans: false,

client/executor/src/wasm_runtime.rs

@@ -317,11 +317,11 @@ where
 		WasmExecutionMethod::Compiled => sc_executor_wasmtime::create_runtime::<H>(
 			blob,
 			sc_executor_wasmtime::Config {
-				heap_pages,
 				max_memory_size: None,
 				allow_missing_func_imports,
 				cache_path: cache_path.map(ToOwned::to_owned),
 				semantics: sc_executor_wasmtime::Semantics {
+					extra_heap_pages: heap_pages,
 					fast_instance_reuse: true,
 					deterministic_stack_limit: None,
 					canonicalize_nans: false,

client/executor/wasmtime/src/runtime.rs

@@ -414,20 +414,22 @@ pub struct Semantics {
 
 	/// Configures wasmtime to use multiple threads for compiling.
 	pub parallel_compilation: bool,
+
+	/// The number of extra WASM pages which will be allocated
+	/// on top of what is requested by the WASM blob itself.
+	pub extra_heap_pages: u64,
 }
 
 pub struct Config {
-	/// The number of wasm pages to be mounted after instantiation.
-	pub heap_pages: u64,
-
 	/// The total amount of memory in bytes an instance can request.
 	///
 	/// If specified, the runtime will be able to allocate only that much of wasm memory.
-	/// This is the total number and therefore the [`Config::heap_pages`] is accounted for.
+	/// This is the total number and therefore the [`Semantics::extra_heap_pages`] is accounted
+	/// for.
 	///
 	/// That means that the initial number of pages of a linear memory plus the
-	/// [`Config::heap_pages`] multiplied by the wasm page size (64KiB) should be less than or
-	/// equal to `max_memory_size`, otherwise the instance won't be created.
+	/// [`Semantics::extra_heap_pages`] multiplied by the wasm page size (64KiB) should be less
+	/// than or equal to `max_memory_size`, otherwise the instance won't be created.
 	///
 	/// Moreover, `memory.grow` will fail (return -1) if the sum of sizes of currently mounted
 	/// and additional pages exceeds `max_memory_size`.
@@ -534,21 +536,7 @@ where
 
 	let (module, snapshot_data) = match code_supply_mode {
 		CodeSupplyMode::Verbatim { blob } => {
-			let mut blob = instrument(blob, &config.semantics)?;
-
-			// We don't actually need the memory to be imported so we can just convert any memory
-			// import into an export with impunity. This simplifies our code since `wasmtime` will
-			// now automatically take care of creating the memory for us, and it also allows us
-			// to potentially enable `wasmtime`'s instance pooling at a later date. (Imported
-			// memories are ineligible for pooling.)
-			blob.convert_memory_import_into_export()?;
-			blob.add_extra_heap_pages_to_memory_section(
-				config
-					.heap_pages
-					.try_into()
-					.map_err(|e| WasmError::Other(format!("invalid `heap_pages`: {}", e)))?,
-			)?;
-
+			let blob = prepare_blob_for_compilation(blob, &config.semantics)?;
 			let serialized_blob = blob.clone().serialize();
 
 			let module = wasmtime::Module::new(&engine, &serialized_blob)
@@ -587,7 +575,7 @@ where
 	Ok(WasmtimeRuntime { engine, instance_pre: Arc::new(instance_pre), snapshot_data, config })
 }
 
-fn instrument(
+fn prepare_blob_for_compilation(
 	mut blob: RuntimeBlob,
 	semantics: &Semantics,
 ) -> std::result::Result<RuntimeBlob, WasmError> {
@@ -600,6 +588,19 @@ fn instrument(
 		blob.expose_mutable_globals();
 	}
 
+	// We don't actually need the memory to be imported so we can just convert any memory
+	// import into an export with impunity. This simplifies our code since `wasmtime` will
+	// now automatically take care of creating the memory for us, and it also allows us
+	// to potentially enable `wasmtime`'s instance pooling at a later date. (Imported
+	// memories are ineligible for pooling.)
+	blob.convert_memory_import_into_export()?;
+	blob.add_extra_heap_pages_to_memory_section(
+		semantics
+			.extra_heap_pages
+			.try_into()
+			.map_err(|e| WasmError::Other(format!("invalid `extra_heap_pages`: {}", e)))?,
+	)?;
+
 	Ok(blob)
 }
 
@@ -609,7 +610,7 @@ pub fn prepare_runtime_artifact(
 	blob: RuntimeBlob,
 	semantics: &Semantics,
 ) -> std::result::Result<Vec<u8>, WasmError> {
-	let blob = instrument(blob, semantics)?;
+	let blob = prepare_blob_for_compilation(blob, semantics)?;
 
 	let engine = Engine::new(&common_config(semantics)?)
 		.map_err(|e| WasmError::Other(format!("cannot create the engine: {}", e)))?;

client/executor/wasmtime/src/tests.rs

@@ -28,8 +28,9 @@ struct RuntimeBuilder {
 	fast_instance_reuse: bool,
 	canonicalize_nans: bool,
 	deterministic_stack: bool,
-	heap_pages: u64,
+	extra_heap_pages: u64,
 	max_memory_size: Option<usize>,
+	precompile_runtime: bool,
 }
 
 impl RuntimeBuilder {
@@ -41,34 +42,44 @@ impl RuntimeBuilder {
 			fast_instance_reuse: false,
 			canonicalize_nans: false,
 			deterministic_stack: false,
-			heap_pages: 1024,
+			extra_heap_pages: 1024,
 			max_memory_size: None,
+			precompile_runtime: false,
 		}
 	}
 
-	fn use_wat(&mut self, code: String) {
+	fn use_wat(&mut self, code: String) -> &mut Self {
 		self.code = Some(code);
+		self
 	}
 
-	fn canonicalize_nans(&mut self, canonicalize_nans: bool) {
+	fn canonicalize_nans(&mut self, canonicalize_nans: bool) -> &mut Self {
 		self.canonicalize_nans = canonicalize_nans;
+		self
 	}
 
-	fn deterministic_stack(&mut self, deterministic_stack: bool) {
+	fn deterministic_stack(&mut self, deterministic_stack: bool) -> &mut Self {
 		self.deterministic_stack = deterministic_stack;
+		self
 	}
 
-	fn max_memory_size(&mut self, max_memory_size: Option<usize>) {
+	fn precompile_runtime(&mut self, precompile_runtime: bool) -> &mut Self {
+		self.precompile_runtime = precompile_runtime;
+		self
+	}
+
+	fn max_memory_size(&mut self, max_memory_size: Option<usize>) -> &mut Self {
 		self.max_memory_size = max_memory_size;
+		self
 	}
 
-	fn build(self) -> Arc<dyn WasmModule> {
+	fn build(&mut self) -> Arc<dyn WasmModule> {
 		let blob = {
 			let wasm: Vec<u8>;
 
 			let wasm = match self.code {
 				None => wasm_binary_unwrap(),
-				Some(wat) => {
+				Some(ref wat) => {
 					wasm = wat::parse_str(wat).expect("wat parsing failed");
 					&wasm
 				},
@@ -78,27 +89,31 @@ impl RuntimeBuilder {
 				.expect("failed to create a runtime blob out of test runtime")
 		};
 
-		let rt = crate::create_runtime::<HostFunctions>(
-			blob,
-			crate::Config {
-				heap_pages: self.heap_pages,
-				max_memory_size: self.max_memory_size,
-				allow_missing_func_imports: true,
-				cache_path: None,
-				semantics: crate::Semantics {
-					fast_instance_reuse: self.fast_instance_reuse,
-					deterministic_stack_limit: match self.deterministic_stack {
-						true => Some(crate::DeterministicStackLimit {
-							logical_max: 65536,
-							native_stack_max: 256 * 1024 * 1024,
-						}),
-						false => None,
-					},
-					canonicalize_nans: self.canonicalize_nans,
-					parallel_compilation: true,
+		let config = crate::Config {
+			max_memory_size: self.max_memory_size,
+			allow_missing_func_imports: true,
+			cache_path: None,
+			semantics: crate::Semantics {
+				fast_instance_reuse: self.fast_instance_reuse,
+				deterministic_stack_limit: match self.deterministic_stack {
+					true => Some(crate::DeterministicStackLimit {
+						logical_max: 65536,
+						native_stack_max: 256 * 1024 * 1024,
+					}),
+					false => None,
 				},
+				canonicalize_nans: self.canonicalize_nans,
+				parallel_compilation: true,
+				extra_heap_pages: self.extra_heap_pages,
 			},
-		)
+		};
+
+		let rt = if self.precompile_runtime {
+			let artifact = crate::prepare_runtime_artifact(blob, &config.semantics).unwrap();
+			unsafe { crate::create_runtime_from_artifact::<HostFunctions>(&artifact, config) }
+		} else {
+			crate::create_runtime::<HostFunctions>(blob, config)
+		}
 		.expect("cannot create runtime");
 
 		Arc::new(rt) as Arc<dyn WasmModule>
@@ -107,11 +122,7 @@ impl RuntimeBuilder {
 
 #[test]
 fn test_nan_canonicalization() {
-	let runtime = {
-		let mut builder = RuntimeBuilder::new_on_demand();
-		builder.canonicalize_nans(true);
-		builder.build()
-	};
+	let runtime = RuntimeBuilder::new_on_demand().canonicalize_nans(true).build();
 
 	let mut instance = runtime.new_instance().expect("failed to instantiate a runtime");
 
@@ -150,12 +161,10 @@ fn test_nan_canonicalization() {
 fn test_stack_depth_reaching() {
 	const TEST_GUARD_PAGE_SKIP: &str = include_str!("test-guard-page-skip.wat");
 
-	let runtime = {
-		let mut builder = RuntimeBuilder::new_on_demand();
-		builder.use_wat(TEST_GUARD_PAGE_SKIP.to_string());
-		builder.deterministic_stack(true);
-		builder.build()
-	};
+	let runtime = RuntimeBuilder::new_on_demand()
+		.use_wat(TEST_GUARD_PAGE_SKIP.to_string())
+		.deterministic_stack(true)
+		.build();
 	let mut instance = runtime.new_instance().expect("failed to instantiate a runtime");
 
 	match instance.call_export("test-many-locals", &[]).unwrap_err() {
@@ -168,26 +177,36 @@ fn test_stack_depth_reaching() {
 }
 
 #[test]
-fn test_max_memory_pages_imported_memory() {
-	test_max_memory_pages(true);
+fn test_max_memory_pages_imported_memory_without_precompilation() {
+	test_max_memory_pages(true, false);
+}
+
+#[test]
+fn test_max_memory_pages_exported_memory_without_precompilation() {
+	test_max_memory_pages(false, false);
 }
 
 #[test]
-fn test_max_memory_pages_exported_memory() {
-	test_max_memory_pages(false);
+fn test_max_memory_pages_imported_memory_with_precompilation() {
+	test_max_memory_pages(true, true);
 }
 
-fn test_max_memory_pages(import_memory: bool) {
+#[test]
+fn test_max_memory_pages_exported_memory_with_precompilation() {
+	test_max_memory_pages(false, true);
+}
+
+fn test_max_memory_pages(import_memory: bool, precompile_runtime: bool) {
 	fn try_instantiate(
 		max_memory_size: Option<usize>,
 		wat: String,
+		precompile_runtime: bool,
 	) -> Result<(), Box<dyn std::error::Error>> {
-		let runtime = {
-			let mut builder = RuntimeBuilder::new_on_demand();
-			builder.use_wat(wat);
-			builder.max_memory_size(max_memory_size);
-			builder.build()
-		};
+		let runtime = RuntimeBuilder::new_on_demand()
+			.use_wat(wat)
+			.max_memory_size(max_memory_size)
+			.precompile_runtime(precompile_runtime)
+			.build();
 		let mut instance = runtime.new_instance()?;
 		let _ = instance.call_export("main", &[])?;
 		Ok(())
@@ -235,6 +254,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			*/
 			memory(64511, None, import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 
@@ -257,6 +277,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			// 1 initial, max is not specified.
 			memory(1, None, import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 
@@ -277,6 +298,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			// Max is 2048.
 			memory(1, Some(2048), import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 
@@ -309,6 +331,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			// Zero starting pages.
 			memory(0, None, import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 
@@ -341,6 +364,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			// Initial=1, meaning after heap pages mount the total will be already 1025.
 			memory(1, None, import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 }
@@ -353,7 +377,6 @@ fn test_instances_without_reuse_are_not_leaked() {
 	let runtime = crate::create_runtime::<HostFunctions>(
 		RuntimeBlob::uncompress_if_needed(wasm_binary_unwrap()).unwrap(),
 		crate::Config {
-			heap_pages: 2048,
 			max_memory_size: None,
 			allow_missing_func_imports: true,
 			cache_path: None,
@@ -362,6 +385,7 @@ fn test_instances_without_reuse_are_not_leaked() {
 				deterministic_stack_limit: None,
 				canonicalize_nans: false,
 				parallel_compilation: true,
+				extra_heap_pages: 2048,
 			},
 		},
 	)