Add Image and Video total gauges

src/main.rs

@@ -6,11 +6,14 @@ use crate::auth::login;
 use actix_web_prom::PrometheusMetrics;
 use database::{DbError, DbErrorKind, FavoriteDao, SqliteFavoriteDao, SqliteUserDao, UserDao};
 use futures::stream::StreamExt;
+use lazy_static::lazy_static;
+use prometheus::{self, Gauge};
 use std::path::{Path, PathBuf};
 use std::sync::mpsc::channel;
 use std::sync::Arc;
 use std::{collections::HashMap, io::prelude::*};
 use std::{env, fs::File};
+use walkdir::{DirEntry, WalkDir};
 
 use actix::prelude::*;
 use actix_files::NamedFile;
@@ -39,6 +42,19 @@ mod database;
 mod files;
 mod video;
 
+lazy_static! {
+    static ref IMAGE_GAUGE: Gauge = Gauge::new(
+        "imageserver_image_total",
+        "Count of the images on the server"
+    )
+    .unwrap();
+    static ref VIDEO_GAUGE: Gauge = Gauge::new(
+        "imageserver_video_total",
+        "Count of the videos on the server"
+    )
+    .unwrap();
+}
+
 #[get("/photos")]
 async fn list_photos(_claims: Claims, req: Query<ThumbnailRequest>) -> impl Responder {
     info!("{}", req.path);
@@ -317,7 +333,7 @@ fn create_thumbnails() {
                     generate_video_thumbnail(entry.path(), &thumb_path);
                     false
                 } else {
-                    ext == "jpg" || ext == "jpeg" || ext == "png" || ext == "nef"
+                    is_image(entry)
                 }
             } else {
                 error!("Unable to get extension for file: {:?}", entry.path());
@@ -349,7 +365,42 @@ fn create_thumbnails() {
         })
         .for_each(drop);
 
-    debug!("Finished");
+    debug!("Finished making thumbnails");
+
+    update_media_counts(&images);
+}
+
+fn update_media_counts(media_dir: &PathBuf) {
+    let mut image_count = 0;
+    let mut video_count = 0;
+    for ref entry in WalkDir::new(media_dir).into_iter().filter_map(|e| e.ok()) {
+        if is_image(entry) {
+            image_count += 1;
+        } else if is_video(entry) {
+            video_count += 1;
+        }
+    }
+
+    IMAGE_GAUGE.set(image_count as f64);
+    VIDEO_GAUGE.set(video_count as f64);
+}
+
+fn is_image(entry: &DirEntry) -> bool {
+    entry
+        .path()
+        .extension()
+        .and_then(|ext| ext.to_str())
+        .map(|ext| ext == "jpg" || ext == "jpeg" || ext == "png" || ext == "nef")
+        .unwrap_or(false)
+}
+
+fn is_video(entry: &DirEntry) -> bool {
+    entry
+        .path()
+        .extension()
+        .and_then(|ext| ext.to_str())
+        .map(|ext| ext == "mp4" || ext == "mov")
+        .unwrap_or(false)
 }
 
 fn main() -> std::io::Result<()> {
@@ -392,6 +443,9 @@ fn main() -> std::io::Result<()> {
                             }
                         }
                     }
+                    DebouncedEvent::Remove(_) => {
+                        update_media_counts(&PathBuf::from(env::var("BASE_PATH").unwrap()))
+                    }
                     _ => continue,
                 };
             }
@@ -407,6 +461,14 @@ fn main() -> std::io::Result<()> {
 
     let labels = HashMap::new();
     let prometheus = PrometheusMetrics::new("", Some("/metrics"), Some(labels));
+    prometheus
+        .registry
+        .register(Box::new(IMAGE_GAUGE.clone()))
+        .unwrap();
+    prometheus
+        .registry
+        .register(Box::new(VIDEO_GAUGE.clone()))
+        .unwrap();
 
     HttpServer::new(move || {
         let user_dao = SqliteUserDao::new();