Tracking/limiting memory allocator

polkadot/Cargo.toml

@@ -34,6 +34,7 @@ polkadot-cli = { path = "cli", features = [
 polkadot-node-core-pvf = { path = "node/core/pvf" }
 polkadot-node-core-pvf-prepare-worker = { path = "node/core/pvf/prepare-worker" }
 polkadot-overseer = { path = "node/overseer" }
+tracking-allocator = { path = "node/tracking-allocator", optional = true }
 
 # Needed for worker binaries.
 polkadot-node-core-pvf-common = { path = "node/core/pvf/common" }
@@ -68,6 +69,12 @@ jemalloc-allocator = [
 	"polkadot-node-core-pvf/jemalloc-allocator",
 	"polkadot-overseer/jemalloc-allocator",
 ]
+tracking-allocator = [
+	"dep:tracking-allocator",
+	"jemalloc-allocator",
+	"polkadot-node-core-pvf-prepare-worker/tracking-allocator",
+	"polkadot-node-core-pvf/tracking-allocator",
+]
 
 # Enables timeout-based tests supposed to be run only in CI environment as they may be flaky
 # when run locally depending on system load

polkadot/node/core/pvf/Cargo.toml

@@ -49,6 +49,10 @@ halt = { package = "test-parachain-halt", path = "../../../parachain/test-parach
 [features]
 ci-only-tests = []
 jemalloc-allocator = [ "polkadot-node-core-pvf-common/jemalloc-allocator" ]
+tracking-allocator = [
+	"polkadot-node-core-pvf-common/tracking-allocator",
+	"polkadot-node-core-pvf-prepare-worker/tracking-allocator",
+]
 # This feature is used to export test code to other crates without putting it in the production build.
 test-utils = [
 	"polkadot-node-core-pvf-execute-worker",

polkadot/node/core/pvf/common/Cargo.toml

@@ -39,3 +39,4 @@ tempfile = "3.3.0"
 # Also used for building the puppet worker.
 test-utils = []
 jemalloc-allocator = []
+tracking-allocator = []

polkadot/node/core/pvf/common/src/prepare.rs

@@ -35,6 +35,9 @@ pub struct MemoryStats {
 	/// `ru_maxrss` from `getrusage`. `None` if an error occurred.
 	#[cfg(target_os = "linux")]
 	pub max_rss: Option<i64>,
+	/// Peak allocation in bytes measured by tracking allocator
+	#[cfg(feature = "tracking-allocator")]
+	pub peak_alloc: u64,
 }
 
 /// Statistics of collected memory metrics.

polkadot/node/core/pvf/prepare-worker/Cargo.toml

@@ -13,6 +13,8 @@ libc = "0.2.139"
 rayon = "1.5.1"
 tikv-jemalloc-ctl = { version = "0.5.0", optional = true }
 tokio = { version = "1.24.2", features = ["fs", "process"] }
+tracking-allocator = { path = "../../../tracking-allocator", optional = true }
+tikv-jemallocator = { version = "0.5.0", optional = true }
 
 parity-scale-codec = { version = "3.6.1", default-features = false, features = ["derive"] }
 
@@ -36,3 +38,16 @@ jemalloc-allocator = [
 	"dep:tikv-jemalloc-ctl",
 	"polkadot-node-core-pvf-common/jemalloc-allocator",
 ]
+tracking-allocator = [
+	"dep:tikv-jemallocator",
+	"dep:tracking-allocator",
+	"polkadot-node-core-pvf-common/tracking-allocator",
+]
+
+[dev-dependencies]
+criterion = { version = "0.4.0", default-features = false, features = ["cargo_bench_support"] }
+staging-kusama-runtime = { path = "../../../../runtime/kusama" }
+
+[[bench]]
+name = "prepare_kusama_runtime"
+harness = false

polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs

@@ -0,0 +1,60 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot 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.
+
+// Polkadot 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 Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+use criterion::{criterion_group, criterion_main, Criterion, SamplingMode};
+use polkadot_node_core_pvf_common::{prepare::PrepareJobKind, pvf::PvfPrepData};
+use polkadot_node_core_pvf_prepare_worker::{prepare, prevalidate};
+use polkadot_primitives::ExecutorParams;
+use std::time::Duration;
+
+fn do_prepare_kusama_runtime(pvf: PvfPrepData) {
+	let blob = match prevalidate(&pvf.code()) {
+		Err(err) => panic!("{:?}", err),
+		Ok(b) => b,
+	};
+
+	match prepare(blob, &pvf.executor_params()) {
+		Ok(_) => (),
+		Err(err) => panic!("{:?}", err),
+	}
+}
+
+fn prepare_kusama_runtime(c: &mut Criterion) {
+	let blob = staging_kusama_runtime::WASM_BINARY.unwrap();
+	let pvf = match sp_maybe_compressed_blob::decompress(&blob, 64 * 1024 * 1024) {
+		Ok(code) => PvfPrepData::from_code(
+			code.into_owned(),
+			ExecutorParams::default(),
+			Duration::from_secs(360),
+			PrepareJobKind::Compilation,
+		),
+		Err(e) => {
+			panic!("Cannot decompress blob: {:?}", e);
+		},
+	};
+
+	let mut group = c.benchmark_group("kusama");
+	group.sampling_mode(SamplingMode::Flat);
+	group.sample_size(20);
+	group.measurement_time(Duration::from_secs(240));
+	group.bench_function("prepare Kusama runtime", |b| {
+		b.iter(|| do_prepare_kusama_runtime(pvf.clone()))
+	});
+	group.finish();
+}
+
+criterion_group!(preparation, prepare_kusama_runtime);
+criterion_main!(preparation);

polkadot/node/core/pvf/prepare-worker/src/lib.rs

@@ -53,6 +53,14 @@ use std::{
 };
 use tokio::{io, net::UnixStream};
 
+#[cfg(feature = "tracking-allocator")]
+use tikv_jemallocator::Jemalloc;
+#[cfg(feature = "tracking-allocator")]
+use tracking_allocator::TrackingAllocator;
+#[cfg(feature = "tracking-allocator")]
+#[global_allocator]
+static ALLOC: TrackingAllocator<Jemalloc> = TrackingAllocator(Jemalloc);
+
 /// Contains the bytes for a successfully compiled artifact.
 pub struct CompiledArtifact(Vec<u8>);
 
@@ -168,13 +176,17 @@ pub fn worker_entrypoint(
 					Arc::clone(&condvar),
 					WaitOutcome::TimedOut,
 				)?;
+
+				#[cfg(feature = "tracking-allocator")]
+				ALLOC.start_tracking(executor_params.prechecking_max_memory().map(|v| v as isize));
+
 				// Spawn another thread for preparation.
 				let prepare_thread = thread::spawn_worker_thread(
 					"prepare thread",
 					move || {
 						// Try to enable landlock.
 						#[cfg(target_os = "linux")]
-					let landlock_status = polkadot_node_core_pvf_common::worker::security::landlock::try_restrict_thread()
+						let landlock_status = polkadot_node_core_pvf_common::worker::security::landlock::try_restrict_thread()
 						.map(LandlockStatus::from_ruleset_status)
 						.map_err(|e| e.to_string());
 						#[cfg(not(target_os = "linux"))]
@@ -208,6 +220,18 @@ pub fn worker_entrypoint(
 
 				let outcome = thread::wait_for_threads(condvar);
 
+				#[cfg(feature = "tracking-allocator")]
+				let peak_alloc = {
+					let peak = ALLOC.end_tracking();
+					gum::debug!(
+						target: LOG_TARGET,
+						%worker_pid,
+						"prepare job peak allocation is {} bytes",
+						peak,
+					);
+					peak
+				};
+
 				let result = match outcome {
 					WaitOutcome::Finished => {
 						let _ = cpu_time_monitor_tx.send(());
@@ -240,6 +264,8 @@ pub fn worker_entrypoint(
 									memory_tracker_stats,
 									#[cfg(target_os = "linux")]
 									max_rss: extract_max_rss_stat(max_rss, worker_pid),
+									#[cfg(feature = "tracking-allocator")]
+									peak_alloc: peak_alloc as u64,
 								};
 
 								// Log if landlock threw an error.

polkadot/node/core/pvf/src/metrics.rs

@@ -93,6 +93,11 @@ impl Metrics {
 				metrics.preparation_max_resident.observe(max_resident_kb);
 				metrics.preparation_max_allocated.observe(max_allocated_kb);
 			}
+
+			#[cfg(feature = "tracking-allocator")]
+			metrics
+				.preparation_peak_allocation
+				.observe((memory_stats.peak_alloc / 1024) as f64);
 		}
 	}
 }
@@ -114,6 +119,8 @@ struct MetricsInner {
 	preparation_max_allocated: prometheus::Histogram,
 	#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 	preparation_max_resident: prometheus::Histogram,
+	#[cfg(feature = "tracking-allocator")]
+	preparation_peak_allocation: prometheus::Histogram,
 }
 
 impl metrics::Metrics for Metrics {
@@ -271,6 +278,19 @@ impl metrics::Metrics for Metrics {
 				)?,
 				registry,
 			)?,
+			#[cfg(feature = "tracking-allocator")]
+			preparation_peak_allocation: prometheus::register(
+				prometheus::Histogram::with_opts(
+					prometheus::HistogramOpts::new(
+						"polkadot_pvf_preparation_peak_allocattion",
+						"peak allocation observed for preparation (in kilobytes)",
+					).buckets(
+						prometheus::exponential_buckets(8192.0, 2.0, 10)
+							.expect("arguments are always valid; qed"),
+					),
+				)?,
+				registry,
+			)?,
 		};
 		Ok(Metrics(Some(inner)))
 	}