Add simple in-memory cache for encoded images

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<minijinja::Environment<'static>>;
 type ImageDir = PathBuf;
 type SecretKey = [u8; 64];
+type ImageCache = Arc<RwLock<HashMap<OsString, (DateTime<Utc>, Vec<u8>)>>>;
 
 #[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<String>,
+    State(image_cache): State<ImageCache>,
     State(image_dir): State<ImageDir>,
     State(secret_key): State<SecretKey>,
     session: ReadableSession,
@@ -175,19 +179,51 @@ 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)?;
 
+    // Check if we have the file already in cache
+    let image_modified: Option<DateTime<Utc>> = 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((
         [
             (header::CONTENT_TYPE, "image/webp"),