Import changes from archived repo

polkadot/Cargo.toml

@@ -33,6 +33,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" }
@@ -139,6 +140,11 @@ 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

polkadot/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"] }
 
@@ -32,4 +33,13 @@ tikv-jemalloc-ctl = "0.5.0"
 
 [features]
 builder = []
-jemalloc-allocator = [ "dep:tikv-jemalloc-ctl" ]
+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

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

@@ -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

polkadot/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"

polkadot/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);

polkadot/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;