From 37cca40192d6e08e122c1e3c23f674d01e260ab3 Mon Sep 17 00:00:00 2001
From: JMARyA <jmarya@hydrar.de>
Date: Fri, 16 Aug 2024 12:19:19 +0200
Subject: [PATCH] cache

---
 Cargo.toml       |  4 ++++
 src/cache.rs     | 38 ++++++++++++++++++++++++++++++++++++++
 src/lib.rs       | 38 ++++++++++++++++++++++++++++++++++++++
 src/model/mod.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 124 insertions(+), 2 deletions(-)
 create mode 100644 src/cache.rs

diff --git a/Cargo.toml b/Cargo.toml
index a0a618b..d7418e1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,6 +3,10 @@ name = "mongod"
 version = "0.2.0"
 edition = "2021"
 
+[features]
+default = []
+cache = []
+
 [dependencies]
 chrono = "0.4.38"
 futures = "0.3.30"
diff --git a/src/cache.rs b/src/cache.rs
new file mode 100644
index 0000000..c636384
--- /dev/null
+++ b/src/cache.rs
@@ -0,0 +1,38 @@
+use std::collections::HashMap;
+
+// todo : implement cache feature
+#[derive(Debug)]
+pub struct ModelCache {
+    cache: HashMap<String, HashMap<String, mongodb::bson::Document>>,
+}
+
+impl ModelCache {
+    pub fn new() -> Self {
+        Self {
+            cache: HashMap::new(),
+        }
+    }
+
+    pub fn insert(&mut self, col: &str, id: &str, model: mongodb::bson::Document) {
+        let collection = self
+            .cache
+            .entry(col.to_string())
+            .or_insert_with(HashMap::new);
+        collection.insert(id.to_string(), model);
+    }
+
+    pub fn invalidate(&mut self, col: &str, id: &str) {
+        if let Some(collection) = self.cache.get_mut(col) {
+            collection.remove(id);
+            if collection.is_empty() {
+                self.cache.remove(col);
+            }
+        }
+    }
+
+    pub fn get(&self, col: &str, id: &str) -> Option<&mongodb::bson::Document> {
+        self.cache
+            .get(col)
+            .and_then(|collection| collection.get(id))
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index b5eda2c..b0f1ee2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,11 +5,49 @@ pub use model::valid::Validate;
 pub use model::Model;
 pub use mongod_derive as derive;
 pub use mongodb;
+#[cfg(feature = "cache")]
+pub mod cache;
+#[cfg(feature = "cache")]
+pub use cache::ModelCache;
 
 use tokio::sync::OnceCell;
 
+#[cfg(feature = "cache")]
+use std::sync::RwLock;
+
 pub static MONGO_CLIENT: OnceCell<mongodb::Client> = OnceCell::const_new();
 
+#[cfg(feature = "cache")]
+pub static DB_CACHE: OnceCell<RwLock<ModelCache>> = OnceCell::const_new();
+
+#[cfg(feature = "cache")]
+#[macro_export]
+macro_rules! cache_read {
+    () => {
+        if let Some(cache) = $crate::DB_CACHE.get() {
+            cache.read().unwrap()
+        } else {
+            let cache = std::sync::RwLock::new($crate::ModelCache::new());
+            $crate::DB_CACHE.set(cache).unwrap();
+            $crate::DB_CACHE.get().unwrap().read().unwrap()
+        }
+    };
+}
+
+#[cfg(feature = "cache")]
+#[macro_export]
+macro_rules! cache_write {
+    () => {
+        if let Some(cache) = $crate::DB_CACHE.get() {
+            cache.write().unwrap()
+        } else {
+            let cache = std::sync::RwLock::new($crate::ModelCache::new());
+            $crate::DB_CACHE.set(cache).unwrap();
+            $crate::DB_CACHE.get().unwrap().write().unwrap()
+        }
+    };
+}
+
 /// Get a `MongoDB` Client from the environment
 #[macro_export]
 macro_rules! get_mongo {
diff --git a/src/model/mod.rs b/src/model/mod.rs
index 1cfa7dc..90df8d3 100644
--- a/src/model/mod.rs
+++ b/src/model/mod.rs
@@ -8,7 +8,7 @@ use serde::de::DeserializeOwned;
 use serde_json::{Map, Value};
 use valid::Validate;
 
-use crate::{col, collect_results, get_mongo, id_of};
+use crate::{cache_read, cache_write, col, collect_results, get_mongo, id_of};
 
 pub mod historic;
 pub mod reference;
@@ -40,6 +40,15 @@ pub trait Model:
     where
         Self: Sync,
     {
+        #[cfg(feature = "cache")]
+        {
+            cache_write!().insert(
+                Self::collection_name(),
+                self.id(),
+                mongodb::bson::to_document(self).unwrap(),
+            );
+        }
+
         async {
             let db = get_mongo!();
             let collection = col!(db, Self::collection_name());
@@ -82,6 +91,12 @@ pub trait Model:
         async move {
             let db = get_mongo!();
             let collection = col!(db, Self::collection_name());
+
+            #[cfg(feature = "cache")]
+            {
+                cache_write!().invalidate(Self::collection_name(), id);
+            }
+
             collection.delete_one(id_of!(id), None).await
         }
     }
@@ -102,10 +117,26 @@ pub trait Model:
     #[must_use]
     fn get(id: &str) -> impl std::future::Future<Output = Option<Self>> {
         async move {
+            #[cfg(feature = "cache")]
+            {
+                if let Some(cached) = cache_read!().get(Self::collection_name(), id) {
+                    return mongodb::bson::from_document(cached.clone()).ok();
+                }
+            }
+
             let db = get_mongo!();
             let collection = col!(db, Self::collection_name());
             let doc = collection.find_one(id_of!(id), None).await.ok()??;
-            mongodb::bson::from_document(doc).ok()
+            let res = mongodb::bson::from_document(doc.clone()).ok();
+
+            #[cfg(feature = "cache")]
+            {
+                if res.is_some() {
+                    cache_write!().insert(Self::collection_name(), id, doc);
+                }
+            }
+
+            res
         }
     }
 
@@ -116,6 +147,11 @@ pub trait Model:
         mut part: serde_json::Value,
     ) -> impl std::future::Future<Output = Option<Self::Partial>> {
         async move {
+            #[cfg(feature = "cache")]
+            {
+                // todo : implement
+            }
+
             let db = get_mongo!();
             let collection = col!(db, Self::collection_name());
 
@@ -272,6 +308,12 @@ pub trait Model:
                     )
                     .await
                     .map_err(UpdateError::Database)?;
+
+                #[cfg(feature = "cache")]
+                {
+                    cache_write!().invalidate(Self::collection_name(), self.id());
+                }
+
                 return Ok(());
             }