chore: Add support for metrics-exporter-dogstatsd

Cargo.toml

@@ -25,6 +25,7 @@ http = "1.1.0"
 http-body-util = "0.1.2"
 libsqlite3-sys = "0.30.1"
 metrics = "0.24.0"
+metrics-exporter-dogstatsd = "0.9.6"
 metrics-exporter-statsd = "0.9.0"
 prost = "0.13"
 prost-types = "0.13.3"

src/config.rs

@@ -39,6 +39,9 @@ pub struct Config {
     /// The statsd address to report metrics to.
     pub statsd_addr: String,
 
+    /// The dogstatsd metrics host.
+    pub metrics_host: String,
+
     /// Default tags to add to all metrics.
     pub default_metrics_tags: BTreeMap<String, String>,
 
@@ -246,12 +249,15 @@ impl Default for Config {
             sentry_dsn: None,
             sentry_env: None,
             traces_sample_rate: Some(0.0),
-            log_filter: "info,librdkafka=warn,h2=off".to_owned(),
+            log_filter:
+                "info,librdkafka=warn,h2=off,metrics_exporter_dogstatsd::forwarder::sync=off"
+                    .to_owned(),
             log_format: LogFormat::Text,
             grpc_addr: "0.0.0.0".to_owned(),
             grpc_port: 50051,
             grpc_shared_secret: vec![],
             statsd_addr: "127.0.0.1:8126".parse().unwrap(),
+            metrics_host: "".to_owned(),
             default_metrics_tags: Default::default(),
             kafka_cluster: "127.0.0.1:9092".to_owned(),
             kafka_consumer_group: "taskworker".to_owned(),
@@ -432,7 +438,10 @@ mod tests {
         };
         assert_eq!(config.sentry_dsn, None);
         assert_eq!(config.sentry_env, None);
-        assert_eq!(config.log_filter, "info,librdkafka=warn,h2=off");
+        assert_eq!(
+            config.log_filter,
+            "info,librdkafka=warn,h2=off,metrics_exporter_dogstatsd::forwarder::sync=off"
+        );
         assert_eq!(config.log_format, LogFormat::Text);
         assert_eq!(config.grpc_port, 50051);
         assert_eq!(config.kafka_topic, "taskworker");

src/metrics.rs

@@ -1,11 +1,13 @@
 use crate::config::Config;
+use metrics_exporter_dogstatsd::DogStatsDBuilder;
 use metrics_exporter_statsd::StatsdBuilder;
 use std::{
     collections::BTreeMap,
     net::{SocketAddr, ToSocketAddrs},
 };
 
 pub struct MetricsConfig {
+    pub metrics_host: String,
     pub statsd_addr: SocketAddr,
     pub default_tags: BTreeMap<String, String>,
 }
@@ -21,25 +23,49 @@ impl MetricsConfig {
         };
         MetricsConfig {
             statsd_addr: *statsd_addr,
+            metrics_host: config.metrics_host.clone(),
             default_tags: config.default_metrics_tags.clone(),
         }
     }
 }
 
 pub fn init(metrics_config: MetricsConfig) {
-    let address = metrics_config.statsd_addr;
+    if metrics_config.metrics_host.is_empty() {
+        // Fallback to statsd forwarder
+        let address = metrics_config.statsd_addr;
 
-    let builder = StatsdBuilder::from(address.ip().to_string(), address.port());
+        let builder = StatsdBuilder::from(address.ip().to_string(), address.port());
 
-    let recorder = metrics_config
-        .default_tags
-        .into_iter()
-        .fold(
-            builder.with_queue_size(5000).with_buffer_size(1024),
-            |builder, (key, value)| builder.with_default_tag(key, value),
-        )
-        .build(Some("taskbroker"))
-        .expect("Could not create StatsdRecorder");
+        let recorder = metrics_config
+            .default_tags
+            .into_iter()
+            .fold(
+                builder.with_queue_size(5000).with_buffer_size(1024),
+                |builder, (key, value)| builder.with_default_tag(key, value),
+            )
+            .build(Some("taskbroker"))
+            .expect("Could not create StatsdRecorder");
 
-    metrics::set_global_recorder(recorder).expect("Could not set global metrics recorder")
+        metrics::set_global_recorder(recorder).expect("Could not set global metrics recorder")
+    } else {
+        let default_tags = metrics_config
+            .default_tags
+            .into_iter()
+            .map(|(key, value)| metrics::Label::new(key, value))
+            .collect();
+
+        // Use dogstatsd exporter if enabled.
+        let builder = DogStatsDBuilder::default()
+            .with_remote_address(metrics_config.metrics_host)
+            .expect("Could not set metrics host address")
+            .with_telemetry(true)
+            .send_histograms_as_distributions(true)
+            .with_histogram_sampling(true)
+            .set_global_prefix("taskbroker")
+            .with_global_labels(default_tags);
+
+        builder
+            .install()
+            .expect("Could not create DogStatsDBuilder");
+    }
 }