tikv_alloc: abstract some new jemalloc code (#4411)

* tikv_alloc: abstract some new jemalloc code

This renames the jealloc_metrics module to allocator_metrics, does
similar for the code inside it. It implements a no-op fetch_stats
method for the system allocator.

It marks the jemalloc-sys and jemalloc-ctl crates as optional,
so that they might not be compiled when jemalloc is disabled.

It's not actually possible disable the jemalloc features in this PR
because there is no way to turn off the default-features of the
tikv_alloc crate. A followup in #4207 will make it possible.

It also moves a `mod` to after some `use` statements.

Signed-off-by: Brian Anderson <andersrb@gmail.com>

components/tikv_alloc/Cargo.toml

@@ -6,7 +6,8 @@ authors = ["Brian Anderson <andersrb@gmail.com>"]
 publish = false
 
 [features]
-default = ["jemallocator"]
+default = ["jemalloc"]
+jemalloc = ["jemallocator", "jemalloc-sys", "jemalloc-ctl"]
 
 # Build jemalloc's profiling features. Without this
 # certain profile functions will return nothing.
@@ -27,9 +28,17 @@ version = "0.1.9"
 optional = true
 features = ["unprefixed_malloc_on_supported_platforms"]
 
+# FIXME: shouldn't need to link to this crate, but the 'stats'
+# feature is missing in jemallocator.
+#
+# https://github.com/gnzlbg/jemallocator/issues/112
+# https://github.com/tikv/tikv/issues/4406
+
 [dependencies.jemalloc-sys]
 version = "0.1.8"
+optional = true
 features = ["stats"]
 
 [dependencies.jemalloc-ctl]
 version = "0.2.0"
+optional = true

components/tikv_alloc/src/lib.rs

@@ -106,9 +106,12 @@ static ALLOC: std::alloc::System = std::alloc::System;
 
 pub use self::imp::*;
 
+pub type AllocStats = Vec<(&'static str, usize)>;
+
 // The implementation of this crate when jemalloc is turned on
 #[cfg(all(unix, not(fuzzing), not(feature = "no-jemalloc")))]
 mod imp {
+    use super::AllocStats;
     use jemalloc_ctl::{stats, Epoch as JeEpoch};
     use jemallocator::ffi::malloc_stats_print;
     use libc::{self, c_char, c_void};
@@ -128,27 +131,18 @@ mod imp {
         String::from_utf8_lossy(&buf).into_owned()
     }
 
-    pub struct JemallocStats {
-        pub allocated: usize,
-        pub active: usize,
-        pub metadata: usize,
-        pub resident: usize,
-        pub mapped: usize,
-        pub retained: usize,
-    }
-
-    pub fn fetch_stats() -> io::Result<JemallocStats> {
+    pub fn fetch_stats() -> io::Result<Option<AllocStats>> {
         // Stats are cached. Need to advance epoch to refresh.
         JeEpoch::new()?.advance()?;
 
-        Ok(JemallocStats {
-            allocated: stats::allocated()?,
-            active: stats::active()?,
-            metadata: stats::metadata()?,
-            resident: stats::resident()?,
-            mapped: stats::mapped()?,
-            retained: stats::retained()?,
-        })
+        Ok(Some(vec![
+            ("allocated", stats::allocated()?),
+            ("active", stats::active()?),
+            ("metadata", stats::metadata()?),
+            ("resident", stats::resident()?),
+            ("mapped", stats::mapped()?),
+            ("retained", stats::retained()?),
+        ]))
     }
 
     #[allow(clippy::cast_ptr_alignment)]
@@ -304,8 +298,16 @@ mod imp {
 
 #[cfg(not(all(unix, not(fuzzing), not(feature = "no-jemalloc"))))]
 mod imp {
+
+    use super::AllocStats;
+    use std::io;
+
     pub fn dump_stats() -> String {
         String::new()
     }
     pub fn dump_prof(_path: Option<&str>) {}
+
+    pub fn fetch_stats() -> io::Result<Option<AllocStats>> {
+        Ok(None)
+    }
 }

src/bin/util/setup.rs

@@ -68,8 +68,8 @@ pub fn initial_logger(config: &TiKvConfig) {
 pub fn initial_metric(cfg: &MetricConfig, node_id: Option<u64>) {
     util::metrics::monitor_threads("tikv")
         .unwrap_or_else(|e| fatal!("failed to start monitor thread: {}", e));
-    util::metrics::monitor_jemalloc_stats("tikv")
-        .unwrap_or_else(|e| fatal!("failed to monitor jemalloc stats: {}", e));
+    util::metrics::monitor_allocator_stats("tikv")
+        .unwrap_or_else(|e| fatal!("failed to monitor allocator stats: {}", e));
 
     if cfg.interval.as_secs() == 0 || cfg.address.is_empty() {
         return;

src/util/metrics/allocator_metrics.rs

@@ -0,0 +1,55 @@
+// Copyright 2017 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use prometheus::core::{Collector, Desc};
+use prometheus::proto::MetricFamily;
+use prometheus::{IntGaugeVec, Opts, Result};
+
+use tikv_alloc;
+
+pub fn monitor_allocator_stats<S: Into<String>>(namespace: S) -> Result<()> {
+    prometheus::register(Box::new(AllocStatsCollector::new(namespace)?))
+}
+
+struct AllocStatsCollector {
+    descs: Vec<Desc>,
+    metrics: IntGaugeVec,
+}
+
+impl AllocStatsCollector {
+    fn new<S: Into<String>>(namespace: S) -> Result<AllocStatsCollector> {
+        let stats = IntGaugeVec::new(
+            Opts::new("allocator_stats", "Allocator stats").namespace(namespace.into()),
+            &["type"],
+        )?;
+        Ok(AllocStatsCollector {
+            descs: stats.desc().into_iter().cloned().collect(),
+            metrics: stats,
+        })
+    }
+}
+
+impl Collector for AllocStatsCollector {
+    fn desc(&self) -> Vec<&Desc> {
+        self.descs.iter().collect()
+    }
+
+    fn collect(&self) -> Vec<MetricFamily> {
+        if let Ok(Some(stats)) = tikv_alloc::fetch_stats() {
+            for stat in stats {
+                self.metrics.with_label_values(&[stat.0]).set(stat.1 as i64);
+            }
+        }
+        self.metrics.collect()
+    }
+}

src/util/metrics/mod.rs

@@ -11,8 +11,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub mod jemalloc_metrics;
-
 use std::thread;
 use std::time::Duration;
 
@@ -28,8 +26,9 @@ mod threads_dummy;
 #[cfg(not(target_os = "linux"))]
 pub use self::threads_dummy::monitor_threads;
 
-#[cfg(all(unix, not(fuzzing), not(feature = "no-jemalloc")))]
-pub use self::jemalloc_metrics::monitor_jemalloc_stats;
+pub use self::allocator_metrics::monitor_allocator_stats;
+
+pub mod allocator_metrics;
 
 /// Runs a background Prometheus client.
 pub fn run_prometheus(