Video decoders: budget repeated-frame and queued output pictures, drop a failed VP9 sample's outputs (#1778)

* Video decoders: budget repeated-frame and queued output pictures, drop a failed VP9 sample's outputs

* address review: size AV1 repeats by the reference format and after layer selection, budget HEVC queued frames at their own size

* address review: charge the larger of the output queue and the in-flight reservation, not both

* address review: count frames held by the old DPB at their own size in the budget check

* address review: charge queued frames outside the DPB and the in-flight set separately

crates/av1/src/decoder.rs

@@ -118,6 +118,8 @@ enum Output {
         grain: FilmGrainParams,
         apply_film_grain: bool,
         pts: i64,
+        /// What copying the picture out costs, in the reference's own format.
+        bytes: u64,
     },
 }
 
@@ -132,7 +134,7 @@ impl Output {
     fn take(self) -> Result<Picture> {
         match self {
             Output::Decoded(s) => s.wait().map(|p| Arc::try_unwrap(p).unwrap_or_else(|p| (*p).clone())),
-            Output::Existing { slot, seq, width, height, grain, apply_film_grain, pts } => {
+            Output::Existing { slot, seq, width, height, grain, apply_film_grain, pts, .. } => {
                 let rf = slot.wait()?;
                 Ok(output_picture(&seq, &rf, width, height, &grain, apply_film_grain, pts))
             }
@@ -379,8 +381,8 @@ impl Decoder {
         // Output policy (7.18.1 note): with scalability, show one frame per temporal unit, the
         // highest spatial layer present.
         let scalable = self.seq.as_ref().is_some_and(|s| s.op_idc != 0);
-        if !scalable {
-            self.out.extend(out.into_iter().map(|(o, _, _)| o));
+        let kept: Vec<Output> = if !scalable {
+            out.into_iter().map(|(o, _, _)| o).collect()
         } else {
             let mut last: Option<(usize, u32)> = None;
             let mut sel: Vec<Output> = Vec::new();
@@ -400,8 +402,27 @@ impl Decoder {
                     }
                 }
             }
-            self.out.extend(sel);
+            sel
+        };
+        // Every repeat is copied out whole, so the repeats this call keeps (after the layer
+        // selection) must not queue more than half the decoder budget; the first one over it and
+        // everything after is dropped.
+        let mut r = r;
+        let mut repeat_bytes = 0u64;
+        let mut ok = Vec::with_capacity(kept.len());
+        for o in kept {
+            if let Output::Existing { bytes, .. } = &o {
+                repeat_bytes = repeat_bytes.saturating_add(*bytes);
+                if repeat_bytes > self.memory_budget / 2 {
+                    if r.is_ok() {
+                        r = Err(Error::Invalid("repeated frame outputs exceed the decoder memory budget"));
+                    }
+                    break;
+                }
+            }
+            ok.push(o);
         }
+        self.out.extend(ok);
         r
     }
 
@@ -503,6 +524,15 @@ impl Decoder {
             let idx = fh.frame_to_show_map_idx;
             let slot = self.ref_frames[idx].clone().ok_or(Error::Invalid("show_existing_frame of an empty slot"))?;
             let info = self.refs[idx].clone();
+            // Sized in the reference's own format: a new sequence header may differ from it.
+            let shape = deckcraft_bitstream::budget::PictureShape {
+                width: u64::from(info.upscaled_width),
+                height: u64::from(info.frame_height),
+                shift_x: u32::from(info.subsampling_x),
+                shift_y: u32::from(info.subsampling_y),
+                chroma_planes: if info.mono_chrome { 0 } else { 2 },
+                bytes_per_sample: 2,
+            };
             let o = Output::Existing {
                 slot: slot.clone(),
                 seq: seq.clone(),
@@ -511,6 +541,7 @@ impl Decoder {
                 grain: fh.film_grain.clone(),
                 apply_film_grain: self.apply_film_grain,
                 pts: self.pts,
+                bytes: shape.bytes().unwrap_or(u64::MAX),
             };
             if fh.frame_type == KEY_FRAME as u8 {
                 // reference frame loading process (7.21) then refresh every slot (7.20)
@@ -629,6 +660,9 @@ impl Decoder {
                 ri.lf_deltas = fh.lf.deltas;
                 ri.seg_features = fh.seg.features;
                 ri.grain = fh.film_grain.clone();
+                ri.subsampling_x = job.seq.color.subsampling_x;
+                ri.subsampling_y = job.seq.color.subsampling_y;
+                ri.mono_chrome = job.seq.color.mono_chrome;
                 self.ref_frames[i] = Some(job.slot.clone());
             }
         }

crates/av1/src/header.rs

@@ -360,6 +360,10 @@ pub struct RefInfo {
     pub lf_deltas: LoopFilterDeltas,
     pub seg_features: SegmentationFeatures,
     pub grain: FilmGrainParams,
+    /// The colour format the frame was decoded in (a later sequence header may change it).
+    pub subsampling_x: u8,
+    pub subsampling_y: u8,
+    pub mono_chrome: bool,
 }
 
 /// A parsed uncompressed frame header with its derived variables.

crates/hevc/src/decoder.rs

@@ -404,10 +404,36 @@ impl Decoder {
             chroma_planes: if chroma == 0 { 0 } else { 2 },
             bytes_per_sample: 2,
         };
+        let frame_bytes = |f: &crate::picture::Frame| {
+            (f.width as u64)
+                .saturating_mul(f.height as u64)
+                .saturating_add((f.cwidth as u64).saturating_mul(f.cheight as u64).saturating_mul(2))
+                .saturating_mul(2)
+        };
+        // Output frames waiting to be returned stay allocated after they leave the DPB, at their
+        // own size: a resolution change in this call can leave pictures of the previous SPS.
+        let queued: u64 = self
+            .out_queue
+            .iter()
+            .filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame)) && !self.in_flight.iter().any(|f| Arc::ptr_eq(f, &o.frame)))
+            .map(|o| frame_bytes(&o.frame))
+            .fold(0, u64::saturating_add);
+        // Frames the DPB holds now count at their own size too (an SPS change flushes them into the
+        // queue before the new pictures are made): the DPB part is the larger of its new slots and that.
+        let held: u64 = self.dpb.entries.iter().map(|e| frame_bytes(&e.frame)).fold(0, u64::saturating_add);
+        // A queued frame still being decoded is one of the in-flight pictures already reserved, so the
+        // queue and the in-flight reservation overlap: charge the larger, not both.
         let pictures = u64::from(sps.max_dec_pic_buffering) + 1 + self.max_in_flight as u64;
+        let slots = shape.store_bytes(u64::from(sps.max_dec_pic_buffering));
+        let flight = shape.store_bytes(self.max_in_flight as u64);
+        let total = shape
+            .store_bytes(1)
+            .zip(slots)
+            .zip(flight)
+            .map(|((one, slots), flight)| slots.max(held).saturating_add(one).saturating_add(flight).saturating_add(queued));
         ensure!(
-            shape.fits(pictures, self.memory_budget),
-            "{}x{} pictures with {pictures} in the decoder exceed the memory budget",
+            total.is_some_and(|b| b <= self.memory_budget),
+            "{}x{} pictures with {pictures} in the decoder and {queued} bytes queued exceed the memory budget",
             sps.width,
             sps.height
         );

crates/vp9/src/decoder.rs

@@ -166,11 +166,17 @@ impl Decoder {
         if frames.len() > 1 {
             self.stats.superframes += 1;
         }
+        let queued = self.pending.len();
         for f in frames {
             if f.is_empty() {
                 continue;
             }
-            self.decode_frame(f, pts)?;
+            if let Err(e) = self.decode_frame(f, pts) {
+                // A failed sample outputs nothing: drop what its earlier subframes queued so a
+                // later drain does not expand each repeat into its own picture.
+                self.pending.truncate(queued);
+                return Err(e);
+            }
         }
         Ok(self.drain(false))
     }