From b6d0b02878e97085dce125bffc9fef49970e5e51 Mon Sep 17 00:00:00 2001 From: mavrick Date: Thu, 12 Nov 2026 23:20:04 -0600 Subject: [PATCH 1/2] perf(data): compress GeoTIFF tiles and strips on several threads (#1111) The COG writer's tiles and the untiled GeoTIFF's strips are each gathered, predicted and Deflated on their own, so on the desktop worker threads make them while the calling thread writes them in their order: the same file as one thread writes, byte for byte. Workers stay at most two pieces a thread ahead of the writer, so memory stays bounded; cancel and progress stay on the writing thread. The browser and one-piece rasters keep one thread. On the 9667 x 11947 DEM the COG export takes 6.6-11.7 s instead of 15.1-23.4 s, the untiled export 4.8-9.9 s instead of 7.7-12.1 s. Co-Authored-By: Claude Opus 5.5 --- crates/data/src/raster/cogwrite.rs | 303 ++++++++++++++++++++++++----- crates/data/src/raster/geotiff.rs | 2 +- 2 files changed, 260 insertions(+), 45 deletions(-) diff --git a/crates/data/src/raster/cogwrite.rs b/crates/data/src/raster/cogwrite.rs index f42540bc3..02f06f4ec 100644 --- a/crates/data/src/raster/cogwrite.rs +++ b/crates/data/src/raster/cogwrite.rs @@ -8,9 +8,10 @@ //! directories' sizes depend only on the tile counts, so they are written first with empty tile //! tables, and the tables are filled in once the tiles' places are known. //! -//! Memory: the raster is read in place. A tile is gathered from the bands, compressed and written -//! before the next one, so no second copy of the pixels is made, and each overview tile is made -//! straight from the full image (a block of 2ⁿ × 2ⁿ pixels), so no pyramid is kept. The price is +//! Memory: the raster is read in place. A tile is gathered from the bands and compressed on its +//! own, and only a few tiles a thread are held before they are written, so no second copy of the +//! pixels is made, and each overview tile is made straight from the full image (a block of +//! 2ⁿ × 2ⁿ pixels), so no pyramid is kept. The price is //! time: every overview level reads the whole full image again, so writing takes about as many //! passes over the pixels as there are levels, where building each level from the one before //! would take about 1⅓ but hold a second copy. @@ -24,11 +25,15 @@ //! don't change smoothly, so differences make them no smaller, and their long runs compress quickly //! at any level, to half the size of the fastest. //! +//! On the desktop the tiles are made on several threads ([`crate::threads::max`]) while one thread +//! writes them in their order, so the file is the same as one thread writes, byte for byte; in the +//! browser they are made one after another. +//! //! Classic TIFF is used when the file can't pass 4 GiB even with every tile left uncompressed, //! BigTIFF otherwise. //! //! The untiled GeoTIFF ([`super::geotiff::write`]) is written here too, by [`write_strips`]: one -//! image in strips of rows, with the same predictors and Deflate effort, and no overviews. +//! image in strips of rows, with the same predictors, Deflate effort and threads, and no overviews. use std::io::{Seek, SeekFrom, Write}; use std::ops::Range; @@ -689,28 +694,39 @@ pub fn write_to(r: &Raster, band: usize, flavour: Flavour, out: let mut pos = head.bytes.len() as u64; // Smallest overview first, the full image last; each in row order. + let order: Vec<(usize, u32, u32)> = levels + .iter() + .enumerate() + .rev() + .flat_map(|(i, lv)| (0..lv.down).flat_map(move |ty| (0..lv.across).map(move |tx| (i, tx, ty)))) + .collect(); let mut places: Vec> = vec![Vec::new(); levels.len()]; - crate::progress::begin(tiles); - let mut row = Vec::new(); let effort = if plan.categorical { DEFLATE_CLASSES } else { DEFLATE_FAST }; - for (i, lv) in levels.iter().enumerate().rev() { - for ty in 0..lv.down { - for tx in 0..lv.across { - crate::progress::step(1)?; - let mut raw = tile_samples(r, &plan, lv, tx, ty); - predict(&mut raw, &plan, &mut row); - let zipped = miniz_oxide::deflate::compress_to_vec_zlib(&raw, effort); - if !big && pos + zipped.len() as u64 > u64::from(u32::MAX) { - return Err(err("the file is too large for classic TIFF")); - } - out.write_all(&zipped).map_err(err)?; - if let Some(p) = places.get_mut(i) { - p.push((pos, zipped.len() as u64)); - } - pos += zipped.len() as u64; + let tile_bytes = (TILE as usize * TILE as usize).saturating_mul(spp * bps); + in_order( + order.len(), + tile_bytes, + err, + Vec::new, + |row, k| { + let (i, tx, ty) = order.get(k).copied().ok_or_else(|| err("a tile past the image"))?; + let lv = levels.get(i).ok_or_else(|| err("a tile past the image"))?; + let mut raw = tile_samples(r, &plan, lv, tx, ty); + predict(&mut raw, &plan, row); + Ok(miniz_oxide::deflate::compress_to_vec_zlib(&raw, effort)) + }, + |k, zipped| { + if !big && pos + zipped.len() as u64 > u64::from(u32::MAX) { + return Err(err("the file is too large for classic TIFF")); } - } - } + out.write_all(&zipped).map_err(err)?; + if let Some(p) = order.get(k).and_then(|(i, _, _)| places.get_mut(*i)) { + p.push((pos, zipped.len() as u64)); + } + pos += zipped.len() as u64; + Ok(()) + }, + )?; fill_tables(out, &head.patches, &places, big, pos, err)?; Ok(CogInfo { big, levels: levels.iter().map(|l| (l.w, l.h)).collect(), bytes: pos }) @@ -768,13 +784,13 @@ const STRIP_BYTES: u64 = 1 << 20; /// Writes an untiled TIFF of one image into `out` (which must be empty: offsets are from its /// start): its directory and tags, then its strips, each made by `rows` (which appends the /// samples of a range of rows, interleaved and little-endian), run through the predictor and -/// compressed with Deflate before the next is made. Classic TIFF when the file can't pass 4 GiB, -/// BigTIFF otherwise. It counts the strips as its progress, and a cancelled run stops it within a -/// strip. +/// compressed with Deflate, on several threads on the desktop ([`in_order`]). Classic TIFF when +/// the file can't pass 4 GiB, BigTIFF otherwise. It counts the strips as its progress, and a +/// cancelled run stops it within the few strips being made. pub(super) fn write_strips( img: Strips, out: &mut W, - mut rows: impl FnMut(Range, &mut Vec) -> Result<()>, + rows: impl Fn(Range, &mut Vec) -> Result<()> + Sync, ) -> Result<()> { if img.w == 0 || img.h == 0 { return Err(strip_err("the raster has no pixels to write")); @@ -820,26 +836,178 @@ pub(super) fn write_strips( let effort = if img.classes { DEFLATE_CLASSES } else { DEFLATE_FAST }; let mut places = vec![Vec::with_capacity(strips)]; - let (mut raw, mut row) = (Vec::new(), Vec::new()); - crate::progress::begin(strips); - for y in (0..img.h as usize).step_by(per_strip as usize) { + let (h, per_strip) = (img.h as usize, per_strip as usize); + in_order( + strips, + line.saturating_mul(per_strip), + strip_err, + || (Vec::new(), Vec::new()), + |(raw, row), k| { + let y = k * per_strip; + raw.clear(); + rows(y..(y + per_strip).min(h), raw)?; + predict_rows(raw, line, spp, img.sample, predictor, row); + Ok(miniz_oxide::deflate::compress_to_vec_zlib(raw, effort)) + }, + |_, zipped| { + if !big && pos + zipped.len() as u64 > u64::from(u32::MAX) { + return Err(strip_err("the file is too large for classic TIFF")); + } + out.write_all(&zipped).map_err(strip_err)?; + if let Some(p) = places.first_mut() { + p.push((pos, zipped.len() as u64)); + } + pos += zipped.len() as u64; + Ok(()) + }, + )?; + fill_tables(out, &head.patches, &places, big, pos, strip_err) +} + +/// Makes `count` pieces of a file (tiles or strips, each about `bytes` uncompressed) with `make` +/// and hands each to `put` in their order, with its index. `make` gets the scratch its thread +/// made with `scratch` once; `err` words a failure of the threads. Every piece stands alone, so +/// on the desktop, with more than one piece, they are made on several threads +/// ([`crate::threads::max`], fewer when [`crate::memory::thread_scratch`] can't hold their pieces) +/// while this thread writes them: the same file as one thread makes, byte for byte. In the browser +/// they are made one after another. +/// +/// It counts the pieces as its progress, and a cancelled run stops it within the few pieces being +/// made ([`AHEAD`] a thread). The first error in the pieces' order comes back, as on one thread. +fn in_order( + count: usize, + bytes: usize, + err: fn(String) -> DataError, + scratch: impl Fn() -> S + Sync, + make: impl Fn(&mut S, usize) -> Result> + Sync, + put: impl FnMut(usize, Vec) -> Result<()>, +) -> Result<()> { + crate::progress::begin(count); + // A piece being made holds about twice its bytes (its samples and their compressed copy), and + // each waiting to be written up to its bytes again. + let affordable = 1 + crate::memory::thread_scratch() / bytes.saturating_mul(2 + AHEAD).max(1); + let threads = crate::threads::for_items(count).min(affordable); + if threads > 1 { + return in_order_parallel(count, threads, err, scratch, make, put); + } + one_by_one(count, scratch, make, put)?; + // A cancel that came while the last piece was made stops the write too. + crate::progress::check() +} + +/// [`in_order`] on this thread alone: each piece made, then written, before the next. +fn one_by_one( + count: usize, + scratch: impl Fn() -> S, + make: impl Fn(&mut S, usize) -> Result>, + mut put: impl FnMut(usize, Vec) -> Result<()>, +) -> Result<()> { + let mut s = scratch(); + for k in 0..count { crate::progress::step(1)?; - raw.clear(); - rows(y..(y + per_strip as usize).min(img.h as usize), &mut raw)?; - predict_rows(&mut raw, line, spp, img.sample, predictor, &mut row); - let zipped = miniz_oxide::deflate::compress_to_vec_zlib(&raw, effort); - if !big && pos + zipped.len() as u64 > u64::from(u32::MAX) { - return Err(strip_err("the file is too large for classic TIFF")); + put(k, make(&mut s, k)?)?; + } + Ok(()) +} + +/// The pieces a thread of [`in_order`] may make ahead of the one being written: enough that the +/// threads rarely wait for a slow piece, few enough that the pieces waiting stay small. +const AHEAD: usize = 2; + +/// [`in_order`] on `threads` threads besides this one, which writes. A thread takes the next piece +/// unless it is more than [`AHEAD`] a thread past the one to write next, so a slow piece holds up +/// only that many. A piece's panic comes back as an error in its place. A thread the OS refuses +/// costs speed, not the write; when it refuses all of them, this thread makes the pieces itself. +fn in_order_parallel( + count: usize, + threads: usize, + err: fn(String) -> DataError, + scratch: impl Fn() -> S + Sync, + make: impl Fn(&mut S, usize) -> Result> + Sync, + mut put: impl FnMut(usize, Vec) -> Result<()>, +) -> Result<()> { + use std::collections::BTreeMap; + use std::sync::{Condvar, Mutex, mpsc}; + + /// The next piece to take, the next to write, and whether to stop taking pieces. + struct Turn { + next: usize, + written: usize, + stop: bool, + } + let ahead = threads.saturating_mul(AHEAD); + let turn = Mutex::new(Turn { next: 0, written: 0, stop: false }); + let moved = Condvar::new(); + let failed = || err("a piece of the image couldn't be made".into()); + // Takes the next piece once it is no more than `ahead` past the one to write; None to stop. + let take = || -> Option { + let mut t = turn.lock().ok()?; + while !t.stop && t.next < count && t.next >= t.written.saturating_add(ahead) { + t = moved.wait(t).ok()?; } - out.write_all(&zipped).map_err(strip_err)?; - if let Some(p) = places.first_mut() { - p.push((pos, zipped.len() as u64)); + if t.stop || t.next >= count { + return None; } - pos += zipped.len() as u64; - } - // A cancel that came while the last strip was made stops the write too. - crate::progress::check()?; - fill_tables(out, &head.patches, &places, big, pos, strip_err) + t.next += 1; + Some(t.next - 1) + }; + let stop = || { + if let Ok(mut t) = turn.lock() { + t.stop = true; + } + moved.notify_all(); + }; + + std::thread::scope(|scope| { + let (tx, rx) = mpsc::channel::<(usize, Result>)>(); + let work = |tx: mpsc::Sender<(usize, Result>)>| { + let mut s = scratch(); + while let Some(k) = take() { + let made = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| make(&mut s, k))); + let panicked = made.is_err(); + // A thread whose piece panicked stops: its scratch may be half changed. + if tx.send((k, made.unwrap_or_else(|_| Err(failed())))).is_err() || panicked { + break; + } + } + }; + let started = (0..threads) + .filter_map(|_| { + let tx = tx.clone(); + std::thread::Builder::new().spawn_scoped(scope, move || work(tx)).ok() + }) + .count(); + drop(tx); + let res = (|| -> Result<()> { + if started == 0 { + return one_by_one(count, &scratch, &make, &mut put); + } + // Pieces made before the ones ahead of them, by their index. + let mut waiting = BTreeMap::new(); + for k in 0..count { + let piece = loop { + if let Some(p) = waiting.remove(&k) { + break p; + } + // Every thread gone before piece `k` came: one panicked outside a piece. + let (i, p) = rx.recv().map_err(|_| failed())?; + waiting.insert(i, p); + }; + crate::progress::step(1)?; + put(k, piece?)?; + if let Ok(mut t) = turn.lock() { + t.written = k + 1; + } + moved.notify_all(); + } + Ok(()) + })(); + // The threads finish the pieces they are making and take no more; the scope waits for them. + stop(); + res + })?; + // A cancel that came while the last piece was made stops the write too. + crate::progress::check() } /// The bands [`write`] and [`write_file`] write of `r` given `band`, from the same plan. @@ -1397,6 +1565,53 @@ mod tests { assert!(geotiff::write(&r, 0).is_ok()); } + /// Pieces made on several threads, the later ones sooner, are written in their order; the + /// first failure in that order comes back, and a piece that panics is an error, not a crash. + #[test] + fn pieces_made_on_threads_are_written_in_order() { + let slow_first = |_: &mut (), k: usize| { + std::thread::sleep(std::time::Duration::from_millis(((40 - k) % 7) as u64)); + Ok(vec![k as u8]) + }; + let mut written = Vec::new(); + in_order_parallel( + 40, + 4, + err, + || (), + slow_first, + |k, p| { + written.push((k, p)); + Ok(()) + }, + ) + .unwrap(); + assert_eq!(written, (0..40).map(|k| (k, vec![k as u8])).collect::>()); + + let failing = |_: &mut (), k: usize| match k { + 9 => Err(err("nine")), + 17 => panic!("seventeen"), + _ => Ok(Vec::new()), + }; + let mut n = 0; + let e = in_order_parallel( + 40, + 4, + err, + || (), + failing, + |_, _| { + n += 1; + Ok(()) + }, + ); + assert!(e.unwrap_err().to_string().contains("nine")); + assert_eq!(n, 9, "the pieces before the failure are written"); + let panics = |_: &mut (), k: usize| if k == 3 { panic!("three") } else { Ok(Vec::new()) }; + let e = in_order_parallel(10, 3, err, || (), panics, |_, _| Ok(())); + assert!(e.unwrap_err().to_string().contains("couldn't be made")); + } + /// A cancel that comes while the untiled writer makes its last strip stops the write, rather /// than the file being finished as if nothing had happened. #[test] diff --git a/crates/data/src/raster/geotiff.rs b/crates/data/src/raster/geotiff.rs index 9f851259f..e8e684766 100644 --- a/crates/data/src/raster/geotiff.rs +++ b/crates/data/src/raster/geotiff.rs @@ -1658,7 +1658,7 @@ pub fn write_all(r: &Raster) -> Result> { /// Writes `data` (samples interleaved as `img` lays them out, a row after another) as `img`'s /// strips, at the end of `out`. -fn put( +fn put( img: super::cogwrite::Strips, data: &[T], le: fn(T) -> [u8; N], From a7196370ad903b213aedebc79ca2d461004dc0c1 Mon Sep 17 00:00:00 2001 From: mavrick Date: Fri, 13 Nov 2026 00:09:35 -0600 Subject: [PATCH 2/2] fix(data): stop the writer's threads however the write ends A panic while a piece was written skipped stopping the threads, so threads waiting for the writer kept the write waiting for good; a guard now stops them on any exit. A piece that panics is an error on one thread too, as on several. Co-Authored-By: Claude Opus 5.5 --- crates/data/src/raster/cogwrite.rs | 68 +++++++++++++++++++++--------- 1 file changed, 48 insertions(+), 20 deletions(-) diff --git a/crates/data/src/raster/cogwrite.rs b/crates/data/src/raster/cogwrite.rs index 02f06f4ec..26b04dc53 100644 --- a/crates/data/src/raster/cogwrite.rs +++ b/crates/data/src/raster/cogwrite.rs @@ -11,10 +11,10 @@ //! Memory: the raster is read in place. A tile is gathered from the bands and compressed on its //! own, and only a few tiles a thread are held before they are written, so no second copy of the //! pixels is made, and each overview tile is made straight from the full image (a block of -//! 2ⁿ × 2ⁿ pixels), so no pyramid is kept. The price is -//! time: every overview level reads the whole full image again, so writing takes about as many -//! passes over the pixels as there are levels, where building each level from the one before -//! would take about 1⅓ but hold a second copy. +//! 2ⁿ × 2ⁿ pixels), so no pyramid is kept. The price is time: every overview level reads the +//! whole full image again, so writing takes about as many passes over the pixels as there are +//! levels, where building each level from the one before would take about 1⅓ but hold a second +//! copy. //! //! Tiles are compressed with Deflate at its fastest level, after a predictor that makes neighbouring //! samples alike: the floating-point one for floats, differences along the row for whole numbers. @@ -890,7 +890,7 @@ fn in_order( if threads > 1 { return in_order_parallel(count, threads, err, scratch, make, put); } - one_by_one(count, scratch, make, put)?; + one_by_one(count, err, scratch, make, put)?; // A cancel that came while the last piece was made stops the write too. crate::progress::check() } @@ -898,6 +898,7 @@ fn in_order( /// [`in_order`] on this thread alone: each piece made, then written, before the next. fn one_by_one( count: usize, + err: fn(String) -> DataError, scratch: impl Fn() -> S, make: impl Fn(&mut S, usize) -> Result>, mut put: impl FnMut(usize, Vec) -> Result<()>, @@ -905,11 +906,25 @@ fn one_by_one( let mut s = scratch(); for k in 0..count { crate::progress::step(1)?; - put(k, make(&mut s, k)?)?; + put(k, caught(&make, &mut s, k, err).0?)?; } Ok(()) } +/// Piece `k` from `make`, a panic in it as an error, so the write fails the same way on any number +/// of threads; and whether it panicked, after which `s` may be half changed. +fn caught( + make: impl Fn(&mut S, usize) -> Result>, + s: &mut S, + k: usize, + err: fn(String) -> DataError, +) -> (Result>, bool) { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| make(s, k))) { + Ok(made) => (made, false), + Err(_) => (Err(err("a piece of the image couldn't be made".into())), true), + } +} + /// The pieces a thread of [`in_order`] may make ahead of the one being written: enough that the /// threads rarely wait for a slow piece, few enough that the pieces waiting stay small. const AHEAD: usize = 2; @@ -951,22 +966,27 @@ fn in_order_parallel( t.next += 1; Some(t.next - 1) }; - let stop = || { - if let Ok(mut t) = turn.lock() { - t.stop = true; + /// Stops the threads when dropped, however the writing ends: done, an error, or a panic in + /// `put`, which would otherwise leave them waiting for the writer for good, and the scope + /// waiting for them. + struct Stop<'a>(&'a Mutex, &'a Condvar); + impl Drop for Stop<'_> { + fn drop(&mut self) { + if let Ok(mut t) = self.0.lock() { + t.stop = true; + } + self.1.notify_all(); } - moved.notify_all(); - }; + } std::thread::scope(|scope| { let (tx, rx) = mpsc::channel::<(usize, Result>)>(); let work = |tx: mpsc::Sender<(usize, Result>)>| { let mut s = scratch(); while let Some(k) = take() { - let made = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| make(&mut s, k))); - let panicked = made.is_err(); + let (made, panicked) = caught(&make, &mut s, k, err); // A thread whose piece panicked stops: its scratch may be half changed. - if tx.send((k, made.unwrap_or_else(|_| Err(failed())))).is_err() || panicked { + if tx.send((k, made)).is_err() || panicked { break; } } @@ -978,9 +998,11 @@ fn in_order_parallel( }) .count(); drop(tx); - let res = (|| -> Result<()> { + // The threads finish the pieces they are making and take no more; the scope waits for them. + let _stop = Stop(&turn, &moved); + (|| -> Result<()> { if started == 0 { - return one_by_one(count, &scratch, &make, &mut put); + return one_by_one(count, err, &scratch, &make, &mut put); } // Pieces made before the ones ahead of them, by their index. let mut waiting = BTreeMap::new(); @@ -1001,10 +1023,7 @@ fn in_order_parallel( moved.notify_all(); } Ok(()) - })(); - // The threads finish the pieces they are making and take no more; the scope waits for them. - stop(); - res + })() })?; // A cancel that came while the last piece was made stops the write too. crate::progress::check() @@ -1610,6 +1629,15 @@ mod tests { let panics = |_: &mut (), k: usize| if k == 3 { panic!("three") } else { Ok(Vec::new()) }; let e = in_order_parallel(10, 3, err, || (), panics, |_, _| Ok(())); assert!(e.unwrap_err().to_string().contains("couldn't be made")); + // A panic while a piece is written, with the threads waiting for it, is passed on rather + // than leaving the write waiting for them. + let caught_write = std::panic::catch_unwind(|| { + in_order_parallel(40, 2, err, || (), |_, _| Ok(Vec::new()), |k, _| if k == 0 { panic!("write") } else { Ok(()) }) + }); + assert!(caught_write.is_err()); + // On one thread too (the browser, a one-core machine). + let e = one_by_one(10, err, || (), panics, |_, _| Ok(())); + assert!(e.unwrap_err().to_string().contains("couldn't be made")); } /// A cancel that comes while the untiled writer makes its last strip stops the write, rather