fix: use jemalloc everywhere we're rusty (#24796)

rust/Cargo.toml

@@ -7,6 +7,7 @@ members = [
   "common/health",
   "common/metrics",
   "common/dns",
+  "common/alloc",
   "feature-flags",
   "hook-api",
   "hook-common",

rust/capture/Cargo.toml

@@ -7,10 +7,6 @@ edition = "2021"
 workspace = true
 
 [dependencies]
-
-[target.'cfg(not(target_env = "msvc"))'.dependencies]
-tikv-jemallocator = "0.6"
-
 anyhow = { workspace = true }
 async-trait = { workspace = true }
 axum = { workspace = true }
@@ -21,6 +17,7 @@ envconfig = { workspace = true }
 flate2 = { workspace = true }
 governor = { workspace = true }
 health = { path = "../common/health" }
+common-alloc = { path = "../common/alloc" }
 metrics = { workspace = true }
 metrics-exporter-prometheus = { workspace = true }
 opentelemetry = { workspace = true }

rust/capture/src/main.rs

@@ -16,12 +16,7 @@ use tracing_subscriber::{EnvFilter, Layer};
 use capture::config::Config;
 use capture::server::serve;
 
-#[cfg(not(target_env = "msvc"))]
-use tikv_jemallocator::Jemalloc;
-
-#[cfg(not(target_env = "msvc"))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+common_alloc::used!();
 
 async fn shutdown() {
     let mut term = signal::unix::signal(signal::unix::SignalKind::terminate())

rust/common/alloc/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "common-alloc"
+version = "0.1.0"
+edition = "2021"
+
+[lints]
+workspace = true
+
+[dependencies]
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = "0.6"

rust/common/alloc/README.md

@@ -0,0 +1,12 @@
+# What is this?
+
+We use jemalloc everywhere we can, for any binary that we expect to run in a long-lived process. The reason for this is that our workloads are:
+- multi-threaded
+- extremely prone to memory fragmentation (due to our heavy use of `serde_json`, or json generally)
+
+jemalloc helps reduce memory fragmentation hugely, to the point of solving production OOMs that would have made use of capture-rs for replay a non-starter with the default system allocator.
+
+At time of writing (2024-09-04), rust workspaces don't have good support for specifying dependencies on a per-target basis, so this crate does the work of pulling in jemalloc only when compiling for supported targets, and then exposes a simple macro to use jemalloc as the global allocator. Anyone writing a binary crate should put this macro at the top of their `main.rs`. Libraries should not make use of this crate.
+
+## Future work
+Functions could be added to this crate to, in situations where jemalloc is in use, report a set of metrics about the allocator, as well as other functionality (health/liveness, a way to specify hooks to execute when memory usage exceeds a certain threshold, etc). Right now, it's prety barebones.

rust/common/alloc/src/lib.rs

@@ -0,0 +1,12 @@
+#[cfg(target_env = "msvc")]
+pub use std::alloc::System as DefaultAllocator;
+#[cfg(not(target_env = "msvc"))]
+pub use tikv_jemallocator::Jemalloc as DefaultAllocator;
+
+#[macro_export]
+macro_rules! used {
+    () => {
+        #[global_allocator]
+        static GLOBAL: $crate::DefaultAllocator = $crate::DefaultAllocator;
+    };
+}

rust/cyclotron-fetch/Cargo.toml

@@ -19,6 +19,7 @@ cyclotron-core = { path = "../cyclotron-core" }
 common-metrics = { path = "../common/metrics" }
 common-dns = { path = "../common/dns" }
 common-kafka = { path = "../common/kafka" }
+common-alloc = { path = "../common/alloc" }
 health = { path = "../common/health" }
 reqwest = { workspace = true }
 serde = { workspace = true }

rust/cyclotron-fetch/src/main.rs

@@ -10,6 +10,8 @@ use health::HealthRegistry;
 use std::{future::ready, sync::Arc};
 use tracing::{error, info};
 
+common_alloc::used!();
+
 async fn listen(app: Router, bind: String) -> Result<(), std::io::Error> {
     let listener = tokio::net::TcpListener::bind(bind).await?;
 

rust/cyclotron-janitor/Cargo.toml

@@ -19,6 +19,7 @@ cyclotron-core = { path = "../cyclotron-core" }
 common-metrics = { path = "../common/metrics" }
 common-kafka = { path = "../common/kafka" }
 health = { path = "../common/health" }
+common-alloc = { path = "../common/alloc" }
 time = { workspace = true }
 rdkafka = { workspace = true }
 

rust/cyclotron-janitor/src/main.rs

@@ -7,8 +7,7 @@ use health::{HealthHandle, HealthRegistry};
 use std::{future::ready, time::Duration};
 use tracing::{error, info};
 
-/// Most of this stuff is stolen pretty shamelessly from the rustyhook janitor. It'll diverge more
-/// once we introduce the management command stuff, but for now it's a good starting point.
+common_alloc::used!();
 
 async fn cleanup_loop(janitor: Janitor, livenes: HealthHandle, interval_secs: u64) -> Result<()> {
     let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));

rust/cyclotron-janitor/tests/janitor.rs

@@ -1,14 +1,12 @@
-use chrono::{DateTime, Duration, Timelike, Utc};
+use chrono::{Duration, Timelike, Utc};
 use common_kafka::kafka_messages::app_metrics2::{
     AppMetric2, Kind as AppMetric2Kind, Source as AppMetric2Source,
 };
 use cyclotron_core::{JobInit, JobState, QueueManager, Worker};
 use cyclotron_janitor::{config::JanitorSettings, janitor::Janitor};
 use rdkafka::consumer::{Consumer, StreamConsumer};
-use rdkafka::types::{RDKafkaApiKey, RDKafkaRespErr};
 use rdkafka::{ClientConfig, Message};
 use sqlx::PgPool;
-use std::str::FromStr;
 use uuid::Uuid;
 
 use common_kafka::{test::create_mock_kafka, APP_METRICS2_TOPIC};
@@ -58,7 +56,7 @@ async fn janitor_test(db: PgPool) {
         queue_name: queue_name.clone(),
         priority: 0,
         scheduled: now,
-        function_id: Some(uuid.clone()),
+        function_id: Some(uuid),
         vm_state: None,
         parameters: None,
         blob: None,

rust/feature-flags/Cargo.toml

@@ -31,6 +31,7 @@ regex = "1.10.4"
 maxminddb = "0.17"
 sqlx = { workspace = true }
 uuid = { workspace = true }
+common-alloc = { path = "../common/alloc" }
 
 [lints]
 workspace = true

rust/feature-flags/src/main.rs

@@ -7,6 +7,8 @@ use tracing_subscriber::{EnvFilter, Layer};
 use feature_flags::config::Config;
 use feature_flags::server::serve;
 
+common_alloc::used!();
+
 async fn shutdown() {
     let mut term = signal::unix::signal(signal::unix::SignalKind::terminate())
         .expect("failed to register SIGTERM handler");

rust/hook-api/Cargo.toml

@@ -23,3 +23,4 @@ tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 url = { workspace = true }
 common-metrics = { path = "../common/metrics" }
+common-alloc = { path = "../common/alloc" }

rust/hook-api/src/main.rs

@@ -9,6 +9,8 @@ use hook_common::pgqueue::PgQueue;
 mod config;
 mod handlers;
 
+common_alloc::used!();
+
 async fn listen(app: Router, bind: String) -> Result<()> {
     let listener = tokio::net::TcpListener::bind(bind).await?;
 

rust/hook-janitor/Cargo.toml

@@ -26,3 +26,4 @@ tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 common-metrics = { path = "../common/metrics" }
 common-kafka = { path = "../common/kafka" }
+common-alloc = { path = "../common/alloc" }

rust/hook-janitor/src/main.rs

@@ -17,6 +17,8 @@ mod config;
 mod handlers;
 mod webhooks;
 
+common_alloc::used!();
+
 async fn listen(app: Router, bind: String) -> Result<()> {
     let listener = tokio::net::TcpListener::bind(bind).await?;
 

rust/hook-worker/Cargo.toml

@@ -28,6 +28,7 @@ url = { version = "2.2" }
 common-metrics = { path = "../common/metrics" }
 common-dns = { path = "../common/dns" }
 common-kafka = { path = "../common/kafka" }
+common-alloc = { path = "../common/alloc" }
 
 [dev-dependencies]
 httpmock = { workspace = true }

rust/hook-worker/src/main.rs

@@ -13,6 +13,8 @@ use hook_worker::config::Config;
 use hook_worker::error::WorkerError;
 use hook_worker::worker::WebhookWorker;
 
+common_alloc::used!();
+
 #[tokio::main]
 async fn main() -> Result<(), WorkerError> {
     tracing_subscriber::fmt::init();

rust/property-defs-rs/Cargo.toml

@@ -21,6 +21,7 @@ metrics = { workspace = true }
 chrono = { workspace = true }
 quick_cache = { workspace = true }
 common-metrics = { path = "../common/metrics" }
+common-alloc = { path = "../common/alloc" }
 ahash = { workspace = true }
 uuid = { workspace = true }
 

rust/property-defs-rs/src/main.rs

@@ -31,6 +31,8 @@ use tokio::{
 use tracing::{info, warn};
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
 
+common_alloc::used!();
+
 fn setup_tracing() {
     let log_layer: tracing_subscriber::filter::Filtered<
         tracing_subscriber::fmt::Layer<tracing_subscriber::Registry>,