rivet/encoder_worker/
invariant.rs

1//! Per-rung codec invariant: types + the validate/set helper.
2
3use anyhow::{Result, anyhow};
4use std::sync::RwLock;
5use codec::frame::VideoCodec;
6use codec::pixel_format::{
7    Av1SequenceHeader, H264SpsInfo, HevcSpsInfo, parse_av1_sequence_header, parse_h264_sps,
8    parse_hevc_sps,
9};
10
11/// Mandatory AV1 sequence-header fields that every encoder
12/// contributing segments to a single rendition MUST agree on.
13///
14/// Why these specific fields: each is part of the codec-init contract
15/// that the player sets up once from `av1C` and expects to hold for
16/// every segment. The decoder re-parses the inline OBU sequence
17/// header in each segment's IDR; if its parsed values disagree with
18/// the av1C from `init.mp4` on any of these fields, strict decoders
19/// (dav1d in conformance mode, Safari AVFoundation, hls.js+libdav1d)
20/// will reject the segment. Optional fields not listed here (timing
21/// info presence, decoder model presence, film grain `present` flag,
22/// operating-point details) are tolerated by every major player; we
23/// deliberately don't check them so that NVENC + QSV + AMF + rav1e
24/// can co-exist on one rendition without cosmetic byte differences
25/// triggering false rejections.
26///
27/// First worker on a rung SETS the invariant. Subsequent workers
28/// (helpers from any vendor) COMPARE; mismatch fails the run loudly
29/// instead of silently corrupting output.
30/// Per-rung codec invariant. Each chunk encoded on a different GPU must agree on
31/// these decode-init fields, or strict players reject the stitched stream. AV1
32/// compares sequence-header fields; H.264/H.265 compare the SPS profile / level
33/// / chroma / bit-depth / dims (the `avcC`/`hvcC` decode-init contract).
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub enum RungCodecInvariant {
36    Av1(Av1Invariant),
37    /// Shared by H.264 + H.265 — a rung is single-codec, so the variant only
38    /// ever compares chunks of the same codec.
39    H26x(H26xInvariant),
40}
41
42impl RungCodecInvariant {
43    /// Human-readable diff for error messages. Empty when the two agree.
44    pub(super) fn describe_diff(&self, other: &Self) -> String {
45        if self == other {
46            return String::new();
47        }
48        match (self, other) {
49            (RungCodecInvariant::Av1(a), RungCodecInvariant::Av1(b)) => a.describe_diff(b),
50            _ => format!("rung={self:?}, this worker={other:?}"),
51        }
52    }
53}
54
55/// H.264 / H.265 decode-init invariant, derived from the SPS.
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub struct H26xInvariant {
58    pub profile_idc: u8,
59    pub level_idc: u8,
60    pub chroma_format_idc: u8,
61    pub bit_depth_luma: u8,
62    pub bit_depth_chroma: u8,
63    pub width: u32,
64    pub height: u32,
65}
66
67impl H26xInvariant {
68    fn from_h264(sps: &H264SpsInfo) -> Self {
69        Self {
70            profile_idc: sps.profile_idc,
71            level_idc: sps.level_idc,
72            chroma_format_idc: sps.chroma_format_idc,
73            bit_depth_luma: sps.bit_depth_luma,
74            bit_depth_chroma: sps.bit_depth_chroma,
75            width: sps.width.unwrap_or(0),
76            height: sps.height.unwrap_or(0),
77        }
78    }
79
80    fn from_h265(sps: &HevcSpsInfo) -> Self {
81        Self {
82            profile_idc: sps.profile_idc,
83            level_idc: sps.level_idc,
84            chroma_format_idc: sps.chroma_format_idc,
85            bit_depth_luma: sps.bit_depth_luma,
86            bit_depth_chroma: sps.bit_depth_chroma,
87            width: sps.width.unwrap_or(0),
88            height: sps.height.unwrap_or(0),
89        }
90    }
91}
92
93/// AV1 sequence-header invariant — every encoder contributing segments to a
94/// single rendition MUST agree on these fields.
95///
96/// Why these specific fields: each is part of the codec-init contract that the
97/// player sets up once from `av1C` and expects to hold for every segment. The
98/// decoder re-parses the inline OBU sequence header in each segment's IDR; if
99/// its parsed values disagree with the av1C from `init.mp4`, strict decoders
100/// (dav1d in conformance mode, Safari AVFoundation, hls.js+libdav1d) reject the
101/// segment. Optional fields (timing info, decoder model, film grain present,
102/// operating points) are tolerated by every major player; we deliberately don't
103/// check them so NVENC + QSV + AMF + rav1e co-exist on one rendition without
104/// cosmetic byte differences triggering false rejections.
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct Av1Invariant {
107    pub seq_profile: u8,
108    pub seq_level_idx_0: u8,
109    pub seq_tier_0: u8,
110    pub bit_depth: u8,
111    pub monochrome: bool,
112    pub chroma_subsampling_x: bool,
113    pub chroma_subsampling_y: bool,
114    pub color_primaries: u8,
115    pub transfer_characteristics: u8,
116    pub matrix_coefficients: u8,
117    pub color_range: bool,
118    pub max_frame_width_minus1: u32,
119    pub max_frame_height_minus1: u32,
120    pub still_picture: bool,
121}
122
123impl Av1Invariant {
124    pub fn from_sequence_header(sh: &Av1SequenceHeader) -> Self {
125        Self {
126            seq_profile: sh.seq_profile,
127            seq_level_idx_0: sh.seq_level_idx_0,
128            seq_tier_0: sh.seq_tier_0,
129            bit_depth: sh.bit_depth,
130            monochrome: sh.monochrome,
131            chroma_subsampling_x: sh.chroma_subsampling_x,
132            chroma_subsampling_y: sh.chroma_subsampling_y,
133            color_primaries: sh.color_primaries,
134            transfer_characteristics: sh.transfer_characteristics,
135            matrix_coefficients: sh.matrix_coefficients,
136            color_range: sh.color_range,
137            max_frame_width_minus1: sh.max_frame_width_minus1,
138            max_frame_height_minus1: sh.max_frame_height_minus1,
139            still_picture: sh.still_picture,
140        }
141    }
142
143    /// Human-readable diff for error messages.
144    fn describe_diff(&self, other: &Self) -> String {
145        let mut diffs = Vec::new();
146        macro_rules! diff_field {
147            ($field:ident) => {
148                if self.$field != other.$field {
149                    diffs.push(format!(
150                        "{}: rung={:?}, this worker={:?}",
151                        stringify!($field),
152                        self.$field,
153                        other.$field
154                    ));
155                }
156            };
157        }
158        diff_field!(seq_profile);
159        diff_field!(seq_level_idx_0);
160        diff_field!(seq_tier_0);
161        diff_field!(bit_depth);
162        diff_field!(monochrome);
163        diff_field!(chroma_subsampling_x);
164        diff_field!(chroma_subsampling_y);
165        diff_field!(color_primaries);
166        diff_field!(transfer_characteristics);
167        diff_field!(matrix_coefficients);
168        diff_field!(color_range);
169        diff_field!(max_frame_width_minus1);
170        diff_field!(max_frame_height_minus1);
171        diff_field!(still_picture);
172        diffs.join("; ")
173    }
174}
175
176/// Outcome of comparing a worker's first packet against the rung's
177/// codec invariant. The caller — `run_encoder_worker_blocking` —
178/// branches on this to decide whether to keep encoding, soft-fail
179/// (requeue the chunk for another worker), or hard-fail (parse error
180/// from a malformed bitstream).
181#[derive(Debug)]
182pub enum InvariantCheck {
183    /// First worker on the rung. Invariant has been recorded.
184    SetByThisWorker,
185    /// Matches the rung's invariant. Proceed to publish.
186    Matched,
187    /// Mandatory fields mismatch. Worker should requeue its chunk and
188    /// exit cleanly; the rung continues with workers whose vendors
189    /// agree with the invariant the first worker set. **Mission-
190    /// critical jobs DO NOT abort on this** — only this one helper's
191    /// contribution is lost, and another worker picks up the chunk.
192    Mismatched { diff: String },
193}
194
195/// Parse a worker's first packet, derive the codec invariant, and
196/// compare-or-set it against the per-rung slot. Returns
197/// [`InvariantCheck`] on a successful parse; an `Err` only on
198/// malformed bitstream (the encoder failed to emit an
199/// `OBU_SEQUENCE_HEADER` at all, which is a configuration bug that
200/// nothing downstream can recover from).
201pub fn validate_or_set_rung_invariant(
202    rung_idx: usize,
203    gpu_vendor: Option<codec::gpu::GpuVendor>,
204    slot: &RwLock<Option<RungCodecInvariant>>,
205    first_packet: &[u8],
206    codec: VideoCodec,
207) -> Result<InvariantCheck> {
208    // Derive the codec invariant from the worker's first encoded packet: AV1
209    // from the OBU sequence header, H.264/H.265 from the SPS in the Annex-B AU.
210    let observed = match codec {
211        VideoCodec::Av1 => {
212            let parsed = parse_av1_sequence_header(first_packet).ok_or_else(|| {
213                anyhow!(
214                    "rung {} (vendor {:?}): could not parse AV1 sequence header from first \
215                     encoded packet; encoder did not emit OBU_SEQUENCE_HEADER as required for \
216                     segment alignment",
217                    rung_idx,
218                    gpu_vendor,
219                )
220            })?;
221            RungCodecInvariant::Av1(Av1Invariant::from_sequence_header(&parsed))
222        }
223        VideoCodec::H264 => {
224            let sps = parse_h264_sps(first_packet).ok_or_else(|| {
225                anyhow!(
226                    "rung {} (vendor {:?}): could not parse H.264 SPS from first encoded packet; \
227                     encoder did not emit an SPS NAL on the first IDR",
228                    rung_idx,
229                    gpu_vendor,
230                )
231            })?;
232            RungCodecInvariant::H26x(H26xInvariant::from_h264(&sps))
233        }
234        VideoCodec::H265 => {
235            let sps = parse_hevc_sps(first_packet).ok_or_else(|| {
236                anyhow!(
237                    "rung {} (vendor {:?}): could not parse H.265 SPS from first encoded packet; \
238                     encoder did not emit an SPS NAL on the first IRAP",
239                    rung_idx,
240                    gpu_vendor,
241                )
242            })?;
243            RungCodecInvariant::H26x(H26xInvariant::from_h265(&sps))
244        }
245    };
246
247    // Fast path: read lock, check if set + matches.
248    if let Some(existing) = &*slot.read().unwrap() {
249        if existing == &observed {
250            return Ok(InvariantCheck::Matched);
251        }
252        return Ok(InvariantCheck::Mismatched {
253            diff: existing.describe_diff(&observed),
254        });
255    }
256    // First worker — write under write-lock with double-check (race
257    // against another worker setting the slot between read and write).
258    let mut w = slot.write().unwrap();
259    match &*w {
260        Some(existing) if existing != &observed => Ok(InvariantCheck::Mismatched {
261            diff: existing.describe_diff(&observed),
262        }),
263        Some(_) => Ok(InvariantCheck::Matched),
264        None => {
265            tracing::info!(
266                rung_idx,
267                gpu_vendor = ?gpu_vendor,
268                ?codec,
269                invariant = ?observed,
270                "rung codec invariant captured from first worker"
271            );
272            *w = Some(observed);
273            Ok(InvariantCheck::SetByThisWorker)
274        }
275    }
276}
rivet/encoder_worker/invariant.rs

rivet/encoder_worker/
invariant.rs
