Decommit instance memory after a runtime call on Linux

bin/node/executor/benches/bench.rs

@@ -47,7 +47,7 @@ const TRANSACTION_VERSION: u32 = node_runtime::VERSION.transaction_version;
 
 const SPEC_VERSION: u32 = node_runtime::VERSION.spec_version;
 
-const HEAP_PAGES: u64 = 20;
+const HEAP_PAGES: u64 = 2048;
 
 type TestExternalities<H> = CoreTestExternalities<H, u64>;
 

client/executor/runtime-test/src/lib.rs

@@ -69,6 +69,36 @@ sp_core::wasm_export_functions! {
 
 	fn test_empty_return() {}
 
+	fn test_dirty_plenty_memory(heap_base: u32, heap_pages: u32) {
+		// This piece of code will dirty multiple pages of memory. The number of pages is given by
+		// the `heap_pages`. It's unit is a wasm page (64KiB). The first page to be cleared
+		// is a wasm page that that follows the one that holds the `heap_base` address.
+		//
+		// This function dirties the **host** pages. I.e. we dirty 4KiB at a time and it will take
+		// 16 writes to process a single wasm page.
+
+		let mut heap_ptr = heap_base as usize;
+
+		// Find the next wasm page boundary.
+		let heap_ptr = round_up_to(heap_ptr, 65536);
+
+		// Make it an actual pointer
+		let heap_ptr = heap_ptr as *mut u8;
+
+		// Traverse the host pages and make each one dirty
+		let host_pages = heap_pages as usize * 16;
+		for i in 0..host_pages {
+			unsafe {
+				// technically this is an UB, but there is no way Rust can find this out.
+				heap_ptr.add(i * 4096).write(0);
+			}
+		}
+
+		fn round_up_to(n: usize, divisor: usize) -> usize {
+			(n + divisor - 1) / divisor
+		}
+	}
+
 	fn test_exhaust_heap() -> Vec<u8> { Vec::with_capacity(16777216) }
 
 	fn test_panic() { panic!("test panic") }

client/executor/src/integration_tests/mod.rs

@@ -64,6 +64,15 @@ macro_rules! test_wasm_execution {
 			}
 		}
 	};
+
+	(compiled_only $method_name:ident) => {
+		paste::item! {
+			#[test]
+			fn [<$method_name _compiled>]() {
+				$method_name(WasmExecutionMethod::Compiled);
+			}
+		}
+	};
 }
 
 fn call_in_wasm<E: Externalities>(
@@ -775,3 +784,82 @@ fn panic_in_spawned_instance_panics_on_joining_its_result(wasm_method: WasmExecu
 
 	assert!(format!("{}", error_result).contains("Spawned task"));
 }
+
+
+#[cfg(target_os = "linux")]
+mod linux {
+	use super::*;
+
+	/// Gets the rollup of for the current process from procfs.
+	///
+	/// See the description here: https://www.kernel.org/doc/Documentation/ABI/testing/procfs-smaps_rollup
+	fn obtain_rss() -> u32 {
+		let smaps_rollup = std::fs::read("/proc/self/smaps_rollup").expect("failed to read smaps");
+		let smaps_rollup = String::from_utf8(smaps_rollup).expect("expected proper utf8");
+
+		for line in smaps_rollup.lines() {
+			// We are looking for a line that looks something like below. We are interested only in the
+			// number there.
+			//
+			// Rss:               63624 kB
+			//                    ^^^^^
+			//                    |
+			//                    = our target
+			//
+			// we don't even care about the suffix since it appears to be hardcoded to kB.
+
+			let line = match line.strip_prefix("Rss:") {
+				None => continue,
+				Some(line) => line,
+			};
+
+			let line = match line.strip_suffix("kB") {
+				None => {
+					panic!("weird, we expected that the smaps output is always terminated with `kB`");
+				}
+				Some(line) => line,
+			};
+
+			let line = line.trim();
+
+			let rss = line.parse::<u32>()
+				.expect("failed to parse the RSS");
+
+			return rss
+		}
+
+		panic!("smaps doesn't contain rss!");
+	}
+
+	test_wasm_execution!(compiled_only memory_consumption);
+	fn memory_consumption(wasm_method: WasmExecutionMethod) {
+		let runtime = mk_test_runtime(wasm_method, 1024);
+
+		let instance = runtime.new_instance().unwrap();
+		let heap_base = instance.get_global_const("__heap_base")
+			.expect("`__heap_base` is valid")
+			.expect("`__heap_base` exists")
+			.as_i32()
+			.expect("`__heap_base` is an `i32`");
+
+		instance.call_export("test_dirty_plenty_memory", &(heap_base as u32, 1u32).encode()).unwrap();
+
+		let rss_pre = obtain_rss();
+		instance.call_export("test_dirty_plenty_memory", &(heap_base as u32, 1024u32).encode()).unwrap();
+		let rss_post = obtain_rss();
+
+		// In case the memory is sucessfully decommited we expect only a slight increase of resident memory
+		// usage.
+		//
+		// However, in practice the memory usage can either go up and sometimes down (hence
+		// `saturating_sub` below). This is because this test is run along with other concurrently
+		// and the other tests may allocate or deallocate memory. This in theory makes this test
+		// flaky. This problem should be excaberated with the number of CPU cores
+		//
+		// However, empirically, this works just fine on the "build host" (numcpus=128) and is able
+		// reliably to tell if the decommiting really works or not.
+		//
+		// That said, if the test fails randomly we should consider turning it off.
+		assert!(rss_post.saturating_sub(rss_pre) < 32768 /* kB */);
+	}
+}

client/executor/wasmtime/Cargo.toml

@@ -13,6 +13,8 @@ readme = "README.md"
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
+libc = "0.2.90"
+cfg-if = "1.0"
 log = "0.4.8"
 scoped-tls = "1.0"
 parity-wasm = "0.42.0"

client/executor/wasmtime/src/instance_wrapper.rs

@@ -415,6 +415,38 @@ impl InstanceWrapper {
 			slice::from_raw_parts_mut(ptr, len)
 		}
 	}
+
+	/// Removes physical backing from the allocated linear memory. This leads to returning the memory
+	/// back to the system. While the memory is zeroed this is considered as a side-effect and is not
+	/// relied upon. Thus this function acts as a hint.
+	pub fn decommit(&self) {
+		if self.memory.data_size() == 0 {
+			return;
+		}
+
+		cfg_if::cfg_if! {
+			if #[cfg(target_os = "linux")] {
+				use std::sync::Once;
+
+				unsafe {
+					let ptr = self.memory.data_ptr();
+					let len = self.memory.data_size();
+
+					// Linux handles MADV_DONTNEED reliably. The result is that the given area
+					// is unmapped and will be zeroed on the next pagefault.
+					if libc::madvise(ptr as _, len, libc::MADV_DONTNEED) != 0 {
+						static LOGGED: Once = Once::new();
+						LOGGED.call_once(|| {
+							log::warn!(
+								"madvise(MADV_DONTNEED) failed: {}",
+								std::io::Error::last_os_error(),
+							);
+						});
+					}
+				}
+			}
+		}
+	}
 }
 
 impl runtime_blob::InstanceGlobals for InstanceWrapper {

client/executor/wasmtime/src/runtime.rs

@@ -150,7 +150,13 @@ impl WasmInstance for WasmtimeInstance {
 				globals_snapshot.apply(&**instance_wrapper);
 				let allocator = FreeingBumpHeapAllocator::new(*heap_base);
 
-				perform_call(data, Rc::clone(&instance_wrapper), entrypoint, allocator)
+				let result = perform_call(data, Rc::clone(&instance_wrapper), entrypoint, allocator);
+
+				// Signal to the OS that we are done with the linear memory and that it can be
+				// reclaimed.
+				instance_wrapper.decommit();
+
+				result
 			}
 			Strategy::RecreateInstance(instance_creator) => {
 				let instance_wrapper = instance_creator.instantiate()?;