Cache thumbnails to disk

src/api.rs

@@ -1,29 +1,108 @@
 use std::{
     collections::HashMap,
+    io::Cursor,
     iter::repeat,
     sync::{Arc, Mutex},
 };
 
 use anyhow::{Context as _, anyhow};
+use image::{
+    DynamicImage, EncodableLayout,
+    codecs::webp::{WebPDecoder, WebPEncoder},
+};
 use immich_sdk::{AssetId, AssetVisibility};
 use slint::{Rgba8Pixel, SharedPixelBuffer};
 
-use crate::thumbhash::thumbhashes_to_pixels;
+use crate::{
+    cachemap::{AsyncFnFetcher, CacheMap, Fetcher},
+    thumbhash::thumbhashes_to_pixels,
+};
 
 pub type TimeBucketKey = String;
 
 pub struct Api {
     client: immich_sdk::Client,
     buckets: Mutex<HashMap<TimeBucketKey, Arc<TimeBucket>>>,
-    thumbnails: Mutex<HashMap<AssetId, Arc<AssetThumbnail>>>,
+    thumbnails: Mutex<CacheMap<AssetId, Arc<AssetThumbnail>>>,
+}
+
+fn serialize_thumbnail(_id: &AssetId, thumbnail: &Arc<AssetThumbnail>) -> Vec<u8> {
+    let image: &SharedPixelBuffer<Rgba8Pixel> = &thumbnail.thumbnail;
+
+    let mut webp = vec![];
+    WebPEncoder::new_lossless(&mut webp)
+        .encode(
+            image.as_bytes(),
+            image.width(),
+            image.height(),
+            image::ExtendedColorType::Rgba8,
+        )
+        .expect("width and height matches image.as_bytes().len()");
+
+    webp
+}
+fn deserialize_thumbnail(&id: &AssetId, bytes: &[u8]) -> anyhow::Result<Arc<AssetThumbnail>> {
+    let image = WebPDecoder::new(Cursor::new(bytes))
+        .and_then(DynamicImage::from_decoder)
+        .context("Failed to decode image")?
+        .to_rgba8();
+
+    let image = SharedPixelBuffer::<Rgba8Pixel>::clone_from_slice(
+        image.as_bytes(),
+        image.width(),
+        image.height(),
+    );
+
+    Ok(Arc::new(AssetThumbnail {
+        id,
+        thumbnail: image,
+    }))
 }
 
 impl Api {
     pub fn new(client: immich_sdk::Client) -> Arc<Self> {
+        let client_ = client.clone();
+        let thumbnail_fetcher = AsyncFnFetcher::new(move |asset_id: &AssetId| {
+            let client = client_.clone();
+            let asset_id = asset_id.clone();
+            tokio::spawn(async move {
+                let response = client
+                    .assets()
+                    .thumbnail(asset_id)
+                    .size(immich_sdk::AssetMediaSize::Thumbnail)
+                    .execute()
+                    .await
+                    .context(anyhow!("Failed to get asset thumbnail for {asset_id}"))?;
+
+                let thumbnail = response
+                    .decode()
+                    .context(anyhow!("Failed to decode asset thumbnail for {asset_id}"))?
+                    .into_rgba8();
+
+                let pixel_buffer = SharedPixelBuffer::<Rgba8Pixel>::clone_from_slice(
+                    &thumbnail,
+                    thumbnail.width(),
+                    thumbnail.height(),
+                );
+
+                Ok(Arc::new(AssetThumbnail {
+                    id: asset_id,
+                    thumbnail: pixel_buffer,
+                }))
+            })
+        });
+        let cache_map = CacheMap::new(
+            Arc::new(thumbnail_fetcher) as Arc<dyn Fetcher<Arc<AssetThumbnail>, Key = AssetId>>,
+            "thumbnails",
+            serialize_thumbnail,
+            deserialize_thumbnail,
+        )
+        .unwrap();
+
         Arc::new(Self {
             client,
             buckets: Default::default(),
-            thumbnails: Default::default(),
+            thumbnails: Mutex::new(cache_map),
         })
     }
 }
@@ -136,40 +215,9 @@ impl Api {
         &self,
         asset_id: AssetId,
     ) -> anyhow::Result<Arc<AssetThumbnail>> {
-        if let Some(thumbnail) = self.thumbnails.lock().unwrap().get(&asset_id).cloned() {
-            return Ok(thumbnail);
-        }
-
-        let response = self
-            .client
-            .assets()
-            .thumbnail(asset_id)
-            .size(immich_sdk::AssetMediaSize::Thumbnail)
-            .execute()
+        let fetch = self.thumbnails.lock().unwrap().get(asset_id);
+        fetch
             .await
-            .context(anyhow!("Failed to get asset thumbnail for {asset_id}"))?;
-
-        let thumbnail = response
-            .decode()
-            .context(anyhow!("Failed to decode asset thumbnail for {asset_id}"))?
-            .into_rgba8();
-
-        let pixel_buffer = SharedPixelBuffer::<Rgba8Pixel>::clone_from_slice(
-            &thumbnail,
-            thumbnail.width(),
-            thumbnail.height(),
-        );
-
-        let thumbnail = Arc::new(AssetThumbnail {
-            id: asset_id,
-            thumbnail: pixel_buffer,
-        });
-
-        self.thumbnails
-            .lock()
-            .unwrap()
-            .insert(asset_id, Arc::clone(&thumbnail));
-
-        Ok(thumbnail)
+            .context(anyhow!("Failed to fetch thumbnail for {asset_id:?}"))
     }
 }

src/cachemap.rs

@@ -0,0 +1,190 @@
+use std::{
+    collections::{HashMap, hash_map::Entry},
+    fmt::Debug,
+    future::ready,
+    hash::Hash,
+    marker::PhantomData,
+    path::PathBuf,
+    pin::Pin,
+    str::FromStr,
+    sync::Arc,
+};
+
+use anyhow::{Context, anyhow};
+use either::Either;
+use futures::{FutureExt, future::WeakShared};
+use tokio::{fs, task::JoinHandle};
+
+pub trait Fetcher<V>: Send + Sync + 'static {
+    type Key: ToString + FromStr;
+
+    fn fetch(
+        &self,
+        key: &Self::Key,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<V>> + Send + Sync>>;
+}
+
+pub struct AsyncFnFetcher<K, V, F> {
+    f: F,
+    _phantom: PhantomData<(K, V)>,
+}
+
+impl<K, V, F> AsyncFnFetcher<K, V, F>
+where
+    K: ToString + FromStr + Send + 'static,
+    F: Fn(&K) -> JoinHandle<anyhow::Result<V>> + Send + 'static,
+    V: Send + 'static,
+{
+    pub fn new(f: F) -> Self {
+        Self {
+            f,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<K, V, F> Fetcher<V> for AsyncFnFetcher<K, V, F>
+where
+    K: ToString + FromStr + Send + Sync + 'static,
+    F: Fn(&K) -> JoinHandle<anyhow::Result<V>> + Send + Sync + 'static,
+    V: Send + Sync + 'static,
+{
+    type Key = K;
+
+    fn fetch(
+        &self,
+        key: &Self::Key,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<V>> + Send + Sync>> {
+        let handle = (self.f)(key);
+        Box::pin(async move { handle.await.context("Fetch task panicked")? })
+    }
+}
+
+type FetchJob<V> = Pin<Box<dyn Future<Output = Option<V>> + Send + Sync>>;
+
+pub struct CacheMap<K, V> {
+    fetcher: Arc<dyn Fetcher<V, Key = K>>,
+    cache_dir: PathBuf,
+    serialize: fn(&K, &V) -> Vec<u8>,
+    deserialize: fn(&K, &[u8]) -> anyhow::Result<V>,
+
+    /// Cache of ongoing [`FetchJob`]s.
+    ///
+    /// If [`CacheMap::get`] is called while a fetch job is ongoing, [`WeakShared::upgrade`]
+    /// will succeed, and the [`FetchJob`] can be cloned.
+    fetch_jobs: HashMap<K, WeakShared<FetchJob<V>>>,
+}
+
+impl<K, V> CacheMap<K, V>
+where
+    K: Debug + FromStr + ToString + Eq + Hash + Clone + Send + Sync + 'static,
+    V: Clone + Send + Sync + 'static,
+{
+    pub fn new(
+        fetcher: Arc<dyn Fetcher<V, Key = K> + 'static>,
+        cache_name: &str,
+        serialize: fn(&K, &V) -> Vec<u8>,
+        deserialize: fn(&K, &[u8]) -> anyhow::Result<V>,
+    ) -> anyhow::Result<Self> {
+        let crate_name = env!("CARGO_CRATE_NAME");
+        let data_dir = xdg::BaseDirectories::new()
+            .create_cache_directory(format!("{crate_name}/thumbnails"))
+            .context(anyhow!(
+                "Failed to create XDG data folder for {cache_name:?}"
+            ))?;
+
+        Ok(Self {
+            fetcher,
+            cache_dir: data_dir,
+            serialize,
+            deserialize,
+            fetch_jobs: HashMap::new(),
+        })
+    }
+
+    fn fetch_from_cache(
+        &self,
+        key: &K,
+    ) -> impl Future<Output = anyhow::Result<V>> + Send + Sync + use<K, V> {
+        let key = key.clone();
+        let key_str = key.to_string();
+        let path = self.cache_dir.join(key_str);
+        let deserialize = self.deserialize;
+        Box::pin(async move {
+            let bytes = fs::read(&path)
+                .await
+                .context(anyhow!("Failed to read {path:?}"))?;
+            deserialize(&key, &bytes).context(anyhow!("Failed to deserialize value at {path:?}"))
+        })
+    }
+
+    pub fn get(&mut self, key: K) -> impl Future<Output = Option<V>> + use<K, V> {
+        // FIXME: creating this future here because lifetimes.
+        let fetch_from_cache = self.fetch_from_cache(&key);
+
+        let entry = match self.fetch_jobs.entry(key.clone()) {
+            Entry::Vacant(entry) => entry,
+            Entry::Occupied(entry) => {
+                if let Some(fetching) = entry.get().upgrade() {
+                    return match fetching.clone().now_or_never() {
+                        // Value fetched
+                        Some(Some(value)) => {
+                            entry.remove();
+                            Either::Left(ready(Some(value)))
+                        }
+
+                        // Failed to fetch
+                        Some(None) => {
+                            entry.remove();
+                            Either::Left(ready(None))
+                        }
+
+                        // Still pending
+                        None => Either::Right(fetching.clone()),
+                    };
+                }
+
+                entry.remove();
+                let Entry::Vacant(entry) = self.fetch_jobs.entry(key.clone()) else {
+                    unreachable!()
+                };
+                entry
+            }
+        };
+
+        let fetcher = self.fetcher.clone();
+
+        let serialize = self.serialize;
+        let file_path = self.cache_dir.join(key.to_string());
+        let fetching = Box::pin(async move {
+            if let Ok(value) = fetch_from_cache.await.inspect_err(|e| {
+                tracing::debug!("Failed to fetch {key:?} from cache: {e}");
+            }) {
+                return Some(value);
+            }
+
+            let value = match fetcher.fetch(&key).await {
+                Ok(value) => value,
+                Err(e) => {
+                    tracing::warn!("Couldn't fetch {key:?}: {e}");
+                    return None;
+                }
+            };
+
+            let data = serialize(&key, &value);
+            if let Err(e) = fs::write(file_path, data).await {
+                tracing::error!("Failed to cahce value for {key:?}: {e}");
+            }
+
+            Some(value)
+        });
+        let fetching = fetching as Pin<Box<dyn Future<Output = Option<V>> + Send + Sync>>;
+        let fetching = fetching.shared();
+
+        if let Some(fetching) = fetching.downgrade() {
+            entry.insert(fetching);
+        }
+
+        Either::Right(fetching)
+    }
+}

src/main.rs

@@ -17,6 +17,7 @@ use crate::{
 };
 
 mod api;
+pub mod cachemap;
 mod thumbhash;
 
 mod ui {