enricobuehler
93c8dc4712
Turn the 1446-line packet.rs into a packet/ directory module (mod.rs facade + header/packetize/reassemble/tests) behind glob re-exports, so every crate::packet::X path stays byte-stable. Pure move: the header consts + PacketHeader -> header.rs; Packetizer -> packetize.rs; the Reassembler cluster (kept WHOLE -- disjoint-borrow hot path) + loss-window consts -> reassemble.rs; the inline #[cfg(test)] block -> tests.rs. Sole visibility change: LOSS_WINDOW_NS -> pub(super) (a test imports it). No behavior change. Verified on both platforms from a clean HEAD snapshot: Linux clippy (--features quic and --no-default-features, --all-targets -D warnings) + full cargo test; Windows clippy (both feature sets) + cargo test --lib (156 pass). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
635 lines
31 KiB
Rust
635 lines
31 KiB
Rust
//! Client side: buffer incoming shards, FEC-recover lost ones, and emit whole access units.
|
||
//! The per-packet [`Reassembler::push`] hot path is kept whole (disjoint field borrows).
|
||
|
||
use super::*;
|
||
use crate::config::Config;
|
||
use crate::error::Result;
|
||
use crate::fec::ErasureCoder;
|
||
use crate::session::Frame;
|
||
use crate::stats::StatsCounters;
|
||
use std::collections::HashMap;
|
||
use zerocopy::FromBytes;
|
||
|
||
/// 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.
|
||
pub(super) 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;
|
||
|
||
/// 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
|
||
/// loss verdict.
|
||
const REORDER_WINDOW: u32 = 64;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Client side: reassembly + FEC recovery
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Per-block reassembly state. The block's DATA bytes live in the owning [`FrameBuf::buf`]
|
||
/// (each shard copied once, straight to its final AU offset); this tracks presence and holds
|
||
/// the received recovery shards until the block resolves.
|
||
struct BlockState {
|
||
/// The block's K/M — pinned by the frame geometry derived from `frame_bytes` and validated
|
||
/// against every packet of the block.
|
||
data_shards: usize,
|
||
recovery_shards: usize,
|
||
/// Per-data-shard presence: which ranges of the frame buffer hold received bytes (also the
|
||
/// FEC input map — the codec reads only present slots).
|
||
have_data: Vec<bool>,
|
||
data_received: usize,
|
||
/// Received recovery shards (pooled shard-sized buffers, reclaimed when the block resolves).
|
||
recovery: Vec<Option<Vec<u8>>>,
|
||
recovery_received: usize,
|
||
/// Terminal — either reconstructed (its buffer range is fully written) or unrecoverable
|
||
/// (corrupt shards; the frame can never complete). Later shards for it are ignored.
|
||
done: bool,
|
||
/// The block resolved by actually consuming parity (`missing > 0` at reconstruct) — the only
|
||
/// case where a data shard arriving after `done` was counted into `fec_recovered_shards` and
|
||
/// must be netted back out as [`fec_late_shards`](crate::stats::Stats::fec_late_shards).
|
||
reconstructed: bool,
|
||
}
|
||
|
||
struct FrameBuf {
|
||
frame_bytes: usize,
|
||
block_count: usize,
|
||
pts_ns: u64,
|
||
user_flags: u32,
|
||
/// The whole frame's data region — `total_data_shards × shard_bytes` zeroed bytes. Data
|
||
/// shards are copied to their final offset on arrival; FEC reconstruction writes only the
|
||
/// missing shards' ranges. On completion this Vec IS [`Frame::data`] (truncated to
|
||
/// `frame_bytes`) — the old shard→block→AU copy chain and its ~per-packet allocations are
|
||
/// gone (the 2026-07-14 sweeps pinned the client pump as the ~1.5 Gbps wall, ~85% userspace).
|
||
buf: Vec<u8>,
|
||
blocks: HashMap<u16, BlockState>,
|
||
/// Blocks fully reconstructed into `buf`. The frame completes when this reaches
|
||
/// `block_count` (a failed block never counts — the frame then ages out as dropped).
|
||
blocks_ok: usize,
|
||
}
|
||
|
||
/// 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,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// One frame-index space's reassembly state: the in-flight frames, the recently-emitted memory,
|
||
/// and the loss-window anchor. The [`Reassembler`] keeps two — video and speed-test probe filler —
|
||
/// because the two ride **separate index counters** on a [`VIDEO_CAP_PROBE_SEQ`]-aware host
|
||
/// (a probe burst must neither advance the video loss window nor be dropped as "stale" against
|
||
/// it). [`VIDEO_CAP_PROBE_SEQ`]: crate::quic::VIDEO_CAP_PROBE_SEQ
|
||
#[derive(Default)]
|
||
struct ReassemblyWindow {
|
||
frames: HashMap<u32, FrameBuf>,
|
||
/// Recently-terminated frames (emitted OR abandoned by the loss window), so stray/late shards
|
||
/// can't resurrect them. The value is the frame's parity-restored data shards (frame-wide
|
||
/// index `block × max_data_shards + shard`, usually empty): each was counted into
|
||
/// `fec_recovered_shards` at reconstruct, so when one ARRIVES after all — late, not lost —
|
||
/// it's removed here and counted into `fec_late_shards` for the loss windows to net out
|
||
/// (reordering alone must not read as packet loss). The removal makes the accounting exact:
|
||
/// a wire duplicate of a shard that did arrive matches nothing and counts nothing. Pruned to
|
||
/// the reorder window alongside `frames`.
|
||
completed: HashMap<u32, Vec<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)>,
|
||
}
|
||
|
||
/// Frame buffers are allocated whole (zeroed) at a frame's first shard, so bound how much a
|
||
/// window of tiny first-shards can commit: the sum of in-flight `FrameBuf::buf` bytes (both index
|
||
/// spaces) may not exceed `IN_FLIGHT_BUF_FACTOR × max_frame_bytes`. Honest streams hold 1–3
|
||
/// partially-arrived frames of ACTUAL size (≪ max); without this cap, [`HARD_LOSS_WINDOW`]
|
||
/// max-sized declarations from one header-sized packet each could commit gigabytes — an
|
||
/// amplification the old sparse per-shard allocation didn't have.
|
||
const IN_FLIGHT_BUF_FACTOR: usize = 4;
|
||
|
||
/// Recovery-shard buffer pool ceiling (shard-sized buffers): enough for several max-recovery
|
||
/// blocks in flight, small enough (~720 KB at a 1408-byte shard) to keep after a loss burst.
|
||
const RECOVERY_POOL_MAX: usize = 512;
|
||
|
||
/// Buffers incoming shards, recovers lost ones via FEC, and emits whole access units.
|
||
/// 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<Frame>,
|
||
/// The video stream's window — its aged-out incomplete frames count into `frames_dropped`
|
||
/// (the client's loss-recovery trigger).
|
||
video: ReassemblyWindow,
|
||
/// Speed-test probe filler ([`FLAG_PROBE`] in `user_flags`). Routed by the flag, so it also
|
||
/// captures an OLD host's probe frames (which still carry video-space indexes — they complete
|
||
/// fine here, and keeping them out of the video window means a burst can no longer advance the
|
||
/// video loss anchor). Aged-out probe frames are NOT `frames_dropped` — probe loss is measured
|
||
/// bytes-wise by the probe accumulator and must not fire video recovery.
|
||
probe: ReassemblyWindow,
|
||
/// Reusable shard-sized buffers for received recovery shards — the only shard bytes that
|
||
/// still need their own storage (data shards land straight in the frame buffer). Capped at
|
||
/// [`RECOVERY_POOL_MAX`].
|
||
recovery_pool: Vec<Vec<u8>>,
|
||
/// Sum of in-flight `FrameBuf::buf` bytes across both windows (see [`IN_FLIGHT_BUF_FACTOR`]).
|
||
in_flight_bytes: usize,
|
||
}
|
||
|
||
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(),
|
||
in_flight_bytes: 0,
|
||
}
|
||
}
|
||
|
||
/// 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<Frame> {
|
||
self.pending_partial.take()
|
||
}
|
||
|
||
/// 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);
|
||
}
|
||
};
|
||
|
||
// Disjoint field borrows: the window (`video`/`probe`), the recovery pool, and the
|
||
// in-flight budget are all touched while a frame entry is mutably borrowed.
|
||
let Reassembler {
|
||
limits,
|
||
deliver_partial,
|
||
pending_partial,
|
||
video,
|
||
probe,
|
||
recovery_pool,
|
||
in_flight_bytes,
|
||
} = self;
|
||
let lim = *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);
|
||
}
|
||
// Derived-geometry firewall: every sender (our Packetizer, any version) slices a frame
|
||
// into consecutive blocks of exactly `max_data_per_block` data shards with only the LAST
|
||
// block smaller, and stamps the exact `frame_bytes` in every header. That makes every
|
||
// data shard's final AU offset computable on arrival —
|
||
// offset = (block_index × max_data_per_block + shard_index) × shard_bytes
|
||
// — which is what lets shards land straight in the frame buffer below. Enforce the
|
||
// invariant so a header lying about its geometry is dropped instead of scribbling into
|
||
// another shard's range.
|
||
let total_data = frame_bytes.div_ceil(shard_bytes).max(1);
|
||
let expect_blocks = total_data.div_ceil(lim.max_data_shards).max(1);
|
||
let block_idx = hdr.block_index as usize;
|
||
let expect_data_shards = if block_idx + 1 == expect_blocks {
|
||
total_data - (expect_blocks - 1) * lim.max_data_shards
|
||
} else {
|
||
lim.max_data_shards
|
||
};
|
||
if block_count != expect_blocks || data_shards != expect_data_shards {
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
let body = &pkt[HEADER_LEN..HEADER_LEN + shard_bytes];
|
||
|
||
// Route by index space: speed-test probe filler (FLAG_PROBE in user_flags) reassembles in
|
||
// its own window so its indexes never interact with the video loss window — a probe burst
|
||
// can neither advance the video anchor nor be dropped as stale against it (and its aged-out
|
||
// frames never count as `frames_dropped`, which would fire video loss recovery).
|
||
let is_probe = hdr.user_flags & (FLAG_PROBE as u32) != 0;
|
||
let win = if is_probe { probe } else { video };
|
||
win.advance_window(
|
||
hdr.frame_index,
|
||
hdr.pts_ns,
|
||
stats,
|
||
!is_probe,
|
||
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
|
||
// frame that completed early via the all-originals-present fast path — or abandoned by
|
||
// the loss window) and for frames that have fallen out of the loss window entirely.
|
||
if let Some(reconstructed) = win.completed.get_mut(&hdr.frame_index) {
|
||
// A data shard the parity reconstruct already restored (and counted into
|
||
// `fec_recovered_shards`) was late, not lost: count the arrival so the loss windows
|
||
// net it out (`recovered - late`), or plain reordering reads as packet loss and
|
||
// spooks adaptive FEC + the bitrate controller. Removing the match keeps it exact —
|
||
// wire duplicates of delivered shards match nothing, recovery shards are never in
|
||
// the list. No probe/video split: `fec_recovered_shards` counts both windows.
|
||
if shard_index < data_shards {
|
||
let fw = block_idx as u32 * lim.max_data_shards as u32 + shard_index as u32;
|
||
if let Some(pos) = reconstructed.iter().position(|&s| s == fw) {
|
||
reconstructed.swap_remove(pos);
|
||
StatsCounters::add(&stats.fec_late_shards, 1);
|
||
}
|
||
}
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
if win.is_stale(hdr.frame_index, hdr.pts_ns) {
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
|
||
// First packet of a frame allocates its whole (zeroed) buffer, budget-gated; later
|
||
// packets must agree with its geometry.
|
||
let buf_len = total_data * shard_bytes;
|
||
let frame = match win.frames.entry(hdr.frame_index) {
|
||
std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
|
||
std::collections::hash_map::Entry::Vacant(e) => {
|
||
if *in_flight_bytes + buf_len > IN_FLIGHT_BUF_FACTOR * lim.max_frame_bytes {
|
||
// Budget exhausted (several max-size frames all partially in flight) — a
|
||
// stream this bites is already deep in loss; dropping the packet is strictly
|
||
// milder than what the loss window would do to the frame moments later.
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
*in_flight_bytes += buf_len;
|
||
e.insert(FrameBuf {
|
||
frame_bytes,
|
||
block_count,
|
||
pts_ns: hdr.pts_ns,
|
||
user_flags: hdr.user_flags,
|
||
buf: vec![0; buf_len],
|
||
blocks: HashMap::new(),
|
||
blocks_ok: 0,
|
||
})
|
||
}
|
||
};
|
||
if frame.block_count != block_count || frame.frame_bytes != frame_bytes {
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
let FrameBuf {
|
||
buf,
|
||
blocks,
|
||
blocks_ok,
|
||
..
|
||
} = frame;
|
||
|
||
// First packet of a block sizes its state; `data_shards` is already pinned by the
|
||
// derived geometry above, but `recovery_shards` is per-block wire input (adaptive FEC
|
||
// varies it per frame) — later packets must match the block's first.
|
||
let block = blocks.entry(hdr.block_index).or_insert_with(|| BlockState {
|
||
data_shards,
|
||
recovery_shards,
|
||
have_data: vec![false; data_shards],
|
||
data_received: 0,
|
||
recovery: vec![None; recovery_shards],
|
||
recovery_received: 0,
|
||
done: false,
|
||
reconstructed: false,
|
||
});
|
||
if block.recovery_shards != recovery_shards {
|
||
drop(stats);
|
||
return Ok(None);
|
||
}
|
||
if block.done {
|
||
// A data shard the parity reconstruct already restored (`!have_data`) was late, not
|
||
// lost — net it out of the `fec_recovered_shards` it was counted into (see the
|
||
// completed-frame twin above; this arm covers multi-block frames whose other blocks
|
||
// are still in flight). `have_data == true` = wire duplicate; a failed reconstruct
|
||
// (`!reconstructed`) never counted its missing shards, so neither do we.
|
||
if block.reconstructed
|
||
&& shard_index < block.data_shards
|
||
&& !block.have_data[shard_index]
|
||
{
|
||
block.have_data[shard_index] = true; // it HAS arrived now — dedups a re-dup
|
||
StatsCounters::add(&stats.fec_late_shards, 1);
|
||
}
|
||
return Ok(None);
|
||
}
|
||
|
||
if shard_index < data_shards {
|
||
// A data shard lands at its final AU offset — the only copy its bytes ever make
|
||
// past decrypt.
|
||
if !block.have_data[shard_index] {
|
||
let off = (block_idx * lim.max_data_shards + shard_index) * shard_bytes;
|
||
buf[off..off + shard_bytes].copy_from_slice(body);
|
||
block.have_data[shard_index] = true;
|
||
block.data_received += 1;
|
||
}
|
||
} else {
|
||
let slot = shard_index - data_shards;
|
||
if block.recovery[slot].is_none() {
|
||
let mut rb = recovery_pool.pop().unwrap_or_default();
|
||
rb.clear();
|
||
rb.extend_from_slice(body);
|
||
block.recovery[slot] = Some(rb);
|
||
block.recovery_received += 1;
|
||
}
|
||
}
|
||
|
||
// Reconstruct as soon as we hold enough shards.
|
||
if block.data_received + block.recovery_received >= block.data_shards {
|
||
let missing = block.data_shards - block.data_received;
|
||
let outcome = if missing == 0 {
|
||
Ok(()) // every original arrived — its bytes are already in place
|
||
} else {
|
||
let base = block_idx * lim.max_data_shards * shard_bytes;
|
||
let region = &mut buf[base..base + block.data_shards * shard_bytes];
|
||
let mut slots: Vec<&mut [u8]> = region.chunks_mut(shard_bytes).collect();
|
||
let parity: Vec<(usize, &[u8])> = block
|
||
.recovery
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(j, s)| s.as_deref().map(|b| (j, b)))
|
||
.collect();
|
||
coder.reconstruct_into(block.recovery_shards, &mut slots, &block.have_data, &parity)
|
||
};
|
||
// The parity buffers are spent either way — reclaim them for the next block.
|
||
for slot in block.recovery.iter_mut() {
|
||
if let Some(rb) = slot.take() {
|
||
if recovery_pool.len() < RECOVERY_POOL_MAX {
|
||
recovery_pool.push(rb);
|
||
}
|
||
}
|
||
}
|
||
block.done = true;
|
||
match outcome {
|
||
Ok(()) => {
|
||
// With in-order delivery `missing` is exactly the block's lost shards; under
|
||
// reordering the early trigger also "recovers" shards that are merely still
|
||
// in flight — their later arrival counts `fec_late_shards` (both arms above)
|
||
// so loss estimators can net the two (`window_loss_ppm`).
|
||
block.reconstructed = missing > 0;
|
||
StatsCounters::add(&stats.fec_recovered_shards, missing as u64);
|
||
*blocks_ok += 1;
|
||
}
|
||
Err(_) => {
|
||
// Corrupt/incompatible shards that slipped past the header checks: discard
|
||
// this block (done, but never counted ok — the frame can't complete and ages
|
||
// out) and keep the session alive; the client recovers at the next
|
||
// keyframe/RFI.
|
||
StatsCounters::add(&stats.packets_dropped, 1);
|
||
return Ok(None);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Whole frame ready?
|
||
if *blocks_ok == block_count {
|
||
let mut done = win.frames.remove(&hdr.frame_index).unwrap();
|
||
win.completed.insert(
|
||
hdr.frame_index,
|
||
reconstructed_shards(&done.blocks, lim.max_data_shards),
|
||
);
|
||
*in_flight_bytes -= done.buf.len();
|
||
done.buf.truncate(done.frame_bytes); // trim trailing-shard zero padding
|
||
return Ok(Some(Frame {
|
||
data: done.buf,
|
||
frame_index: hdr.frame_index,
|
||
pts_ns: done.pts_ns,
|
||
flags: done.user_flags,
|
||
complete: true,
|
||
}));
|
||
}
|
||
Ok(None)
|
||
}
|
||
|
||
/// Drop all in-flight state — every partially-assembled frame and the completed/abandoned
|
||
/// index memory, in both index spaces — 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 anchors (`newest_frame`)
|
||
/// point into the discarded past.
|
||
pub fn reset(&mut self) {
|
||
self.video = ReassemblyWindow::default();
|
||
self.probe = ReassemblyWindow::default();
|
||
// The dropped frames' buffers (and their parity bufs) go back to the allocator, not the
|
||
// pool — a flush is the rare path. The budget resets with them.
|
||
self.in_flight_bytes = 0;
|
||
}
|
||
}
|
||
|
||
/// The data shards of a terminating frame that only exist because parity restored them
|
||
/// (`reconstructed` blocks' still-absent originals), as frame-wide indexes
|
||
/// (`block × max_data_shards + shard`) for the [`ReassemblyWindow::completed`] late-shard
|
||
/// memory. Empty (no allocation) for the overwhelmingly common clean frame.
|
||
fn reconstructed_shards(blocks: &HashMap<u16, BlockState>, max_data_shards: usize) -> Vec<u32> {
|
||
let mut v = Vec::new();
|
||
for (&bi, b) in blocks {
|
||
if b.reconstructed {
|
||
for (i, have) in b.have_data.iter().enumerate() {
|
||
if !have {
|
||
v.push(bi as u32 * max_data_shards as u32 + i as u32);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
v
|
||
}
|
||
|
||
impl ReassemblyWindow {
|
||
/// 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 — for the
|
||
/// video window (`count_drops`) counting them dropped, which is what drives the client's
|
||
/// recovery-keyframe request — and prune the completed-index memory to [`REORDER_WINDOW`].
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn advance_window(
|
||
&mut self,
|
||
frame_index: u32,
|
||
pts_ns: u64,
|
||
stats: &StatsCounters,
|
||
count_drops: bool,
|
||
recovery_pool: &mut Vec<Vec<u8>>,
|
||
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<Frame>>,
|
||
) {
|
||
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;
|
||
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) <= 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. Blocks that reconstructed
|
||
// before the frame died still counted `fec_recovered_shards`, so their restored
|
||
// shards join the late-shard memory exactly like an emitted frame's.
|
||
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() {
|
||
if recovery_pool.len() < RECOVERY_POOL_MAX {
|
||
recovery_pool.push(rb);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
keep
|
||
});
|
||
let pruned = before - self.frames.len();
|
||
if pruned > 0 && count_drops {
|
||
StatsCounters::add(&stats.frames_dropped, pruned as u64);
|
||
}
|
||
self.completed
|
||
.retain(|&idx, _| newest.wrapping_sub(idx) <= REORDER_WINDOW);
|
||
}
|
||
|
||
/// 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`](Self::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,
|
||
}
|
||
}
|
||
}
|