Move memory stats gathering from polkadot to parity-util-mem (#588)

* Move memory stats gathering from `polkadot` to `parity-util-mem`

* Bump version to 0.10.1

* Update the CHANGELOG

parity-util-mem/CHANGELOG.md

@@ -4,7 +4,10 @@ The format is based on [Keep a Changelog].
 
 [Keep a Changelog]: http://keepachangelog.com/en/1.0.0/
 
-## [Unreleased]
+## [0.10.1] - 2021-09-15
+- Added support for memory stats gathering, ported over from `polkadot`. [#588](https://github.com/paritytech/parity-common/pull/588)
+
+## [0.10.0] - 2021-07-02
 - Fixed `malloc_usable_size` for FreeBSD. [#553](https://github.com/paritytech/parity-common/pull/553)
 
 ### Breaking

parity-util-mem/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "parity-util-mem"
-version = "0.10.0"
+version = "0.10.1"
 authors = ["Parity Technologies <admin@parity.io>"]
 repository = "https://github.com/paritytech/parity-common"
 description = "Collection of memory related utilities"
@@ -37,6 +37,10 @@ winapi = { version = "0.3.8", features = ["heapapi"] }
 version = "0.3.2"
 optional = true
 
+[target.'cfg(not(target_os = "windows"))'.dependencies.jemalloc-ctl]
+version = "0.3.3"
+optional = true
+
 [features]
 default = ["std", "ethereum-impls", "lru", "hashbrown", "smallvec", "primitive-types"]
 std = ["parking_lot"]
@@ -45,7 +49,7 @@ dlmalloc-global = ["dlmalloc", "estimate-heapsize"]
 # use wee_alloc as global allocator
 weealloc-global = ["wee_alloc", "estimate-heapsize"]
 # use jemalloc as global allocator
-jemalloc-global = ["jemallocator"]
+jemalloc-global = ["jemallocator", "jemalloc-ctl"]
 # use mimalloc as global allocator
 mimalloc-global = ["mimalloc", "libmimalloc-sys"]
 # implement additional types

parity-util-mem/src/lib.rs

@@ -24,23 +24,37 @@ cfg_if::cfg_if! {
 		/// Global allocator
 		#[global_allocator]
 		pub static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
+		mod memory_stats_jemalloc;
+		use memory_stats_jemalloc as memory_stats;
 	} else if #[cfg(feature = "dlmalloc-global")] {
 		/// Global allocator
 		#[global_allocator]
 		pub static ALLOC: dlmalloc::GlobalDlmalloc = dlmalloc::GlobalDlmalloc;
+
+		mod memory_stats_noop;
+		use memory_stats_noop as memory_stats;
 	} else if #[cfg(feature = "weealloc-global")] {
 		/// Global allocator
 		#[global_allocator]
 		pub static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
+
+		mod memory_stats_noop;
+		use memory_stats_noop as memory_stats;
 	} else if #[cfg(all(
 			feature = "mimalloc-global",
 			not(target_arch = "wasm32")
 		))] {
 		/// Global allocator
 		#[global_allocator]
 		pub static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+		mod memory_stats_noop;
+		use memory_stats_noop as memory_stats;
 	} else {
 		// default allocator used
+		mod memory_stats_noop;
+		use memory_stats_noop as memory_stats;
 	}
 }
 
@@ -78,6 +92,48 @@ pub fn malloc_size<T: MallocSizeOf + ?Sized>(t: &T) -> usize {
 	MallocSizeOf::size_of(t, &mut allocators::new_malloc_size_ops())
 }
 
+/// An error related to the memory stats gathering.
+#[derive(Clone, Debug)]
+pub struct MemoryStatsError(memory_stats::Error);
+
+#[cfg(feature = "std")]
+impl std::fmt::Display for MemoryStatsError {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		self.0.fmt(fmt)
+	}
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for MemoryStatsError {}
+
+/// 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,
+}
+
+/// Accessor to the allocator internals.
+#[derive(Clone)]
+pub struct MemoryAllocationTracker(self::memory_stats::MemoryAllocationTracker);
+
+impl MemoryAllocationTracker {
+	/// Create an instance of an allocation tracker.
+	pub fn new() -> Result<Self, MemoryStatsError> {
+		self::memory_stats::MemoryAllocationTracker::new()
+			.map(MemoryAllocationTracker)
+			.map_err(MemoryStatsError)
+	}
+
+	/// Create an allocation snapshot.
+	pub fn snapshot(&self) -> Result<MemoryAllocationSnapshot, MemoryStatsError> {
+		self.0.snapshot().map_err(MemoryStatsError)
+	}
+}
+
 #[cfg(feature = "std")]
 #[cfg(test)]
 mod test {

parity-util-mem/src/memory_stats_jemalloc.rs

@@ -0,0 +1,32 @@
+// Copyright 2021 Parity Technologies
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+pub use jemalloc_ctl::Error;
+use jemalloc_ctl::{epoch, stats};
+
+#[derive(Clone)]
+pub struct MemoryAllocationTracker {
+	epoch: 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<crate::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(crate::MemoryAllocationSnapshot { allocated, resident })
+	}
+}

parity-util-mem/src/memory_stats_noop.rs

@@ -0,0 +1,31 @@
+// Copyright 2021 Parity Technologies
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#[derive(Clone, Debug)]
+pub struct Unimplemented;
+pub use Unimplemented as Error;
+
+#[cfg(feature = "std")]
+impl std::fmt::Display for Unimplemented {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		fmt.write_str("unimplemented")
+	}
+}
+
+#[derive(Clone)]
+pub struct MemoryAllocationTracker {}
+
+impl MemoryAllocationTracker {
+	pub fn new() -> Result<Self, Error> {
+		Err(Error)
+	}
+
+	pub fn snapshot(&self) -> Result<crate::MemoryAllocationSnapshot, Error> {
+		unimplemented!();
+	}
+}