wl-webrtc/src/bin/common/mod.rs

use std::time::Instant;

use anyhow::{bail, Result};
use ffmpeg_next as ff;
use ffmpeg_next::ffi;
use ffmpeg_next::packet::Mut;
use wl_webrtc::cap_portal::{CapPortal, PwCtrlEvent};

pub fn receive_first_frame(cap: &CapPortal) -> Result<wl_webrtc::cap_portal::PwDmaBufFrame> {
    // Drain-and-wait loop that mirrors production's repeated-poll semantics
    // (state_portal.rs::poll_and_encode driven by main.rs's outer loop), but with
    // a single bounded 10s total deadline appropriate for a bench tool. Unlike a
    // single 10s blocking wait, this loop actually iterates: each turn drains ALL
    // pending control events (the ctrl channel is bounded to 8 — a single
    // if-let would silently miss backlog) and then waits a short slice for a
    // frame, so StreamEnded/Error arriving mid-wait are observed within ~200ms.
    const TOTAL_DEADLINE: std::time::Duration = std::time::Duration::from_secs(10);
    const WAIT_SLICE: std::time::Duration = std::time::Duration::from_millis(200);
    let deadline = Instant::now() + TOTAL_DEADLINE;
    loop {
        while let Ok(ctrl) = cap.event_receiver().try_recv() {
            match ctrl {
                PwCtrlEvent::StreamEnded => bail!("PipeWire stream ended before first frame"),
                PwCtrlEvent::FormatChanged { .. } => {}
                PwCtrlEvent::Error(e) => bail!("PipeWire error: {e}"),
            }
        }
        let remaining = match deadline.checked_duration_since(Instant::now()) {
            Some(r) if !r.is_zero() => r,
            _ => bail!("Timeout waiting for first frame (10s)"),
        };
        let slice = remaining.min(WAIT_SLICE);
        match cap.frame_receiver().recv_timeout(slice) {
            Ok(frame) => return Ok(frame),
            Err(crossbeam_channel::RecvTimeoutError::Timeout) => continue,
            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
                bail!("PipeWire frame channel disconnected");
            }
        }
    }
}

pub fn drain_encoder(
    enc_video: &mut ff::encoder::video::Video,
    octx: &mut ff::format::context::Output,
) -> Result<()> {
    loop {
        let mut pkt = ff::Packet::empty();
        // SAFETY: enc_video is the opened encoder; pkt is an empty Packet whose
        // inner AVPacket pointer is valid. avcodec_receive_packet fills pkt with
        // the next encoded packet, or returns EAGAIN/EOF when drained.
        let ret = unsafe { ffi::avcodec_receive_packet(enc_video.as_mut_ptr(), pkt.as_mut_ptr()) };
        if ret < 0 {
            if ret == ffi::AVERROR(ffi::EAGAIN) || ret == ffi::AVERROR_EOF {
                break;
            }
            eprintln!("avcodec_receive_packet failed: {ret}");
            break;
        }

        let enc_tb = enc_video.time_base();
        // SAFETY: octx.as_ptr() is a valid AVFormatContext; streams is a NULL-terminated
        // array of AVStream*. We index [0] which exists because we created exactly one
        // stream in setup. Reading time_base is a plain AVRational field access.
        let stream_tb = unsafe {
            let streams = (*octx.as_ptr()).streams;
            let st = *streams.add(0);
            ff::Rational::from((*st).time_base)
        };
        pkt.rescale_ts(enc_tb, stream_tb);
        pkt.set_stream(0);
        pkt.write_interleaved(octx)
            .map_err(|e| anyhow::anyhow!("write packet failed: {e}"))?;
    }
    Ok(())
}