diff --git a/Cargo.lock b/Cargo.lock index 6ee9349..e28eb68 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1661,6 +1661,7 @@ dependencies = [ "image", "jiff", "libc", + "memmap2 0.9.11", "notify", "serde", "serde_json", diff --git a/compositor/Cargo.toml b/compositor/Cargo.toml index 1de664e..e36acd2 100644 --- a/compositor/Cargo.toml +++ b/compositor/Cargo.toml @@ -22,6 +22,7 @@ jiff.workspace = true notify = "8.2.0" zbus = "5.19.0" image = { version = "0.25.10", default-features = false, features = ["png", "jpeg", "webp"] } +memmap2 = "0.9.11" smithay.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/compositor/src/wallpaper.rs b/compositor/src/wallpaper.rs index 3f5ed89..f21dcda 100644 --- a/compositor/src/wallpaper.rs +++ b/compositor/src/wallpaper.rs @@ -18,6 +18,9 @@ use smithay::utils::{Buffer, Physical, Rectangle, Size}; use crate::presentation::NativeTextureElement; +mod pixels; +use pixels::Pixels; + const UPLOAD_RETRY_DELAY: Duration = Duration::from_secs(5); #[derive(Clone, Debug, Deserialize, PartialEq)] @@ -59,8 +62,8 @@ pub(crate) struct WallpaperState { commit: CommitCounter, owned: bool, mode: WallpaperMode, - receiver: Option>>, - pixels: Option, + receiver: Option>>, + pixels: Option, textures: HashMap, upload_retries: HashMap, wakeup: Option, @@ -82,25 +85,7 @@ impl WallpaperState { let path = config.path.unwrap(); let wakeup = wakeup.clone(); std::thread::spawn(move || { - let load = || -> Result { - use image::ImageDecoder; - let mut reader = image::ImageReader::open(&path) - .map_err(|e| e.to_string())? - .with_guessed_format() - .map_err(|e| e.to_string())?; - let mut limits = image::Limits::default(); - limits.max_alloc = Some(256 * 1024 * 1024); - reader.limits(limits); - let decoder = reader.into_decoder().map_err(|e| e.to_string())?; - let (width, height) = decoder.dimensions(); - if u64::from(width) * u64::from(height) * 4 > 256 * 1024 * 1024 { - return Err("wallpaper exceeds the 256 MiB decode limit".into()); - } - image::DynamicImage::from_decoder(decoder) - .map(|image| image.into_rgba8()) - .map_err(|e| e.to_string()) - }; - let _ = sender.send(load()); + let _ = sender.send(pixels::load(&path)); if let Some(wakeup) = wakeup { wakeup.wakeup(); } @@ -356,7 +341,8 @@ mod tests { .recv_timeout(Duration::from_secs(5)) .unwrap() .unwrap(); - assert_eq!(decoded, expected); + assert_eq!(decoded.dimensions(), expected.dimensions()); + assert_eq!(&decoded.as_raw()[..], expected.as_raw()); } } @@ -439,7 +425,9 @@ mod tests { state.upload_retries.insert(context.clone(), now + UPLOAD_RETRY_DELAY); let (sender, receiver) = mpsc::channel(); state.receiver = Some(receiver); - sender.send(Ok(image::RgbaImage::new(2, 2))).unwrap(); + sender + .send(Ok(pixels::from_rgba(image::RgbaImage::new(2, 2)).unwrap())) + .unwrap(); assert!(state.poll()); assert!(state.upload_ready(&context, now)); } @@ -454,7 +442,9 @@ mod tests { }); assert!(state.receiver.is_some()); let before = state.commit; - sender.send(Ok(image::RgbaImage::new(2, 2))).unwrap(); + sender + .send(Ok(pixels::from_rgba(image::RgbaImage::new(2, 2)).unwrap())) + .unwrap(); assert!(state.poll()); assert_eq!(state.mode, WallpaperMode::Fit); assert!(state.pixels.is_some()); @@ -468,7 +458,7 @@ mod tests { #[test] fn reverting_a_pending_reload_keeps_latest_request_and_failed_decode_keeps_pixels() { let mut state = WallpaperState::new(WallpaperConfig::default()); - state.pixels = Some(image::RgbaImage::new(2, 2)); + state.pixels = Some(pixels::from_rgba(image::RgbaImage::new(2, 2)).unwrap()); let (sender, receiver) = mpsc::channel(); state.receiver = Some(receiver); state.reload(WallpaperConfig { diff --git a/compositor/src/wallpaper/pixels.rs b/compositor/src/wallpaper/pixels.rs new file mode 100644 index 0000000..bc869e7 --- /dev/null +++ b/compositor/src/wallpaper/pixels.rs @@ -0,0 +1,202 @@ +use std::path::Path; + +use image::{ColorType, ImageDecoder}; +use memmap2::MmapMut; + +pub(super) type Pixels = image::ImageBuffer, MmapMut>; + +const DECODE_LIMIT: u64 = 256 * 1024 * 1024; + +pub(super) fn load(path: &Path) -> Result { + let result = decode(path); + + // Decoders can leave large freed scratch buffers in glibc arenas. Reclaim + // them here, after decode has dropped its temporaries, on the worker only. + // The returned pixels have their own mapping and are unaffected by trimming. + #[cfg(all(target_os = "linux", target_env = "gnu"))] + unsafe { + libc::malloc_trim(0); + } + + result +} + +fn decode(path: &Path) -> Result { + let mut reader = image::ImageReader::open(path) + .map_err(|error| error.to_string())? + .with_guessed_format() + .map_err(|error| error.to_string())?; + let mut limits = image::Limits::default(); + limits.max_alloc = Some(DECODE_LIMIT); + reader.limits(limits); + let decoder = reader.into_decoder().map_err(|error| error.to_string())?; + let (width, height) = decoder.dimensions(); + let bytes = rgba_bytes(width, height)?; + + // Give both the retained pixels and RGB conversion scratch their own + // mappings. Dropping either releases its pages even when another thread + // owns the image, rather than leaving them in the decoder thread's arena. + match decoder.color_type() { + ColorType::Rgba8 => { + let mut data = MmapMut::map_anon(bytes).map_err(|error| error.to_string())?; + decoder.read_image(&mut data).map_err(|error| error.to_string())?; + mapped_image(width, height, data) + } + + ColorType::Rgb8 => { + let mut rgb = MmapMut::map_anon(bytes / 4 * 3).map_err(|error| error.to_string())?; + decoder.read_image(&mut rgb).map_err(|error| error.to_string())?; + let mut rgba = MmapMut::map_anon(bytes).map_err(|error| error.to_string())?; + for (source, target) in rgb.chunks_exact(3).zip(rgba.chunks_exact_mut(4)) { + target[..3].copy_from_slice(source); + target[3] = 255; + } + + mapped_image(width, height, rgba) + } + + _ => { + let decoded = image::DynamicImage::from_decoder(decoder).map_err(|error| error.to_string())?; + from_rgba(decoded.into_rgba8()) + } + } +} + +fn rgba_bytes(width: u32, height: u32) -> Result { + let pixels = u64::from(width) * u64::from(height); + if pixels == 0 || pixels > DECODE_LIMIT / 4 { + return Err("wallpaper must be nonempty and within the 256 MiB decode limit".into()); + } + + Ok(pixels as usize * 4) +} + +pub(super) fn from_rgba(image: image::RgbaImage) -> Result { + let (width, height) = image.dimensions(); + let mut data = MmapMut::map_anon(rgba_bytes(width, height)?).map_err(|error| error.to_string())?; + data.copy_from_slice(image.as_raw()); + mapped_image(width, height, data) +} + +fn mapped_image(width: u32, height: u32, data: MmapMut) -> Result { + Pixels::from_raw(width, height, data).ok_or_else(|| "invalid wallpaper pixel dimensions".into()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn mapped_pixels_match_existing_decoder_for_png_jpeg_and_webp() { + use image::{DynamicImage, ImageBuffer, ImageFormat, Luma, LumaA, Rgb, Rgba}; + let directory = tempfile::tempdir().unwrap(); + let images = [ + DynamicImage::ImageRgb8(ImageBuffer::from_pixel(7, 5, Rgb([31, 81, 159]))), + DynamicImage::ImageRgba8(ImageBuffer::from_pixel(7, 5, Rgba([31, 81, 159, 127]))), + DynamicImage::ImageLuma8(ImageBuffer::from_pixel(7, 5, Luma([81]))), + DynamicImage::ImageLumaA8(ImageBuffer::from_pixel(7, 5, LumaA([81, 127]))), + DynamicImage::ImageRgb16(ImageBuffer::from_pixel(7, 5, Rgb([8000, 21000, 41000]))), + DynamicImage::ImageRgba16(ImageBuffer::from_pixel(7, 5, Rgba([8000, 21000, 41000, 32000]))), + ]; + for (index, image) in images.iter().enumerate() { + let path = directory.path().join(format!("image-{index}.png")); + image.save(&path).unwrap(); + let expected = image::open(&path).unwrap().into_rgba8(); + let actual = load(&path).unwrap(); + assert_eq!(actual.dimensions(), expected.dimensions()); + assert_eq!(&actual.as_raw()[..], expected.as_raw()); + } + + for format in [ImageFormat::Jpeg, ImageFormat::WebP] { + let path = directory.path().join("without-extension"); + images[0].save_with_format(&path, format).unwrap(); + let expected = image::ImageReader::open(&path) + .unwrap() + .with_guessed_format() + .unwrap() + .decode() + .unwrap() + .into_rgba8(); + let actual = load(&path).unwrap(); + assert_eq!(&actual.as_raw()[..], expected.as_raw()); + } + } + + #[test] + fn empty_and_oversized_dimensions_are_rejected_without_overflow() { + assert!(rgba_bytes(0, 1).is_err()); + assert!(rgba_bytes(u32::MAX, u32::MAX).is_err()); + assert!(rgba_bytes(8192, 8193).is_err()); + assert_eq!(rgba_bytes(8192, 8192).unwrap(), DECODE_LIMIT as usize); + } + + #[test] + fn corrupt_file_does_not_publish_pixels() { + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(file.path(), b"not an image").unwrap(); + assert!(load(file.path()).is_err()); + } + + #[test] + #[ignore = "isolated process RSS regression; run in release with --ignored"] + fn repeated_wallpaper_loads_release_resident_memory() { + const CHILD: &str = "FERESE_WALLPAPER_MEMORY_TEST_CHILD"; + if std::env::var_os(CHILD).is_none() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "wallpaper::pixels::tests::repeated_wallpaper_loads_release_resident_memory", + "--ignored", + "--nocapture", + ]) + .env(CHILD, "1") + .output() + .unwrap(); + println!("{}", String::from_utf8_lossy(&output.stdout)); + assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr)); + return; + } + + fn rss_kib() -> usize { + std::fs::read_to_string("/proc/self/status") + .unwrap() + .lines() + .find(|line| line.starts_with("VmRSS:")) + .unwrap() + .split_whitespace() + .nth(1) + .unwrap() + .parse() + .unwrap() + } + + let paths = [ + Path::new(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../assets/wallpapers/ferese-wallpaper-dark.jpg" + )), + Path::new(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../assets/wallpapers/ferese-wallpaper-light.png" + )), + ]; + drop(std::thread::spawn(move || load(paths[0]).unwrap()).join().unwrap()); + let baseline = rss_kib(); + let mut current = None; + for index in 0..20 { + // Transfer ownership out of the decoder thread, as WallpaperState does. + let path = paths[index % paths.len()]; + current = Some(std::thread::spawn(move || load(path).unwrap()).join().unwrap()); + } + + let with_wallpaper = rss_kib(); + drop(current); + let released = rss_kib(); + println!("RSS KiB: baseline={baseline}, wallpaper={with_wallpaper}, released={released}"); + assert!( + released <= baseline + 16 * 1024, + "repeated decoder threads retained {} KiB after all pixels were dropped", + released.saturating_sub(baseline) + ); + } +}