From db036500650862e55b786167654378bc680a047d Mon Sep 17 00:00:00 2001
From: Fl1tzi <git@fl1tzi.com>
Date: Mon, 26 Feb 2024 01:22:02 +0100
Subject: [PATCH] major memory improvement by moving the image cache to fs

---
 Cargo.toml             |   5 +-
 src/device.rs          |  25 ++-----
 src/image_rendering.rs |   8 ++-
 src/main.rs            |  10 +++
 src/modules.rs         | 147 +++++++++++++++++++++--------------------
 src/modules/blank.rs   |   6 +-
 src/modules/counter.rs |   8 +--
 src/modules/image.rs   |  27 +++-----
 src/modules/space.rs   |   6 +-
 9 files changed, 115 insertions(+), 127 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 9fe9511..bc69a6b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,5 +19,6 @@ rusttype = "0.9.3"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 font-loader = "0.11.0"
-clru = "0.6.1"
-once_cell = "1.19.0"
\ No newline at end of file
+once_cell = "1.19.0"
+ring = "0.17"
+base64 = "0.21"
diff --git a/src/device.rs b/src/device.rs
index a41a72f..6e24bdf 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -3,11 +3,9 @@ use crate::{
     modules::{start_module, HostEvent, MODULE_REGISTRY},
     unwrap_or_error,
 };
-use clru::CLruCache;
 use deck_driver as streamdeck;
 use hidapi::HidApi;
-use image::DynamicImage;
-use std::{collections::HashMap, fmt::Display, num::NonZeroUsize, sync::Arc};
+use std::{collections::HashMap, fmt::Display, sync::Arc};
 use streamdeck::{
     asynchronous::{AsyncStreamDeck, ButtonStateUpdate},
     info::Kind,
@@ -16,10 +14,7 @@ use streamdeck::{
 use tokio::{
     process::Command,
     runtime::Runtime,
-    sync::{
-        mpsc::{self, error::TrySendError},
-        Mutex,
-    },
+    sync::mpsc::{self, error::TrySendError},
 };
 use tracing::{debug, error, info_span, trace, warn};
 
@@ -46,11 +41,6 @@ impl Display for Device {
     }
 }
 
-/// image cache
-///
-/// (button, key), image data
-pub type ImageCache = CLruCache<(u8, u32), Arc<DynamicImage>>;
-
 /// Handles everything related to a single device
 pub struct Device {
     modules: HashMap<u8, ModuleController>,
@@ -61,7 +51,6 @@ pub struct Device {
     selected_space: Option<String>,
     is_dead_handle: Arc<std::sync::Mutex<bool>>,
     serial: String,
-    image_cache: Arc<Mutex<ImageCache>>,
 }
 
 impl Device {
@@ -104,9 +93,6 @@ impl Device {
             selected_space: None,
             is_dead_handle,
             serial,
-            image_cache: Arc::new(Mutex::new(CLruCache::new(
-                NonZeroUsize::new(button_count.into()).unwrap(),
-            ))),
         })
     }
 
@@ -147,11 +133,10 @@ impl Device {
                 let ser = self.serial.clone();
                 let dev = self.device.clone();
                 let b = btn.clone();
-                let image_cache = self.image_cache.clone();
 
-                runtime.spawn(async move {
-                    start_module(ser, b, *module, dev, module_receiver, image_cache).await
-                });
+                runtime.spawn(
+                    async move { start_module(ser, b, *module, dev, module_receiver).await },
+                );
             }
             // if the receiver already dropped the listener then just directly insert none.
             // Optimizes performance because the key_listener just does not try to send the event.
diff --git a/src/image_rendering.rs b/src/image_rendering.rs
index e643244..8b00df9 100644
--- a/src/image_rendering.rs
+++ b/src/image_rendering.rs
@@ -5,11 +5,13 @@ use image::imageops;
 use image::{io::Reader, DynamicImage, ImageBuffer, Rgb, RgbImage};
 use imageproc::drawing::draw_text_mut;
 use rusttype::Scale;
-use std::io;
+use std::{io, path::Path};
+use tracing::trace;
 
-/// Loads an image from a path
+/// Retrieve an image from a path
 #[allow(dead_code)]
-pub fn load_image(path: String) -> io::Result<DynamicImage> {
+pub fn retrieve_image(path: &Path) -> io::Result<DynamicImage> {
+    trace!("Retrieving image from filesystem");
     Ok(Reader::open(path)?
         .decode()
         .expect("Unable to decode image"))
diff --git a/src/main.rs b/src/main.rs
index a34adc7..0d7cc71 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,6 +6,7 @@ use hidapi::HidApi;
 use rusttype::Font;
 use std::{
     collections::HashMap,
+    fs,
     process::exit,
     sync::{Arc, Mutex, OnceLock},
     time::Duration,
@@ -24,6 +25,7 @@ mod modules;
 mod type_definition;
 
 pub static GLOBAL_FONT: OnceLock<Font> = OnceLock::new();
+pub static CACHE_DIR: &'static str = "microdeck";
 
 #[macro_export]
 macro_rules! skip_if_none {
@@ -64,6 +66,14 @@ fn main() {
         }
     };
 
+    // create cache dir
+    if let Some(mut path) = dirs::cache_dir() {
+        path.push(CACHE_DIR);
+        if path.exists() == false {
+            fs::create_dir(path).expect("Could not create cache directory");
+        }
+    }
+
     // load font
     let font = match config.global.font_family {
         Some(ref f) => font_loader::system_fonts::get(
diff --git a/src/modules.rs b/src/modules.rs
index 91b83f7..e2c0aa8 100644
--- a/src/modules.rs
+++ b/src/modules.rs
@@ -10,21 +10,24 @@ use self::space::Space;
 
 // other things
 use crate::config::{Button, ButtonConfigError};
-use crate::device::ImageCache;
-use crate::image_rendering::{load_image, ImageBuilder};
+use crate::image_rendering::{retrieve_image, ImageBuilder};
 use ::image::imageops::FilterType;
+use ::image::io::Reader as ImageReader;
 use ::image::DynamicImage;
 use async_trait::async_trait;
+use base64::engine::{general_purpose, Engine};
 pub use deck_driver as streamdeck;
+use dirs::cache_dir;
 use futures_util::Future;
 use once_cell::sync::Lazy;
+use ring::digest;
 use std::collections::HashMap;
-use std::{error::Error, pin::Pin, sync::Arc};
+use std::{error::Error, path::PathBuf, pin::Pin, sync::Arc};
 use streamdeck::info::ImageFormat;
 use streamdeck::info::Kind;
 use streamdeck::AsyncStreamDeck;
 use streamdeck::StreamDeckError;
-use tokio::sync::{mpsc, Mutex};
+use tokio::sync::mpsc;
 use tracing::{debug, error, trace};
 
 pub static MODULE_REGISTRY: Lazy<ModuleRegistry> = Lazy::new(|| ModuleRegistry::default());
@@ -41,7 +44,7 @@ pub enum HostEvent {
 pub type ModuleObject = Box<dyn Module + Send + Sync>;
 pub type ModuleFuture =
     Pin<Box<dyn Future<Output = Result<ModuleObject, ButtonConfigError>> + Send>>;
-pub type ModuleInitFunction = fn(Arc<Button>, ModuleCache) -> ModuleFuture;
+pub type ModuleInitFunction = fn(Arc<Button>) -> ModuleFuture;
 
 pub type ModuleMap = HashMap<&'static str, ModuleInitFunction>;
 
@@ -83,21 +86,15 @@ pub async fn start_module(
     module_init_function: ModuleInitFunction,
     device: Arc<AsyncStreamDeck>,
     br: ChannelReceiver,
-    image_cache: Arc<Mutex<ImageCache>>,
 ) {
     debug!("STARTED");
-    let mc = ModuleCache::new(
-        image_cache,
-        button.index,
-        device.kind().key_image_format().size,
-    );
     let da = DeviceAccess::new(device.clone(), button.index).await;
 
     // init
     //
     // This function should be called after the config was checked,
     // otherwise it will panic and the module wont be started.
-    let mut module = match module_init_function(button.clone(), mc).await {
+    let mut module = match module_init_function(button.clone()).await {
         Ok(m) => m,
         Err(e) => panic!("{}", e),
     };
@@ -121,63 +118,76 @@ pub async fn start_module(
     }
 }
 
-/// A wrapper around [ImageCache] to provide easy access to values in the device cache
-pub struct ModuleCache {
-    image_cache: Arc<Mutex<ImageCache>>,
-    button_index: u8,
-    /// Resolution of the deck (required for optimization of storage space)
-    resolution: (usize, usize),
+/// Loads an image from the system or retrieves it from the cache. If
+/// the provided image is not already in the cache it will be inserted.
+#[allow(dead_code)]
+pub async fn load_image(path: PathBuf, resolution: (usize, usize)) -> Option<DynamicImage> {
+    // hash the image
+    let mut image = tokio::task::spawn_blocking(move || retrieve_image(&path))
+        .await
+        .unwrap()
+        .ok()?;
+
+    let image_hash = hash_image(image.as_bytes());
+
+    if let Some(image) = get_image_from_cache(&image_hash, resolution) {
+        trace!("Cached image is available");
+        return Some(image);
+    }
+
+    let image = tokio::task::spawn_blocking(move || {
+        trace!("Resizing image");
+        image = image.resize_exact(
+            resolution.0 as u32,
+            resolution.1 as u32,
+            FilterType::Lanczos3,
+        );
+        trace!("Resizing finished");
+        let mut path = match cache_dir() {
+            Some(dir) => dir,
+            None => return None, // System does not provide cache
+        };
+        path.push("microdeck");
+        path.push(image_cache_file_name(&image_hash, resolution));
+        // let mut file = File::create(path);
+        // file.write_all(image.as_bytes());
+        // currently just ignore if saving is not possible
+        image.save(path).ok()?;
+        Some(image)
+    })
+    .await
+    .unwrap()?;
+    Some(image.into())
 }
 
-impl ModuleCache {
-    pub fn new(
-        image_cache: Arc<Mutex<ImageCache>>,
-        button_index: u8,
-        resolution: (usize, usize),
-    ) -> Self {
-        ModuleCache {
-            image_cache,
-            button_index,
-            resolution,
-        }
-    }
+/// File name for a cached image
+///
+/// `<hash>-<height>x<width>`
+pub fn image_cache_file_name(image_hash: &str, resolution: (usize, usize)) -> String {
+    format!("{}-{}x{}.png", image_hash, resolution.0, resolution.1)
+}
 
-    /// Load an image from the [ImageCache] or create a new one and insert it into the [ImageCache].
-    /// Returns None if no image was found.
-    ///
-    /// index: Provide an index where your data is cached. With this number the value can be
-    /// accessed again. Use [DeviceAccess::get_image_cached()] for just getting the data.
-    #[allow(dead_code)]
-    pub async fn load_image(&mut self, path: String, index: u32) -> Option<Arc<DynamicImage>> {
-        if let Some(image) = self.get_image(index).await {
-            Some(image)
-        } else {
-            trace!("Decoding image");
-            let mut image = tokio::task::spawn_blocking(move || load_image(path))
-                .await
-                .unwrap()
-                .ok()?;
-            image = image.resize_exact(
-                self.resolution.0 as u32,
-                self.resolution.1 as u32,
-                FilterType::Lanczos3,
-            );
-            trace!("Decoding finished");
-            let image = Arc::new(image);
-            let mut data = self.image_cache.lock().await;
-            data.put((self.button_index, index), image.clone());
-            trace!("Wrote data into cache (new size: {})", data.len());
-            drop(data);
-            Some(image.into())
-        }
-    }
+pub fn hash_image(data: &[u8]) -> String {
+    let mut context = digest::Context::new(&digest::SHA256);
+    context.update(data);
+    let hash = context.finish();
+    general_purpose::STANDARD.encode(hash)
+}
 
-    /// Just try to retrieve a value from the key (index) in the [ImageCache].
-    #[allow(dead_code)]
-    pub async fn get_image(&self, index: u32) -> Option<Arc<DynamicImage>> {
-        let mut data = self.image_cache.lock().await;
-        data.get(&(self.button_index, index)).cloned()
-    }
+/// Try to retrieve an image from the cache. Will return None if
+/// the image was not cached yet (or is not accessible)
+/// or if the system does not provide a [dirs::cache_dir].
+#[allow(dead_code)]
+pub fn get_image_from_cache(image_hash: &str, resolution: (usize, usize)) -> Option<DynamicImage> {
+    let mut path = match cache_dir() {
+        Some(dir) => dir,
+        None => return None, // System does not provide cache
+    };
+
+    path.push("microdeck");
+    path.push(image_cache_file_name(image_hash, resolution));
+
+    Some(ImageReader::open(path).ok()?.decode().ok()?)
 }
 
 /// Wrapper to provide easier access to the Deck
@@ -240,14 +250,11 @@ pub trait Module: Sync + Send {
     ///
     /// This function should **not** panic as the panic will not be catched and therefore would be
     /// not noticed.
-    async fn new(
-        config: Arc<Button>,
-        mut cache: ModuleCache,
-    ) -> Result<ModuleObject, ButtonConfigError>
+    async fn new(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError>
     where
         Self: Sized;
     /// Function for actually running the module and interacting with the device. Errors that
-    /// happen here should be mostly prevented.
+    /// happen here should be mostly prevented as they are not properly handled.
     async fn run(
         &mut self,
         device: DeviceAccess,
diff --git a/src/modules/blank.rs b/src/modules/blank.rs
index 7d9afc7..451c122 100644
--- a/src/modules/blank.rs
+++ b/src/modules/blank.rs
@@ -3,7 +3,6 @@ use super::ButtonConfigError;
 use super::ChannelReceiver;
 use super::DeviceAccess;
 use super::Module;
-use super::ModuleCache;
 use super::ModuleObject;
 use super::ReturnError;
 use async_trait::async_trait;
@@ -13,10 +12,7 @@ pub struct Blank;
 
 #[async_trait]
 impl Module for Blank {
-    async fn new(
-        _config: Arc<Button>,
-        _cache: ModuleCache,
-    ) -> Result<ModuleObject, ButtonConfigError> {
+    async fn new(_config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
         Ok(Box::new(Blank {}))
     }
 
diff --git a/src/modules/counter.rs b/src/modules/counter.rs
index 730b4c9..edb439a 100644
--- a/src/modules/counter.rs
+++ b/src/modules/counter.rs
@@ -3,8 +3,7 @@ use std::sync::Arc;
 use crate::config::Button;
 
 use super::{
-    ButtonConfigError, ChannelReceiver, DeviceAccess, HostEvent, Module, ModuleCache, ModuleObject,
-    ReturnError,
+    ButtonConfigError, ChannelReceiver, DeviceAccess, HostEvent, Module, ModuleObject, ReturnError,
 };
 
 use crate::image_rendering::wrap_text;
@@ -25,10 +24,7 @@ pub struct Counter {
 
 #[async_trait]
 impl Module for Counter {
-    async fn new(
-        config: Arc<Button>,
-        _cache: ModuleCache,
-    ) -> Result<ModuleObject, ButtonConfigError> {
+    async fn new(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
         let title = config.parse_module("title", " ".to_string()).res()?;
         let title_size = config.parse_module("title_size", 15.0).res()?;
         let number_size = config.parse_module("number_size", 25.0).res()?;
diff --git a/src/modules/image.rs b/src/modules/image.rs
index a0258ff..1847681 100644
--- a/src/modules/image.rs
+++ b/src/modules/image.rs
@@ -1,38 +1,34 @@
+use super::load_image;
 use super::Button;
 use super::ButtonConfigError;
 use super::ChannelReceiver;
 use super::DeviceAccess;
 use super::Module;
-use super::ModuleCache;
 use super::ModuleObject;
 use super::ReturnError;
 use crate::image_rendering::ImageBuilder;
 use async_trait::async_trait;
-use image::DynamicImage;
-use std::sync::Arc;
+use std::{path::PathBuf, sync::Arc};
 
 pub struct Image {
-    image: Arc<DynamicImage>,
     scale: f32,
+    path: PathBuf,
 }
 
 #[async_trait]
 impl Module for Image {
-    async fn new(
-        config: Arc<Button>,
-        mut cache: ModuleCache,
-    ) -> Result<ModuleObject, ButtonConfigError> {
+    async fn new(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
         let path = config.parse_module("path", String::new()).required()?;
         let scale = config.parse_module("scale", 100.0).res()?;
 
-        let image = cache
-            .load_image(path, 1)
-            .await
-            .ok_or(ButtonConfigError::General(
+        let path = PathBuf::from(path);
+        if path.exists() == false {
+            return Err(ButtonConfigError::General(
                 "Image was not found".to_string(),
-            ))?;
+            ));
+        };
 
-        Ok(Box::new(Image { image, scale }))
+        Ok(Box::new(Image { scale, path }))
     }
 
     async fn run(
@@ -41,9 +37,8 @@ impl Module for Image {
         _button_receiver: ChannelReceiver,
     ) -> Result<(), ReturnError> {
         let (h, w) = streamdeck.resolution();
-        let img = (*self.image).clone();
         let img = ImageBuilder::new(h, w)
-            .set_image(img)
+            .set_image(load_image(self.path.clone(), (h, w)).await.unwrap())
             .set_image_scale(self.scale)
             .build();
         streamdeck.write_img(img).await.unwrap();
diff --git a/src/modules/space.rs b/src/modules/space.rs
index 65199d1..03f5d25 100644
--- a/src/modules/space.rs
+++ b/src/modules/space.rs
@@ -3,7 +3,6 @@ use super::ButtonConfigError;
 use super::ChannelReceiver;
 use super::DeviceAccess;
 use super::Module;
-use super::ModuleCache;
 use super::ModuleObject;
 use super::ReturnError;
 use crate::image_rendering::{create_error_image, ImageBuilder};
@@ -17,10 +16,7 @@ pub struct Space {
 
 #[async_trait]
 impl Module for Space {
-    async fn new(
-        config: Arc<Button>,
-        _cache: ModuleCache,
-    ) -> Result<ModuleObject, ButtonConfigError> {
+    async fn new(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
         let name = config.parse_module("name", "Unknown".to_string()).res()?;
         Ok(Box::new(Space { name }))
     }