[DNM] Wrapper allocator PoC

Cargo.toml

@@ -37,6 +37,7 @@ polkadot-cli = { path = "cli", features = [ "polkadot-native", "kusama-native",
 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", features = ["test-utils"] }
@@ -121,6 +122,7 @@ members = [
 	"node/subsystem-types",
 	"node/subsystem-test-helpers",
 	"node/subsystem-util",
+	"node/tracking-allocator",
 	"node/jaeger",
 	"node/gum",
 	"node/gum/proc-macro",
@@ -227,6 +229,8 @@ fast-runtime = [ "polkadot-cli/fast-runtime" ]
 runtime-metrics = [ "polkadot-cli/runtime-metrics" ]
 pyroscope = ["polkadot-cli/pyroscope"]
 jemalloc-allocator = ["polkadot-node-core-pvf-prepare-worker/jemalloc-allocator", "polkadot-overseer/jemalloc-allocator"]
+tracking-allocator = ["jemalloc-allocator", "dep:tracking-allocator", "polkadot-node-core-pvf-prepare-worker/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
 ci-only-tests = ["polkadot-node-core-pvf/ci-only-tests"]

cli/src/cli.rs

@@ -99,11 +99,6 @@ pub struct RunCmd {
 	#[arg(long = "grandpa-pause", num_args = 2)]
 	pub grandpa_pause: Vec<u32>,
 
-	/// Disable the BEEFY gadget
-	/// (currently enabled by default on Rococo, Wococo and Versi).
-	#[arg(long)]
-	pub no_beefy: bool,
-
 	/// Add the destination address to the jaeger agent.
 	///
 	/// Must be valid socket address, of format `IP:Port`

cli/src/command.rs

@@ -235,15 +235,11 @@ fn run_node_inner<F>(
 where
 	F: FnOnce(&mut sc_cli::LoggerBuilder, &sc_service::Configuration),
 {
-	let runner = cli
+	let mut runner = cli
 		.create_runner_with_logger_hook::<sc_cli::RunCmd, F>(&cli.run.base, logger_hook)
 		.map_err(Error::from)?;
 	let chain_spec = &runner.config().chain_spec;
 
-	// By default, enable BEEFY on test networks.
-	let enable_beefy = (chain_spec.is_rococo() || chain_spec.is_wococo() || chain_spec.is_versi()) &&
-		!cli.run.no_beefy;
-
 	set_default_ss58_version(chain_spec);
 
 	let grandpa_pause = if cli.run.grandpa_pause.is_empty() {
@@ -259,6 +255,10 @@ where
 		info!("     KUSAMA FOUNDATION      ");
 		info!("----------------------------");
 	}
+	// BEEFY allowed only on test networks.
+	if !(chain_spec.is_rococo() || chain_spec.is_wococo() || chain_spec.is_versi()) {
+		runner.config_mut().disable_beefy = true;
+	}
 
 	let jaeger_agent = if let Some(ref jaeger_agent) = cli.run.jaeger_agent {
 		Some(
@@ -289,7 +289,6 @@ where
 			service::NewFullParams {
 				is_parachain_node: service::IsParachainNode::No,
 				grandpa_pause,
-				enable_beefy,
 				jaeger_agent,
 				telemetry_worker_handle: None,
 				node_version,

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

@@ -13,6 +13,7 @@ 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 }
 
 parity-scale-codec = { version = "3.6.1", default-features = false, features = ["derive"] }
 
@@ -33,3 +34,12 @@ tikv-jemalloc-ctl = "0.5.0"
 [features]
 builder = []
 jemalloc-allocator = ["dep:tikv-jemalloc-ctl"]
+tracking-allocator = ["dep:tracking-allocator"]
+
+[dev-dependencies]
+criterion = { version = "0.4.0", default-features = false, features = ["cargo_bench_support"] }
+kusama-runtime = { path = "../../../../runtime/kusama" }
+
+[[bench]]
+name = "prepare_kusama_runtime"
+harness = false

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 = 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);

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

@@ -52,6 +52,8 @@ use std::{
 	time::Duration,
 };
 use tokio::{io, net::UnixStream};
+#[cfg(feature = "tracking-allocator")]
+use tracking_allocator::ALLOC;
 
 /// Contains the bytes for a successfully compiled artifact.
 pub struct CompiledArtifact(Vec<u8>);
@@ -180,9 +182,23 @@ pub fn worker_entrypoint(
 						#[cfg(not(target_os = "linux"))]
 						let landlock_status: Result<LandlockStatus, String> = Ok(LandlockStatus::NotEnforced);
 
+						#[cfg(feature = "tracking-allocator")]
+						ALLOC.start_tracking();
+
 						#[allow(unused_mut)]
 						let mut result = prepare_artifact(pvf, cpu_time_start);
 
+						#[cfg(feature = "tracking-allocator")]
+						{
+							let peak = ALLOC.end_tracking();
+							gum::debug!(
+								target: LOG_TARGET,
+								%worker_pid,
+								"prepare job peak allocation is {} bytes",
+								peak,
+							);
+						}
+
 						// Get the `ru_maxrss` stat. If supported, call getrusage for the thread.
 						#[cfg(target_os = "linux")]
 						let mut result = result

node/service/src/lib.rs

@@ -629,7 +629,6 @@ where
 pub struct NewFullParams<OverseerGenerator: OverseerGen> {
 	pub is_parachain_node: IsParachainNode,
 	pub grandpa_pause: Option<(u32, u32)>,
-	pub enable_beefy: bool,
 	pub jaeger_agent: Option<std::net::SocketAddr>,
 	pub telemetry_worker_handle: Option<TelemetryWorkerHandle>,
 	/// The version of the node. TESTING ONLY: `None` can be passed to skip the node/worker version
@@ -720,7 +719,6 @@ pub fn new_full<OverseerGenerator: OverseerGen>(
 	NewFullParams {
 		is_parachain_node,
 		grandpa_pause,
-		enable_beefy,
 		jaeger_agent,
 		telemetry_worker_handle,
 		node_version,
@@ -754,6 +752,7 @@ pub fn new_full<OverseerGenerator: OverseerGen>(
 		Some(backoff)
 	};
 
+	let enable_beefy = !config.disable_beefy;
 	// If not on a known test network, warn the user that BEEFY is still experimental.
 	if enable_beefy &&
 		!config.chain_spec.is_rococo() &&

node/test/service/src/lib.rs

@@ -81,7 +81,6 @@ pub fn new_full(
 		polkadot_service::NewFullParams {
 			is_parachain_node,
 			grandpa_pause: None,
-			enable_beefy: true,
 			jaeger_agent: None,
 			telemetry_worker_handle: None,
 			node_version: None,
@@ -187,6 +186,7 @@ pub fn node_config(
 		offchain_worker: Default::default(),
 		force_authoring: false,
 		disable_grandpa: false,
+		disable_beefy: false,
 		dev_key_seed: Some(key_seed),
 		tracing_targets: None,
 		tracing_receiver: Default::default(),

node/tracking-allocator/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "tracking-allocator"
+description = "Tracking allocator to control amount of memory consumed by PVF preparation process"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+
+[dependencies]
+tikv-jemallocator = "0.5.0"

node/tracking-allocator/src/lib.rs

@@ -0,0 +1,136 @@
+// 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/>.
+
+//! Tracking global allocator. Calculates the peak allocation between two checkpoints.
+
+use core::alloc::{GlobalAlloc, Layout};
+use std::sync::atomic::{AtomicBool, Ordering};
+use tikv_jemallocator::Jemalloc;
+
+struct TrackingAllocatorData {
+	lock: AtomicBool,
+	current: isize,
+	peak: isize,
+}
+
+impl TrackingAllocatorData {
+	#[inline]
+	fn lock(&self) {
+		loop {
+			// Try to acquire the lock.
+			if self
+				.lock
+				.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
+				.is_ok()
+			{
+				break
+			}
+			// We failed to acquire the lock; wait until it's unlocked.
+			//
+			// In theory this should result in less coherency traffic as unlike `compare_exchange`
+			// it is a read-only operation, so multiple cores can execute it simultaneously
+			// without taking an exclusive lock over the cache line.
+			while self.lock.load(Ordering::Relaxed) {
+				std::hint::spin_loop();
+			}
+		}
+	}
+
+	#[inline]
+	fn unlock(&self) {
+		self.lock.store(false, Ordering::Release);
+	}
+
+	fn start_tracking(&mut self) {
+		self.lock();
+		self.current = 0;
+		self.peak = 0;
+		self.unlock();
+	}
+
+	fn end_tracking(&self) -> isize {
+		self.lock();
+		let peak = self.peak;
+		self.unlock();
+		peak
+	}
+
+	#[inline]
+	fn track(&mut self, alloc: isize) {
+		self.lock();
+		self.current += alloc;
+		if self.current > self.peak {
+			self.peak = self.current;
+		}
+		self.unlock();
+	}
+}
+
+static mut ALLOCATOR_DATA: TrackingAllocatorData =
+	TrackingAllocatorData { lock: AtomicBool::new(false), current: 0, peak: 0 };
+
+pub struct TrackingAllocator<A: GlobalAlloc>(A);
+
+impl<A: GlobalAlloc> TrackingAllocator<A> {
+	// SAFETY:
+	// * The following functions write to `static mut`. That is safe as the critical section
+	//   inside is isolated by an exclusive lock.
+
+	/// Start tracking
+	pub fn start_tracking(&self) {
+		unsafe {
+			ALLOCATOR_DATA.start_tracking();
+		}
+	}
+
+	/// End tracking and return the peak allocation value in bytes (as `isize`). Peak allocation
+	/// value is not guaranteed to be neither non-zero nor positive.
+	pub fn end_tracking(&self) -> isize {
+		unsafe { ALLOCATOR_DATA.end_tracking() }
+	}
+}
+
+unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
+	// SAFETY:
+	// * The wrapped methods are as safe as the underlying allocator implementation is
+
+	#[inline]
+	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+		ALLOCATOR_DATA.track(layout.size() as isize);
+		self.0.alloc(layout)
+	}
+
+	#[inline]
+	unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+		ALLOCATOR_DATA.track(layout.size() as isize);
+		self.0.alloc_zeroed(layout)
+	}
+
+	#[inline]
+	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) -> () {
+		ALLOCATOR_DATA.track(-(layout.size() as isize));
+		self.0.dealloc(ptr, layout)
+	}
+
+	#[inline]
+	unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+		ALLOCATOR_DATA.track((new_size as isize) - (layout.size() as isize));
+		self.0.realloc(ptr, layout, new_size)
+	}
+}
+
+#[global_allocator]
+pub static ALLOC: TrackingAllocator<Jemalloc> = TrackingAllocator(Jemalloc);

src/main.rs

@@ -22,7 +22,10 @@ use color_eyre::eyre;
 
 /// Global allocator. Changing it to another allocator will require changing
 /// `memory_stats::MemoryAllocationTracker`.
-#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
+#[cfg(all(
+	any(target_os = "linux", feature = "jemalloc-allocator"),
+	not(feature = "wrapper-allocator")
+))]
 #[global_allocator]
 pub static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;