notedeck

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

imgcache.rs (16794B)


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