From a79179c5c308ea3468d95df5cc4b286c174f29bb Mon Sep 17 00:00:00 2001
From: Cameron Cordes <cameronc.dev@gmail.com>
Date: Fri, 30 Apr 2021 23:53:10 -0400
Subject: [PATCH] Add Image and Video total gauges

---
 Cargo.lock  |  9 ++++++++
 Cargo.toml  |  2 ++
 src/main.rs | 66 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ded651d..0de7d06 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1160,9 +1160,11 @@ dependencies = [
  "hmac",
  "image",
  "jsonwebtoken",
+ "lazy_static",
  "log",
  "notify",
  "path-absolutize",
+ "prometheus",
  "rayon",
  "serde",
  "serde_json",
@@ -1753,10 +1755,17 @@ dependencies = [
  "fnv",
  "lazy_static",
  "parking_lot",
+ "protobuf",
  "regex",
  "thiserror",
 ]
 
+[[package]]
+name = "protobuf"
+version = "2.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "45604fc7a88158e7d514d8e22e14ac746081e7a70d7690074dd0029ee37458d6"
+
 [[package]]
 name = "quick-error"
 version = "1.2.3"
diff --git a/Cargo.toml b/Cargo.toml
index f1f92ec..df48e1e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -33,3 +33,5 @@ path-absolutize = "3.0.6"
 log="0.4"
 env_logger="0.8"
 actix-web-prom = "0.5.1"
+prometheus = "0.11"
+lazy_static = "1.1"
diff --git a/src/main.rs b/src/main.rs
index b549c60..d52700c 100644
--- a/src/main.rs
+++ b/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();