Auto merge of #478 - Zeegomo:prometheus, r=pietroalbini

add metrics endpoint Fixes #477 Record metrics about jobs completed by each agent and make them available to prometheus
rust-lang · Oct 14, 2019 · c551d9a · c551d9a
2 parents 82e890a + 68ea1dd
commit c551d9a
Show file tree

Hide file tree

Showing 7 changed files with 91 additions and 0 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -67,6 +67,7 @@ rustwide = { version = "0.3.1", features = ["unstable"] }
 percent-encoding = "2.1.0"
 remove_dir_all = "0.5.2"
 ctrlc = "3.1.3"
+prometheus = "0.7.0"
 
 [dev-dependencies]
 assert_cmd = "0.10.1"

diff --git a/src/server/metrics.rs b/src/server/metrics.rs
@@ -0,0 +1,24 @@
+use crate::prelude::*;
+use prometheus::{IntCounterVec, __register_counter_vec};
+
+#[derive(Clone)]
+pub struct Metrics {
+    crater_completed_jobs_total: IntCounterVec,
+}
+
+impl Metrics {
+    pub fn new() -> Fallible<Self> {
+        let opts = prometheus::opts!("crater_completed_jobs_total", "total completed jobs");
+        let crater_completed_jobs_total =
+            prometheus::register_int_counter_vec!(opts, &["agent", "experiment"])?;
+        Ok(Metrics {
+            crater_completed_jobs_total,
+        })
+    }
+
+    pub fn record_completed_jobs(&self, agent: &str, experiment: &str, amount: i64) {
+        self.crater_completed_jobs_total
+            .with_label_values(&[&agent, &experiment])
+            .inc_by(amount);
+    }
+}
diff --git a/src/server/mod.rs b/src/server/mod.rs
@@ -3,6 +3,7 @@ pub mod api_types;
 mod auth;
 mod github;
 mod messages;
+mod metrics;
 mod reports;
 mod routes;
 pub mod tokens;
@@ -17,6 +18,7 @@ use crate::server::github::{GitHub, GitHubApi};
 use crate::server::tokens::Tokens;
 use http::{self, header::HeaderValue, Response};
 use hyper::Body;
+use metrics::Metrics;
 use std::sync::{Arc, Mutex};
 use warp::{self, Filter};
 
@@ -43,6 +45,7 @@ pub struct Data {
     pub db: Database,
     pub reports_worker: reports::ReportsWorker,
     pub acl: ACL,
+    pub metrics: Metrics,
 }
 
 pub fn run(config: Config) -> Fallible<()> {
@@ -52,6 +55,7 @@ pub fn run(config: Config) -> Fallible<()> {
     let agents = Agents::new(db.clone(), &tokens)?;
     let bot_username = github.username()?;
     let acl = ACL::new(&config, &github)?;
+    let metrics = Metrics::new()?;
 
     info!("bot username: {}", bot_username);
 
@@ -64,6 +68,7 @@ pub fn run(config: Config) -> Fallible<()> {
         db: db.clone(),
         reports_worker: reports::ReportsWorker::new(),
         acl,
+        metrics,
     };
 
     let mutex = Arc::new(Mutex::new(data.clone()));
@@ -80,6 +85,8 @@ pub fn run(config: Config) -> Fallible<()> {
                 .and(warp::path("webhooks").and(routes::webhooks::routes(data.clone())))
                 .or(warp::path("agent-api").and(routes::agent::routes(data.clone(), mutex.clone())))
                 .unify()
+                .or(warp::path("metrics").and(routes::metrics::routes()))
+                .unify()
                 .or(routes::ui::routes(data.clone()))
                 .unify(),
         )

diff --git a/src/server/routes/agent.rs b/src/server/routes/agent.rs
@@ -166,6 +166,9 @@ fn endpoint_record_progress(
         ex.name, auth.name,
     );
 
+    data.metrics
+        .record_completed_jobs(&auth.name, &ex.name, result.data.results.len() as i64);
+
     let db = DatabaseDB::new(&data.db);
     db.store(&ex, &result.data, EncodingType::Gzip)?;
 

diff --git a/src/server/routes/metrics.rs b/src/server/routes/metrics.rs
@@ -0,0 +1,28 @@
+use crate::prelude::*;
+use http::{Response, StatusCode};
+use hyper::Body;
+use prometheus::{Encoder, TextEncoder};
+use warp::{self, Filter, Rejection};
+
+pub fn routes() -> impl Filter<Extract = (Response<Body>,), Error = Rejection> + Clone {
+    warp::get2()
+        .and(warp::path::end())
+        .map(|| match endpoint_metrics() {
+            Ok(resp) => resp,
+            Err(err) => {
+                error!("error while processing metrics");
+                crate::utils::report_failure(&err);
+
+                let mut resp = Response::new(format!("Error: {}\n", err).into());
+                *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
+                resp
+            }
+        })
+}
+
+fn endpoint_metrics() -> Fallible<Response<Body>> {
+    let mut buffer = Vec::new();
+    let families = prometheus::gather();
+    TextEncoder::new().encode(&families, &mut buffer)?;
+    Ok(Response::new(Body::from(buffer)))
+}
diff --git a/src/server/routes/mod.rs b/src/server/routes/mod.rs
@@ -1,3 +1,4 @@
 pub mod agent;
+pub mod metrics;
 pub mod ui;
 pub mod webhooks;