fix CDN invalidation queue metric when the queue is empty

Author: syphar

src/cdn.rs

@@ -470,9 +470,19 @@ pub(crate) fn queued_or_active_crate_invalidations(
 /// Return the count of queued or active invalidations, per distribution id
 pub(crate) fn queued_or_active_crate_invalidation_count_by_distribution(
     conn: &mut impl postgres::GenericClient,
+    config: &Config,
 ) -> Result<HashMap<String, i64>> {
-    Ok(conn
-        .query(
+    let mut result: HashMap<String, i64> = HashMap::from_iter(
+        config
+            .cloudfront_distribution_id_web
+            .iter()
+            .chain(config.cloudfront_distribution_id_static.iter())
+            .cloned()
+            .map(|id| (id, 0)),
+    );
+
+    result.extend(
+        conn.query(
             r#"
              SELECT
                 cdn_distribution_id, 
@@ -482,8 +492,10 @@ pub(crate) fn queued_or_active_crate_invalidation_count_by_distribution(
             &[],
         )?
         .iter()
-        .map(|row| (row.get(0), row.get(1)))
-        .collect())
+        .map(|row| (row.get(0), row.get(1))),
+    );
+
+    Ok(result)
 }
 
 #[cfg(test)]
@@ -527,6 +539,27 @@ mod tests {
         })
     }
 
+    #[test]
+    fn invalidation_counts_are_zero_with_empty_queue() {
+        crate::test::wrapper(|env| {
+            env.override_config(|config| {
+                config.cloudfront_distribution_id_web = Some("distribution_id_web".into());
+                config.cloudfront_distribution_id_static = Some("distribution_id_static".into());
+            });
+
+            let config = env.config();
+            let mut conn = env.db().conn();
+            assert!(queued_or_active_crate_invalidations(&mut *conn)?.is_empty());
+
+            let counts =
+                queued_or_active_crate_invalidation_count_by_distribution(&mut *conn, &config)?;
+            assert_eq!(counts.len(), 2);
+            assert_eq!(*counts.get("distribution_id_web").unwrap(), 0);
+            assert_eq!(*counts.get("distribution_id_static").unwrap(), 0);
+            Ok(())
+        })
+    }
+
     #[test]
     fn invalidate_a_crate() {
         crate::test::wrapper(|env| {
@@ -536,6 +569,7 @@ mod tests {
             });
 
             let cdn = env.cdn();
+            let config = env.config();
             let mut conn = env.db().conn();
             assert!(queued_or_active_crate_invalidations(&mut *conn)?.is_empty());
 
@@ -574,7 +608,8 @@ mod tests {
                 ]
             );
 
-            let counts = queued_or_active_crate_invalidation_count_by_distribution(&mut *conn)?;
+            let counts =
+                queued_or_active_crate_invalidation_count_by_distribution(&mut *conn, &config)?;
             assert_eq!(counts.len(), 2);
             assert_eq!(*counts.get("distribution_id_web").unwrap(), 2);
             assert_eq!(*counts.get("distribution_id_static").unwrap(), 1);

src/metrics/mod.rs

@@ -2,7 +2,7 @@
 mod macros;
 
 use self::macros::MetricFromOpts;
-use crate::{cdn, db::Pool, target::TargetAtom, BuildQueue};
+use crate::{cdn, db::Pool, target::TargetAtom, BuildQueue, Config};
 use anyhow::Error;
 use dashmap::DashMap;
 use prometheus::proto::MetricFamily;
@@ -185,6 +185,7 @@ impl Metrics {
         &self,
         pool: &Pool,
         queue: &BuildQueue,
+        config: &Config,
     ) -> Result<Vec<MetricFamily>, Error> {
         self.idle_db_connections.set(pool.idle_connections() as i64);
         self.used_db_connections.set(pool.used_connections() as i64);
@@ -205,7 +206,7 @@ impl Metrics {
 
         let mut conn = pool.get()?;
         for (distribution_id, count) in
-            cdn::queued_or_active_crate_invalidation_count_by_distribution(&mut *conn)?
+            cdn::queued_or_active_crate_invalidation_count_by_distribution(&mut *conn, config)?
         {
             self.queued_cdn_invalidations_by_distribution
                 .with_label_values(&[&distribution_id])

src/web/metrics.rs

@@ -1,6 +1,6 @@
 use crate::{
     db::Pool, metrics::duration_to_seconds, utils::spawn_blocking, web::error::AxumResult,
-    BuildQueue, Metrics,
+    BuildQueue, Config, Metrics,
 };
 use anyhow::Context as _;
 use axum::{
@@ -21,10 +21,11 @@ use tracing::debug;
 
 pub(super) async fn metrics_handler(
     Extension(pool): Extension<Pool>,
+    Extension(config): Extension<Arc<Config>>,
     Extension(metrics): Extension<Arc<Metrics>>,
     Extension(queue): Extension<Arc<BuildQueue>>,
 ) -> AxumResult<impl IntoResponse> {
-    let families = spawn_blocking(move || metrics.gather(&pool, &queue)).await?;
+    let families = spawn_blocking(move || metrics.gather(&pool, &queue, &config)).await?;
 
     let mut buffer = Vec::new();
     TextEncoder::new()
@@ -199,7 +200,8 @@ mod tests {
             }
 
             // this shows what the routes were *actually* recorded as, making it easier to update ROUTES if the name changes.
-            let metrics_serialized = metrics.gather(&env.pool()?, &env.build_queue())?;
+            let metrics_serialized =
+                metrics.gather(&env.pool()?, &env.build_queue(), &env.config())?;
             let all_routes_visited = metrics_serialized
                 .iter()
                 .find(|x| x.get_name() == "docsrs_routes_visited")