untested in-memory cache

1 files changed, 94 insertions, 2 deletions

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");
+}