add cache

src/cache.rs

@@ -0,0 +1,116 @@
+use std::collections::HashMap;
+
+use rocket::tokio::sync::RwLock;
+
+#[macro_export]
+macro_rules! use_api_cache {
+    ($route:literal, $id:ident, $cache:ident) => {
+        if let Some(ret) = $cache.get_only($route, $id).await {
+            return Ok(serde_json::from_str(&ret).unwrap());
+        }
+    };
+}
+
+pub struct RouteCache {
+    inner: RwLock<HashMap<String, HashMap<String, Option<String>>>>,
+}
+
+impl RouteCache {
+    pub fn new() -> Self {
+        Self {
+            inner: RwLock::new(HashMap::new()),
+        }
+    }
+
+    pub async fn get<F, Fut>(&self, route: &str, id: &str, generator: F) -> String
+    where
+        F: FnOnce() -> Fut,
+        Fut: std::future::Future<Output = String>,
+    {
+        {
+            // Try to get a read lock first.
+            let lock = self.inner.read().await;
+            if let Some(inner_map) = lock.get(route) {
+                if let Some(cached_value) = inner_map.get(id) {
+                    log::trace!("Using cached value for {route} / {id}");
+                    return cached_value.clone().unwrap();
+                }
+            }
+        }
+
+        // If the value was not found, acquire a write lock to insert the computed value.
+        let mut lock = self.inner.write().await;
+
+        log::trace!("Computing value for {route} / {id}");
+        let computed = generator().await;
+
+        lock.entry(route.to_string())
+            .or_insert_with(HashMap::new)
+            .insert(id.to_string(), Some(computed.clone()));
+
+        computed
+    }
+
+    pub async fn get_only(&self, route: &str, id: &str) -> Option<String> {
+        let lock = self.inner.read().await;
+        if let Some(inner_map) = lock.get(route) {
+            if let Some(cached_value) = inner_map.get(id) {
+                log::trace!("Using cached value for {route} / {id}");
+                return cached_value.clone();
+            }
+        }
+
+        None
+    }
+
+    pub async fn get_option<F, Fut>(&self, route: &str, id: &str, generator: F) -> Option<String>
+    where
+        F: FnOnce() -> Fut,
+        Fut: std::future::Future<Output = Option<String>>,
+    {
+        {
+            // Try to get a read lock first.
+            let lock = self.inner.read().await;
+            if let Some(inner_map) = lock.get(route) {
+                if let Some(cached_value) = inner_map.get(id) {
+                    log::trace!("Using cached value for {route} / {id}");
+                    return cached_value.clone();
+                }
+            }
+        }
+
+        // If the value was not found, acquire a write lock to insert the computed value.
+        let mut lock = self.inner.write().await;
+
+        log::trace!("Computing value for {route} / {id}");
+        let computed = generator().await;
+
+        lock.entry(route.to_string())
+            .or_insert_with(HashMap::new)
+            .insert(id.to_string(), computed.clone());
+
+        computed
+    }
+
+    pub async fn insert(&self, route: &str, id: &str, value: String) {
+        let mut lock = self.inner.write().await;
+
+        log::trace!("Inserting value for {route} / {id}");
+
+        lock.entry(route.to_string())
+            .or_insert_with(HashMap::new)
+            .insert(id.to_string(), Some(value));
+    }
+
+    pub async fn invalidate(&self, route: &str, id: &str) {
+        let mut lock = self.inner.write().await;
+        if let Some(inner_map) = lock.get_mut(route) {
+            inner_map.remove(id);
+
+            // If the inner map is empty, remove the route entry as well.
+            if inner_map.is_empty() {
+                lock.remove(route);
+            }
+        }
+    }
+}