Add back jemalloc

Cargo.toml

@@ -15,6 +15,7 @@ readme = "README.md"
 [dependencies]
 polkadot-cli = { path = "cli", features = [ "kusama-native", "westend-native", "rococo-native" ]  }
 color-eyre = { version = "0.6.1", default-features = false }
+tikv-jemallocator = "0.5.0"
 
 [dev-dependencies]
 assert_cmd = "2.0.4"

node/overseer/Cargo.toml

@@ -20,6 +20,7 @@ gum = { package = "tracing-gum", path = "../gum" }
 lru = "0.8"
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "master" }
 async-trait = "0.1.57"
+tikv-jemalloc-ctl = "0.5.0"
 
 [dev-dependencies]
 metered = { package = "prioritized-metered-channel", version = "0.2.0" }

node/overseer/src/lib.rs

@@ -116,11 +116,14 @@ pub const KNOWN_LEAVES_CACHE_SIZE: NonZeroUsize = match NonZeroUsize::new(2 * 24
 	None => panic!("Known leaves cache size must be non-zero"),
 };
 
+mod memory_stats;
 #[cfg(test)]
 mod tests;
 
 use sp_core::traits::SpawnNamed;
 
+use memory_stats::MemoryAllocationTracker;
+
 /// Glue to connect `trait orchestra::Spawner` and `SpawnNamed` from `substrate`.
 pub struct SpawnGlue<S>(pub S);
 
@@ -649,7 +652,34 @@ where
 	}
 	let subsystem_meters = overseer.map_subsystems(ExtractNameAndMeters);
 
+	let collect_memory_stats: Box<dyn Fn(&OverseerMetrics) + Send> =
+		match MemoryAllocationTracker::new() {
+			Ok(memory_stats) =>
+				Box::new(move |metrics: &OverseerMetrics| match memory_stats.snapshot() {
+					Ok(memory_stats_snapshot) => {
+						gum::trace!(
+							target: LOG_TARGET,
+							"memory_stats: {:?}",
+							&memory_stats_snapshot
+						);
+						metrics.memory_stats_snapshot(memory_stats_snapshot);
+					},
+					Err(e) =>
+						gum::debug!(target: LOG_TARGET, "Failed to obtain memory stats: {:?}", e),
+				}),
+			Err(_) => {
+				gum::debug!(
+					target: LOG_TARGET,
+					"Memory allocation tracking is not supported by the allocator.",
+				);
+
+				Box::new(|_| {})
+			},
+		};
+
 	let metronome = Metronome::new(std::time::Duration::from_millis(950)).for_each(move |_| {
+		collect_memory_stats(&metronome_metrics);
+
 		// We combine the amount of messages from subsystems to the overseer
 		// as well as the amount of messages from external sources to the overseer
 		// into one `to_overseer` value.

node/overseer/src/memory_stats.rs

@@ -0,0 +1,37 @@
+use tikv_jemalloc_ctl::{epoch, stats, Error};
+
+#[derive(Clone)]
+pub struct MemoryAllocationTracker {
+	epoch: tikv_jemalloc_ctl::epoch_mib,
+	allocated: stats::allocated_mib,
+	resident: stats::resident_mib,
+}
+
+impl MemoryAllocationTracker {
+	pub fn new() -> Result<Self, Error> {
+		Ok(Self {
+			epoch: epoch::mib()?,
+			allocated: stats::allocated::mib()?,
+			resident: stats::resident::mib()?,
+		})
+	}
+
+	pub fn snapshot(&self) -> Result<MemoryAllocationSnapshot, Error> {
+		// update stats by advancing the allocation epoch
+		self.epoch.advance()?;
+
+		let allocated: u64 = self.allocated.read()? as _;
+		let resident: u64 = self.resident.read()? as _;
+		Ok(MemoryAllocationSnapshot { allocated, resident })
+	}
+}
+
+/// Snapshot of collected memory metrics.
+#[non_exhaustive]
+#[derive(Debug, Clone)]
+pub struct MemoryAllocationSnapshot {
+	/// Total resident memory, in bytes.
+	pub resident: u64,
+	/// Total allocated memory, in bytes.
+	pub allocated: u64,
+}

node/overseer/src/metrics.rs

@@ -19,6 +19,8 @@
 use super::*;
 pub use polkadot_node_metrics::metrics::{self, prometheus, Metrics as MetricsTrait};
 
+use memory_stats::MemoryAllocationSnapshot;
+
 /// Overseer Prometheus metrics.
 #[derive(Clone)]
 struct MetricsInner {
@@ -65,6 +67,13 @@ impl Metrics {
 		}
 	}
 
+	pub(crate) fn memory_stats_snapshot(&self, memory_stats: MemoryAllocationSnapshot) {
+		if let Some(metrics) = &self.0 {
+			metrics.memory_stats_allocated.set(memory_stats.allocated);
+			metrics.memory_stats_resident.set(memory_stats.resident);
+		}
+	}
+
 	pub(crate) fn channel_metrics_snapshot(
 		&self,
 		collection: impl IntoIterator<Item = (&'static str, SubsystemMeterReadouts)>,

src/main.rs

@@ -20,6 +20,10 @@
 
 use color_eyre::eyre;
 
+/// Global allocator
+#[global_allocator]
+pub static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 fn main() -> eyre::Result<()> {
 	color_eyre::install()?;
 	polkadot_cli::run()?;