enricobuehler
68a863866a
Stage B of the zero-copy host packetize path (networking-audit deferred plan §1): Packetizer::packetize_each yields (header, shard) pairs in exact wire order; Session::seal_frame writes seq(8) ‖ header(40) ‖ shard ‖ tag scratch directly into the pooled wire buffer and seals [8..] in place. The per-packet intermediate Vec (header ++ body) and its extra memcpy are gone — with Stage A, every data byte is now copied once (frame → wire) instead of three times, and the ~2 transient allocs/packet on the send thread are zero after pool warmup (~180k allocs/s at 1 Gbps rates). packetize() stays as a thin wrapper over packetize_each — the reference implementation used by tests and the loss harness. - wire-equivalence test: pooled path vs wrapper path byte-identical across multi-block/partial-tail/exact-multiple/empty frames, fec 0%/50%, both schemes, crypto on/off - loss-harness sweep: recovery rates identical to the pre-item-1 baseline - bench pipeline (end-to-end incl. client half) vs pre-item-1 baseline, stages A+B cumulative: gf16/64K -3.6%, gf16/1M -3.2%; gf8 cases are Cauchy-math-bound and unchanged within noise - cargo ndk check (arm64-v8a) green Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
778 lines
31 KiB
Rust
778 lines
31 KiB
Rust
//! Zero-copy wire framing: split an access unit into FEC blocks of MTU-sized shards,
|
||
//! and reassemble + FEC-recover them on the far side.
|
||
//!
|
||
//! ## Wire layout
|
||
//!
|
||
//! Each packet is a fixed [`PacketHeader`] followed by one FEC shard's payload. Fields
|
||
//! are host-endian for now (every target platform is little-endian); the `punktfunk/1` (P2)
|
||
//! spec will pin byte order explicitly when we talk to non-LE peers.
|
||
//!
|
||
//! ## GameStream mapping (P1)
|
||
//!
|
||
//! `frame_index`↔`frameIndex`, `stream_seq`↔`streamPacketIndex`,
|
||
//! (`block_index`,`block_count`)↔the `multiFecBlocks` nibbles, and
|
||
//! (`data_shards`,`recovery_shards`,`shard_index`)↔the `fecInfo` bitfield. We carry them
|
||
//! as explicit fields rather than bit-packing; full GameStream wire-exactness is a GameStream-host
|
||
//! concern (it also needs RTP framing + RTSP), this is the coherent internal format.
|
||
|
||
use crate::config::Config;
|
||
use crate::error::{PunktfunkError, Result};
|
||
use crate::fec::ErasureCoder;
|
||
use crate::session::Frame;
|
||
use crate::stats::StatsCounters;
|
||
use std::collections::{BTreeMap, HashMap, HashSet};
|
||
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
|
||
|
||
/// Identifies a punktfunk video packet (vs. an input datagram, see [`crate::input`]).
|
||
pub const PUNKTFUNK_MAGIC: u8 = 0xC9;
|
||
|
||
// Frame flags (mirroring GameStream's FLAG_*).
|
||
pub const FLAG_PIC: u8 = 0x1;
|
||
pub const FLAG_EOF: u8 = 0x2;
|
||
pub const FLAG_SOF: u8 = 0x4;
|
||
/// Bandwidth-probe filler, not decodable video: a [`crate::quic::ProbeRequest`] speed test makes
|
||
/// the host burst access units carrying this flag so the client measures throughput/loss without
|
||
/// feeding them to the decoder. Punktfunk/1 only (GameStream never sets it).
|
||
pub const FLAG_PROBE: u8 = 0x8;
|
||
|
||
/// Crypto framing overhead [`Session`](crate::session::Session) adds when encrypting:
|
||
/// an 8-byte sequence prefix plus the GCM tag.
|
||
pub const CRYPTO_OVERHEAD: usize = 8 + crate::crypto::TAG_LEN;
|
||
|
||
/// Largest UDP datagram the core will send or accept. `Config::validate` bounds
|
||
/// `shard_payload` so `HEADER_LEN + shard_payload + CRYPTO_OVERHEAD ≤ MAX_DATAGRAM_BYTES`.
|
||
pub const MAX_DATAGRAM_BYTES: usize = 2048;
|
||
|
||
/// How far behind the newest frame's capture pts an INCOMPLETE frame may sit before it is
|
||
/// declared lost (counted in `frames_dropped`, which triggers the client's recovery-keyframe
|
||
/// request). TIME-based, not frame-count-based, so the fuse is the same at every refresh rate: a
|
||
/// fixed index window is refresh-relative (4 frames = 66 ms at 60 fps but only 33 ms at 120 fps —
|
||
/// inside normal Wi-Fi retry/block-ack reorder timescales, where a delayed-not-lost shard can
|
||
/// trail newer frames). Observed live at 120 fps: the too-tight fuse declared merely-late frames
|
||
/// dead every few seconds, and each false loss cost a recovery-IDR burst + an inflated loss report
|
||
/// (FEC churn) — a self-sustaining latency/bitrate oscillation. 120 ms rides safely above radio
|
||
/// retry jitter while still detecting a real loss ~2× faster than the original 16-frame window did
|
||
/// at 60 fps.
|
||
const LOSS_WINDOW_NS: u64 = 120_000_000;
|
||
|
||
/// Hard cap on how many frame INDICES behind the newest an incomplete frame may sit, whatever its
|
||
/// pts claims — bounds the reassembler's memory against a corrupt/hostile pts (which
|
||
/// [`LOSS_WINDOW_NS`] alone would trust) and against pathologically high frame rates. At 120 fps,
|
||
/// 120 ms ≈ 14 indices, so 64 leaves ample slack up to ~500 fps.
|
||
const HARD_LOSS_WINDOW: u32 = 64;
|
||
|
||
/// 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
|
||
/// loss verdict.
|
||
const REORDER_WINDOW: u32 = 64;
|
||
|
||
/// Fixed per-packet header. `#[repr(C)]`, no padding, zero-copy (de)serializable.
|
||
#[repr(C)]
|
||
#[derive(Clone, Copy, Debug, FromBytes, IntoBytes, KnownLayout, Immutable)]
|
||
pub struct PacketHeader {
|
||
pub pts_ns: u64,
|
||
pub frame_index: u32,
|
||
pub stream_seq: u32,
|
||
pub frame_bytes: u32,
|
||
pub user_flags: u32,
|
||
pub block_index: u16,
|
||
pub block_count: u16,
|
||
pub data_shards: u16,
|
||
pub recovery_shards: u16,
|
||
pub shard_index: u16,
|
||
pub shard_bytes: u16,
|
||
pub magic: u8,
|
||
pub version: u8,
|
||
pub fec_scheme: u8,
|
||
pub flags: u8,
|
||
}
|
||
|
||
/// Size of [`PacketHeader`] on the wire (40 bytes).
|
||
pub const HEADER_LEN: usize = std::mem::size_of::<PacketHeader>();
|
||
|
||
const _: () = assert!(HEADER_LEN == 40, "PacketHeader must be 40 bytes / unpadded");
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Host side: packetization
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Splits encoded access units into FEC-protected shard packets. Host-side only.
|
||
pub struct Packetizer {
|
||
next_frame_index: u32,
|
||
next_seq: u32,
|
||
shard_payload: usize,
|
||
fec: crate::config::FecConfig,
|
||
version: u8,
|
||
/// Reusable zero-padded scratch for the frame's final data shard when the frame isn't an
|
||
/// exact `shard_payload` multiple (and for the single all-zero shard of an empty frame).
|
||
/// Every other data shard is a `shard_payload`-sized slice straight into the frame buffer —
|
||
/// blocks are consecutive shard ranges, so only the frame's last shard can be partial.
|
||
tail: Vec<u8>,
|
||
}
|
||
|
||
impl Packetizer {
|
||
pub fn new(config: &Config) -> Self {
|
||
Packetizer {
|
||
next_frame_index: 0,
|
||
next_seq: 0,
|
||
shard_payload: config.shard_payload,
|
||
fec: config.fec,
|
||
version: config.phase as u8,
|
||
tail: Vec::new(),
|
||
}
|
||
}
|
||
|
||
/// Live-adjust the FEC recovery percentage (adaptive FEC). Takes effect on the next
|
||
/// [`packetize`](Self::packetize); the wire is self-describing (each packet carries its block's
|
||
/// data/recovery counts), so the receiver needs no notification. Clamped to ≤ 90.
|
||
pub fn set_fec_percent(&mut self, pct: u8) {
|
||
self.fec.fec_percent = pct.min(90);
|
||
}
|
||
|
||
/// The current FEC recovery percentage.
|
||
pub fn fec_percent(&self) -> u8 {
|
||
self.fec.fec_percent
|
||
}
|
||
|
||
/// Packetize one access unit into owned wire packets (header ++ shard payload each).
|
||
/// Thin wrapper over [`packetize_each`](Self::packetize_each) — the allocation-free
|
||
/// streaming path's reference implementation (tests and the loss harness use this).
|
||
pub fn packetize(
|
||
&mut self,
|
||
frame: &[u8],
|
||
pts_ns: u64,
|
||
user_flags: u32,
|
||
coder: &dyn ErasureCoder,
|
||
) -> Result<Vec<Vec<u8>>> {
|
||
let mut packets = Vec::new();
|
||
self.packetize_each(frame, pts_ns, user_flags, coder, |hdr, body| {
|
||
let mut pkt = Vec::with_capacity(HEADER_LEN + body.len());
|
||
pkt.extend_from_slice(hdr.as_bytes());
|
||
pkt.extend_from_slice(body);
|
||
packets.push(pkt);
|
||
Ok(())
|
||
})?;
|
||
Ok(packets)
|
||
}
|
||
|
||
/// Packetize one access unit, yielding each packet to `emit` as a `(header, shard bytes)`
|
||
/// pair — in exact wire order, which is also the order the session's nonce counter
|
||
/// advances. No per-packet allocation happens here, so the caller can write header and
|
||
/// shard straight into a pooled wire buffer and seal in place
|
||
/// ([`Session::seal_frame`](crate::session::Session::seal_frame)). An `emit` error aborts
|
||
/// the frame mid-way (packet numbering has already advanced — callers treat it as fatal).
|
||
pub fn packetize_each(
|
||
&mut self,
|
||
frame: &[u8],
|
||
pts_ns: u64,
|
||
user_flags: u32,
|
||
coder: &dyn ErasureCoder,
|
||
mut emit: impl FnMut(&PacketHeader, &[u8]) -> Result<()>,
|
||
) -> Result<()> {
|
||
let payload = self.shard_payload;
|
||
let frame_index = self.next_frame_index;
|
||
self.next_frame_index = self.next_frame_index.wrapping_add(1);
|
||
|
||
// At least one (zero-padded) data shard even for an empty frame.
|
||
let total_data = frame.len().div_ceil(payload).max(1);
|
||
let max_block = self.fec.max_data_per_block as usize;
|
||
let block_count = total_data.div_ceil(max_block).max(1);
|
||
let frame_bytes = frame.len() as u32;
|
||
|
||
// Defend the u16 wire fields against silent truncation. `Config::validate`
|
||
// already rejects configs that could reach these for valid frame sizes; this is
|
||
// the belt-and-suspenders for a frame larger than the negotiated maximum.
|
||
if payload > u16::MAX as usize {
|
||
return Err(PunktfunkError::InvalidArg("shard_payload exceeds u16"));
|
||
}
|
||
if block_count > u16::MAX as usize {
|
||
return Err(PunktfunkError::Unsupported(
|
||
"frame too large: block count exceeds u16",
|
||
));
|
||
}
|
||
|
||
// Stage the frame's one possibly-partial shard (the last) in the reusable
|
||
// zero-padded scratch; every full shard is referenced in place below.
|
||
let full_shards = frame.len() / payload;
|
||
self.tail.clear();
|
||
self.tail.resize(payload, 0);
|
||
let rem = frame.len() % payload;
|
||
if rem > 0 {
|
||
self.tail[..rem].copy_from_slice(&frame[full_shards * payload..]);
|
||
}
|
||
let tail = &self.tail;
|
||
let shard_at = |s: usize| -> &[u8] {
|
||
if s < full_shards {
|
||
&frame[s * payload..(s + 1) * payload]
|
||
} else {
|
||
tail.as_slice()
|
||
}
|
||
};
|
||
|
||
for b in 0..block_count {
|
||
let first = b * max_block;
|
||
let last = ((b + 1) * max_block).min(total_data);
|
||
let block_data_count = last - first;
|
||
|
||
// This block's data shards: references into `frame` (plus the staged tail).
|
||
let data_shards: Vec<&[u8]> = (first..last).map(shard_at).collect();
|
||
|
||
let recovery_count = self.fec.recovery_for(block_data_count);
|
||
let recovery = coder.encode(&data_shards, recovery_count)?;
|
||
let total_shards = block_data_count + recovery_count;
|
||
if total_shards > u16::MAX as usize {
|
||
return Err(PunktfunkError::Unsupported("block shard count exceeds u16"));
|
||
}
|
||
|
||
for shard_index in 0..total_shards {
|
||
let body: &[u8] = if shard_index < block_data_count {
|
||
data_shards[shard_index]
|
||
} else {
|
||
&recovery[shard_index - block_data_count]
|
||
};
|
||
|
||
let seq = self.next_seq;
|
||
self.next_seq = self.next_seq.wrapping_add(1);
|
||
|
||
let mut flags = FLAG_PIC;
|
||
if b == 0 && shard_index == 0 {
|
||
flags |= FLAG_SOF;
|
||
}
|
||
if b + 1 == block_count && shard_index + 1 == total_shards {
|
||
flags |= FLAG_EOF;
|
||
}
|
||
|
||
let hdr = PacketHeader {
|
||
pts_ns,
|
||
frame_index,
|
||
stream_seq: seq,
|
||
frame_bytes,
|
||
user_flags,
|
||
block_index: b as u16,
|
||
block_count: block_count as u16,
|
||
data_shards: block_data_count as u16,
|
||
recovery_shards: recovery_count as u16,
|
||
shard_index: shard_index as u16,
|
||
shard_bytes: payload as u16,
|
||
magic: PUNKTFUNK_MAGIC,
|
||
version: self.version,
|
||
fec_scheme: coder.scheme() as u8,
|
||
flags,
|
||
};
|
||
emit(&hdr, body)?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Client side: reassembly + FEC recovery
|
||
// ---------------------------------------------------------------------------
|
||
|
||
struct BlockBuf {
|
||
data_shards: usize,
|
||
recovery_shards: usize,
|
||
shard_bytes: usize,
|
||
/// Length `data_shards + recovery_shards`; `Some` = received.
|
||
shards: Vec<Option<Vec<u8>>>,
|
||
received: usize,
|
||
done: bool,
|
||
}
|
||
|
||
struct FrameBuf {
|
||
frame_bytes: usize,
|
||
block_count: usize,
|
||
pts_ns: u64,
|
||
user_flags: u32,
|
||
blocks: HashMap<u16, BlockBuf>,
|
||
/// Reconstructed payload per completed block, ordered by block index.
|
||
block_data: BTreeMap<u16, Vec<u8>>,
|
||
}
|
||
|
||
/// Per-session bounds the reassembler enforces on every packet header *before*
|
||
/// allocating, so a hostile or corrupt header cannot drive unbounded memory use. All
|
||
/// derived from the negotiated [`Config`].
|
||
#[derive(Clone, Copy, Debug)]
|
||
pub struct ReassemblerLimits {
|
||
/// Expected shard payload length; every shard in the stream must match exactly.
|
||
pub shard_bytes: usize,
|
||
/// Max data shards per block (the negotiated `max_data_per_block`).
|
||
pub max_data_shards: usize,
|
||
/// Max total shards per block (data + recovery), capped by the FEC scheme ceiling.
|
||
pub max_total_shards: usize,
|
||
/// Max FEC blocks per frame.
|
||
pub max_blocks: usize,
|
||
/// Max accepted access-unit size.
|
||
pub max_frame_bytes: usize,
|
||
}
|
||
|
||
impl ReassemblerLimits {
|
||
pub fn from_config(c: &Config) -> Self {
|
||
let max_data = c.fec.max_data_per_block as usize;
|
||
let max_total =
|
||
(max_data + c.fec.recovery_for(max_data)).min(c.fec.scheme.max_total_shards());
|
||
let total_data = c.max_frame_bytes.div_ceil(c.shard_payload.max(1)).max(1);
|
||
ReassemblerLimits {
|
||
shard_bytes: c.shard_payload,
|
||
max_data_shards: max_data,
|
||
max_total_shards: max_total,
|
||
max_blocks: total_data.div_ceil(max_data).max(1),
|
||
max_frame_bytes: c.max_frame_bytes,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Buffers incoming shards, recovers lost ones via FEC, and emits whole access units.
|
||
/// Client-side only.
|
||
pub struct Reassembler {
|
||
limits: ReassemblerLimits,
|
||
frames: HashMap<u32, FrameBuf>,
|
||
/// Recently-emitted frames, so stray/late shards can't resurrect them. Pruned to
|
||
/// the reorder window alongside `frames`.
|
||
completed: HashSet<u32>,
|
||
/// The newest frame seen, as `(frame_index, capture pts)` — the loss-window anchor: an
|
||
/// incomplete frame is declared lost once it sits [`LOSS_WINDOW_NS`] behind this pts (or
|
||
/// [`HARD_LOSS_WINDOW`] indices, whichever trips first).
|
||
newest_frame: Option<(u32, u64)>,
|
||
}
|
||
|
||
impl Reassembler {
|
||
pub fn new(limits: ReassemblerLimits) -> Self {
|
||
Reassembler {
|
||
limits,
|
||
frames: HashMap::new(),
|
||
completed: HashSet::new(),
|
||
newest_frame: None,
|
||
}
|
||
}
|
||
|
||
/// Ingest one (already-decrypted) packet. Returns the access unit when its last
|
||
/// block completes, otherwise `None`.
|
||
pub fn push(
|
||
&mut self,
|
||
pkt: &[u8],
|
||
coder: &dyn ErasureCoder,
|
||
stats: &StatsCounters,
|
||
) -> Result<Option<Frame>> {
|
||
// On a lossy datagram link a malformed or non-video packet is dropped, never
|
||
// fatal: it must not abort `poll_frame`. A FEC reconstruction failure (corrupt or
|
||
// incompatible shards that passed the header checks) likewise drops the block rather
|
||
// than killing the whole session — the stream recovers at the next keyframe/RFI.
|
||
if pkt.len() < HEADER_LEN {
|
||
StatsCounters::add(&stats.packets_dropped, 1);
|
||
return Ok(None);
|
||
}
|
||
let hdr = match PacketHeader::read_from_bytes(&pkt[..HEADER_LEN]) {
|
||
Ok(h) => h,
|
||
Err(_) => {
|
||
StatsCounters::add(&stats.packets_dropped, 1);
|
||
return Ok(None);
|
||
}
|
||
};
|
||
|
||
let lim = self.limits;
|
||
let shard_bytes = hdr.shard_bytes as usize;
|
||
let data_shards = hdr.data_shards as usize;
|
||
let recovery_shards = hdr.recovery_shards as usize;
|
||
let total = data_shards + recovery_shards;
|
||
let shard_index = hdr.shard_index as usize;
|
||
let block_count = hdr.block_count as usize;
|
||
let frame_bytes = hdr.frame_bytes as usize;
|
||
|
||
// Bound every attacker-controllable header field against the negotiated limits
|
||
// BEFORE allocating anything keyed on it — this is the firewall against a tiny
|
||
// datagram triggering a huge `vec![None; total]` / `Vec::with_capacity`.
|
||
let drop = |stats: &StatsCounters| {
|
||
StatsCounters::add(&stats.packets_dropped, 1);
|
||
};
|
||
if hdr.magic != PUNKTFUNK_MAGIC
|
||
|| shard_bytes != lim.shard_bytes
|
||
|| pkt.len() < HEADER_LEN + shard_bytes
|
||
|| data_shards == 0
|
||
|| data_shards > lim.max_data_shards
|
||
|| total == 0
|
||
|| total > lim.max_total_shards
|
||
|| shard_index >= total
|
||
|| block_count == 0
|
||
|| block_count > lim.max_blocks
|
||
|| hdr.block_index as usize >= block_count
|
||
|| frame_bytes > lim.max_frame_bytes
|
||
{
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
let payload = pkt[HEADER_LEN..HEADER_LEN + shard_bytes].to_vec();
|
||
|
||
self.advance_window(hdr.frame_index, hdr.pts_ns, stats);
|
||
|
||
// Drop shards for frames we've already emitted (e.g. the recovery shards of a
|
||
// frame that completed early via the all-originals-present fast path) or that
|
||
// have fallen out of the loss window.
|
||
if self.completed.contains(&hdr.frame_index) || self.is_stale(hdr.frame_index, hdr.pts_ns) {
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
|
||
// First packet of a frame establishes its geometry; later packets must agree.
|
||
let frame = self
|
||
.frames
|
||
.entry(hdr.frame_index)
|
||
.or_insert_with(|| FrameBuf {
|
||
frame_bytes,
|
||
block_count,
|
||
pts_ns: hdr.pts_ns,
|
||
user_flags: hdr.user_flags,
|
||
blocks: HashMap::new(),
|
||
block_data: BTreeMap::new(),
|
||
});
|
||
if frame.block_count != block_count || frame.frame_bytes != frame_bytes {
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
|
||
if frame.block_data.contains_key(&hdr.block_index) {
|
||
return Ok(None); // block already reconstructed; late/duplicate shard
|
||
}
|
||
|
||
// First packet of a block sizes its shard vector; later packets must match its
|
||
// (data, recovery, shard_bytes) geometry, so `shard_index` is always in bounds.
|
||
frame
|
||
.blocks
|
||
.entry(hdr.block_index)
|
||
.or_insert_with(|| BlockBuf {
|
||
data_shards,
|
||
recovery_shards,
|
||
shard_bytes,
|
||
shards: vec![None; total],
|
||
received: 0,
|
||
done: false,
|
||
});
|
||
let block = frame.blocks.get_mut(&hdr.block_index).unwrap();
|
||
if block.data_shards != data_shards
|
||
|| block.recovery_shards != recovery_shards
|
||
|| block.shard_bytes != shard_bytes
|
||
{
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
|
||
if block.shards[shard_index].is_none() {
|
||
block.shards[shard_index] = Some(payload);
|
||
block.received += 1;
|
||
}
|
||
|
||
// Reconstruct as soon as we hold enough shards.
|
||
if !block.done && block.received >= block.data_shards {
|
||
let present_data = block.shards[..block.data_shards]
|
||
.iter()
|
||
.filter(|s| s.is_some())
|
||
.count();
|
||
let recovered = match coder.reconstruct(
|
||
block.data_shards,
|
||
block.recovery_shards,
|
||
&mut block.shards,
|
||
) {
|
||
Ok(r) => r,
|
||
Err(_) => {
|
||
// Corrupt/incompatible shards that slipped past the header checks: discard this
|
||
// block (mark done so later shards for it are ignored) and keep the session
|
||
// alive — a lossy link must not be torn down by one unrecoverable block; the
|
||
// frame stays incomplete and the client recovers at the next keyframe/RFI.
|
||
block.done = true;
|
||
StatsCounters::add(&stats.packets_dropped, 1);
|
||
return Ok(None);
|
||
}
|
||
};
|
||
block.done = true;
|
||
StatsCounters::add(
|
||
&stats.fec_recovered_shards,
|
||
(block.data_shards - present_data) as u64,
|
||
);
|
||
|
||
// Concatenate the block's data shards into its contiguous payload.
|
||
let mut block_payload = Vec::with_capacity(block.data_shards * block.shard_bytes);
|
||
for shard in &recovered {
|
||
block_payload.extend_from_slice(shard);
|
||
}
|
||
frame.block_data.insert(hdr.block_index, block_payload);
|
||
frame.blocks.remove(&hdr.block_index);
|
||
}
|
||
|
||
// Whole frame ready?
|
||
if frame.block_data.len() == frame.block_count {
|
||
let frame = self.frames.remove(&hdr.frame_index).unwrap();
|
||
self.completed.insert(hdr.frame_index);
|
||
// Reserve based on the bytes we actually hold, not the (already-bounded but
|
||
// still caller-supplied) frame_bytes, so a small frame can't over-reserve.
|
||
let actual: usize = frame.block_data.values().map(|b| b.len()).sum();
|
||
let mut data = Vec::with_capacity(actual);
|
||
for (_, block_payload) in frame.block_data.into_iter() {
|
||
data.extend_from_slice(&block_payload);
|
||
}
|
||
data.truncate(frame.frame_bytes); // trim trailing-shard zero padding
|
||
return Ok(Some(Frame {
|
||
data,
|
||
frame_index: hdr.frame_index,
|
||
pts_ns: frame.pts_ns,
|
||
flags: frame.user_flags,
|
||
}));
|
||
}
|
||
Ok(None)
|
||
}
|
||
|
||
/// Track the newest frame, declare incomplete frames that fell out of the loss window
|
||
/// ([`LOSS_WINDOW_NS`] behind the newest pts, or [`HARD_LOSS_WINDOW`] indices) lost — counting
|
||
/// them dropped, which is what drives the client's recovery-keyframe request — and prune the
|
||
/// completed-index memory to [`REORDER_WINDOW`].
|
||
fn advance_window(&mut self, frame_index: u32, pts_ns: u64, stats: &StatsCounters) {
|
||
let (newest, newest_pts) = match self.newest_frame {
|
||
// `frame_index` is newer iff it's within the forward half of the index space.
|
||
Some((n, p)) if frame_index.wrapping_sub(n) > u32::MAX / 2 => (n, p),
|
||
_ => (frame_index, pts_ns),
|
||
};
|
||
self.newest_frame = Some((newest, newest_pts));
|
||
|
||
let before = self.frames.len();
|
||
let completed = &mut self.completed;
|
||
self.frames.retain(|&idx, f| {
|
||
let keep = newest.wrapping_sub(idx) <= HARD_LOSS_WINDOW
|
||
&& newest_pts.saturating_sub(f.pts_ns) <= LOSS_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
|
||
// and double-count the drop when it aged out again.
|
||
completed.insert(idx);
|
||
}
|
||
keep
|
||
});
|
||
let pruned = before - self.frames.len();
|
||
if pruned > 0 {
|
||
StatsCounters::add(&stats.frames_dropped, pruned as u64);
|
||
}
|
||
self.completed
|
||
.retain(|&idx| newest.wrapping_sub(idx) <= REORDER_WINDOW);
|
||
}
|
||
|
||
/// Drop all in-flight state — every partially-assembled frame and the completed/abandoned
|
||
/// index memory — as if the session just started. Used by the client's backlog flush
|
||
/// ([`Session::flush_backlog`](crate::session::Session::flush_backlog)): after the socket
|
||
/// backlog is discarded wholesale, the partial frames here can never complete (their remaining
|
||
/// shards were just thrown away) and the window anchor (`newest_frame`) points into the
|
||
/// discarded past.
|
||
pub fn reset(&mut self) {
|
||
self.frames.clear();
|
||
self.completed.clear();
|
||
self.newest_frame = None;
|
||
}
|
||
|
||
/// True if this packet's frame lies outside the loss window (behind the newest frame by more
|
||
/// than [`LOSS_WINDOW_NS`] of capture time or [`HARD_LOSS_WINDOW`] indices) — its shards
|
||
/// arrive too late to be useful, and accepting one would only create a frame buffer the next
|
||
/// [`advance_window`] immediately declares lost.
|
||
fn is_stale(&self, frame_index: u32, pts_ns: u64) -> bool {
|
||
match self.newest_frame {
|
||
Some((n, newest_pts)) => {
|
||
let behind = n.wrapping_sub(frame_index);
|
||
behind <= u32::MAX / 2
|
||
&& (behind > HARD_LOSS_WINDOW
|
||
|| newest_pts.saturating_sub(pts_ns) > LOSS_WINDOW_NS)
|
||
}
|
||
None => false,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::config::FecScheme;
|
||
use crate::fec::coder_for;
|
||
|
||
fn limits() -> ReassemblerLimits {
|
||
ReassemblerLimits {
|
||
shard_bytes: 16,
|
||
max_data_shards: 8,
|
||
max_total_shards: 12,
|
||
max_blocks: 4,
|
||
max_frame_bytes: 4096,
|
||
}
|
||
}
|
||
|
||
fn base_header() -> PacketHeader {
|
||
PacketHeader {
|
||
pts_ns: 0,
|
||
frame_index: 0,
|
||
stream_seq: 0,
|
||
frame_bytes: 16,
|
||
user_flags: 0,
|
||
block_index: 0,
|
||
block_count: 1,
|
||
data_shards: 1,
|
||
recovery_shards: 0,
|
||
shard_index: 0,
|
||
shard_bytes: 16,
|
||
magic: PUNKTFUNK_MAGIC,
|
||
version: 1,
|
||
fec_scheme: 0,
|
||
flags: FLAG_PIC,
|
||
}
|
||
}
|
||
|
||
fn packet(h: PacketHeader) -> Vec<u8> {
|
||
let mut p = Vec::new();
|
||
p.extend_from_slice(h.as_bytes());
|
||
p.extend_from_slice(&vec![0xAB; h.shard_bytes as usize]);
|
||
p
|
||
}
|
||
|
||
/// A header advertising 65535+65535 shards must be dropped, not allocate gigabytes.
|
||
#[test]
|
||
fn rejects_oversized_shard_counts() {
|
||
let mut r = Reassembler::new(limits());
|
||
let coder = coder_for(FecScheme::Gf8);
|
||
let stats = StatsCounters::default();
|
||
let mut h = base_header();
|
||
h.data_shards = 65535;
|
||
h.recovery_shards = 65535;
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
assert_eq!(stats.snapshot().packets_dropped, 1);
|
||
}
|
||
|
||
/// A second packet for a block whose geometry differs from the first must be dropped
|
||
/// — never index past the block's allocated shard vector (the old OOB panic).
|
||
#[test]
|
||
fn rejects_inconsistent_block_geometry_without_panicking() {
|
||
let mut r = Reassembler::new(limits());
|
||
let coder = coder_for(FecScheme::Gf8);
|
||
let stats = StatsCounters::default();
|
||
|
||
let mut h1 = base_header();
|
||
h1.data_shards = 4;
|
||
h1.recovery_shards = 2; // block sized to 6 slots
|
||
h1.frame_bytes = 64;
|
||
assert!(r
|
||
.push(&packet(h1), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
|
||
// Same block, different geometry, shard_index valid for ITS total (8) but past
|
||
// the established block's 6 slots.
|
||
let mut h2 = base_header();
|
||
h2.data_shards = 6;
|
||
h2.recovery_shards = 2;
|
||
h2.shard_index = 7;
|
||
h2.frame_bytes = 64;
|
||
assert!(r
|
||
.push(&packet(h2), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
assert_eq!(stats.snapshot().packets_dropped, 1);
|
||
}
|
||
|
||
/// The loss window is TIME-based: an incomplete frame survives newer frames arriving within
|
||
/// [`LOSS_WINDOW_NS`] of its capture pts (a 33 ms-late shard at 120 fps is late, not lost —
|
||
/// the old 4-INDEX window wrongly killed it), is declared lost once the newest pts moves past
|
||
/// the window (`frames_dropped`), and a straggler shard can't resurrect it afterwards.
|
||
#[test]
|
||
fn incomplete_frames_age_out_by_capture_time_not_frame_count() {
|
||
let mut r = Reassembler::new(limits());
|
||
let coder = coder_for(FecScheme::Gf8);
|
||
let stats = StatsCounters::default();
|
||
const FRAME_NS: u64 = 8_333_333; // 120 fps
|
||
|
||
// Frame 0: one of its two shards arrives — incomplete.
|
||
let mut h = base_header();
|
||
h.data_shards = 2;
|
||
h.frame_bytes = 32;
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
|
||
// Frames 1..=8 complete around it (well past the old 4-index window, inside 120 ms):
|
||
// frame 0 must still be alive — no drop counted.
|
||
for i in 1..=8u32 {
|
||
let mut h = base_header();
|
||
h.frame_index = i;
|
||
h.pts_ns = i as u64 * FRAME_NS;
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_some());
|
||
}
|
||
assert_eq!(stats.snapshot().frames_dropped, 0);
|
||
|
||
// Frame 0's second shard arrives 8 frames late (~66 ms at 120 fps) — completes fine.
|
||
let mut h = base_header();
|
||
h.data_shards = 2;
|
||
h.frame_bytes = 32;
|
||
h.shard_index = 1;
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_some());
|
||
|
||
// Frame 20: incomplete again; then a frame lands past the 120 ms window → declared lost.
|
||
let mut h = base_header();
|
||
h.frame_index = 20;
|
||
h.pts_ns = 20 * FRAME_NS;
|
||
h.data_shards = 2;
|
||
h.frame_bytes = 32;
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
let mut h = base_header();
|
||
h.frame_index = 21;
|
||
h.pts_ns = 20 * FRAME_NS + LOSS_WINDOW_NS + 1;
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_some());
|
||
assert_eq!(stats.snapshot().frames_dropped, 1);
|
||
|
||
// A straggler shard for the abandoned frame 20 is dropped, never resurrected.
|
||
let mut h = base_header();
|
||
h.frame_index = 20;
|
||
h.pts_ns = 20 * FRAME_NS;
|
||
h.data_shards = 2;
|
||
h.frame_bytes = 32;
|
||
h.shard_index = 1;
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
assert_eq!(stats.snapshot().frames_dropped, 1, "no double-count");
|
||
}
|
||
|
||
#[test]
|
||
fn rejects_wrong_shard_bytes_and_oversized_frame() {
|
||
let coder = coder_for(FecScheme::Gf8);
|
||
|
||
let mut r = Reassembler::new(limits());
|
||
let stats = StatsCounters::default();
|
||
let mut h = base_header();
|
||
h.shard_bytes = 8; // != negotiated 16
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
assert_eq!(stats.snapshot().packets_dropped, 1);
|
||
|
||
let mut r = Reassembler::new(limits());
|
||
let stats = StatsCounters::default();
|
||
let mut h = base_header();
|
||
h.frame_bytes = 1_000_000; // > max_frame_bytes
|
||
assert!(r
|
||
.push(&packet(h), coder.as_ref(), &stats)
|
||
.unwrap()
|
||
.is_none());
|
||
assert_eq!(stats.snapshot().packets_dropped, 1);
|
||
}
|
||
}
|