Add extra WASM heap pages when precompiling the runtime blob

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,21 @@ 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 [`Config::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.
+	/// [`Config::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 +535,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 +574,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 +587,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 +609,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,8 +42,9 @@ 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,
 		}
 	}
 
@@ -58,6 +60,10 @@ impl RuntimeBuilder {
 		self.deterministic_stack = deterministic_stack;
 	}
 
+	fn precompile_runtime(&mut self, precompile_runtime: bool) {
+		self.precompile_runtime = precompile_runtime;
+	}
+
 	fn max_memory_size(&mut self, max_memory_size: Option<usize>) {
 		self.max_memory_size = max_memory_size;
 	}
@@ -78,27 +84,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>
@@ -168,24 +178,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_imported_memory_with_precompilation() {
+	test_max_memory_pages(true, true);
 }
 
 #[test]
-fn test_max_memory_pages_exported_memory() {
-	test_max_memory_pages(false);
+fn test_max_memory_pages_exported_memory_with_precompilation() {
+	test_max_memory_pages(false, true);
 }
 
-fn test_max_memory_pages(import_memory: bool) {
+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.precompile_runtime(precompile_runtime);
 			builder.build()
 		};
 		let mut instance = runtime.new_instance()?;
@@ -235,6 +257,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			*/
 			memory(64511, None, import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 
@@ -257,6 +280,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 +301,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			// Max is 2048.
 			memory(1, Some(2048), import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 
@@ -309,6 +334,7 @@ fn test_max_memory_pages(import_memory: bool) {
 			// Zero starting pages.
 			memory(0, None, import_memory)
 		),
+		precompile_runtime,
 	)
 	.unwrap();
 
@@ -341,6 +367,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 +380,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 +388,7 @@ fn test_instances_without_reuse_are_not_leaked() {
 				deterministic_stack_limit: None,
 				canonicalize_nans: false,
 				parallel_compilation: true,
+				extra_heap_pages: 2048,
 			},
 		},
 	)