fix(backend): prometheus histogram for worker job timer (#312)

* prometheus histogram for worker job timer

hosts on :8001

* some new metrics in worker + field

adds start_time_seconds, job_duration_seconds & jobs_failed

* use tokio task_local to count job failures

* METRICS_ADDR environment variable off by default

true defaults to 0.0.0.0:8001 otherwise expects a socket address

* pass metrics as args instead of task local

backend/Cargo.lock

@@ -2405,6 +2405,20 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "prometheus"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cface98dfa6d645ea4c789839f176e4b072265d085bfcc48eaa8d137f58d3c39"
+dependencies = [
+ "cfg-if 1.0.0",
+ "fnv",
+ "lazy_static",
+ "memchr",
+ "parking_lot 0.12.1",
+ "thiserror",
+]
+
 [[package]]
 name = "prost"
 version = "0.10.4"
@@ -4538,6 +4552,7 @@ dependencies = [
  "lettre",
  "magic-crypt",
  "mime_guess",
+ "prometheus",
  "rand 0.8.5",
  "rand_core 0.6.3",
  "regex",

backend/Cargo.toml

@@ -24,6 +24,7 @@ chrono = { version = "^0", features = ["serde"]}
 tracing = "^0"
 tracing-subscriber = { version = "^0", features = ["env-filter", "json"]}
 console-subscriber = "^0"
+prometheus = { version = "^0", default-features = false }
 
 rust-embed = "^6"
 mime_guess = "^2"

backend/src/jobs.rs

@@ -24,6 +24,7 @@ use crate::{
     scripts::{get_hub_script_by_path, ScriptHash, ScriptLang},
     users::{owner_to_token_owner, Authed},
     utils::{require_admin, Pagination, StripPath},
+    worker,
     worker_flow::init_flow_status,
 };
 use axum::{
@@ -1257,7 +1258,9 @@ pub async fn add_completed_job_error<E: ToString + std::fmt::Debug>(
     queued_job: &QueuedJob,
     logs: String,
     e: E,
+    metrics: &worker::Metrics,
 ) -> Result<(Uuid, Map<String, Value>), Error> {
+    metrics.jobs_failed.inc();
     let mut output_map = serde_json::Map::new();
     output_map.insert(
         "error".to_string(),

backend/src/lib.rs

@@ -9,6 +9,7 @@
 use argon2::Argon2;
 use axum::{handler::Handler, middleware::from_extractor, routing::get, Extension, Router};
 use db::DB;
+use futures::FutureExt;
 use git_version::git_version;
 use slack_http_verifier::SlackVerifier;
 use std::{net::SocketAddr, sync::Arc};
@@ -287,3 +288,24 @@ pub async fn shutdown_signal(tx: tokio::sync::broadcast::Sender<()>) -> anyhow::
     let _ = tx.send(());
     Ok(())
 }
+
+pub async fn serve_metrics(
+    addr: SocketAddr,
+    mut rx: tokio::sync::broadcast::Receiver<()>,
+) -> Result<(), hyper::Error> {
+    axum::Server::bind(&addr)
+        .serve(
+            Router::new()
+                .route("/metrics", get(metrics))
+                .into_make_service(),
+        )
+        .with_graceful_shutdown(rx.recv().map(drop))
+        .await
+}
+
+async fn metrics() -> Result<String, Error> {
+    let metric_families = prometheus::gather();
+    Ok(prometheus::TextEncoder::new()
+        .encode_to_string(&metric_families)
+        .map_err(anyhow::Error::from)?)
+}

backend/src/main.rs

@@ -23,6 +23,16 @@ async fn main() -> anyhow::Result<()> {
         .and_then(|x| x.parse::<i32>().ok())
         .unwrap_or(windmill::DEFAULT_NUM_WORKERS as i32);
 
+    let metrics_addr: Option<SocketAddr> = std::env::var("METRICS_ADDR")
+        .ok()
+        .map(|s| {
+            s.parse::<bool>()
+                .map(|b| b.then(|| SocketAddr::from(([0, 0, 0, 0], 8001))))
+                .or_else(|_| s.parse::<SocketAddr>().map(Some))
+        })
+        .transpose()?
+        .flatten();
+
     let (server_mode, monitor_mode, migrate_db) = (true, true, true);
 
     if migrate_db {
@@ -104,7 +114,16 @@ async fn main() -> anyhow::Result<()> {
             Ok(()) as anyhow::Result<()>
         };
 
-        futures::try_join!(shutdown_signal, server_f, workers_f, monitor_f)?;
+        let metrics_f = async {
+            match metrics_addr {
+                Some(addr) => windmill::serve_metrics(addr, tx.subscribe())
+                    .await
+                    .map_err(anyhow::Error::from),
+                None => Ok(()),
+            }
+        };
+
+        futures::try_join!(shutdown_signal, server_f, workers_f, monitor_f, metrics_f)?;
     }
 
     Ok(())

backend/src/scripts.rs

@@ -67,6 +67,16 @@ pub enum ScriptLang {
     Deno,
     Python3,
 }
+
+impl ScriptLang {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            ScriptLang::Deno => "deno",
+            ScriptLang::Python3 => "python3",
+        }
+    }
+}
+
 #[derive(sqlx::Type, PartialEq, Debug, Hash, Clone, Copy)]
 #[sqlx(transparent)]
 pub struct ScriptHash(pub i64);

backend/src/worker.rs

@@ -39,12 +39,11 @@ use tokio::{
     fs::{DirBuilder, File},
     io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader},
     process::{Child, Command},
-    sync::Mutex,
+    sync::{mpsc, Mutex},
     time::Instant,
 };
 
 use async_recursion::async_recursion;
-use tokio::sync::mpsc;
 
 const TMP_DIR: &str = "/tmp/windmill";
 const PIP_CACHE_DIR: &str = "/tmp/windmill/cache/pip";
@@ -57,6 +56,10 @@ const NSJAIL_CONFIG_RUN_PYTHON3_CONTENT: &str =
     include_str!("../../nsjail/run.python3.config.proto");
 const NSJAIL_CONFIG_RUN_DENO_CONTENT: &str = include_str!("../../nsjail/run.deno.config.proto");
 
+pub struct Metrics {
+    pub jobs_failed: prometheus::IntCounter,
+}
+
 pub async fn run_worker(
     db: &DB,
     timeout: i32,
@@ -89,6 +92,37 @@ pub async fn run_worker(
 
     insert_initial_ping(worker_instance, &worker_name, ip, db).await;
 
+    prometheus::register_int_gauge!(prometheus::Opts::new(
+        "start_time_seconds",
+        "Start time of worker as seconds since unix epoch",
+    )
+    .const_label("name", &worker_name))
+    .expect("register prometheus metric")
+    .set(
+        std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .ok()
+            .map(|duration| duration.as_secs() as i64)
+            .unwrap_or(0),
+    );
+
+    let job_duration_seconds = prometheus::register_histogram_vec!(
+        prometheus::HistogramOpts::new(
+            "job_duration_seconds",
+            "Duration between receiving a job and completing it",
+        )
+        .const_label("name", &worker_name),
+        &["workspace_id", "language"],
+    )
+    .expect("register prometheus metric");
+
+    let jobs_failed = prometheus::register_int_counter_vec!(
+        prometheus::Opts::new("jobs_failed", "Number of failed jobs",)
+            .const_label("name", &worker_name),
+        &["workspace_id", "language"],
+    )
+    .expect("register prometheus metric");
+
     let mut jobs_executed = 0;
     let mut rx = tx.subscribe();
     loop {
@@ -108,12 +142,24 @@ pub async fn run_worker(
 
         match pull(db).await {
             Ok(Some(job)) => {
+                let label_values = [
+                    &job.workspace_id,
+                    job.language.as_ref().map(|l| l.as_str()).unwrap_or(""),
+                ];
+
+                let _timer = job_duration_seconds
+                    .with_label_values(label_values.as_slice())
+                    .start_timer();
+
                 jobs_executed += 1;
 
+                let metrics =
+                    Metrics { jobs_failed: jobs_failed.with_label_values(label_values.as_slice()) };
+
                 tracing::info!(worker = %worker_name, id = %job.id, "Fetched job");
-                let job2 = job.clone();
+
                 if let Some(err) = handle_queued_job(
-                    job,
+                    job.clone(),
                     db,
                     timeout,
                     &worker_name,
@@ -121,39 +167,43 @@ pub async fn run_worker(
                     base_url,
                     disable_nuser,
                     disable_nsjail,
+                    &metrics,
                 )
                 .await
                 .err()
                 {
-                    let err_string = err.to_string().clone();
                     let m = add_completed_job_error(
                         db,
-                        &job2,
+                        &job,
                         "Unexpected error during job execution:\n".to_string(),
-                        err,
+                        &err,
+                        &metrics,
                     )
                     .await
                     .map(|(_, m)| m)
                     .unwrap_or_else(|_| Map::new());
 
-                    {
-                        let job = job2.clone();
-                        let _ = postprocess_queued_job(
-                            job.is_flow_step,
-                            job.schedule_path,
-                            job.script_path,
-                            &job2.workspace_id,
-                            job2.id,
+                    let _ = postprocess_queued_job(
+                        job.is_flow_step,
+                        job.schedule_path.clone(),
+                        job.script_path.clone(),
+                        &job.workspace_id,
+                        job.id,
+                        db,
+                    )
+                    .await;
+
+                    if job.parent_job.is_some() {
+                        let _ = update_flow_status_after_job_completion(
                             db,
+                            &job,
+                            false,
+                            Some(m),
+                            &metrics,
                         )
                         .await;
                     }
-
-                    if job2.parent_job.is_some() {
-                        let _ = update_flow_status_after_job_completion(db, &job2, false, Some(m))
-                            .await;
-                    }
-                    tracing::error!(job_id = %job2.id, "Error handling job: {err_string}");
+                    tracing::error!(job_id = %job.id, "Error handling job: {err}");
                 };
             }
             Ok(None) => (),
@@ -193,6 +243,7 @@ async fn handle_queued_job(
     base_url: &str,
     disable_nuser: bool,
     disable_nsjail: bool,
+    metrics: &Metrics,
 ) -> crate::error::Result<()> {
     let job_id = job.id;
     let w_id = &job.workspace_id.clone();
@@ -238,14 +289,22 @@ async fn handle_queued_job(
                 Ok(r) => {
                     add_completed_job(db, &job, true, false, r.result.clone(), logs).await?;
                     if job.is_flow_step {
-                        update_flow_status_after_job_completion(db, &job, true, r.result).await?;
+                        update_flow_status_after_job_completion(db, &job, true, r.result, metrics)
+                            .await?;
                     }
                 }
                 Err(e) => {
-                    let (_, output_map) = add_completed_job_error(db, &job, logs, e).await?;
+                    let (_, output_map) =
+                        add_completed_job_error(db, &job, logs, e, &metrics).await?;
                     if job.is_flow_step {
-                        update_flow_status_after_job_completion(db, &job, false, Some(output_map))
-                            .await?;
+                        update_flow_status_after_job_completion(
+                            db,
+                            &job,
+                            false,
+                            Some(output_map),
+                            metrics,
+                        )
+                        .await?;
                     }
                 }
             };

backend/src/worker_flow.rs

@@ -10,6 +10,7 @@ use crate::{
     },
     js_eval::{eval_timeout, EvalCreds},
     users::create_token_for_owner,
+    worker,
 };
 use async_recursion::async_recursion;
 use serde::{Deserialize, Serialize};
@@ -49,6 +50,7 @@ pub async fn update_flow_status_after_job_completion(
     job: &QueuedJob,
     success: bool,
     result: Option<Map<String, Value>>,
+    metrics: &worker::Metrics,
 ) -> error::Result<()> {
     tracing::debug!("HANDLE FLOW: {job:?} {success} {result:?}");
 
@@ -197,6 +199,7 @@ pub async fn update_flow_status_after_job_completion(
                     &flow_job,
                     "Unexpected error during flow chaining:\n".to_string(),
                     err,
+                    metrics,
                 )
                 .await;
                 true
@@ -217,9 +220,10 @@ pub async fn update_flow_status_after_job_completion(
         .await?;
 
         if flow_job.parent_job.is_some() {
-            return Ok(
-                update_flow_status_after_job_completion(db, &flow_job, success, result).await?,
-            );
+            return Ok(update_flow_status_after_job_completion(
+                db, &flow_job, success, result, metrics,
+            )
+            .await?);
         }
     }