Use jemalloc in standalone and add stats.

Cargo.toml

@@ -273,6 +273,8 @@ wasmbin = "0.6"
 webbrowser = "1.0.2"
 windows-sys = "0.59"
 xdg = "2.5"
+tikv-jemallocator = { version = "0.6.0", features = ["profiling", "stats"] }
+tikv-jemalloc-ctl = { version = "0.6.0", features = ["stats"]}
 
 # Vendor the openssl we rely on, rather than depend on a
 # potentially very old system version.

crates/bench/Cargo.toml

@@ -74,6 +74,10 @@ tracing-subscriber.workspace = true
 walkdir.workspace = true
 itertools.workspace = true
 
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = { workspace = true }
+tikv-jemalloc-ctl = { workspace = true }
+
 [target.'cfg(target_os = "linux")'.dependencies]
 # only try to build these on linux
 

crates/bench/benches/generic.rs

@@ -4,7 +4,6 @@ use criterion::{
     Bencher, BenchmarkGroup, Criterion,
 };
 use lazy_static::lazy_static;
-use mimalloc::MiMalloc;
 use spacetimedb_bench::{
     database::BenchDatabase,
     schemas::{create_sequential, u32_u64_str, u32_u64_u64, BenchTable, IndexStrategy, RandomTable},
@@ -14,9 +13,17 @@ use spacetimedb_lib::sats::AlgebraicType;
 use spacetimedb_primitives::ColId;
 use spacetimedb_testing::modules::{Csharp, Rust};
 
+#[cfg(target_env = "msvc")]
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
 lazy_static! {
     static ref RUN_ONE_MILLION: bool = std::env::var("RUN_ONE_MILLION").is_ok();
 }

crates/bench/benches/special.rs

@@ -1,6 +1,5 @@
 use criterion::async_executor::AsyncExecutor;
 use criterion::{criterion_group, criterion_main, Criterion, SamplingMode};
-use mimalloc::MiMalloc;
 use spacetimedb_bench::{
     database::BenchDatabase,
     schemas::{create_sequential, u32_u64_str, u32_u64_u64, u64_u64_u32, BenchTable, RandomTable},
@@ -13,9 +12,17 @@ use spacetimedb_testing::modules::{Csharp, ModuleLanguage, Rust};
 use std::sync::Arc;
 use std::sync::OnceLock;
 
+#[cfg(target_env = "msvc")]
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
 fn criterion_benchmark(c: &mut Criterion) {
     serialize_benchmarks::<u32_u64_str>(c);
     serialize_benchmarks::<u32_u64_u64>(c);

crates/cli/Cargo.toml

@@ -75,6 +75,10 @@ wasmtime.workspace = true
 webbrowser.workspace = true
 clap-markdown.workspace = true
 
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = { workspace = true }
+tikv-jemalloc-ctl = { workspace = true }
+
 [target.'cfg(windows)'.dependencies]
 windows-sys = { workspace = true, features = ["Win32_System_Console"] }
 

crates/cli/src/main.rs

@@ -1,11 +1,23 @@
 use std::process::ExitCode;
 
 use clap::{Arg, Command};
-use mimalloc::MiMalloc;
 use spacetimedb_cli::*;
 use spacetimedb_paths::cli::CliTomlPath;
 use spacetimedb_paths::{RootDir, SpacetimePaths};
 
+// Note that the standalone server is invoked through standaline/src/main.rs, so you will
+// also want to set the allocator there.
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
+#[cfg(target_env = "msvc")]
+use mimalloc::MiMalloc;
+
+#[cfg(target_env = "msvc")]
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 

crates/core/Cargo.toml

@@ -111,6 +111,10 @@ jwks = { package = "spacetimedb-jwks", version = "0.1.3" }
 async_cache = "0.3.1"
 faststr = "0.2.23"
 
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = {workspace = true}
+tikv-jemalloc-ctl = {workspace = true}
+
 [features]
 # Print a warning when doing an unindexed `iter_by_col_range` on a large table.
 unindexed_iter_by_col_range_warn = []

crates/core/src/worker_metrics/mod.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use crate::execution_context::WorkloadType;
 use crate::hash::Hash;
 use once_cell::sync::Lazy;
@@ -22,6 +24,21 @@ metrics_group!(
         #[labels(replica_id: u64, protocol: str)]
         pub websocket_request_msg_size: HistogramVec,
 
+        #[name = jemalloc_active_bytes]
+        #[help = "Number of bytes in jemallocs heap"]
+        #[labels(node_id: str)]
+        pub jemalloc_active_bytes: IntGaugeVec,
+
+        #[name = jemalloc_allocated_bytes]
+        #[help = "Number of bytes in use by the application"]
+        #[labels(node_id: str)]
+        pub jemalloc_allocated_bytes: IntGaugeVec,
+
+        #[name = jemalloc_resident_bytes]
+        #[help = "Total memory used by jemalloc"]
+        #[labels(node_id: str)]
+        pub jemalloc_resident_bytes: IntGaugeVec,
+
         #[name = spacetime_websocket_sent_msg_size_bytes]
         #[help = "The size of messages sent to connected sessions"]
         #[labels(db: Identity, workload: WorkloadType)]
@@ -113,3 +130,38 @@ metrics_group!(
 );
 
 pub static WORKER_METRICS: Lazy<WorkerMetrics> = Lazy::new(WorkerMetrics::new);
+
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemalloc_ctl::{epoch, stats};
+
+use std::sync::Once;
+use tokio::{spawn, time::sleep};
+static SPAWN_JEMALLOC_GUARD: Once = Once::new();
+pub fn spawn_jemalloc_stats(node_id: String) {
+    #[cfg(not(target_env = "msvc"))]
+    SPAWN_JEMALLOC_GUARD.call_once(|| {
+        spawn(async move {
+            let e = epoch::mib().unwrap();
+            loop {
+                e.advance().unwrap();
+                let allocated = stats::allocated::read().unwrap();
+                WORKER_METRICS
+                    .jemalloc_allocated_bytes
+                    .with_label_values(&node_id)
+                    .set(allocated as i64);
+                let resident = stats::resident::read().unwrap();
+                WORKER_METRICS
+                    .jemalloc_resident_bytes
+                    .with_label_values(&node_id)
+                    .set(resident as i64);
+                let active = stats::active::read().unwrap();
+                WORKER_METRICS
+                    .jemalloc_active_bytes
+                    .with_label_values(&node_id)
+                    .set(active as i64);
+
+                sleep(Duration::from_secs(10)).await;
+            }
+        });
+    });
+}

crates/standalone/Cargo.toml

@@ -43,6 +43,10 @@ tower-http.workspace = true
 toml.workspace = true
 tracing = { workspace = true, features = ["release_max_level_debug"] }
 
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = {workspace = true}
+tikv-jemalloc-ctl = {workspace = true}
+
 [dev-dependencies]
 once_cell.workspace = true
 tempfile.workspace = true

crates/standalone/src/main.rs

@@ -39,6 +39,13 @@ Example usage:
         )
 }
 
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
 fn main() -> anyhow::Result<()> {
     // take_hook() returns the default hook in case when a custom one is not set
     let orig_hook = panic::take_hook();

crates/standalone/src/subcommands/start.rs

@@ -8,6 +8,7 @@ use clap::{Arg, ArgMatches};
 use spacetimedb::config::{CertificateAuthority, ConfigFile};
 use spacetimedb::db::{Config, Storage};
 use spacetimedb::startup::{self, TracingOptions};
+use spacetimedb::worker_metrics;
 use spacetimedb_client_api::routes::database::DatabaseRoutes;
 use spacetimedb_client_api::routes::router;
 use spacetimedb_paths::cli::{PrivKeyPath, PubKeyPath};
@@ -133,6 +134,7 @@ pub async fn exec(args: &ArgMatches) -> anyhow::Result<()> {
 
     let data_dir = Arc::new(data_dir.clone());
     let ctx = StandaloneEnv::init(db_config, &certs, data_dir).await?;
+    worker_metrics::spawn_jemalloc_stats(listen_addr.clone());
 
     let mut db_routes = DatabaseRoutes::default();
     db_routes.root_post = db_routes.root_post.layer(DefaultBodyLimit::disable());