trakt import

2 files changed, 72 insertions, 4 deletions

diff --git a/base/src/cache.rs b/base/src/cache.rs
index e580130..d1c3e4d 100644
--- a/base/src/cache.rs
+++ b/base/src/cache.rs
@@ -24,7 +24,10 @@ use std::{
     },
     time::Instant,
 };
-use tokio::sync::Mutex;
+use tokio::{
+    io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt},
+    sync::Mutex,
+};
 
 pub fn cache_location(seed: &[&str]) -> (usize, AssetLocation) {
     use sha2::Digest;
@@ -169,6 +172,67 @@ where
     Ok(object)
 }
 
+pub async fn async_cache_memory<Fun, Fut, T>(
+    seed: &[&str],
+    mut generate: Fun,
+) -> Result<Arc<T>, anyhow::Error>
+where
+    Fun: FnMut() -> Fut,
+    Fut: Future<Output = 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 = async_cache_file(seed, move |mut file| async move {
+        let object = generate().await?;
+        let data = bincode::encode_to_vec(&object, bincode::config::standard())
+            .context("encoding cache object")?;
+
+        file.write_all(&data).await?;
+
+        Ok(())
+    })
+    .await?;
+    let mut file = tokio::fs::File::open(location.path()).await?;
+    let mut data = Vec::new();
+    file.read_to_end(&mut data)
+        .await
+        .context("reading cache object")?;
+    let (object, _) = bincode::decode_from_slice::<T, _>(&data, bincode::config::standard())
+        .context("decoding cache object")?;
+    let object = Arc::new(object);
+    let size = file.stream_position().await? 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 {
diff --git a/base/src/lib.rs b/base/src/lib.rs
index 0001caa..48d3b37 100644
--- a/base/src/lib.rs
+++ b/base/src/lib.rs
@@ -19,9 +19,13 @@ use std::{fs::File, path::PathBuf, sync::LazyLock};
 pub static CONF: LazyLock<GlobalConfig> = LazyLock::new(|| {
     serde_yaml::from_reader(
         File::open(std::env::var("JELLYTHING_CONFIG").unwrap_or_else(|_| {
-            std::env::args().nth(1).expect(
-                "First argument or JELLYTHING_CONFIG must specify the configuration to use.",
-            )
+            if std::env::args().nth(0) == Some("jellything".to_string()) {
+                std::env::args().nth(1).expect(
+                    "First argument or JELLYTHING_CONFIG must specify the configuration to use.",
+                )
+            } else {
+                panic!("JELLYTHING_CONFIG variable is required.")
+            }
         }))
         .expect("config cannot be read"),
     )