perf: bound the memory photo preparation holds, not just its count

`PREP_SLOTS` caps how many items are prepared at once (6) but says nothing about what they hold: one photo's decode buffer can be up to `MAX_DECODE_BYTES` (512 MiB) and that guard is *per photo*, so six of them — an album of large scans, two chats at once — could peak near 3 GiB on a host sized for a fraction of it. The upload fallback is the only path that allocates like this; nothing downstream notices until the kernel does. Photo preparation now charges a process-wide memory budget (`MEMORY_UNITS` × 64 MiB = 512 MiB) for what it actually holds: the downloaded bytes plus the decode buffer the *header* predicts — the same prediction the per-photo guards apply, now shared (`decode_bytes`, `decode_budget_bytes`) so the reservation and the guard cannot drift. A photo that is already within Telegram's limits is billed only its download, so an ordinary 10-image album still runs several at a time; two photos near the per-photo cap serialize (each takes the whole budget). The request is clamped to the budget so a single huge photo runs alone instead of waiting for permits that cannot exist. Verified with a throwaway harness against a locally served 9999x9999 PNG (126 KB on the wire, ~100 MB decoded) driven through the real `prepare_upload_item`: with the budget held the preparation waits — "after 1516ms the prep is still waiting on the budget" — and finishes in 11.8s the moment it is released, so the accounting binds in the real path and not just in the semaphore. Kept as permanent tests instead: the unit math (rounding, clamp, and that a max-size photo still gets the whole budget rather than waiting forever), the budget sharing (huge decodes cannot overlap, ordinary ones do not queue), and the prediction agreeing with the processing decision (over-sized PNG/JPEG charged, within-limits and unknown formats free). `cargo fmt --check`, `cargo clippy --workspace --all-targets --locked -- -D warnings` and `cargo test --workspace --locked` clean (122 bot + 91 x-media).
2026-09-23 23:32:05 +00:00 · 2026-09-21 14:31:50 +08:00
parent 667f523c8b
commit 62507d3a01
2 changed files with 198 additions and 2 deletions
@@ -12,6 +12,7 @@
 //! and 24-bit RGB have no alpha channel).

 use std::io::Write;
+use std::sync::LazyLock;

 use fast_image_resize as fir;
 use tempfile::NamedTempFile;
@@ -36,6 +37,90 @@ pub(crate) const MAX_DECODE_BYTES: u64 = 512 * 1024 * 1024;
 /// the cap the item degrades to the smaller URL
 /// (`FallbackError::MediaTooLarge`), it is never an error.
 pub(crate) const MAX_PHOTO_DOWNLOAD_BYTES: u64 = 32 * 1024 * 1024;
+
+/// Size of one memory-budget unit. Small enough that ordinary photos do not
+/// queue behind each other, coarse enough that the semaphore is not a counter
+/// per megabyte.
+const MEMORY_UNIT_BYTES: u64 = 64 * 1024 * 1024;
+
+/// Process-wide memory budget for photo preparation, in [`MEMORY_UNIT_BYTES`]
+/// units: 512 MiB. `PREP_SLOTS` bounds how many items are prepared at once but
+/// not how much memory they hold — one photo's decode buffer can be up to
+/// [`MAX_DECODE_BYTES`] (512 MiB), and the guard that refuses a bigger one is
+/// per photo, so six concurrent photos could peak near 3 GiB on a host sized
+/// for a fraction of that. Each item charges what it actually holds (its
+/// downloaded bytes plus the decode buffer its header predicts), so a 10-image
+/// album of ordinary photos still runs several at a time while huge ones
+/// serialize.
+const MEMORY_UNITS: u32 = 8;
+
+static MEMORY_BUDGET: LazyLock<std::sync::Arc<tokio::sync::Semaphore>> =
+    LazyLock::new(|| std::sync::Arc::new(tokio::sync::Semaphore::new(MEMORY_UNITS as usize)));
+
+/// The buffer `w`×`h` needs in `channels` output channels — the one number the
+/// per-photo guards and the reservation below both use, so they cannot drift.
+fn decode_bytes(w: u32, h: u32, channels: usize) -> u64 {
+    (w as u64) * (h as u64) * channels as u64
+}
+
+/// Units to charge for `bytes`, clamped to the whole budget: an item must never
+/// ask for more than exists, or it would wait for itself forever.
+fn memory_units(bytes: u64) -> u32 {
+    bytes
+        .div_ceil(MEMORY_UNIT_BYTES)
+        .clamp(1, MEMORY_UNITS as u64) as u32
+}
+
+/// Reserves `bytes` of the preparation budget until the returned permit drops.
+pub(crate) async fn reserve_memory(bytes: u64) -> tokio::sync::OwnedSemaphorePermit {
+    reserve(std::sync::Arc::clone(&MEMORY_BUDGET), bytes).await
+}
+
+/// [`reserve_memory`] against a caller-chosen budget; the tests pass their own
+/// so they do not fight over the process-wide one.
+async fn reserve(
+    budget: std::sync::Arc<tokio::sync::Semaphore>,
+    bytes: u64,
+) -> tokio::sync::OwnedSemaphorePermit {
+    budget
+        .acquire_many_owned(memory_units(bytes))
+        .await
+        .expect("memory budget semaphore closed")
+}
+
+/// The decode buffer a downloaded photo will allocate, from its header alone —
+/// zero when it is already within Telegram's limits and is uploaded as-is, zero
+/// for a format [`prepare_photo`] does not decode. Mirrors the early return and
+/// the guard of the two branches below.
+pub(crate) fn decode_budget_bytes(bytes: &[u8]) -> u64 {
+    if let Some((w, h, _depth, color)) = parse_png_header(bytes) {
+        if within_limits(w, h, bytes) {
+            return 0;
+        }
+        return decode_bytes(w, h, output_channels(color));
+    }
+    if let Some((w, h)) = jpeg_dims(bytes) {
+        if within_limits(w, h, bytes) {
+            return 0;
+        }
+        return decode_bytes(w, h, 3);
+    }
+    0
+}
+
+/// Whether a photo is uploaded untouched (Telegram's dimension sum, and the
+/// upload cap its bytes are compared against).
+fn within_limits(w: u32, h: u32, bytes: &[u8]) -> bool {
+    w + h <= PHOTO_MAX_DIMENSION_SUM && bytes.len() as u64 <= MAX_UPLOAD_BYTES
+}
+
+/// JPEG dimensions from the headers, without decoding any pixels.
+fn jpeg_dims(bytes: &[u8]) -> Option<(u32, u32)> {
+    let mut decoder = zune_jpeg::JpegDecoder::new(std::io::Cursor::new(bytes));
+    decoder.decode_headers().ok()?;
+    let info = decoder.info()?;
+    Some((info.width as u32, info.height as u32))
+}
 /// JPEG output quality (1-100).
 const JPEG_QUALITY: u8 = 90;

@@ -239,7 +324,7 @@ fn prepare_png(file: NamedTempFile, bytes: &[u8]) -> Result<PhotoPrep, String> {
    );

    let channels = output_channels(color_type);
-    if (w as u64) * (h as u64) * channels as u64 > MAX_DECODE_BYTES {
+    if decode_bytes(w, h, channels) > MAX_DECODE_BYTES {
        log::warn!("photo decode buffer exceeds the memory budget; falling back to smaller media");
        return Ok(PhotoPrep::UseFallback);
    }
@@ -310,7 +395,7 @@ fn prepare_jpeg(file: NamedTempFile, bytes: &[u8]) -> Result<PhotoPrep, String>
    if w + h <= PHOTO_MAX_DIMENSION_SUM && !size_over {
        return Ok(PhotoPrep::Upload(file));
    }
-    if (w as u64) * (h as u64) * 3 > MAX_DECODE_BYTES {
+    if decode_bytes(w, h, 3) > MAX_DECODE_BYTES {
        log::warn!("photo decode buffer exceeds the memory budget; falling back to smaller media");
        return Ok(PhotoPrep::UseFallback);
    }
@@ -334,6 +419,7 @@ fn prepare_jpeg(file: NamedTempFile, bytes: &[u8]) -> Result<PhotoPrep, String>
 #[cfg(test)]
 mod tests {
    use super::*;
+    use std::time::Duration;

    fn png_header(w: u32, h: u32, depth: u8, color: u8) -> Vec<u8> {
        let mut bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR".to_vec();
@@ -343,6 +429,100 @@ mod tests {
        bytes
    }

+    /// The budget is a *process-wide* memory bound: `PREP_SLOTS` (6) caps how
+    /// many photos are prepared at once, but six max-size photos would still
+    /// hold six decode buffers of up to 512 MiB each.
+    #[tokio::test]
+    async fn huge_decodes_cannot_overlap_but_do_run_alone() {
+        let budget = std::sync::Arc::new(tokio::sync::Semaphore::new(MEMORY_UNITS as usize));
+        let max_photo = MAX_DECODE_BYTES + MAX_PHOTO_DOWNLOAD_BYTES;
+
+        // One max-size photo fits (clamped to the whole budget), so it can
+        // never wait for budget that cannot exist.
+        let first = tokio::time::timeout(
+            Duration::from_millis(50),
+            reserve(budget.clone(), max_photo),
+        )
+        .await
+        .expect("a max-size photo must not wait");
+        // A second one of the same size has to wait for the first to finish.
+        assert!(
+            tokio::time::timeout(
+                Duration::from_millis(50),
+                reserve(budget.clone(), max_photo)
+            )
+            .await
+            .is_err(),
+            "two max-size decodes overlapped"
+        );
+        drop(first);
+        assert!(
+            tokio::time::timeout(
+                Duration::from_millis(50),
+                reserve(budget.clone(), max_photo)
+            )
+            .await
+            .is_ok(),
+            "the budget was not released"
+        );
+    }
+
+    /// A 10-image album of ordinary photos must not serialize: they charge
+    /// their real (small) buffers, not a fixed heavyweight slot.
+    #[tokio::test]
+    async fn ordinary_photos_share_the_budget() {
+        let budget = std::sync::Arc::new(tokio::sync::Semaphore::new(MEMORY_UNITS as usize));
+        // A 4 MiB photo that decodes to ~36 MiB (4000x3000 RGB).
+        let ordinary = 4 * 1024 * 1024 + 36 * 1024 * 1024;
+        let mut held = Vec::new();
+        for i in 0..MEMORY_UNITS {
+            held.push(
+                tokio::time::timeout(Duration::from_millis(50), reserve(budget.clone(), ordinary))
+                    .await
+                    .unwrap_or_else(|_| panic!("ordinary photo {i} waited for budget")),
+            );
+        }
+    }
+
+    #[test]
+    fn memory_units_round_up_and_clamp() {
+        assert_eq!(memory_units(1), 1);
+        assert_eq!(memory_units(MEMORY_UNIT_BYTES), 1);
+        assert_eq!(memory_units(MEMORY_UNIT_BYTES + 1), 2);
+        // Never more than exists, or the item waits for itself forever.
+        assert_eq!(memory_units(u64::MAX), MEMORY_UNITS);
+        // One item's worst case (a max download plus a max decode) takes the
+        // whole budget by itself.
+        assert_eq!(
+            memory_units(MAX_DECODE_BYTES + MAX_PHOTO_DOWNLOAD_BYTES),
+            MEMORY_UNITS
+        );
+    }
+
+    /// What the reservation is charged is decided by the header, and it has to
+    /// agree with what the pipeline does: a photo uploaded as-is costs nothing,
+    /// one that gets processed costs its decoded buffer.
+    #[test]
+    fn decode_budget_follows_the_processing_decision() {
+        // 9999x2 (sum 10001) is over the dimension cap → processed → charged.
+        let oversized = png_header(9999, 2, 8, 2); // 8-bit RGB
+        assert_eq!(decode_budget_bytes(&oversized), 9999 * 2 * 3);
+        // Inside the limits (dimensions *and* bytes) → uploaded as-is.
+        let small = png_header(100, 100, 8, 2);
+        assert_eq!(decode_budget_bytes(&small), 0);
+        // A format the pipeline does not decode costs nothing either.
+        assert_eq!(decode_budget_bytes(b"GIF89a not a photo"), 0);
+
+        // JPEG: 9999x2 is over the cap, so its RGB decode buffer is charged.
+        let (w, h) = (9999u16, 2u16);
+        let rgb = vec![90u8; w as usize * h as usize * 3];
+        let mut bytes = Vec::new();
+        jpeg_encoder::Encoder::new(&mut bytes, 90)
+            .encode(&rgb, w, h, jpeg_encoder::ColorType::Rgb)
+            .unwrap();
+        assert_eq!(decode_budget_bytes(&bytes), 9999 * 2 * 3);
+    }
+
    #[test]
    fn parses_png_header() {
        let bytes = png_header(8979, 5316, 16, 6); // 16-bit RGBA
@@ -227,6 +227,22 @@ pub(super) async fn prepare_upload_item(
                // before uploading; photos that cannot be brought within the
                // limits degrade to the smaller URL. CPU-heavy work runs off
                // the async executor thread.
+                //
+                // The header decides what that will cost in memory, so the
+                // probe travels with the downloaded bytes (both stay alive
+                // through the decode) and the reservation covers their sum:
+                // `PREP_SLOTS` bounds how many photos are prepared at once,
+                // this bounds what they hold between them — 512 MiB, whatever
+                // the batch looks like.
+                let (bytes, decode) = tokio::task::spawn_blocking(move || {
+                    let decode = photo::decode_budget_bytes(&bytes);
+                    (bytes, decode)
+                })
+                .await
+                .map_err(|e| FallbackError::Permanent {
+                    message: format!("photo worker panicked: {e}"),
+                })?;
+                let _budget = photo::reserve_memory(bytes.len() as u64 + decode).await;
                let prep = tokio::task::spawn_blocking(move || photo::prepare_photo(file, &bytes))
                    .await
                    .map_err(|e| FallbackError::Permanent {