Export job state metrics for Prometheus

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;
@@ -51,6 +53,63 @@ pub mod schedulers;
 mod states;
 
 const TTL_PIPELINE_CLEANUP_TIME: Duration = Duration::from_secs(60 * 60);
+const JOB_STATES: [&str; 6] = [
+    "running",
+    "transitioning",
+    "stopped",
+    "failed",
+    "user_failed",
+    "unknown",
+];
+
+lazy_static! {
+    static ref JOBS_BY_STATE: IntGaugeVec = register_int_gauge_vec!(
+        "arroyo_controller_jobs",
+        "Current number of jobs by operational state",
+        &["state"]
+    )
+    .unwrap();
+}
+
+fn metric_job_state(
+    state: Option<&str>,
+    running_desired: bool,
+    failure_domain: Option<&str>,
+) -> &'static str {
+    match (state.unwrap_or("Created"), running_desired) {
+        ("Failed", _) if failure_domain == Some("user") => "user_failed",
+        ("Failed", _) => "failed",
+        ("Running", true) => "running",
+        ("Created" | "Stopped" | "Finished", false) => "stopped",
+        (
+            "Created" | "Compiling" | "Scheduling" | "Running" | "Rescaling" | "CheckpointStopping"
+            | "Recovering" | "Restarting" | "Stopping" | "Stopped" | "Finishing" | "Finished"
+            | "Failing",
+            _,
+        ) => "transitioning",
+        _ => "unknown",
+    }
+}
+
+fn job_state_counts<'a>(
+    jobs: impl Iterator<Item = (Option<&'a str>, bool, Option<&'a str>)>,
+) -> HashMap<&'static str, i64> {
+    let mut counts = HashMap::new();
+    for (state, running_desired, failure_domain) in jobs {
+        *counts
+            .entry(metric_job_state(state, running_desired, failure_domain))
+            .or_default() += 1;
+    }
+    counts
+}
+
+fn update_job_state_metrics(counts: &HashMap<&'static str, i64>) {
+    for state in JOB_STATES {
+        JOBS_BY_STATE
+            .with_label_values(&[state])
+            .set(counts.get(state).copied().unwrap_or_default());
+    }
+}
 
 include!(concat!(env!("OUT_DIR"), "/controller-sql.rs"));
 
@@ -645,6 +704,15 @@ 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.stop == StopMode::none,
+                        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 +867,77 @@ impl ControllerServer {
         Ok(local_addr.port())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn metric_job_states_use_operational_buckets() {
+        assert_eq!(metric_job_state(Some("Running"), true, None), "running");
+        assert_eq!(
+            metric_job_state(Some("Running"), false, None),
+            "transitioning"
+        );
+        assert_eq!(metric_job_state(Some("Created"), false, None), "stopped");
+        assert_eq!(
+            metric_job_state(Some("Created"), true, None),
+            "transitioning"
+        );
+        assert_eq!(metric_job_state(Some("Stopped"), false, None), "stopped");
+        assert_eq!(
+            metric_job_state(Some("Stopped"), true, None),
+            "transitioning"
+        );
+        assert_eq!(metric_job_state(Some("Finished"), false, None), "stopped");
+        assert_eq!(
+            metric_job_state(Some("Finished"), true, None),
+            "transitioning"
+        );
+        assert_eq!(metric_job_state(None, false, None), "stopped");
+
+        for state in [
+            "Compiling",
+            "Scheduling",
+            "Rescaling",
+            "CheckpointStopping",
+            "Recovering",
+            "Restarting",
+            "Stopping",
+            "Finishing",
+            "Failing",
+        ] {
+            assert_eq!(metric_job_state(Some(state), true, None), "transitioning");
+        }
+
+        assert_eq!(
+            metric_job_state(Some("Failed"), false, Some("user")),
+            "user_failed"
+        );
+        assert_eq!(
+            metric_job_state(Some("Failed"), true, Some("internal")),
+            "failed"
+        );
+        assert_eq!(metric_job_state(Some("Unexpected"), true, None), "unknown");
+    }
+
+    #[test]
+    fn job_state_metrics_clear_absent_states() {
+        let counts = job_state_counts(
+            [
+                (Some("Running"), true, None),
+                (Some("Running"), true, None),
+                (Some("Unexpected"), true, None),
+            ]
+            .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(&["unknown"]).get(), 1);
+
+        update_job_state_metrics(&HashMap::new());
+        for state in JOB_STATES {
+            assert_eq!(JOBS_BY_STATE.with_label_values(&[state]).get(), 0);
+        }
+    }
+}

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);
+    }
+}