From 532b313b8cfd485eb03f06b405a5393cbb83d302 Mon Sep 17 00:00:00 2001 From: enricobuehler Date: Thu, 16 Jul 2026 12:21:49 +0200 Subject: [PATCH] =?UTF-8?q?refactor(host):=20extract=20encode/codec.rs=20?= =?UTF-8?q?=E2=80=94=20the=20encoder=20contract?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move the Tier-1 encoder contract out of the stuffed encode.rs facade into a new encode/codec.rs submodule (plan §7 / W2): EncodedFrame, Codec (all methods except host_wire_caps), ChromaFormat, EncoderCaps, the Encoder trait, validate_dimensions, vbv_frames_env, and the dimension + wire-roundtrip contract tests. host_wire_caps stays in encode.rs alongside the backend-selection probes it depends on; CodecSupport and its wire-mask test stay too. encode.rs gains `mod codec;` + `pub(crate) use codec::*;` so every existing crate::encode::X path — crate::encode::vbv_frames_env, ::Codec, ::Encoder, … — stays byte-stable. Pure relocation: no call sites touched. Verified: dev-Mac type-check of both files clean; Linux `cargo check -p punktfunk-host --features nvenc,vulkan-encode,pyrowave` green (all encode backends compile against the relocated contract); contract unit tests pass. Co-Authored-By: Claude Fable 5 --- crates/punktfunk-host/src/encode.rs | 401 +-------------------- crates/punktfunk-host/src/encode/codec.rs | 415 ++++++++++++++++++++++ 2 files changed, 417 insertions(+), 399 deletions(-) create mode 100644 crates/punktfunk-host/src/encode/codec.rs diff --git a/crates/punktfunk-host/src/encode.rs b/crates/punktfunk-host/src/encode.rs index 5aa314b6..03bf241a 100644 --- a/crates/punktfunk-host/src/encode.rs +++ b/crates/punktfunk-host/src/encode.rs @@ -10,92 +10,10 @@ use crate::capture::{CapturedFrame, PixelFormat}; use anyhow::Result; -/// An encoded access unit (one NAL/AU) to hand to `punktfunk_core` for FEC + packetization. -/// `data` is in-band Annex-B (the encoder is opened without a global header), so each -/// keyframe carries its own VPS/SPS/PPS — the bytes are both a playable elementary -/// stream and a self-contained AU for the wire. -pub struct EncodedFrame { - pub data: Vec, - pub pts_ns: u64, - /// True for IDR/keyframes (sets the SOF/keyframe wire flags). - pub keyframe: bool, - /// True when this AU is a **reference-frame-invalidation recovery frame** — a clean P-frame the - /// encoder coded against a known-good reference in response to - /// [`invalidate_ref_frames`](Encoder::invalidate_ref_frames). The pump tags it - /// [`punktfunk_core::packet::USER_FLAG_RECOVERY_ANCHOR`] so the client lifts its post-loss - /// freeze on it without an IDR. Set by BOTH RFI backends: native AMF (the LTR force-reference - /// frame) and Windows direct-NVENC (the first frame encoded after `nvEncInvalidateRefFrames` — - /// the invalidation applies at the next `encode_picture`, so that AU is by construction the - /// clean re-anchor). Without it the client's freeze can only lift on an IDR — which the host - /// suppresses after a successful RFI (the cooldown), a ~1 s frozen stall per loss event. - pub recovery_anchor: bool, - /// The AU is shard-aligned self-delimiting chunks (see [`Encoder::set_wire_chunking`]); - /// the session stamps [`punktfunk_core::packet::USER_FLAG_CHUNK_ALIGNED`] so the client - /// windows its parse and may opt into partial delivery. Only the PyroWave backend sets it. - pub chunk_aligned: bool, -} - -/// Codec selection negotiated with the client. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Codec { - H264, - H265, - Av1, - /// PyroWave — the opt-in wired-LAN intra-only wavelet codec (design/pyrowave-codec-plan.md). - /// Only ever negotiated via the client's explicit `preferred_codec` (never the precedence - /// ladder) and only emitted by the `pyrowave`-feature backend; every AU is a keyframe. - PyroWave, -} - -/// Chroma subsampling the encoder emits, negotiated with the client (the `PUNKTFUNK_444` gate + the -/// client's `VIDEO_CAP_444` + a GPU probe). `Yuv420` is the universal default; `Yuv444` is HEVC-only, -/// native-protocol-only (GameStream stays 4:2:0), and the host only ever passes it after -/// [`can_encode_444`] confirmed the active backend supports it. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] -pub enum ChromaFormat { - #[default] - Yuv420, - Yuv444, -} - -impl ChromaFormat { - /// The HEVC `chroma_format_idc` this maps to: `1` (4:2:0) or `3` (4:4:4). Also the wire value - /// echoed in [`punktfunk_core::quic::Welcome::chroma_format`]. - pub fn idc(self) -> u8 { - match self { - ChromaFormat::Yuv420 => punktfunk_core::quic::CHROMA_IDC_420, - ChromaFormat::Yuv444 => punktfunk_core::quic::CHROMA_IDC_444, - } - } - - /// True for full-chroma 4:4:4. - pub fn is_444(self) -> bool { - matches!(self, ChromaFormat::Yuv444) - } -} +mod codec; +pub(crate) use codec::*; impl Codec { - /// Map a negotiated `quic` codec bit ([`punktfunk_core::quic::CODEC_H264`] etc.) to the encoder - /// [`Codec`]. Unknown / `0` → HEVC (the pre-negotiation default). Inverse of [`Codec::to_wire`]. - pub fn from_wire(bit: u8) -> Codec { - match bit { - punktfunk_core::quic::CODEC_H264 => Codec::H264, - punktfunk_core::quic::CODEC_AV1 => Codec::Av1, - punktfunk_core::quic::CODEC_PYROWAVE => Codec::PyroWave, - _ => Codec::H265, - } - } - - /// The single `quic` codec bit for this codec (echoed in [`punktfunk_core::quic::Welcome::codec`]). - pub fn to_wire(self) -> u8 { - match self { - Codec::H264 => punktfunk_core::quic::CODEC_H264, - Codec::H265 => punktfunk_core::quic::CODEC_HEVC, - Codec::Av1 => punktfunk_core::quic::CODEC_AV1, - Codec::PyroWave => punktfunk_core::quic::CODEC_PYROWAVE, - } - } - /// The `quic` codec bitfield the host can currently **emit** on the punktfunk/1 native path, /// given the resolved encode backend — the same GPU-aware advertisement GameStream builds for /// Moonlight ([`crate::gamestream::serverinfo`]), in `quic::CODEC_*` bits. The GPU-less software @@ -173,275 +91,6 @@ impl Codec { })(); base | pyro } - - /// Lowercase stats/console label (`"h264"` / `"hevc"` / `"av1"`) — the codec string seeded into - /// the web console's session meta ([`crate::stats_recorder::StatsRecorder::register_session`]). - pub fn label(self) -> &'static str { - match self { - Codec::H264 => "h264", - Codec::H265 => "hevc", - Codec::Av1 => "av1", - Codec::PyroWave => "pyrowave", - } - } - - /// The FFmpeg NVENC encoder name (selected by name, not codec id — the latter would - /// pick the software encoder). - pub fn nvenc_name(self) -> &'static str { - match self { - Codec::H264 => "h264_nvenc", - Codec::H265 => "hevc_nvenc", - Codec::Av1 => "av1_nvenc", - // Guarded by the open_video dispatch: a PyroWave session never reaches a - // libavcodec backend. - Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), - } - } - - /// The FFmpeg VAAPI encoder name (AMD via Mesa `radeonsi`, Intel via `iHD`/`i965`). One - /// libavcodec encoder per codec covers both vendors — the kernel driver differs, the libva - /// userspace API is identical. Selected by name (the codec id would pick the SW encoder). - /// AV1 VAAPI encode is narrow (Intel Arc/Xe2+, AMD RDNA3+/RDNA4) — gate it on a capability - /// probe, never assume it (see [`open_video`]). - pub fn vaapi_name(self) -> &'static str { - match self { - Codec::H264 => "h264_vaapi", - Codec::H265 => "hevc_vaapi", - Codec::Av1 => "av1_vaapi", - // Guarded by the open_video dispatch: a PyroWave session never reaches a - // libavcodec backend. - Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), - } - } - - /// The FFmpeg AMD **AMF** encoder name (the Windows AMD backend). Selected by name (the codec id - /// would pick the software encoder). AV1 (`av1_amf`) is RDNA3+/RX 7000+ — probe, never assume. - pub fn amf_name(self) -> &'static str { - match self { - Codec::H264 => "h264_amf", - Codec::H265 => "hevc_amf", - Codec::Av1 => "av1_amf", - // Guarded by the open_video dispatch: a PyroWave session never reaches a - // libavcodec backend. - Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), - } - } - - /// The FFmpeg Intel **QSV** encoder name (the Windows Intel backend). Selected by name. AV1 - /// (`av1_qsv`) is Arc/Xe2+; HEVC Main10 is Gen9.5+ — probe, never assume. - pub fn qsv_name(self) -> &'static str { - match self { - Codec::H264 => "h264_qsv", - Codec::H265 => "hevc_qsv", - Codec::Av1 => "av1_qsv", - // Guarded by the open_video dispatch: a PyroWave session never reaches a - // libavcodec backend. - Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), - } - } -} - -/// Static capabilities an [`Encoder`] declares so the session glue routes loss-recovery and HDR -/// plumbing by *query* rather than relying on a method's no-op/`false` default. Cheap `Copy`; fixed -/// for the session (an HDR toggle re-initialises the encoder — re-query if that matters). -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub struct EncoderCaps { - /// The encoder can perform real reference-frame invalidation — i.e. - /// [`invalidate_ref_frames`](Encoder::invalidate_ref_frames) can return `true`. When `false` - /// the caller skips that always-`false` call and forces a keyframe directly on loss recovery. - /// Two backends implement RFI: Windows direct-NVENC (`nvEncInvalidateRefFrames`) and native - /// AMF (user-LTR force-reference, when the driver accepted the LTR slots at open). The - /// libavcodec paths (Linux NVENC, VAAPI, QSV) can't express it and always keyframe. - pub supports_rfi: bool, - /// The encoder emits in-band HDR mastering/CLL SEI from [`set_hdr_meta`](Encoder::set_hdr_meta). - /// When `false`, `set_hdr_meta` is a no-op and no in-band grade reaches the client. Only the - /// Windows direct-NVENC path attaches it today. - pub supports_hdr_metadata: bool, - /// The opened encoder is actually producing a full-chroma 4:4:4 (`chroma_format_idc = 3`) stream. - /// `false` on every 4:2:0 session (the default) and on a backend that declined 4:4:4. Set by the - /// NVENC backends (Linux + Windows). The chroma is committed to the wire (`Welcome::chroma_format`) - /// from the pre-open probe, so this is a *post-open cross-check*: the session glue logs loudly if - /// the encoder's real chroma disagrees with what was negotiated (the in-band SPS is authoritative - /// for the decoder either way). - pub chroma_444: bool, - /// The encoder runs a periodic **intra-refresh wave** — a moving band of intra blocks that - /// re-codes the whole picture over ~0.5 s, no periodic IDR. FEC-unrecoverable loss self-heals as - /// the band sweeps, so the session glue rate-limits client keyframe requests instead of answering - /// each with a full IDR (the 20-40× frame-size spike that cascades under loss). Linux NVENC / AMF - /// set it when `PUNKTFUNK_INTRA_REFRESH` opened the encoder in that mode; VAAPI/QSV/software never - /// do. NOTE — the wave carries NO decoder-visible clean-point: FFmpeg never sets `AV_FRAME_FLAG_KEY` - /// at a recovery point (H.264 flags key only when `recovery_frame_cnt == 0`; HEVC only on IRAP), - /// and AMF emits no recovery-point SEI at all. So this cap ALONE does not let the client lift its - /// post-loss freeze without an IDR — that needs [`intra_refresh_recovery`](Self::intra_refresh_recovery). - pub intra_refresh: bool, - /// The intra-refresh wave is a *validated constrained GDR* — verified on real hardware to fully - /// heal a lost picture within one wave period with no residual artifacts. Only then does the host - /// tag each wave-boundary AU with [`USER_FLAG_RECOVERY_POINT`](punktfunk_core::packet::USER_FLAG_RECOVERY_POINT), - /// so the client can lift its freeze on the second mark (a proven clean re-anchor) instead of - /// waiting out its backstop and forcing a full IDR. Default `false` on every backend until on-glass - /// validation flips it — an un-validated encoder keeps the IDR recovery path, so this is inert and - /// cannot regress. Meaningless unless [`intra_refresh`](Self::intra_refresh) is also set. - pub intra_refresh_recovery: bool, - /// Length of the intra-refresh wave in frames — the boundary period the host marks on (it sets - /// `USER_FLAG_RECOVERY_POINT` on every Nth emitted AU, re-phased at each IDR). 0 when intra-refresh - /// is off. Only consulted when [`intra_refresh_recovery`](Self::intra_refresh_recovery) is set. - pub intra_refresh_period: u32, -} - -/// A hardware encoder. One per session; runs on the encode thread. -pub trait Encoder: Send { - fn submit(&mut self, frame: &CapturedFrame) -> Result<()>; - /// [`submit`](Self::submit) with the **wire frame index** this frame's AU will carry — the - /// number the packetizer stamps on it and the client's loss reports/RFI requests name. The - /// session glue predicts it exactly as `AUs sent so far + frames in flight` (AUs are emitted - /// FIFO, one per submission; anything that would break the prediction — an in-place reset, a - /// device-change teardown, an encoder rebuild — forfeits the in-flight frames on BOTH sides - /// and clears the encoder's reference state, so stale predictions die with it). The RFI - /// backends pin their frame numbering (LTR marks, DPB timestamps) to this so - /// [`invalidate_ref_frames`](Self::invalidate_ref_frames) compares client frame numbers - /// against the same domain — an encoder-internal counter desyncs from the wire on the first - /// mid-stream rebuild (adaptive bitrate steps do this under congestion, exactly when losses - /// happen). Default: ignore the index and delegate to `submit` (backends without per-frame - /// reference bookkeeping don't care). - fn submit_indexed(&mut self, frame: &CapturedFrame, wire_index: u32) -> Result<()> { - let _ = wire_index; - self.submit(frame) - } - /// This encoder's static [capabilities](EncoderCaps) (RFI, HDR SEI), so the session glue can - /// route by query rather than rely on the no-op/`false` defaults of - /// [`invalidate_ref_frames`](Self::invalidate_ref_frames) / [`set_hdr_meta`](Self::set_hdr_meta). - /// Default: no optional capabilities (the SDR / libavcodec backends) — only the direct-NVENC - /// path overrides it. - fn caps(&self) -> EncoderCaps { - EncoderCaps::default() - } - /// Force the next submitted frame to be an IDR keyframe (e.g. after a client - /// reference-frame-invalidation request). Default: no-op. - fn request_keyframe(&mut self) {} - /// Set the source's static HDR mastering metadata (from the capturer). An HDR encoder emits it - /// as in-band SEI (`mastering_display_colour_volume` + `content_light_level_info`) on each - /// keyframe so any decoder — including stock Moonlight — tone-maps from the source's real grade. - /// Default: no-op (SDR encoders / libavcodec paths that don't attach it yet). Cheap to call - /// every frame; only the direct-NVENC path consumes it. - fn set_hdr_meta(&mut self, _meta: Option) {} - /// Invalidate a contiguous range of previously-encoded reference frames (client frame numbers - /// — WIRE frame indexes, the domain [`submit_indexed`](Self::submit_indexed) pins the encoder's - /// bookkeeping to) so the encoder re-references an older still-valid frame instead of emitting - /// a full IDR. Returns `true` if a real reference invalidation was performed; `false` means the - /// encoder couldn't (range older than the DPB/LTR history, or the backend has no RFI) and the - /// caller should fall back to [`request_keyframe`](Self::request_keyframe). Default: `false` — - /// the Windows direct-NVENC path (`nvEncInvalidateRefFrames`) and native AMF (LTR - /// force-reference) implement true RFI; the libavcodec paths can't express it, so they keyframe. - fn invalidate_ref_frames(&mut self, _first_frame: i64, _last_frame: i64) -> bool { - false - } - /// Pull the next encoded AU if one is ready. - fn poll(&mut self) -> Result>; - /// Tear the underlying hardware encoder down and rebuild it in place, keeping the session's - /// negotiated parameters — the encode-stall watchdog's recovery lever (a wedged AMF/QSV - /// driver stops emitting AUs or accepting frames without ever returning an error). Returns - /// `true` when the encoder was rebuilt: every submitted-but-unpolled frame is forfeited and - /// the next submitted frame starts a fresh stream (IDR). Default `false`: the backend has no - /// in-place rebuild and the caller must treat the stall as fatal instead. - fn reset(&mut self) -> bool { - false - } - /// Retarget the encoder's rate control to `bps` (average == max, CBR) **in place** — same - /// codec/resolution/fps, only the bitrate and its derived VBV move. Returns `true` when the - /// live encoder accepted the change: the reference chain, the in-flight frames and the - /// caller's wire-index prediction all survive, so an adaptive-bitrate step costs *nothing* on - /// the wire (no IDR, no in-flight forfeit — the whole point vs. a rebuild). `false` = the - /// backend can't (or the driver rejected the new rate, e.g. above the codec-level ceiling) — - /// the caller falls back to its full rebuild path, which also owns the bitrate clamping. - /// Default: no in-place retarget (the libavcodec/software paths). - fn reconfigure_bitrate(&mut self, _bps: u64) -> bool { - false - } - /// Wire-chunk the encoder's AUs at the session's shard payload size (the PyroWave - /// datagram-aligned mode, plan §4.4): every `shard_payload` window of the emitted AU - /// starts a fresh self-delimiting codec packet, zero-padded to the window — so a lost - /// datagram costs a few coefficient blocks, not the frame. AUs produced this way are - /// flagged [`EncodedFrame::chunk_aligned`] and the session marks them on the wire. - /// Default: no-op (the H.26x backends' bitstreams cannot be cut losslessly). - fn set_wire_chunking(&mut self, _shard_payload: usize) {} - /// Signal end-of-stream. After this, drain the remaining AUs with [`poll`](Self::poll) - /// until it returns `None` — NVENC buffers frames internally even at `delay=0`. - fn flush(&mut self) -> Result<()>; -} - -impl Codec { - /// Maximum encodable dimension (px) per side for this codec on NVENC. H.264 tops out at - /// 4096 (level constraint); HEVC and AV1 allow 8192. Used to reject out-of-range client - /// modes up front (see [`validate_dimensions`]). - pub fn max_dimension(self) -> u32 { - match self { - Codec::H264 => 4096, - // PyroWave has no codec-level dimension cap (arbitrary even sizes); 8192 matches the - // buffer-math guard the other codecs get. - Codec::H265 | Codec::Av1 | Codec::PyroWave => 8192, - } - } - - /// The codec's *spec* top level/tier bitrate (bits/s) — the usual boundary at which NVENC - /// starts rejecting `avcodec_open2` with EINVAL. NOT a hard cap: [`open_video`](crate::encode:: - /// open_video) probes the actual GPU ceiling by stepping DOWN from the requested bitrate only on - /// EINVAL, and uses this purely as the first step-down candidate (so a card that accepts more — - /// an RTX 5070 Ti does >1 Gbps HEVC where a 4090 caps at ~800 Mbps — is never clamped to it). - /// HEVC Level 6.2 High tier = 800 Mbps; H.264 High level 6.2 ≈ 480 Mbps; AV1's levels allow more. - pub fn max_bitrate_bps(self) -> u64 { - match self { - Codec::H264 => 480_000_000, - Codec::H265 => 800_000_000, - Codec::Av1 => 1_200_000_000, - // No spec level/tier: the rate is a plain per-frame byte budget. Use the protocol's - // own bitrate clamp so the step-down probe logic never binds below it. - Codec::PyroWave => 8_000_000_000, - } - } -} - -/// `PUNKTFUNK_VBV_FRAMES` — HRD/VBV size in frame intervals (default 1.0, the strict low-latency -/// shape every backend ships: each frame must fit its rate share, keeping frame sizes uniform for -/// the pacer). The AMF/VAAPI/QSV paths parse the same variable locally; this helper brings the -/// direct-NVENC paths (which used to hardwire 1 frame) to parity. Larger values let complex -/// frames borrow bits — better rate utilization at the cost of per-frame size variance. -pub(crate) fn vbv_frames_env() -> f64 { - std::env::var("PUNKTFUNK_VBV_FRAMES") - .ok() - .and_then(|s| s.parse::().ok()) - .filter(|v| v.is_finite() && *v > 0.0) - .unwrap_or(1.0) -} - -/// Validate a requested encode resolution before we allocate buffers or open NVENC. Rejects -/// zero/odd-sized and out-of-range modes with a clear error instead of letting buffer math -/// overflow or the encoder open fail with an opaque NVENC code. A client can request any -/// `mode=WxHxFPS`, so this is the gate on attacker/typo-controlled dimensions. -pub fn validate_dimensions(codec: Codec, width: u32, height: u32) -> Result<()> { - if width == 0 || height == 0 { - anyhow::bail!("invalid encode resolution {width}x{height}: dimensions must be non-zero"); - } - // NVENC requires even dimensions for the chroma subsampling it does internally. - if width % 2 != 0 || height % 2 != 0 { - anyhow::bail!("invalid encode resolution {width}x{height}: dimensions must be even"); - } - // PyroWave's 5-level wavelet decomposition needs ≥ 4·2⁵ px per axis (upstream - // `MinimumImageSize` — the band mirroring breaks below it); reject a tiny mode here - // (e.g. a match-window resize dragged to a sliver) instead of failing the encoder - // rebuild after the switch was acked. - if codec == Codec::PyroWave && (width < 128 || height < 128) { - anyhow::bail!( - "invalid PyroWave resolution {width}x{height}: the wavelet needs at least 128px per axis" - ); - } - let max = codec.max_dimension(); - if width > max || height > max { - anyhow::bail!( - "{codec:?} max dimension is {max}px; requested {width}x{height} \ - (use HEVC/AV1 above 4096, or lower the client resolution)" - ); - } - Ok(()) } /// Open a hardware video encoder for frames of the given `format` and mode, selecting the GPU @@ -1425,41 +1074,6 @@ mod pyrowave; mod tests { use super::*; - #[test] - fn rejects_zero_and_odd_dimensions() { - assert!(validate_dimensions(Codec::H265, 0, 1080).is_err()); - assert!(validate_dimensions(Codec::H265, 1920, 0).is_err()); - assert!(validate_dimensions(Codec::H265, 1921, 1080).is_err()); // odd width - assert!(validate_dimensions(Codec::H265, 1920, 1081).is_err()); // odd height - } - - #[test] - fn h264_capped_at_4096() { - assert!(validate_dimensions(Codec::H264, 3840, 2160).is_ok()); // 4K fits (width < 4096) - assert!(validate_dimensions(Codec::H264, 4096, 4096).is_ok()); // exactly at the limit - assert!(validate_dimensions(Codec::H264, 4098, 2160).is_err()); - assert!(validate_dimensions(Codec::H264, 3840, 4098).is_err()); - } - - #[test] - fn hevc_and_av1_allow_up_to_8192() { - for c in [Codec::H265, Codec::Av1] { - assert!(validate_dimensions(c, 3840, 2160).is_ok()); - assert!(validate_dimensions(c, 7680, 4320).is_ok()); // 8K fits - assert!(validate_dimensions(c, 8192, 8192).is_ok()); - assert!(validate_dimensions(c, 8194, 4320).is_err()); - } - } - - #[test] - fn common_modes_accepted() { - for c in [Codec::H264, Codec::H265, Codec::Av1] { - for (w, h) in [(1280, 720), (1920, 1080), (2560, 1440)] { - assert!(validate_dimensions(c, w, h).is_ok(), "{c:?} {w}x{h}"); - } - } - } - /// The probed-capability → wire-bitfield mapping the native codec advertisement is built from. #[cfg(any(target_os = "linux", target_os = "windows"))] #[test] @@ -1486,15 +1100,4 @@ mod tests { }; assert_eq!(none.wire_mask(), None); } - - /// Wire round-trip and the stats label stay in lockstep with the `quic::CODEC_*` bits. - #[test] - fn codec_wire_roundtrip_and_label() { - for c in [Codec::H264, Codec::H265, Codec::Av1] { - assert_eq!(Codec::from_wire(c.to_wire()), c); - } - assert_eq!(Codec::H264.label(), "h264"); - assert_eq!(Codec::H265.label(), "hevc"); - assert_eq!(Codec::Av1.label(), "av1"); - } } diff --git a/crates/punktfunk-host/src/encode/codec.rs b/crates/punktfunk-host/src/encode/codec.rs new file mode 100644 index 00000000..776aeb75 --- /dev/null +++ b/crates/punktfunk-host/src/encode/codec.rs @@ -0,0 +1,415 @@ +//! The encoder contract (plan §7, Tier 1): the [`Encoder`] trait plus the plain-data value types its +//! signatures use — [`EncodedFrame`], [`Codec`], [`ChromaFormat`], [`EncoderCaps`] — and the +//! dimension/VBV helpers [`validate_dimensions`] and [`vbv_frames_env`]. Backend selection, the +//! capability probes that mirror it, and `Codec::host_wire_caps` stay in the parent [`crate::encode`] +//! facade, which re-exports this module (`pub(crate) use codec::*;`) so every `crate::encode::*` path +//! is unchanged. +use crate::capture::CapturedFrame; +use anyhow::Result; + +/// An encoded access unit (one NAL/AU) to hand to `punktfunk_core` for FEC + packetization. +/// `data` is in-band Annex-B (the encoder is opened without a global header), so each +/// keyframe carries its own VPS/SPS/PPS — the bytes are both a playable elementary +/// stream and a self-contained AU for the wire. +pub struct EncodedFrame { + pub data: Vec, + pub pts_ns: u64, + /// True for IDR/keyframes (sets the SOF/keyframe wire flags). + pub keyframe: bool, + /// True when this AU is a **reference-frame-invalidation recovery frame** — a clean P-frame the + /// encoder coded against a known-good reference in response to + /// [`invalidate_ref_frames`](Encoder::invalidate_ref_frames). The pump tags it + /// [`punktfunk_core::packet::USER_FLAG_RECOVERY_ANCHOR`] so the client lifts its post-loss + /// freeze on it without an IDR. Set by BOTH RFI backends: native AMF (the LTR force-reference + /// frame) and Windows direct-NVENC (the first frame encoded after `nvEncInvalidateRefFrames` — + /// the invalidation applies at the next `encode_picture`, so that AU is by construction the + /// clean re-anchor). Without it the client's freeze can only lift on an IDR — which the host + /// suppresses after a successful RFI (the cooldown), a ~1 s frozen stall per loss event. + pub recovery_anchor: bool, + /// The AU is shard-aligned self-delimiting chunks (see [`Encoder::set_wire_chunking`]); + /// the session stamps [`punktfunk_core::packet::USER_FLAG_CHUNK_ALIGNED`] so the client + /// windows its parse and may opt into partial delivery. Only the PyroWave backend sets it. + pub chunk_aligned: bool, +} + +/// Codec selection negotiated with the client. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Codec { + H264, + H265, + Av1, + /// PyroWave — the opt-in wired-LAN intra-only wavelet codec (design/pyrowave-codec-plan.md). + /// Only ever negotiated via the client's explicit `preferred_codec` (never the precedence + /// ladder) and only emitted by the `pyrowave`-feature backend; every AU is a keyframe. + PyroWave, +} + +/// Chroma subsampling the encoder emits, negotiated with the client (the `PUNKTFUNK_444` gate + the +/// client's `VIDEO_CAP_444` + a GPU probe). `Yuv420` is the universal default; `Yuv444` is HEVC-only, +/// native-protocol-only (GameStream stays 4:2:0), and the host only ever passes it after +/// [`can_encode_444`] confirmed the active backend supports it. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum ChromaFormat { + #[default] + Yuv420, + Yuv444, +} + +impl ChromaFormat { + /// The HEVC `chroma_format_idc` this maps to: `1` (4:2:0) or `3` (4:4:4). Also the wire value + /// echoed in [`punktfunk_core::quic::Welcome::chroma_format`]. + pub fn idc(self) -> u8 { + match self { + ChromaFormat::Yuv420 => punktfunk_core::quic::CHROMA_IDC_420, + ChromaFormat::Yuv444 => punktfunk_core::quic::CHROMA_IDC_444, + } + } + + /// True for full-chroma 4:4:4. + pub fn is_444(self) -> bool { + matches!(self, ChromaFormat::Yuv444) + } +} + +impl Codec { + /// Map a negotiated `quic` codec bit ([`punktfunk_core::quic::CODEC_H264`] etc.) to the encoder + /// [`Codec`]. Unknown / `0` → HEVC (the pre-negotiation default). Inverse of [`Codec::to_wire`]. + pub fn from_wire(bit: u8) -> Codec { + match bit { + punktfunk_core::quic::CODEC_H264 => Codec::H264, + punktfunk_core::quic::CODEC_AV1 => Codec::Av1, + punktfunk_core::quic::CODEC_PYROWAVE => Codec::PyroWave, + _ => Codec::H265, + } + } + + /// The single `quic` codec bit for this codec (echoed in [`punktfunk_core::quic::Welcome::codec`]). + pub fn to_wire(self) -> u8 { + match self { + Codec::H264 => punktfunk_core::quic::CODEC_H264, + Codec::H265 => punktfunk_core::quic::CODEC_HEVC, + Codec::Av1 => punktfunk_core::quic::CODEC_AV1, + Codec::PyroWave => punktfunk_core::quic::CODEC_PYROWAVE, + } + } + + /// Lowercase stats/console label (`"h264"` / `"hevc"` / `"av1"`) — the codec string seeded into + /// the web console's session meta ([`crate::stats_recorder::StatsRecorder::register_session`]). + pub fn label(self) -> &'static str { + match self { + Codec::H264 => "h264", + Codec::H265 => "hevc", + Codec::Av1 => "av1", + Codec::PyroWave => "pyrowave", + } + } + + /// The FFmpeg NVENC encoder name (selected by name, not codec id — the latter would + /// pick the software encoder). + pub fn nvenc_name(self) -> &'static str { + match self { + Codec::H264 => "h264_nvenc", + Codec::H265 => "hevc_nvenc", + Codec::Av1 => "av1_nvenc", + // Guarded by the open_video dispatch: a PyroWave session never reaches a + // libavcodec backend. + Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), + } + } + + /// The FFmpeg VAAPI encoder name (AMD via Mesa `radeonsi`, Intel via `iHD`/`i965`). One + /// libavcodec encoder per codec covers both vendors — the kernel driver differs, the libva + /// userspace API is identical. Selected by name (the codec id would pick the SW encoder). + /// AV1 VAAPI encode is narrow (Intel Arc/Xe2+, AMD RDNA3+/RDNA4) — gate it on a capability + /// probe, never assume it (see [`open_video`]). + pub fn vaapi_name(self) -> &'static str { + match self { + Codec::H264 => "h264_vaapi", + Codec::H265 => "hevc_vaapi", + Codec::Av1 => "av1_vaapi", + // Guarded by the open_video dispatch: a PyroWave session never reaches a + // libavcodec backend. + Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), + } + } + + /// The FFmpeg AMD **AMF** encoder name (the Windows AMD backend). Selected by name (the codec id + /// would pick the software encoder). AV1 (`av1_amf`) is RDNA3+/RX 7000+ — probe, never assume. + pub fn amf_name(self) -> &'static str { + match self { + Codec::H264 => "h264_amf", + Codec::H265 => "hevc_amf", + Codec::Av1 => "av1_amf", + // Guarded by the open_video dispatch: a PyroWave session never reaches a + // libavcodec backend. + Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), + } + } + + /// The FFmpeg Intel **QSV** encoder name (the Windows Intel backend). Selected by name. AV1 + /// (`av1_qsv`) is Arc/Xe2+; HEVC Main10 is Gen9.5+ — probe, never assume. + pub fn qsv_name(self) -> &'static str { + match self { + Codec::H264 => "h264_qsv", + Codec::H265 => "hevc_qsv", + Codec::Av1 => "av1_qsv", + // Guarded by the open_video dispatch: a PyroWave session never reaches a + // libavcodec backend. + Codec::PyroWave => unreachable!("PyroWave has no FFmpeg encoder"), + } + } +} + +/// Static capabilities an [`Encoder`] declares so the session glue routes loss-recovery and HDR +/// plumbing by *query* rather than relying on a method's no-op/`false` default. Cheap `Copy`; fixed +/// for the session (an HDR toggle re-initialises the encoder — re-query if that matters). +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct EncoderCaps { + /// The encoder can perform real reference-frame invalidation — i.e. + /// [`invalidate_ref_frames`](Encoder::invalidate_ref_frames) can return `true`. When `false` + /// the caller skips that always-`false` call and forces a keyframe directly on loss recovery. + /// Two backends implement RFI: Windows direct-NVENC (`nvEncInvalidateRefFrames`) and native + /// AMF (user-LTR force-reference, when the driver accepted the LTR slots at open). The + /// libavcodec paths (Linux NVENC, VAAPI, QSV) can't express it and always keyframe. + pub supports_rfi: bool, + /// The encoder emits in-band HDR mastering/CLL SEI from [`set_hdr_meta`](Encoder::set_hdr_meta). + /// When `false`, `set_hdr_meta` is a no-op and no in-band grade reaches the client. Only the + /// Windows direct-NVENC path attaches it today. + pub supports_hdr_metadata: bool, + /// The opened encoder is actually producing a full-chroma 4:4:4 (`chroma_format_idc = 3`) stream. + /// `false` on every 4:2:0 session (the default) and on a backend that declined 4:4:4. Set by the + /// NVENC backends (Linux + Windows). The chroma is committed to the wire (`Welcome::chroma_format`) + /// from the pre-open probe, so this is a *post-open cross-check*: the session glue logs loudly if + /// the encoder's real chroma disagrees with what was negotiated (the in-band SPS is authoritative + /// for the decoder either way). + pub chroma_444: bool, + /// The encoder runs a periodic **intra-refresh wave** — a moving band of intra blocks that + /// re-codes the whole picture over ~0.5 s, no periodic IDR. FEC-unrecoverable loss self-heals as + /// the band sweeps, so the session glue rate-limits client keyframe requests instead of answering + /// each with a full IDR (the 20-40× frame-size spike that cascades under loss). Linux NVENC / AMF + /// set it when `PUNKTFUNK_INTRA_REFRESH` opened the encoder in that mode; VAAPI/QSV/software never + /// do. NOTE — the wave carries NO decoder-visible clean-point: FFmpeg never sets `AV_FRAME_FLAG_KEY` + /// at a recovery point (H.264 flags key only when `recovery_frame_cnt == 0`; HEVC only on IRAP), + /// and AMF emits no recovery-point SEI at all. So this cap ALONE does not let the client lift its + /// post-loss freeze without an IDR — that needs [`intra_refresh_recovery`](Self::intra_refresh_recovery). + pub intra_refresh: bool, + /// The intra-refresh wave is a *validated constrained GDR* — verified on real hardware to fully + /// heal a lost picture within one wave period with no residual artifacts. Only then does the host + /// tag each wave-boundary AU with [`USER_FLAG_RECOVERY_POINT`](punktfunk_core::packet::USER_FLAG_RECOVERY_POINT), + /// so the client can lift its freeze on the second mark (a proven clean re-anchor) instead of + /// waiting out its backstop and forcing a full IDR. Default `false` on every backend until on-glass + /// validation flips it — an un-validated encoder keeps the IDR recovery path, so this is inert and + /// cannot regress. Meaningless unless [`intra_refresh`](Self::intra_refresh) is also set. + pub intra_refresh_recovery: bool, + /// Length of the intra-refresh wave in frames — the boundary period the host marks on (it sets + /// `USER_FLAG_RECOVERY_POINT` on every Nth emitted AU, re-phased at each IDR). 0 when intra-refresh + /// is off. Only consulted when [`intra_refresh_recovery`](Self::intra_refresh_recovery) is set. + pub intra_refresh_period: u32, +} + +/// A hardware encoder. One per session; runs on the encode thread. +pub trait Encoder: Send { + fn submit(&mut self, frame: &CapturedFrame) -> Result<()>; + /// [`submit`](Self::submit) with the **wire frame index** this frame's AU will carry — the + /// number the packetizer stamps on it and the client's loss reports/RFI requests name. The + /// session glue predicts it exactly as `AUs sent so far + frames in flight` (AUs are emitted + /// FIFO, one per submission; anything that would break the prediction — an in-place reset, a + /// device-change teardown, an encoder rebuild — forfeits the in-flight frames on BOTH sides + /// and clears the encoder's reference state, so stale predictions die with it). The RFI + /// backends pin their frame numbering (LTR marks, DPB timestamps) to this so + /// [`invalidate_ref_frames`](Self::invalidate_ref_frames) compares client frame numbers + /// against the same domain — an encoder-internal counter desyncs from the wire on the first + /// mid-stream rebuild (adaptive bitrate steps do this under congestion, exactly when losses + /// happen). Default: ignore the index and delegate to `submit` (backends without per-frame + /// reference bookkeeping don't care). + fn submit_indexed(&mut self, frame: &CapturedFrame, wire_index: u32) -> Result<()> { + let _ = wire_index; + self.submit(frame) + } + /// This encoder's static [capabilities](EncoderCaps) (RFI, HDR SEI), so the session glue can + /// route by query rather than rely on the no-op/`false` defaults of + /// [`invalidate_ref_frames`](Self::invalidate_ref_frames) / [`set_hdr_meta`](Self::set_hdr_meta). + /// Default: no optional capabilities (the SDR / libavcodec backends) — only the direct-NVENC + /// path overrides it. + fn caps(&self) -> EncoderCaps { + EncoderCaps::default() + } + /// Force the next submitted frame to be an IDR keyframe (e.g. after a client + /// reference-frame-invalidation request). Default: no-op. + fn request_keyframe(&mut self) {} + /// Set the source's static HDR mastering metadata (from the capturer). An HDR encoder emits it + /// as in-band SEI (`mastering_display_colour_volume` + `content_light_level_info`) on each + /// keyframe so any decoder — including stock Moonlight — tone-maps from the source's real grade. + /// Default: no-op (SDR encoders / libavcodec paths that don't attach it yet). Cheap to call + /// every frame; only the direct-NVENC path consumes it. + fn set_hdr_meta(&mut self, _meta: Option) {} + /// Invalidate a contiguous range of previously-encoded reference frames (client frame numbers + /// — WIRE frame indexes, the domain [`submit_indexed`](Self::submit_indexed) pins the encoder's + /// bookkeeping to) so the encoder re-references an older still-valid frame instead of emitting + /// a full IDR. Returns `true` if a real reference invalidation was performed; `false` means the + /// encoder couldn't (range older than the DPB/LTR history, or the backend has no RFI) and the + /// caller should fall back to [`request_keyframe`](Self::request_keyframe). Default: `false` — + /// the Windows direct-NVENC path (`nvEncInvalidateRefFrames`) and native AMF (LTR + /// force-reference) implement true RFI; the libavcodec paths can't express it, so they keyframe. + fn invalidate_ref_frames(&mut self, _first_frame: i64, _last_frame: i64) -> bool { + false + } + /// Pull the next encoded AU if one is ready. + fn poll(&mut self) -> Result>; + /// Tear the underlying hardware encoder down and rebuild it in place, keeping the session's + /// negotiated parameters — the encode-stall watchdog's recovery lever (a wedged AMF/QSV + /// driver stops emitting AUs or accepting frames without ever returning an error). Returns + /// `true` when the encoder was rebuilt: every submitted-but-unpolled frame is forfeited and + /// the next submitted frame starts a fresh stream (IDR). Default `false`: the backend has no + /// in-place rebuild and the caller must treat the stall as fatal instead. + fn reset(&mut self) -> bool { + false + } + /// Retarget the encoder's rate control to `bps` (average == max, CBR) **in place** — same + /// codec/resolution/fps, only the bitrate and its derived VBV move. Returns `true` when the + /// live encoder accepted the change: the reference chain, the in-flight frames and the + /// caller's wire-index prediction all survive, so an adaptive-bitrate step costs *nothing* on + /// the wire (no IDR, no in-flight forfeit — the whole point vs. a rebuild). `false` = the + /// backend can't (or the driver rejected the new rate, e.g. above the codec-level ceiling) — + /// the caller falls back to its full rebuild path, which also owns the bitrate clamping. + /// Default: no in-place retarget (the libavcodec/software paths). + fn reconfigure_bitrate(&mut self, _bps: u64) -> bool { + false + } + /// Wire-chunk the encoder's AUs at the session's shard payload size (the PyroWave + /// datagram-aligned mode, plan §4.4): every `shard_payload` window of the emitted AU + /// starts a fresh self-delimiting codec packet, zero-padded to the window — so a lost + /// datagram costs a few coefficient blocks, not the frame. AUs produced this way are + /// flagged [`EncodedFrame::chunk_aligned`] and the session marks them on the wire. + /// Default: no-op (the H.26x backends' bitstreams cannot be cut losslessly). + fn set_wire_chunking(&mut self, _shard_payload: usize) {} + /// Signal end-of-stream. After this, drain the remaining AUs with [`poll`](Self::poll) + /// until it returns `None` — NVENC buffers frames internally even at `delay=0`. + fn flush(&mut self) -> Result<()>; +} + +impl Codec { + /// Maximum encodable dimension (px) per side for this codec on NVENC. H.264 tops out at + /// 4096 (level constraint); HEVC and AV1 allow 8192. Used to reject out-of-range client + /// modes up front (see [`validate_dimensions`]). + pub fn max_dimension(self) -> u32 { + match self { + Codec::H264 => 4096, + // PyroWave has no codec-level dimension cap (arbitrary even sizes); 8192 matches the + // buffer-math guard the other codecs get. + Codec::H265 | Codec::Av1 | Codec::PyroWave => 8192, + } + } + + /// The codec's *spec* top level/tier bitrate (bits/s) — the usual boundary at which NVENC + /// starts rejecting `avcodec_open2` with EINVAL. NOT a hard cap: [`open_video`](crate::encode:: + /// open_video) probes the actual GPU ceiling by stepping DOWN from the requested bitrate only on + /// EINVAL, and uses this purely as the first step-down candidate (so a card that accepts more — + /// an RTX 5070 Ti does >1 Gbps HEVC where a 4090 caps at ~800 Mbps — is never clamped to it). + /// HEVC Level 6.2 High tier = 800 Mbps; H.264 High level 6.2 ≈ 480 Mbps; AV1's levels allow more. + pub fn max_bitrate_bps(self) -> u64 { + match self { + Codec::H264 => 480_000_000, + Codec::H265 => 800_000_000, + Codec::Av1 => 1_200_000_000, + // No spec level/tier: the rate is a plain per-frame byte budget. Use the protocol's + // own bitrate clamp so the step-down probe logic never binds below it. + Codec::PyroWave => 8_000_000_000, + } + } +} + +/// `PUNKTFUNK_VBV_FRAMES` — HRD/VBV size in frame intervals (default 1.0, the strict low-latency +/// shape every backend ships: each frame must fit its rate share, keeping frame sizes uniform for +/// the pacer). The AMF/VAAPI/QSV paths parse the same variable locally; this helper brings the +/// direct-NVENC paths (which used to hardwire 1 frame) to parity. Larger values let complex +/// frames borrow bits — better rate utilization at the cost of per-frame size variance. +pub(crate) fn vbv_frames_env() -> f64 { + std::env::var("PUNKTFUNK_VBV_FRAMES") + .ok() + .and_then(|s| s.parse::().ok()) + .filter(|v| v.is_finite() && *v > 0.0) + .unwrap_or(1.0) +} + +/// Validate a requested encode resolution before we allocate buffers or open NVENC. Rejects +/// zero/odd-sized and out-of-range modes with a clear error instead of letting buffer math +/// overflow or the encoder open fail with an opaque NVENC code. A client can request any +/// `mode=WxHxFPS`, so this is the gate on attacker/typo-controlled dimensions. +pub fn validate_dimensions(codec: Codec, width: u32, height: u32) -> Result<()> { + if width == 0 || height == 0 { + anyhow::bail!("invalid encode resolution {width}x{height}: dimensions must be non-zero"); + } + // NVENC requires even dimensions for the chroma subsampling it does internally. + if width % 2 != 0 || height % 2 != 0 { + anyhow::bail!("invalid encode resolution {width}x{height}: dimensions must be even"); + } + // PyroWave's 5-level wavelet decomposition needs ≥ 4·2⁵ px per axis (upstream + // `MinimumImageSize` — the band mirroring breaks below it); reject a tiny mode here + // (e.g. a match-window resize dragged to a sliver) instead of failing the encoder + // rebuild after the switch was acked. + if codec == Codec::PyroWave && (width < 128 || height < 128) { + anyhow::bail!( + "invalid PyroWave resolution {width}x{height}: the wavelet needs at least 128px per axis" + ); + } + let max = codec.max_dimension(); + if width > max || height > max { + anyhow::bail!( + "{codec:?} max dimension is {max}px; requested {width}x{height} \ + (use HEVC/AV1 above 4096, or lower the client resolution)" + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_zero_and_odd_dimensions() { + assert!(validate_dimensions(Codec::H265, 0, 1080).is_err()); + assert!(validate_dimensions(Codec::H265, 1920, 0).is_err()); + assert!(validate_dimensions(Codec::H265, 1921, 1080).is_err()); // odd width + assert!(validate_dimensions(Codec::H265, 1920, 1081).is_err()); // odd height + } + + #[test] + fn h264_capped_at_4096() { + assert!(validate_dimensions(Codec::H264, 3840, 2160).is_ok()); // 4K fits (width < 4096) + assert!(validate_dimensions(Codec::H264, 4096, 4096).is_ok()); // exactly at the limit + assert!(validate_dimensions(Codec::H264, 4098, 2160).is_err()); + assert!(validate_dimensions(Codec::H264, 3840, 4098).is_err()); + } + + #[test] + fn hevc_and_av1_allow_up_to_8192() { + for c in [Codec::H265, Codec::Av1] { + assert!(validate_dimensions(c, 3840, 2160).is_ok()); + assert!(validate_dimensions(c, 7680, 4320).is_ok()); // 8K fits + assert!(validate_dimensions(c, 8192, 8192).is_ok()); + assert!(validate_dimensions(c, 8194, 4320).is_err()); + } + } + + #[test] + fn common_modes_accepted() { + for c in [Codec::H264, Codec::H265, Codec::Av1] { + for (w, h) in [(1280, 720), (1920, 1080), (2560, 1440)] { + assert!(validate_dimensions(c, w, h).is_ok(), "{c:?} {w}x{h}"); + } + } + } + + /// Wire round-trip and the stats label stay in lockstep with the `quic::CODEC_*` bits. + #[test] + fn codec_wire_roundtrip_and_label() { + for c in [Codec::H264, Codec::H265, Codec::Av1] { + assert_eq!(Codec::from_wire(c.to_wire()), c); + } + assert_eq!(Codec::H264.label(), "h264"); + assert_eq!(Codec::H265.label(), "hevc"); + assert_eq!(Codec::Av1.label(), "av1"); + } +}