Cargo.lock

@@ -326,6 +326,20 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "actix-web-prom"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec994b7dd52e686b9ae8f63bc485e6027fb44b936ed09bdf648495b9966c36ce"
+dependencies = [
+ "actix-http",
+ "actix-service",
+ "actix-web",
+ "futures",
+ "pin-project 1.0.6",
+ "prometheus",
+]
+
 [[package]]
 name = "actix_derive"
 version = "0.5.0"
@@ -1136,6 +1150,7 @@ dependencies = [
  "actix-multipart",
  "actix-rt",
  "actix-web",
+ "actix-web-prom",
  "bcrypt",
  "chrono",
  "diesel",
@@ -1145,9 +1160,11 @@ dependencies = [
  "hmac",
  "image",
  "jsonwebtoken",
+ "lazy_static",
  "log",
  "notify",
  "path-absolutize",
+ "prometheus",
  "rayon",
  "serde",
  "serde_json",
@@ -1728,6 +1745,27 @@ dependencies = [
  "unicode-xid",
 ]
 
+[[package]]
+name = "prometheus"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c8425533e7122f0c3cc7a37e6244b16ad3a2cc32ae7ac6276e2a75da0d9c200d"
+dependencies = [
+ "cfg-if 1.0.0",
+ "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"

Cargo.toml

@@ -32,3 +32,6 @@ notify = "4.0"
 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"

src/main.rs

@@ -3,13 +3,17 @@ extern crate diesel;
 extern crate rayon;
 
 use crate::auth::login;
+use actix_web_prom::PrometheusMetrics;
 use database::{DbError, DbErrorKind, FavoriteDao, SqliteFavoriteDao, SqliteUserDao, UserDao};
 use futures::stream::StreamExt;
-use std::io::prelude::*;
+use lazy_static::lazy_static;
+use prometheus::{self, IntGauge};
 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;
@@ -38,6 +42,19 @@ mod database;
 mod files;
 mod video;
 
+lazy_static! {
+    static ref IMAGE_GAUGE: IntGauge = IntGauge::new(
+        "imageserver_image_total",
+        "Count of the images on the server"
+    )
+    .unwrap();
+    static ref VIDEO_GAUGE: IntGauge = IntGauge::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);
@@ -316,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());
@@ -348,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);
+    VIDEO_GAUGE.set(video_count);
+}
+
+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<()> {
@@ -391,6 +443,9 @@ fn main() -> std::io::Result<()> {
                             }
                         }
                     }
+                    DebouncedEvent::Remove(_) => {
+                        update_media_counts(&PathBuf::from(env::var("BASE_PATH").unwrap()))
+                    }
                     _ => continue,
                 };
             }
@@ -404,6 +459,17 @@ fn main() -> std::io::Result<()> {
         stream_manager: Arc::new(act),
     });
 
+    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();
         let favorites_dao = SqliteFavoriteDao::new();
@@ -422,6 +488,7 @@ fn main() -> std::io::Result<()> {
             .app_data(app_data.clone())
             .data::<Box<dyn UserDao>>(Box::new(user_dao))
             .data::<Box<dyn FavoriteDao>>(Box::new(favorites_dao))
+            .wrap(prometheus.clone())
     })
     .bind(dotenv::var("BIND_URL").unwrap())?
     .bind("localhost:8088")?