Support jemalloc profiling and heap dumps (#3373)

apps/labrinth/Cargo.toml

@@ -125,7 +125,9 @@ lettre = "0.11.3"
 derive-new = "0.6.0"
 rust_iso3166 = "0.1.11"
 
-jemallocator = { version = "0.5.4", optional = true }
+tikv-jemallocator = { version = "0.6.0", features = ["profiling", "unprefixed_malloc_on_supported_platforms"] }
+tikv-jemalloc-ctl = { version = "0.6.0", features = ["stats"] }
+jemalloc_pprof = { version = "0.7.0", features = ["flamegraph"] }
 
 async-stripe = { version = "0.39.1", features = ["runtime-tokio-hyper-rustls"] }
 rusty-money = "0.4.1"
@@ -135,6 +137,3 @@ ariadne = { path = "../../packages/ariadne" }
 
 [dev-dependencies]
 actix-http = "3.4.0"
-
-[features]
-jemalloc = ["jemallocator"]

apps/labrinth/src/lib.rs

@@ -352,6 +352,7 @@ pub fn app_config(
     .app_data(labrinth_config.active_sockets.clone())
     .app_data(labrinth_config.automated_moderation_queue.clone())
     .app_data(web::Data::new(labrinth_config.stripe_client.clone()))
+    .configure(routes::debug::config)
     .configure(routes::v2::config)
     .configure(routes::v3::config)
     .configure(routes::internal::config)

apps/labrinth/src/main.rs

@@ -9,9 +9,14 @@ use std::sync::Arc;
 use tracing::{error, info};
 use tracing_actix_web::TracingLogger;
 
-#[cfg(feature = "jemalloc")]
+#[cfg(not(target_env = "msvc"))]
 #[global_allocator]
-static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+#[allow(non_upper_case_globals)]
+#[export_name = "malloc_conf"]
+pub static malloc_conf: &[u8] =
+    b"prof:true,prof_active:true,lg_prof_sample:19\0";
 
 #[derive(Clone)]
 pub struct Pepper {
@@ -105,6 +110,8 @@ async fn main() -> std::io::Result<()> {
         .register_and_set_metrics(&prometheus.registry)
         .await
         .expect("Failed to register redis metrics");
+    labrinth::routes::debug::jemalloc_mmeory_stats(&prometheus.registry)
+        .expect("Failed to register jemalloc metrics");
 
     let search_config = search::SearchConfig::new(None);
 

apps/labrinth/src/routes/debug/mod.rs

@@ -0,0 +1,87 @@
+use crate::routes::ApiError;
+use crate::util::cors::default_cors;
+use crate::util::guards::admin_key_guard;
+use actix_web::{get, HttpResponse};
+use prometheus::{IntGauge, Registry};
+use std::time::Duration;
+
+pub fn config(cfg: &mut actix_web::web::ServiceConfig) {
+    cfg.service(
+        actix_web::web::scope("/debug")
+            .wrap(default_cors())
+            .service(heap)
+            .service(flame_graph),
+    );
+}
+
+#[get("pprof/heap", guard = "admin_key_guard")]
+pub async fn heap() -> Result<HttpResponse, ApiError> {
+    let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await;
+    require_profiling_activated(&prof_ctl)?;
+    let pprof = prof_ctl
+        .dump_pprof()
+        .map_err(|err| ApiError::InvalidInput(err.to_string()))?;
+
+    Ok(HttpResponse::Ok()
+        .content_type("application/octet-stream")
+        .body(pprof))
+}
+
+#[get("pprof/heap/flamegraph", guard = "admin_key_guard")]
+pub async fn flame_graph() -> Result<HttpResponse, ApiError> {
+    let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await;
+    require_profiling_activated(&prof_ctl)?;
+    let svg = prof_ctl
+        .dump_flamegraph()
+        .map_err(|err| ApiError::InvalidInput(err.to_string()))?;
+
+    Ok(HttpResponse::Ok().content_type("image/svg+xml").body(svg))
+}
+
+fn require_profiling_activated(
+    prof_ctl: &jemalloc_pprof::JemallocProfCtl,
+) -> Result<(), ApiError> {
+    if prof_ctl.activated() {
+        Ok(())
+    } else {
+        Err(ApiError::InvalidInput(
+            "Profiling is not activated".to_string(),
+        ))
+    }
+}
+
+pub fn jemalloc_mmeory_stats(
+    registry: &Registry,
+) -> Result<(), prometheus::Error> {
+    let allocated_mem = IntGauge::new(
+        "labrinth_memory_allocated",
+        "Labrinth allocated memory",
+    )?;
+    let resident_mem =
+        IntGauge::new("labrinth_resident_memory", "Labrinth resident memory")?;
+
+    registry.register(Box::new(allocated_mem.clone()))?;
+    registry.register(Box::new(resident_mem.clone()))?;
+
+    tokio::spawn(async move {
+        let e = tikv_jemalloc_ctl::epoch::mib().unwrap();
+        let allocated = tikv_jemalloc_ctl::stats::allocated::mib().unwrap();
+        let resident = tikv_jemalloc_ctl::stats::resident::mib().unwrap();
+
+        loop {
+            e.advance().unwrap();
+
+            if let Ok(allocated) = allocated.read() {
+                allocated_mem.set(allocated as i64);
+            }
+
+            if let Ok(resident) = resident.read() {
+                resident_mem.set(resident as i64);
+            }
+
+            tokio::time::sleep(Duration::from_secs(5)).await;
+        }
+    });
+
+    Ok(())
+}

apps/labrinth/src/routes/mod.rs

@@ -12,6 +12,8 @@ pub mod internal;
 pub mod v2;
 pub mod v3;
 
+pub mod debug;
+
 pub mod v2_reroute;
 
 mod analytics;