feat: Support prometheus remote write protocol for metrics (#563)

* fix: make sure to post correct registry to vm * feat: Add support for remote writing prometheus --------- Co-authored-by: Simon Hornby <[email protected]>
Unleash · Oct 16, 2024 · 3e1f60d · 3e1f60d
1 parent 90c9ae7
commit 3e1f60d
Show file tree

Hide file tree

Showing 8 changed files with 211 additions and 8 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/server/Cargo.toml b/server/Cargo.toml
@@ -34,6 +34,7 @@ anyhow = "1.0.89"
 async-trait = "0.1.83"
 aws-config = { version = "1.5.7", features = ["behavior-version-latest"] }
 aws-sdk-s3 = { version = "1.53.0", features = ["behavior-version-latest"] }
+base64 = "0.22.1"
 chrono = { version = "0.4.38", features = ["serde"] }
 cidr = "0.3.0"
 clap = { version = "4.5.19", features = ["derive", "env"] }
@@ -54,7 +55,8 @@ opentelemetry_sdk = { version = "0.24.0", features = [
     "serde_json",
     "logs",
 ] }
-prometheus = { version = "0.13.4", features = ["process"] }
+prometheus = { version = "0.13.4", features = ["process", "push"] }
+prometheus-reqwest-remote-write = { version = "0.1.1" }
 prometheus-static-metric = "0.5.1"
 rand = "0.8.5"
 redis = { version = "0.27.0", features = [

diff --git a/server/src/builder.rs b/server/src/builder.rs
@@ -358,6 +358,11 @@ mod tests {
             },
             upstream_certificate_file: Default::default(),
             token_revalidation_interval_seconds: Default::default(),
+            prometheus_push_interval: 60,
+            prometheus_remote_write_url: None,
+            prometheus_user_id: None,
+            prometheus_password: None,
+            prometheus_username: None,
         };
 
         let result = build_edge(&args, "test-app").await;

diff --git a/server/src/cli.rs b/server/src/cli.rs
@@ -136,6 +136,11 @@ pub struct ClientIdentity {
     pub pkcs12_passphrase: Option<String>,
 }
 
+pub enum PromAuth {
+    None,
+    Basic(String, String),
+}
+
 #[derive(Args, Debug, Clone)]
 #[command(group(
     ArgGroup::new("data-provider")
@@ -207,6 +212,23 @@ pub struct EdgeArgs {
     /// If set to true, Edge starts with dynamic behavior. Dynamic behavior means that Edge will accept tokens outside the scope of the startup tokens
     #[clap(long, env, default_value_t = false, conflicts_with = "strict")]
     pub dynamic: bool,
+
+    /// Sets a remote write url for prometheus metrics, if this is set, prometheus metrics will be written upstream
+    #[clap(long, env)]
+    pub prometheus_remote_write_url: Option<String>,
+
+    /// Sets the interval for prometheus push metrics, only relevant if `prometheus_remote_write_url` is set. Defaults to 60 seconds
+    #[clap(long, env, default_value_t = 60)]
+    pub prometheus_push_interval: u64,
+
+    #[clap(long, env)]
+    pub prometheus_username: Option<String>,
+
+    #[clap(long, env)]
+    pub prometheus_password: Option<String>,
+
+    #[clap(long, env)]
+    pub prometheus_user_id: Option<String>,
 }
 
 pub fn string_to_header_tuple(s: &str) -> Result<(String, String), String> {

diff --git a/server/src/main.rs b/server/src/main.rs
@@ -26,6 +26,8 @@ use unleash_edge::{internal_backstage, tls};
 #[cfg(not(tarpaulin_include))]
 #[actix_web::main]
 async fn main() -> Result<(), anyhow::Error> {
+    use unleash_edge::metrics::metrics_pusher;
+
     let args = CliArgs::parse();
     let disable_all_endpoint = args.disable_all_endpoint;
     if args.markdown_help {
@@ -70,6 +72,7 @@ async fn main() -> Result<(), anyhow::Error> {
 
     let openapi = openapi::ApiDoc::openapi();
     let refresher_for_app_data = feature_refresher.clone();
+    let prom_registry_for_write = metrics_handler.registry.clone();
     let server = HttpServer::new(move || {
         let qs_config =
             serde_qs::actix::QsQueryConfig::default().qs_config(serde_qs::Config::new(5, false));
@@ -165,6 +168,9 @@ async fn main() -> Result<(), anyhow::Error> {
                 _ = validator.schedule_revalidation_of_startup_tokens(edge.tokens, lazy_feature_refresher) => {
                     tracing::info!("Token validator validation of startup tokens was unexpectedly shut down");
                 }
+                _ = metrics_pusher::prometheus_remote_write(prom_registry_for_write, edge.prometheus_remote_write_url, edge.prometheus_push_interval, edge.prometheus_username, edge.prometheus_password) => {
+                    tracing::info!("Prometheus push unexpectedly shut down");
+                }
             }
         }
         cli::EdgeMode::Offline(offline_args) if offline_args.reload_interval > 0 => {

diff --git a/server/src/metrics/actix_web_metrics.rs b/server/src/metrics/actix_web_metrics.rs
@@ -387,7 +387,7 @@ where
 
 #[derive(Clone, Debug)]
 pub struct PrometheusMetricsHandler {
-    registry: prometheus::Registry,
+    pub registry: prometheus::Registry,
 }
 
 impl PrometheusMetricsHandler {