diff --git a/Cargo.lock b/Cargo.lock index 3c1e301..0de7d06 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index a8f3ed8..df48e1e 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/main.rs b/src/main.rs index 7718412..16ee93b 100644 --- a/src/main.rs +++ b/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) -> 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::new(user_dao)) .data::>(Box::new(favorites_dao)) + .wrap(prometheus.clone()) }) .bind(dotenv::var("BIND_URL").unwrap())? .bind("localhost:8088")?