rivet/encoder_worker/
chunk_worker.rs

1//! Single-file chunked encode: workers collect packets (instead of writing CMAF
2//! segments) so the orchestrator can stitch them, in segment order, into one MP4.
3
4use anyhow::{Context, Result};
5use std::sync::Arc;
6use tokio::sync::mpsc;
7use codec::encode::{self, EncoderConfig};
8use crate::frame_queue::{SegmentChunk, SegmentChunkQueue};
9use super::{EncoderWorkerConfig, InvariantCheck, validate_or_set_rung_invariant};
10
11/// One chunk's encoded packets, in encode (= display, no B-frames) order.
12#[derive(Debug)]
13pub struct ChunkPackets {
14    pub segment_idx: usize,
15    pub packets: Vec<encode::EncodedPacket>,
16}
17
18/// Encoder worker that COLLECTS packets per chunk (single-file path). Each
19/// chunk is encoded by a fresh encoder (first frame an IDR); the cross-vendor
20/// codec invariant is enforced on the first packet (mismatch → requeue + exit,
21/// exactly like the CMAF worker). Ordered `ChunkPackets` are pushed to `out`.
22#[allow(clippy::too_many_arguments)]
23pub fn run_chunk_encoder_worker_blocking(
24    cfg: EncoderWorkerConfig,
25    queue: Arc<SegmentChunkQueue>,
26    rt: tokio::runtime::Handle,
27    shared_frames_encoded: Arc<std::sync::atomic::AtomicU64>,
28    progress_tx: mpsc::Sender<u64>,
29    out: Arc<std::sync::Mutex<Vec<ChunkPackets>>>,
30) -> Result<()> {
31    let enc_config = super::build_enc_config(&cfg);
32    loop {
33        let chunk = match rt.block_on(queue.pop()) {
34            Some(c) => c,
35            None => break,
36        };
37        match encode_chunk_to_packets(&cfg, &enc_config, chunk, &shared_frames_encoded, &progress_tx)?
38        {
39            ChunkOutcome::Encoded(c) => out.lock().unwrap().push(c),
40            ChunkOutcome::RequeuedOnMismatch { chunk, diff } => {
41                tracing::warn!(
42                    rung_idx = cfg.rung_idx,
43                    gpu_vendor = ?cfg.gpu_vendor,
44                    diff = %diff,
45                    "chunk worker: codec invariant mismatch — requeuing chunk and exiting"
46                );
47                let _ = queue.push_front(chunk);
48                break;
49            }
50        }
51    }
52    Ok(())
53}
54
55enum ChunkOutcome {
56    Encoded(ChunkPackets),
57    RequeuedOnMismatch { chunk: SegmentChunk, diff: String },
58}
59
60fn encode_chunk_to_packets(
61    cfg: &EncoderWorkerConfig,
62    enc_config: &EncoderConfig,
63    chunk: SegmentChunk,
64    shared_frames_encoded: &std::sync::atomic::AtomicU64,
65    progress_tx: &mpsc::Sender<u64>,
66) -> Result<ChunkOutcome> {
67    let mut encoder =
68        encode::select_encoder(enc_config.clone(), None).context("creating encoder for chunk")?;
69    let segment_idx = chunk.segment_idx;
70    let mut packets: Vec<encode::EncodedPacket> = Vec::new();
71    let mut pending: Vec<encode::EncodedPacket> = Vec::new();
72    let mut decided = false;
73
74    for frame in &chunk.frames {
75        encoder.send_frame(frame).context("send_frame in chunk worker")?;
76        while let Some(packet) = encoder.receive_packet().context("receive_packet in chunk worker")? {
77            if !decided {
78                match validate_or_set_rung_invariant(
79                    cfg.rung_idx,
80                    cfg.gpu_vendor,
81                    &cfg.rung_invariant,
82                    &packet.data,
83                    cfg.codec,
84                )? {
85                    InvariantCheck::Matched | InvariantCheck::SetByThisWorker => decided = true,
86                    InvariantCheck::Mismatched { diff } => {
87                        return Ok(ChunkOutcome::RequeuedOnMismatch { chunk, diff });
88                    }
89                }
90                pending.push(packet);
91                continue;
92            }
93            packets.append(&mut pending);
94            packets.push(packet);
95        }
96        let n = shared_frames_encoded.fetch_add(1, std::sync::atomic::Ordering::AcqRel) + 1;
97        let _ = progress_tx.try_send(n);
98    }
99    if decided {
100        packets.append(&mut pending);
101    }
102    encoder.flush().context("flush in chunk worker")?;
103    while let Some(packet) = encoder
104        .receive_packet()
105        .context("receive_packet after flush in chunk worker")?
106    {
107        packets.push(packet);
108    }
109    Ok(ChunkOutcome::Encoded(ChunkPackets { segment_idx, packets }))
110}
rivet/encoder_worker/chunk_worker.rs

rivet/encoder_worker/
chunk_worker.rs
