Move memory usage statistics to ra_prof

3 files changed, 72 insertions, 0 deletions

diff --git a/crates/ra_prof/Cargo.toml b/crates/ra_prof/Cargo.toml
index 787e18385..e986019ca 100644
--- a/crates/ra_prof/Cargo.toml
+++ b/crates/ra_prof/Cargo.toml
@@ -10,3 +10,9 @@ once_cell = "0.2.0"
 itertools = "0.8.0"
 backtrace = "0.3.28"
 cpuprofiler = { version = "0.0.3", optional = true }
+jemallocator = { version = "0.1.9", optional = true }
+jemalloc-ctl = { version = "0.2.0", optional = true }
+
+
+[features]
+jemalloc = [ "jemallocator", "jemalloc-ctl" ]
diff --git a/crates/ra_prof/src/lib.rs b/crates/ra_prof/src/lib.rs
index 1e8d780ab..6f7918745 100644
--- a/crates/ra_prof/src/lib.rs
+++ b/crates/ra_prof/src/lib.rs
@@ -1,3 +1,5 @@
+mod memory_usage;
+
 use std::{
     cell::RefCell,
     time::{Duration, Instant},
@@ -11,6 +13,14 @@ use std::{
 use once_cell::sync::Lazy;
 use itertools::Itertools;
 
+pub use crate::memory_usage::{MemoryUsage, Bytes};
+
+// We use jemalloc mainly to get heap usage statistics, actual performance
+// difference is not measures.
+#[cfg(feature = "jemalloc")]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
 /// Set profiling filter. It specifies descriptions allowed to profile.
 /// This is helpful when call stack has too many nested profiling scopes.
 /// Additionally filter can specify maximum depth of profiling scopes nesting.
@@ -288,6 +298,10 @@ impl Drop for CpuProfiler {
     }
 }
 
+pub fn memory_usage() -> MemoryUsage {
+    MemoryUsage::current()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/crates/ra_prof/src/memory_usage.rs b/crates/ra_prof/src/memory_usage.rs
new file mode 100644
index 000000000..2bde8fb5f
--- /dev/null
+++ b/crates/ra_prof/src/memory_usage.rs
@@ -0,0 +1,52 @@
+use std::fmt;
+
+pub struct MemoryUsage {
+    pub allocated: Bytes,
+    pub resident: Bytes,
+}
+
+impl MemoryUsage {
+    #[cfg(feature = "jemalloc")]
+    pub fn current() -> MemoryUsage {
+        jemalloc_ctl::epoch().unwrap();
+        MemoryUsage {
+            allocated: Bytes(jemalloc_ctl::stats::allocated().unwrap()),
+            resident: Bytes(jemalloc_ctl::stats::resident().unwrap()),
+        }
+    }
+
+    #[cfg(not(feature = "jemalloc"))]
+    pub fn current() -> MemoryUsage {
+        MemoryUsage { allocated: Bytes(0), resident: Bytes(0) }
+    }
+}
+
+impl fmt::Display for MemoryUsage {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        write!(fmt, "{} allocated {} resident", self.allocated, self.resident,)
+    }
+}
+
+#[derive(Default)]
+pub struct Bytes(usize);
+
+impl fmt::Display for Bytes {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let bytes = self.0;
+        if bytes < 4096 {
+            return write!(f, "{} bytes", bytes);
+        }
+        let kb = bytes / 1024;
+        if kb < 4096 {
+            return write!(f, "{}kb", kb);
+        }
+        let mb = kb / 1024;
+        write!(f, "{}mb", mb)
+    }
+}
+
+impl std::ops::AddAssign<usize> for Bytes {
+    fn add_assign(&mut self, x: usize) {
+        self.0 += x;
+    }
+}