diff --git a/crates/pf-client-core/src/session.rs b/crates/pf-client-core/src/session.rs index 03078f81..5cb000fb 100644 --- a/crates/pf-client-core/src/session.rs +++ b/crates/pf-client-core/src/session.rs @@ -289,7 +289,12 @@ fn pump( let built = if connector.codec == punktfunk_core::quic::CODEC_PYROWAVE { let mode = connector.mode(); match params.vulkan.as_ref() { - Some(vk) => Decoder::new_pyrowave(vk, mode.width, mode.height), + Some(vk) => Decoder::new_pyrowave( + vk, + mode.width, + mode.height, + connector.shard_payload as usize, + ), None => Err(anyhow::anyhow!( "pyrowave session without a presenter device" )), @@ -324,6 +329,8 @@ fn pump( // step, drift) keep the capture-clock latency stats honest — never cached at session start. let clock_offset_live = connector.clock_offset_shared(); let mut total_frames = 0u64; + // Newest frame index handed to the decoder — the staleness bar for late partials. + let mut newest_decoded_idx: Option = None; let mut window_start = Instant::now(); let mut frames_n = 0u32; let mut bytes_n = 0u64; @@ -432,7 +439,21 @@ fn pump( } None => next_expected_index = Some(frame.frame_index.wrapping_add(1)), } - match decoder.decode(&frame.data) { + // A PARTIAL that lost the race (a newer frame already decoded) is pure + // time travel — skip it; each PyroWave frame is independent, so nothing + // downstream needs it. Completes keep the normal path (reorder is handled + // by the continuity gate). + if !frame.complete + && newest_decoded_idx + .is_some_and(|n: u32| n.wrapping_sub(frame.frame_index) <= u32::MAX / 2) + { + continue; + } + newest_decoded_idx = Some(match newest_decoded_idx { + Some(n) if frame.frame_index.wrapping_sub(n) > u32::MAX / 2 => n, + _ => frame.frame_index, + }); + match decoder.decode_frame(&frame.data, frame.flags, frame.complete) { Ok(Some(image)) => { // Fold this decoded frame through the shared freeze gate: it reads the AU's // re-anchor wire flags (FLAG_SOF IDR marker / RECOVERY_ANCHOR / RECOVERY_POINT), diff --git a/crates/pf-client-core/src/video.rs b/crates/pf-client-core/src/video.rs index c3c47cb3..0866cc18 100644 --- a/crates/pf-client-core/src/video.rs +++ b/crates/pf-client-core/src/video.rs @@ -572,10 +572,18 @@ impl Decoder { /// compute on the presenter's device, no FFmpeg. `codec_id` is irrelevant (kept as /// HEVC so an — impossible — demotion path stays well-formed). #[cfg(all(target_os = "linux", feature = "pyrowave"))] - pub fn new_pyrowave(vk: &VulkanDecodeDevice, width: u32, height: u32) -> Result { + pub fn new_pyrowave( + vk: &VulkanDecodeDevice, + width: u32, + height: u32, + shard_payload: usize, + ) -> Result { Ok(Decoder { backend: Backend::PyroWave(Box::new(crate::video_pyrowave::PyroWaveDecoder::new( - vk, width, height, + vk, + width, + height, + shard_payload, )?)), codec_id: ffmpeg::codec::Id::HEVC, vaapi_fails: 0, @@ -610,8 +618,24 @@ impl Decoder { /// pump asks the host for a fresh IDR — under the infinite GOP nothing else resyncs a /// rebuilt/erroring decoder, so skipping this leaves the picture gray/frozen for good. pub fn decode(&mut self, au: &[u8]) -> Result> { + self.decode_frame(au, 0, true) + } + + /// [`decode`](Self::decode) with the AU's wire facts: `user_flags` (chunk-aligned AUs + /// are parsed in shard windows — [`punktfunk_core::packet::USER_FLAG_CHUNK_ALIGNED`]) + /// and completeness (`false` = a partial delivery; only the PyroWave backend decodes + /// those — as one frame of localized blur, plan §4.4). + pub fn decode_frame( + &mut self, + au: &[u8], + user_flags: u32, + complete: bool, + ) -> Result> { let result = match &mut self.backend { - Backend::Vulkan(v) => v.decode(au).map(|f| f.map(DecodedImage::VkFrame)), + Backend::Vulkan(v) => { + debug_assert!(complete, "partial AUs are pyrowave-only"); + v.decode(au).map(|f| f.map(DecodedImage::VkFrame)) + } #[cfg(target_os = "linux")] Backend::Vaapi(v) => v.decode(au).map(|f| f.map(DecodedImage::Dmabuf)), #[cfg(windows)] @@ -620,7 +644,12 @@ impl Decoder { // error; the pump surfaces it and the session falls back to HEVC by // renegotiation (plan §4.6), not by decoder swap. #[cfg(all(target_os = "linux", feature = "pyrowave"))] - Backend::PyroWave(p) => return Ok(p.decode(au)?.map(DecodedImage::PyroWave)), + Backend::PyroWave(p) => { + let aligned = user_flags & punktfunk_core::packet::USER_FLAG_CHUNK_ALIGNED != 0; + return Ok(p + .decode_frame(au, aligned, complete)? + .map(DecodedImage::PyroWave)); + } Backend::Software(s) => return Ok(s.decode(au)?.map(DecodedImage::Cpu)), }; match result { diff --git a/crates/pf-client-core/src/video_pyrowave.rs b/crates/pf-client-core/src/video_pyrowave.rs index 2ebf6ff9..74cce622 100644 --- a/crates/pf-client-core/src/video_pyrowave.rs +++ b/crates/pf-client-core/src/video_pyrowave.rs @@ -190,6 +190,9 @@ pub struct PyroWaveDecoder { fence: vk::Fence, width: u32, height: u32, + /// The wire shard payload — the parse-window size for chunk-aligned AUs (§4.4): each + /// window holds whole self-delimiting codec packets, zero-padded to the window. + wire_window: usize, } // SAFETY: used only from the single decode thread; the shared-queue accesses go through @@ -197,7 +200,12 @@ pub struct PyroWaveDecoder { unsafe impl Send for PyroWaveDecoder {} impl PyroWaveDecoder { - pub fn new(vkd: &VulkanDecodeDevice, width: u32, height: u32) -> Result { + pub fn new( + vkd: &VulkanDecodeDevice, + width: u32, + height: u32, + shard_payload: usize, + ) -> Result { if !vkd.pyrowave_decode { bail!("presenter device lacks the PyroWave compute feature set"); } @@ -207,13 +215,14 @@ impl PyroWaveDecoder { // SAFETY: the handles in `vkd` are the presenter's live instance/device (it // outlives the decoder — same contract the FFmpeg Vulkan backend relies on); // `Hold` pins the reconstructed create-infos for the pyrowave device's lifetime. - unsafe { Self::new_inner(vkd, width, height) } + unsafe { Self::new_inner(vkd, width, height, shard_payload) } } unsafe fn new_inner( vkd: &VulkanDecodeDevice, width: u32, height: u32, + shard_payload: usize, ) -> Result { let static_fn = ash::StaticFn { get_instance_proc_addr: std::mem::transmute::( @@ -380,25 +389,113 @@ impl PyroWaveDecoder { fence, width, height, + wire_window: shard_payload.max(64), }) } - /// One AU in → one frame out (the AU is a complete pyrowave frame: one packet). - pub fn decode(&mut self, au: &[u8]) -> Result> { + /// One AU in → one frame out. `aligned` = the AU is shard-window chunked (each + /// `wire_window` holds whole self-delimiting packets, zero-padded — walk and strip); + /// `complete` = every shard arrived (a partial decodes anyway: missing blocks are + /// localized blur for exactly this frame, §4.4). + pub fn decode_frame( + &mut self, + au: &[u8], + aligned: bool, + complete: bool, + ) -> Result> { // SAFETY: single decode thread; all handles owned/pinned by `self`; queue access // serialized under the device-wide QueueLock; the fence bounds GPU completion // before the frame is handed to the presenter. - unsafe { self.decode_inner(au) } + unsafe { self.decode_inner(au, aligned, complete) } } - unsafe fn decode_inner(&mut self, au: &[u8]) -> Result> { - pw_check( - pw::pyrowave_decoder_push_packet(self.pw_dec, au.as_ptr() as *const c_void, au.len()), - "push_packet", - )?; - // The reassembler delivers complete AUs only, so a frame is ready per push; a - // stale/duplicate packet (sequence rewind) simply isn't — skip, no error. - if !pw::pyrowave_decoder_decode_is_ready(self.pw_dec, false) { + /// Consume one framed shard window (§4.4): a 4-byte prefix (u16 used-length + u16 + /// kind) then either WHOLE self-delimiting codec packets (PACKED) or one fragment of + /// an oversized packet (FRAG chain). A lost shard arrives as a zeroed window + /// (used = 0) — skipped, and it breaks any fragment chain it interrupts (that + /// packet's blocks are unusable without their end; dropping them is the §4.4 blur). + unsafe fn push_window(&mut self, win: &[u8], frag: &mut Vec) -> Result<()> { + if win.len() < 4 { + return Ok(()); + } + let used = u16::from_le_bytes([win[0], win[1]]) as usize; + let kind = u16::from_le_bytes([win[2], win[3]]); + if used == 0 || 4 + used > win.len() { + frag.clear(); // missing / garbage window — drop any chain in progress + return Ok(()); + } + let body = &win[4..4 + used]; + match kind { + 0 => { + frag.clear(); + pw_check( + pw::pyrowave_decoder_push_packet( + self.pw_dec, + body.as_ptr() as *const c_void, + body.len(), + ), + "push_packet", + ) + } + 1 => { + frag.clear(); + frag.extend_from_slice(body); + Ok(()) + } + 2 => { + if !frag.is_empty() { + frag.extend_from_slice(body); + } + Ok(()) + } + 3 => { + if !frag.is_empty() { + frag.extend_from_slice(body); + let r = pw_check( + pw::pyrowave_decoder_push_packet( + self.pw_dec, + frag.as_ptr() as *const c_void, + frag.len(), + ), + "push_packet (fragmented)", + ); + frag.clear(); + return r; + } + Ok(()) + } + _ => { + frag.clear(); + Ok(()) + } + } + } + + unsafe fn decode_inner( + &mut self, + au: &[u8], + aligned: bool, + complete: bool, + ) -> Result> { + if aligned { + let mut frag: Vec = Vec::new(); + for win in au.chunks(self.wire_window) { + self.push_window(win, &mut frag)?; + } + } else { + pw_check( + pw::pyrowave_decoder_push_packet( + self.pw_dec, + au.as_ptr() as *const c_void, + au.len(), + ), + "push_packet", + )?; + } + // A complete AU that isn't ready is a stale/duplicate (sequence rewind) — skip. + // A PARTIAL is decoded regardless: missing wavelet blocks reconstruct as zeros, + // i.e. localized blur for exactly this one frame (the next is complete again). + if complete && !pw::pyrowave_decoder_decode_is_ready(self.pw_dec, false) { return Ok(None); } diff --git a/crates/pf-presenter/src/vk.rs b/crates/pf-presenter/src/vk.rs index eff7528e..abbb8dc5 100644 --- a/crates/pf-presenter/src/vk.rs +++ b/crates/pf-presenter/src/vk.rs @@ -1254,7 +1254,7 @@ impl Presenter { .csc_planar .as_ref() .context("PyroWave frame but the device failed the pyrowave probe")?; - planar.bind_planes_planar(&self.device, f.views.map(|v| vk::ImageView::from_raw(v))); + planar.bind_planes_planar(&self.device, f.views.map(vk::ImageView::from_raw)); } if let Some(o) = overlay { // Point the composite at this overlay image (same fence-wait safety). diff --git a/crates/punktfunk-core/src/client.rs b/crates/punktfunk-core/src/client.rs index a7dcb212..fe9bb093 100644 --- a/crates/punktfunk-core/src/client.rs +++ b/crates/punktfunk-core/src/client.rs @@ -71,6 +71,8 @@ enum CtrlRequest { #[derive(Clone, Copy)] struct Negotiated { mode: Mode, + /// Wire shard payload — the chunk-aligned parse window (plan §4.4). + shard_payload: u16, compositor: CompositorPref, gamepad: GamepadPref, /// SHA-256 of the certificate the host actually presented (TOFU callers persist this). @@ -513,6 +515,9 @@ pub struct NativeClient { /// requested rate clamped to the host's range, or its default if we requested `0`. `0` = an /// older host that didn't report it. pub resolved_bitrate_kbps: u32, + /// The session's wire shard payload (bytes of AU per datagram) — the parse-window size + /// for chunk-aligned AUs ([`crate::packet::USER_FLAG_CHUNK_ALIGNED`], plan §4.4). + pub shard_payload: u16, /// Host clock minus client clock (ns), from the connect-time skew handshake. Add it to a local /// receive/present timestamp to express it in the host's capture clock (the AU `pts_ns`), making /// glass-to-glass latency valid across machines. `0` = no correction (an old host that didn't @@ -778,6 +783,7 @@ impl NativeClient { resolved_compositor: negotiated.compositor, resolved_gamepad: negotiated.gamepad, resolved_bitrate_kbps: negotiated.bitrate_kbps, + shard_payload: negotiated.shard_payload, clock_offset_ns: negotiated.clock_offset_ns, bit_depth: negotiated.bit_depth, color: negotiated.color, @@ -1540,7 +1546,14 @@ async fn worker_main(args: WorkerArgs) { if let Ok(sock) = transport.try_clone_socket() { crate::transport::spawn_data_punch(sock, shutdown.clone()); } - let session = Session::new(welcome.session_config(Role::Client), Box::new(transport))?; + let mut session = + Session::new(welcome.session_config(Role::Client), Box::new(transport))?; + // PyroWave sessions opt into partial delivery (plan §4.4): an aged-out lossy + // frame arrives as blocks-with-holes instead of vanishing — the all-intra codec + // renders it as one frame of localized blur, strictly better than a freeze. + if welcome.codec == crate::quic::CODEC_PYROWAVE { + session.set_deliver_partial_frames(true); + } Ok::<_, PunktfunkError>(( session, send, @@ -1558,6 +1571,7 @@ async fn worker_main(args: WorkerArgs) { chroma_format: welcome.chroma_format, audio_channels: welcome.audio_channels, codec: welcome.codec, + shard_payload: welcome.shard_payload, }, welcome.host_caps, )) @@ -2490,6 +2504,7 @@ mod frame_channel_tests { frame_index: i, pts_ns: i as u64, flags: 0, + complete: true, } } diff --git a/crates/punktfunk-core/src/packet.rs b/crates/punktfunk-core/src/packet.rs index fb23e505..67b55ff3 100644 --- a/crates/punktfunk-core/src/packet.rs +++ b/crates/punktfunk-core/src/packet.rs @@ -54,6 +54,14 @@ pub const USER_FLAG_RECOVERY_POINT: u32 = 0x10; /// `AV_FRAME_FLAG_KEY` — this host flag is the only signal. pub const USER_FLAG_RECOVERY_ANCHOR: u32 = 0x20; +/// `user_flags` bit: the AU's content is **shard-aligned self-delimiting chunks** — every +/// `shard_payload`-sized window of the frame buffer starts a fresh codec packet, padded to the +/// window with zeros (PyroWave datagram-aligned mode, design/pyrowave-codec-plan.md §4.4). Two +/// consequences: a receiver that opted into partial delivery can use an aged-out frame's buffer +/// AS-IS (missing shards stay zeroed; the codec's block walk skips zero windows), and even a +/// COMPLETE frame must be consumed window-by-window (the padding is not part of the stream). +pub const USER_FLAG_CHUNK_ALIGNED: u32 = 0x40; + /// Widest lost-frame range (frames, wrapping `last - first`) a reference-frame-invalidation /// recovery may be asked to repair; anything wider goes straight to the keyframe path on BOTH /// ends. RFI can only re-reference history the encoder still holds — NVENC keeps a 5-frame DPB, @@ -89,6 +97,13 @@ const LOSS_WINDOW_NS: u64 = 120_000_000; /// 120 ms ≈ 14 indices, so 64 leaves ample slack up to ~500 fps. const HARD_LOSS_WINDOW: u32 = 64; +/// The much tighter fuse for PARTIAL-deliverable frames (chunk-aligned AUs with a consumer +/// that opted in): once anything newer exists and this much capture time passed, the frame +/// is delivered as-is — its stragglers can only make it less late, and each frame is +/// independently decodable, so waiting the full loss window (120 ms) would inject ancient +/// frames into a live stream. ~2 frame periods at 60 fps rides out normal reorder. +const PARTIAL_WINDOW_NS: u64 = 30_000_000; + /// How many frames behind the newest the reassembler remembers emitted/abandoned frame indices /// (`completed`), so a straggler shard can neither resurrect an abandoned frame nor re-open an /// emitted one. Must cover at least [`HARD_LOSS_WINDOW`]: stragglers can trickle in later than the @@ -451,6 +466,13 @@ const RECOVERY_POOL_MAX: usize = 512; /// Client-side only. pub struct Reassembler { limits: ReassemblerLimits, + /// Deliver aged-out incomplete frames whose AUs are [`USER_FLAG_CHUNK_ALIGNED`] instead of + /// silently dropping them (client opt-in — the PyroWave decode path): the frame buffer is + /// already the right shape (received shards at their final offsets, zeros elsewhere). + /// They still count into `frames_dropped` — a partial IS lost data for the loss reports. + deliver_partial: bool, + /// The newest such partial awaiting pickup (newest-wins: partials are a lossy byproduct). + pending_partial: Option, /// The video stream's window — its aged-out incomplete frames count into `frames_dropped` /// (the client's loss-recovery trigger). video: ReassemblyWindow, @@ -472,6 +494,8 @@ impl Reassembler { pub fn new(limits: ReassemblerLimits) -> Self { Reassembler { limits, + deliver_partial: false, + pending_partial: None, video: ReassemblyWindow::default(), probe: ReassemblyWindow::default(), recovery_pool: Vec::new(), @@ -479,6 +503,19 @@ impl Reassembler { } } + /// Opt into partial delivery of chunk-aligned frames (see [`Reassembler::deliver_partial`]). + pub fn set_deliver_partial(&mut self, on: bool) { + self.deliver_partial = on; + if !on { + self.pending_partial = None; + } + } + + /// Take the newest aged-out partial frame, if one is pending (see `set_deliver_partial`). + pub fn take_partial(&mut self) -> Option { + self.pending_partial.take() + } + /// Ingest one (already-decrypted) packet. Returns the access unit when its last /// block completes, otherwise `None`. pub fn push( @@ -507,6 +544,8 @@ impl Reassembler { // in-flight budget are all touched while a frame entry is mutably borrowed. let Reassembler { limits, + deliver_partial, + pending_partial, video, probe, recovery_pool, @@ -579,6 +618,7 @@ impl Reassembler { recovery_pool, in_flight_bytes, lim.max_data_shards, + (*deliver_partial && !is_probe).then_some(pending_partial), ); // Drop shards for frames already terminated (emitted — e.g. the recovery shards of a @@ -756,6 +796,7 @@ impl Reassembler { frame_index: hdr.frame_index, pts_ns: done.pts_ns, flags: done.user_flags, + complete: true, })); } Ok(None) @@ -809,6 +850,8 @@ impl ReassemblyWindow { recovery_pool: &mut Vec>, in_flight_bytes: &mut usize, max_data_shards: usize, + // `Some(sink)` = deliver aged-out CHUNK_ALIGNED frames instead of only dropping them. + mut partial_sink: Option<&mut Option>, ) { let (newest, newest_pts) = match self.newest_frame { // `frame_index` is newer iff it's within the forward half of the index space. @@ -819,9 +862,17 @@ impl ReassemblyWindow { let before = self.frames.len(); let completed = &mut self.completed; + let partial_on = partial_sink.is_some(); self.frames.retain(|&idx, f| { + // Partial-deliverable frames age out on the TIGHT fuse (see PARTIAL_WINDOW_NS); + // everything else keeps the full loss window. + let window_ns = if partial_on && f.user_flags & USER_FLAG_CHUNK_ALIGNED != 0 { + PARTIAL_WINDOW_NS + } else { + LOSS_WINDOW_NS + }; let keep = newest.wrapping_sub(idx) <= HARD_LOSS_WINDOW - && newest_pts.saturating_sub(f.pts_ns) <= LOSS_WINDOW_NS; + && newest_pts.saturating_sub(f.pts_ns) <= window_ns; if !keep { // Remember the abandoned index so a straggler shard is dropped (below, and in // `push`) instead of resurrecting the frame — which would re-allocate its buffers @@ -831,6 +882,28 @@ impl ReassemblyWindow { completed.insert(idx, reconstructed_shards(&f.blocks, max_data_shards)); // Release its buffer budget and reclaim its parity bufs for the pool. *in_flight_bytes -= f.buf.len(); + // Partial delivery (chunk-aligned AUs only): the buffer is already exactly + // what the consumer needs — received shards at their final offsets, zeros + // where shards are missing (the codec's block walk skips zero windows). + // Newest-wins if several age out in one prune. Still counted dropped below. + if let Some(sink) = partial_sink.as_deref_mut() { + if f.user_flags & USER_FLAG_CHUNK_ALIGNED != 0 { + let mut buf = std::mem::take(&mut f.buf); + buf.truncate(f.frame_bytes); + let newer = sink + .as_ref() + .is_none_or(|p| idx.wrapping_sub(p.frame_index) <= u32::MAX / 2); + if newer { + *sink = Some(Frame { + data: buf, + frame_index: idx, + pts_ns: f.pts_ns, + flags: f.user_flags, + complete: false, + }); + } + } + } for block in f.blocks.values_mut() { for slot in block.recovery.iter_mut() { if let Some(rb) = slot.take() { diff --git a/crates/punktfunk-core/src/session.rs b/crates/punktfunk-core/src/session.rs index 35caaaac..42c9af0e 100644 --- a/crates/punktfunk-core/src/session.rs +++ b/crates/punktfunk-core/src/session.rs @@ -24,6 +24,12 @@ pub struct Frame { pub frame_index: u32, pub pts_ns: u64, pub flags: u32, + /// `false` = a partial delivery: the frame aged out of the loss window with shards + /// missing, and the session opted into receiving it anyway + /// ([`Session::set_deliver_partial_frames`]). Only chunk-aligned AUs + /// ([`crate::packet::USER_FLAG_CHUNK_ALIGNED`]) are ever delivered partial; missing + /// shard ranges are zero-filled at their exact offsets. + pub complete: bool, } /// One end of a stream. Constructed for a single [`Role`]; calling the other role's @@ -634,6 +640,20 @@ impl Session { /// Client: drain the transport until a whole access unit is recovered, or no more /// packets are pending ([`PunktfunkError::NoFrame`]). + /// Client opt-in: deliver aged-out incomplete chunk-aligned frames as + /// [`Frame`]`{ complete: false }` instead of only dropping them (the PyroWave + /// datagram-aligned mode, plan §4.4 — a lost datagram costs a few blocks of blur, + /// not the frame). No effect on other codecs' AUs (they never carry the flag). + pub fn set_deliver_partial_frames(&mut self, on: bool) { + self.reassembler.set_deliver_partial(on); + } + + /// The session's negotiated wire shard payload size (bytes of AU per datagram) — + /// the window size for chunk-aligned AUs (`USER_FLAG_CHUNK_ALIGNED`). + pub fn shard_payload(&self) -> usize { + self.config.shard_payload + } + pub fn poll_frame(&mut self) -> Result { if self.config.role != Role::Client { return Err(PunktfunkError::InvalidArg( @@ -661,6 +681,11 @@ impl Session { } self.recv_idx = 0; if self.recv_count == 0 { + // Nothing new on the wire — hand over an aged-out partial if one is + // waiting (it can only get staler). + if let Some(p) = self.reassembler.take_partial() { + return Ok(p); + } return Err(PunktfunkError::NoFrame); } } @@ -725,6 +750,11 @@ impl Session { StatsCounters::add(&self.stats.frames_completed, 1); return Ok(frame); } + // A push that completed nothing may still have aged a partial out — deliver it + // ahead of further draining (its successors are already arriving). + if let Some(p) = self.reassembler.take_partial() { + return Ok(p); + } } } @@ -980,6 +1010,101 @@ mod wire_equivalence_tests { fn pattern(len: usize) -> Vec { (0..len).map(|i| (i * 31 + 7) as u8).collect() } + + /// Partial delivery (plan §4.4): a chunk-aligned frame that loses shards past FEC's + /// reach is DELIVERED once it ages out — `complete: false`, received shards at their + /// exact offsets, missing ranges zero-filled — instead of silently dropping. Plain + /// AUs (no flag) keep the drop behavior even with the opt-in enabled. + #[test] + fn partial_delivery_of_chunk_aligned_frames() { + use crate::packet::USER_FLAG_CHUNK_ALIGNED; + let mk = |role| Config { + role, + phase: ProtocolPhase::P2Punktfunk, + fec: FecConfig { + scheme: FecScheme::Gf16, + fec_percent: 0, // no parity — any drop leaves a hole + max_data_per_block: 64, + }, + shard_payload: 1024, + max_frame_bytes: 8 * 1024 * 1024, + encrypt: false, + key: [0u8; 16], + salt: [0u8; 4], + loopback_drop_period: 0, + }; + // Drop exactly one datagram (the 3rd) out of the first frame's shards. + let (h, c) = crate::transport::loopback_pair(3, 1); + let mut host = Session::new(mk(Role::Host), Box::new(h)).unwrap(); + let mut client = Session::new(mk(Role::Client), Box::new(c)).unwrap(); + client.set_deliver_partial_frames(true); + + // 8 shards of chunk-aligned payload. + let frame = pattern(8 * 1024); + host.submit_frame(&frame, 1_000, USER_FLAG_CHUNK_ALIGNED) + .unwrap(); + // The incomplete frame ages out on the HARD index window — push enough newer + // (complete) frames past it. Collect everything the client emits. + let mut got_partial = None; + let mut completes = 0; + for i in 0..80u64 { + let filler = pattern(1024); + host.submit_frame(&filler, 2_000 + i, USER_FLAG_CHUNK_ALIGNED) + .unwrap(); + loop { + match client.poll_frame() { + Ok(f) if !f.complete => got_partial = Some(f), + Ok(_) => completes += 1, + Err(PunktfunkError::NoFrame) => break, + Err(e) => panic!("unexpected: {e}"), + } + } + } + let p = got_partial.expect("the lossy frame must be delivered partial"); + assert_eq!(p.pts_ns, 1_000); + assert_eq!(p.data.len(), frame.len()); + assert!(p.flags & USER_FLAG_CHUNK_ALIGNED != 0); + // Exactly one 1024-byte shard is zeroed; every other offset matches the original. + let mut zero_windows = 0; + for w in 0..8 { + let win = &p.data[w * 1024..(w + 1) * 1024]; + if win.iter().all(|&b| b == 0) { + zero_windows += 1; + } else { + assert_eq!(win, &frame[w * 1024..(w + 1) * 1024], "window {w} corrupt"); + } + } + // loopback_pair(3, _) drops every 3rd datagram, so several of the 8 shards are + // gone — the exact count depends on phase; what matters is that SOME are zeroed + // and every survivor is intact. + assert!( + (1..8).contains(&zero_windows), + "dropped shards zero-filled (got {zero_windows})" + ); + assert!(completes > 40, "surviving filler frames flow normally"); + + // Control: WITHOUT the flag the same loss is a plain drop, opt-in or not. + let (h2, c2) = crate::transport::loopback_pair(3, 1); + let mut host2 = Session::new(mk(Role::Host), Box::new(h2)).unwrap(); + let mut client2 = Session::new(mk(Role::Client), Box::new(c2)).unwrap(); + client2.set_deliver_partial_frames(true); + host2.submit_frame(&pattern(8 * 1024), 1_000, 0).unwrap(); + let mut saw_partial = false; + for i in 0..80u64 { + host2.submit_frame(&pattern(1024), 2_000 + i, 0).unwrap(); + loop { + match client2.poll_frame() { + Ok(f) => saw_partial |= !f.complete, + Err(PunktfunkError::NoFrame) => break, + Err(e) => panic!("unexpected: {e}"), + } + } + } + assert!( + !saw_partial, + "unflagged AUs must never be delivered partial" + ); + } } #[cfg(test)] diff --git a/crates/punktfunk-host/src/encode.rs b/crates/punktfunk-host/src/encode.rs index 6b69ba8a..7049e831 100644 --- a/crates/punktfunk-host/src/encode.rs +++ b/crates/punktfunk-host/src/encode.rs @@ -29,6 +29,10 @@ pub struct EncodedFrame { /// 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. @@ -356,6 +360,13 @@ pub trait Encoder: Send { 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<()>; @@ -519,6 +530,12 @@ impl Encoder for TrackedEncoder { fn invalidate_ref_frames(&mut self, first_frame: i64, last_frame: i64) -> bool { self.inner.invalidate_ref_frames(first_frame, last_frame) } + // The classic TrackedEncoder trap: a defaulted trait method that isn't forwarded + // silently no-ops through the wrapper (bit the direct-NVENC work, then THIS — the + // §4.4 chunking probe run hit the default while the plan said Some(1408)). + fn set_wire_chunking(&mut self, shard_payload: usize) { + self.inner.set_wire_chunking(shard_payload) + } fn poll(&mut self) -> Result> { self.inner.poll() } diff --git a/crates/punktfunk-host/src/encode/linux/mod.rs b/crates/punktfunk-host/src/encode/linux/mod.rs index 15991caf..4a592d66 100644 --- a/crates/punktfunk-host/src/encode/linux/mod.rs +++ b/crates/punktfunk-host/src/encode/linux/mod.rs @@ -652,6 +652,7 @@ impl Encoder for NvencEncoder { pts_ns, keyframe: pkt.is_key(), recovery_anchor: false, + chunk_aligned: false, })) } // No packet ready yet (need another input frame). diff --git a/crates/punktfunk-host/src/encode/linux/nvenc_cuda.rs b/crates/punktfunk-host/src/encode/linux/nvenc_cuda.rs index 8835bc6c..cba087ff 100644 --- a/crates/punktfunk-host/src/encode/linux/nvenc_cuda.rs +++ b/crates/punktfunk-host/src/encode/linux/nvenc_cuda.rs @@ -1148,6 +1148,7 @@ impl Encoder for NvencCudaEncoder { pts_ns, keyframe, recovery_anchor: anchor, + chunk_aligned: false, })) } } diff --git a/crates/punktfunk-host/src/encode/linux/pyrowave.rs b/crates/punktfunk-host/src/encode/linux/pyrowave.rs index 4793c22d..cf68f6f8 100644 --- a/crates/punktfunk-host/src/encode/linux/pyrowave.rs +++ b/crates/punktfunk-host/src/encode/linux/pyrowave.rs @@ -43,6 +43,13 @@ const IMPORT_CACHE_CAP: usize = 16; /// Headroom over the per-frame rate budget for the packetized bitstream (block headers + meta; /// the rate controller itself never exceeds the budget). const BS_SLACK: usize = 256 * 1024; +/// Chunked-mode window framing (§4.4): 4-byte prefix per shard-sized window. +const WINDOW_PREFIX: usize = 4; +/// Window kinds: whole packets / an oversized packet's fragments. +const WIN_PACKED: u16 = 0; +const WIN_FRAG_FIRST: u16 = 1; +const WIN_FRAG_CONT: u16 = 2; +const WIN_FRAG_LAST: u16 = 3; /// The DRM modifiers the PyroWave device can import as a SAMPLED image of the capture's /// packed-RGB format. The capture advertises these for the pyrowave passthrough instead of @@ -178,6 +185,10 @@ pub struct PyroWaveEncoder { fps: u32, /// Per-frame bitstream budget (hard CBR): `bitrate / (8 * fps)`. frame_budget: usize, + /// Datagram-aligned mode (plan §4.4): packetize at this boundary and pad every codec + /// packet to it, so each wire shard carries whole self-delimiting packets. `None` = + /// one packet per AU (the dense MVP shape). + wire_chunk: Option, bitstream: Vec, pending: VecDeque, frame_count: u64, @@ -532,6 +543,7 @@ impl PyroWaveEncoder { height: h, fps, frame_budget, + wire_chunk: None, bitstream: Vec::new(), pending: VecDeque::new(), frame_count: 0, @@ -865,31 +877,100 @@ impl PyroWaveEncoder { dev.wait_for_fences(&[self.fence], true, 5_000_000_000) .context("pyrowave encode fence")?; - // ---- packetize: boundary = whole buffer, so the AU is exactly one pyrowave packet ---- + // ---- packetize ---- + // Dense (default): boundary = whole buffer → the AU is exactly one pyrowave packet. + // Datagram-aligned (§4.4, `set_wire_chunking`): boundary = the wire shard payload; + // each codec packet is zero-padded to the boundary so every shard carries whole + // self-delimiting packets — the client windows its parse and a lost shard costs + // only those blocks. Padding cost is small: the packetizer fills close to the + // boundary by design. let cap = self.frame_budget + BS_SLACK; self.bitstream.resize(cap, 0); + // Chunked mode reserves 4 bytes per window for the framing prefix. + let boundary = self.wire_chunk.map(|c| c - WINDOW_PREFIX).unwrap_or(cap); let mut n: usize = 0; pw_check( - pw::pyrowave_encoder_compute_num_packets(self.pw_enc, cap, &mut n), + pw::pyrowave_encoder_compute_num_packets(self.pw_enc, boundary, &mut n), "compute_num_packets", )?; - if n != 1 { - bail!("pyrowave: expected a single packet at boundary {cap}, got {n}"); + if n == 0 || (self.wire_chunk.is_none() && n != 1) { + bail!("pyrowave: unexpected packet count {n} at boundary {boundary}"); } - let mut packet = pw::pyrowave_packet { offset: 0, size: 0 }; + let mut packets = vec![pw::pyrowave_packet { offset: 0, size: 0 }; n]; let mut out_n: usize = 0; pw_check( pw::pyrowave_encoder_packetize( self.pw_enc, - &mut packet, - cap, + packets.as_mut_ptr(), + boundary, &mut out_n, self.bitstream.as_mut_ptr() as *mut std::ffi::c_void, cap, ), "packetize", )?; - let au = self.bitstream[packet.offset..packet.offset + packet.size].to_vec(); + packets.truncate(out_n.max(1)); + let au = if let Some(chunk) = self.wire_chunk { + // Window framing (§4.4): each `chunk`-sized window opens with a 4-byte prefix + // (u16 used-length + u16 kind) and carries either WHOLE self-delimiting codec + // packets (PACKED — several small ones share a window) or one fragment of an + // oversized packet (FRAG chain — pyrowave 32×32 blocks are atomic and may + // exceed a shard). A lost shard zeroes its window (used = 0) — the receiver + // skips it and drops any fragment chain it interrupts. + let payload_max = chunk - WINDOW_PREFIX; + let mut au: Vec = Vec::with_capacity((packets.len() + 1) * chunk); + // The currently-open PACKED window: (start offset of its prefix, bytes used). + let mut open: Option<(usize, usize)> = None; + let close = |au: &mut Vec, open: &mut Option<(usize, usize)>, chunk: usize| { + if let Some((start, used)) = open.take() { + au[start..start + 2].copy_from_slice(&(used as u16).to_le_bytes()); + au[start + 2..start + 4].copy_from_slice(&WIN_PACKED.to_le_bytes()); + au.resize(start + chunk, 0); + } + }; + for p in &packets { + let bytes = &self.bitstream[p.offset..p.offset + p.size]; + if p.size <= payload_max { + let fits = open.is_some_and(|(_, used)| used + p.size <= payload_max); + if !fits { + close(&mut au, &mut open, chunk); + let start = au.len(); + au.resize(start + WINDOW_PREFIX, 0); + open = Some((start, 0)); + } + au.extend_from_slice(bytes); + if let Some((_, used)) = open.as_mut() { + *used += p.size; + } + } else { + // Oversized packet: its own FRAG chain of full windows. + close(&mut au, &mut open, chunk); + let mut off = 0usize; + while off < p.size { + let take = (p.size - off).min(payload_max); + let kind = if off == 0 { + WIN_FRAG_FIRST + } else if off + take == p.size { + WIN_FRAG_LAST + } else { + WIN_FRAG_CONT + }; + let start = au.len(); + au.resize(start + WINDOW_PREFIX, 0); + au[start..start + 2].copy_from_slice(&(take as u16).to_le_bytes()); + au[start + 2..start + 4].copy_from_slice(&kind.to_le_bytes()); + au.extend_from_slice(&bytes[off..off + take]); + au.resize(start + chunk, 0); + off += take; + } + } + } + close(&mut au, &mut open, chunk); + au + } else { + let p = &packets[0]; + self.bitstream[p.offset..p.offset + p.size].to_vec() + }; self.frame_count += 1; self.pending.push_back(EncodedFrame { data: au, @@ -898,6 +979,7 @@ impl PyroWaveEncoder { // whole recovery story (plan §1.2). keyframe: true, recovery_anchor: false, + chunk_aligned: self.wire_chunk.is_some(), }); Ok(()) } @@ -960,6 +1042,17 @@ impl Encoder for PyroWaveEncoder { true } + fn set_wire_chunking(&mut self, shard_payload: usize) { + // Sanity floor: a boundary below one block header + payload word is meaningless. + if shard_payload >= 64 { + self.wire_chunk = Some(shard_payload); + tracing::info!( + shard_payload, + "pyrowave: datagram-aligned packetization on (partial-frame loss mode)" + ); + } + } + fn flush(&mut self) -> Result<()> { // Synchronous per-frame encode: nothing buffered beyond `pending`. Ok(()) @@ -1127,6 +1220,83 @@ mod tests { ); } + // Datagram-aligned mode (§4.4): every emitted AU is a whole number of framed + // windows — 4-byte prefix (used-length + kind), whole packets or FRAG chains for + // oversized atomic blocks, zero padding after `used`. Walking + reassembling the + // fragments must reproduce a decodable packet stream. + enc.set_wire_chunking(1408); + enc.submit(&cpu_frame(w, h, 500, [90, 60, 30, 255])) + .expect("chunked submit"); + let au = enc.poll().expect("poll").expect("chunked AU"); + assert!(au.chunk_aligned); + assert_eq!(au.data.len() % 1408, 0, "AU is a whole number of windows"); + // SAFETY: test-only FFI with locally-owned buffers. + unsafe { + let mut dev: pw::pyrowave_device = std::ptr::null_mut(); + assert_eq!( + pw::pyrowave_create_default_device(&mut dev), + pw::pyrowave_result_PYROWAVE_SUCCESS + ); + let dinfo = pw::pyrowave_decoder_create_info { + device: dev, + width: w as i32, + height: h as i32, + chroma: pw::pyrowave_chroma_subsampling_PYROWAVE_CHROMA_SUBSAMPLING_420, + fragment_path: false, + }; + let mut dec: pw::pyrowave_decoder = std::ptr::null_mut(); + assert_eq!( + pw::pyrowave_decoder_create(&dinfo, &mut dec), + pw::pyrowave_result_PYROWAVE_SUCCESS + ); + let mut frag: Vec = Vec::new(); + let mut pushed = 0usize; + for win in au.data.chunks(1408) { + let used = u16::from_le_bytes([win[0], win[1]]) as usize; + let kind = u16::from_le_bytes([win[2], win[3]]); + assert!(4 + used <= win.len(), "window overrun"); + assert!(win[4 + used..].iter().all(|&b| b == 0), "non-zero padding"); + let body = &win[4..4 + used]; + match kind { + 0 => { + assert_eq!( + pw::pyrowave_decoder_push_packet( + dec, + body.as_ptr() as *const _, + body.len() + ), + pw::pyrowave_result_PYROWAVE_SUCCESS + ); + pushed += body.len(); + } + 1 => frag = body.to_vec(), + 2 => frag.extend_from_slice(body), + 3 => { + frag.extend_from_slice(body); + assert_eq!( + pw::pyrowave_decoder_push_packet( + dec, + frag.as_ptr() as *const _, + frag.len() + ), + pw::pyrowave_result_PYROWAVE_SUCCESS + ); + pushed += frag.len(); + frag.clear(); + } + k => panic!("unknown window kind {k}"), + } + } + assert!(pushed > 0, "chunked AU carries real packets"); + assert!( + pw::pyrowave_decoder_decode_is_ready(dec, false), + "chunked AU incomplete after framed walk" + ); + pw::pyrowave_decoder_destroy(dec); + pw::pyrowave_device_destroy(dev); + } + enc.set_wire_chunking(0); // below the floor — back to dense + // In-place rate retarget + encoder rebuild both keep encoding. assert!(enc.reconfigure_bitrate(100_000_000)); assert!(enc.reset()); diff --git a/crates/punktfunk-host/src/encode/linux/vaapi.rs b/crates/punktfunk-host/src/encode/linux/vaapi.rs index eeb6d03d..5385b4f6 100644 --- a/crates/punktfunk-host/src/encode/linux/vaapi.rs +++ b/crates/punktfunk-host/src/encode/linux/vaapi.rs @@ -297,6 +297,7 @@ fn poll_encoder(enc: &mut encoder::video::Encoder, fps: u32) -> Result Result Result<()> { // path now reads this typed `SessionPlan` instead of re-deriving from config at each dispatch site // (the latent "capture and encode disagree on the backend" hazard, plan §2.4). `bit_depth` is the // only per-session input — capture/topology/encoder are otherwise pure functions of `HostConfig`. - let plan = crate::session_plan::SessionPlan::resolve(ctx.bit_depth, ctx.chroma, ctx.codec); + let mut plan = crate::session_plan::SessionPlan::resolve(ctx.bit_depth, ctx.chroma, ctx.codec); + // PyroWave rides the datagram-aligned wire mode (§4.4): every encoder this session opens + // packetizes at the negotiated shard payload, so a lost datagram costs blocks, not frames. + if ctx.codec == crate::encode::Codec::PyroWave { + plan.wire_chunk = Some(ctx.session.shard_payload()); + } tracing::info!(?plan, "resolved session plan"); // Single-process path: unpack the context into the locals the loop below uses (names unchanged, so the // body is byte-for-byte the same; the receivers are now owned but `try_recv()` is identical). @@ -4430,12 +4435,15 @@ fn virtual_stream(ctx: SessionContext) -> Result<()> { bit_depth, plan.chroma, ) { - Ok(new_enc) => { + Ok(mut new_enc) => { tracing::info!( from_kbps = bitrate_kbps, to_kbps = new_kbps, "encoder rebuilt at new bitrate (adaptive bitrate)" ); + if let Some(c) = plan.wire_chunk { + new_enc.set_wire_chunking(c); + } enc = new_enc; bitrate_kbps = new_kbps; live_bitrate.store(new_kbps, Ordering::Relaxed); @@ -4866,6 +4874,11 @@ fn virtual_stream(ctx: SessionContext) -> Result<()> { if au.recovery_anchor { flags |= punktfunk_core::packet::USER_FLAG_RECOVERY_ANCHOR; } + // Datagram-aligned PyroWave AU (plan §4.4): the client windows its parse at the + // shard payload and may opt into partial delivery of lossy frames. + if au.chunk_aligned { + flags |= punktfunk_core::packet::USER_FLAG_CHUNK_ALIGNED; + } // Re-send the HDR mastering metadata (0xCE) on each keyframe (a decoder-resync point) and // whenever it changed, so a client that dropped the best-effort datagram re-converges. if let Some(m) = last_hdr_meta { @@ -5194,7 +5207,7 @@ fn build_pipeline( }; // `bit_depth` is the handshake-negotiated value (8, or 10 = HEVC Main10 when the client // advertised VIDEO_CAP_10BIT and the host opted in). Threaded down from the Welcome. - let enc = crate::encode::open_video( + let mut enc = crate::encode::open_video( plan.codec, frame.format, frame.width, @@ -5206,6 +5219,9 @@ fn build_pipeline( plan.chroma, ) .context("open video encoder")?; + if let Some(c) = plan.wire_chunk { + enc.set_wire_chunking(c); + } // Post-open cross-check: the Welcome already committed `chroma_format` from the pre-open probe, so // warn loudly if the encoder actually opened a different chroma than negotiated (the in-band SPS is // authoritative for the decoder, but a mismatch means the probe and the live open disagreed). diff --git a/crates/punktfunk-host/src/session_plan.rs b/crates/punktfunk-host/src/session_plan.rs index c4168c7b..f203bffb 100644 --- a/crates/punktfunk-host/src/session_plan.rs +++ b/crates/punktfunk-host/src/session_plan.rs @@ -97,6 +97,10 @@ pub struct SessionPlan { /// Handshake-negotiated video codec the encoder emits — HEVC by default, H.264 for a GPU-less /// software host (`resolve_codec` over the client's advertised codecs ∩ the host's capability). pub codec: crate::encode::Codec, + /// Datagram-aligned wire chunking for the encoder (plan §4.4): `Some(shard_payload)` on a + /// PyroWave session — applied to EVERY encoder this plan opens (initial + all rebuilds) so + /// AUs stay shard-aligned across mode/bitrate/stall rebuilds. `None` for the H.26x codecs. + pub wire_chunk: Option, } impl SessionPlan { @@ -115,6 +119,7 @@ impl SessionPlan { hdr: bit_depth >= 10, chroma, codec, + wire_chunk: None, } } diff --git a/include/punktfunk_core.h b/include/punktfunk_core.h index d7545271..5af622be 100644 --- a/include/punktfunk_core.h +++ b/include/punktfunk_core.h @@ -295,6 +295,14 @@ // `AV_FRAME_FLAG_KEY` — this host flag is the only signal. #define USER_FLAG_RECOVERY_ANCHOR 32 +// `user_flags` bit: the AU's content is **shard-aligned self-delimiting chunks** — every +// `shard_payload`-sized window of the frame buffer starts a fresh codec packet, padded to the +// window with zeros (PyroWave datagram-aligned mode, design/pyrowave-codec-plan.md §4.4). Two +// consequences: a receiver that opted into partial delivery can use an aged-out frame's buffer +// AS-IS (missing shards stay zeroed; the codec's block walk skips zero windows), and even a +// COMPLETE frame must be consumed window-by-window (the padding is not part of the stream). +#define USER_FLAG_CHUNK_ALIGNED 64 + // Widest lost-frame range (frames, wrapping `last - first`) a reference-frame-invalidation // recovery may be asked to repair; anything wider goes straight to the keyframe path on BOTH // ends. RFI can only re-reference history the encoder still holds — NVENC keeps a 5-frame DPB,