Added color cache

.env.example

@@ -1,3 +1,10 @@
 DISCORD_TOKEN=
 LASTFM_API_KEY=
 DATABASE_URL=sqlite:watcat.db
+
+# optional, for spotify support
+RSPOTIFY_CLIENT_ID=
+RSPOTIFY_CLIENT_SECRET=
+
+# optional, to set cache size, set to 0 to disable
+MAX_CACHE_SIZE=

Cargo.lock

@@ -4270,7 +4270,7 @@ dependencies = [
 
 [[package]]
 name = "watcat"
-version = "0.5.0"
+version = "0.6.0"
 dependencies = [
  "arabic_reshaper",
  "dotenvy",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "watcat"
-version = "0.5.0"
+version = "0.6.0"
 edition = "2024"
 
 [dependencies]

README.md

@@ -33,7 +33,7 @@ TODO
 + ~~refactor to combine repeated parts of art, url, and fmi.~~
 + use macro for commands
 + use Option instead of empty string sentinel value
-+ add cache for album art colors
++ ~~add cache for album art colors~~
 + ~~add spotify link fetching~~
 + make various parameters user-configurable
 + use tiny-skia instead of magick-rust?

src/main.rs

@@ -16,6 +16,7 @@ use magick_rust::{
 };
 
 use std::env;
+use std::collections::HashMap;
 use std::sync::{Arc, Once};
 use tokio::sync::Mutex;
 
@@ -69,8 +70,9 @@ macro_rules! log {
 macro_rules! handle_magick_result {
     ($a: expr, $b: literal) => {
         match $a {
-            Ok(_) => {
+            Ok(a) => {
                 log!("{} run successfully", stringify!($a));
+                a
             }
             Err(e) => return Reply::Text(format!("Error: {} {}", $b, e)),
         }
@@ -179,7 +181,6 @@ async fn spot(ctx: &Context, arg: &str, id: UserId) -> Reply {
 		return Reply::Text("Unable to use Spotify command: Contact bot Administrator.".to_owned())
 	}
 	let spotify_arc_mutex = spotify_option.as_mut().unwrap();
-	//let spotify_arc_mutex_ref = Arc::clone(&spotify_arc_mutex);
 	let spotify_client = spotify_arc_mutex.lock().await;
 	spotify_client.request_token().await.unwrap();
 	let search = spotify_client.search(format!("{} {}", track.name, track.artist.name).as_str(), SearchType::Track, None, None, None, None).await;
@@ -209,11 +210,8 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
         None => return Reply::Text("Error: getting image uri failed".to_owned()),
     };
 
-    let image = match get_image(ctx, image_uri.as_str()).await {
-        Ok(i) => i,
-        Err(e) => return Reply::Text(format!("{e}")),
-    };
     let mut base_color = PixelWand::new();
+    let mut accent_color = PixelWand::new();
     let mut white = PixelWand::new();
     let main_wand = MagickWand::new();
     let art_wand = MagickWand::new();
@@ -221,67 +219,94 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
     let mask_wand = MagickWand::new();
     let text_image_wand = MagickWand::new();
 
-    handle_magick_result!(
+	let image = match get_image(ctx, image_uri.as_str()).await {
-        base_color.set_color("#7f7f7f"),
+	    Ok(i) => i,
-        "Failed to set init base color"
+	    Err(e) => return Reply::Text(format!("{e}")),
-    );
+	};
-    handle_magick_result!(
+	handle_magick_result!(
-        art_wand.read_image_blob(image),
+	    art_wand.read_image_blob(image),
-        "Failed to read image from uri"
+	    "Failed to read image from uri"
-    );
+	);
-    handle_magick_result!(
+	handle_magick_result!(
-        art_wand.transform_image_colorspace(ColorspaceType::Lab),
+	    art_wand.transform_image_colorspace(ColorspaceType::Lab),
-        "Failed to set art colorspace"
+	    "Failed to set art colorspace"
-    );
+	);
 	
-    let (art_width, art_height) = (art_wand.get_image_width(), art_wand.get_image_height());
+	handle_magick_result!(
-    handle_magick_result!(
+	    base_color.set_color("#7f7f7f"),
-        art_wand_cluster.new_image(art_width, art_height, &base_color),
+	    "Failed to set init base color"
-        "Failed to init cluster image"
+	);
-    );
-    handle_magick_result!(
-        art_wand_cluster.compose_images(&art_wand, CompositeOperator::SrcOver, true, 0, 0),
-        "Failed to copy onto cluster image"
-    );
-    // least stupid way I could figure out to do this, since it seems that MagickWand::new_from_wand() doesn't work
-    handle_magick_result!(
-        art_wand_cluster.set_image_colorspace(ColorspaceType::Lab),
-        "Failed to set cluster colorspace"
-    );
-    handle_magick_result!(
-        art_wand_cluster.kmeans(10, 200, 0.001),
-        "Failed to run kmeans"
-    );
-    let mut colors = handle_magick_option!(
-        art_wand_cluster.get_image_histogram(),
-        "Failed to get color histogram"
-    );
-    colors.sort_by_cached_key(|color| -(color.get_color_count() as isize));
-    let base_lab = validate_color(&colors[0]).expect("Failed to validate base color");
-    let other_colors = colors
-        .iter()
-        .filter_map(|col| {
-            let testing_color = validate_color(col)?;
-            match testing_color.improved_difference(base_lab) {
-                10.0.. => Some(testing_color),
-                _ => None,
-            }
-        })
-        .collect::<Vec<_>>();
-    let mut accent_color = PixelWand::new();
-    handle_magick_result!(
-        accent_color.set_color("cielab(0.0,0.0,0.0)"),
-        "Failed to init accent color"
-    );
-    if let Some(color) = other_colors.first() {
-        let col = format!("cielab{:?}", color.into_components());
-        handle_magick_result!(
-            accent_color.set_color(col.as_str()),
-            "Failed to set accent color"
-        );
-    }
 	
-    let use_dark_color = match validate_color(&colors[0]) {
+	let color_cache_entry = {
+    	let data = ctx.data.write().await;
+		let color_cache = data.get::<ColorsHaver>().expect("Failed to have colors option");
+		if color_cache.max > 0 {
+			let color_cache_map = color_cache.map.lock().await;
+			color_cache_map.get(image_uri.as_str()).map(|s| s.to_owned())
+		} else {
+			None
+		}
+	};
+		
+
+    let (base_lab, accent_lab) = match color_cache_entry {
+    	Some(value) => value.to_owned(),
+    	None => {
+		    let (art_width, art_height) = (art_wand.get_image_width(), art_wand.get_image_height());
+	    	handle_magick_result!(
+	    	    art_wand_cluster.new_image(art_width, art_height, &base_color),
+	    	    "Failed to init cluster image"
+	    	);
+	    	handle_magick_result!(
+	    	    art_wand_cluster.compose_images(&art_wand, CompositeOperator::SrcOver, true, 0, 0),
+	    	    "Failed to copy onto cluster image"
+	    	);
+	    	// least stupid way I could figure out to do this, since it seems that MagickWand::new_from_wand() doesn't work
+	    	handle_magick_result!(
+	    	    art_wand_cluster.set_image_colorspace(ColorspaceType::Lab),
+	    	    "Failed to set cluster colorspace"
+	    	);
+	    	handle_magick_result!(
+	    	    art_wand_cluster.kmeans(10, 200, 0.001),
+	    	    "Failed to run kmeans"
+	    	);
+	    	let mut colors = handle_magick_option!(
+	    	    art_wand_cluster.get_image_histogram(),
+	    	    "Failed to get color histogram"
+	    	);
+	    	colors.sort_by_cached_key(|color| -(color.get_color_count() as isize));
+	    	let base_lab = validate_color(&colors[0]).expect("Failed to validate base color");
+	    	let other_colors = colors
+	    	    .iter()
+	    	    .filter_map(|col| {
+	    	        let testing_color = validate_color(col)?;
+	    	        match testing_color.improved_difference(base_lab) {
+	    	            10.0.. => Some(testing_color),
+	    	            _ => None,
+	    	        }
+	    	    })
+	    	    .collect::<Vec<_>>();
+	    	let mut accent_lab = "cielab(0.0,0.0,0.0)".to_owned();
+			if let Some(color) = other_colors.first() {
+			    accent_lab = format!("cielab{:?}", color.into_components());
+			}
+			let base_lab = handle_magick_result!(colors[0].get_color_as_string(), "Failed to serialize base color");
+	    	let data = ctx.data.write().await;
+			let color_cache = data.get::<ColorsHaver>().expect("Failed to have colors option");
+			let mut color_cache_map = color_cache.map.lock().await;
+			if color_cache_map.keys().len() > color_cache.max {
+				let key: String = color_cache_map.iter().next().unwrap().0.to_owned();
+				color_cache_map.remove(&key);
+			}
+			color_cache_map.insert(image_uri, (base_lab.clone(), accent_lab.clone()));
+	    	(base_lab, accent_lab)
+    	}
+    };
+
+    handle_magick_result!(base_color.set_color(base_lab.as_str()), "Failed to deserialize base color");
+    handle_magick_result!(accent_color.set_color(accent_lab.as_str()), "Failed to deserialize accent color");
+
+    let use_dark_color = match validate_color(&base_color) {
         Some(color) => {
             let black_delta_e = Lab::<D65>::new(0.0, 0.0, 0.0).improved_difference(color);
             let white_delta_e = Lab::<D65>::new(100.0, 0.0, 0.0).improved_difference(color);
@@ -362,7 +387,7 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
     );
 
     handle_magick_result!(
-        main_wand.new_image(image_width as usize, image_height as usize, &colors[0]),
+        main_wand.new_image(image_width as usize, image_height as usize, &base_color),
         "Failed to create wand"
     );
     handle_magick_result!(
@@ -373,10 +398,10 @@ async fn fmi(ctx: &Context, arg: &str, id: UserId, avatar: Option<String>) -> Re
         art_wand.adaptive_resize_image(124, 124),
         "Failed to resize art"
     );
-    handle_magick_result!(
+    /*handle_magick_result!(
         art_wand_cluster.adaptive_resize_image(124, 124),
         "Failed to resize art_cluster"
-    );
+    );*/
     handle_magick_result!(
         main_wand.compose_images(&art_wand, CompositeOperator::SrcOver, false, 12, 12),
         "Failed to combine art image"
@@ -519,6 +544,15 @@ impl TypeMapKey for SpotifyHaver {
     type Value = Option<Arc<Mutex<ClientCredsSpotify>>>;
 }
 
+struct ColorCache {
+	map: Arc<Mutex<HashMap<String, (String, String)>>>,
+	max: usize,
+}
+struct ColorsHaver;
+impl TypeMapKey for ColorsHaver {
+	type Value = ColorCache;
+}
+
 struct DBResponse {
     lastfm_username: String,
 }
@@ -552,7 +586,7 @@ async fn get_lastfm_username(ctx: &Context, id: UserId) -> Option<String> {
 
 async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
     log!("set db user {} {}", id, user);
-    let data = ctx.data.write().await;
+    let data = ctx.data.read().await;
     let pool = data
         .get::<PoolHaver>()
         .expect("Failed to get pool container");
@@ -571,8 +605,8 @@ async fn set_lastfm_username(ctx: &Context, id: UserId, user: String) {
 }
 
 async fn get_image(ctx: &Context, url: &str) -> Result<Vec<u8>, reqwest_middleware::Error> {
-    log!("get {}", url);
+	log!("get {}", url);
-    let data = ctx.data.write().await;
+    let data = ctx.data.read().await;
     let http = data.get::<HttpHaver>().expect("Failed to get http client");
     match http.get(url).send().await {
         Ok(resp) => {
@@ -652,6 +686,11 @@ async fn main() {
         let _ = &spotify_client.request_token().await.unwrap();
     }
 
+    let colors_cache: ColorCache = ColorCache {
+    	map: Arc::new(Mutex::new(HashMap::new())),
+    	max: env::var("MAX_CACHE_SIZE").unwrap_or("".to_owned()).parse::<usize>().unwrap_or(1000000),
+	};
+
     let mut discord_client = Client::builder(&token, intents)
         .event_handler(Handler)
         .await
@@ -663,6 +702,7 @@ async fn main() {
         data.insert::<HttpHaver>(http);
         data.insert::<FontsHaver>((regular_fonts, bold_fonts));
         data.insert::<SpotifyHaver>(spotify.map(Mutex::new).map(Arc::new));
+        data.insert::<ColorsHaver>(colors_cache);
     }
 
     if let Err(why) = discord_client.start().await {