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 }