From a713c20aa45cf3ba1803ea59f2e991c23d267dae Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@protonmail.com>
Date: Sun, 28 Nov 2021 16:48:04 +0100
Subject: [PATCH] Add size limit to cache buffer

Signed-off-by: trivernis <trivernis@protonmail.com>
---
 .../src/tauri_plugin/commands/file.rs         |  2 -
 mediarepo-api/src/tauri_plugin/mod.rs         |  3 ++
 mediarepo-api/src/tauri_plugin/state.rs       | 37 +++++++++++++++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/mediarepo-api/src/tauri_plugin/commands/file.rs b/mediarepo-api/src/tauri_plugin/commands/file.rs
index 98753bd..f3d1c20 100644
--- a/mediarepo-api/src/tauri_plugin/commands/file.rs
+++ b/mediarepo-api/src/tauri_plugin/commands/file.rs
@@ -104,7 +104,6 @@ pub async fn read_file(
     api_state: ApiAccess<'_>,
     buffer_state: BufferAccess<'_>,
     hash: String,
-    mime_type: String,
 ) -> PluginResult<Vec<u8>> {
     if let Some(buffer) = buffer_state.get_entry(&hash) {
         Ok(buffer.buf)
@@ -114,7 +113,6 @@ pub async fn read_file(
             .file
             .read_file(FileIdentifier::Hash(hash.clone()))
             .await?;
-        buffer_state.add_entry(hash, mime_type, content.clone());
 
         Ok(content)
     }
diff --git a/mediarepo-api/src/tauri_plugin/mod.rs b/mediarepo-api/src/tauri_plugin/mod.rs
index 97bc43b..c7b06b5 100644
--- a/mediarepo-api/src/tauri_plugin/mod.rs
+++ b/mediarepo-api/src/tauri_plugin/mod.rs
@@ -16,6 +16,8 @@ mod utils;
 
 use commands::*;
 
+const MAX_BUFFER_SIZE: usize = 2 * 1024 * 1024 * 1024;
+
 pub fn register_plugin<R: Runtime>(builder: Builder<R>) -> Builder<R> {
     let repo_plugin = MediarepoPlugin::new();
 
@@ -84,6 +86,7 @@ impl<R: Runtime> Plugin<R> for MediarepoPlugin<R> {
         thread::spawn(move || loop {
             thread::sleep(Duration::from_secs(10));
             buffer_state.clear_expired();
+            buffer_state.trim_to_size(MAX_BUFFER_SIZE);
         });
 
         Ok(())
diff --git a/mediarepo-api/src/tauri_plugin/state.rs b/mediarepo-api/src/tauri_plugin/state.rs
index 6df107d..890437a 100644
--- a/mediarepo-api/src/tauri_plugin/state.rs
+++ b/mediarepo-api/src/tauri_plugin/state.rs
@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::mem;
+use std::ops::Deref;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -120,6 +121,42 @@ impl BufferState {
             }
         }
     }
+
+    /// Trims the buffer to the given target size
+    pub fn trim_to_size(&self, target_size: usize) {
+        let mut size = self.get_size();
+        if size < target_size {
+            return;
+        }
+
+        let mut keys: Vec<String> = {
+            let buffer = self.buffer.read();
+            buffer.keys().cloned().collect()
+        };
+        keys.reverse();
+
+        while size > target_size && keys.len() > 0 {
+            let key = keys.pop().unwrap();
+            let mut buffers = self.buffer.write();
+
+            if let Some(entry) = buffers.remove(&key) {
+                size -= entry.lock().buf.len();
+            }
+        }
+    }
+
+    /// Calculates the size of the whole buffer
+    pub fn get_size(&self) -> usize {
+        let buffer = self.buffer.read();
+        let mut size = 0;
+
+        for value in buffer.deref().values() {
+            let value = value.lock();
+            size += value.buf.len();
+        }
+
+        size
+    }
 }
 
 pub struct AppState {