Export job state metrics for Prometheus (#1114)

Author: QnJ1c2kNCg

crates/arroyo-controller/src/lib.rs

@@ -28,6 +28,8 @@ use arroyo_server_common::wrap_start;
 use arroyo_types::{MachineId, PipelineId, WorkerId, from_micros};
 use arroyo_worker::job_controller::job_metrics::JobMetrics;
 use cornucopia_async::DatabaseSource;
+use lazy_static::lazy_static;
+use prometheus::{IntGaugeVec, register_int_gauge_vec};
 use states::{Created, State, StateMachine};
 use std::collections::{HashMap, HashSet};
 use std::env;
@@ -52,6 +54,43 @@ mod states;
 
 const TTL_PIPELINE_CLEANUP_TIME: Duration = Duration::from_secs(60 * 60);
 
+lazy_static! {
+    static ref JOBS_BY_STATE: IntGaugeVec = register_int_gauge_vec!(
+        "arroyo_controller_jobs",
+        "Current number of jobs by controller state",
+        &["state"]
+    )
+    .unwrap();
+}
+
+fn metric_job_state<'a>(state: Option<&'a str>, failure_domain: Option<&str>) -> &'a str {
+    let state = state.unwrap_or("Created");
+    if state == "Failed" && failure_domain == Some("user") {
+        "UserFailed"
+    } else {
+        state
+    }
+}
+
+fn job_state_counts<'a>(
+    jobs: impl Iterator<Item = (Option<&'a str>, Option<&'a str>)>,
+) -> HashMap<&'a str, i64> {
+    let mut counts = HashMap::new();
+    for (state, failure_domain) in jobs {
+        *counts
+            .entry(metric_job_state(state, failure_domain))
+            .or_default() += 1;
+    }
+    counts
+}
+
+fn update_job_state_metrics(counts: &HashMap<&str, i64>) {
+    JOBS_BY_STATE.reset();
+    for (state, count) in counts {
+        JOBS_BY_STATE.with_label_values(&[state]).set(*count);
+    }
+}
+
 include!(concat!(env!("OUT_DIR"), "/controller-sql.rs"));
 
 use crate::schedulers::{ManualScheduler, NodeScheduler, ProcessScheduler, Scheduler};
@@ -645,6 +684,12 @@ impl ControllerServer {
             while !token.is_cancelled() {
                 let client = db.client().await?;
                 let res = queries::controller_queries::fetch_all_jobs(&client).await?;
+                let state_counts = job_state_counts(
+                    res.iter()
+                        .map(|p| (p.state.as_deref(), p.failure_domain.as_deref())),
+                );
+                update_job_state_metrics(&state_counts);
+
                 for p in res {
                     let id = Arc::new(p.id);
                     let config = JobConfig {
@@ -799,3 +844,39 @@ impl ControllerServer {
         Ok(local_addr.port())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use prometheus::core::Collector;
+
+    use super::*;
+
+    #[test]
+    fn metric_job_states_preserve_raw_states() {
+        for state in ["Created", "Running", "Finished", "Unexpected"] {
+            assert_eq!(metric_job_state(Some(state), None), state);
+        }
+        assert_eq!(metric_job_state(None, None), "Created");
+
+        assert_eq!(metric_job_state(Some("Failed"), Some("user")), "UserFailed");
+        assert_eq!(metric_job_state(Some("Failed"), Some("internal")), "Failed");
+    }
+
+    #[test]
+    fn job_state_metrics_clear_absent_states() {
+        let counts = job_state_counts(
+            [
+                (Some("Running"), None),
+                (Some("Running"), None),
+                (Some("Failed"), Some("user")),
+            ]
+            .into_iter(),
+        );
+        update_job_state_metrics(&counts);
+        assert_eq!(JOBS_BY_STATE.with_label_values(&["Running"]).get(), 2);
+        assert_eq!(JOBS_BY_STATE.with_label_values(&["UserFailed"]).get(), 1);
+
+        update_job_state_metrics(&HashMap::new());
+        assert!(JOBS_BY_STATE.collect()[0].get_metric().is_empty());
+    }
+}

crates/arroyo-rpc/src/config.rs

@@ -524,6 +524,9 @@ pub struct AdminConfig {
 
     #[serde(default)]
     pub auth_mode: ApiAuthMode,
+
+    #[serde(default)]
+    pub allow_unauthenticated_metrics: bool,
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone, Default)]
@@ -1095,6 +1098,11 @@ mod tests {
             jail.set_env("ARROYO__ADMIN__HTTP_PORT", 9111);
             let config: Config = load_config(&[]).extract().unwrap();
             assert_eq!(config.admin.http_port, 9111);
+            assert!(!config.admin.allow_unauthenticated_metrics);
+
+            jail.set_env("ARROYO__ADMIN__ALLOW_UNAUTHENTICATED_METRICS", true);
+            let config: Config = load_config(&[]).extract().unwrap();
+            assert!(config.admin.allow_unauthenticated_metrics);
 
             Ok(())
         });

crates/arroyo-server-common/Cargo.toml

@@ -53,3 +53,5 @@ serde = { workspace = true }
 [build-dependencies]
 vergen = { version = "8.0.0", features = ["build", "cargo", "git", "gitcl"] }
 
+[dev-dependencies]
+tower = { workspace = true, features = ["util"] }

crates/arroyo-server-common/src/lib.rs

@@ -344,28 +344,48 @@ fn require_profiling_activated(
     }
 }
 
-pub async fn start_admin_server(service: &str) -> anyhow::Result<()> {
-    let config = config();
-    let addr = SocketAddr::new(config.admin.bind_address, config.admin.http_port);
-
+fn admin_router(
+    service: &str,
+    auth_mode: &ApiAuthMode,
+    allow_unauthenticated_metrics: bool,
+) -> Router {
     let state = Arc::new(AdminState {
         name: format!("arroyo-{service}"),
     });
-    let mut app = Router::new()
-        .route("/status", get(status))
+    let mut protected = Router::new()
         .route("/name", get(root))
-        .route("/metrics", get(metrics))
         .route("/metrics.pb", get(metrics_proto))
         .route("/details", get(details))
         .route("/config", get(config_route))
         .route("/debug/pprof/heap", get(handle_get_heap))
         .route("/debug/pprof/profile", get(handle_get_profile))
         .with_state(state);
 
-    if let ApiAuthMode::StaticApiKey { api_key } = &config.admin.auth_mode {
-        app = app.layer(ValidateRequestHeaderLayer::bearer(api_key));
+    if !allow_unauthenticated_metrics {
+        protected = protected.route("/metrics", get(metrics));
+    }
+
+    if let ApiAuthMode::StaticApiKey { api_key } = auth_mode {
+        protected = protected.layer(ValidateRequestHeaderLayer::bearer(api_key));
     };
 
+    // /status is always reachable without auth (e.g. for liveness probes).
+    let mut public = Router::new().route("/status", get(status));
+    if allow_unauthenticated_metrics {
+        public = public.route("/metrics", get(metrics));
+    }
+    public.merge(protected)
+}
+
+pub async fn start_admin_server(service: &str) -> anyhow::Result<()> {
+    let config = config();
+    let addr = SocketAddr::new(config.admin.bind_address, config.admin.http_port);
+    let app = admin_router(
+        service,
+        &config.admin.auth_mode,
+        config.admin.allow_unauthenticated_metrics,
+    );
+
     let tls_config =
         tls::create_http_tls_config(&config.admin.auth_mode, &config.admin.tls).await?;
     if let Some(tls_config) = tls_config {
@@ -535,3 +555,45 @@ pub async fn wrap_start(
         )
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use axum::body::Body;
+    use axum::http::Request;
+    use tower::ServiceExt;
+
+    use super::*;
+
+    fn static_auth() -> ApiAuthMode {
+        toml::from_str(
+            r#"
+                type = "static-api-key"
+                api-key = "secret"
+            "#,
+        )
+        .unwrap()
+    }
+
+    #[tokio::test]
+    async fn unauthenticated_metrics_are_configurable() {
+        let response = admin_router("test", &static_auth(), false)
+            .oneshot(Request::get("/metrics").body(Body::empty()).unwrap())
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
+
+        let app = admin_router("test", &static_auth(), true);
+        let response = app
+            .clone()
+            .oneshot(Request::get("/metrics").body(Body::empty()).unwrap())
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let response = app
+            .oneshot(Request::get("/details").body(Body::empty()).unwrap())
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
+    }
+}