Decommit instance memory after a runtime call on Linux (#8998)

* Decommit instance memory after a runtime call on Linux

* Update documentation for the test

* Remove unfinished comment

* Use saturating_sub.

Also update the doc comment.

* Precise RSS tracking in the test

Instead of tracking RSS for the whole process we just look at the particular mapping that is associated with the linear memory of the runtime instance

* Remove unused import

* Fix unused imports

* Fix the unused imports error for good

* Rollback an accidental change to benches

* Fix the test

* Remove now unneeded code

client/executor/Cargo.toml

@@ -50,6 +50,7 @@ sc-tracing = { version = "3.0.0", path = "../tracing" }
 tracing = "0.1.25"
 tracing-subscriber = "0.2.18"
 paste = "1.0"
+regex = "1"
 
 [features]
 default = [ "std" ]

client/executor/common/src/wasm_runtime.rs

@@ -93,4 +93,13 @@ pub trait WasmInstance: Send {
 	///
 	/// This method is only suitable for getting immutable globals.
 	fn get_global_const(&self, name: &str) -> Result<Option<Value>, Error>;
+
+	/// **Testing Only**. This function returns the base address of the linear memory.
+	///
+	/// This is meant to be the starting address of the memory mapped area for the linear memory.
+	///
+	/// This function is intended only for a specific test that measures physical memory consumption.
+	fn linear_memory_base_ptr(&self) -> Option<*const u8> {
+		None
+	}
 }

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/linux.rs

@@ -0,0 +1,73 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program 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.
+
+// This program 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 this program. If not, see <https://www.gnu.org/licenses/>.
+
+//! Tests that are only relevant for Linux.
+
+// Constrain this only to wasmtime for the time being. Without this rustc will complain on unused
+// imports and items. The alternative is to plop `cfg(feature = wasmtime)` everywhere which seems
+// borthersome.
+#![cfg(feature = "wasmtime")]
+
+use crate::WasmExecutionMethod;
+use super::mk_test_runtime;
+use codec::Encode as _;
+
+mod smaps;
+
+use self::smaps::Smaps;
+
+#[test]
+fn memory_consumption_compiled() {
+	// This aims to see if linear memory stays backed by the physical memory after a runtime call.
+	//
+	// For that we make a series of runtime calls, probing the RSS for the VMA matching the linear
+	// memory. After the call we expect RSS to be equal to 0.
+
+	let runtime = mk_test_runtime(WasmExecutionMethod::Compiled, 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`");
+
+	fn probe_rss(instance: &dyn sc_executor_common::wasm_runtime::WasmInstance) -> usize {
+		let base_addr = instance.linear_memory_base_ptr().unwrap() as usize;
+		Smaps::new().get_rss(base_addr).expect("failed to get rss")
+	}
+
+	instance
+		.call_export(
+			"test_dirty_plenty_memory",
+			&(heap_base as u32, 1u32).encode(),
+		)
+		.unwrap();
+	let probe_1 = probe_rss(&*instance);
+	instance
+		.call_export(
+			"test_dirty_plenty_memory",
+			&(heap_base as u32, 1024u32).encode(),
+		)
+		.unwrap();
+	let probe_2 = probe_rss(&*instance);
+
+	assert_eq!(probe_1, 0);
+	assert_eq!(probe_2, 0);
+}

client/executor/src/integration_tests/linux/smaps.rs

@@ -0,0 +1,82 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2017-2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program 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.
+
+// This program 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 this program. If not, see <https://www.gnu.org/licenses/>.
+
+//! A tool for extracting information about the memory consumption of the current process from
+//! the procfs.
+
+use std::ops::Range;
+use std::collections::BTreeMap;
+
+/// An interface to the /proc/self/smaps
+///
+/// See docs about [procfs on kernel.org][procfs]
+///
+/// [procfs]: https://www.kernel.org/doc/html/latest/filesystems/proc.html
+pub struct Smaps(Vec<(Range<usize>, BTreeMap<String, usize>)>);
+
+impl Smaps {
+	pub fn new() -> Self {
+		let regex_start = regex::RegexBuilder::new("^([0-9a-f]+)-([0-9a-f]+)")
+			.multi_line(true)
+			.build()
+			.unwrap();
+		let regex_kv = regex::RegexBuilder::new(r#"^([^:]+):\s*(\d+) kB"#)
+			.multi_line(true)
+			.build()
+			.unwrap();
+		let smaps = std::fs::read_to_string("/proc/self/smaps").unwrap();
+		let boundaries: Vec<_> = regex_start
+			.find_iter(&smaps)
+			.map(|matched| matched.start())
+			.chain(std::iter::once(smaps.len()))
+			.collect();
+
+		let mut output = Vec::new();
+		for window in boundaries.windows(2) {
+			let chunk = &smaps[window[0]..window[1]];
+			let caps = regex_start.captures(chunk).unwrap();
+			let start = usize::from_str_radix(caps.get(1).unwrap().as_str(), 16).unwrap();
+			let end = usize::from_str_radix(caps.get(2).unwrap().as_str(), 16).unwrap();
+
+			let values = regex_kv
+				.captures_iter(chunk)
+				.map(|cap| {
+					let key = cap.get(1).unwrap().as_str().to_owned();
+					let value = cap.get(2).unwrap().as_str().parse().unwrap();
+					(key, value)
+				})
+				.collect();
+
+			output.push((start..end, values));
+		}
+
+		Self(output)
+	}
+
+	fn get_map(&self, addr: usize) -> &BTreeMap<String, usize> {
+		&self.0
+			.iter()
+			.find(|(range, _)| addr >= range.start && addr < range.end)
+			.unwrap()
+			.1
+	}
+
+	pub fn get_rss(&self, addr: usize) -> Option<usize> {
+		self.get_map(addr).get("Rss").cloned()
+	}
+}

client/executor/src/integration_tests/mod.rs

@@ -15,6 +15,9 @@
 
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+#[cfg(target_os = "linux")]
+mod linux;
 mod sandbox;
 
 use std::sync::Arc;

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,43 @@ impl InstanceWrapper {
 			slice::from_raw_parts_mut(ptr, len)
 		}
 	}
+
+	/// Returns the pointer to the first byte of the linear memory for this instance.
+	pub fn base_ptr(&self) -> *const u8 {
+		self.memory.data_ptr()
+	}
+
+	/// 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()?;
@@ -173,6 +179,19 @@ impl WasmInstance for WasmtimeInstance {
 			}
 		}
 	}
+
+	fn linear_memory_base_ptr(&self) -> Option<*const u8> {
+		match &self.strategy {
+			Strategy::RecreateInstance(_) => {
+				// We do not keep the wasm instance around, therefore there is no linear memory
+				// associated with it.
+				None
+			}
+			Strategy::FastInstanceReuse {
+				instance_wrapper, ..
+			} => Some(instance_wrapper.base_ptr()),
+		}
+	}
 }
 
 /// Prepare a directory structure and a config file to enable wasmtime caching.