From 47ecf58b983020c6d195e7e817a2c9123b8d0b53 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Sun, 27 Aug 2023 13:56:26 +0200
Subject: [PATCH 01/33] Import changes from archived repo

---
 Cargo.lock                                    |  11 ++
 polkadot/Cargo.toml                           |   6 +
 .../node/core/pvf/prepare-worker/Cargo.toml   |  12 +-
 .../benches/prepare_kusama_runtime.rs         |  60 ++++++++
 .../node/core/pvf/prepare-worker/src/lib.rs   |  16 +++
 polkadot/node/tracking-allocator/Cargo.toml   |   9 ++
 polkadot/node/tracking-allocator/src/lib.rs   | 136 ++++++++++++++++++
 polkadot/src/main.rs                          |   5 +-
 8 files changed, 253 insertions(+), 2 deletions(-)
 create mode 100644 polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
 create mode 100644 polkadot/node/tracking-allocator/Cargo.toml
 create mode 100644 polkadot/node/tracking-allocator/src/lib.rs

diff --git a/Cargo.lock b/Cargo.lock
index c755be63042b..f250fd2ff25b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -11631,6 +11631,7 @@ dependencies = [
  "tempfile",
  "tikv-jemallocator",
  "tokio",
+ "tracking-allocator",
 ]
 
 [[package]]
@@ -12309,7 +12310,9 @@ dependencies = [
 name = "polkadot-node-core-pvf-prepare-worker"
 version = "1.0.0"
 dependencies = [
+ "criterion 0.4.0",
  "futures",
+ "kusama-runtime",
  "libc",
  "parity-scale-codec",
  "polkadot-node-core-pvf-common",
@@ -12325,6 +12328,7 @@ dependencies = [
  "tikv-jemalloc-ctl",
  "tokio",
  "tracing-gum",
+ "tracking-allocator",
 ]
 
 [[package]]
@@ -19152,6 +19156,13 @@ dependencies = [
  "tracing-log",
 ]
 
+[[package]]
+name = "tracking-allocator"
+version = "1.0.0"
+dependencies = [
+ "tikv-jemallocator",
+]
+
 [[package]]
 name = "trie-bench"
 version = "0.37.0"
diff --git a/polkadot/Cargo.toml b/polkadot/Cargo.toml
index 561b49ab4204..7694411b46ab 100644
--- a/polkadot/Cargo.toml
+++ b/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
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index 61d2fd971564..f52798780f46 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/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
diff --git a/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs b/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
new file mode 100644
index 000000000000..7e5d4e968a96
--- /dev/null
+++ b/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);
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index caa7d33df12a..ac116cf78631 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/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
diff --git a/polkadot/node/tracking-allocator/Cargo.toml b/polkadot/node/tracking-allocator/Cargo.toml
new file mode 100644
index 000000000000..81f95b923398
--- /dev/null
+++ b/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"
diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
new file mode 100644
index 000000000000..5b51f26184ab
--- /dev/null
+++ b/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);
diff --git a/polkadot/src/main.rs b/polkadot/src/main.rs
index 5986d8cea7bb..9f614507f6fd 100644
--- a/polkadot/src/main.rs
+++ b/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;
 

From 58330dee61c7e554b4853d5841eafe968a98be28 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Wed, 30 Aug 2023 11:08:20 +0200
Subject: [PATCH 02/33] Fix dependencies

---
 Cargo.lock                                       | 4 +---
 polkadot/node/core/pvf/prepare-worker/Cargo.toml | 3 ++-
 polkadot/node/core/pvf/prepare-worker/src/lib.rs | 8 +++++++-
 polkadot/node/tracking-allocator/Cargo.toml      | 3 ---
 polkadot/node/tracking-allocator/src/lib.rs      | 6 +-----
 5 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f250fd2ff25b..dd60250b5431 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -12326,6 +12326,7 @@ dependencies = [
  "sp-maybe-compressed-blob",
  "sp-tracing",
  "tikv-jemalloc-ctl",
+ "tikv-jemallocator",
  "tokio",
  "tracing-gum",
  "tracking-allocator",
@@ -19159,9 +19160,6 @@ dependencies = [
 [[package]]
 name = "tracking-allocator"
 version = "1.0.0"
-dependencies = [
- "tikv-jemallocator",
-]
 
 [[package]]
 name = "trie-bench"
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index f52798780f46..0721ae9c475c 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
@@ -14,6 +14,7 @@ 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"] }
 
@@ -34,7 +35,7 @@ tikv-jemalloc-ctl = "0.5.0"
 [features]
 builder = []
 jemalloc-allocator = ["dep:tikv-jemalloc-ctl"]
-tracking-allocator = ["dep:tracking-allocator"]
+tracking-allocator = ["dep:tracking-allocator", "dep:tikv-jemallocator"]
 
 [dev-dependencies]
 criterion = { version = "0.4.0", default-features = false, features = ["cargo_bench_support"] }
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index ac116cf78631..43b7e60220e8 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -52,8 +52,14 @@ use std::{
 	time::Duration,
 };
 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")]
-use tracking_allocator::ALLOC;
+#[global_allocator]
+static ALLOC: TrackingAllocator<Jemalloc> = TrackingAllocator(Jemalloc);
 
 /// Contains the bytes for a successfully compiled artifact.
 pub struct CompiledArtifact(Vec<u8>);
diff --git a/polkadot/node/tracking-allocator/Cargo.toml b/polkadot/node/tracking-allocator/Cargo.toml
index 81f95b923398..0edbb5fcc741 100644
--- a/polkadot/node/tracking-allocator/Cargo.toml
+++ b/polkadot/node/tracking-allocator/Cargo.toml
@@ -4,6 +4,3 @@ description = "Tracking allocator to control amount of memory consumed by PVF pr
 version.workspace = true
 authors.workspace = true
 edition.workspace = true
-
-[dependencies]
-tikv-jemallocator = "0.5.0"
diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 5b51f26184ab..40d0022803f6 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -18,7 +18,6 @@
 
 use core::alloc::{GlobalAlloc, Layout};
 use std::sync::atomic::{AtomicBool, Ordering};
-use tikv_jemallocator::Jemalloc;
 
 struct TrackingAllocatorData {
 	lock: AtomicBool,
@@ -82,7 +81,7 @@ impl TrackingAllocatorData {
 static mut ALLOCATOR_DATA: TrackingAllocatorData =
 	TrackingAllocatorData { lock: AtomicBool::new(false), current: 0, peak: 0 };
 
-pub struct TrackingAllocator<A: GlobalAlloc>(A);
+pub struct TrackingAllocator<A: GlobalAlloc>(pub A);
 
 impl<A: GlobalAlloc> TrackingAllocator<A> {
 	// SAFETY:
@@ -131,6 +130,3 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 		self.0.realloc(ptr, layout, new_size)
 	}
 }
-
-#[global_allocator]
-pub static ALLOC: TrackingAllocator<Jemalloc> = TrackingAllocator(Jemalloc);

From 6ee5ff354a7284e8349fcd128e6c3a08084a3f1a Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Thu, 14 Sep 2023 14:58:10 +0200
Subject: [PATCH 03/33] Format manifests

---
 polkadot/Cargo.toml                              | 4 ++--
 polkadot/node/core/pvf/prepare-worker/Cargo.toml | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/polkadot/Cargo.toml b/polkadot/Cargo.toml
index b7b6023d73e4..e729b79961bc 100644
--- a/polkadot/Cargo.toml
+++ b/polkadot/Cargo.toml
@@ -70,9 +70,9 @@ jemalloc-allocator = [
 	"polkadot-overseer/jemalloc-allocator",
 ]
 tracking-allocator = [
-	"jemalloc-allocator",
 	"dep:tracking-allocator",
-	"polkadot-node-core-pvf-prepare-worker/tracking-allocator"
+	"jemalloc-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
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index c6b314508b60..fb6f349b20a8 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
@@ -38,7 +38,7 @@ jemalloc-allocator = [
 	"dep:tikv-jemalloc-ctl",
 	"polkadot-node-core-pvf-common/jemalloc-allocator",
 ]
-tracking-allocator = ["dep:tracking-allocator", "dep:tikv-jemallocator"]
+tracking-allocator = [ "dep:tikv-jemallocator", "dep:tracking-allocator" ]
 
 [dev-dependencies]
 criterion = { version = "0.4.0", default-features = false, features = ["cargo_bench_support"] }

From f85c753d5514c1431c99dfc27334eda894986641 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Fri, 15 Sep 2023 20:01:29 +0200
Subject: [PATCH 04/33] Make peak allocation value observable

---
 polkadot/Cargo.toml                           |  1 +
 polkadot/node/core/pvf/Cargo.toml             |  1 +
 polkadot/node/core/pvf/common/Cargo.toml      |  1 +
 polkadot/node/core/pvf/common/src/prepare.rs  |  3 ++
 .../node/core/pvf/prepare-worker/Cargo.toml   |  6 +++-
 .../node/core/pvf/prepare-worker/src/lib.rs   | 34 +++++++++++--------
 polkadot/node/core/pvf/src/metrics.rs         | 20 +++++++++++
 7 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/polkadot/Cargo.toml b/polkadot/Cargo.toml
index e729b79961bc..e1832c68ae16 100644
--- a/polkadot/Cargo.toml
+++ b/polkadot/Cargo.toml
@@ -73,6 +73,7 @@ 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
diff --git a/polkadot/node/core/pvf/Cargo.toml b/polkadot/node/core/pvf/Cargo.toml
index 478d1952d9d9..7f158ce82443 100644
--- a/polkadot/node/core/pvf/Cargo.toml
+++ b/polkadot/node/core/pvf/Cargo.toml
@@ -49,6 +49,7 @@ 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" ]
 # 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",
diff --git a/polkadot/node/core/pvf/common/Cargo.toml b/polkadot/node/core/pvf/common/Cargo.toml
index 621f7e24f72b..53f891678e63 100644
--- a/polkadot/node/core/pvf/common/Cargo.toml
+++ b/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 = []
diff --git a/polkadot/node/core/pvf/common/src/prepare.rs b/polkadot/node/core/pvf/common/src/prepare.rs
index c205eddfb8b1..704f8d06aaae 100644
--- a/polkadot/node/core/pvf/common/src/prepare.rs
+++ b/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.
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index fb6f349b20a8..52ecb60d3eb1 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
@@ -38,7 +38,11 @@ jemalloc-allocator = [
 	"dep:tikv-jemalloc-ctl",
 	"polkadot-node-core-pvf-common/jemalloc-allocator",
 ]
-tracking-allocator = [ "dep:tikv-jemallocator", "dep:tracking-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"] }
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index 43b7e60220e8..292733e43fa8 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -176,35 +176,25 @@ pub fn worker_entrypoint(
 					Arc::clone(&condvar),
 					WaitOutcome::TimedOut,
 				)?;
+
+				#[cfg(feature = "tracking-allocator")]
+				ALLOC.start_tracking();
+
 				// 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"))]
 						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
@@ -230,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(());
@@ -262,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.
diff --git a/polkadot/node/core/pvf/src/metrics.rs b/polkadot/node/core/pvf/src/metrics.rs
index 3d792793498b..fdf3dacbe9ff 100644
--- a/polkadot/node/core/pvf/src/metrics.rs
+++ b/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)))
 	}

From edfa2e18f41b5c4cd8e5df0e140745572e3110f8 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Sat, 16 Sep 2023 14:31:44 +0200
Subject: [PATCH 05/33] Enforce memory limits

---
 polkadot/node/core/pvf/Cargo.toml             |  5 +-
 .../node/core/pvf/prepare-worker/src/lib.rs   |  2 +-
 polkadot/node/core/pvf/tests/it/main.rs       | 88 +++++++++++++++++--
 polkadot/node/tracking-allocator/src/lib.rs   | 39 +++++---
 polkadot/primitives/src/v5/executor_params.rs | 10 +++
 5 files changed, 121 insertions(+), 23 deletions(-)

diff --git a/polkadot/node/core/pvf/Cargo.toml b/polkadot/node/core/pvf/Cargo.toml
index 7f158ce82443..a5c1fa7389c6 100644
--- a/polkadot/node/core/pvf/Cargo.toml
+++ b/polkadot/node/core/pvf/Cargo.toml
@@ -49,7 +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" ]
+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",
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index 292733e43fa8..bbdb7445c2da 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -178,7 +178,7 @@ pub fn worker_entrypoint(
 				)?;
 
 				#[cfg(feature = "tracking-allocator")]
-				ALLOC.start_tracking();
+				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(
diff --git a/polkadot/node/core/pvf/tests/it/main.rs b/polkadot/node/core/pvf/tests/it/main.rs
index dc8f00098ec5..13d71c9b751f 100644
--- a/polkadot/node/core/pvf/tests/it/main.rs
+++ b/polkadot/node/core/pvf/tests/it/main.rs
@@ -21,12 +21,20 @@ use polkadot_node_core_pvf::{
 	start, Config, InvalidCandidate, Metrics, PrepareJobKind, PvfPrepData, ValidationError,
 	ValidationHost, JOB_TIMEOUT_WALL_CLOCK_FACTOR,
 };
-use polkadot_parachain_primitives::primitives::{BlockData, ValidationParams, ValidationResult};
+use polkadot_parachain_primitives::primitives::{
+	BlockData as GenericBlockData, ValidationParams, ValidationResult,
+};
 use polkadot_primitives::ExecutorParams;
 
-#[cfg(feature = "ci-only-tests")]
+#[cfg(any(feature = "ci-only-tests", feature = "tracking-allocator"))]
 use polkadot_primitives::ExecutorParam;
 
+#[cfg(feature = "tracking-allocator")]
+use ::adder::{hash_state, BlockData, HeadData};
+
+#[cfg(feature = "tracking-allocator")]
+use polkadot_primitives::HeadData as GenericHeadData;
+
 use std::time::Duration;
 use tokio::sync::Mutex;
 
@@ -111,7 +119,7 @@ async fn terminates_on_timeout() {
 		.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: BlockData(Vec::new()),
+				block_data: GenericBlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -138,7 +146,7 @@ async fn ensure_parallel_execution() {
 	let execute_pvf_future_1 = host.validate_candidate(
 		halt::wasm_binary_unwrap(),
 		ValidationParams {
-			block_data: BlockData(Vec::new()),
+			block_data: GenericBlockData(Vec::new()),
 			parent_head: Default::default(),
 			relay_parent_number: 1,
 			relay_parent_storage_root: Default::default(),
@@ -148,7 +156,7 @@ async fn ensure_parallel_execution() {
 	let execute_pvf_future_2 = host.validate_candidate(
 		halt::wasm_binary_unwrap(),
 		ValidationParams {
-			block_data: BlockData(Vec::new()),
+			block_data: GenericBlockData(Vec::new()),
 			parent_head: Default::default(),
 			relay_parent_number: 1,
 			relay_parent_storage_root: Default::default(),
@@ -191,7 +199,7 @@ async fn execute_queue_doesnt_stall_if_workers_died() {
 		host.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: BlockData(Vec::new()),
+				block_data: GenericBlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -233,7 +241,7 @@ async fn execute_queue_doesnt_stall_with_varying_executor_params() {
 		host.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: BlockData(Vec::new()),
+				block_data: GenericBlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -273,7 +281,7 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 		.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: BlockData(Vec::new()),
+				block_data: GenericBlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -303,7 +311,7 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 		.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: BlockData(Vec::new()),
+				block_data: GenericBlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -317,3 +325,65 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 		r => panic!("{:?}", r),
 	}
 }
+
+// This test checks if the adder parachain runtime can be prepared with 10Mb preparation memory
+// limit enforced. At the moment of writing, the limit if far enough to prepare the PVF. If it
+// starts failing, either Wasmtime version has changed, or the PVF code itself has changed, and
+// more memory is required now. Multi-threaded preparation, if ever enabled, may also affect
+// memory consumption.
+#[cfg(feature = "tracking-allocator")]
+#[tokio::test]
+async fn prechecking_within_memory_limits() {
+	let host = TestHost::new();
+	let parent_head = HeadData { number: 0, parent_hash: [0; 32], post_state: hash_state(0) };
+	let block_data = BlockData { state: 0, add: 512 };
+	let result = host
+		.validate_candidate(
+			::adder::wasm_binary_unwrap(),
+			ValidationParams {
+				parent_head: GenericHeadData(parent_head.encode()),
+				block_data: GenericBlockData(block_data.encode()),
+				relay_parent_number: 1,
+				relay_parent_storage_root: Default::default(),
+			},
+			ExecutorParams::from(&[ExecutorParam::PrecheckingMaxMemory(10 * 1024 * 1024)][..]),
+		)
+		.await;
+
+	match result {
+		Ok(_) => (),
+		r => panic!("{:?}", r),
+	}
+}
+
+// This test checks if the adder parachain runtime can be prepared with 512Kb preparation memory
+// limit enforced. At the moment of writing, the limit if not enough to prepare the PVF, and the
+// preparation is supposed to generate an error. If the test starts failing, either Wasmtime
+// version has changed, or the PVF code itself has changed, and less memory is required now.
+#[cfg(feature = "tracking-allocator")]
+#[tokio::test]
+async fn prechecking_out_of_memory() {
+	use polkadot_node_core_pvf::ValidationError::InternalError;
+	use polkadot_node_core_pvf_common::error::InternalValidationError::NonDeterministicPrepareError;
+
+	let host = TestHost::new();
+	let parent_head = HeadData { number: 0, parent_hash: [0; 32], post_state: hash_state(0) };
+	let block_data = BlockData { state: 0, add: 512 };
+	let result = host
+		.validate_candidate(
+			::adder::wasm_binary_unwrap(),
+			ValidationParams {
+				parent_head: GenericHeadData(parent_head.encode()),
+				block_data: GenericBlockData(block_data.encode()),
+				relay_parent_number: 1,
+				relay_parent_storage_root: Default::default(),
+			},
+			ExecutorParams::from(&[ExecutorParam::PrecheckingMaxMemory(512 * 1024)][..]),
+		)
+		.await;
+
+	match result {
+		Err(InternalError(NonDeterministicPrepareError(_))) => (),
+		r => panic!("{:?}", r),
+	}
+}
diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 386488e871df..2f719a09cf50 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -18,11 +18,13 @@
 
 use core::alloc::{GlobalAlloc, Layout};
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::ptr::null_mut;
 
 struct TrackingAllocatorData {
 	lock: AtomicBool,
 	current: isize,
 	peak: isize,
+	limit: isize,
 }
 
 impl TrackingAllocatorData {
@@ -53,33 +55,37 @@ impl TrackingAllocatorData {
 		self.lock.store(false, Ordering::Release);
 	}
 
-	fn start_tracking(&mut self) {
+	fn start_tracking(&mut self, limit: isize) {
 		self.lock();
 		self.current = 0;
 		self.peak = 0;
+		self.limit = limit;
 		self.unlock();
 	}
 
-	fn end_tracking(&self) -> isize {
+	fn end_tracking(&mut self) -> isize {
 		self.lock();
 		let peak = self.peak;
+		self.limit = 0;
 		self.unlock();
 		peak
 	}
 
 	#[inline]
-	fn track(&mut self, alloc: isize) {
+	fn track(&mut self, alloc: isize) -> bool {
 		self.lock();
 		self.current += alloc;
 		if self.current > self.peak {
 			self.peak = self.current;
 		}
+		let within_limits = self.limit == 0 || self.peak <= self.limit;
 		self.unlock();
+		within_limits
 	}
 }
 
 static mut ALLOCATOR_DATA: TrackingAllocatorData =
-	TrackingAllocatorData { lock: AtomicBool::new(false), current: 0, peak: 0 };
+	TrackingAllocatorData { lock: AtomicBool::new(false), current: 0, peak: 0, limit: 0 };
 
 pub struct TrackingAllocator<A: GlobalAlloc>(pub A);
 
@@ -89,9 +95,9 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 	//   is isolated by an exclusive lock.
 
 	/// Start tracking
-	pub fn start_tracking(&self) {
+	pub fn start_tracking(&self, limit: Option<isize>) {
 		unsafe {
-			ALLOCATOR_DATA.start_tracking();
+			ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0));
 		}
 	}
 
@@ -108,14 +114,20 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 
 	#[inline]
 	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-		ALLOCATOR_DATA.track(layout.size() as isize);
-		self.0.alloc(layout)
+		if ALLOCATOR_DATA.track(layout.size() as isize) {
+			self.0.alloc(layout)
+		} else {
+			null_mut()
+		}
 	}
 
 	#[inline]
 	unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
-		ALLOCATOR_DATA.track(layout.size() as isize);
-		self.0.alloc_zeroed(layout)
+		if ALLOCATOR_DATA.track(layout.size() as isize) {
+			self.0.alloc_zeroed(layout)
+		} else {
+			null_mut()
+		}
 	}
 
 	#[inline]
@@ -126,7 +138,10 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 
 	#[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)
+		if ALLOCATOR_DATA.track((new_size as isize) - (layout.size() as isize)) {
+			self.0.realloc(ptr, layout, new_size)
+		} else {
+			null_mut()
+		}
 	}
 }
diff --git a/polkadot/primitives/src/v5/executor_params.rs b/polkadot/primitives/src/v5/executor_params.rs
index 6fbf3037fd6c..7b11a8f06fc2 100644
--- a/polkadot/primitives/src/v5/executor_params.rs
+++ b/polkadot/primitives/src/v5/executor_params.rs
@@ -130,6 +130,16 @@ impl ExecutorParams {
 		}
 		None
 	}
+
+	/// Returns pre-checking memory limit, if any
+	pub fn prechecking_max_memory(&self) -> Option<u64> {
+		for param in &self.0 {
+			if let ExecutorParam::PrecheckingMaxMemory(limit) = param {
+				return Some(*limit)
+			}
+		}
+		None
+	}
 }
 
 impl Deref for ExecutorParams {

From 994273aeb129dd8980a7e16f92aaaa28e704c273 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Sat, 16 Sep 2023 14:57:18 +0200
Subject: [PATCH 06/33] Fix the node allocator declaration

---
 polkadot/src/main.rs | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/polkadot/src/main.rs b/polkadot/src/main.rs
index 9f614507f6fd..8b125d3d8409 100644
--- a/polkadot/src/main.rs
+++ b/polkadot/src/main.rs
@@ -20,12 +20,8 @@
 
 use color_eyre::eyre;
 
-/// Global allocator. Changing it to another allocator will require changing
-/// `memory_stats::MemoryAllocationTracker`.
-#[cfg(all(
-	any(target_os = "linux", feature = "jemalloc-allocator"),
-	not(feature = "wrapper-allocator")
-))]
+#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
+#[doc(hidden)]
 #[global_allocator]
 pub static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 

From 72ded39b2f297a3c47de0b1d13e57e68939ab9a0 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Sat, 16 Sep 2023 15:03:58 +0200
Subject: [PATCH 07/33] `cargo fmt`

---
 polkadot/node/tracking-allocator/src/lib.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 2f719a09cf50..37c7921c6f1a 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -17,8 +17,10 @@
 //! Tracking global allocator. Calculates the peak allocation between two checkpoints.
 
 use core::alloc::{GlobalAlloc, Layout};
-use std::sync::atomic::{AtomicBool, Ordering};
-use std::ptr::null_mut;
+use std::{
+	ptr::null_mut,
+	sync::atomic::{AtomicBool, Ordering},
+};
 
 struct TrackingAllocatorData {
 	lock: AtomicBool,

From dfa0daffbed5af4de01e4ec7d8249d49753d8327 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Thu, 21 Sep 2023 15:55:00 +0200
Subject: [PATCH 08/33] Implement abort handler

---
 Cargo.lock                                    | 34 +++++++++++++++
 polkadot/node/core/pvf/common/src/error.rs    | 14 ++++++-
 .../node/core/pvf/prepare-worker/Cargo.toml   |  2 +
 .../node/core/pvf/prepare-worker/src/lib.rs   | 41 ++++++++++++++++++-
 polkadot/node/core/pvf/src/metrics.rs         |  2 +-
 polkadot/node/core/pvf/src/prepare/pool.rs    | 14 +++++++
 .../node/core/pvf/src/prepare/worker_intf.rs  | 10 +++++
 polkadot/node/core/pvf/tests/it/main.rs       |  7 ++--
 8 files changed, 118 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e15ea3b4074d..cc703adfe487 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3024,6 +3024,30 @@ dependencies = [
  "wasmtime-types",
 ]
 
+[[package]]
+name = "crash-context"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b85cef661eeca0c6675116310936972c520ebb0a33ddef16fd7efc957f4c1288"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "mach2",
+]
+
+[[package]]
+name = "crash-handler"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c43ddb8d457c6c6c6d0ebeeadb3b3f4b246c39ed83b0d3393f91bf06a5b79b05"
+dependencies = [
+ "cfg-if",
+ "crash-context",
+ "libc",
+ "mach2",
+ "parking_lot 0.12.1",
+]
+
 [[package]]
 name = "crc"
 version = "3.0.1"
@@ -7603,6 +7627,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "mach2"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d0d1830bcd151a6fc4aea1369af235b36c1528fe976b8ff678683c9995eade8"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "macro_magic"
 version = "0.4.2"
@@ -12089,6 +12122,7 @@ dependencies = [
 name = "polkadot-node-core-pvf-prepare-worker"
 version = "1.0.0"
 dependencies = [
+ "crash-handler",
  "criterion 0.4.0",
  "futures",
  "libc",
diff --git a/polkadot/node/core/pvf/common/src/error.rs b/polkadot/node/core/pvf/common/src/error.rs
index 6eb0d9b7df42..2c355a7a5cb4 100644
--- a/polkadot/node/core/pvf/common/src/error.rs
+++ b/polkadot/node/core/pvf/common/src/error.rs
@@ -26,25 +26,36 @@ pub type PrepareResult = Result<PrepareStats, PrepareError>;
 #[derive(Debug, Clone, Encode, Decode)]
 pub enum PrepareError {
 	/// During the prevalidation stage of preparation an issue was found with the PVF.
+	#[codec(index = 0)]
 	Prevalidation(String),
 	/// Compilation failed for the given PVF.
+	#[codec(index = 1)]
 	Preparation(String),
 	/// Instantiation of the WASM module instance failed.
+	#[codec(index = 2)]
 	RuntimeConstruction(String),
 	/// An unexpected panic has occurred in the preparation worker.
+	#[codec(index = 3)]
 	Panic(String),
 	/// Failed to prepare the PVF due to the time limit.
+	#[codec(index = 4)]
 	TimedOut,
 	/// An IO error occurred. This state is reported by either the validation host or by the
 	/// worker.
+	#[codec(index = 5)]
 	IoErr(String),
 	/// The temporary file for the artifact could not be created at the given cache path. This
 	/// state is reported by the validation host (not by the worker).
+	#[codec(index = 6)]
 	CreateTmpFileErr(String),
 	/// The response from the worker is received, but the file cannot be renamed (moved) to the
 	/// final destination location. This state is reported by the validation host (not by the
 	/// worker).
+	#[codec(index = 7)]
 	RenameTmpFileErr(String),
+	/// Memory limit reached
+	#[codec(index = 8)]
+	OutOfMemory,
 }
 
 impl PrepareError {
@@ -57,7 +68,7 @@ impl PrepareError {
 	pub fn is_deterministic(&self) -> bool {
 		use PrepareError::*;
 		match self {
-			Prevalidation(_) | Preparation(_) | Panic(_) => true,
+			Prevalidation(_) | Preparation(_) | Panic(_) | OutOfMemory => true,
 			TimedOut | IoErr(_) | CreateTmpFileErr(_) | RenameTmpFileErr(_) => false,
 			// Can occur due to issues with the PVF, but also due to local errors.
 			RuntimeConstruction(_) => false,
@@ -77,6 +88,7 @@ impl fmt::Display for PrepareError {
 			IoErr(err) => write!(f, "prepare: io error while receiving response: {}", err),
 			CreateTmpFileErr(err) => write!(f, "prepare: error creating tmp file: {}", err),
 			RenameTmpFileErr(err) => write!(f, "prepare: error renaming tmp file: {}", err),
+			OutOfMemory => write!(f, "prepare: out of memory"),
 		}
 	}
 }
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index 52ecb60d3eb1..284b31a0c4c4 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
@@ -15,6 +15,7 @@ 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 }
+crash-handler = { version = "0.6.0", optional = true }
 
 parity-scale-codec = { version = "3.6.1", default-features = false, features = ["derive"] }
 
@@ -39,6 +40,7 @@ jemalloc-allocator = [
 	"polkadot-node-core-pvf-common/jemalloc-allocator",
 ]
 tracking-allocator = [
+	"dep:crash-handler",
 	"dep:tikv-jemallocator",
 	"dep:tracking-allocator",
 	"polkadot-node-core-pvf-common/tracking-allocator",
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index bbdb7445c2da..0a8a780e5e90 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -60,6 +60,10 @@ use tracking_allocator::TrackingAllocator;
 #[cfg(feature = "tracking-allocator")]
 #[global_allocator]
 static ALLOC: TrackingAllocator<Jemalloc> = TrackingAllocator(Jemalloc);
+#[cfg(feature = "tracking-allocator")]
+use crash_handler::{CrashEventResult, CrashHandler};
+#[cfg(feature = "tracking-allocator")]
+use std::os::fd::AsRawFd;
 
 /// Contains the bytes for a successfully compiled artifact.
 pub struct CompiledArtifact(Vec<u8>);
@@ -177,6 +181,37 @@ pub fn worker_entrypoint(
 					WaitOutcome::TimedOut,
 				)?;
 
+				#[cfg(feature = "tracking-allocator")]
+				let fd = stream.as_raw_fd();
+				#[cfg(feature = "tracking-allocator")]
+				let crash_event = unsafe {
+					crash_handler::make_crash_event(move |context| {
+						ALLOC.end_tracking();
+						if context.siginfo.ssi_signo as i32 == libc::SIGABRT {
+							// We're inside `SIGABRT` handler. The process state is compromised so
+							// no sudden movements please. Avoid allocations and make no assumptions
+							// based on the state of the stack.
+
+							// Send pre-encoded `Err(PrepareError::OutOfMemory)` to the host, close
+							// the connection, bail out.
+							libc::write(
+								fd,
+								b"\x02\x00\x00\x00\x00\x00\x00\x00\x01\x08" as *const _
+									as *const libc::c_void,
+								10,
+							);
+							libc::close(fd);
+							std::process::exit(1);
+						}
+						CrashEventResult::Handled(true)
+					})
+				};
+
+				#[cfg(feature = "tracking-allocator")]
+				let _crash_handler = CrashHandler::attach(crash_event).map_err(|_| {
+					io::Error::new(io::ErrorKind::Other, "Failed to install crash handler")
+				})?;
+
 				#[cfg(feature = "tracking-allocator")]
 				ALLOC.start_tracking(executor_params.prechecking_max_memory().map(|v| v as isize));
 
@@ -265,7 +300,11 @@ pub fn worker_entrypoint(
 									#[cfg(target_os = "linux")]
 									max_rss: extract_max_rss_stat(max_rss, worker_pid),
 									#[cfg(feature = "tracking-allocator")]
-									peak_alloc: peak_alloc as u64,
+									peak_alloc: if peak_alloc > 0 {
+										peak_alloc as u64
+									} else {
+										0u64
+									},
 								};
 
 								// Log if landlock threw an error.
diff --git a/polkadot/node/core/pvf/src/metrics.rs b/polkadot/node/core/pvf/src/metrics.rs
index fdf3dacbe9ff..264aa16232d5 100644
--- a/polkadot/node/core/pvf/src/metrics.rs
+++ b/polkadot/node/core/pvf/src/metrics.rs
@@ -282,7 +282,7 @@ impl metrics::Metrics for Metrics {
 			preparation_peak_allocation: prometheus::register(
 				prometheus::Histogram::with_opts(
 					prometheus::HistogramOpts::new(
-						"polkadot_pvf_preparation_peak_allocattion",
+						"polkadot_pvf_preparation_peak_allocation",
 						"peak allocation observed for preparation (in kilobytes)",
 					).buckets(
 						prometheus::exponential_buckets(8192.0, 2.0, 10)
diff --git a/polkadot/node/core/pvf/src/prepare/pool.rs b/polkadot/node/core/pvf/src/prepare/pool.rs
index 92aa4896c263..a6a0d05f075b 100644
--- a/polkadot/node/core/pvf/src/prepare/pool.rs
+++ b/polkadot/node/core/pvf/src/prepare/pool.rs
@@ -363,6 +363,20 @@ fn handle_mux(
 						)?;
 					}
 
+					Ok(())
+				},
+				Outcome::OutOfMemory => {
+					if attempt_retire(metrics, spawned, worker) {
+						reply(
+							from_pool,
+							FromPool::Concluded {
+								worker,
+								rip: true,
+								result: Err(PrepareError::OutOfMemory),
+							},
+						)?;
+					}
+
 					Ok(())
 				},
 			}
diff --git a/polkadot/node/core/pvf/src/prepare/worker_intf.rs b/polkadot/node/core/pvf/src/prepare/worker_intf.rs
index 5280ab6b42a2..6b42d77d09b5 100644
--- a/polkadot/node/core/pvf/src/prepare/worker_intf.rs
+++ b/polkadot/node/core/pvf/src/prepare/worker_intf.rs
@@ -73,6 +73,8 @@ pub enum Outcome {
 	///
 	/// This doesn't return an idle worker instance, thus this worker is no longer usable.
 	IoErr(String),
+	/// The worker ran out of memory and is aborting. The worker should be ripped.
+	OutOfMemory,
 }
 
 /// Given the idle token of a worker and parameters of work, communicates with the worker and
@@ -122,6 +124,14 @@ pub async fn start_work(
 
 		match result {
 			// Received bytes from worker within the time limit.
+			Ok(Ok(Err(PrepareError::OutOfMemory))) => {
+				gum::warn!(
+					target: LOG_TARGET,
+					worker_pid = %pid,
+					"worker is out of memory",
+				);
+				Outcome::OutOfMemory
+			},
 			Ok(Ok(prepare_result)) =>
 				handle_response(
 					metrics,
diff --git a/polkadot/node/core/pvf/tests/it/main.rs b/polkadot/node/core/pvf/tests/it/main.rs
index 13d71c9b751f..a5ecab998be5 100644
--- a/polkadot/node/core/pvf/tests/it/main.rs
+++ b/polkadot/node/core/pvf/tests/it/main.rs
@@ -363,8 +363,7 @@ async fn prechecking_within_memory_limits() {
 #[cfg(feature = "tracking-allocator")]
 #[tokio::test]
 async fn prechecking_out_of_memory() {
-	use polkadot_node_core_pvf::ValidationError::InternalError;
-	use polkadot_node_core_pvf_common::error::InternalValidationError::NonDeterministicPrepareError;
+	use polkadot_node_core_pvf::{InvalidCandidate, ValidationError};
 
 	let host = TestHost::new();
 	let parent_head = HeadData { number: 0, parent_hash: [0; 32], post_state: hash_state(0) };
@@ -383,7 +382,9 @@ async fn prechecking_out_of_memory() {
 		.await;
 
 	match result {
-		Err(InternalError(NonDeterministicPrepareError(_))) => (),
+		Err(ValidationError::InvalidCandidate(InvalidCandidate::PrepareError(err)))
+			if err == "prepare: out of memory" =>
+			(),
 		r => panic!("{:?}", r),
 	}
 }

From 48cde3ec25bd1ba10e2416af9acfa51cbafb8481 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Thu, 21 Sep 2023 17:54:33 +0200
Subject: [PATCH 09/33] Address some discussions

---
 .../benches/prepare_kusama_runtime.rs         |  6 +-
 polkadot/node/core/pvf/tests/it/main.rs       | 64 ++++++++++---------
 polkadot/src/main.rs                          |  3 +-
 3 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs b/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
index 8df1df17d2f0..ac4213c0476c 100644
--- a/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
+++ b/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
@@ -20,7 +20,7 @@ use polkadot_node_core_pvf_prepare_worker::{prepare, prevalidate};
 use polkadot_primitives::ExecutorParams;
 use std::time::Duration;
 
-fn do_prepare_kusama_runtime(pvf: PvfPrepData) {
+fn do_prepare_runtime(pvf: PvfPrepData) {
 	let blob = match prevalidate(&pvf.code()) {
 		Err(err) => panic!("{:?}", err),
 		Ok(b) => b,
@@ -51,7 +51,9 @@ fn prepare_kusama_runtime(c: &mut Criterion) {
 	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()))
+		// `PvfPrepData` is designed to be cheap to clone, so cloning shouldn't affect the
+		// benchmark accuracy
+		b.iter(|| do_prepare_runtime(pvf.clone()))
 	});
 	group.finish();
 }
diff --git a/polkadot/node/core/pvf/tests/it/main.rs b/polkadot/node/core/pvf/tests/it/main.rs
index a5ecab998be5..241065132c3b 100644
--- a/polkadot/node/core/pvf/tests/it/main.rs
+++ b/polkadot/node/core/pvf/tests/it/main.rs
@@ -18,8 +18,8 @@
 use assert_matches::assert_matches;
 use parity_scale_codec::Encode as _;
 use polkadot_node_core_pvf::{
-	start, Config, InvalidCandidate, Metrics, PrepareJobKind, PvfPrepData, ValidationError,
-	ValidationHost, JOB_TIMEOUT_WALL_CLOCK_FACTOR,
+	start, Config, InvalidCandidate, Metrics, PrepareError, PrepareJobKind, PrepareStats,
+	PvfPrepData, ValidationError, ValidationHost, JOB_TIMEOUT_WALL_CLOCK_FACTOR,
 };
 use polkadot_parachain_primitives::primitives::{
 	BlockData as GenericBlockData, ValidationParams, ValidationResult,
@@ -29,12 +29,6 @@ use polkadot_primitives::ExecutorParams;
 #[cfg(any(feature = "ci-only-tests", feature = "tracking-allocator"))]
 use polkadot_primitives::ExecutorParam;
 
-#[cfg(feature = "tracking-allocator")]
-use ::adder::{hash_state, BlockData, HeadData};
-
-#[cfg(feature = "tracking-allocator")]
-use polkadot_primitives::HeadData as GenericHeadData;
-
 use std::time::Duration;
 use tokio::sync::Mutex;
 
@@ -78,6 +72,33 @@ impl TestHost {
 		Self { cache_dir, host: Mutex::new(host) }
 	}
 
+	async fn precheck_pvf(
+		&self,
+		code: &[u8],
+		executor_params: ExecutorParams,
+	) -> Result<PrepareStats, PrepareError> {
+		let (result_tx, result_rx) = futures::channel::oneshot::channel();
+
+		let code = sp_maybe_compressed_blob::decompress(code, 16 * 1024 * 1024)
+			.expect("Compression works");
+
+		self.host
+			.lock()
+			.await
+			.precheck_pvf(
+				PvfPrepData::from_code(
+					code.into(),
+					executor_params,
+					TEST_PREPARATION_TIMEOUT,
+					PrepareJobKind::Prechecking,
+				),
+				result_tx,
+			)
+			.await
+			.unwrap();
+		result_rx.await.unwrap()
+	}
+
 	async fn validate_candidate(
 		&self,
 		code: &[u8],
@@ -335,17 +356,9 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 #[tokio::test]
 async fn prechecking_within_memory_limits() {
 	let host = TestHost::new();
-	let parent_head = HeadData { number: 0, parent_hash: [0; 32], post_state: hash_state(0) };
-	let block_data = BlockData { state: 0, add: 512 };
 	let result = host
-		.validate_candidate(
+		.precheck_pvf(
 			::adder::wasm_binary_unwrap(),
-			ValidationParams {
-				parent_head: GenericHeadData(parent_head.encode()),
-				block_data: GenericBlockData(block_data.encode()),
-				relay_parent_number: 1,
-				relay_parent_storage_root: Default::default(),
-			},
 			ExecutorParams::from(&[ExecutorParam::PrecheckingMaxMemory(10 * 1024 * 1024)][..]),
 		)
 		.await;
@@ -363,28 +376,17 @@ async fn prechecking_within_memory_limits() {
 #[cfg(feature = "tracking-allocator")]
 #[tokio::test]
 async fn prechecking_out_of_memory() {
-	use polkadot_node_core_pvf::{InvalidCandidate, ValidationError};
+	use polkadot_node_core_pvf::PrepareError;
 
 	let host = TestHost::new();
-	let parent_head = HeadData { number: 0, parent_hash: [0; 32], post_state: hash_state(0) };
-	let block_data = BlockData { state: 0, add: 512 };
 	let result = host
-		.validate_candidate(
+		.precheck_pvf(
 			::adder::wasm_binary_unwrap(),
-			ValidationParams {
-				parent_head: GenericHeadData(parent_head.encode()),
-				block_data: GenericBlockData(block_data.encode()),
-				relay_parent_number: 1,
-				relay_parent_storage_root: Default::default(),
-			},
 			ExecutorParams::from(&[ExecutorParam::PrecheckingMaxMemory(512 * 1024)][..]),
 		)
 		.await;
-
 	match result {
-		Err(ValidationError::InvalidCandidate(InvalidCandidate::PrepareError(err)))
-			if err == "prepare: out of memory" =>
-			(),
+		Err(PrepareError::OutOfMemory) => (),
 		r => panic!("{:?}", r),
 	}
 }
diff --git a/polkadot/src/main.rs b/polkadot/src/main.rs
index 8b125d3d8409..fd133adfce51 100644
--- a/polkadot/src/main.rs
+++ b/polkadot/src/main.rs
@@ -21,9 +21,8 @@
 use color_eyre::eyre;
 
 #[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
-#[doc(hidden)]
 #[global_allocator]
-pub static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
 fn main() -> eyre::Result<()> {
 	color_eyre::install()?;

From acf149e403ff395c308313495179105857cad5cd Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Thu, 21 Sep 2023 19:44:57 +0200
Subject: [PATCH 10/33] Remove feature gate

---
 polkadot/Cargo.toml                           |  8 +-----
 polkadot/node/core/pvf/Cargo.toml             |  4 ---
 polkadot/node/core/pvf/common/Cargo.toml      |  1 -
 polkadot/node/core/pvf/common/src/prepare.rs  |  1 -
 .../node/core/pvf/prepare-worker/Cargo.toml   | 14 ++++------
 .../node/core/pvf/prepare-worker/src/lib.rs   | 26 +++++++------------
 polkadot/node/core/pvf/src/metrics.rs         |  3 ---
 polkadot/node/core/pvf/tests/it/main.rs       |  7 +----
 8 files changed, 17 insertions(+), 47 deletions(-)

diff --git a/polkadot/Cargo.toml b/polkadot/Cargo.toml
index dc59dbd1d83d..0e6530914926 100644
--- a/polkadot/Cargo.toml
+++ b/polkadot/Cargo.toml
@@ -30,7 +30,7 @@ polkadot-cli = { path = "cli", features = [ "westend-native", "rococo-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 }
+tracking-allocator = { path = "node/tracking-allocator" }
 
 # Needed for worker binaries.
 polkadot-node-core-pvf-common = { path = "node/core/pvf/common" }
@@ -65,12 +65,6 @@ 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
diff --git a/polkadot/node/core/pvf/Cargo.toml b/polkadot/node/core/pvf/Cargo.toml
index a5c1fa7389c6..478d1952d9d9 100644
--- a/polkadot/node/core/pvf/Cargo.toml
+++ b/polkadot/node/core/pvf/Cargo.toml
@@ -49,10 +49,6 @@ 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",
diff --git a/polkadot/node/core/pvf/common/Cargo.toml b/polkadot/node/core/pvf/common/Cargo.toml
index 53f891678e63..621f7e24f72b 100644
--- a/polkadot/node/core/pvf/common/Cargo.toml
+++ b/polkadot/node/core/pvf/common/Cargo.toml
@@ -39,4 +39,3 @@ tempfile = "3.3.0"
 # Also used for building the puppet worker.
 test-utils = []
 jemalloc-allocator = []
-tracking-allocator = []
diff --git a/polkadot/node/core/pvf/common/src/prepare.rs b/polkadot/node/core/pvf/common/src/prepare.rs
index 704f8d06aaae..6a826acca94d 100644
--- a/polkadot/node/core/pvf/common/src/prepare.rs
+++ b/polkadot/node/core/pvf/common/src/prepare.rs
@@ -36,7 +36,6 @@ pub struct MemoryStats {
 	#[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,
 }
 
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index 284b31a0c4c4..b031b3a78b59 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
@@ -11,11 +11,11 @@ futures = "0.3.21"
 gum = { package = "tracing-gum", path = "../../../gum" }
 libc = "0.2.139"
 rayon = "1.5.1"
-tikv-jemalloc-ctl = { version = "0.5.0", optional = true }
+crash-handler = { version = "0.6.0" }
 tokio = { version = "1.24.2", features = ["fs", "process"] }
-tracking-allocator = { path = "../../../tracking-allocator", optional = true }
+tracking-allocator = { path = "../../../tracking-allocator" }
+tikv-jemalloc-ctl = { version = "0.5.0", optional = true }
 tikv-jemallocator = { version = "0.5.0", optional = true }
-crash-handler = { version = "0.6.0", optional = true }
 
 parity-scale-codec = { version = "3.6.1", default-features = false, features = ["derive"] }
 
@@ -31,19 +31,15 @@ sp-maybe-compressed-blob = { path = "../../../../../substrate/primitives/maybe-c
 sp-tracing = { path = "../../../../../substrate/primitives/tracing" }
 
 [target.'cfg(target_os = "linux")'.dependencies]
+tikv-jemallocator = "0.5.0"
 tikv-jemalloc-ctl = "0.5.0"
 
 [features]
 builder = []
 jemalloc-allocator = [
 	"dep:tikv-jemalloc-ctl",
-	"polkadot-node-core-pvf-common/jemalloc-allocator",
-]
-tracking-allocator = [
-	"dep:crash-handler",
 	"dep:tikv-jemallocator",
-	"dep:tracking-allocator",
-	"polkadot-node-core-pvf-common/tracking-allocator",
+	"polkadot-node-core-pvf-common/jemalloc-allocator",
 ]
 
 [dev-dependencies]
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index 0a8a780e5e90..31e9fed85b6a 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -29,6 +29,7 @@ const LOG_TARGET: &str = "parachain::pvf-prepare-worker";
 use crate::memory_stats::max_rss_stat::{extract_max_rss_stat, get_max_rss_thread};
 #[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 use crate::memory_stats::memory_tracker::{get_memory_tracker_loop_stats, memory_tracker_loop};
+use crash_handler::{CrashEventResult, CrashHandler};
 use parity_scale_codec::{Decode, Encode};
 use polkadot_node_core_pvf_common::{
 	error::{PrepareError, PrepareResult},
@@ -47,23 +48,22 @@ use polkadot_node_core_pvf_common::{
 };
 use polkadot_primitives::ExecutorParams;
 use std::{
+	os::fd::AsRawFd,
 	path::PathBuf,
 	sync::{mpsc::channel, Arc},
 	time::Duration,
 };
 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")]
+
+#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 #[global_allocator]
-static ALLOC: TrackingAllocator<Jemalloc> = TrackingAllocator(Jemalloc);
-#[cfg(feature = "tracking-allocator")]
-use crash_handler::{CrashEventResult, CrashHandler};
-#[cfg(feature = "tracking-allocator")]
-use std::os::fd::AsRawFd;
+static ALLOC: TrackingAllocator<tikv_jemallocator::Jemalloc> =
+	TrackingAllocator(tikv_jemallocator::Jemalloc);
+
+#[cfg(not(any(target_os = "linux", feature = "jemalloc-allocator")))]
+#[global_allocator]
+static ALLOC: TrackingAllocator<std::alloc::System> = TrackingAllocator(std::alloc::System);
 
 /// Contains the bytes for a successfully compiled artifact.
 pub struct CompiledArtifact(Vec<u8>);
@@ -181,9 +181,7 @@ pub fn worker_entrypoint(
 					WaitOutcome::TimedOut,
 				)?;
 
-				#[cfg(feature = "tracking-allocator")]
 				let fd = stream.as_raw_fd();
-				#[cfg(feature = "tracking-allocator")]
 				let crash_event = unsafe {
 					crash_handler::make_crash_event(move |context| {
 						ALLOC.end_tracking();
@@ -207,12 +205,10 @@ pub fn worker_entrypoint(
 					})
 				};
 
-				#[cfg(feature = "tracking-allocator")]
 				let _crash_handler = CrashHandler::attach(crash_event).map_err(|_| {
 					io::Error::new(io::ErrorKind::Other, "Failed to install crash handler")
 				})?;
 
-				#[cfg(feature = "tracking-allocator")]
 				ALLOC.start_tracking(executor_params.prechecking_max_memory().map(|v| v as isize));
 
 				// Spawn another thread for preparation.
@@ -255,7 +251,6 @@ 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!(
@@ -299,7 +294,6 @@ 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: if peak_alloc > 0 {
 										peak_alloc as u64
 									} else {
diff --git a/polkadot/node/core/pvf/src/metrics.rs b/polkadot/node/core/pvf/src/metrics.rs
index 264aa16232d5..fa062c1518bb 100644
--- a/polkadot/node/core/pvf/src/metrics.rs
+++ b/polkadot/node/core/pvf/src/metrics.rs
@@ -94,7 +94,6 @@ impl Metrics {
 				metrics.preparation_max_allocated.observe(max_allocated_kb);
 			}
 
-			#[cfg(feature = "tracking-allocator")]
 			metrics
 				.preparation_peak_allocation
 				.observe((memory_stats.peak_alloc / 1024) as f64);
@@ -119,7 +118,6 @@ 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,
 }
 
@@ -278,7 +276,6 @@ impl metrics::Metrics for Metrics {
 				)?,
 				registry,
 			)?,
-			#[cfg(feature = "tracking-allocator")]
 			preparation_peak_allocation: prometheus::register(
 				prometheus::Histogram::with_opts(
 					prometheus::HistogramOpts::new(
diff --git a/polkadot/node/core/pvf/tests/it/main.rs b/polkadot/node/core/pvf/tests/it/main.rs
index 241065132c3b..e76db4229bdd 100644
--- a/polkadot/node/core/pvf/tests/it/main.rs
+++ b/polkadot/node/core/pvf/tests/it/main.rs
@@ -24,10 +24,7 @@ use polkadot_node_core_pvf::{
 use polkadot_parachain_primitives::primitives::{
 	BlockData as GenericBlockData, ValidationParams, ValidationResult,
 };
-use polkadot_primitives::ExecutorParams;
-
-#[cfg(any(feature = "ci-only-tests", feature = "tracking-allocator"))]
-use polkadot_primitives::ExecutorParam;
+use polkadot_primitives::{ExecutorParam, ExecutorParams};
 
 use std::time::Duration;
 use tokio::sync::Mutex;
@@ -352,7 +349,6 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 // starts failing, either Wasmtime version has changed, or the PVF code itself has changed, and
 // more memory is required now. Multi-threaded preparation, if ever enabled, may also affect
 // memory consumption.
-#[cfg(feature = "tracking-allocator")]
 #[tokio::test]
 async fn prechecking_within_memory_limits() {
 	let host = TestHost::new();
@@ -373,7 +369,6 @@ async fn prechecking_within_memory_limits() {
 // limit enforced. At the moment of writing, the limit if not enough to prepare the PVF, and the
 // preparation is supposed to generate an error. If the test starts failing, either Wasmtime
 // version has changed, or the PVF code itself has changed, and less memory is required now.
-#[cfg(feature = "tracking-allocator")]
 #[tokio::test]
 async fn prechecking_out_of_memory() {
 	use polkadot_node_core_pvf::PrepareError;

From c7fa464593a66d30a1776070fea10d875f49312f Mon Sep 17 00:00:00 2001
From: s0me0ne-unkn0wn <48632512+s0me0ne-unkn0wn@users.noreply.github.com>
Date: Fri, 22 Sep 2023 15:51:31 +0200
Subject: [PATCH 11/33] Update crate description

Co-authored-by: Marcin S. <marcin@realemail.net>
---
 polkadot/node/tracking-allocator/Cargo.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/polkadot/node/tracking-allocator/Cargo.toml b/polkadot/node/tracking-allocator/Cargo.toml
index 293b29d80e69..2ea3f69c7d0d 100644
--- a/polkadot/node/tracking-allocator/Cargo.toml
+++ b/polkadot/node/tracking-allocator/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "tracking-allocator"
-description = "Tracking allocator to control amount of memory consumed by PVF preparation process"
+description = "Tracking allocator to control the amount of memory consumed by the process"
 version = "1.0.0"
 authors.workspace = true
 edition.workspace = true

From edeb9010c15db5a42a4793af5412e77d31cc62d2 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Fri, 22 Sep 2023 20:34:46 +0200
Subject: [PATCH 12/33] Implement failure callback

---
 Cargo.lock                                    | 34 --------------
 .../node/core/pvf/prepare-worker/Cargo.toml   |  1 -
 .../node/core/pvf/prepare-worker/src/lib.rs   | 44 +++++++------------
 polkadot/node/tracking-allocator/src/lib.rs   | 34 +++++++++++---
 4 files changed, 43 insertions(+), 70 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f2bf38fd8cce..3be8aaa9bf5a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3026,30 +3026,6 @@ dependencies = [
  "wasmtime-types",
 ]
 
-[[package]]
-name = "crash-context"
-version = "0.6.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b85cef661eeca0c6675116310936972c520ebb0a33ddef16fd7efc957f4c1288"
-dependencies = [
- "cfg-if",
- "libc",
- "mach2",
-]
-
-[[package]]
-name = "crash-handler"
-version = "0.6.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c43ddb8d457c6c6c6d0ebeeadb3b3f4b246c39ed83b0d3393f91bf06a5b79b05"
-dependencies = [
- "cfg-if",
- "crash-context",
- "libc",
- "mach2",
- "parking_lot 0.12.1",
-]
-
 [[package]]
 name = "crc"
 version = "3.0.1"
@@ -7634,15 +7610,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "mach2"
-version = "0.4.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6d0d1830bcd151a6fc4aea1369af235b36c1528fe976b8ff678683c9995eade8"
-dependencies = [
- "libc",
-]
-
 [[package]]
 name = "macro_magic"
 version = "0.4.2"
@@ -12146,7 +12113,6 @@ dependencies = [
 name = "polkadot-node-core-pvf-prepare-worker"
 version = "1.0.0"
 dependencies = [
- "crash-handler",
  "criterion 0.4.0",
  "futures",
  "libc",
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index b031b3a78b59..edc91769ae2d 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
@@ -11,7 +11,6 @@ futures = "0.3.21"
 gum = { package = "tracing-gum", path = "../../../gum" }
 libc = "0.2.139"
 rayon = "1.5.1"
-crash-handler = { version = "0.6.0" }
 tokio = { version = "1.24.2", features = ["fs", "process"] }
 tracking-allocator = { path = "../../../tracking-allocator" }
 tikv-jemalloc-ctl = { version = "0.5.0", optional = true }
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index 31e9fed85b6a..a79c35228c9d 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -29,7 +29,6 @@ const LOG_TARGET: &str = "parachain::pvf-prepare-worker";
 use crate::memory_stats::max_rss_stat::{extract_max_rss_stat, get_max_rss_thread};
 #[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 use crate::memory_stats::memory_tracker::{get_memory_tracker_loop_stats, memory_tracker_loop};
-use crash_handler::{CrashEventResult, CrashHandler};
 use parity_scale_codec::{Decode, Encode};
 use polkadot_node_core_pvf_common::{
 	error::{PrepareError, PrepareResult},
@@ -182,34 +181,21 @@ pub fn worker_entrypoint(
 				)?;
 
 				let fd = stream.as_raw_fd();
-				let crash_event = unsafe {
-					crash_handler::make_crash_event(move |context| {
-						ALLOC.end_tracking();
-						if context.siginfo.ssi_signo as i32 == libc::SIGABRT {
-							// We're inside `SIGABRT` handler. The process state is compromised so
-							// no sudden movements please. Avoid allocations and make no assumptions
-							// based on the state of the stack.
-
-							// Send pre-encoded `Err(PrepareError::OutOfMemory)` to the host, close
-							// the connection, bail out.
-							libc::write(
-								fd,
-								b"\x02\x00\x00\x00\x00\x00\x00\x00\x01\x08" as *const _
-									as *const libc::c_void,
-								10,
-							);
-							libc::close(fd);
-							std::process::exit(1);
-						}
-						CrashEventResult::Handled(true)
-					})
-				};
-
-				let _crash_handler = CrashHandler::attach(crash_event).map_err(|_| {
-					io::Error::new(io::ErrorKind::Other, "Failed to install crash handler")
-				})?;
-
-				ALLOC.start_tracking(executor_params.prechecking_max_memory().map(|v| v as isize));
+				ALLOC.start_tracking(
+					executor_params.prechecking_max_memory().map(|v| v as isize),
+					Some(Box::new(move || unsafe {
+						// Inside the failure handler, the allocator is locked and no allocations
+						// are possible
+						libc::write(
+							fd,
+							b"\x02\x00\x00\x00\x00\x00\x00\x00\x01\x08" as *const _
+								as *const libc::c_void,
+							10,
+						);
+						libc::close(fd);
+						std::process::exit(1);
+					})),
+				);
 
 				// Spawn another thread for preparation.
 				let prepare_thread = thread::spawn_worker_thread(
diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 37c7921c6f1a..4debea683d9e 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -27,6 +27,7 @@ struct TrackingAllocatorData {
 	current: isize,
 	peak: isize,
 	limit: isize,
+	failure_handler: Option<Box<dyn Fn()>>,
 }
 
 impl TrackingAllocatorData {
@@ -57,11 +58,12 @@ impl TrackingAllocatorData {
 		self.lock.store(false, Ordering::Release);
 	}
 
-	fn start_tracking(&mut self, limit: isize) {
+	fn start_tracking(&mut self, limit: isize, failure_handler: Option<Box<dyn Fn()>>) {
 		self.lock();
 		self.current = 0;
 		self.peak = 0;
 		self.limit = limit;
+		self.failure_handler = failure_handler;
 		self.unlock();
 	}
 
@@ -69,6 +71,7 @@ impl TrackingAllocatorData {
 		self.lock();
 		let peak = self.peak;
 		self.limit = 0;
+		self.failure_handler = None;
 		self.unlock();
 		peak
 	}
@@ -81,13 +84,20 @@ impl TrackingAllocatorData {
 			self.peak = self.current;
 		}
 		let within_limits = self.limit == 0 || self.peak <= self.limit;
-		self.unlock();
+		if within_limits {
+			self.unlock()
+		}
 		within_limits
 	}
 }
 
-static mut ALLOCATOR_DATA: TrackingAllocatorData =
-	TrackingAllocatorData { lock: AtomicBool::new(false), current: 0, peak: 0, limit: 0 };
+static mut ALLOCATOR_DATA: TrackingAllocatorData = TrackingAllocatorData {
+	lock: AtomicBool::new(false),
+	current: 0,
+	peak: 0,
+	limit: 0,
+	failure_handler: None,
+};
 
 pub struct TrackingAllocator<A: GlobalAlloc>(pub A);
 
@@ -97,9 +107,9 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 	//   is isolated by an exclusive lock.
 
 	/// Start tracking
-	pub fn start_tracking(&self, limit: Option<isize>) {
+	pub fn start_tracking(&self, limit: Option<isize>, failure_handler: Option<Box<dyn Fn()>>) {
 		unsafe {
-			ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0));
+			ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0), failure_handler);
 		}
 	}
 
@@ -119,6 +129,10 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 		if ALLOCATOR_DATA.track(layout.size() as isize) {
 			self.0.alloc(layout)
 		} else {
+			if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
+				failure_handler()
+			}
+			ALLOCATOR_DATA.unlock();
 			null_mut()
 		}
 	}
@@ -128,6 +142,10 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 		if ALLOCATOR_DATA.track(layout.size() as isize) {
 			self.0.alloc_zeroed(layout)
 		} else {
+			if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
+				failure_handler()
+			}
+			ALLOCATOR_DATA.unlock();
 			null_mut()
 		}
 	}
@@ -143,6 +161,10 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 		if ALLOCATOR_DATA.track((new_size as isize) - (layout.size() as isize)) {
 			self.0.realloc(ptr, layout, new_size)
 		} else {
+			if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
+				failure_handler()
+			}
+			ALLOCATOR_DATA.unlock();
 			null_mut()
 		}
 	}

From 41958067daa369c3594b59c90b932d643464e8af Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Fri, 22 Sep 2023 22:21:03 +0200
Subject: [PATCH 13/33] Address discussions

---
 Cargo.lock                                    |  1 -
 polkadot/Cargo.toml                           |  1 -
 polkadot/node/core/pvf/common/src/error.rs    | 12 +++++++++
 polkadot/node/core/pvf/common/src/prepare.rs  |  4 +--
 .../node/core/pvf/prepare-worker/src/lib.rs   | 12 +++++----
 polkadot/node/core/pvf/src/metrics.rs         | 11 +++++---
 .../node/core/pvf/src/prepare/worker_intf.rs  |  9 +------
 polkadot/node/core/pvf/tests/it/adder.rs      |  2 ++
 polkadot/node/core/pvf/tests/it/main.rs       | 26 +++++++------------
 polkadot/src/main.rs                          |  2 ++
 10 files changed, 43 insertions(+), 37 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3be8aaa9bf5a..4280b3c3d14c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -11435,7 +11435,6 @@ dependencies = [
  "tempfile",
  "tikv-jemallocator",
  "tokio",
- "tracking-allocator",
 ]
 
 [[package]]
diff --git a/polkadot/Cargo.toml b/polkadot/Cargo.toml
index 0e6530914926..aacc6ad405cc 100644
--- a/polkadot/Cargo.toml
+++ b/polkadot/Cargo.toml
@@ -30,7 +30,6 @@ polkadot-cli = { path = "cli", features = [ "westend-native", "rococo-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" }
 
 # Needed for worker binaries.
 polkadot-node-core-pvf-common = { path = "node/core/pvf/common" }
diff --git a/polkadot/node/core/pvf/common/src/error.rs b/polkadot/node/core/pvf/common/src/error.rs
index 2c355a7a5cb4..c82bc93852fa 100644
--- a/polkadot/node/core/pvf/common/src/error.rs
+++ b/polkadot/node/core/pvf/common/src/error.rs
@@ -23,6 +23,7 @@ use std::fmt;
 pub type PrepareResult = Result<PrepareStats, PrepareError>;
 
 /// An error that occurred during the prepare part of the PVF pipeline.
+// Codec indexes are intended to stabilize pre-encoded payloads (see `OOM_PAYLOAD` below)
 #[derive(Debug, Clone, Encode, Decode)]
 pub enum PrepareError {
 	/// During the prevalidation stage of preparation an issue was found with the PVF.
@@ -58,6 +59,9 @@ pub enum PrepareError {
 	OutOfMemory,
 }
 
+/// Pre-encoded length-prefixed `PrepareResult::Err(PrepareError::OutOfMemory)`
+pub const OOM_PAYLOAD: &[u8] = b"\x02\x00\x00\x00\x00\x00\x00\x00\x01\x08";
+
 impl PrepareError {
 	/// Returns whether this is a deterministic error, i.e. one that should trigger reliably. Those
 	/// errors depend on the PVF itself and the sc-executor/wasmtime logic.
@@ -124,3 +128,11 @@ impl fmt::Display for InternalValidationError {
 		}
 	}
 }
+
+#[test]
+fn pre_encoded_payloads() {
+	let oom_enc = PrepareResult::Err(PrepareError::OutOfMemory).encode();
+	let mut oom_payload = oom_enc.len().to_le_bytes().to_vec();
+	oom_payload.extend(oom_enc);
+	assert_eq!(oom_payload, OOM_PAYLOAD);
+}
diff --git a/polkadot/node/core/pvf/common/src/prepare.rs b/polkadot/node/core/pvf/common/src/prepare.rs
index 6a826acca94d..cefb906b6ac9 100644
--- a/polkadot/node/core/pvf/common/src/prepare.rs
+++ b/polkadot/node/core/pvf/common/src/prepare.rs
@@ -29,14 +29,14 @@ pub struct PrepareStats {
 /// supported by the OS, `ru_maxrss`.
 #[derive(Clone, Debug, Default, Encode, Decode)]
 pub struct MemoryStats {
-	/// Memory stats from `tikv_jemalloc_ctl`.
+	/// Memory stats from `tikv_jemalloc_ctl`, polling-based and not vary precise.
 	#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 	pub memory_tracker_stats: Option<MemoryAllocationStats>,
 	/// `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
-	pub peak_alloc: u64,
+	pub peak_tracked_alloc: u64,
 }
 
 /// Statistics of collected memory metrics.
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index a79c35228c9d..e6c9ef2c09b0 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -31,7 +31,7 @@ use crate::memory_stats::max_rss_stat::{extract_max_rss_stat, get_max_rss_thread
 use crate::memory_stats::memory_tracker::{get_memory_tracker_loop_stats, memory_tracker_loop};
 use parity_scale_codec::{Decode, Encode};
 use polkadot_node_core_pvf_common::{
-	error::{PrepareError, PrepareResult},
+	error::{PrepareError, PrepareResult, OOM_PAYLOAD},
 	executor_intf::Executor,
 	framed_recv, framed_send,
 	prepare::{MemoryStats, PrepareJobKind, PrepareStats},
@@ -188,9 +188,8 @@ pub fn worker_entrypoint(
 						// are possible
 						libc::write(
 							fd,
-							b"\x02\x00\x00\x00\x00\x00\x00\x00\x01\x08" as *const _
-								as *const libc::c_void,
-							10,
+							OOM_PAYLOAD as *const _ as *const libc::c_void,
+							OOM_PAYLOAD.len(),
 						);
 						libc::close(fd);
 						std::process::exit(1);
@@ -280,7 +279,10 @@ pub fn worker_entrypoint(
 									memory_tracker_stats,
 									#[cfg(target_os = "linux")]
 									max_rss: extract_max_rss_stat(max_rss, worker_pid),
-									peak_alloc: if peak_alloc > 0 {
+									// Negative peak allocation values are legit; they are narrow
+									// corner cases and shouldn't affect overall statistics
+									// significantly
+									peak_tracked_alloc: if peak_alloc > 0 {
 										peak_alloc as u64
 									} else {
 										0u64
diff --git a/polkadot/node/core/pvf/src/metrics.rs b/polkadot/node/core/pvf/src/metrics.rs
index fa062c1518bb..cc41745ba277 100644
--- a/polkadot/node/core/pvf/src/metrics.rs
+++ b/polkadot/node/core/pvf/src/metrics.rs
@@ -95,8 +95,8 @@ impl Metrics {
 			}
 
 			metrics
-				.preparation_peak_allocation
-				.observe((memory_stats.peak_alloc / 1024) as f64);
+				.preparation_peak_tracked_allocation
+				.observe((memory_stats.peak_tracked_alloc / 1024) as f64);
 		}
 	}
 }
@@ -114,11 +114,14 @@ struct MetricsInner {
 	execution_time: prometheus::Histogram,
 	#[cfg(target_os = "linux")]
 	preparation_max_rss: prometheus::Histogram,
+	// Max. allocated memory, tracked by Jemallocator, polling-based
 	#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 	preparation_max_allocated: prometheus::Histogram,
+	// Max. resident memory, tracked by Jemallocator, polling-based
 	#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 	preparation_max_resident: prometheus::Histogram,
-	preparation_peak_allocation: prometheus::Histogram,
+	// Peak allocation value, tracked by tracking-allocator
+	preparation_peak_tracked_allocation: prometheus::Histogram,
 }
 
 impl metrics::Metrics for Metrics {
@@ -276,7 +279,7 @@ impl metrics::Metrics for Metrics {
 				)?,
 				registry,
 			)?,
-			preparation_peak_allocation: prometheus::register(
+			preparation_peak_tracked_allocation: prometheus::register(
 				prometheus::Histogram::with_opts(
 					prometheus::HistogramOpts::new(
 						"polkadot_pvf_preparation_peak_allocation",
diff --git a/polkadot/node/core/pvf/src/prepare/worker_intf.rs b/polkadot/node/core/pvf/src/prepare/worker_intf.rs
index 6b42d77d09b5..8c7a7e4ad475 100644
--- a/polkadot/node/core/pvf/src/prepare/worker_intf.rs
+++ b/polkadot/node/core/pvf/src/prepare/worker_intf.rs
@@ -124,14 +124,6 @@ pub async fn start_work(
 
 		match result {
 			// Received bytes from worker within the time limit.
-			Ok(Ok(Err(PrepareError::OutOfMemory))) => {
-				gum::warn!(
-					target: LOG_TARGET,
-					worker_pid = %pid,
-					"worker is out of memory",
-				);
-				Outcome::OutOfMemory
-			},
 			Ok(Ok(prepare_result)) =>
 				handle_response(
 					metrics,
@@ -184,6 +176,7 @@ async fn handle_response(
 		Ok(result) => result,
 		// Timed out on the child. This should already be logged by the child.
 		Err(PrepareError::TimedOut) => return Outcome::TimedOut,
+		Err(PrepareError::OutOfMemory) => return Outcome::OutOfMemory,
 		Err(_) => return Outcome::Concluded { worker, result },
 	};
 
diff --git a/polkadot/node/core/pvf/tests/it/adder.rs b/polkadot/node/core/pvf/tests/it/adder.rs
index bad7a66054c9..8622a4c88b36 100644
--- a/polkadot/node/core/pvf/tests/it/adder.rs
+++ b/polkadot/node/core/pvf/tests/it/adder.rs
@@ -14,6 +14,8 @@
 // You should have received a copy of the GNU General Public License
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
+//! PVF host integration tests checking the chain production pipeline.
+
 use super::TestHost;
 use adder::{hash_state, BlockData, HeadData};
 use parity_scale_codec::{Decode, Encode};
diff --git a/polkadot/node/core/pvf/tests/it/main.rs b/polkadot/node/core/pvf/tests/it/main.rs
index e76db4229bdd..c16b40a55aa9 100644
--- a/polkadot/node/core/pvf/tests/it/main.rs
+++ b/polkadot/node/core/pvf/tests/it/main.rs
@@ -14,16 +14,15 @@
 // You should have received a copy of the GNU General Public License
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
-#[cfg(feature = "ci-only-tests")]
+//! General PVF host integration tests checking the functionality of the PVF host itself.
+
 use assert_matches::assert_matches;
 use parity_scale_codec::Encode as _;
 use polkadot_node_core_pvf::{
 	start, Config, InvalidCandidate, Metrics, PrepareError, PrepareJobKind, PrepareStats,
 	PvfPrepData, ValidationError, ValidationHost, JOB_TIMEOUT_WALL_CLOCK_FACTOR,
 };
-use polkadot_parachain_primitives::primitives::{
-	BlockData as GenericBlockData, ValidationParams, ValidationResult,
-};
+use polkadot_parachain_primitives::primitives::{BlockData, ValidationParams, ValidationResult};
 use polkadot_primitives::{ExecutorParam, ExecutorParams};
 
 use std::time::Duration;
@@ -137,7 +136,7 @@ async fn terminates_on_timeout() {
 		.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: GenericBlockData(Vec::new()),
+				block_data: BlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -217,7 +216,7 @@ async fn execute_queue_doesnt_stall_if_workers_died() {
 		host.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: GenericBlockData(Vec::new()),
+				block_data: BlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -299,7 +298,7 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 		.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: GenericBlockData(Vec::new()),
+				block_data: BlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -329,7 +328,7 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 		.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: GenericBlockData(Vec::new()),
+				block_data: BlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),
@@ -359,10 +358,7 @@ async fn prechecking_within_memory_limits() {
 		)
 		.await;
 
-	match result {
-		Ok(_) => (),
-		r => panic!("{:?}", r),
-	}
+	assert_matches!(result, Ok(_));
 }
 
 // This test checks if the adder parachain runtime can be prepared with 512Kb preparation memory
@@ -380,8 +376,6 @@ async fn prechecking_out_of_memory() {
 			ExecutorParams::from(&[ExecutorParam::PrecheckingMaxMemory(512 * 1024)][..]),
 		)
 		.await;
-	match result {
-		Err(PrepareError::OutOfMemory) => (),
-		r => panic!("{:?}", r),
-	}
+
+	assert_matches!(result, Err(PrepareError::OutOfMemory));
 }
diff --git a/polkadot/src/main.rs b/polkadot/src/main.rs
index fd133adfce51..1a96bf8fb00f 100644
--- a/polkadot/src/main.rs
+++ b/polkadot/src/main.rs
@@ -20,6 +20,8 @@
 
 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"))]
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

From 02ad5367b1a5f2f715d3fba43b7719310c99a9f3 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Sat, 23 Sep 2023 01:42:53 +0200
Subject: [PATCH 14/33] Fix tests

---
 polkadot/node/core/pvf/tests/it/main.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/polkadot/node/core/pvf/tests/it/main.rs b/polkadot/node/core/pvf/tests/it/main.rs
index c16b40a55aa9..2da7c9a8055c 100644
--- a/polkadot/node/core/pvf/tests/it/main.rs
+++ b/polkadot/node/core/pvf/tests/it/main.rs
@@ -163,7 +163,7 @@ async fn ensure_parallel_execution() {
 	let execute_pvf_future_1 = host.validate_candidate(
 		halt::wasm_binary_unwrap(),
 		ValidationParams {
-			block_data: GenericBlockData(Vec::new()),
+			block_data: BlockData(Vec::new()),
 			parent_head: Default::default(),
 			relay_parent_number: 1,
 			relay_parent_storage_root: Default::default(),
@@ -173,7 +173,7 @@ async fn ensure_parallel_execution() {
 	let execute_pvf_future_2 = host.validate_candidate(
 		halt::wasm_binary_unwrap(),
 		ValidationParams {
-			block_data: GenericBlockData(Vec::new()),
+			block_data: BlockData(Vec::new()),
 			parent_head: Default::default(),
 			relay_parent_number: 1,
 			relay_parent_storage_root: Default::default(),
@@ -258,7 +258,7 @@ async fn execute_queue_doesnt_stall_with_varying_executor_params() {
 		host.validate_candidate(
 			halt::wasm_binary_unwrap(),
 			ValidationParams {
-				block_data: GenericBlockData(Vec::new()),
+				block_data: BlockData(Vec::new()),
 				parent_head: Default::default(),
 				relay_parent_number: 1,
 				relay_parent_storage_root: Default::default(),

From e1a2b631af629cf7ac5da57c8f6797d9add3119d Mon Sep 17 00:00:00 2001
From: s0me0ne-unkn0wn <48632512+s0me0ne-unkn0wn@users.noreply.github.com>
Date: Mon, 25 Sep 2023 16:59:38 +0200
Subject: [PATCH 15/33] Fix typo

Co-authored-by: Koute <koute@users.noreply.github.com>
---
 polkadot/node/core/pvf/common/src/prepare.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/polkadot/node/core/pvf/common/src/prepare.rs b/polkadot/node/core/pvf/common/src/prepare.rs
index cefb906b6ac9..4436ebe4861e 100644
--- a/polkadot/node/core/pvf/common/src/prepare.rs
+++ b/polkadot/node/core/pvf/common/src/prepare.rs
@@ -29,7 +29,7 @@ pub struct PrepareStats {
 /// supported by the OS, `ru_maxrss`.
 #[derive(Clone, Debug, Default, Encode, Decode)]
 pub struct MemoryStats {
-	/// Memory stats from `tikv_jemalloc_ctl`, polling-based and not vary precise.
+	/// Memory stats from `tikv_jemalloc_ctl`, polling-based and not very precise.
 	#[cfg(any(target_os = "linux", feature = "jemalloc-allocator"))]
 	pub memory_tracker_stats: Option<MemoryAllocationStats>,
 	/// `ru_maxrss` from `getrusage`. `None` if an error occurred.

From 1918d2a51f8c66534bf2a907d3be416eca105a42 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Tue, 26 Sep 2023 12:17:30 +0200
Subject: [PATCH 16/33] Address discussions

---
 polkadot/node/core/pvf/Cargo.toml             |  1 -
 .../node/core/pvf/prepare-worker/src/lib.rs   | 54 ++++++++++++++-----
 polkadot/node/malus/Cargo.toml                |  2 +-
 polkadot/node/tracking-allocator/src/lib.rs   | 44 ++++++++-------
 4 files changed, 67 insertions(+), 34 deletions(-)

diff --git a/polkadot/node/core/pvf/Cargo.toml b/polkadot/node/core/pvf/Cargo.toml
index 478d1952d9d9..571e9c8a928a 100644
--- a/polkadot/node/core/pvf/Cargo.toml
+++ b/polkadot/node/core/pvf/Cargo.toml
@@ -40,7 +40,6 @@ substrate-build-script-utils = { path = "../../../../substrate/utils/build-scrip
 assert_matches = "1.4.0"
 hex-literal = "0.4.1"
 polkadot-node-core-pvf-common = { path = "common", features = ["test-utils"] }
-# For the puppet worker, depend on ourselves with the test-utils feature.
 polkadot-node-core-pvf = { path = ".", features = ["test-utils"] }
 
 adder = { package = "test-parachain-adder", path = "../../../parachain/test-parachains/adder" }
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index e6c9ef2c09b0..582b110e830e 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -181,20 +181,46 @@ pub fn worker_entrypoint(
 				)?;
 
 				let fd = stream.as_raw_fd();
-				ALLOC.start_tracking(
-					executor_params.prechecking_max_memory().map(|v| v as isize),
-					Some(Box::new(move || unsafe {
-						// Inside the failure handler, the allocator is locked and no allocations
-						// are possible
-						libc::write(
-							fd,
-							OOM_PAYLOAD as *const _ as *const libc::c_void,
-							OOM_PAYLOAD.len(),
-						);
-						libc::close(fd);
-						std::process::exit(1);
-					})),
-				);
+				unsafe {
+					// SAFETY: Inside the failure handler, the allocator is locked and no
+					// allocations or deallocations are possible. For Linux, that always holds for
+					// the code below, so it's safe. For MacOS, that technically holds at the time
+					// of writing, but there's no future guarantees.
+					// The arguments of unsafe `libc` calls are valid, the payload validity is
+					// covered with a test.
+					ALLOC.start_tracking(
+						executor_params.prechecking_max_memory().map(|v| {
+							v.try_into().unwrap_or_else(|_| {
+								gum::warn!(
+									LOG_TARGET,
+									%worker_pid,
+									"Illegal pre-checking max memory value {} discarded",
+									v,
+								);
+								0
+							})
+						}),
+						Some(Box::new(move || {
+							#[cfg(target_os = "linux")]
+							{
+								libc::syscall(
+									libc::SYS_write,
+									fd,
+									OOM_PAYLOAD.as_ptr(),
+									OOM_PAYLOAD.len(),
+								);
+								libc::syscall(libc::SYS_close, fd);
+								libc::syscall(libc::SYS_exit, 1);
+							}
+							#[cfg(not(target_os = "linux"))]
+							{
+								libc::write(fd, OOM_PAYLOAD.as_ptr().cast(), OOM_PAYLOAD.len());
+								libc::close(fd);
+								std::process::exit(1);
+							}
+						})),
+					);
+				}
 
 				// Spawn another thread for preparation.
 				let prepare_thread = thread::spawn_worker_thread(
diff --git a/polkadot/node/malus/Cargo.toml b/polkadot/node/malus/Cargo.toml
index 42dd4af73c13..1135389019b7 100644
--- a/polkadot/node/malus/Cargo.toml
+++ b/polkadot/node/malus/Cargo.toml
@@ -48,7 +48,7 @@ erasure = { package = "polkadot-erasure-coding", path = "../../erasure-coding" }
 rand = "0.8.5"
 
 # Required for worker binaries to build.
-polkadot-node-core-pvf-common = { path = "../core/pvf/common", features = ["test-utils"] }
+polkadot-node-core-pvf-common = { path = "../core/pvf/common" }
 polkadot-node-core-pvf-execute-worker = { path = "../core/pvf/execute-worker" }
 polkadot-node-core-pvf-prepare-worker = { path = "../core/pvf/prepare-worker" }
 
diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 4debea683d9e..3131ef0647ca 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -14,7 +14,9 @@
 // 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.
+//! Tracking/limiting global allocator. Calculates the peak allocation between two checkpoints for
+//! the whole process. Accepts an optional limit and a failure handler which is called if the limit
+//! is overflown.
 
 use core::alloc::{GlobalAlloc, Layout};
 use std::{
@@ -63,6 +65,7 @@ impl TrackingAllocatorData {
 		self.current = 0;
 		self.peak = 0;
 		self.limit = limit;
+		let _old_handler = self.failure_handler.take();
 		self.failure_handler = failure_handler;
 		self.unlock();
 	}
@@ -71,7 +74,7 @@ impl TrackingAllocatorData {
 		self.lock();
 		let peak = self.peak;
 		self.limit = 0;
-		self.failure_handler = None;
+		let _old_handler = self.failure_handler.take();
 		self.unlock();
 		peak
 	}
@@ -107,7 +110,14 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 	//   is isolated by an exclusive lock.
 
 	/// Start tracking
-	pub fn start_tracking(&self, limit: Option<isize>, failure_handler: Option<Box<dyn Fn()>>) {
+	/// SAFETY: Failure handler is called with the allocator being in the locked state. Thus, no
+	/// allocations or deallocations are possible inside the failure handler; otherwise, a
+	/// deadlock will occur.
+	pub unsafe fn start_tracking(
+		&self,
+		limit: Option<isize>,
+		failure_handler: Option<Box<dyn Fn()>>,
+	) {
 		unsafe {
 			ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0), failure_handler);
 		}
@@ -120,6 +130,16 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 	}
 }
 
+#[cold]
+#[inline(never)]
+unsafe fn fail_allocation() -> *mut u8 {
+	if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
+		failure_handler()
+	}
+	ALLOCATOR_DATA.unlock();
+	null_mut()
+}
+
 unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 	// SAFETY:
 	// * The wrapped methods are as safe as the underlying allocator implementation is
@@ -129,11 +149,7 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 		if ALLOCATOR_DATA.track(layout.size() as isize) {
 			self.0.alloc(layout)
 		} else {
-			if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
-				failure_handler()
-			}
-			ALLOCATOR_DATA.unlock();
-			null_mut()
+			fail_allocation()
 		}
 	}
 
@@ -142,11 +158,7 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 		if ALLOCATOR_DATA.track(layout.size() as isize) {
 			self.0.alloc_zeroed(layout)
 		} else {
-			if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
-				failure_handler()
-			}
-			ALLOCATOR_DATA.unlock();
-			null_mut()
+			fail_allocation()
 		}
 	}
 
@@ -161,11 +173,7 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 		if ALLOCATOR_DATA.track((new_size as isize) - (layout.size() as isize)) {
 			self.0.realloc(ptr, layout, new_size)
 		} else {
-			if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
-				failure_handler()
-			}
-			ALLOCATOR_DATA.unlock();
-			null_mut()
+			fail_allocation()
 		}
 	}
 }

From 007054ad08c8690e8c7e53ecb236e16aabed3528 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Tue, 26 Sep 2023 13:57:26 +0200
Subject: [PATCH 17/33] Address discussions

---
 .../node/core/pvf/prepare-worker/src/lib.rs   | 97 +++++++++++--------
 polkadot/node/core/pvf/src/metrics.rs         |  2 +-
 polkadot/node/tracking-allocator/src/lib.rs   |  8 +-
 3 files changed, 61 insertions(+), 46 deletions(-)

diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index 582b110e830e..95d42fd12b64 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -47,7 +47,7 @@ use polkadot_node_core_pvf_common::{
 };
 use polkadot_primitives::ExecutorParams;
 use std::{
-	os::fd::AsRawFd,
+	os::fd::{AsRawFd, RawFd},
 	path::PathBuf,
 	sync::{mpsc::channel, Arc},
 	time::Duration,
@@ -102,6 +102,44 @@ async fn send_response(stream: &mut UnixStream, result: PrepareResult) -> io::Re
 	framed_send(stream, &result.encode()).await
 }
 
+fn start_memory_tracking(fd: RawFd, limit: Option<isize>) {
+	unsafe {
+		// SAFETY: Inside the failure handler, the allocator is locked and no allocations or
+		// deallocations are possible. For Linux, that always holds for the code below, so it's
+		// safe. For MacOS, that technically holds at the time of writing, but there are no future
+		// guarantees.
+		// The arguments of unsafe `libc` calls are valid, the payload validity is covered with
+		// a test.
+		ALLOC.start_tracking(
+			limit,
+			Some(Box::new(move || {
+				#[cfg(target_os = "linux")]
+				{
+					// Syscalls never allocate or deallocate, so this is safe.
+					libc::syscall(libc::SYS_write, fd, OOM_PAYLOAD.as_ptr(), OOM_PAYLOAD.len());
+					libc::syscall(libc::SYS_close, fd);
+					libc::syscall(libc::SYS_exit, 1);
+				}
+				#[cfg(not(target_os = "linux"))]
+				{
+					// Syscalls are not available on MacOS, so we have to use `libc` wrappers.
+					// Technicaly, there may be allocations inside, although they shouldn't be
+					// there. In that case, we'll see deadlocks on MacOS after the OOM condition
+					// triggered. As we consider running a validator on MacOS unsafe, and this
+					// code is only run by a validator, it's a lesser evil.
+					libc::write(fd, OOM_PAYLOAD.as_ptr().cast(), OOM_PAYLOAD.len());
+					libc::close(fd);
+					std::process::exit(1);
+				}
+			})),
+		);
+	}
+}
+
+fn end_memory_tracking() -> isize {
+	ALLOC.end_tracking()
+}
+
 /// The entrypoint that the spawned prepare worker should start with.
 ///
 /// # Parameters
@@ -180,47 +218,20 @@ pub fn worker_entrypoint(
 					WaitOutcome::TimedOut,
 				)?;
 
-				let fd = stream.as_raw_fd();
-				unsafe {
-					// SAFETY: Inside the failure handler, the allocator is locked and no
-					// allocations or deallocations are possible. For Linux, that always holds for
-					// the code below, so it's safe. For MacOS, that technically holds at the time
-					// of writing, but there's no future guarantees.
-					// The arguments of unsafe `libc` calls are valid, the payload validity is
-					// covered with a test.
-					ALLOC.start_tracking(
-						executor_params.prechecking_max_memory().map(|v| {
-							v.try_into().unwrap_or_else(|_| {
-								gum::warn!(
-									LOG_TARGET,
-									%worker_pid,
-									"Illegal pre-checking max memory value {} discarded",
-									v,
-								);
-								0
-							})
-						}),
-						Some(Box::new(move || {
-							#[cfg(target_os = "linux")]
-							{
-								libc::syscall(
-									libc::SYS_write,
-									fd,
-									OOM_PAYLOAD.as_ptr(),
-									OOM_PAYLOAD.len(),
-								);
-								libc::syscall(libc::SYS_close, fd);
-								libc::syscall(libc::SYS_exit, 1);
-							}
-							#[cfg(not(target_os = "linux"))]
-							{
-								libc::write(fd, OOM_PAYLOAD.as_ptr().cast(), OOM_PAYLOAD.len());
-								libc::close(fd);
-								std::process::exit(1);
-							}
-						})),
-					);
-				}
+				start_memory_tracking(
+					stream.as_raw_fd(),
+					executor_params.prechecking_max_memory().map(|v| {
+						v.try_into().unwrap_or_else(|_| {
+							gum::warn!(
+								LOG_TARGET,
+								%worker_pid,
+								"Illegal pre-checking max memory value {} discarded",
+								v,
+							);
+							0
+						})
+					}),
+				);
 
 				// Spawn another thread for preparation.
 				let prepare_thread = thread::spawn_worker_thread(
@@ -263,7 +274,7 @@ pub fn worker_entrypoint(
 				let outcome = thread::wait_for_threads(condvar);
 
 				let peak_alloc = {
-					let peak = ALLOC.end_tracking();
+					let peak = end_memory_tracking();
 					gum::debug!(
 						target: LOG_TARGET,
 						%worker_pid,
diff --git a/polkadot/node/core/pvf/src/metrics.rs b/polkadot/node/core/pvf/src/metrics.rs
index cc41745ba277..7fd876cf1740 100644
--- a/polkadot/node/core/pvf/src/metrics.rs
+++ b/polkadot/node/core/pvf/src/metrics.rs
@@ -282,7 +282,7 @@ impl metrics::Metrics for Metrics {
 			preparation_peak_tracked_allocation: prometheus::register(
 				prometheus::Histogram::with_opts(
 					prometheus::HistogramOpts::new(
-						"polkadot_pvf_preparation_peak_allocation",
+						"polkadot_pvf_preparation_peak_tracked_allocation",
 						"peak allocation observed for preparation (in kilobytes)",
 					).buckets(
 						prometheus::exponential_buckets(8192.0, 2.0, 10)
diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 3131ef0647ca..dc4a9e93f1d4 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -65,17 +65,21 @@ impl TrackingAllocatorData {
 		self.current = 0;
 		self.peak = 0;
 		self.limit = limit;
-		let _old_handler = self.failure_handler.take();
+		// Cannot drop it yet, as it would trigger a deallocation
+		let old_handler = self.failure_handler.take();
 		self.failure_handler = failure_handler;
 		self.unlock();
+		core::mem::drop(old_handler);
 	}
 
 	fn end_tracking(&mut self) -> isize {
 		self.lock();
 		let peak = self.peak;
 		self.limit = 0;
-		let _old_handler = self.failure_handler.take();
+		// Cannot drop it yet, as it would trigger a deallocation
+		let old_handler = self.failure_handler.take();
 		self.unlock();
+		core::mem::drop(old_handler);
 		peak
 	}
 

From c81e8aab0059261cd3f85bfd9ddb4f94830bf665 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Tue, 26 Sep 2023 20:28:30 +0200
Subject: [PATCH 18/33] Try to fix test pipeline

---
 Cargo.lock                                    |  2 +-
 .../node/core/pvf/common/src/executor_intf.rs | 21 ++++++++++
 .../node/core/pvf/execute-worker/src/lib.rs   |  2 +-
 .../pvf/prepare-worker/src/executor_intf.rs   | 42 -------------------
 .../node/core/pvf/prepare-worker/src/lib.rs   |  3 +-
 polkadot/node/core/pvf/src/testing.rs         |  3 +-
 .../node/test/performance-test/Cargo.toml     |  2 +-
 .../node/test/performance-test/src/lib.rs     |  4 +-
 8 files changed, 28 insertions(+), 51 deletions(-)
 delete mode 100644 polkadot/node/core/pvf/prepare-worker/src/executor_intf.rs

diff --git a/Cargo.lock b/Cargo.lock
index 443312aa2421..9d156e967e9f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -12472,7 +12472,7 @@ dependencies = [
  "env_logger 0.9.3",
  "log",
  "polkadot-erasure-coding",
- "polkadot-node-core-pvf-prepare-worker",
+ "polkadot-node-core-pvf-common",
  "polkadot-node-primitives",
  "polkadot-primitives",
  "quote",
diff --git a/polkadot/node/core/pvf/common/src/executor_intf.rs b/polkadot/node/core/pvf/common/src/executor_intf.rs
index 79839149ebdf..9e624709cd0c 100644
--- a/polkadot/node/core/pvf/common/src/executor_intf.rs
+++ b/polkadot/node/core/pvf/common/src/executor_intf.rs
@@ -120,6 +120,27 @@ pub fn params_to_wasmtime_semantics(par: &ExecutorParams) -> Result<Semantics, S
 	Ok(sem)
 }
 
+/// Runs the prevalidation on the given code. Returns a [`RuntimeBlob`] if it succeeds.
+pub fn prevalidate(code: &[u8]) -> Result<RuntimeBlob, sc_executor_common::error::WasmError> {
+	let blob = RuntimeBlob::new(code)?;
+	// It's assumed this function will take care of any prevalidation logic
+	// that needs to be done.
+	//
+	// Do nothing for now.
+	Ok(blob)
+}
+
+/// Runs preparation on the given runtime blob. If successful, it returns a serialized compiled
+/// artifact which can then be used to pass into `Executor::execute` after writing it to the disk.
+pub fn prepare(
+	blob: RuntimeBlob,
+	executor_params: &ExecutorParams,
+) -> Result<Vec<u8>, sc_executor_common::error::WasmError> {
+	let semantics = params_to_wasmtime_semantics(executor_params)
+		.map_err(|e| sc_executor_common::error::WasmError::Other(e))?;
+	sc_executor_wasmtime::prepare_runtime_artifact(blob, &semantics)
+}
+
 /// A WASM executor with a given configuration. It is instantiated once per execute worker and is
 /// specific to that worker.
 #[derive(Clone)]
diff --git a/polkadot/node/core/pvf/execute-worker/src/lib.rs b/polkadot/node/core/pvf/execute-worker/src/lib.rs
index 36793a5c71ec..f23c998f4289 100644
--- a/polkadot/node/core/pvf/execute-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/execute-worker/src/lib.rs
@@ -16,7 +16,7 @@
 
 //! Contains the logic for executing PVFs. Used by the polkadot-execute-worker binary.
 
-pub use polkadot_node_core_pvf_common::executor_intf::Executor;
+use polkadot_node_core_pvf_common::executor_intf::Executor;
 
 // NOTE: Initializing logging in e.g. tests will not have an effect in the workers, as they are
 //       separate spawned processes. Run with e.g. `RUST_LOG=parachain::pvf-execute-worker=trace`.
diff --git a/polkadot/node/core/pvf/prepare-worker/src/executor_intf.rs b/polkadot/node/core/pvf/prepare-worker/src/executor_intf.rs
deleted file mode 100644
index 1f88f6a6dd6e..000000000000
--- a/polkadot/node/core/pvf/prepare-worker/src/executor_intf.rs
+++ /dev/null
@@ -1,42 +0,0 @@
-// 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/>.
-
-//! Interface to the Substrate Executor
-
-use polkadot_node_core_pvf_common::executor_intf::params_to_wasmtime_semantics;
-use polkadot_primitives::ExecutorParams;
-use sc_executor_common::runtime_blob::RuntimeBlob;
-
-/// Runs the prevalidation on the given code. Returns a [`RuntimeBlob`] if it succeeds.
-pub fn prevalidate(code: &[u8]) -> Result<RuntimeBlob, sc_executor_common::error::WasmError> {
-	let blob = RuntimeBlob::new(code)?;
-	// It's assumed this function will take care of any prevalidation logic
-	// that needs to be done.
-	//
-	// Do nothing for now.
-	Ok(blob)
-}
-
-/// Runs preparation on the given runtime blob. If successful, it returns a serialized compiled
-/// artifact which can then be used to pass into `Executor::execute` after writing it to the disk.
-pub fn prepare(
-	blob: RuntimeBlob,
-	executor_params: &ExecutorParams,
-) -> Result<Vec<u8>, sc_executor_common::error::WasmError> {
-	let semantics = params_to_wasmtime_semantics(executor_params)
-		.map_err(|e| sc_executor_common::error::WasmError::Other(e))?;
-	sc_executor_wasmtime::prepare_runtime_artifact(blob, &semantics)
-}
diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index 95d42fd12b64..2d06ab0a6bb1 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -16,10 +16,9 @@
 
 //! Contains the logic for preparing PVFs. Used by the polkadot-prepare-worker binary.
 
-mod executor_intf;
 mod memory_stats;
 
-pub use executor_intf::{prepare, prevalidate};
+use polkadot_node_core_pvf_common::executor_intf::{prepare, prevalidate};
 
 // NOTE: Initializing logging in e.g. tests will not have an effect in the workers, as they are
 //       separate spawned processes. Run with e.g. `RUST_LOG=parachain::pvf-prepare-worker=trace`.
diff --git a/polkadot/node/core/pvf/src/testing.rs b/polkadot/node/core/pvf/src/testing.rs
index 4301afc3cc7e..29dc2c509347 100644
--- a/polkadot/node/core/pvf/src/testing.rs
+++ b/polkadot/node/core/pvf/src/testing.rs
@@ -29,8 +29,7 @@ pub fn validate_candidate(
 	code: &[u8],
 	params: &[u8],
 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
-	use polkadot_node_core_pvf_execute_worker::Executor;
-	use polkadot_node_core_pvf_prepare_worker::{prepare, prevalidate};
+	use polkadot_node_core_pvf_common::executor_intf::{prepare, prevalidate, Executor};
 
 	let code = sp_maybe_compressed_blob::decompress(code, 10 * 1024 * 1024)
 		.expect("Decompressing code failed");
diff --git a/polkadot/node/test/performance-test/Cargo.toml b/polkadot/node/test/performance-test/Cargo.toml
index 5747ac88b1e4..d04757790ae0 100644
--- a/polkadot/node/test/performance-test/Cargo.toml
+++ b/polkadot/node/test/performance-test/Cargo.toml
@@ -12,7 +12,7 @@ quote = "1.0.28"
 env_logger = "0.9"
 log = "0.4"
 
-polkadot-node-core-pvf-prepare-worker = { path = "../../core/pvf/prepare-worker" }
+polkadot-node-core-pvf-common = { path = "../../core/pvf/common" }
 polkadot-erasure-coding = { path = "../../../erasure-coding" }
 polkadot-node-primitives = { path = "../../primitives" }
 polkadot-primitives = { path = "../../../primitives" }
diff --git a/polkadot/node/test/performance-test/src/lib.rs b/polkadot/node/test/performance-test/src/lib.rs
index 15073912654a..2a0a5e737901 100644
--- a/polkadot/node/test/performance-test/src/lib.rs
+++ b/polkadot/node/test/performance-test/src/lib.rs
@@ -65,9 +65,9 @@ pub fn measure_pvf_prepare(wasm_code: &[u8]) -> Result<Duration, PerfCheckError>
 		.or(Err(PerfCheckError::CodeDecompressionFailed))?;
 
 	// Recreate the pipeline from the pvf prepare worker.
-	let blob = polkadot_node_core_pvf_prepare_worker::prevalidate(code.as_ref())
+	let blob = polkadot_node_core_pvf_common::executor_intf::prevalidate(code.as_ref())
 		.map_err(PerfCheckError::from)?;
-	polkadot_node_core_pvf_prepare_worker::prepare(blob, &ExecutorParams::default())
+	polkadot_node_core_pvf_common::executor_intf::prepare(blob, &ExecutorParams::default())
 		.map_err(PerfCheckError::from)?;
 
 	Ok(start.elapsed())

From 5920260f70755973238d8fd71734b1934fda73f7 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Tue, 26 Sep 2023 20:35:22 +0200
Subject: [PATCH 19/33] Fix PVF preparation benchmark

---
 .../pvf/prepare-worker/benches/prepare_kusama_runtime.rs   | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs b/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
index ac4213c0476c..4602d08a6068 100644
--- a/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
+++ b/polkadot/node/core/pvf/prepare-worker/benches/prepare_kusama_runtime.rs
@@ -15,8 +15,11 @@
 // 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_node_core_pvf_common::{
+	executor_intf::{prepare, prevalidate},
+	prepare::PrepareJobKind,
+	pvf::PvfPrepData,
+};
 use polkadot_primitives::ExecutorParams;
 use std::time::Duration;
 

From 96d71c960829e1398fcfec6158253b624bbe59ba Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Wed, 27 Sep 2023 10:27:32 +0200
Subject: [PATCH 20/33] Minor fixes

---
 polkadot/node/core/pvf/prepare-worker/src/lib.rs | 4 ++--
 polkadot/node/tracking-allocator/src/lib.rs      | 6 ++----
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/polkadot/node/core/pvf/prepare-worker/src/lib.rs b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
index 2d06ab0a6bb1..9dce321da9f1 100644
--- a/polkadot/node/core/pvf/prepare-worker/src/lib.rs
+++ b/polkadot/node/core/pvf/prepare-worker/src/lib.rs
@@ -239,8 +239,8 @@ pub fn worker_entrypoint(
 						// Try to enable landlock.
 						#[cfg(target_os = "linux")]
 						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());
+							.map(LandlockStatus::from_ruleset_status)
+							.map_err(|e| e.to_string());
 						#[cfg(not(target_os = "linux"))]
 						let landlock_status: Result<LandlockStatus, String> = Ok(LandlockStatus::NotEnforced);
 
diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index dc4a9e93f1d4..e9f1a99b221c 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -115,16 +115,14 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 
 	/// Start tracking
 	/// SAFETY: Failure handler is called with the allocator being in the locked state. Thus, no
-	/// allocations or deallocations are possible inside the failure handler; otherwise, a
+	/// allocations or deallocations are allowed inside the failure handler; otherwise, a
 	/// deadlock will occur.
 	pub unsafe fn start_tracking(
 		&self,
 		limit: Option<isize>,
 		failure_handler: Option<Box<dyn Fn()>>,
 	) {
-		unsafe {
-			ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0), failure_handler);
-		}
+		ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0), failure_handler);
 	}
 
 	/// End tracking and return the peak allocation value in bytes (as `isize`). Peak allocation

From 93b9d616cfe6815379fb559c429e7ccdc6ba1e7e Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Thu, 28 Sep 2023 19:42:11 +0200
Subject: [PATCH 21/33] Eliminate redundant branch

---
 polkadot/node/tracking-allocator/src/lib.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index e9f1a99b221c..53126d144b0b 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -90,7 +90,7 @@ impl TrackingAllocatorData {
 		if self.current > self.peak {
 			self.peak = self.current;
 		}
-		let within_limits = self.limit == 0 || self.peak <= self.limit;
+		let within_limits = self.peak <= self.limit;
 		if within_limits {
 			self.unlock()
 		}
@@ -122,7 +122,7 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 		limit: Option<isize>,
 		failure_handler: Option<Box<dyn Fn()>>,
 	) {
-		ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0), failure_handler);
+		ALLOCATOR_DATA.start_tracking(limit.unwrap_or(isize::MAX), failure_handler);
 	}
 
 	/// End tracking and return the peak allocation value in bytes (as `isize`). Peak allocation

From 7a3875a441766189b64ecfa0f2b6e15bbc390e4e Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Thu, 28 Sep 2023 20:52:13 +0200
Subject: [PATCH 22/33] Revert "Eliminate redundant branch"

This reverts commit 93b9d616cfe6815379fb559c429e7ccdc6ba1e7e.
---
 polkadot/node/tracking-allocator/src/lib.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 53126d144b0b..e9f1a99b221c 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -90,7 +90,7 @@ impl TrackingAllocatorData {
 		if self.current > self.peak {
 			self.peak = self.current;
 		}
-		let within_limits = self.peak <= self.limit;
+		let within_limits = self.limit == 0 || self.peak <= self.limit;
 		if within_limits {
 			self.unlock()
 		}
@@ -122,7 +122,7 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 		limit: Option<isize>,
 		failure_handler: Option<Box<dyn Fn()>>,
 	) {
-		ALLOCATOR_DATA.start_tracking(limit.unwrap_or(isize::MAX), failure_handler);
+		ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0), failure_handler);
 	}
 
 	/// End tracking and return the peak allocation value in bytes (as `isize`). Peak allocation

From 83e6b992f2392f62d1287cc6d3ec2f81017fbf3b Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Thu, 28 Sep 2023 21:01:34 +0200
Subject: [PATCH 23/33] Remove redundant todo comment

---
 polkadot/node/core/pvf/common/src/executor_intf.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/polkadot/node/core/pvf/common/src/executor_intf.rs b/polkadot/node/core/pvf/common/src/executor_intf.rs
index 9e624709cd0c..cf6dc6649691 100644
--- a/polkadot/node/core/pvf/common/src/executor_intf.rs
+++ b/polkadot/node/core/pvf/common/src/executor_intf.rs
@@ -111,9 +111,9 @@ pub fn params_to_wasmtime_semantics(par: &ExecutorParams) -> Result<Semantics, S
 			ExecutorParam::StackLogicalMax(slm) => stack_limit.logical_max = *slm,
 			ExecutorParam::StackNativeMax(snm) => stack_limit.native_stack_max = *snm,
 			ExecutorParam::WasmExtBulkMemory => sem.wasm_bulk_memory = true,
-			// TODO: Not implemented yet; <https://github.com/paritytech/polkadot/issues/6472>.
-			ExecutorParam::PrecheckingMaxMemory(_) => (),
-			ExecutorParam::PvfPrepTimeout(_, _) | ExecutorParam::PvfExecTimeout(_, _) => (), /* Not used here */
+			ExecutorParam::PrecheckingMaxMemory(_) |
+			ExecutorParam::PvfPrepTimeout(_, _) |
+			ExecutorParam::PvfExecTimeout(_, _) => (), /* Not used here */
 		}
 	}
 	sem.deterministic_stack_limit = Some(stack_limit);

From d74aa205ff184cd1169b76646a828dab25f6accd Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Wed, 11 Oct 2023 14:09:20 +0200
Subject: [PATCH 24/33] Remove stale file

---
 polkadot/node/test/performance-test/src/lib.rs | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 delete mode 100644 polkadot/node/test/performance-test/src/lib.rs

diff --git a/polkadot/node/test/performance-test/src/lib.rs b/polkadot/node/test/performance-test/src/lib.rs
deleted file mode 100644
index e69de29bb2d1..000000000000

From 8c58a681e3c8634f367d51eaea9e6745290c88f8 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Mon, 23 Oct 2023 18:47:34 +0200
Subject: [PATCH 25/33] Abstract `Spinlock`

---
 polkadot/node/tracking-allocator/src/lib.rs | 167 +++++++++++++-------
 1 file changed, 106 insertions(+), 61 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index e9f1a99b221c..6bb84cf466c0 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -18,23 +18,36 @@
 //! the whole process. Accepts an optional limit and a failure handler which is called if the limit
 //! is overflown.
 
-use core::alloc::{GlobalAlloc, Layout};
+use core::{
+	alloc::{GlobalAlloc, Layout},
+	ops::{Deref, DerefMut},
+};
 use std::{
+	cell::UnsafeCell,
 	ptr::null_mut,
 	sync::atomic::{AtomicBool, Ordering},
 };
 
-struct TrackingAllocatorData {
+struct Spinlock<T: Send> {
 	lock: AtomicBool,
-	current: isize,
-	peak: isize,
-	limit: isize,
-	failure_handler: Option<Box<dyn Fn()>>,
+	data: UnsafeCell<T>,
 }
 
-impl TrackingAllocatorData {
+struct SpinlockGuard<'a, T: 'a + Send> {
+	lock: &'a Spinlock<T>,
+}
+
+// SAFETY: Data under `UnsafeCell<T>` are protected by an exclusive lock, so they may be shared
+// between threads safely
+unsafe impl<T: Send> Sync for Spinlock<T> {}
+
+impl<T: Send> Spinlock<T> {
+	pub const fn new(t: T) -> Spinlock<T> {
+		Spinlock { lock: AtomicBool::new(false), data: UnsafeCell::new(t) }
+	}
+
 	#[inline]
-	fn lock(&self) {
+	pub fn lock(&self) -> SpinlockGuard<T> {
 		loop {
 			// Try to acquire the lock.
 			if self
@@ -42,7 +55,7 @@ impl TrackingAllocatorData {
 				.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
 				.is_ok()
 			{
-				break
+				return SpinlockGuard { lock: self }
 			}
 			// We failed to acquire the lock; wait until it's unlocked.
 			//
@@ -59,60 +72,83 @@ impl TrackingAllocatorData {
 	fn unlock(&self) {
 		self.lock.store(false, Ordering::Release);
 	}
+}
+
+impl<T: Send> Deref for SpinlockGuard<'_, T> {
+	type Target = T;
+
+	fn deref(&self) -> &T {
+		// SAFETY: It is safe to dereference a guard to the `UnsafeCell` underlying data as the
+		// presence of the guard means the data are already locked.
+		unsafe { &*self.lock.data.get() }
+	}
+}
+
+impl<T: Send> DerefMut for SpinlockGuard<'_, T> {
+	fn deref_mut(&mut self) -> &mut T {
+		unsafe { &mut *self.lock.data.get() }
+	}
+}
+
+impl<T: Send> Drop for SpinlockGuard<'_, T> {
+	fn drop(&mut self) {
+		self.lock.unlock();
+	}
+}
+
+struct TrackingAllocatorData {
+	current: isize,
+	peak: isize,
+	limit: isize,
+	failure_handler: Option<Box<dyn Fn() + Send>>,
+}
 
-	fn start_tracking(&mut self, limit: isize, failure_handler: Option<Box<dyn Fn()>>) {
-		self.lock();
-		self.current = 0;
-		self.peak = 0;
-		self.limit = limit;
+impl TrackingAllocatorData {
+	fn start_tracking(
+		mut guard: SpinlockGuard<Self>,
+		limit: isize,
+		failure_handler: Option<Box<dyn Fn() + Send>>,
+	) {
+		guard.current = 0;
+		guard.peak = 0;
+		guard.limit = limit;
 		// Cannot drop it yet, as it would trigger a deallocation
-		let old_handler = self.failure_handler.take();
-		self.failure_handler = failure_handler;
-		self.unlock();
-		core::mem::drop(old_handler);
+		let old_handler = guard.failure_handler.take();
+		guard.failure_handler = failure_handler;
+		drop(guard);
+		drop(old_handler);
 	}
 
-	fn end_tracking(&mut self) -> isize {
-		self.lock();
-		let peak = self.peak;
-		self.limit = 0;
+	fn end_tracking(mut guard: SpinlockGuard<Self>) -> isize {
+		let peak = guard.peak;
+		guard.limit = 0;
 		// Cannot drop it yet, as it would trigger a deallocation
-		let old_handler = self.failure_handler.take();
-		self.unlock();
-		core::mem::drop(old_handler);
+		let old_handler = guard.failure_handler.take();
+		drop(guard);
+		drop(old_handler);
 		peak
 	}
 
 	#[inline]
-	fn track(&mut self, alloc: isize) -> bool {
-		self.lock();
-		self.current += alloc;
-		if self.current > self.peak {
-			self.peak = self.current;
+	fn track(mut guard: SpinlockGuard<Self>, alloc: isize) -> Option<SpinlockGuard<Self>> {
+		guard.current += alloc;
+		if guard.current > guard.peak {
+			guard.peak = guard.current;
 		}
-		let within_limits = self.limit == 0 || self.peak <= self.limit;
-		if within_limits {
-			self.unlock()
+		if guard.limit == 0 || guard.peak <= guard.limit {
+			None
+		} else {
+			Some(guard)
 		}
-		within_limits
 	}
 }
 
-static mut ALLOCATOR_DATA: TrackingAllocatorData = TrackingAllocatorData {
-	lock: AtomicBool::new(false),
-	current: 0,
-	peak: 0,
-	limit: 0,
-	failure_handler: None,
-};
+static ALLOCATOR_DATA: Spinlock<TrackingAllocatorData> =
+	Spinlock::new(TrackingAllocatorData { current: 0, peak: 0, limit: 0, failure_handler: None });
 
 pub struct TrackingAllocator<A: GlobalAlloc>(pub 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
 	/// SAFETY: Failure handler is called with the allocator being in the locked state. Thus, no
 	/// allocations or deallocations are allowed inside the failure handler; otherwise, a
@@ -120,25 +156,28 @@ impl<A: GlobalAlloc> TrackingAllocator<A> {
 	pub unsafe fn start_tracking(
 		&self,
 		limit: Option<isize>,
-		failure_handler: Option<Box<dyn Fn()>>,
+		failure_handler: Option<Box<dyn Fn() + Send>>,
 	) {
-		ALLOCATOR_DATA.start_tracking(limit.unwrap_or(0), failure_handler);
+		TrackingAllocatorData::start_tracking(
+			ALLOCATOR_DATA.lock(),
+			limit.unwrap_or(0),
+			failure_handler,
+		);
 	}
 
 	/// 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() }
+		TrackingAllocatorData::end_tracking(ALLOCATOR_DATA.lock())
 	}
 }
 
 #[cold]
 #[inline(never)]
-unsafe fn fail_allocation() -> *mut u8 {
-	if let Some(failure_handler) = &ALLOCATOR_DATA.failure_handler {
+unsafe fn fail_allocation(guard: SpinlockGuard<TrackingAllocatorData>) -> *mut u8 {
+	if let Some(failure_handler) = &guard.failure_handler {
 		failure_handler()
 	}
-	ALLOCATOR_DATA.unlock();
 	null_mut()
 }
 
@@ -148,34 +187,40 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 
 	#[inline]
 	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-		if ALLOCATOR_DATA.track(layout.size() as isize) {
-			self.0.alloc(layout)
+		let guard = ALLOCATOR_DATA.lock();
+		if let Some(guard) = TrackingAllocatorData::track(guard, layout.size() as isize) {
+			fail_allocation(guard)
 		} else {
-			fail_allocation()
+			self.0.alloc(layout)
 		}
 	}
 
 	#[inline]
 	unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
-		if ALLOCATOR_DATA.track(layout.size() as isize) {
-			self.0.alloc_zeroed(layout)
+		let guard = ALLOCATOR_DATA.lock();
+		if let Some(guard) = TrackingAllocatorData::track(guard, layout.size() as isize) {
+			fail_allocation(guard)
 		} else {
-			fail_allocation()
+			self.0.alloc_zeroed(layout)
 		}
 	}
 
 	#[inline]
 	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) -> () {
-		ALLOCATOR_DATA.track(-(layout.size() as isize));
+		let guard = ALLOCATOR_DATA.lock();
+		TrackingAllocatorData::track(guard, -(layout.size() as isize));
 		self.0.dealloc(ptr, layout)
 	}
 
 	#[inline]
 	unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
-		if ALLOCATOR_DATA.track((new_size as isize) - (layout.size() as isize)) {
-			self.0.realloc(ptr, layout, new_size)
+		let guard = ALLOCATOR_DATA.lock();
+		if let Some(guard) =
+			TrackingAllocatorData::track(guard, (new_size as isize) - (layout.size() as isize))
+		{
+			fail_allocation(guard)
 		} else {
-			fail_allocation()
+			self.0.realloc(ptr, layout, new_size)
 		}
 	}
 }

From 778513cd7ba6604c6f1afc42682cb300ea6fe79a Mon Sep 17 00:00:00 2001
From: s0me0ne-unkn0wn <48632512+s0me0ne-unkn0wn@users.noreply.github.com>
Date: Sun, 29 Oct 2023 09:05:22 +0100
Subject: [PATCH 26/33] Update comment

Co-authored-by: Koute <koute@users.noreply.github.com>
---
 polkadot/node/tracking-allocator/src/lib.rs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 6bb84cf466c0..4881d94aa9b2 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -37,8 +37,12 @@ struct SpinlockGuard<'a, T: 'a + Send> {
 	lock: &'a Spinlock<T>,
 }
 
-// SAFETY: Data under `UnsafeCell<T>` are protected by an exclusive lock, so they may be shared
-// between threads safely
+// SAFETY: We require that the data inside of the `SpinLock` is `Send`, so it can be sent
+// and accessed by any thread as long as it's accessed by only one thread at a time.
+// The `SpinLock` provides an exclusive lock over it, so it guarantees that multiple
+// threads cannot access it at the same time, hence it implements `Sync` (that is, it can be
+// accessed concurrently from multiple threads, even though the `T` itself might not
+// necessarily be `Sync` too).
 unsafe impl<T: Send> Sync for Spinlock<T> {}
 
 impl<T: Send> Spinlock<T> {

From 9c3ea9bda845ffbbe0beae16e5228f9610c4507f Mon Sep 17 00:00:00 2001
From: s0me0ne-unkn0wn <48632512+s0me0ne-unkn0wn@users.noreply.github.com>
Date: Sun, 29 Oct 2023 09:06:05 +0100
Subject: [PATCH 27/33] Add safety comment

Co-authored-by: Koute <koute@users.noreply.github.com>
---
 polkadot/node/tracking-allocator/src/lib.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 4881d94aa9b2..80c302487f92 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -90,6 +90,7 @@ impl<T: Send> Deref for SpinlockGuard<'_, T> {
 
 impl<T: Send> DerefMut for SpinlockGuard<'_, T> {
 	fn deref_mut(&mut self) -> &mut T {
+		// SAFETY: Same as for `Deref::deref`.
 		unsafe { &mut *self.lock.data.get() }
 	}
 }

From 54686a1cba8d86acb6c7c98be717efcf8cce7d57 Mon Sep 17 00:00:00 2001
From: s0me0ne-unkn0wn <48632512+s0me0ne-unkn0wn@users.noreply.github.com>
Date: Sun, 29 Oct 2023 09:06:31 +0100
Subject: [PATCH 28/33] Fix comment formatting

Co-authored-by: Koute <koute@users.noreply.github.com>
---
 polkadot/node/tracking-allocator/src/lib.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 80c302487f92..7b6335edda0d 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -154,8 +154,11 @@ static ALLOCATOR_DATA: Spinlock<TrackingAllocatorData> =
 pub struct TrackingAllocator<A: GlobalAlloc>(pub A);
 
 impl<A: GlobalAlloc> TrackingAllocator<A> {
-	/// Start tracking
-	/// SAFETY: Failure handler is called with the allocator being in the locked state. Thus, no
+	/// Start tracking memory allocations and deallocations.
+	///
+	/// # Safety
+	///
+	/// Failure handler is called with the allocator being in the locked state. Thus, no
 	/// allocations or deallocations are allowed inside the failure handler; otherwise, a
 	/// deadlock will occur.
 	pub unsafe fn start_tracking(

From 07eb6a9cc1cb50b83010756c79b7c3d82f873417 Mon Sep 17 00:00:00 2001
From: s0me0ne-unkn0wn <48632512+s0me0ne-unkn0wn@users.noreply.github.com>
Date: Sun, 29 Oct 2023 09:27:08 +0100
Subject: [PATCH 29/33] Fix typo

Co-authored-by: Koute <koute@users.noreply.github.com>
---
 polkadot/node/tracking-allocator/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index 7b6335edda0d..d2bdb13dc6e8 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -83,7 +83,7 @@ impl<T: Send> Deref for SpinlockGuard<'_, T> {
 
 	fn deref(&self) -> &T {
 		// SAFETY: It is safe to dereference a guard to the `UnsafeCell` underlying data as the
-		// presence of the guard means the data are already locked.
+		// presence of the guard means the data is already locked.
 		unsafe { &*self.lock.data.get() }
 	}
 }

From 8be635954bbfbede60c1505fc10f60eec39cc5eb Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Sun, 29 Oct 2023 10:10:39 +0100
Subject: [PATCH 30/33] Remove unneeded trait bound

---
 polkadot/node/tracking-allocator/src/lib.rs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index d2bdb13dc6e8..a1bf21e0d5dc 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -28,12 +28,12 @@ use std::{
 	sync::atomic::{AtomicBool, Ordering},
 };
 
-struct Spinlock<T: Send> {
+struct Spinlock<T> {
 	lock: AtomicBool,
 	data: UnsafeCell<T>,
 }
 
-struct SpinlockGuard<'a, T: 'a + Send> {
+struct SpinlockGuard<'a, T: 'a> {
 	lock: &'a Spinlock<T>,
 }
 
@@ -45,7 +45,7 @@ struct SpinlockGuard<'a, T: 'a + Send> {
 // necessarily be `Sync` too).
 unsafe impl<T: Send> Sync for Spinlock<T> {}
 
-impl<T: Send> Spinlock<T> {
+impl<T> Spinlock<T> {
 	pub const fn new(t: T) -> Spinlock<T> {
 		Spinlock { lock: AtomicBool::new(false), data: UnsafeCell::new(t) }
 	}
@@ -78,7 +78,7 @@ impl<T: Send> Spinlock<T> {
 	}
 }
 
-impl<T: Send> Deref for SpinlockGuard<'_, T> {
+impl<T> Deref for SpinlockGuard<'_, T> {
 	type Target = T;
 
 	fn deref(&self) -> &T {
@@ -88,14 +88,14 @@ impl<T: Send> Deref for SpinlockGuard<'_, T> {
 	}
 }
 
-impl<T: Send> DerefMut for SpinlockGuard<'_, T> {
+impl<T> DerefMut for SpinlockGuard<'_, T> {
 	fn deref_mut(&mut self) -> &mut T {
 		// SAFETY: Same as for `Deref::deref`.
 		unsafe { &mut *self.lock.data.get() }
 	}
 }
 
-impl<T: Send> Drop for SpinlockGuard<'_, T> {
+impl<T> Drop for SpinlockGuard<'_, T> {
 	fn drop(&mut self) {
 		self.lock.unlock();
 	}

From 8b0a722f05db7ea098efd0d31f208b584e89e770 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Sun, 29 Oct 2023 10:12:19 +0100
Subject: [PATCH 31/33] Remove stale file

---
 polkadot/node/test/performance-test/Cargo.toml | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 delete mode 100644 polkadot/node/test/performance-test/Cargo.toml

diff --git a/polkadot/node/test/performance-test/Cargo.toml b/polkadot/node/test/performance-test/Cargo.toml
deleted file mode 100644
index e69de29bb2d1..000000000000

From 331ed0150821c82c57248862440563e283b92e72 Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Wed, 1 Nov 2023 18:07:36 +0100
Subject: [PATCH 32/33] Rename tracking fn

---
 polkadot/node/tracking-allocator/src/lib.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/polkadot/node/tracking-allocator/src/lib.rs b/polkadot/node/tracking-allocator/src/lib.rs
index a1bf21e0d5dc..d7d94f4e4f2c 100644
--- a/polkadot/node/tracking-allocator/src/lib.rs
+++ b/polkadot/node/tracking-allocator/src/lib.rs
@@ -135,7 +135,7 @@ impl TrackingAllocatorData {
 	}
 
 	#[inline]
-	fn track(mut guard: SpinlockGuard<Self>, alloc: isize) -> Option<SpinlockGuard<Self>> {
+	fn track_and_check_limits(mut guard: SpinlockGuard<Self>, alloc: isize) -> Option<SpinlockGuard<Self>> {
 		guard.current += alloc;
 		if guard.current > guard.peak {
 			guard.peak = guard.current;
@@ -196,7 +196,7 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 	#[inline]
 	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
 		let guard = ALLOCATOR_DATA.lock();
-		if let Some(guard) = TrackingAllocatorData::track(guard, layout.size() as isize) {
+		if let Some(guard) = TrackingAllocatorData::track_and_check_limits(guard, layout.size() as isize) {
 			fail_allocation(guard)
 		} else {
 			self.0.alloc(layout)
@@ -206,7 +206,7 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 	#[inline]
 	unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
 		let guard = ALLOCATOR_DATA.lock();
-		if let Some(guard) = TrackingAllocatorData::track(guard, layout.size() as isize) {
+		if let Some(guard) = TrackingAllocatorData::track_and_check_limits(guard, layout.size() as isize) {
 			fail_allocation(guard)
 		} else {
 			self.0.alloc_zeroed(layout)
@@ -216,7 +216,7 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 	#[inline]
 	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) -> () {
 		let guard = ALLOCATOR_DATA.lock();
-		TrackingAllocatorData::track(guard, -(layout.size() as isize));
+		TrackingAllocatorData::track_and_check_limits(guard, -(layout.size() as isize));
 		self.0.dealloc(ptr, layout)
 	}
 
@@ -224,7 +224,7 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 	unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
 		let guard = ALLOCATOR_DATA.lock();
 		if let Some(guard) =
-			TrackingAllocatorData::track(guard, (new_size as isize) - (layout.size() as isize))
+			TrackingAllocatorData::track_and_check_limits(guard, (new_size as isize) - (layout.size() as isize))
 		{
 			fail_allocation(guard)
 		} else {

From 10e74433ff0aa9c1c100aa9a35b8edc42cbb670a Mon Sep 17 00:00:00 2001
From: Dmitry Sinyavin <dmitry.sinyavin@parity.io>
Date: Wed, 1 Nov 2023 19:25:52 +0100
Subject: [PATCH 33/33] Fix clippy and tests

---
 Cargo.lock                                       | 1 +
 polkadot/node/core/pvf/prepare-worker/Cargo.toml | 1 +
 polkadot/node/core/pvf/tests/it/main.rs          | 4 ++--
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 65257710a508..e48227890f80 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -12278,6 +12278,7 @@ dependencies = [
  "rococo-runtime",
  "sc-executor-common",
  "sc-executor-wasmtime",
+ "sp-maybe-compressed-blob",
  "tikv-jemalloc-ctl",
  "tikv-jemallocator",
  "tracing-gum",
diff --git a/polkadot/node/core/pvf/prepare-worker/Cargo.toml b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
index 7adf3335fa24..dae88afc555d 100644
--- a/polkadot/node/core/pvf/prepare-worker/Cargo.toml
+++ b/polkadot/node/core/pvf/prepare-worker/Cargo.toml
@@ -38,6 +38,7 @@ jemalloc-allocator = [
 [dev-dependencies]
 criterion = { version = "0.4.0", default-features = false, features = ["cargo_bench_support"] }
 rococo-runtime = { path = "../../../../runtime/rococo" }
+sp-maybe-compressed-blob = { path = "../../../../../substrate/primitives/maybe-compressed-blob" }
 
 [[bench]]
 name = "prepare_rococo_runtime"
diff --git a/polkadot/node/core/pvf/tests/it/main.rs b/polkadot/node/core/pvf/tests/it/main.rs
index 9523e02bd5ca..f4fd7f802f5e 100644
--- a/polkadot/node/core/pvf/tests/it/main.rs
+++ b/polkadot/node/core/pvf/tests/it/main.rs
@@ -440,7 +440,7 @@ async fn deleting_prepared_artifact_does_not_dispute() {
 // memory consumption.
 #[tokio::test]
 async fn prechecking_within_memory_limits() {
-	let host = TestHost::new();
+	let host = TestHost::new().await;
 	let result = host
 		.precheck_pvf(
 			::adder::wasm_binary_unwrap(),
@@ -459,7 +459,7 @@ async fn prechecking_within_memory_limits() {
 async fn prechecking_out_of_memory() {
 	use polkadot_node_core_pvf::PrepareError;
 
-	let host = TestHost::new();
+	let host = TestHost::new().await;
 	let result = host
 		.precheck_pvf(
 			::adder::wasm_binary_unwrap(),