From c7b02ce8ec1c6e64782438cdc35700f19ca93219 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Wed, 31 Jan 2024 13:51:11 +0000 Subject: [PATCH] proxy: use jemalloc (#6531) ## Summary of changes Experiment with jemalloc in proxy --- Cargo.lock | 33 +++++++++++++ Cargo.toml | 2 + proxy/Cargo.toml | 2 + proxy/src/bin/proxy.rs | 10 ++++ proxy/src/jemalloc.rs | 100 ++++++++++++++++++++++++++++++++++++++ proxy/src/lib.rs | 1 + workspace_hack/Cargo.toml | 4 +- 7 files changed, 150 insertions(+), 2 deletions(-) create mode 100644 proxy/src/jemalloc.rs diff --git a/Cargo.lock b/Cargo.lock index e14196350b..28ec84be1f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4080,6 +4080,8 @@ dependencies = [ "sync_wrapper", "task-local-extensions", "thiserror", + "tikv-jemalloc-ctl", + "tikv-jemallocator", "tls-listener", "tokio", "tokio-postgres", @@ -5530,6 +5532,37 @@ dependencies = [ "ordered-float 2.10.1", ] +[[package]] +name = "tikv-jemalloc-ctl" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "619bfed27d807b54f7f776b9430d4f8060e66ee138a28632ca898584d462c31c" +dependencies = [ + "libc", + "paste", + "tikv-jemalloc-sys", +] + +[[package]] +name = "tikv-jemalloc-sys" +version = "0.5.4+5.3.0-patched" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9402443cb8fd499b6f327e40565234ff34dbda27460c5b47db0db77443dd85d1" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "965fe0c26be5c56c94e38ba547249074803efd52adfb66de62107d95aab3eaca" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.21" diff --git a/Cargo.toml b/Cargo.toml index 29618ca328..26cf604a91 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -149,6 +149,8 @@ tar = "0.4" task-local-extensions = "0.1.4" test-context = "0.1" thiserror = "1.0" +tikv-jemallocator = "0.5" +tikv-jemalloc-ctl = "0.5" tls-listener = { version = "0.7", features = ["rustls", "hyper-h1"] } tokio = { version = "1.17", features = ["macros"] } tokio-epoll-uring = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "main" } diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index f075c718a7..79abe639ed 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -62,6 +62,8 @@ socket2.workspace = true sync_wrapper.workspace = true task-local-extensions.workspace = true thiserror.workspace = true +tikv-jemallocator.workspace = true +tikv-jemalloc-ctl = { workspace = true, features = ["use_std"] } tls-listener.workspace = true tokio-postgres.workspace = true tokio-rustls.workspace = true diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index ba113a89eb..3960b080be 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -32,6 +32,9 @@ project_build_tag!(BUILD_TAG); use clap::{Parser, ValueEnum}; +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + #[derive(Clone, Debug, ValueEnum)] enum AuthBackend { Console, @@ -187,6 +190,13 @@ async fn main() -> anyhow::Result<()> { info!("Build_tag: {BUILD_TAG}"); ::metrics::set_build_info_metric(GIT_VERSION, BUILD_TAG); + match proxy::jemalloc::MetricRecorder::new(prometheus::default_registry()) { + Ok(t) => { + t.start(); + } + Err(e) => tracing::error!(error = ?e, "could not start jemalloc metrics loop"), + } + let args = ProxyCliArgs::parse(); let config = build_config(&args)?; diff --git a/proxy/src/jemalloc.rs b/proxy/src/jemalloc.rs new file mode 100644 index 0000000000..ed20798d56 --- /dev/null +++ b/proxy/src/jemalloc.rs @@ -0,0 +1,100 @@ +use std::time::Duration; + +use metrics::IntGauge; +use prometheus::{register_int_gauge_with_registry, Registry}; +use tikv_jemalloc_ctl::{config, epoch, epoch_mib, stats, version}; + +pub struct MetricRecorder { + epoch: epoch_mib, + active: stats::active_mib, + active_gauge: IntGauge, + allocated: stats::allocated_mib, + allocated_gauge: IntGauge, + mapped: stats::mapped_mib, + mapped_gauge: IntGauge, + metadata: stats::metadata_mib, + metadata_gauge: IntGauge, + resident: stats::resident_mib, + resident_gauge: IntGauge, + retained: stats::retained_mib, + retained_gauge: IntGauge, +} + +impl MetricRecorder { + pub fn new(registry: &Registry) -> Result { + tracing::info!( + config = config::malloc_conf::read()?, + version = version::read()?, + "starting jemalloc recorder" + ); + + Ok(Self { + epoch: epoch::mib()?, + active: stats::active::mib()?, + active_gauge: register_int_gauge_with_registry!( + "jemalloc_active_bytes", + "Total number of bytes in active pages allocated by the process", + registry + )?, + allocated: stats::allocated::mib()?, + allocated_gauge: register_int_gauge_with_registry!( + "jemalloc_allocated_bytes", + "Total number of bytes allocated by the process", + registry + )?, + mapped: stats::mapped::mib()?, + mapped_gauge: register_int_gauge_with_registry!( + "jemalloc_mapped_bytes", + "Total number of bytes in active extents mapped by the allocator", + registry + )?, + metadata: stats::metadata::mib()?, + metadata_gauge: register_int_gauge_with_registry!( + "jemalloc_metadata_bytes", + "Total number of bytes dedicated to jemalloc metadata", + registry + )?, + resident: stats::resident::mib()?, + resident_gauge: register_int_gauge_with_registry!( + "jemalloc_resident_bytes", + "Total number of bytes in physically resident data pages mapped by the allocator", + registry + )?, + retained: stats::retained::mib()?, + retained_gauge: register_int_gauge_with_registry!( + "jemalloc_retained_bytes", + "Total number of bytes in virtual memory mappings that were retained rather than being returned to the operating system", + registry + )?, + }) + } + + fn _poll(&self) -> Result<(), anyhow::Error> { + self.epoch.advance()?; + self.active_gauge.set(self.active.read()? as i64); + self.allocated_gauge.set(self.allocated.read()? as i64); + self.mapped_gauge.set(self.mapped.read()? as i64); + self.metadata_gauge.set(self.metadata.read()? as i64); + self.resident_gauge.set(self.resident.read()? as i64); + self.retained_gauge.set(self.retained.read()? as i64); + Ok(()) + } + + #[inline] + pub fn poll(&self) { + if let Err(error) = self._poll() { + tracing::warn!(%error, "Failed to poll jemalloc stats"); + } + } + + pub fn start(self) -> tokio::task::JoinHandle<()> { + tokio::task::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(15)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + self.poll(); + interval.tick().await; + } + }) + } +} diff --git a/proxy/src/lib.rs b/proxy/src/lib.rs index a9e4a38302..db6256d611 100644 --- a/proxy/src/lib.rs +++ b/proxy/src/lib.rs @@ -16,6 +16,7 @@ pub mod console; pub mod context; pub mod error; pub mod http; +pub mod jemalloc; pub mod logging; pub mod metrics; pub mod parse; diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index c29f8b422f..8fd49956cc 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -45,7 +45,7 @@ hmac = { version = "0.12", default-features = false, features = ["reset"] } hyper = { version = "0.14", features = ["full"] } indexmap = { version = "1", default-features = false, features = ["std"] } itertools = { version = "0.10" } -libc = { version = "0.2", features = ["extra_traits"] } +libc = { version = "0.2", features = ["extra_traits", "use_std"] } log = { version = "0.4", default-features = false, features = ["std"] } memchr = { version = "2" } nom = { version = "7" } @@ -94,7 +94,7 @@ getrandom = { version = "0.2", default-features = false, features = ["std"] } hashbrown = { version = "0.14", default-features = false, features = ["raw"] } indexmap = { version = "1", default-features = false, features = ["std"] } itertools = { version = "0.10" } -libc = { version = "0.2", features = ["extra_traits"] } +libc = { version = "0.2", features = ["extra_traits", "use_std"] } log = { version = "0.4", default-features = false, features = ["std"] } memchr = { version = "2" } nom = { version = "7" }