From c427dcf5bb271c53aaff3300400a25173fd8f128 Mon Sep 17 00:00:00 2001
From: Jeffrey Dallatezza <jeffreydallatezza@gmail.com>
Date: Wed, 5 Feb 2025 21:12:38 -0800
Subject: [PATCH] Add jemalloc stats

---
 Cargo.lock                            |  1 +
 crates/core/Cargo.toml                |  1 +
 crates/core/src/worker_metrics/mod.rs | 39 +++++++++++++++++++++++++--
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6fd8727882..add6b550a5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4964,6 +4964,7 @@ dependencies = [
  "tempfile",
  "thin-vec",
  "thiserror",
+ "tikv-jemalloc-ctl",
  "tokio",
  "tokio-stream",
  "tokio-util",
diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml
index c9ef401248..485112e062 100644
--- a/crates/core/Cargo.toml
+++ b/crates/core/Cargo.toml
@@ -86,6 +86,7 @@ sha1.workspace = true
 similar.workspace = true
 slab.workspace = true
 sled.workspace = true
+tikv-jemalloc-ctl.workspace = true
 smallvec.workspace = true
 sqlparser.workspace = true
 strum.workspace = true
diff --git a/crates/core/src/worker_metrics/mod.rs b/crates/core/src/worker_metrics/mod.rs
index 2f7359d287..1f8104269c 100644
--- a/crates/core/src/worker_metrics/mod.rs
+++ b/crates/core/src/worker_metrics/mod.rs
@@ -1,4 +1,4 @@
-use std::thread::available_parallelism;
+use std::{thread::available_parallelism, time::Duration};
 
 use crate::execution_context::WorkloadType;
 use crate::hash::Hash;
@@ -112,10 +112,46 @@ metrics_group!(
         #[help = "The cumulative number of bytes sent to clients"]
         #[labels(txn_type: WorkloadType, db: Identity)]
         pub bytes_sent_to_clients: IntCounterVec,
+
+        #[name = jemalloc_active_bytes]
+        #[help = "Number of bytes in jemallocs heap"]
+        #[labels(db: Identity)]
+        pub jemalloc_active_bytes: IntGaugeVec,
+
+        #[name = jemalloc_allocated_bytes]
+        #[help = "Number of bytes in use"]
+        #[labels(db: Identity)]
+        pub jemalloc_allocated_bytes: IntGaugeVec,
+
+        #[name = jemalloc_resident_bytes]
+        #[help = "Total memory used by jemalloc"]
+        #[labels(db: Identity)]
+        pub jemalloc_resident_bytes: IntGaugeVec,
     }
 );
 
 pub static WORKER_METRICS: Lazy<WorkerMetrics> = Lazy::new(WorkerMetrics::new);
+use tikv_jemalloc_ctl::{epoch, stats};
+use tokio::{spawn, time::sleep};
+pub fn spawn_jemalloc_stats() {
+
+    spawn(async {
+        let e = epoch::mib().unwrap();
+        loop {
+            e.advance().unwrap();
+            let allocated = stats::allocated::read().unwrap();
+            WORKER_METRICS.jemalloc_active_bytes.with_label_values(&Identity::ZERO).set(allocated as i64);
+            let resident = stats::resident::read().unwrap();
+            WORKER_METRICS.jemalloc_resident_bytes.with_label_values(&Identity::ZERO).set(resident as i64);
+            let active = stats::active::read().unwrap();
+            WORKER_METRICS.jemalloc_active_bytes.with_label_values(&Identity::ZERO).set(active as i64);
+
+            // Make a gauge for each of these.
+
+            sleep(Duration::from_secs(10)).await;
+        }
+    });
+}
 
 fn build_global_metrics() -> Vec<Box<dyn Collector>> {
     let mut metrics: Vec<Box<dyn Collector>> = Vec::new();
@@ -127,7 +163,6 @@ fn build_global_metrics() -> Vec<Box<dyn Collector>> {
         )
         .unwrap(),
     ));
-
     metrics.push(Box::new(
         PullingGauge::new(
             "physical_memory_bytes",