untested in-memory cache

5 files changed, 100 insertions, 4 deletions

diff --git a/Cargo.lock b/Cargo.lock
index d74439f..6889be0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1389,6 +1389,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "base64",
+ "bincode 2.0.0-rc.3",
  "jellycommon",
  "log",
  "serde",
diff --git a/base/Cargo.toml b/base/Cargo.toml
index 3081262..03b5055 100644
--- a/base/Cargo.toml
+++ b/base/Cargo.toml
@@ -12,3 +12,4 @@ sha2 = "0.10.8"
 base64 = "0.21.4"
 tokio = { workspace = true }
 anyhow = "1.0.75"
+bincode = "2.0.0-rc.3"
diff --git a/base/src/cache.rs b/base/src/cache.rs
index 2fa2680..d840c00 100644
--- a/base/src/cache.rs
+++ b/base/src/cache.rs
@@ -1,7 +1,20 @@
-use crate::AssetLocationExt;
+use crate::{AssetLocationExt, CONF};
+use anyhow::anyhow;
 use base64::Engine;
+use bincode::{Decode, Encode};
 use jellycommon::AssetLocation;
-use std::{future::Future, sync::LazyLock};
+use log::info;
+use std::{
+    any::Any,
+    collections::{BTreeMap, HashMap},
+    future::Future,
+    io::Seek,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc, LazyLock, RwLock,
+    },
+    time::Instant,
+};
 use tokio::sync::Mutex;
 
 pub fn cache_location(seed: &[&str]) -> (usize, AssetLocation) {
@@ -58,3 +71,82 @@ where
     Ok(location)
 }
 
+pub struct InMemoryCacheEntry {
+    size: usize,
+    last_access: Instant,
+    object: Arc<dyn Any + Send + Sync + 'static>,
+}
+pub static CACHE_IN_MEMORY_OBJECTS: LazyLock<RwLock<HashMap<AssetLocation, InMemoryCacheEntry>>> =
+    LazyLock::new(|| RwLock::new(HashMap::new()));
+pub static CACHE_IN_MEMORY_SIZE: AtomicUsize = AtomicUsize::new(0);
+
+pub fn cache_memory<Fun, T>(seed: &[&str], mut generate: Fun) -> Result<Arc<T>, anyhow::Error>
+where
+    Fun: FnMut() -> Result<T, anyhow::Error>,
+    T: Encode + Decode + Send + Sync + 'static,
+{
+    let (_, location) = cache_location(seed);
+    {
+        let mut g = CACHE_IN_MEMORY_OBJECTS.write().unwrap();
+        if let Some(entry) = g.get_mut(&location) {
+            entry.last_access = Instant::now();
+            let object = entry
+                .object
+                .clone()
+                .downcast::<T>()
+                .map_err(|_| anyhow!("inconsistent types for in-memory cache"))?;
+            return Ok(object);
+        }
+    }
+
+    let location = cache_file(seed, move |mut file| {
+        let object = generate()?;
+        bincode::encode_into_std_write(&object, &mut file, bincode::config::standard())?;
+        Ok(())
+    })?;
+    let mut file = std::fs::File::open(location.path())?;
+    let object = bincode::decode_from_std_read::<T, _, _>(&mut file, bincode::config::standard())?;
+    let object = Arc::new(object);
+    let size = file.stream_position()? as usize; // this is an approximation mainly since varint is used in bincode
+
+    {
+        let mut g = CACHE_IN_MEMORY_OBJECTS.write().unwrap();
+        g.insert(
+            location,
+            InMemoryCacheEntry {
+                size,
+                last_access: Instant::now(),
+                object: object.clone(),
+            },
+        );
+        CACHE_IN_MEMORY_SIZE.fetch_add(size, Ordering::Relaxed);
+    }
+
+    cleanup_cache();
+
+    Ok(object)
+}
+
+pub fn cleanup_cache() {
+    let current_size = CACHE_IN_MEMORY_SIZE.load(Ordering::Relaxed);
+    if current_size < CONF.max_in_memory_cache_size {
+        return;
+    }
+    info!("running cache eviction");
+    let mut g = CACHE_IN_MEMORY_OBJECTS.write().unwrap();
+
+    // TODO: if two entries have *exactly* the same size, only one of the will be remove; this is fine for now
+    let mut k = BTreeMap::new();
+    for (loc, entry) in g.iter() {
+        k.insert(entry.last_access.elapsed(), (loc.to_owned(), entry.size));
+    }
+    let mut reduction = 0;
+    for (loc, size) in k.values().rev().take(k.len().div_ceil(2)) {
+        g.remove(loc);
+        reduction += size;
+    }
+    CACHE_IN_MEMORY_SIZE.fetch_sub(reduction, Ordering::Relaxed);
+    drop(g);
+
+    info!("done");
+}
diff --git a/common/src/config.rs b/common/src/config.rs
index 467dc02..a9e3e0e 100644
--- a/common/src/config.rs
+++ b/common/src/config.rs
@@ -19,6 +19,7 @@ pub struct GlobalConfig {
     #[serde(default = "default::library_path")] pub library_path: PathBuf,
     #[serde(default = "default::cache_path")] pub cache_path: PathBuf,
     #[serde(default = "default::admin_username")] pub admin_username: String,
+    #[serde(default = "default::max_in_memory_cache_size")] pub max_in_memory_cache_size: usize,
     pub admin_password: String,
     #[serde(default)] pub cookie_key: Option<String>,
     #[serde(default)] pub session_key: Option<String>,
@@ -37,4 +38,5 @@ mod default {
     pub fn database_path() -> PathBuf { "data/database".into() }
     pub fn library_path() -> PathBuf { "data/library".into() }
     pub fn cache_path() -> PathBuf { "data/cache".into() }
+    pub fn max_in_memory_cache_size() -> usize { 50_000_000 }
 }
diff --git a/common/src/lib.rs b/common/src/lib.rs
index 6403469..a7e1c3f 100644
--- a/common/src/lib.rs
+++ b/common/src/lib.rs
@@ -6,10 +6,10 @@
 pub mod config;
 pub mod helpers;
 pub mod r#impl;
+pub mod jhls;
 pub mod seek_index;
 pub mod stream;
 pub mod user;
-pub mod jhls;
 
 #[cfg(feature = "rocket")]
 use rocket::{FromFormField, UriDisplayQuery};
@@ -56,7 +56,7 @@ pub struct RemoteImportOptions {
     #[serde(default)] pub prefix: Option<String>,
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize)]
+#[derive(Debug, Clone, Deserialize, Serialize, Hash, PartialEq, Eq)]
 #[serde(rename_all = "snake_case")]
 pub enum AssetLocation {
     Cache(PathBuf),