From ae216b24104e828a252b5ce0ea73a9987ad3772e Mon Sep 17 00:00:00 2001
From: Edward Shen <code@eddie.sh>
Date: Sun, 18 Apr 2021 17:38:33 -0400
Subject: [PATCH] add pruning

---
 src/cache/low_mem.rs |  7 +++++--
 src/cache/mod.rs     |  6 ++++++
 src/config.rs        |  8 +++++++-
 src/main.rs          | 33 ++++++++++++++++++++++-----------
 4 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/src/cache/low_mem.rs b/src/cache/low_mem.rs
index e7f637b..caf23ac 100644
--- a/src/cache/low_mem.rs
+++ b/src/cache/low_mem.rs
@@ -3,6 +3,7 @@
 use std::path::PathBuf;
 
 use async_trait::async_trait;
+use log::warn;
 use lru::LruCache;
 
 use super::{BoxedImageStream, Cache, CacheError, CacheKey, CacheStream, ImageMetadata};
@@ -25,8 +26,6 @@ impl LowMemCache {
     }
 }
 
-// todo: schedule eviction
-
 #[async_trait]
 impl Cache for LowMemCache {
     async fn get(
@@ -55,4 +54,8 @@ impl Cache for LowMemCache {
             .map(move |stream| (stream, self.on_disk.get(&key).unwrap()))
             .map_err(Into::into)
     }
+
+    async fn prune(&mut self) {
+        warn!("Trimming has not been implemented yet. Cache is unbounded!");
+    }
 }
diff --git a/src/cache/mod.rs b/src/cache/mod.rs
index 22e7f0d..5b097a0 100644
--- a/src/cache/mod.rs
+++ b/src/cache/mod.rs
@@ -10,6 +10,7 @@ use bytes::Bytes;
 use chrono::{DateTime, FixedOffset};
 use fs::FsStream;
 use futures::{Stream, StreamExt};
+use log::debug;
 use thiserror::Error;
 
 pub use fs::UpstreamError;
@@ -148,12 +149,17 @@ pub trait Cache: Send + Sync {
         &mut self,
         key: &CacheKey,
     ) -> Option<Result<(CacheStream, &ImageMetadata), CacheError>>;
+
     async fn put(
         &mut self,
         key: CacheKey,
         image: BoxedImageStream,
         metadata: ImageMetadata,
     ) -> Result<(CacheStream, &ImageMetadata), CacheError>;
+
+    async fn prune(&mut self) {
+        debug!("Would trim but cache does not implement trimming!");
+    }
 }
 
 pub enum CacheStream {
diff --git a/src/config.rs b/src/config.rs
index 6d936e9..6ecd731 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -35,7 +35,13 @@ pub struct CliArgs {
     /// reasons.
     #[clap(long, env = "ENABLE_SERVER_STRING", takes_value = false)]
     pub enable_server_string: bool,
-    #[clap(short, long, conflicts_with("memory-quota"), env = "LOW_MEMORY_MODE")]
+    #[clap(
+        short,
+        long,
+        conflicts_with("memory-quota"),
+        env = "LOW_MEMORY_MODE",
+        takes_value = false
+    )]
     pub low_memory: bool,
     #[clap(short, long, parse(from_occurrences))]
     pub verbose: usize,
diff --git a/src/main.rs b/src/main.rs
index 5d20e51..ba9b1c6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -130,24 +130,35 @@ async fn main() -> Result<(), std::io::Error> {
         }
     });
 
+    let cache: Box<dyn Cache> = if low_mem_mode {
+        Box::new(LowMemCache::new(disk_quota, cache_path.clone()))
+    } else {
+        Box::new(GenerationalCache::new(
+            memory_max_size,
+            disk_quota,
+            cache_path.clone(),
+        ))
+    };
+    let cache = Arc::new(Mutex::new(cache));
+    let cache1 = Arc::clone(&cache);
+
+    // Spawn periodic cache trimming
+    spawn(async move {
+        let mut interval = time::interval(Duration::from_secs(3 * 60));
+        loop {
+            interval.tick().await;
+            cache.lock().prune().await;
+        }
+    });
+
     // Start HTTPS server
     HttpServer::new(move || {
-        let cache: Box<dyn Cache> = if low_mem_mode {
-            Box::new(LowMemCache::new(disk_quota, cache_path.clone()))
-        } else {
-            Box::new(GenerationalCache::new(
-                memory_max_size,
-                disk_quota,
-                cache_path.clone(),
-            ))
-        };
-
         App::new()
             .service(routes::token_data)
             .service(routes::token_data_saver)
             .route("{tail:.*}", web::get().to(routes::default))
             .app_data(Data::from(Arc::clone(&data_1)))
-            .app_data(Data::new(Mutex::new(cache)))
+            .app_data(Data::from(Arc::clone(&cache1)))
     })
     .shutdown_timeout(60)
     .bind_rustls(format!("0.0.0.0:{}", port), tls_config)?