feat(core): zero-copy pooled reassembly — shards land at their final AU offset
Rewrite the client Reassembler around one whole-frame buffer per frame: frame_bytes rides in every header and packetize geometry is deterministic (every non-final block is exactly max_data_per_block data shards), so a data shard's final AU offset is computable on arrival — copy it there once, straight from the decrypt ring. New ErasureCoder::reconstruct_into decodes ONLY the missing shards directly into the frame buffer's holes (gf16 native; gf8 legacy shim); received recovery shards ride pooled shard-sized buffers. The completed buffer IS Frame::data. Deletes the per-shard to_vec + per-block concat + final AU concat (~178k allocs and a double copy of every byte per second at 2 Gbps — the pump wall the 2026-07-14 sweeps measured at 98.9% of an M3 Ultra core). Reassembly now costs ~0.4 µs/packet in-stream. The eager buffer changes the hostile-header exposure, so two new firewalls: derived-geometry validation (a header lying about its data_shards/block_count vs its own frame_bytes is dropped before it can scribble across another shard's range) and an in-flight allocation budget (IN_FLIGHT_BUF_FACTOR × max_frame_bytes) so a window of tiny first-shards can't commit gigabytes. Behavior parity pinned by the existing suite (all green unchanged) plus new end-to-end roundtrips through the real Packetizer (multi-block + partial tail, loss within budget, reversed delivery, duplicates, empty frame, unrecoverable block ages out, budget enforcement). loss-harness recovery curve identical; pipeline bench: gf8/1MB +42%, gf16 neutral (host-encode dominated). Known pre-existing quirk kept as-is: reversed delivery reconstructs early (data+recovery ≥ k) and counts late-not-lost shards into fec_recovered_shards. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -2,7 +2,9 @@
|
||||
//! shards/block — this is what removes the GameStream 255-shard / ~1 Gbps wall.
|
||||
//! Shard length must be even.
|
||||
|
||||
use super::{validate_block_shape, validate_encode_shape, ErasureCoder, FecError};
|
||||
use super::{
|
||||
validate_block_shape, validate_encode_shape, validate_into_shape, ErasureCoder, FecError,
|
||||
};
|
||||
use crate::config::FecScheme;
|
||||
|
||||
pub struct Gf16Coder;
|
||||
@@ -81,4 +83,46 @@ impl ErasureCoder for Gf16Coder {
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn reconstruct_into(
|
||||
&self,
|
||||
recovery_count: usize,
|
||||
data: &mut [&mut [u8]],
|
||||
have: &[bool],
|
||||
recovery: &[(usize, &[u8])],
|
||||
) -> Result<(), FecError> {
|
||||
validate_into_shape(data, have, recovery, recovery_count)?;
|
||||
if have.iter().all(|h| *h) {
|
||||
return Ok(()); // nothing missing — no codec work, no copies
|
||||
}
|
||||
if data[0].len() % 2 != 0 {
|
||||
return Err(FecError::Config("GF(2^16) shard length must be even"));
|
||||
}
|
||||
let data_count = data.len();
|
||||
// Present originals as indexed refs (shared reborrows of the caller's slots); the decoder
|
||||
// returns the restored shards owned, so the borrows end before the write-back below.
|
||||
let original_in: Vec<(usize, &[u8])> = data
|
||||
.iter()
|
||||
.zip(have)
|
||||
.enumerate()
|
||||
.filter(|(_, (_, &h))| h)
|
||||
.map(|(i, (s, _))| (i, &**s))
|
||||
.collect();
|
||||
let restored = reed_solomon_simd::decode(
|
||||
data_count,
|
||||
recovery_count,
|
||||
original_in,
|
||||
recovery.iter().copied(),
|
||||
)
|
||||
.map_err(|_| FecError::Backend("gf16 decode"))?;
|
||||
for (i, h) in have.iter().enumerate() {
|
||||
if !*h {
|
||||
let shard = restored
|
||||
.get(&i)
|
||||
.ok_or(FecError::Backend("gf16 decode left an original missing"))?;
|
||||
data[i].copy_from_slice(shard);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,7 +4,9 @@
|
||||
//! client (unlike Vandermonde RS, whose parity is not interoperable). Hard ceiling: data +
|
||||
//! recovery ≤ 255 shards/block.
|
||||
|
||||
use super::{validate_block_shape, validate_encode_shape, ErasureCoder, FecError};
|
||||
use super::{
|
||||
validate_block_shape, validate_encode_shape, validate_into_shape, ErasureCoder, FecError,
|
||||
};
|
||||
use crate::config::FecScheme;
|
||||
use fec_rs::ReedSolomon;
|
||||
|
||||
@@ -56,6 +58,44 @@ impl ErasureCoder for Gf8Coder {
|
||||
.map_err(|_| FecError::Backend("gf8 reconstruct"))?;
|
||||
collect_originals(received, data_count)
|
||||
}
|
||||
|
||||
fn reconstruct_into(
|
||||
&self,
|
||||
recovery_count: usize,
|
||||
data: &mut [&mut [u8]],
|
||||
have: &[bool],
|
||||
recovery: &[(usize, &[u8])],
|
||||
) -> Result<(), FecError> {
|
||||
validate_into_shape(data, have, recovery, recovery_count)?;
|
||||
if have.iter().all(|h| *h) {
|
||||
return Ok(());
|
||||
}
|
||||
// Legacy-scheme shim: fec-rs reconstructs through owned `Option<Vec<u8>>` slots, so copy
|
||||
// the present shards into that shape and the recovered ones back out. Only P1/gf8
|
||||
// sessions on loss pay this — the hot gf16 path decodes straight into the caller's slots.
|
||||
let data_count = data.len();
|
||||
let mut received: Vec<Option<Vec<u8>>> = Vec::with_capacity(data_count + recovery_count);
|
||||
for (s, h) in data.iter().zip(have) {
|
||||
received.push(h.then(|| s.to_vec()));
|
||||
}
|
||||
received.resize(data_count + recovery_count, None);
|
||||
for &(j, bytes) in recovery {
|
||||
received[data_count + j] = Some(bytes.to_vec());
|
||||
}
|
||||
let rs = ReedSolomon::new(data_count, recovery_count)
|
||||
.map_err(|_| FecError::Config("invalid GF(2^8) shard counts"))?;
|
||||
rs.reconstruct_data(&mut received)
|
||||
.map_err(|_| FecError::Backend("gf8 reconstruct"))?;
|
||||
for (i, h) in have.iter().enumerate() {
|
||||
if !*h {
|
||||
let shard = received[i]
|
||||
.as_ref()
|
||||
.ok_or(FecError::Backend("reconstruction left an original missing"))?;
|
||||
data[i].copy_from_slice(shard);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_originals(
|
||||
|
||||
@@ -43,6 +43,25 @@ pub trait ErasureCoder: Send + Sync {
|
||||
recovery_count: usize,
|
||||
received: &mut [Option<Vec<u8>>],
|
||||
) -> Result<Vec<Vec<u8>>, FecError>;
|
||||
|
||||
/// Reconstruct ONLY the missing data shards of a block, writing each straight into its final
|
||||
/// slot in the caller's buffer — the receive-side half of [`encode`](Self::encode)'s ref-based
|
||||
/// contract (the reassembler's slots are slices of one contiguous frame buffer, so recovery
|
||||
/// lands at its final AU offset with no per-shard `Vec`s and no block/AU concat copies).
|
||||
///
|
||||
/// `data` holds the block's K equal-length shard slots; `have[i]` marks the slots whose bytes
|
||||
/// were received (valid codec input — a missing slot's contents are unspecified on entry).
|
||||
/// `recovery` is the received parity as `(recovery_index, bytes)` with `recovery_index <
|
||||
/// recovery_count` (the block's declared M, which the codec math needs even when not all M
|
||||
/// arrived). On success every missing slot has been filled; on error missing slots are
|
||||
/// unspecified and the caller must discard the block.
|
||||
fn reconstruct_into(
|
||||
&self,
|
||||
recovery_count: usize,
|
||||
data: &mut [&mut [u8]],
|
||||
have: &[bool],
|
||||
recovery: &[(usize, &[u8])],
|
||||
) -> Result<(), FecError>;
|
||||
}
|
||||
|
||||
/// Construct the coder for a scheme.
|
||||
@@ -80,6 +99,43 @@ pub(crate) fn validate_block_shape(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validate the shape [`ErasureCoder::reconstruct_into`] promises: `have` matches `data`, one
|
||||
/// shard length across data slots and recovery shards, recovery indices within the declared M,
|
||||
/// and enough shards present to reconstruct at all. Both backends call this first.
|
||||
pub(crate) fn validate_into_shape(
|
||||
data: &[&mut [u8]],
|
||||
have: &[bool],
|
||||
recovery: &[(usize, &[u8])],
|
||||
recovery_count: usize,
|
||||
) -> Result<(), FecError> {
|
||||
if data.is_empty() {
|
||||
return Err(FecError::Config("no data shards"));
|
||||
}
|
||||
if have.len() != data.len() {
|
||||
return Err(FecError::Config("have length must equal data length"));
|
||||
}
|
||||
let len = data[0].len();
|
||||
if data.iter().any(|s| s.len() != len) {
|
||||
return Err(FecError::Config("shards in a block must be equal length"));
|
||||
}
|
||||
for &(j, bytes) in recovery {
|
||||
if j >= recovery_count {
|
||||
return Err(FecError::Config("recovery index out of range"));
|
||||
}
|
||||
if bytes.len() != len {
|
||||
return Err(FecError::Config("shards in a block must be equal length"));
|
||||
}
|
||||
}
|
||||
let present = have.iter().filter(|h| **h).count();
|
||||
if present + recovery.len() < data.len() {
|
||||
return Err(FecError::TooFewShards {
|
||||
have: present + recovery.len(),
|
||||
need: data.len(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validate `encode` inputs: at least one data shard, all of equal length.
|
||||
pub(crate) fn validate_encode_shape(data: &[&[u8]]) -> Result<(), FecError> {
|
||||
let first = data
|
||||
@@ -117,6 +173,93 @@ mod tests {
|
||||
assert_eq!(restored, data);
|
||||
}
|
||||
|
||||
/// Round-trip through `reconstruct_into`: encode, zero out `lose_data` slots in a contiguous
|
||||
/// buffer (the reassembler's frame-buffer shape), drop `lose_recovery` parity shards, and
|
||||
/// assert the missing slots are restored in place while the present ones are untouched.
|
||||
fn roundtrip_into(
|
||||
coder: &dyn ErasureCoder,
|
||||
k: usize,
|
||||
m: usize,
|
||||
shard_len: usize,
|
||||
lose_data: &[usize],
|
||||
lose_recovery: &[usize],
|
||||
) {
|
||||
let src: Vec<Vec<u8>> = (0..k)
|
||||
.map(|i| (0..shard_len).map(|b| (i * 31 + b * 7) as u8).collect())
|
||||
.collect();
|
||||
let refs: Vec<&[u8]> = src.iter().map(|s| s.as_slice()).collect();
|
||||
let parity = coder.encode(&refs, m).unwrap();
|
||||
|
||||
let mut buf = vec![0u8; k * shard_len];
|
||||
let mut have = vec![true; k];
|
||||
for (i, s) in src.iter().enumerate() {
|
||||
if lose_data.contains(&i) {
|
||||
have[i] = false; // slot stays zeroed — codec must fill it
|
||||
} else {
|
||||
buf[i * shard_len..(i + 1) * shard_len].copy_from_slice(s);
|
||||
}
|
||||
}
|
||||
let recovery: Vec<(usize, &[u8])> = parity
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(j, _)| !lose_recovery.contains(j))
|
||||
.map(|(j, p)| (j, p.as_slice()))
|
||||
.collect();
|
||||
|
||||
let mut slots: Vec<&mut [u8]> = buf.chunks_mut(shard_len).collect();
|
||||
coder
|
||||
.reconstruct_into(m, &mut slots, &have, &recovery)
|
||||
.unwrap();
|
||||
for (i, s) in src.iter().enumerate() {
|
||||
assert_eq!(
|
||||
&buf[i * shard_len..(i + 1) * shard_len],
|
||||
s.as_slice(),
|
||||
"shard {i}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gf16_reconstruct_into_fills_only_the_holes() {
|
||||
roundtrip_into(&Gf16Coder, 16, 4, 256, &[1, 9], &[3]);
|
||||
roundtrip_into(&Gf16Coder, 4, 2, 16, &[0, 3], &[]);
|
||||
roundtrip_into(&Gf16Coder, 4, 2, 16, &[], &[0, 1]); // nothing missing, no parity needed
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gf8_reconstruct_into_fills_only_the_holes() {
|
||||
roundtrip_into(&Gf8Coder, 16, 4, 256, &[0, 7], &[1]);
|
||||
roundtrip_into(&Gf8Coder, 4, 2, 16, &[2], &[1]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconstruct_into_rejects_bad_shapes() {
|
||||
let mut buf = [0u8; 4 * 8];
|
||||
// Too few shards: 2 of 4 data present, no recovery.
|
||||
let mut slots: Vec<&mut [u8]> = buf.chunks_mut(8).collect();
|
||||
let have = [true, true, false, false];
|
||||
assert!(Gf16Coder
|
||||
.reconstruct_into(2, &mut slots, &have, &[])
|
||||
.is_err());
|
||||
// Recovery index out of the declared range.
|
||||
let parity = [0u8; 8];
|
||||
let mut slots: Vec<&mut [u8]> = buf.chunks_mut(8).collect();
|
||||
assert!(Gf16Coder
|
||||
.reconstruct_into(2, &mut slots, &have, &[(2, &parity), (3, &parity)])
|
||||
.is_err());
|
||||
// Mismatched recovery shard length.
|
||||
let short = [0u8; 6];
|
||||
let mut slots: Vec<&mut [u8]> = buf.chunks_mut(8).collect();
|
||||
assert!(Gf8Coder
|
||||
.reconstruct_into(2, &mut slots, &have, &[(0, &short), (1, &parity)])
|
||||
.is_err());
|
||||
// `have` length disagreeing with `data`.
|
||||
let mut slots: Vec<&mut [u8]> = buf.chunks_mut(8).collect();
|
||||
assert!(Gf8Coder
|
||||
.reconstruct_into(2, &mut slots, &[true; 3], &[(0, &parity)])
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gf8_recovers_within_budget() {
|
||||
// 16 data + 4 recovery; lose 2 data + 2 recovery (== budget).
|
||||
|
||||
@@ -20,7 +20,7 @@ use crate::error::{PunktfunkError, Result};
|
||||
use crate::fec::ErasureCoder;
|
||||
use crate::session::Frame;
|
||||
use crate::stats::StatsCounters;
|
||||
use std::collections::{BTreeMap, HashMap, HashSet};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
|
||||
|
||||
/// Identifies a punktfunk video packet (vs. an input datagram, see [`crate::input`]).
|
||||
@@ -331,13 +331,23 @@ impl Packetizer {
|
||||
// Client side: reassembly + FEC recovery
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct BlockBuf {
|
||||
/// 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,
|
||||
shard_bytes: usize,
|
||||
/// Length `data_shards + recovery_shards`; `Some` = received.
|
||||
shards: Vec<Option<Vec<u8>>>,
|
||||
received: 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,
|
||||
}
|
||||
|
||||
@@ -346,9 +356,16 @@ struct FrameBuf {
|
||||
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>>,
|
||||
/// 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*
|
||||
@@ -401,6 +418,18 @@ struct ReassemblyWindow {
|
||||
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 {
|
||||
@@ -414,6 +443,12 @@ pub struct Reassembler {
|
||||
/// 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 {
|
||||
@@ -422,6 +457,8 @@ impl Reassembler {
|
||||
limits,
|
||||
video: ReassemblyWindow::default(),
|
||||
probe: ReassemblyWindow::default(),
|
||||
recovery_pool: Vec::new(),
|
||||
in_flight_bytes: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -449,7 +486,16 @@ impl Reassembler {
|
||||
}
|
||||
};
|
||||
|
||||
let lim = self.limits;
|
||||
// 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,
|
||||
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;
|
||||
@@ -480,19 +526,42 @@ impl Reassembler {
|
||||
drop(stats);
|
||||
return Ok(None);
|
||||
}
|
||||
let payload = pkt[HEADER_LEN..HEADER_LEN + shard_bytes].to_vec();
|
||||
// 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 {
|
||||
&mut self.probe
|
||||
} else {
|
||||
&mut self.video
|
||||
};
|
||||
win.advance_window(hdr.frame_index, hdr.pts_ns, stats, !is_probe);
|
||||
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,
|
||||
);
|
||||
|
||||
// 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
|
||||
@@ -502,108 +571,135 @@ impl Reassembler {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// First packet of a frame establishes its geometry; later packets must agree.
|
||||
let frame = win
|
||||
.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(),
|
||||
});
|
||||
// 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;
|
||||
|
||||
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
|
||||
{
|
||||
// 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,
|
||||
});
|
||||
if block.recovery_shards != recovery_shards {
|
||||
drop(stats);
|
||||
return Ok(None);
|
||||
}
|
||||
if block.done {
|
||||
return Ok(None); // late/duplicate shard for a resolved block — silent, like before
|
||||
}
|
||||
|
||||
if block.shards[shard_index].is_none() {
|
||||
block.shards[shard_index] = Some(payload);
|
||||
block.received += 1;
|
||||
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.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,
|
||||
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(()) => {
|
||||
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 (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;
|
||||
// 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);
|
||||
}
|
||||
};
|
||||
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 = win.frames.remove(&hdr.frame_index).unwrap();
|
||||
if *blocks_ok == block_count {
|
||||
let mut done = win.frames.remove(&hdr.frame_index).unwrap();
|
||||
win.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
|
||||
*in_flight_bytes -= done.buf.len();
|
||||
done.buf.truncate(done.frame_bytes); // trim trailing-shard zero padding
|
||||
return Ok(Some(Frame {
|
||||
data,
|
||||
data: done.buf,
|
||||
frame_index: hdr.frame_index,
|
||||
pts_ns: frame.pts_ns,
|
||||
flags: frame.user_flags,
|
||||
pts_ns: done.pts_ns,
|
||||
flags: done.user_flags,
|
||||
}));
|
||||
}
|
||||
Ok(None)
|
||||
@@ -618,6 +714,9 @@ impl Reassembler {
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -632,6 +731,8 @@ impl ReassemblyWindow {
|
||||
pts_ns: u64,
|
||||
stats: &StatsCounters,
|
||||
count_drops: bool,
|
||||
recovery_pool: &mut Vec<Vec<u8>>,
|
||||
in_flight_bytes: &mut usize,
|
||||
) {
|
||||
let (newest, newest_pts) = match self.newest_frame {
|
||||
// `frame_index` is newer iff it's within the forward half of the index space.
|
||||
@@ -650,6 +751,17 @@ impl ReassemblyWindow {
|
||||
// `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);
|
||||
// Release its buffer budget and reclaim its parity bufs for the pool.
|
||||
*in_flight_bytes -= f.buf.len();
|
||||
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
|
||||
});
|
||||
@@ -957,6 +1069,195 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// Build a host config for the end-to-end roundtrips: 16-byte shards, 4-data-shard blocks.
|
||||
fn e2e_config(scheme: FecScheme, fec_percent: u8) -> Config {
|
||||
use crate::config::{FecConfig, ProtocolPhase, Role};
|
||||
Config {
|
||||
role: Role::Host,
|
||||
phase: ProtocolPhase::P2Punktfunk,
|
||||
fec: FecConfig {
|
||||
scheme,
|
||||
fec_percent,
|
||||
max_data_per_block: 4,
|
||||
},
|
||||
shard_payload: 16,
|
||||
max_frame_bytes: 4096,
|
||||
encrypt: false,
|
||||
key: [0u8; 16],
|
||||
salt: [0u8; 4],
|
||||
loopback_drop_period: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Packetize a synthetic AU, deliver a mangled subset (losses within the FEC budget,
|
||||
/// optionally reversed, with a duplicate), and assert the reassembled AU is byte-identical
|
||||
/// to the source — the shards landed straight in the frame buffer at the right offsets and
|
||||
/// FEC filled the holes.
|
||||
///
|
||||
/// `fec_recovered_shards` accounting: with in-order delivery it equals the kill count
|
||||
/// exactly. With reversed delivery parity arrives first, so the `data + recovery ≥ k`
|
||||
/// trigger reconstructs EARLY and restores late-not-lost shards too — longstanding behavior
|
||||
/// (the trigger predates this rewrite); assert `≥` there.
|
||||
fn e2e_roundtrip(
|
||||
scheme: FecScheme,
|
||||
frame_len: usize,
|
||||
fec_percent: u8,
|
||||
kill: &[usize],
|
||||
reverse: bool,
|
||||
) {
|
||||
let cfg = e2e_config(scheme, fec_percent);
|
||||
let coder = coder_for(scheme);
|
||||
let mut pk = Packetizer::new(&cfg);
|
||||
let src: Vec<u8> = (0..frame_len).map(|i| (i * 131 + 7) as u8).collect();
|
||||
let pkts = pk.packetize(&src, 12345, 0, coder.as_ref()).unwrap();
|
||||
|
||||
let mut delivery: Vec<Vec<u8>> = pkts
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(i, _)| !kill.contains(i))
|
||||
.map(|(_, p)| p.clone())
|
||||
.collect();
|
||||
if reverse {
|
||||
delivery.reverse(); // recovery shards (and the tail) arrive first
|
||||
}
|
||||
if let Some(dup) = delivery.first().cloned() {
|
||||
delivery.push(dup); // a duplicate must be ignored, not double-counted
|
||||
}
|
||||
|
||||
let mut r = Reassembler::new(ReassemblerLimits::from_config(&cfg));
|
||||
let stats = StatsCounters::default();
|
||||
let mut got = None;
|
||||
for p in &delivery {
|
||||
if let Some(f) = r.push(p, coder.as_ref(), &stats).unwrap() {
|
||||
assert!(got.is_none(), "frame must complete exactly once");
|
||||
got = Some(f);
|
||||
}
|
||||
}
|
||||
let f = got.expect("frame must complete within the FEC budget");
|
||||
assert_eq!(f.data, src, "reassembled AU must be byte-identical");
|
||||
assert_eq!(f.pts_ns, 12345);
|
||||
let recovered = stats.snapshot().fec_recovered_shards;
|
||||
if reverse {
|
||||
assert!(
|
||||
recovered >= kill.len() as u64,
|
||||
"early reconstruct counts more"
|
||||
);
|
||||
} else {
|
||||
assert_eq!(recovered, kill.len() as u64);
|
||||
}
|
||||
}
|
||||
|
||||
/// Multi-block frame with a partial tail shard, heavy loss, both delivery orders + dups.
|
||||
/// 100 bytes / 16 = 7 shards → blocks of (4 data + 2 rec) and (3 data + 2 rec).
|
||||
#[test]
|
||||
fn e2e_multiblock_loss_reorder_dup_gf16() {
|
||||
// Packet order: blk0 = idx 0..6 (4 data + 2 rec), blk1 = idx 6..11 (3 data + 2 rec).
|
||||
// Kill 2 data in block 0 and 1 data in block 1 — all within the 50% budget.
|
||||
e2e_roundtrip(FecScheme::Gf16, 100, 50, &[0, 2, 7], false);
|
||||
e2e_roundtrip(FecScheme::Gf16, 100, 50, &[0, 2, 7], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn e2e_multiblock_loss_reorder_dup_gf8() {
|
||||
e2e_roundtrip(FecScheme::Gf8, 100, 50, &[1, 3, 8], false);
|
||||
e2e_roundtrip(FecScheme::Gf8, 100, 50, &[1, 3, 8], true);
|
||||
}
|
||||
|
||||
/// Zero losses, in order: the pure fast path (no codec call, recovered == 0) must still
|
||||
/// emit an identical AU.
|
||||
#[test]
|
||||
fn e2e_clean_delivery_gf16() {
|
||||
e2e_roundtrip(FecScheme::Gf16, 100, 50, &[], false);
|
||||
}
|
||||
|
||||
/// An empty AU rides one zero-padded shard and reassembles to zero bytes.
|
||||
#[test]
|
||||
fn e2e_empty_frame() {
|
||||
let cfg = e2e_config(FecScheme::Gf16, 0);
|
||||
let coder = coder_for(FecScheme::Gf16);
|
||||
let mut pk = Packetizer::new(&cfg);
|
||||
let pkts = pk.packetize(&[], 7, 0, coder.as_ref()).unwrap();
|
||||
assert_eq!(pkts.len(), 1);
|
||||
let mut r = Reassembler::new(ReassemblerLimits::from_config(&cfg));
|
||||
let stats = StatsCounters::default();
|
||||
let f = r
|
||||
.push(&pkts[0], coder.as_ref(), &stats)
|
||||
.unwrap()
|
||||
.expect("empty frame completes");
|
||||
assert!(f.data.is_empty());
|
||||
}
|
||||
|
||||
/// Loss beyond the FEC budget: the frame never emits, ages out as dropped, and the
|
||||
/// unrecoverable-block path must not fire (block never gathers k shards at all).
|
||||
#[test]
|
||||
fn e2e_unrecoverable_loss_ages_out() {
|
||||
let cfg = e2e_config(FecScheme::Gf16, 50);
|
||||
let coder = coder_for(FecScheme::Gf16);
|
||||
let mut pk = Packetizer::new(&cfg);
|
||||
let src = vec![0x5Au8; 64]; // one block: 4 data + 2 recovery
|
||||
let pkts = pk.packetize(&src, 1_000, 0, coder.as_ref()).unwrap();
|
||||
let mut r = Reassembler::new(ReassemblerLimits::from_config(&cfg));
|
||||
let stats = StatsCounters::default();
|
||||
// Deliver only 3 of 6 shards (k=4): can never reconstruct.
|
||||
for p in &pkts[..3] {
|
||||
assert!(r.push(p, coder.as_ref(), &stats).unwrap().is_none());
|
||||
}
|
||||
// A newer frame past the loss window ages it out as a video drop.
|
||||
let next = pk
|
||||
.packetize(&src, 1_000 + LOSS_WINDOW_NS + 1, 0, coder.as_ref())
|
||||
.unwrap();
|
||||
let mut done = false;
|
||||
for p in &next {
|
||||
done |= r.push(p, coder.as_ref(), &stats).unwrap().is_some();
|
||||
}
|
||||
assert!(done);
|
||||
assert_eq!(stats.snapshot().frames_dropped, 1);
|
||||
}
|
||||
|
||||
/// The in-flight buffer budget: a window of tiny first-shards all declaring max-size frames
|
||||
/// stops allocating at [`IN_FLIGHT_BUF_FACTOR`] × max_frame_bytes instead of committing
|
||||
/// gigabytes (the eager whole-frame buffer's amplification defense).
|
||||
#[test]
|
||||
fn in_flight_buffer_budget_bounds_allocation() {
|
||||
let lim = limits(); // max_frame_bytes 4096, shards 16 B, ≤8 data shards × ≤4 blocks
|
||||
let mut r = Reassembler::new(lim);
|
||||
let coder = coder_for(FecScheme::Gf8);
|
||||
let stats = StatsCounters::default();
|
||||
// Largest geometry-consistent frame: 4 blocks × 8 shards × 16 B = 512 B per buffer.
|
||||
// Budget = 4 × 4096 = 16384 B → exactly 32 such frames fit; the 33rd must be refused.
|
||||
for i in 0..33u32 {
|
||||
let mut h = base_header();
|
||||
h.frame_index = i;
|
||||
h.frame_bytes = 512;
|
||||
h.block_count = 4;
|
||||
h.data_shards = 8;
|
||||
r.push(&packet(h), coder.as_ref(), &stats).unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
stats.snapshot().packets_dropped,
|
||||
1,
|
||||
"the frame past the budget is dropped, everything under it accepted"
|
||||
);
|
||||
}
|
||||
|
||||
/// A header whose (data_shards, block_count) disagree with the geometry derived from its own
|
||||
/// frame_bytes is dropped — the derived-offset invariant that lets shards land directly in
|
||||
/// the frame buffer.
|
||||
#[test]
|
||||
fn rejects_geometry_inconsistent_with_frame_bytes() {
|
||||
let mut r = Reassembler::new(limits());
|
||||
let coder = coder_for(FecScheme::Gf8);
|
||||
let stats = StatsCounters::default();
|
||||
let mut h = base_header();
|
||||
h.frame_bytes = 16; // exactly one shard…
|
||||
h.data_shards = 2; // …but claims two
|
||||
assert!(r
|
||||
.push(&packet(h), coder.as_ref(), &stats)
|
||||
.unwrap()
|
||||
.is_none());
|
||||
assert_eq!(stats.snapshot().packets_dropped, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_wrong_shard_bytes_and_oversized_frame() {
|
||||
let coder = coder_for(FecScheme::Gf8);
|
||||
|
||||
Reference in New Issue
Block a user