feat(metrics): Add support of customize labels

zuston · Dec 10, 2024 · d67201b · d67201b
1 parent d764bd4
commit d67201b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/src/config.rs b/src/config.rs
@@ -16,6 +16,7 @@
 // under the License.
 
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use std::fs;
 use std::path::Path;
 
@@ -368,6 +369,8 @@ pub struct MetricsConfig {
 
     #[serde(default = "as_default_push_interval_sec")]
     pub push_interval_sec: u32,
+
+    pub labels: Option<HashMap<String, String>>,
 }
 
 fn as_default_push_interval_sec() -> u32 {
@@ -553,6 +556,10 @@ mod test {
         [hdfs_store.kerberos_security_config]
         keytab_path = "/tmp/a.keytab"
         principal = "a@xxx"
+
+        [metrics]
+        push_gateway_endpoint = "http://localhost:5000"
+        labels = { l1 = "k1", l2 = "k2" }
         "#;
 
         let decoded: Config = toml::from_str(toml_str).unwrap();
@@ -576,5 +583,9 @@ mod test {
         let hdfs = decoded.hdfs_store.unwrap();
         let kerberos_config = hdfs.kerberos_security_config.unwrap();
         assert_eq!(kerberos_config.principal, "a@xxx");
+
+        // check labels of metrics
+        let metrics_labels = decoded.metrics.unwrap().labels;
+        assert_eq!(2, metrics_labels.unwrap().len());
     }
 }
diff --git a/src/metric.rs b/src/metric.rs
@@ -27,6 +27,7 @@ use prometheus::{
     register_int_counter_vec, register_int_gauge_vec, GaugeVec, Histogram, HistogramOpts,
     HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, Registry,
 };
+use std::collections::HashMap;
 use std::time::Duration;
 
 const DEFAULT_BUCKETS: &[f64] = &[
@@ -853,9 +854,19 @@ impl MetricService {
                     metrics.extend_from_slice(&custom_metrics);
                     metrics.extend_from_slice(&general_metrics);
 
+                    let mut all_labels = HashMap::new();
+                    all_labels.insert(
+                        "worker_id".to_owned(),
+                        SHUFFLE_SERVER_ID.get().unwrap().to_string(),
+                    );
+                    if let Some(labels) = &cfg.labels {
+                        let labels = labels.clone();
+                        all_labels.extend(labels);
+                    }
+
                     let pushed_result = prometheus::push_add_metrics(
                         job_name,
-                        labels! {"worker_id".to_owned() => SHUFFLE_SERVER_ID.get().unwrap().to_string(),},
+                        all_labels,
                         &push_gateway_endpoint.to_owned().unwrap().to_owned(),
                         metrics,
                         None,