diff --git a/src/main.rs b/src/main.rs index 5c90463..0139e07 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use std::{net::SocketAddr, path::{Path, PathBuf}, ffi::{OsStr, OsString}, fs::File, io::{BufReader, Seek, Cursor, BufRead}, env::args_os, os::unix::prelude::OsStrExt, cmp::Reverse}; +use std::{net::SocketAddr, path::{Path, PathBuf}, ffi::{OsStr, OsString}, fs::File, io::{BufReader, Seek, Cursor, BufRead}, env::args_os, os::unix::prelude::OsStrExt, cmp::Reverse, sync::Arc, collections::HashMap}; use anyhow::{Context, Result, anyhow}; use axum::{Router, routing::{get, post}, response::{IntoResponse, Redirect}, http::{StatusCode, header}, extract::{self, State}, Form, handler::Handler}; @@ -10,7 +10,7 @@ use error::AppError; use exif::Exif; use image::{imageops::FilterType, DynamicImage}; use serde::{Serialize, Deserialize}; -use tokio::task; +use tokio::{task, sync::RwLock}; use tower_http::{trace::{self, TraceLayer}, compression::CompressionLayer}; use tracing::Level; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; @@ -29,10 +29,12 @@ static GLOBAL: jemallocator::Jemalloc = jemallocator::Jemalloc; type TemplateEngine = Engine>; type ImageDir = PathBuf; type SecretKey = [u8; 64]; +type ImageCache = Arc, Vec)>>>; #[derive(Clone, extract::FromRef)] struct ApplicationState { engine: TemplateEngine, + image_cache: ImageCache, image_dir: ImageDir, secret_key: SecretKey, } @@ -74,6 +76,7 @@ async fn main() { .layer(session_layer) .with_state(ApplicationState { engine: Engine::from(jinja), + image_cache: Default::default(), image_dir: PathBuf::from(image_path), secret_key, }); @@ -158,6 +161,7 @@ async fn authenticate( async fn converted_image( extract::Path(encrypted_filename_hex): extract::Path, + State(image_cache): State, State(image_dir): State, State(secret_key): State, session: ReadableSession, @@ -175,18 +179,50 @@ async fn converted_image( let chacha = XChaCha20Poly1305::new(chacha_key); chacha.decrypt(nonce, ciphertext) }.with_context(|| format!("Could not decrypt filename {}", encrypted_filename_hex))?; + let image_name = OsStr::from_bytes(&image_name); - let image_path = image_dir.join(OsStr::from_bytes(&image_name)); + let image_path = image_dir.join(image_name); image_path.exists() .then_some(()) .ok_or(anyhow!("Requested image not found!")) .context(StatusCode::NOT_FOUND)?; - let image_buffer = task::spawn_blocking(move || { - convert_image(&image_path) - .with_context(|| format!("Could not convert image {:?}", image_path)) - }).await??; + // Check if we have the file already in cache + let image_modified: Option> = std::fs::metadata(&image_path) + .and_then(|m| m.modified()) + .map(DateTime::from) + .ok(); + let cached_buffer = if let Some(image_modified) = image_modified { + image_cache + .read().await + .get(image_name) + .filter(|(cache_modified, _)| *cache_modified == image_modified) + .map(|(_, image_buffer)| image_buffer) + .cloned() + } else { None }; + + let image_buffer = match cached_buffer { + Some(image_buffer) => { + tracing::debug!("Read {:?} from cache", image_name); + image_buffer + } + None => { + let image_buffer = task::spawn_blocking(move || { + convert_image(&image_path) + .with_context(|| format!("Could not convert image {:?}", image_path)) + }).await??; + + if let Some(image_modified) = image_modified { + tracing::debug!("Add {:?} ({}k) to cache", image_name, image_buffer.len() / 1024); + image_cache + .write().await + .insert(image_name.to_owned(), (image_modified, image_buffer.clone())); + } + + image_buffer + } + }; Ok(( [