notedeck

One damus client to rule them all
git clone git://jb55.com/notedeck
Log | Files | Refs | README | LICENSE

imgcache.rs (13916B)

      1 use crate::urls::{UrlCache, UrlMimes};
      2 use crate::Result;
      3 use egui::TextureHandle;
      4 use image::{Delay, Frame};
      5 use poll_promise::Promise;
      6 
      7 use egui::ColorImage;
      8 
      9 use std::collections::HashMap;
     10 use std::fs::{self, create_dir_all, File};
     11 use std::sync::mpsc::Receiver;
     12 use std::sync::{Arc, Mutex};
     13 use std::time::{Duration, Instant, SystemTime};
     14 use std::{io, thread};
     15 
     16 use hex::ToHex;
     17 use sha2::Digest;
     18 use std::path::PathBuf;
     19 use std::path::{self, Path};
     20 use tracing::warn;
     21 
     22 #[derive(Default)]
     23 pub struct TexturesCache {
     24     cache: hashbrown::HashMap<String, TextureStateInternal>,
     25 }
     26 
     27 impl TexturesCache {
     28     pub fn handle_and_get_or_insert_loadable(
     29         &mut self,
     30         url: &str,
     31         closure: impl FnOnce() -> Promise<Option<Result<TexturedImage>>>,
     32     ) -> LoadableTextureState {
     33         let internal = self.handle_and_get_state_internal(url, true, closure);
     34 
     35         internal.into()
     36     }
     37 
     38     pub fn handle_and_get_or_insert(
     39         &mut self,
     40         url: &str,
     41         closure: impl FnOnce() -> Promise<Option<Result<TexturedImage>>>,
     42     ) -> TextureState {
     43         let internal = self.handle_and_get_state_internal(url, false, closure);
     44 
     45         internal.into()
     46     }
     47 
     48     fn handle_and_get_state_internal(
     49         &mut self,
     50         url: &str,
     51         use_loading: bool,
     52         closure: impl FnOnce() -> Promise<Option<Result<TexturedImage>>>,
     53     ) -> &mut TextureStateInternal {
     54         let state = match self.cache.raw_entry_mut().from_key(url) {
     55             hashbrown::hash_map::RawEntryMut::Occupied(entry) => {
     56                 let state = entry.into_mut();
     57                 handle_occupied(state, use_loading);
     58 
     59                 state
     60             }
     61             hashbrown::hash_map::RawEntryMut::Vacant(entry) => {
     62                 let res = closure();
     63                 let (_, state) = entry.insert(url.to_owned(), TextureStateInternal::Pending(res));
     64 
     65                 state
     66             }
     67         };
     68 
     69         state
     70     }
     71 
     72     pub fn insert_pending(&mut self, url: &str, promise: Promise<Option<Result<TexturedImage>>>) {
     73         self.cache
     74             .insert(url.to_owned(), TextureStateInternal::Pending(promise));
     75     }
     76 
     77     pub fn move_to_loaded(&mut self, url: &str) {
     78         let hashbrown::hash_map::RawEntryMut::Occupied(entry) =
     79             self.cache.raw_entry_mut().from_key(url)
     80         else {
     81             return;
     82         };
     83 
     84         entry.replace_entry_with(|_, v| {
     85             let TextureStateInternal::Loading(textured) = v else {
     86                 return None;
     87             };
     88 
     89             Some(TextureStateInternal::Loaded(textured))
     90         });
     91     }
     92 
     93     pub fn get_and_handle(&mut self, url: &str) -> Option<LoadableTextureState> {
     94         self.cache.get_mut(url).map(|state| {
     95             handle_occupied(state, true);
     96             state.into()
     97         })
     98     }
     99 }
    100 
    101 fn handle_occupied(state: &mut TextureStateInternal, use_loading: bool) {
    102     let TextureStateInternal::Pending(promise) = state else {
    103         return;
    104     };
    105 
    106     let Some(res) = promise.ready_mut() else {
    107         return;
    108     };
    109 
    110     let Some(res) = res.take() else {
    111         tracing::error!("Failed to take the promise");
    112         *state =
    113             TextureStateInternal::Error(crate::Error::Generic("Promise already taken".to_owned()));
    114         return;
    115     };
    116 
    117     match res {
    118         Ok(textured) => {
    119             *state = if use_loading {
    120                 TextureStateInternal::Loading(textured)
    121             } else {
    122                 TextureStateInternal::Loaded(textured)
    123             }
    124         }
    125         Err(e) => *state = TextureStateInternal::Error(e),
    126     }
    127 }
    128 
    129 pub enum LoadableTextureState<'a> {
    130     Pending,
    131     Error(&'a crate::Error),
    132     Loading {
    133         actual_image_tex: &'a mut TexturedImage,
    134     }, // the texture is in the loading state, for transitioning between the pending and loaded states
    135     Loaded(&'a mut TexturedImage),
    136 }
    137 
    138 pub enum TextureState<'a> {
    139     Pending,
    140     Error(&'a crate::Error),
    141     Loaded(&'a mut TexturedImage),
    142 }
    143 
    144 impl<'a> From<&'a mut TextureStateInternal> for TextureState<'a> {
    145     fn from(value: &'a mut TextureStateInternal) -> Self {
    146         match value {
    147             TextureStateInternal::Pending(_) => TextureState::Pending,
    148             TextureStateInternal::Error(error) => TextureState::Error(error),
    149             TextureStateInternal::Loading(textured_image) => TextureState::Loaded(textured_image),
    150             TextureStateInternal::Loaded(textured_image) => TextureState::Loaded(textured_image),
    151         }
    152     }
    153 }
    154 
    155 pub enum TextureStateInternal {
    156     Pending(Promise<Option<Result<TexturedImage>>>),
    157     Error(crate::Error),
    158     Loading(TexturedImage), // the image is in the loading state, for transitioning between blur and image
    159     Loaded(TexturedImage),
    160 }
    161 
    162 impl<'a> From<&'a mut TextureStateInternal> for LoadableTextureState<'a> {
    163     fn from(value: &'a mut TextureStateInternal) -> Self {
    164         match value {
    165             TextureStateInternal::Pending(_) => LoadableTextureState::Pending,
    166             TextureStateInternal::Error(error) => LoadableTextureState::Error(error),
    167             TextureStateInternal::Loading(textured_image) => LoadableTextureState::Loading {
    168                 actual_image_tex: textured_image,
    169             },
    170             TextureStateInternal::Loaded(textured_image) => {
    171                 LoadableTextureState::Loaded(textured_image)
    172             }
    173         }
    174     }
    175 }
    176 
    177 pub enum TexturedImage {
    178     Static(TextureHandle),
    179     Animated(Animation),
    180 }
    181 
    182 impl TexturedImage {
    183     pub fn get_first_texture(&self) -> &TextureHandle {
    184         match self {
    185             TexturedImage::Static(texture_handle) => texture_handle,
    186             TexturedImage::Animated(animation) => &animation.first_frame.texture,
    187         }
    188     }
    189 }
    190 
    191 pub struct Animation {
    192     pub first_frame: TextureFrame,
    193     pub other_frames: Vec<TextureFrame>,
    194     pub receiver: Option<Receiver<TextureFrame>>,
    195 }
    196 
    197 impl Animation {
    198     pub fn get_frame(&self, index: usize) -> Option<&TextureFrame> {
    199         if index == 0 {
    200             Some(&self.first_frame)
    201         } else {
    202             self.other_frames.get(index - 1)
    203         }
    204     }
    205 
    206     pub fn num_frames(&self) -> usize {
    207         self.other_frames.len() + 1
    208     }
    209 }
    210 
    211 pub struct TextureFrame {
    212     pub delay: Duration,
    213     pub texture: TextureHandle,
    214 }
    215 
    216 pub struct ImageFrame {
    217     pub delay: Duration,
    218     pub image: ColorImage,
    219 }
    220 
    221 pub struct MediaCache {
    222     pub cache_dir: path::PathBuf,
    223     pub textures_cache: TexturesCache,
    224     pub cache_type: MediaCacheType,
    225     pub cache_size: Arc<Mutex<Option<u64>>>,
    226 }
    227 
    228 #[derive(Debug, Eq, PartialEq, Clone, Copy)]
    229 pub enum MediaCacheType {
    230     Image,
    231     Gif,
    232 }
    233 
    234 impl MediaCache {
    235     pub fn new(parent_dir: &Path, cache_type: MediaCacheType) -> Self {
    236         let cache_dir = parent_dir.join(Self::rel_dir(cache_type));
    237 
    238         let cache_dir_clone = cache_dir.clone();
    239         let cache_size = Arc::new(Mutex::new(None));
    240         let cache_size_clone = Arc::clone(&cache_size);
    241 
    242         thread::spawn(move || {
    243             let mut last_checked = Instant::now() - Duration::from_secs(999);
    244             loop {
    245                 // check cache folder size every 60 s
    246                 if last_checked.elapsed() >= Duration::from_secs(60) {
    247                     let size = compute_folder_size(&cache_dir_clone);
    248                     *cache_size_clone.lock().unwrap() = Some(size);
    249                     last_checked = Instant::now();
    250                 }
    251                 thread::sleep(Duration::from_secs(5));
    252             }
    253         });
    254 
    255         Self {
    256             cache_dir,
    257             textures_cache: TexturesCache::default(),
    258             cache_type,
    259             cache_size,
    260         }
    261     }
    262 
    263     pub fn rel_dir(cache_type: MediaCacheType) -> &'static str {
    264         match cache_type {
    265             MediaCacheType::Image => "img",
    266             MediaCacheType::Gif => "gif",
    267         }
    268     }
    269 
    270     pub fn write(cache_dir: &path::Path, url: &str, data: ColorImage) -> Result<()> {
    271         let file = Self::create_file(cache_dir, url)?;
    272         let encoder = image::codecs::webp::WebPEncoder::new_lossless(file);
    273 
    274         encoder.encode(
    275             data.as_raw(),
    276             data.size[0] as u32,
    277             data.size[1] as u32,
    278             image::ColorType::Rgba8.into(),
    279         )?;
    280 
    281         Ok(())
    282     }
    283 
    284     fn create_file(cache_dir: &path::Path, url: &str) -> Result<File> {
    285         let file_path = cache_dir.join(Self::key(url));
    286         if let Some(p) = file_path.parent() {
    287             create_dir_all(p)?;
    288         }
    289         Ok(File::options()
    290             .write(true)
    291             .create(true)
    292             .truncate(true)
    293             .open(file_path)?)
    294     }
    295 
    296     pub fn write_gif(cache_dir: &path::Path, url: &str, data: Vec<ImageFrame>) -> Result<()> {
    297         let file = Self::create_file(cache_dir, url)?;
    298 
    299         let mut encoder = image::codecs::gif::GifEncoder::new(file);
    300         for img in data {
    301             let buf = color_image_to_rgba(img.image);
    302             let frame = Frame::from_parts(buf, 0, 0, Delay::from_saturating_duration(img.delay));
    303             if let Err(e) = encoder.encode_frame(frame) {
    304                 tracing::error!("problem encoding frame: {e}");
    305             }
    306         }
    307 
    308         Ok(())
    309     }
    310 
    311     pub fn key(url: &str) -> String {
    312         let k: String = sha2::Sha256::digest(url.as_bytes()).encode_hex();
    313         PathBuf::from(&k[0..2])
    314             .join(&k[2..4])
    315             .join(k)
    316             .to_string_lossy()
    317             .to_string()
    318     }
    319 
    320     /// Migrate from base32 encoded url to sha256 url + sub-dir structure
    321     pub fn migrate_v0(&self) -> Result<()> {
    322         for file in std::fs::read_dir(&self.cache_dir)? {
    323             let file = if let Ok(f) = file {
    324                 f
    325             } else {
    326                 // not sure how this could fail, skip entry
    327                 continue;
    328             };
    329             if !file.path().is_file() {
    330                 continue;
    331             }
    332             let old_filename = file.file_name().to_string_lossy().to_string();
    333             let old_url = if let Some(u) =
    334                 base32::decode(base32::Alphabet::Crockford, &old_filename)
    335                     .and_then(|s| String::from_utf8(s).ok())
    336             {
    337                 u
    338             } else {
    339                 warn!("Invalid base32 filename: {}", &old_filename);
    340                 continue;
    341             };
    342             let new_path = self.cache_dir.join(Self::key(&old_url));
    343             if let Some(p) = new_path.parent() {
    344                 create_dir_all(p)?;
    345             }
    346 
    347             if let Err(e) = std::fs::rename(file.path(), &new_path) {
    348                 warn!(
    349                     "Failed to migrate file from {} to {}: {:?}",
    350                     file.path().display(),
    351                     new_path.display(),
    352                     e
    353                 );
    354             }
    355         }
    356 
    357         Ok(())
    358     }
    359 
    360     fn clear(&mut self) {
    361         self.textures_cache.cache.clear();
    362         *self.cache_size.try_lock().unwrap() = Some(0);
    363     }
    364 }
    365 
    366 fn color_image_to_rgba(color_image: ColorImage) -> image::RgbaImage {
    367     let width = color_image.width() as u32;
    368     let height = color_image.height() as u32;
    369 
    370     let rgba_pixels: Vec<u8> = color_image
    371         .pixels
    372         .iter()
    373         .flat_map(|color| color.to_array()) // Convert Color32 to `[u8; 4]`
    374         .collect();
    375 
    376     image::RgbaImage::from_raw(width, height, rgba_pixels)
    377         .expect("Failed to create RgbaImage from ColorImage")
    378 }
    379 
    380 fn compute_folder_size<P: AsRef<Path>>(path: P) -> u64 {
    381     fn walk(path: &Path) -> u64 {
    382         let mut size = 0;
    383         if let Ok(entries) = fs::read_dir(path) {
    384             for entry in entries.flatten() {
    385                 let path = entry.path();
    386                 if let Ok(metadata) = entry.metadata() {
    387                     if metadata.is_file() {
    388                         size += metadata.len();
    389                     } else if metadata.is_dir() {
    390                         size += walk(&path);
    391                     }
    392                 }
    393             }
    394         }
    395         size
    396     }
    397     walk(path.as_ref())
    398 }
    399 
    400 pub struct Images {
    401     pub base_path: path::PathBuf,
    402     pub static_imgs: MediaCache,
    403     pub gifs: MediaCache,
    404     pub urls: UrlMimes,
    405     pub gif_states: GifStateMap,
    406 }
    407 
    408 impl Images {
    409     /// path to directory to place [`MediaCache`]s
    410     pub fn new(path: path::PathBuf) -> Self {
    411         Self {
    412             base_path: path.clone(),
    413             static_imgs: MediaCache::new(&path, MediaCacheType::Image),
    414             gifs: MediaCache::new(&path, MediaCacheType::Gif),
    415             urls: UrlMimes::new(UrlCache::new(path.join(UrlCache::rel_dir()))),
    416             gif_states: Default::default(),
    417         }
    418     }
    419 
    420     pub fn migrate_v0(&self) -> Result<()> {
    421         self.static_imgs.migrate_v0()?;
    422         self.gifs.migrate_v0()
    423     }
    424 
    425     pub fn get_cache(&self, cache_type: MediaCacheType) -> &MediaCache {
    426         match cache_type {
    427             MediaCacheType::Image => &self.static_imgs,
    428             MediaCacheType::Gif => &self.gifs,
    429         }
    430     }
    431 
    432     pub fn get_cache_mut(&mut self, cache_type: MediaCacheType) -> &mut MediaCache {
    433         match cache_type {
    434             MediaCacheType::Image => &mut self.static_imgs,
    435             MediaCacheType::Gif => &mut self.gifs,
    436         }
    437     }
    438 
    439     pub fn clear_folder_contents(&mut self) -> io::Result<()> {
    440         for entry in fs::read_dir(self.base_path.clone())? {
    441             let entry = entry?;
    442             let path = entry.path();
    443 
    444             if path.is_dir() {
    445                 fs::remove_dir_all(path)?;
    446             } else {
    447                 fs::remove_file(path)?;
    448             }
    449         }
    450 
    451         self.urls.cache.clear();
    452         self.static_imgs.clear();
    453         self.gifs.clear();
    454         self.gif_states.clear();
    455 
    456         Ok(())
    457     }
    458 }
    459 
    460 pub type GifStateMap = HashMap<String, GifState>;
    461 
    462 pub struct GifState {
    463     pub last_frame_rendered: Instant,
    464     pub last_frame_duration: Duration,
    465     pub next_frame_time: Option<SystemTime>,
    466     pub last_frame_index: usize,
    467 }