diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs index 514a8fe7e..19726df2e 100644 --- a/crates/h264/src/decoder.rs +++ b/crates/h264/src/decoder.rs @@ -451,17 +451,19 @@ impl Decoder { chroma_planes: if chroma == 0 { 0 } else { 2 }, bytes_per_sample: 1, }; - // Pictures waiting for output that already left the DPB are held too (one decode call may - // carry many access units). A buffer shared with a DPB entry is counted once. - let queued = self.out_queue.iter().filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame))).count(); - let dpb = sps.max_dpb_frames().max(self.dpb.entries.len()); - let pictures = (dpb + 1 + self.max_in_flight.max(queued)) as u64; - ensure!( - shape.fits(pictures, self.memory_budget), - "{}x{} pictures with {pictures} in the decoder exceed the memory budget", - sps.width(), - sps.height() - ); + // Pictures already held count at their own size, not the incoming SPS's: after a + // resolution change, larger pictures of the old size can still be in the DPB or waiting + // for output (one decode call may carry many access units). A buffer shared with a DPB + // entry is counted once. + let pic = shape.bytes().unwrap_or(u64::MAX); + let held = |frames: &mut dyn Iterator>| frames.fold(0u64, |n, f| n.saturating_add(f.bytes())); + let queued = held(&mut self.out_queue.iter().filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame))).map(|o| &o.frame)); + let in_dpb = held(&mut self.dpb.entries.iter().map(|e| &e.frame)); + let dpb = sps.max_dpb_frames().max(self.dpb.entries.len()) as u64; + // The DPB (the larger of its new slots and what it holds now), the picture being decoded, + // and what is in flight or queued. + let bytes = pic.saturating_mul(dpb).max(in_dpb).saturating_add(pic).saturating_add(pic.saturating_mul(self.max_in_flight as u64).max(queued)); + ensure!(bytes <= self.memory_budget, "{}x{} pictures with {bytes} bytes in the decoder exceed the memory budget", sps.width(), sps.height()); Ok(()) } diff --git a/crates/h264/src/picture.rs b/crates/h264/src/picture.rs index 39add6a08..83331cb95 100644 --- a/crates/h264/src/picture.rs +++ b/crates/h264/src/picture.rs @@ -160,6 +160,14 @@ impl Frame { self.rows.len() } + /// Sample bytes this picture holds (luma and both chroma planes), as the decoder's memory + /// budget counts a picture. + pub fn bytes(&self) -> u64 { + (self.width as u64) + .saturating_mul(self.height as u64) + .saturating_add((self.cwidth as u64).saturating_mul(self.cheight as u64).saturating_mul(2)) + } + /// Publish MB row `r` (ignored if already published). pub fn publish(&self, r: usize, row: FrameRow) { let _ = self.rows[r].set(row); diff --git a/crates/h264/src/synth_tests.rs b/crates/h264/src/synth_tests.rs index 8d99685ab..199f5e69d 100644 --- a/crates/h264/src/synth_tests.rs +++ b/crates/h264/src/synth_tests.rs @@ -715,3 +715,11 @@ fn output_queue_counts_against_the_memory_budget() { let err = dec.decode(&[head, ps.concat()].concat(), 0).expect_err("output queue is over budget"); assert!(err.to_string().contains("memory budget"), "{err}"); } + +/// A held picture is charged at its own size (luma and both 4:2:0 chroma planes), the same +/// count the budget uses for the incoming SPS's pictures. +#[test] +fn a_held_frame_is_charged_at_its_own_size() { + let f = crate::picture::Frame::new(0, 0, 2, 3); + assert_eq!(f.bytes(), (2 * 3 * 256 + 2 * 2 * 3 * 64) as u64); +}