From 6f0e1e7feb3f874e7cfb674819a7f91e15af8fcf Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Sat, 24 Aug 2024 21:53:24 -0700
Subject: [PATCH] switch to jemalloc and add memory stats

---
 Cargo.lock         | 39 +++++++++++++++++++++++++++++++++++++++
 server/Cargo.toml  |  2 ++
 server/src/main.rs | 41 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1434e61..b6e0b7d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -573,6 +573,8 @@ dependencies = [
  "serde_json",
  "serde_yaml",
  "shell-words",
+ "tikv-jemalloc-ctl",
+ "tikv-jemallocator",
  "tokio",
  "tokio-util",
  "toml",
@@ -1205,6 +1207,12 @@ dependencies = [
  "windows-targets",
 ]
 
+[[package]]
+name = "paste"
+version = "1.0.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
+
 [[package]]
 name = "percent-encoding"
 version = "2.3.1"
@@ -1726,6 +1734,37 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "tikv-jemalloc-ctl"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f21f216790c8df74ce3ab25b534e0718da5a1916719771d3fec23315c99e468b"
+dependencies = [
+ "libc",
+ "paste",
+ "tikv-jemalloc-sys",
+]
+
+[[package]]
+name = "tikv-jemalloc-sys"
+version = "0.6.0+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd3c60906412afa9c2b5b5a48ca6a5abe5736aec9eb48ad05037a677e52e4e2d"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "tikv-jemallocator"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4cec5ff18518d81584f477e9bfdf957f5bb0979b0bac3af4ca30b5b3ae2d2865"
+dependencies = [
+ "libc",
+ "tikv-jemalloc-sys",
+]
+
 [[package]]
 name = "tokio"
 version = "1.39.3"
diff --git a/server/Cargo.toml b/server/Cargo.toml
index 1c2b00f..76dfd83 100644
--- a/server/Cargo.toml
+++ b/server/Cargo.toml
@@ -27,6 +27,8 @@ serde = { version = "1.0.208", features = ["derive"] }
 serde_json = { version = "1.0.125", optional = true }
 serde_yaml = { version = "0.9.34", optional = true }
 shell-words = { version = "1.1.0", optional = true }
+tikv-jemalloc-ctl = { version = "0.6.0", features = ["stats", "use_std"] }
+tikv-jemallocator = "0.6.0"
 tokio = { version = "1.39.3", features = ["full"] }
 tokio-util = { version = "0.7.11", features = ["codec", "compat", "io-util", "net"] }
 toml = { version = "0.8.19", optional = true }
diff --git a/server/src/main.rs b/server/src/main.rs
index 9df0430..c497adc 100644
--- a/server/src/main.rs
+++ b/server/src/main.rs
@@ -40,9 +40,45 @@ fn format_stream_type(stream_type: StreamType) -> &'static str {
 	}
 }
 
-fn generate_stats() -> Result<String, std::fmt::Error> {
+fn generate_stats() -> anyhow::Result<String> {
 	let mut out = String::new();
 	let len = CLIENTS.len();
+
+	use tikv_jemalloc_ctl::stats::{active, allocated, mapped, metadata, resident, retained};
+	tikv_jemalloc_ctl::epoch::advance()?;
+
+	writeln!(&mut out, "Memory usage:")?;
+	writeln!(
+		&mut out,
+		"\tActive: {:?} MiB",
+		active::read()? as f64 / (1024 * 1024) as f64
+	)?;
+	writeln!(
+		&mut out,
+		"\tAllocated: {:?} MiB",
+		allocated::read()? as f64 / (1024 * 1024) as f64
+	)?;
+	writeln!(
+		&mut out,
+		"\tMapped: {:?} MiB",
+		mapped::read()? as f64 / (1024 * 1024) as f64
+	)?;
+	writeln!(
+		&mut out,
+		"\tMetadata: {:?} MiB",
+		metadata::read()? as f64 / (1024 * 1024) as f64
+	)?;
+	writeln!(
+		&mut out,
+		"\tResident: {:?} MiB",
+		resident::read()? as f64 / (1024 * 1024) as f64
+	)?;
+	writeln!(
+		&mut out,
+		"\tRetained: {:?} MiB",
+		retained::read()? as f64 / (1024 * 1024) as f64
+	)?;
+
 	writeln!(
 		&mut out,
 		"{} clients connected{}",
@@ -108,6 +144,9 @@ fn handle_stream(stream: ServerRouteResult, id: String) {
 	});
 }
 
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 #[tokio::main(flavor = "multi_thread")]
 async fn main() -> anyhow::Result<()> {
 	if CLI.default_config {