[DNM] Wrapper allocator PoC

Cargo.toml

@@ -26,6 +26,7 @@ tikv-jemallocator = "0.5.0"
 polkadot-cli = { path = "cli", features = [ "kusama-native", "westend-native", "rococo-native" ]  }
 polkadot-node-core-pvf-worker = { path = "node/core/pvf/worker" }
 polkadot-overseer = { path = "node/overseer" }
+wrapper-allocator = { path = "node/wrapper-allocator", optional = true }
 
 [dev-dependencies]
 assert_cmd = "2.0.4"
@@ -101,6 +102,7 @@ members = [
 	"node/subsystem-types",
 	"node/subsystem-test-helpers",
 	"node/subsystem-util",
+	"node/wrapper-allocator",
 	"node/jaeger",
 	"node/gum",
 	"node/gum/proc-macro",
@@ -208,6 +210,7 @@ fast-runtime = [ "polkadot-cli/fast-runtime" ]
 runtime-metrics = [ "polkadot-cli/runtime-metrics" ]
 pyroscope = ["polkadot-cli/pyroscope"]
 jemalloc-allocator = ["polkadot-node-core-pvf-worker/jemalloc-allocator", "polkadot-overseer/jemalloc-allocator"]
+wrapper-allocator = ["jemalloc-allocator", "dep:wrapper-allocator", "polkadot-node-core-pvf-worker/wrapper-allocator"]
 
 # Configuration for building a .deb package - for use with `cargo-deb`
 [package.metadata.deb]

node/core/pvf/worker/Cargo.toml

@@ -18,6 +18,7 @@ rayon = "1.5.1"
 tempfile = "3.3.0"
 tikv-jemalloc-ctl = { version = "0.5.0", optional = true }
 tokio = "1.24.2"
+wrapper-allocator = { path = "../../../wrapper-allocator", optional = true }
 
 parity-scale-codec = { version = "3.4.0", default-features = false, features = ["derive"] }
 
@@ -47,3 +48,4 @@ tempfile = "3.3.0"
 
 [features]
 jemalloc-allocator = ["dep:tikv-jemalloc-ctl"]
+wrapper-allocator = ["dep:wrapper-allocator"]

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

@@ -18,6 +18,8 @@
 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};
+#[cfg(feature = "wrapper-allocator")]
+use wrapper_allocator::ALLOCATOR_DATA;
 use crate::{
 	common::{bytes_to_path, cpu_time_monitor_loop, worker_event_loop},
 	prepare, prevalidate, LOG_TARGET,
@@ -109,8 +111,22 @@ pub fn worker_entrypoint(socket_path: &str, node_version: Option<&str>) {
 			// Spawn another thread for preparation.
 			let prepare_fut = rt_handle
 				.spawn_blocking(move || {
+					#[cfg(feature = "wrapper-allocator")]
+					ALLOCATOR_DATA.checkpoint();
+
 					let result = prepare_artifact(pvf);
 
+					#[cfg(feature = "wrapper-allocator")]
+					{
+						let peak = ALLOCATOR_DATA.checkpoint();
+						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 result = result.map(|artifact| (artifact, get_max_rss_thread()));

node/wrapper-allocator/Cargo.toml

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

node/wrapper-allocator/src/lib.rs

@@ -0,0 +1,78 @@
+// 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 core::sync::atomic::{ Ordering::SeqCst, AtomicUsize };
+use core::alloc::{GlobalAlloc, Layout};
+use tikv_jemallocator::Jemalloc;
+
+pub struct WrapperAllocatorData {
+    allocated: AtomicUsize,
+    checkpoint: AtomicUsize,
+    peak: AtomicUsize,
+    // limit: AtomicUsize, // Should we introduce a checkpoint limit and fail allocation if the limit is hit?
+}
+
+impl WrapperAllocatorData {
+    /// Marks a new checkpoint. Returns peak allocation, in bytes, since the last checkpoint.
+    pub fn checkpoint(&self) -> usize {
+        let allocated = ALLOCATOR_DATA.allocated.load(SeqCst);
+        let old_cp = ALLOCATOR_DATA.checkpoint.swap(allocated, SeqCst);
+        ALLOCATOR_DATA.peak.swap(allocated, SeqCst).saturating_sub(old_cp)
+    }
+}
+
+pub static ALLOCATOR_DATA: WrapperAllocatorData = WrapperAllocatorData { allocated: AtomicUsize::new(0), checkpoint: AtomicUsize::new(0), peak: AtomicUsize::new(0) };
+
+struct WrapperAllocator<A: GlobalAlloc>(A);
+
+unsafe impl<A: GlobalAlloc> GlobalAlloc for WrapperAllocator<A> {
+
+    // SAFETY: The wrapped methods are as safe as the underlying allocator implementation is.
+
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let old_alloc = ALLOCATOR_DATA.allocated.fetch_add(layout.size(), SeqCst);
+        ALLOCATOR_DATA.peak.fetch_max(old_alloc + layout.size(), SeqCst);
+        self.0.alloc(layout)
+    }
+
+    #[inline]
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let old_alloc = ALLOCATOR_DATA.allocated.fetch_add(layout.size(), SeqCst);
+        ALLOCATOR_DATA.peak.fetch_max(old_alloc + layout.size(), SeqCst);
+        self.0.alloc_zeroed(layout)
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) -> () {
+        ALLOCATOR_DATA.allocated.fetch_sub(layout.size(), SeqCst);
+        self.0.dealloc(ptr, layout)
+    }
+
+    #[inline]
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        if new_size > layout.size() {
+            let old_alloc = ALLOCATOR_DATA.allocated.fetch_add(new_size - layout.size(), SeqCst);
+            ALLOCATOR_DATA.peak.fetch_max(old_alloc + new_size - layout.size(), SeqCst);
+        } else {
+            ALLOCATOR_DATA.allocated.fetch_sub(layout.size() - new_size, SeqCst);
+        }
+        self.0.realloc(ptr, layout, new_size)
+    }
+}
+
+#[global_allocator]
+static ALLOC: WrapperAllocator<Jemalloc> = WrapperAllocator(Jemalloc);

src/main.rs

@@ -22,7 +22,7 @@ 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;