use std::{
io::ErrorKind,
net::{SocketAddr, UdpSocket},
path::Path,
thread::sleep,
time::Duration,
};
use s0_cc::{AntiAmp, MAX_TACK_US};
use s0_cdc::{Rd, SliceWr as Wr};
use s0_const::{
ACK_BURST_COUNT, DEFAULT_PKT, DIR_HDR_LEN, DIR_MAGIC, FIN_MAGIC, MAGIC_LEN, MAX_BLOCK,
MAX_DATAGRAM, Role, WINDOW_BLOCKS,
};
use s0_wire::{Ack, AckOut, Kind, SessionHeader, StreamDataHeader as DataHeader, encode_ack};
use crate::{
AckGen, Config, Outcome, Progress, Receiver, UdpError,
bench::{
IDLE_US,
manifest::parse_manifest,
manifest_net::ManifestRx,
probe::handle_probe_observed,
sink::Sink,
sink_basis::{Dec, INIT_DEC_MBPS},
stats::{Stats, done_bytes},
},
bind_std_udp, now_us, send_udp, wait_until,
};
pub fn recv(listen: SocketAddr, out: &Path, show: bool) -> Result<Stats, UdpError> {
recv_opt(listen, out, show, false)
}
pub fn recv_opt(
listen: SocketAddr,
out: &Path,
show: bool,
delete: bool,
) -> Result<Stats, UdpError> {
let sock = bind_std_udp(listen, false)?;
recv_on_sock_opt(&sock, out, show, delete)
}
#[derive(Clone, Copy, Default)]
struct RxTs {
pn: Option<u32>,
ts: u32,
rx_us: u64,
}
impl RxTs {
const fn on_data(&mut self, pn: u32, ts: u32, rx_us: u64) {
if match self.pn {
Some(prev) => pn.wrapping_sub(prev) < (1 << 31),
None => true,
} {
self.pn = Some(pn);
self.ts = ts;
self.rx_us = rx_us;
}
}
const fn delay_units(&self, now_us: u64) -> u16 {
if self.pn.is_none() || self.ts == 0 {
return 0;
}
let d = now_us.saturating_sub(self.rx_us) / Ack::DELAY_US;
if d > u16::MAX as u64 {
u16::MAX
} else {
d as u16
}
}
const fn clear_echo(&mut self) {
self.ts = 0;
}
}
const ACK_RTT_US: u64 = 16_000;
struct AckCtx {
ackgen: AckGen,
rx_ts: RxTs,
amp: AntiAmp<SocketAddr>,
rtt_us: u64,
dec_mbps: u16,
buf: [u8; MAX_DATAGRAM],
}
impl AckCtx {
fn new(first_us: u64) -> Self {
let mut ackgen = AckGen::new();
ackgen.mark_sent(first_us);
Self {
ackgen,
rx_ts: RxTs::default(),
amp: AntiAmp::default(),
rtt_us: ACK_RTT_US,
dec_mbps: INIT_DEC_MBPS,
buf: [0u8; MAX_DATAGRAM],
}
}
fn send_guarded(&mut self, sock: &UdpSocket, pkt: &[u8], peer: SocketAddr) {
if self.amp.can_send(pkt.len()) && send_udp(sock, pkt, peer).is_ok() {
self.amp.on_send(pkt.len());
}
}
#[inline]
fn on_data_pkt(&mut self, h: &DataHeader, rx_us: u64) -> bool {
self.ackgen.on_packet(h.pn);
self.rx_ts.on_data(h.pn, h.ts, rx_us);
h.pn == 0 || h.idx + 1 == h.cnt
}
fn send_ack(
&mut self,
sock: &UdpSocket,
peer: SocketAddr,
recv: &Receiver,
) -> Result<Option<u64>, UdpError> {
let now = now_us();
if !self.ackgen.should_send(now, self.rtt_us) {
return Ok(None);
}
self.send_ack_at(sock, peer, recv, now, 1).map(Some)
}
fn send_ack_burst(
&mut self,
sock: &UdpSocket,
peer: SocketAddr,
recv: &Receiver,
copies: usize,
) -> Result<u64, UdpError> {
self.send_ack_at(sock, peer, recv, now_us(), copies)
}
fn send_ack_at(
&mut self,
sock: &UdpSocket,
peer: SocketAddr,
recv: &Receiver,
now: u64,
copies: usize,
) -> Result<u64, UdpError> {
let (largest, flag, ranges) = self
.ackgen
.build()
.map_or((0, 0, &[][..]), |v| (v.largest, v.flag, v.ranges));
let (blocks, nb) = recv.collect_ack_blocks();
let hdr = SessionHeader {
kind: Kind::Ack,
flag: 0,
x: 0,
sid: 1,
pn: 0,
};
let mut w = Wr::new(&mut self.buf);
encode_ack(
&mut w,
&hdr,
&AckOut {
largest,
delay: self.rx_ts.delay_units(now),
ts_echo: self.rx_ts.ts,
rx_pkts: 0,
credit: recv.credit(),
done_lo: recv.done_lo(),
dec: self.dec_mbps,
flag,
ranges,
blocks: &blocks[..nb],
},
)?;
let n = w.pos();
let mut sent = false;
for _ in 0..copies {
if !self.amp.can_send(n) {
break;
}
send_udp(sock, &self.buf[..n], peer)?;
self.amp.on_send(n);
sent = true;
}
if sent {
self.ackgen.mark_sent(now);
self.rx_ts.clear_echo();
}
Ok(now)
}
}
fn parse_dir_chunk<'a>(
pkt: &'a [u8],
dir_le: &[u8; MAGIC_LEN],
) -> Option<(usize, usize, &'a [u8])> {
if pkt.len() < DIR_HDR_LEN || pkt[..MAGIC_LEN] != *dir_le {
return None;
}
let mut rd = Rd::new(&pkt[MAGIC_LEN..DIR_HDR_LEN]);
let total_len = rd.u32().ok()? as usize;
let idx = rd.u32().ok()? as usize;
let payload = &pkt[DIR_HDR_LEN..];
ManifestRx::valid_chunk(total_len, idx, payload.len()).then_some((total_len, idx, payload))
}
fn recv_on_sock_opt(
sock: &UdpSocket,
out: &Path,
show: bool,
delete: bool,
) -> Result<Stats, UdpError> {
let mut recv =
Receiver::new(Config::new(WINDOW_BLOCKS as usize, MAX_BLOCK, DEFAULT_PKT)?.with_zstd())?;
let first_us = now_us();
let mut ctx = AckCtx::new(first_us);
let mut manifest: Option<ManifestRx> = None;
let mut dir_ack_done: Option<Vec<u8>> = None;
let mut total: Option<u64> = None;
let mut logical_bytes = 0u64;
let mut dec = Dec::default();
let mut bar = show.then(|| Progress::new(Role::Recv, 0));
let mut buf = [0u8; MAX_DATAGRAM];
let fin_le = FIN_MAGIC.to_le_bytes();
let dir_le = DIR_MAGIC.to_le_bytes();
let (mut packets, mut last_data, mut last_ack_us) = (0u64, first_us, first_us);
let (mut first_dir_ack_us, mut dir_ack_wait_us) = (0u64, 100_000u64);
let (mut unacked, mut complete_at) = (false, None);
loop {
match sock.recv_from(&mut buf) {
Ok((n, addr)) => {
if let Some(ack) = handle_probe_observed(&buf[..n], addr) {
let _ = send_udp(sock, &ack, addr);
continue;
}
if n == MAGIC_LEN && buf[..MAGIC_LEN] == fin_le {
if ctx.amp.primary() == Some(addr) && (complete_at.is_some() || packets == 0) {
break;
}
continue;
}
if let Some((total_len, idx, payload)) = parse_dir_chunk(&buf[..n], &dir_le) {
if !ctx.amp.on_recv(addr, n) {
continue;
}
last_data = now_us();
if let Some(ack) = &dir_ack_done {
ctx.send_guarded(sock, ack, addr);
last_ack_us = last_data;
if total == Some(0) {
complete_at = Some(last_data);
}
continue;
}
let m = match &mut manifest {
Some(m) if m.data.len() == total_len => m,
_ => {
let Some(rx) = ManifestRx::new(total_len) else {
continue;
};
manifest.insert(rx)
}
};
m.on_chunk(idx, payload);
let mut polls = 0u8;
while m.rem > 0 {
match sock.recv_from(&mut buf) {
Ok((n2, addr2)) => {
polls = 0;
if let Some((t2, i2, p2)) = parse_dir_chunk(&buf[..n2], &dir_le)
&& t2 == total_len
&& ctx.amp.on_recv(addr2, n2)
{
last_data = now_us();
m.on_chunk(i2, p2);
}
}
Err(_) if polls < 100 => {
polls += 1;
wait_until(now_us() + 50);
}
Err(_) => break,
}
}
if let Some(m) = manifest.take_if(|m| m.rem == 0) {
ctx.send_guarded(sock, &m.ack, addr);
let (entries, l_total, mf_del) = parse_manifest(&m.data)?;
let (mut sk, delta_total, final_ack) =
Sink::prepare(out, entries, &m.ack, mf_del || delete)?;
for _ in 0..ACK_BURST_COUNT {
ctx.send_guarded(sock, &final_ack, addr);
}
last_ack_us = now_us();
first_dir_ack_us = last_ack_us;
sk.stage_copies()?;
if delta_total == 0 {
sk.finalize()?;
complete_at = Some(now_us());
} else {
dec.bind(sk, delta_total);
}
logical_bytes = l_total;
total = Some(delta_total);
if let Some(b) = bar.as_mut() {
b.set_total(delta_total);
}
dir_ack_done = Some(final_ack);
} else if let Some(m) = &manifest {
ctx.send_guarded(sock, &m.ack, addr);
last_ack_us = now_us();
}
continue;
}
let Ok((h, payload)) = DataHeader::parse(&buf[..n]) else {
continue;
};
let Some(t) = total else {
continue;
};
if !dec.is_active() && complete_at.is_none() {
continue;
}
if !ctx.amp.on_recv(addr, n) {
continue;
}
last_data = now_us();
if packets == 0 && first_dir_ack_us > 0 {
ctx.rtt_us = last_data
.saturating_sub(first_dir_ack_us)
.clamp(1_000, 30_000_000);
}
let out_res = recv.on_data(&h, payload)?;
let mut force_ack = ctx.on_data_pkt(&h, last_data);
unacked = true;
packets += 1;
if let Outcome::BlockReady { blk, .. } = out_res {
ctx.amp.mark_validated();
let mut ready = [(0u32, false); 32];
ready[0] = (blk, h.is_zstd());
let mut n_ready = 1usize;
let mut drained = 0u32;
while n_ready < ready.len() && drained < AckGen::MAX_PACKETS {
match sock.recv_from(&mut buf) {
Ok((n2, addr2)) => {
if let Some(ack) = handle_probe_observed(&buf[..n2], addr2) {
let _ = send_udp(sock, &ack, addr2);
continue;
}
let Ok((h2, payload2)) = DataHeader::parse(&buf[..n2]) else {
continue;
};
if !ctx.amp.on_recv(addr2, n2) {
continue;
}
last_data = now_us();
let res2 = recv.on_data(&h2, payload2)?;
force_ack |= ctx.on_data_pkt(&h2, last_data);
unacked = true;
packets += 1;
drained += 1;
if let Outcome::BlockReady { blk: blk2, .. } = res2 {
ready[n_ready] = (blk2, h2.is_zstd());
n_ready += 1;
}
}
Err(_) => break,
}
}
if n_ready > 1 {
ready[..n_ready].sort_unstable_by_key(|&(b, _)| b);
}
for &(b, zstd) in &ready[..n_ready] {
dec.write(&mut recv, b, zstd)?;
}
ctx.dec_mbps = dec.dec_mbps();
}
if let Some(b) = bar.as_mut() {
b.tick(done_bytes(recv.done_lo(), t), 0, packets);
}
if u64::from(recv.done_lo()) >= t.div_ceil(MAX_BLOCK as u64) {
let done_now = ctx.send_ack_burst(sock, addr, &recv, 4)?;
unacked = false;
last_ack_us = done_now;
if complete_at.is_none() {
dec.finalize()?;
let _ = ctx.send_ack_burst(sock, addr, &recv, 2);
complete_at = Some(now_us());
if let Some(b) = bar.as_mut() {
b.mark_done();
}
}
last_data = now_us();
} else if unacked {
let ack_now = if force_ack {
Some(ctx.send_ack_burst(sock, addr, &recv, 1)?)
} else {
ctx.send_ack(sock, addr, &recv)?
};
if let Some(ack_t) = ack_now {
unacked = false;
last_ack_us = ack_t;
}
}
}
Err(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {
let now = now_us();
if let Some(addr) = ctx.amp.primary() {
if let Some(m) = &manifest {
if now.saturating_sub(last_data) >= 8_000 && now.saturating_sub(last_ack_us) >= 25_000 {
ctx.send_guarded(sock, &m.ack, addr);
last_ack_us = now;
}
} else if packets == 0
&& let Some(ack) = &dir_ack_done
{
if now.saturating_sub(last_ack_us) >= dir_ack_wait_us {
ctx.send_guarded(sock, ack, addr);
last_ack_us = now;
dir_ack_wait_us = (dir_ack_wait_us * 2).min(4_000_000);
}
} else if complete_at.is_some() && packets > 0 {
let tail_ack_iv = (ctx.rtt_us / 2).clamp(100_000, 2_000_000);
if now.saturating_sub(last_ack_us) >= tail_ack_iv {
last_ack_us = ctx.send_ack_burst(sock, addr, &recv, 1)?;
}
} else if packets > 0
&& total.is_some_and(|t| t > 0)
&& (unacked || now.saturating_sub(last_ack_us) >= MAX_TACK_US)
{
last_ack_us = ctx.send_ack_burst(sock, addr, &recv, 1)?;
unacked = false;
}
}
if let Some(at) = complete_at {
let linger_us = (ctx.rtt_us.saturating_mul(4)).clamp(1_000_000, 60_000_000);
if now.saturating_sub(at) > linger_us {
break;
}
} else {
let idle_us = (ctx.rtt_us.saturating_mul(16)).clamp(IDLE_US, 180_000_000);
if now.saturating_sub(last_data) > idle_us && packets > 0 {
break;
}
}
sleep(Duration::from_micros(200));
}
Err(e) => return Err(e.into()),
}
}
let delta_bytes = total.ok_or(UdpError::Corrupt)?;
if u64::from(recv.done_lo()) < delta_bytes.div_ceil(MAX_BLOCK as u64) {
return Err(UdpError::Corrupt);
}
if let Some(b) = bar.as_mut() {
b.set_total(delta_bytes);
b.finish(delta_bytes, 0, packets);
}
Ok(Stats {
bytes: logical_bytes,
us: complete_at.unwrap_or_else(now_us).saturating_sub(first_us),
packets,
retx: 0,
blocks: delta_bytes.div_ceil(MAX_BLOCK as u64) as u32,
srtt_us: 0,
rtt_min_us: 0,
})
}