From c30503784332a3004ecbb0243efdb22b52ec3aa8 Mon Sep 17 00:00:00 2001
From: Louis-Marie Baer <lm@baermail.fr>
Date: Fri, 14 Jun 2024 15:26:12 +0200
Subject: [PATCH] feat: limit cache size by configuration

---
 Cargo.toml    |  1 +
 README.md     |  2 +-
 src/config.rs |  2 ++
 src/main.rs   | 17 +++++++++++++++++
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 4b3d968..2f4970b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,3 +20,4 @@ uuid = {version="1.8", features=["v4", "fast-rng"]}
 nohash = "0.2"
 derive_more = {version="0.99", default-features=false, features=["deref", "deref_mut"]}
 enclose = "1.2"
+typesize = "0.1"
diff --git a/README.md b/README.md
index bdd6ea2..069d489 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@ Mnemosyne is a http request caching in memory proxy API.
 ## Status of development
 Work in progress, not functional.
 ### TODO
-- [ ] allows to limit cache by size
+- [x] allows to limit cache by size
 - [ ] remove allocation when possible
 - [ ] organize code in modules
 - [ ] tracing
diff --git a/src/config.rs b/src/config.rs
index dc7a8c8..35244a1 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -38,4 +38,6 @@ impl Default for Config {
 pub struct CacheConfig {
     /// cache expiration after last request
     pub expiration: Duration,
+    /// in megabytes, the maximum size of memory the cache can take.
+    pub size_limit: u64,
 }
diff --git a/src/main.rs b/src/main.rs
index 3c50bb0..574bf40 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,6 +18,7 @@ use reqwest::{
 };
 use std::{str::FromStr, sync::Arc};
 use tokio::{spawn, sync::Mutex};
+use typesize::TypeSize;
 use url::Url;
 use uuid::Uuid;
 
@@ -45,6 +46,22 @@ impl Cache {
             MokaCache::builder()
                 .name("mnemosyne")
                 .time_to_idle(config.cache.expiration)
+                .weigher(
+                    |_key: &Uuid, (s, h, b): &(StatusCode, HeaderMap, Bytes)| -> u32 {
+                        let s = s.to_string().get_size() as u32;
+                        let h = h.iter().fold(0, |acc, x| {
+                            acc + (x.0.to_string().get_size()
+                                + x.1.to_str().unwrap().to_string().get_size())
+                                as u32
+                        });
+                        let b = b.len() as u32;
+                        // note that the size overhead of the index cache is not taken into account.
+                        // could take about 100B per entry.
+                        s + h + b
+                    },
+                )
+                // This cache will hold up to 32MiB of values.
+                .max_capacity(config.cache.size_limit * 1024 * 1024)
                 .build_with_hasher(ahash::RandomState::new()),
         )
     }