use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicU8, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::sync::{mpsc, watch, OwnedSemaphorePermit, Semaphore};
use crate::conn::Frame;
use crate::gen::{
self, Bundle, BundleRecord, Chunk, Credit, Durable, FileRetry, FileRoot, JobDone, JobMap,
JobOpen, JobOpenAck, ManifestEnd, Received, Status,
};
use crate::governor::{self, Class, Governor, GovernorOptions, JobSummary, Mode, Sample};
use crate::manifest::Manifest;
use crate::ranges::{from_runs, Need, RangeSet};
use crate::router::{BoxFut, ConnTx, Inbound, JobLink, LaneOpener, LaneTx};
use crate::source::{read_full_at, Source};
use crate::verify::{self, FileHasher, Outboard, GROUP};
use crate::wire::{FrameMessage, Message, Writer};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct SpaceFigures {
pub job_bytes: u64,
pub durable_bytes: u64,
pub held_bytes: u64,
pub unfinished_bytes: u64,
pub files_total: u64,
pub files_done: u64,
}
impl SpaceFigures {
pub fn of(m: &Manifest, need: &Need, cutoff: u64) -> Result<Self, SendError> {
let mut f = SpaceFigures::default();
let mut large_unfinished = 0u64;
let mut partial_durable = 0u64;
for (i, e) in m.entries.iter().enumerate() {
if e.kind != gen::ENTRY_FILE {
continue;
}
f.files_total += 1;
f.job_bytes = f.job_bytes.saturating_add(e.size);
if need.done.contains(&(i as u32)) {
f.files_done += 1;
f.durable_bytes = f.durable_bytes.saturating_add(e.size);
continue;
}
f.unfinished_bytes = f.unfinished_bytes.saturating_add(e.size);
if e.size >= cutoff {
large_unfinished = large_unfinished.saturating_add(e.size);
if let Some(d) = need.partial.get(&(i as u32)) {
large_remaining(e.size, d)?;
partial_durable = partial_durable.saturating_add(d.covered());
}
}
}
f.durable_bytes = f.durable_bytes.saturating_add(partial_durable);
f.held_bytes = need.held.min(large_unfinished).max(partial_durable);
Ok(f)
}
pub fn to_allocate(&self) -> u64 {
self.unfinished_bytes.saturating_sub(self.held_bytes)
}
}
pub type SpaceGate = Arc<dyn Fn(&SpaceFigures) -> Result<(), String> + Send + Sync>;
pub struct SendOptions {
pub kind: u8,
pub policy: u8,
pub flags: u32,
pub root: String,
pub cutoff: u64,
pub readers: usize,
pub persist: Option<PathBuf>,
pub progress: Arc<Progress>,
pub cancel: Arc<AtomicBool>,
pub bandwidth_cap: Option<u64>,
pub seq: Option<Arc<dyn crate::seq::SeqSource>>,
pub settle_max: Option<Duration>,
pub open_ack_timeout: Option<Duration>,
pub space_gate: Option<SpaceGate>,
pub receiver_silent_timeout: Option<Duration>,
}
impl SendOptions {
pub fn upload(root: &str) -> Self {
Self {
kind: gen::JOB_UPLOAD,
policy: gen::POLICY_REPLACE,
flags: 0,
root: root.into(),
cutoff: gen::LARGE_CUTOFF as u64,
readers: 8,
persist: None,
progress: Arc::default(),
cancel: Arc::default(),
bandwidth_cap: None,
seq: None,
settle_max: None,
open_ack_timeout: None,
space_gate: None,
receiver_silent_timeout: None,
}
}
}
#[derive(Debug, Default)]
pub struct Progress {
pub bytes_total: AtomicU64,
pub files_total: AtomicU64,
pub bytes_sent: AtomicU64,
pub bytes_durable: AtomicU64,
pub files_durable: AtomicU64,
pub resent_bytes: AtomicU64,
pub lanes: AtomicU8,
pub bottleneck: AtomicU8,
pub sequential: AtomicBool,
pub skipped_files: AtomicU64,
pub skipped_bytes: AtomicU64,
pub skip_recorded: AtomicBool,
pub phase: AtomicU8,
pub skip_done_bytes: AtomicU64,
pub skip_total_bytes: AtomicU64,
pub settling: AtomicBool,
pub unswept: AtomicU32,
pub attempts: AtomicU32,
pub unswept_peak: AtomicU32,
pub settle_ms: AtomicU64,
pub telemetry: Mutex<Telemetry>,
}
#[derive(Debug, Default, Clone)]
pub struct Telemetry {
pub shares: JobSummary,
pub console_line: Option<String>,
pub peer_key: Option<[u8; 32]>,
}
#[derive(Debug, Clone)]
pub struct SendReport {
pub status: u16,
pub message: Option<String>,
pub files: u32,
pub bytes: u64,
pub resent: u64,
pub max_lanes: u8,
pub bottleneck: u8,
pub sequential: bool,
}
#[derive(Debug, thiserror::Error)]
pub enum SendError {
#[error("connection lost: {0}")]
Disconnected(String),
#[error("refused ({status}): {message}")]
Refused {
status: u16,
message: String,
},
#[error("cancelled")]
Cancelled,
#[error("the console did not answer the job open within {0:?}")]
OpenTimeout(Duration),
#[error("reading the source: {0}")]
Source(#[from] std::io::Error),
#[error("protocol: {0}")]
Protocol(String),
#[error("not enough free space: {0}")]
NoRoom(String),
}
impl From<crate::Ava1Error> for SendError {
fn from(e: crate::Ava1Error) -> Self {
SendError::Disconnected(e.to_string())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Piece {
pub offset: u64,
pub len: u64,
pub send: bool,
}
pub fn pieces(
size: u64,
durable: &RangeSet,
have_cv: &dyn Fn(u64) -> bool,
chunk: u64,
) -> Vec<Piece> {
let mut out: Vec<Piece> = Vec::new();
let chunk = chunk.max(GROUP) / GROUP * GROUP;
for g in 0..verify::groups(size) {
let off = g * GROUP;
let len = (size - off).min(GROUP);
let send = !durable.covers(off, off + len);
if !send && have_cv(g) {
continue;
}
match out.last_mut() {
Some(p)
if p.send == send
&& p.offset + p.len == off
&& p.len + len <= chunk
&& (send || p.len < chunk) =>
{
p.len += len
}
_ => out.push(Piece {
offset: off,
len,
send,
}),
}
}
out
}
pub fn skip_set(m: &Manifest, durable: &Need, persist: Option<&std::path::Path>) -> Need {
let mut skip = Need::default();
for (i, e) in m.entries.iter().enumerate() {
if e.kind != gen::ENTRY_FILE {
continue;
}
let id = i as u32;
if durable.done.contains(&id) {
skip.done.insert(id);
continue;
}
if e.size < gen::LARGE_CUTOFF as u64 {
continue;
}
let Some(ranges) = durable.partial.get(&id) else {
continue;
};
let ob = persist.and_then(|dir| {
Outboard::open(&dir.join(format!("{id}.ob")), verify::groups(e.size)).ok()
});
let Some(ob) = ob else { continue };
let mut known = RangeSet::new();
for g in 0..verify::groups(e.size) {
let off = g * GROUP;
let end = off + (e.size - off).min(GROUP);
if ranges.covers(off, end) && ob.get(g).is_some() {
known.insert(off, end);
}
}
if known.is_full(e.size) {
skip.done.insert(id);
} else if known.covered() > 0 {
skip.partial.insert(id, known);
}
}
skip
}
type LaneFrames = HashMap<u16, (u64, BTreeMap<u32, u64>)>;
pub struct Window {
credit: u64,
initial: u64,
lanes: LaneFrames,
refunded: HashMap<u32, u64>,
}
impl Window {
pub fn new(credit: u64) -> Self {
Self {
credit,
initial: credit,
lanes: HashMap::new(),
refunded: HashMap::new(),
}
}
pub fn can_send(&self, lane: u16, len: u64, cap: u64) -> bool {
let inflight = self.lanes.get(&lane).map_or(0, |l| l.0);
len <= self.credit && (inflight == 0 || inflight + len <= cap)
}
pub fn sent(&mut self, lane: u16, seq: u32, len: u64) -> bool {
if len > self.credit {
return false;
}
self.credit -= len;
let l = self.lanes.entry(lane).or_default();
l.0 += len;
l.1.insert(seq, len);
true
}
pub fn received(&mut self, seq: u32) -> Option<u64> {
for l in self.lanes.values_mut() {
if let Some(len) = l.1.remove(&seq) {
l.0 -= len;
return Some(len);
}
}
self.refunded.remove(&seq);
None
}
pub fn credit(&mut self, n: u64) {
self.credit = self.credit.saturating_add(n);
}
pub fn lane_down(&mut self, lane: u16) -> Vec<u32> {
let Some((_, frames)) = self.lanes.remove(&lane) else {
return Vec::new();
};
frames
.into_iter()
.map(|(seq, len)| {
self.refunded.insert(seq, len);
seq
})
.collect()
}
pub fn release(&mut self, seq: u32) -> Option<u64> {
let len = self.refunded.remove(&seq)?;
self.credit = self.credit.saturating_add(len);
Some(len)
}
pub fn available(&self) -> u64 {
self.credit
}
pub fn outstanding(&self) -> u64 {
self.initial.saturating_sub(self.credit)
}
}
struct OutFrame {
ty: u8,
body: Arc<Vec<u8>>,
class: Class,
payload: u64,
resend: bool,
_budget: Vec<OwnedSemaphorePermit>,
taken: Option<Arc<AtomicBool>>,
}
#[derive(Default)]
struct Sched {
bundles: VecDeque<OutFrame>,
chunks: VecDeque<OutFrame>,
requeue: VecDeque<OutFrame>,
inflight: HashMap<u32, (u16, OutFrame)>,
bundles_inflight: usize,
next_seq: u32,
decision: Option<governor::Decision>,
floor: usize,
lane_rate: HashMap<u16, f64>,
lane_bytes: HashMap<u16, u64>,
credit_starved: bool,
source_starved: bool,
acked_tick: u64,
small_durable_tick: u64,
large_durable_tick: u64,
stalls: u32,
}
#[derive(Clone, Copy)]
struct Stall {
since: Instant,
grant: u64,
smallest: u64,
}
struct Shared {
sched: Mutex<Sched>,
window: Mutex<Window>,
wake_tx: watch::Sender<u64>,
chunk: AtomicU32,
bundle: AtomicU32,
bytes_budget: Arc<Semaphore>,
stall: Mutex<Option<Stall>>,
}
impl Shared {
fn wake(&self) {
let v = *self.wake_tx.borrow();
self.wake_tx.send_replace(v.wrapping_add(1));
}
}
pub(crate) enum Read {
Record {
file_id: u32,
root: [u8; 32],
data: Vec<u8>,
budget: OwnedSemaphorePermit,
},
Chunk {
file_id: u32,
offset: u64,
data: Vec<u8>,
budget: OwnedSemaphorePermit,
},
Root {
file_id: u32,
root: [u8; 32],
},
Failed(std::io::Error),
}
const READ_AHEAD_KIB: u32 = 96 * 1024;
const OUTBOARD_SYNC_EVERY: Duration = Duration::from_secs(1);
const CHUNK_HDR: u64 = 34;
const BUNDLE_HDR: usize = 22;
const BUNDLE_REC_HDR: usize = 46;
const STALL_FATAL: Duration = Duration::from_secs(10);
const STALL_SLOW_RECEIVER_FATAL: Duration = Duration::from_secs(600);
fn stall_limit(outstanding: u64) -> Duration {
if outstanding == 0 {
STALL_FATAL
} else {
STALL_SLOW_RECEIVER_FATAL
}
}
pub const RECEIVER_SILENT_FATAL: Duration = Duration::from_secs(600);
pub const SKIP_PROGRESS_EVERY: Duration = Duration::from_secs(5);
const _: () = assert!(SKIP_PROGRESS_EVERY.as_secs() <= 10, "SPEC.md §12.8");
fn skip_progress_due(skipped: u64, last_sent: u64, since_last: Duration) -> bool {
skipped != last_sent && since_last >= SKIP_PROGRESS_EVERY
}
fn receiver_silent(limit: Duration) -> SendError {
SendError::Protocol(format!(
"the receiver said nothing about the job for {limit:?} although the connection is alive"
))
}
async fn tell_stalled(link: &JobLink) {
let cancel = gen::JobCancel {
job_id: link.job_id,
reason: gen::ERR_STALLED,
};
let _ = tokio::time::timeout(Duration::from_secs(2), link.control.send(&cancel)).await;
}
pub(crate) async fn next_ctl(link: &mut JobLink) -> Result<Frame, SendError> {
next_ctl_heard(link, &mut Instant::now()).await
}
async fn next_ctl_heard(link: &mut JobLink, heard: &mut Instant) -> Result<Frame, SendError> {
loop {
match link.rx.recv().await {
Some(Inbound::Control(f)) if f.ty == Status::TYPE => {
f.decode::<Status>()
.map_err(|e| SendError::Protocol(e.to_string()))?;
*heard = Instant::now();
continue;
}
Some(Inbound::Control(f)) => {
*heard = Instant::now();
return Ok(f);
}
Some(Inbound::Closed(why)) => return Err(SendError::Disconnected(why)),
None => return Err(SendError::Disconnected("the session ended".into())),
Some(_) => {} }
}
}
pub const OPEN_ACK_TIMEOUT: Duration = Duration::from_secs(30);
pub async fn open_upload(
link: &mut JobLink,
m: &Manifest,
o: &SendOptions,
) -> Result<(u64, Need), SendError> {
let job_id = link.job_id;
link.control
.send(&JobOpen {
job_id,
kind: o.kind,
policy: o.policy,
flags: o.flags,
root: o.root.clone(),
src: None,
credit: None,
})
.await?;
let wait = o.open_ack_timeout.unwrap_or(OPEN_ACK_TIMEOUT);
let deadline = Instant::now() + wait;
let ack: JobOpenAck = loop {
if o.cancel.load(Ordering::Relaxed) {
return Err(SendError::Cancelled);
}
let left = deadline.saturating_duration_since(Instant::now());
if left.is_zero() {
return Err(SendError::OpenTimeout(wait));
}
match tokio::time::timeout(left.min(Duration::from_millis(100)), next_ctl(link)).await {
Err(_) => continue, Ok(f) => {
let f = f?;
if f.ty == JobOpenAck::TYPE {
break f.decode().map_err(|e| SendError::Protocol(e.to_string()))?;
}
}
}
};
if ack.status != gen::STATUS_OK {
return Err(SendError::Refused {
status: ack.status,
message: ack.message.unwrap_or_default(),
});
}
for p in m.pages(job_id) {
link.control.send(&p).await?;
}
link.control
.send(&ManifestEnd {
job_id,
files: m.files(),
bytes: m.bytes(),
manifest_hash: m.hash(),
})
.await?;
let mut need = Need::default();
let mut credit = ack.credit;
let silent = o.receiver_silent_timeout.unwrap_or(RECEIVER_SILENT_FATAL);
let mut heard = Instant::now();
loop {
if o.cancel.load(Ordering::Relaxed) {
let _ = link
.control
.send(&gen::JobCancel {
job_id,
reason: gen::ERR_CANCELLED,
})
.await;
return Err(SendError::Cancelled);
}
if heard.elapsed() >= silent {
tell_stalled(link).await;
return Err(receiver_silent(silent));
}
let Ok(f) =
tokio::time::timeout(Duration::from_millis(100), next_ctl_heard(link, &mut heard))
.await
else {
continue; };
let f = f?;
if f.ty == JobDone::TYPE {
let d: JobDone = f.decode().map_err(|e| SendError::Protocol(e.to_string()))?;
return Err(SendError::Refused {
status: d.status,
message: d.message.unwrap_or_default(),
});
}
if f.ty == Credit::TYPE {
let c: Credit = f.decode().map_err(|e| SendError::Protocol(e.to_string()))?;
credit = credit.saturating_add(c.bytes);
continue;
}
if f.ty != JobMap::TYPE {
continue;
}
let map: JobMap = f.decode().map_err(|e| SendError::Protocol(e.to_string()))?;
if map.status != gen::STATUS_OK {
return Err(SendError::Refused {
status: map.status,
message: map.message.unwrap_or_default(),
});
}
need.add_page(&map);
if map.last == 1 {
return Ok((credit, need));
}
}
}
#[derive(Default)]
struct LiveReaders {
small: Arc<AtomicUsize>,
large: Arc<AtomicUsize>,
}
struct Live {
n: Arc<AtomicUsize>,
counted: bool,
}
impl Live {
fn new(n: &Arc<AtomicUsize>) -> Self {
n.fetch_add(1, Ordering::SeqCst);
Self {
n: n.clone(),
counted: true,
}
}
fn leave(&mut self) {
if std::mem::take(&mut self.counted) {
self.n.fetch_sub(1, Ordering::SeqCst);
}
}
}
impl Drop for Live {
fn drop(&mut self) {
self.leave();
}
}
#[allow(clippy::too_many_arguments)]
fn spawn_readers(
m: Arc<Manifest>,
src: Arc<dyn Source>,
small: Arc<Mutex<VecDeque<u32>>>,
large: Arc<Mutex<VecDeque<(u32, RangeSet)>>>,
sh: Arc<Shared>,
o: &SendOptions,
tx: mpsc::UnboundedSender<Read>,
stop: Arc<AtomicBool>,
live: &LiveReaders,
small_n: usize,
with_large: bool,
) -> Vec<tokio::task::JoinHandle<()>> {
let rt = tokio::runtime::Handle::current();
let mut handles = Vec::new();
let cancel = o.cancel.clone();
for _ in 0..small_n {
let (m, src, small, sh, tx, rt, stop, cancel) = (
m.clone(),
src.clone(),
small.clone(),
sh.clone(),
tx.clone(),
rt.clone(),
stop.clone(),
cancel.clone(),
);
let mut me = Live::new(&live.small);
handles.push(tokio::task::spawn_blocking(move || loop {
if stop.load(Ordering::Relaxed) || cancel.load(Ordering::Relaxed) {
return;
}
let Some(id) = ({
let mut q = small.lock().unwrap();
let id = q.pop_front();
if id.is_none() {
me.leave();
}
id
}) else {
return;
};
let Some(e) = m.entry(id) else {
return;
};
let kib = (e.size / 1024 + 1).min(READ_AHEAD_KIB as u64) as u32;
let Ok(budget) = rt.block_on(sh.bytes_budget.clone().acquire_many_owned(kib)) else {
return;
};
let mut data = vec![0u8; e.size as usize];
let r = src
.open(&e.path)
.and_then(|mut f| read_full_at(f.as_mut(), 0, &mut data));
match r {
Ok(n) if n as u64 == e.size => {
let root = *blake3::hash(&data).as_bytes();
let _ = tx.send(Read::Record {
file_id: id,
root,
data,
budget,
});
}
Ok(_) => {
let _ = tx.send(Read::Failed(std::io::Error::other(format!(
"{} changed while it was being sent",
e.path
))));
return;
}
Err(err) => {
let _ = tx.send(Read::Failed(std::io::Error::new(
err.kind(),
format!("{}: {err}", e.path),
)));
return;
}
}
}));
}
if !with_large {
return handles;
}
let persist = o.persist.clone();
let mut me = Live::new(&live.large);
handles.push(tokio::task::spawn_blocking(move || loop {
if stop.load(Ordering::Relaxed) || cancel.load(Ordering::Relaxed) {
return;
}
let Some((id, durable)) = ({
let mut q = large.lock().unwrap();
let next = q.pop_front();
if next.is_none() {
me.leave();
}
next
}) else {
return;
};
let Some(e) = m.entry(id).cloned() else {
return;
};
let mut hasher = FileHasher::new(e.size);
let mut ob = persist.as_ref().and_then(|d| {
std::fs::create_dir_all(d).ok()?;
Outboard::open(&d.join(format!("{id}.ob")), verify::groups(e.size)).ok()
});
if let Some(ob) = &ob {
for g in 0..verify::groups(e.size) {
if let Some(cv) = ob.get(g) {
hasher.set_cv(g, cv);
}
}
}
let mut f = match src.open(&e.path) {
Ok(f) => f,
Err(err) => {
let _ = tx.send(Read::Failed(err));
return;
}
};
let chunk = sh.chunk.load(Ordering::Relaxed) as u64;
let grant = sh.window.lock().unwrap().available();
let chunk = (chunk.min(grant.saturating_sub(CHUNK_HDR)) / GROUP * GROUP).max(GROUP);
let chunk = if durable.covered() == 0 { chunk } else { GROUP };
let plan = pieces(e.size, &durable, &|g| hasher.cv(g).is_some(), chunk);
let last_piece = plan.len().saturating_sub(1);
let mut ob_synced = Instant::now();
for (pi, p) in plan.into_iter().enumerate() {
if stop.load(Ordering::Relaxed) || cancel.load(Ordering::Relaxed) {
return;
}
let budget = if p.send {
match rt.block_on(
sh.bytes_budget
.clone()
.acquire_many_owned((p.len / 1024 + 1) as u32),
) {
Ok(b) => Some(b),
Err(_) => return,
}
} else {
None
};
let mut data = vec![0u8; p.len as usize];
match read_full_at(f.as_mut(), p.offset, &mut data) {
Ok(n) if n as u64 == p.len => {}
Ok(_) => {
let _ = tx.send(Read::Failed(std::io::Error::other(format!(
"{} changed while it was being sent",
e.path
))));
return;
}
Err(err) => {
let _ = tx.send(Read::Failed(err));
return;
}
}
for (k, g) in data.chunks(GROUP as usize).enumerate() {
let gi = p.offset / GROUP + k as u64;
hasher.add_group(gi, g);
if let (Some(ob), Some(cv)) = (ob.as_mut(), hasher.cv(gi)) {
let _ = ob.put(gi, &cv);
}
}
if pi == last_piece || ob_synced.elapsed() >= OUTBOARD_SYNC_EVERY {
if let Some(ob) = ob.as_mut() {
let _ = ob.sync();
}
ob_synced = Instant::now();
}
if p.send {
let budget = budget.expect("send pieces hold a read-ahead permit");
let _ = tx.send(Read::Chunk {
file_id: id,
offset: p.offset,
data,
budget,
});
}
}
if let Some(root) = hasher.root() {
let _ = tx.send(Read::Root { file_id: id, root });
}
}));
handles
}
struct ReleaseReaders {
stop: Arc<AtomicBool>,
sh: Arc<Shared>,
source: Arc<dyn Source>,
seq: Option<Arc<dyn crate::seq::SeqSource>>,
}
impl Drop for ReleaseReaders {
fn drop(&mut self) {
self.stop.store(true, Ordering::Relaxed);
self.sh.bytes_budget.close();
self.sh.wake();
self.source.close();
if let Some(seq) = &self.seq {
seq.close();
}
}
}
fn spawn_decoder(ctx: crate::seq::DecodeCtx) -> tokio::task::JoinHandle<()> {
tokio::task::spawn_blocking(move || crate::seq::run(ctx))
}
fn bundle_frame(job_id: [u8; 16], recs: Vec<(BundleRecord, OwnedSemaphorePermit)>) -> OutFrame {
let payload = recs.iter().map(|r| r.0.data.len() as u64).sum();
let (records, budget): (Vec<BundleRecord>, Vec<_>) = recs.into_iter().unzip();
let size = BUNDLE_HDR
+ records
.iter()
.map(|r| BUNDLE_REC_HDR + r.data.len())
.sum::<usize>();
let mut w = Writer::with_capacity(size);
Bundle { job_id, records }
.encode_into(&mut w)
.expect("bundle encodes");
let body = w.buf;
OutFrame {
ty: Bundle::TYPE,
body: Arc::new(body),
class: Class::Bundle,
payload,
resend: false,
_budget: budget,
taken: None,
}
}
fn chunk_frame(
job_id: [u8; 16],
file_id: u32,
offset: u64,
data: Vec<u8>,
budget: OwnedSemaphorePermit,
) -> OutFrame {
let payload = data.len() as u64;
let body = Chunk {
job_id,
file_id,
offset,
data,
}
.to_bytes()
.expect("chunk encodes");
OutFrame {
ty: Chunk::TYPE,
body: Arc::new(body),
class: Class::Stream,
payload,
resend: false,
_budget: vec![budget],
taken: None,
}
}
fn pick_any(s: &mut Sched, lane: u16, w: &Window, cap: u64) -> Option<OutFrame> {
let fits = |f: &OutFrame| w.can_send(lane, f.body.len() as u64, cap);
if !s.requeue.is_empty() {
if let Some(i) = s.requeue.iter().position(fits) {
return s.requeue.remove(i);
}
}
let d = s.decision?;
let bundle_first = match d.mode {
Mode::BundleOnly => true,
Mode::StreamOnly => false,
Mode::Mixed => s.bundles_inflight < s.floor || d.prefer == Class::Bundle,
};
let order: [&mut VecDeque<OutFrame>; 2] = if bundle_first {
[&mut s.bundles, &mut s.chunks]
} else {
[&mut s.chunks, &mut s.bundles]
};
for q in order {
if let Some(i) = q.iter().position(&fits) {
return q.remove(i);
}
}
None
}
fn smallest_queued(s: &Sched) -> Option<u64> {
s.requeue
.iter()
.chain(s.bundles.iter())
.chain(s.chunks.iter())
.map(|f| f.body.len() as u64)
.min()
}
fn smooth_rates(lane_bytes: &mut HashMap<u16, u64>, lane_rate: &mut HashMap<u16, f64>, secs: f64) {
let secs = secs.max(0.001);
for (lane, bytes) in lane_bytes.drain() {
let rate = bytes as f64 / secs;
let slot = lane_rate.entry(lane).or_default();
*slot += (rate - *slot) * 0.5;
}
}
type Picked = (u32, u8, Arc<Vec<u8>>, Arc<AtomicBool>);
enum Pick {
Frame(Picked),
Idle,
Stopped,
}
fn pick_for_lane(sh: &Shared, lane: u16, stop: &AtomicBool) -> Pick {
let mut s = sh.sched.lock().unwrap();
if stop.load(Ordering::Relaxed) {
return Pick::Stopped;
}
let mut w = sh.window.lock().unwrap();
let chunk = sh.chunk.load(Ordering::Relaxed);
let rate = s.lane_rate.get(&lane).copied().unwrap_or(0.0);
let cap = governor::inflight_cap(chunk, rate);
match pick_any(&mut s, lane, &w, cap) {
Some(mut f) => {
let len = f.body.len() as u64;
s.next_seq += 1;
let seq = s.next_seq;
let ty = f.ty;
let body = f.body.clone();
if f.class == Class::Bundle {
s.bundles_inflight += 1;
}
assert!(w.sent(lane, seq, len), "can_send passed");
let taken = Arc::new(AtomicBool::new(false));
f.taken = Some(taken.clone());
s.inflight.insert(seq, (lane, f));
drop(w);
drop(s);
*sh.stall.lock().unwrap() = None;
Pick::Frame((seq, ty, body, taken))
}
None => {
let pending = !s.requeue.is_empty() || !s.bundles.is_empty() || !s.chunks.is_empty();
if pending {
let grant = w.available();
let smallest = smallest_queued(&s).unwrap_or(0);
let mut stall = sh.stall.lock().unwrap();
if smallest > grant {
s.credit_starved = true;
if stall.is_none() {
*stall = Some(Stall {
since: Instant::now(),
grant,
smallest,
});
}
} else {
*stall = None;
}
} else {
s.source_starved = true;
}
Pick::Idle
}
}
}
pub(crate) struct Pacer {
bps: f64,
next: Mutex<Option<tokio::time::Instant>>,
}
impl Pacer {
pub(crate) fn new(cap_bps: Option<u64>) -> Option<Arc<Self>> {
match cap_bps {
Some(bps) if bps > 0 => Some(Arc::new(Self {
bps: bps as f64,
next: Mutex::new(None),
})),
_ => None,
}
}
fn reserve(&self, len: u64) -> (tokio::time::Instant, Duration) {
let now = tokio::time::Instant::now();
let mut next = self.next.lock().unwrap();
let start = next.map_or(now, |n| n.max(now));
let span = Duration::from_secs_f64(len as f64 / self.bps);
let due = start + span;
*next = Some(due);
(due, span)
}
pub(crate) async fn pace(&self, len: u64) {
struct Release<'a> {
pacer: &'a Pacer,
span: Duration,
sent: bool,
}
impl Drop for Release<'_> {
fn drop(&mut self) {
if !self.sent {
let mut next = self.pacer.next.lock().unwrap_or_else(|e| e.into_inner());
if let Some(n) = *next {
*next = Some(n.checked_sub(self.span).unwrap_or(n));
}
}
}
}
let (due, span) = self.reserve(len);
let mut release = Release {
pacer: self,
span,
sent: false,
};
tokio::time::sleep_until(due).await;
release.sent = true;
}
}
async fn lane_task(
lane: LaneTx,
sh: Arc<Shared>,
pacer: Option<Arc<Pacer>>,
stop: Arc<AtomicBool>,
) {
let mut wake = sh.wake_tx.subscribe();
loop {
if stop.load(Ordering::Relaxed) {
return;
}
let _ = *wake.borrow_and_update();
let next = match pick_for_lane(&sh, lane.id, &stop) {
Pick::Frame(p) => Some(p),
Pick::Idle => None,
Pick::Stopped => return,
};
let Some((seq, ty, body, taken)) = next else {
if wake.changed().await.is_err() {
return;
}
continue;
};
if let Some(p) = &pacer {
p.pace(body.len() as u64).await; }
if lane
.tx
.send_raw_marked(ty, 0, seq, body, taken)
.await
.is_err()
{
return;
}
}
}
#[cfg(test)]
async fn lane_death(sh: Arc<Shared>, id: u16, tx: ConnTx) {
let untaken = lane_sweep(&sh, id);
lane_settle(sh, id, tx, untaken).await;
}
type Untaken = Vec<(u32, Option<Arc<AtomicBool>>)>;
fn lane_sweep(sh: &Shared, id: u16) -> Untaken {
let seqs = sh.window.lock().unwrap().lane_down(id);
{
let mut s = sh.sched.lock().unwrap();
let mut untaken = Vec::new();
for seq in seqs {
if let Some((_, mut f)) = s.inflight.remove(&seq) {
if f.class == Class::Bundle {
s.bundles_inflight -= 1;
}
f.resend = true;
untaken.push((seq, f.taken.clone()));
s.requeue.push_back(f);
}
}
s.stalls += 1;
untaken
}
}
async fn lane_settle(sh: Arc<Shared>, id: u16, tx: ConnTx, untaken: Untaken) {
let deadline = Instant::now() + Duration::from_millis(250);
while !tx.writer_dead() && Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(5)).await;
}
if tx.writer_dead() {
let mut w = sh.window.lock().unwrap();
for (seq, taken) in untaken {
if taken.as_ref().is_none_or(|t| !t.load(Ordering::SeqCst)) {
let _ = w.release(seq);
}
}
}
let strays = sh.window.lock().unwrap().lane_down(id);
if !strays.is_empty() {
let mut s = sh.sched.lock().unwrap();
for seq in strays {
if let Some((_, mut f)) = s.inflight.remove(&seq) {
if f.class == Class::Bundle {
s.bundles_inflight -= 1;
}
f.resend = true;
let taken = f.taken.as_ref().is_some_and(|t| t.load(Ordering::SeqCst));
if !taken {
let _ = sh.window.lock().unwrap().release(seq);
}
s.requeue.push_back(f);
}
}
}
}
fn large_remaining(size: u64, have: &RangeSet) -> Result<u64, SendError> {
if have.iter().any(|(_, end)| end > size) || have.covered() > size {
return Err(SendError::Protocol(
"the receiver reports durable ranges beyond the end of a file".into(),
));
}
Ok(size - have.covered())
}
pub async fn run_upload(
link: &mut JobLink,
manifest: Arc<Manifest>,
source: Arc<dyn Source>,
opts: SendOptions,
opened: (u64, Need),
) -> Result<SendReport, SendError> {
let (credit, need) = opened;
let job_id = link.job_id;
let pg = opts.progress.clone();
pg.attempts.fetch_add(1, Ordering::Relaxed);
pg.bytes_total.store(manifest.bytes(), Ordering::Relaxed);
pg.files_total
.store(manifest.files() as u64, Ordering::Relaxed);
let (mut small, mut large) = (VecDeque::new(), VecDeque::new());
let (mut small_left, mut large_left) = (0u64, 0u64);
let mut durable_files: HashSet<u32> = need.done.iter().copied().collect();
for (i, e) in manifest.entries.iter().enumerate() {
let id = i as u32;
if e.kind != gen::ENTRY_FILE || durable_files.contains(&id) {
continue;
}
if e.size < opts.cutoff {
small.push_back(id);
small_left += e.size;
} else {
let d = need.partial.get(&id).cloned().unwrap_or_default();
large_left = large_left.saturating_add(large_remaining(e.size, &d)?);
large.push_back((id, d));
}
}
let done_bytes: u64 = durable_files
.iter()
.map(|i| manifest.entry(*i).map_or(0, |e| e.size))
.sum::<u64>()
+ need
.partial
.values()
.fold(0u64, |a, r| a.saturating_add(r.covered()));
if let Some(gate) = opts.space_gate.clone() {
let figures = SpaceFigures::of(&manifest, &need, opts.cutoff)?;
let verdict = tokio::task::spawn_blocking(move || gate(&figures))
.await
.map_err(|e| SendError::Protocol(format!("the space check failed: {e}")))?;
if let Err(why) = verdict {
let _ = link
.control
.send(&gen::JobCancel {
job_id,
reason: gen::ERR_NO_SPACE,
})
.await;
return Err(SendError::NoRoom(why));
}
}
if !pg.skip_recorded.swap(true, Ordering::Relaxed) {
let done_files = need
.done
.iter()
.filter_map(|i| manifest.entry(*i))
.filter(|e| e.kind == gen::ENTRY_FILE);
let (n, b) = done_files.fold((0u64, 0u64), |(n, b), e| (n + 1, b + e.size));
pg.skipped_files.store(n, Ordering::Relaxed);
pg.skipped_bytes.store(b, Ordering::Relaxed);
}
pg.bytes_durable.store(done_bytes, Ordering::Relaxed);
pg.files_durable
.store(durable_files.len() as u64, Ordering::Relaxed);
let sh = Arc::new(Shared {
sched: Mutex::new(Sched {
floor: 4,
..Default::default()
}),
window: Mutex::new(Window::new(credit)),
wake_tx: watch::channel(0).0,
chunk: AtomicU32::new(governor::START_CHUNK),
bundle: AtomicU32::new(governor::START_BUNDLE),
bytes_budget: Arc::new(Semaphore::new(READ_AHEAD_KIB as usize)),
stall: Mutex::new(None),
});
let mut gov = Governor::with_options(GovernorOptions::from_env());
let first = gov.tick(&Sample::default());
sh.chunk.store(first.chunk, Ordering::Relaxed);
let mut summary = JobSummary::default();
sh.sched.lock().unwrap().decision = Some(first);
let small_q = Arc::new(Mutex::new(small));
let large_q = Arc::new(Mutex::new(large));
let (rtx, mut rrx) = mpsc::unbounded_channel();
let stop = Arc::new(AtomicBool::new(false));
let _release_readers = ReleaseReaders {
stop: stop.clone(),
sh: sh.clone(),
source: source.clone(),
seq: opts.seq.clone(),
};
let seq_retries = Arc::new(crate::seq::Retries::default());
let budget_wait: Arc<crate::seq::BudgetWait> = Arc::default();
let mut budget_wait_seen = 0u64;
let seq_skipped = Arc::new(AtomicU64::new(0));
let skip_base = opts.seq.as_ref().map_or(0, |s| s.bytes_skipped());
let skipped_now = |seq: &Arc<dyn crate::seq::SeqSource>| {
seq.bytes_skipped()
.saturating_sub(skip_base)
.saturating_add(seq_skipped.load(Ordering::Relaxed))
};
let (mut skip_sent, mut skip_sent_at) = (0u64, Instant::now());
let mut skip_next: Option<u64> = None;
let live_readers = LiveReaders::default();
let mut readers = if let Some(seq) = opts.seq.clone() {
let chunk_sh = sh.clone();
vec![spawn_decoder(crate::seq::DecodeCtx {
seq,
manifest: manifest.clone(),
need: need.clone(),
cutoff: opts.cutoff,
persist: opts.persist.clone(),
budget: sh.bytes_budget.clone(),
budget_wait: budget_wait.clone(),
chunk: Box::new(move || {
let grant = chunk_sh.window.lock().unwrap().available();
let chunk = chunk_sh.chunk.load(Ordering::Relaxed) as u64;
(chunk.min(grant.saturating_sub(CHUNK_HDR)) / GROUP * GROUP).max(GROUP)
}),
tx: rtx.clone(),
stop: stop.clone(),
cancel: opts.cancel.clone(),
retries: seq_retries.clone(),
rt: tokio::runtime::Handle::current(),
skipped: seq_skipped.clone(),
})]
} else {
spawn_readers(
manifest.clone(),
source.clone(),
small_q.clone(),
large_q.clone(),
sh.clone(),
&opts,
rtx.clone(),
stop.clone(),
&live_readers,
opts.readers.max(1),
true,
)
};
let pacer = Pacer::new(opts.bandwidth_cap);
let mut lane_tasks: HashMap<u16, tokio::task::JoinHandle<()>> = HashMap::new();
let mut lane_links: HashMap<u16, (Arc<AtomicBool>, ConnTx)> = HashMap::new();
let spawn_lane = |id: u16,
tasks: &mut HashMap<u16, tokio::task::JoinHandle<()>>,
links: &mut HashMap<u16, (Arc<AtomicBool>, ConnTx)>,
link: &JobLink| {
if tasks.contains_key(&id) {
return;
}
if let Some(l) = link.lane(id) {
let stop = Arc::new(AtomicBool::new(false));
let h = tokio::spawn(lane_task(
l.clone(),
sh.clone(),
pacer.clone(),
stop.clone(),
));
tasks.insert(id, h);
links.insert(id, (stop, l.tx.clone()));
}
};
for l in link.lanes() {
spawn_lane(l.id, &mut lane_tasks, &mut lane_links, link);
}
let mut pending: Vec<(BundleRecord, OwnedSemaphorePermit)> = Vec::new();
let mut pending_bytes = 0usize;
let mut retries: HashMap<u32, u32> = HashMap::new();
let mut tick = tokio::time::interval(Duration::from_secs(1));
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut last_tick = Instant::now();
let (mut max_lanes, mut receiver_bn, mut sequential) = (0u8, gen::BN_NONE, false);
let mut last_bn = gen::BN_NONE;
let mut root_next: VecDeque<(u32, [u8; 32])> = VecDeque::new();
let mut io_fail: Option<std::io::Error> = None;
let mut opened_start = false;
let mut opening: Option<(BoxFut<'static, Result<u16, crate::Ava1Error>>, bool)> = None;
let start_open = |op: &Arc<dyn LaneOpener>| -> BoxFut<'static, Result<u16, crate::Ava1Error>> {
let op = op.clone();
Box::pin(async move { op.open().await })
};
let mut dying: HashMap<u16, tokio::task::JoinHandle<()>> = HashMap::new();
let mut up_after: HashSet<u16> = HashSet::new();
let (settled_tx, mut settled_rx) = mpsc::unbounded_channel::<u16>();
let mut settle_after = false;
let silent = opts
.receiver_silent_timeout
.unwrap_or(RECEIVER_SILENT_FATAL);
let mut heard = Instant::now();
let result = loop {
if opts.cancel.load(Ordering::Relaxed) {
let _ = link
.control
.send(&gen::JobCancel {
job_id,
reason: gen::ERR_CANCELLED,
})
.await;
break Err(SendError::Cancelled);
}
if !opened_start {
opened_start = true;
if let Some(op) = link.opener() {
if link.lanes().len() < first.lanes as usize {
opening = Some((start_open(op), true));
}
}
}
tokio::select! {
r = rrx.recv() => { match r {
Some(Read::Record { file_id, root, data, budget }) => {
pending_bytes += data.len() + 48;
pending.push((BundleRecord { file_id, root, data }, budget));
}
Some(Read::Chunk { file_id, offset, data, budget }) => {
let f = chunk_frame(job_id, file_id, offset, data, budget);
sh.sched.lock().unwrap().chunks.push_back(f);
sh.wake();
}
Some(Read::Root { file_id, root }) => {
root_next.push_back((file_id, root));
}
Some(Read::Failed(e)) if io_fail.is_none() => {
io_fail = Some(e);
}
Some(Read::Failed(_)) => {} None => {}
}
if !pending.is_empty()
&& (pending_bytes >= sh.bundle.load(Ordering::Relaxed) as usize || rrx.is_empty())
{
let f = bundle_frame(job_id, std::mem::take(&mut pending));
pending_bytes = 0;
sh.sched.lock().unwrap().bundles.push_back(f);
sh.wake();
}
},
root_sent = async {
let (file_id, root) = root_next
.front()
.copied()
.expect("guarded by `if !root_next.is_empty()`");
link.control.send(&FileRoot { job_id, file_id, root }).await
}, if !root_next.is_empty() => {
if let Err(e) = root_sent {
break Err(SendError::Disconnected(e.to_string()));
}
root_next.pop_front();
},
skip_sent_r = async {
let bytes_skipped = skip_next.expect("guarded by `if skip_next.is_some()`");
link.control
.send(&gen::SkipProgress { job_id, bytes_skipped })
.await
}, if skip_next.is_some() => {
if let Err(e) = skip_sent_r {
break Err(SendError::Disconnected(e.to_string()));
}
skip_sent = skip_next.take().expect("guarded by `if skip_next.is_some()`");
skip_sent_at = Instant::now();
},
cancel_sent = async {
link.control
.send(&gen::JobCancel { job_id, reason: gen::ERR_IO })
.await
}, if io_fail.is_some() => {
let _ = cancel_sent;
break Err(SendError::Source(
io_fail.take().expect("guarded by `if io_fail.is_some()`"),
));
},
opened = async {
(&mut opening.as_mut().expect("guarded by `if opening.is_some()`").0).await
}, if opening.is_some() => {
let (_, starting) = opening.take().expect("guarded by `if opening.is_some()`");
match opened {
Ok(_) if starting => {
if let Some(op) = link.opener() {
if link.lanes().len() < first.lanes as usize {
opening = Some((start_open(op), true));
}
}
}
Err(e) if starting => break Err(SendError::Disconnected(e.to_string())),
_ => {}
}
},
settled = settled_rx.recv() => {
if let Some(id) = settled {
dying.remove(&id);
if up_after.remove(&id) {
spawn_lane(id, &mut lane_tasks, &mut lane_links, link);
}
sh.wake();
}
},
ev = link.rx.recv() => match ev {
None => break Err(SendError::Disconnected("the session ended".into())),
Some(Inbound::Closed(why)) => break Err(SendError::Disconnected(why)),
Some(Inbound::LaneUp(id)) => {
if dying.contains_key(&id) {
up_after.insert(id); } else {
spawn_lane(id, &mut lane_tasks, &mut lane_links, link)
}
}
Some(Inbound::LaneDown(id)) => {
if let Some((stop, tx)) = lane_links.remove(&id) {
stop.store(true, Ordering::Relaxed);
lane_tasks.remove(&id);
let untaken = lane_sweep(&sh, id);
let (sh, done) = (sh.clone(), settled_tx.clone());
dying.insert(id, tokio::spawn(async move {
lane_settle(sh, id, tx, untaken).await;
let _ = done.send(id);
}));
}
up_after.remove(&id);
sh.wake();
}
Some(Inbound::Lane { .. }) => {} Some(Inbound::Control(f)) => { heard = Instant::now(); match f.ty {
Received::TYPE => match f.decode::<Received>() {
Ok(r) => { let got = sh.window.lock().unwrap().received(r.seq);
let mut s = sh.sched.lock().unwrap();
if let (Some(len), Some((lane, fr))) = (got, s.inflight.remove(&r.seq)) {
if fr.class == Class::Bundle {
s.bundles_inflight -= 1;
}
s.acked_tick += len;
*s.lane_bytes.entry(lane).or_default() += len;
pg.bytes_sent.fetch_add(fr.payload, Ordering::Relaxed);
if fr.resend {
pg.resent_bytes.fetch_add(fr.payload, Ordering::Relaxed);
}
}
drop(s);
*sh.stall.lock().unwrap() = None; sh.wake();
}
Err(e) => break Err(SendError::Protocol(e.to_string())),
},
Credit::TYPE => match f.decode::<Credit>() {
Ok(c) => { sh.window.lock().unwrap().credit(c.bytes);
*sh.stall.lock().unwrap() = None; sh.wake();
}
Err(e) => break Err(SendError::Protocol(e.to_string())),
},
Durable::TYPE => match f.decode::<Durable>() {
Ok(d) => {
let mut s = sh.sched.lock().unwrap();
for id in from_runs(&d.files) {
if durable_files.insert(id) {
let size = manifest.entry(id).map_or(0, |e| e.size);
pg.files_durable.fetch_add(1, Ordering::Relaxed);
if size < opts.cutoff {
pg.bytes_durable.fetch_add(size, Ordering::Relaxed);
s.small_durable_tick += size;
small_left = small_left.saturating_sub(size);
}
}
}
for r in &d.ranges {
pg.bytes_durable.fetch_add(r.len, Ordering::Relaxed);
s.large_durable_tick += r.len;
large_left = large_left.saturating_sub(r.len);
}
}
Err(e) => break Err(SendError::Protocol(e.to_string())),
},
FileRetry::TYPE => match f.decode::<FileRetry>() {
Ok(r) => {
let Some(entry) = manifest.entry(r.file_id) else {
break Err(SendError::Protocol(format!(
"FileRetry for a file not in the manifest: id {}",
r.file_id
)));
};
let n = retries.entry(r.file_id).or_default();
*n += 1;
if *n > 3 {
let _ = link.control.send(&gen::JobCancel { job_id, reason: gen::ERR_VERIFY }).await;
break Err(SendError::Protocol(format!(
"{} failed verification 3 times",
entry.path
)));
}
if let Some(d) = &opts.persist {
let _ = std::fs::remove_file(d.join(format!("{}.ob", r.file_id)));
}
let (small_n, with_large) = if opts.seq.is_some() {
seq_retries.push(r.file_id);
(0, false)
} else if entry.size < opts.cutoff {
let mut q = small_q.lock().unwrap();
q.push_back(r.file_id);
let idle = live_readers.small.load(Ordering::SeqCst) == 0;
(if idle { opts.readers.max(1) } else { 0 }, false)
} else {
let mut q = large_q.lock().unwrap();
q.push_back((r.file_id, RangeSet::new()));
(0, live_readers.large.load(Ordering::SeqCst) == 0)
};
if small_n > 0 || with_large {
readers.extend(spawn_readers(manifest.clone(), source.clone(), small_q.clone(), large_q.clone(), sh.clone(), &opts, rtx.clone(), stop.clone(), &live_readers, small_n, with_large));
}
}
Err(e) => break Err(SendError::Protocol(e.to_string())),
},
Status::TYPE => match f.decode::<Status>() {
Ok(st) => {
receiver_bn = st.bottleneck;
sh.sched.lock().unwrap().floor = st.workers.max(1) as usize;
}
Err(e) => break Err(SendError::Protocol(e.to_string())),
},
JobDone::TYPE => match f.decode::<JobDone>() {
Ok(d) => {
if d.status == gen::STATUS_OK {
let s = sh.sched.lock().unwrap();
for (_, fr) in s.inflight.values() {
pg.bytes_sent.fetch_add(fr.payload, Ordering::Relaxed);
}
}
settle_after = d.settling == Some(1);
pg.telemetry.lock().unwrap().console_line = d.message.clone();
break Ok(SendReport {
status: d.status,
message: d.message,
files: d.files,
bytes: d.bytes,
resent: pg.resent_bytes.load(Ordering::Relaxed),
max_lanes,
bottleneck: last_bn,
sequential,
});
}
Err(e) => break Err(SendError::Protocol(e.to_string())),
},
gen::JobCancel::TYPE => match f.decode::<gen::JobCancel>() {
Ok(c) => break Err(SendError::Refused {
status: c.reason,
message: "the receiver ended the job".into(),
}),
Err(e) => break Err(SendError::Protocol(e.to_string())),
},
_ => {}
} },
},
_ = tick.tick() => {
if heard.elapsed() >= silent {
tell_stalled(link).await;
break Err(receiver_silent(silent));
}
let stalled = *sh.stall.lock().unwrap();
if let Some(st) = stalled {
let limit = stall_limit(sh.window.lock().unwrap().outstanding());
if st.since.elapsed() >= limit {
let _ = link
.control
.send(&gen::JobCancel {
job_id,
reason: gen::ERR_STALLED,
})
.await;
break Err(SendError::Protocol(format!(
"no queued frame fits the receiver's window ({} bytes granted, the smallest queued frame is {} bytes) and nothing was sent, received or credited for {:?}",
st.grant, st.smallest, limit
)));
}
}
if let Some(seq) = opts.seq.as_ref() {
let n = skipped_now(seq);
if skip_next.is_none() && skip_progress_due(n, skip_sent, skip_sent_at.elapsed()) {
skip_next = Some(n);
}
}
let secs = last_tick.elapsed().as_secs_f64();
last_tick = Instant::now();
let lanes_now = link.lanes().len() as u8;
let budget_waited = budget_wait.waited_since(&mut budget_wait_seen);
let sample = {
let mut s = sh.sched.lock().unwrap();
let mut lane_bytes = std::mem::take(&mut s.lane_bytes);
smooth_rates(&mut lane_bytes, &mut s.lane_rate, secs);
Sample {
secs,
bytes_acked: std::mem::take(&mut s.acked_tick),
lanes: lanes_now,
stalls: std::mem::take(&mut s.stalls),
credit_starved: std::mem::take(&mut s.credit_starved),
source_starved: std::mem::take(&mut s.source_starved)
&& s.requeue.is_empty()
&& !budget_waited,
receiver_bottleneck: receiver_bn,
small_durable: std::mem::take(&mut s.small_durable_tick),
large_durable: std::mem::take(&mut s.large_durable_tick),
small_left,
large_left,
}
};
let d = gov.tick(&sample);
summary.observe(&sample, &d);
sh.chunk.store(d.chunk, Ordering::Relaxed);
sh.bundle.store(d.bundle, Ordering::Relaxed);
sh.sched.lock().unwrap().decision = Some(d);
sequential = d.sequential;
last_bn = d.bottleneck;
pg.bottleneck.store(d.bottleneck, Ordering::Relaxed);
pg.sequential.store(d.sequential, Ordering::Relaxed);
pg.lanes.store(lanes_now, Ordering::Relaxed);
max_lanes = max_lanes.max(lanes_now);
if let Some(op) = link.opener() {
if lanes_now < d.lanes {
if opening.is_none() {
opening = Some((start_open(op), false));
}
} else if lanes_now > d.lanes {
if let Some(l) = link.lanes().last() {
op.close(l.id); }
}
}
sh.wake();
}
}
};
drop(opening);
let mut result = result;
if settle_after && result.is_ok() {
let max = opts.settle_max.unwrap_or(SETTLE_MAX);
let why = settle_wait(link, &pg, &opts.cancel, max).await;
match (&why, result.as_mut()) {
(Settled::Failed(code, msg), Ok(_)) => {
crate::diag::line(format_args!(
"[ava1] upload: the console cannot make its files durable: {msg}"
));
result = Err(SendError::Refused {
status: *code,
message: if msg.is_empty() {
"the console cannot make its files durable".into()
} else {
msg.clone()
},
});
}
(_, Ok(rep)) => {
if let Some(w) = settle_warning(&why) {
rep.message = Some(match rep.message.take() {
Some(m) => format!("{m}; {w}"),
None => w,
});
}
}
_ => {}
}
}
if result.is_err() {
let s = sh.sched.lock().unwrap();
crate::diag::line(format_args!(
"[ava1] sender state at failure: small queue {}, large queue {}, pending records {}, \
bundles {}, chunks {}, requeue {}, inflight {}, roots waiting {}, budget free {} KiB, \
window {} B, lanes {}, readers running {}",
small_q.lock().unwrap().len(),
large_q.lock().unwrap().len(),
pending.len(),
s.bundles.len(),
s.chunks.len(),
s.requeue.len(),
s.inflight.len(),
root_next.len(),
sh.bytes_budget.available_permits(),
sh.window.lock().unwrap().available(),
link.lanes().len(),
readers.iter().filter(|h| !h.is_finished()).count(),
));
}
stop.store(true, Ordering::Relaxed);
sh.wake();
source.close();
if let Some(seq) = &opts.seq {
seq.close();
}
for (_, h) in lane_tasks.drain() {
h.abort();
let _ = h.await;
}
for (_, h) in dying.drain() {
let _ = h.await;
}
sh.bytes_budget.close();
{
let mut s = sh.sched.lock().unwrap();
s.bundles.clear();
s.chunks.clear();
s.requeue.clear();
s.inflight.clear();
s.bundles_inflight = 0;
}
for h in readers {
let _ = h.await;
}
pg.telemetry.lock().unwrap().shares.merge(&summary);
if let Some(line) = summary.line() {
crate::diag::line(format_args!("{line}"));
}
result
}
pub const SETTLE_MAX: Duration = Duration::from_secs(30);
#[derive(Debug, PartialEq, Eq)]
enum Settled {
Done,
TimedOut,
Cancelled,
Closed,
Failed(u16, String),
}
struct SettleTimer<'a>(&'a Progress, Instant);
impl Drop for SettleTimer<'_> {
fn drop(&mut self) {
let ms = self.1.elapsed().as_millis() as u64;
self.0.settle_ms.fetch_add(ms, Ordering::Relaxed);
}
}
async fn settle_wait(
link: &mut JobLink,
pg: &Progress,
cancel: &AtomicBool,
max: Duration,
) -> Settled {
pg.settling.store(true, Ordering::Relaxed);
let t0 = Instant::now();
let _timed = SettleTimer(pg, t0);
let out = loop {
if cancel.load(Ordering::Relaxed) {
break Settled::Cancelled;
}
let Some(left) = max.checked_sub(t0.elapsed()) else {
break Settled::TimedOut;
};
match tokio::time::timeout(left.min(Duration::from_millis(100)), link.rx.recv()).await {
Ok(Some(Inbound::Control(f))) if f.ty == Status::TYPE => {
let Ok(st) = f.decode::<Status>() else {
continue;
};
if let Some(code) = st.code.filter(|c| *c != 0) {
break Settled::Failed(code, st.current.unwrap_or_default());
}
let n = st.unswept.unwrap_or(0);
pg.unswept.store(n, Ordering::Relaxed);
pg.unswept_peak.fetch_max(n, Ordering::Relaxed);
if n == 0 {
break Settled::Done;
}
}
Ok(Some(Inbound::Closed(_))) | Ok(None) => break Settled::Closed,
Ok(Some(_)) | Err(_) => {}
}
};
pg.unswept.store(0, Ordering::Relaxed);
pg.settling.store(false, Ordering::Relaxed);
out
}
fn settle_warning(why: &Settled) -> Option<String> {
let w = match why {
Settled::Done => return None,
Settled::TimedOut => "files are still being made durable on the console",
Settled::Cancelled => {
"cancelled while files were still being made durable on the console (they are safe in its log)"
}
Settled::Closed => {
"the console closed the session while files were still being made durable (they are safe in its log)"
}
Settled::Failed(..) => return None,
};
crate::diag::line(format_args!("[ava1] upload: {w}"));
Some(w.to_string())
}
pub async fn send_job(
link: &mut JobLink,
manifest: Arc<Manifest>,
source: Arc<dyn Source>,
opts: SendOptions,
) -> Result<SendReport, SendError> {
let opened = open_upload(link, &manifest, &opts).await?;
run_upload(link, manifest, source, opts, opened).await
}
async fn open_as_responder(
link: &mut JobLink,
open: &JobOpen,
m: &Manifest,
) -> Result<(u64, Need), SendError> {
let job_id = link.job_id;
link.control
.send(&JobOpenAck {
job_id,
status: gen::STATUS_OK,
credit: 0,
staged: 0,
workers: 0,
message: None,
})
.await
.map_err(|e| SendError::Disconnected(e.to_string()))?;
for p in m.pages(job_id) {
link.control
.send(&p)
.await
.map_err(|e| SendError::Disconnected(e.to_string()))?;
}
link.control
.send(&ManifestEnd {
job_id,
files: m.files(),
bytes: m.bytes(),
manifest_hash: m.hash(),
})
.await
.map_err(|e| SendError::Disconnected(e.to_string()))?;
let mut need = Need::default();
loop {
let f = next_ctl(link).await?;
if f.ty != JobMap::TYPE {
continue;
}
let map: JobMap = f.decode().map_err(|e| SendError::Protocol(e.to_string()))?;
if map.status != gen::STATUS_OK {
return Err(SendError::Refused {
status: map.status,
message: map.message.unwrap_or_default(),
});
}
need.add_page(&map);
if map.last == 1 {
return Ok((open.credit.unwrap_or(16 << 20), need));
}
}
}
pub async fn serve_download(
link: &mut JobLink,
open: JobOpen,
manifest: Arc<Manifest>,
source: Arc<dyn Source>,
mut opts: SendOptions,
) -> Result<SendReport, SendError> {
if open.flags & gen::JF_ORDERED != 0 {
opts.readers = 1;
opts.cutoff = 0;
}
let opened = open_as_responder(link, &open, &manifest).await?;
run_upload(link, manifest, source, opts, opened).await
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conn::{FrameReader, FrameWriter};
use crate::manifest::Entry;
use crate::router::{is_data_type, BoxFut, ConnTx, LaneOpener, Router};
use crate::session::Timing;
use crate::source::SourceMeta;
use crate::wire::SplitMix;
use std::io;
use tokio::io::{duplex, split};
fn mf(sizes: &[u64]) -> Manifest {
let mut m = Manifest::default();
for (i, s) in sizes.iter().enumerate() {
m.entries.push(crate::manifest::Entry {
path: format!("f{i}"),
size: *s,
kind: gen::ENTRY_FILE,
mode: 0o644,
mtime: 1,
root: None,
});
}
m
}
#[test]
fn a_bundle_frame_encodes_the_same_bytes_into_a_buffer_sized_exactly() {
let sem = Arc::new(tokio::sync::Semaphore::new(16));
let recs: Vec<BundleRecord> = [0usize, 1, 5000]
.iter()
.enumerate()
.map(|(i, n)| BundleRecord {
file_id: i as u32,
root: [i as u8; 32],
data: vec![i as u8 + 1; *n],
})
.collect();
let want = Bundle {
job_id: [3; 16],
records: recs.clone(),
}
.to_bytes()
.unwrap();
let f = bundle_frame(
[3; 16],
recs.into_iter()
.map(|r| (r, sem.clone().try_acquire_owned().unwrap()))
.collect(),
);
assert_eq!(*f.body, want, "same bytes on the wire");
assert_eq!(f.body.capacity(), want.len());
assert_eq!(f._budget.len(), 3);
}
#[test]
fn skip_progress_is_sent_only_when_the_skip_moved_and_the_interval_passed() {
let every = SKIP_PROGRESS_EVERY;
assert!(!skip_progress_due(0, 0, every * 10), "nothing skipped");
assert!(
!skip_progress_due(5, 5, every * 10),
"no new skip since the last"
);
assert!(!skip_progress_due(9, 5, every / 2), "too soon");
assert!(skip_progress_due(9, 5, every));
}
#[test]
fn a_fresh_job_needs_every_byte_and_a_resumed_one_only_the_rest() {
let cutoff = 1 << 20;
let m = mf(&[100, 5 << 20, 8 << 20]);
let total = 100 + (13 << 20);
let f = SpaceFigures::of(&m, &Need::default(), cutoff).unwrap();
assert_eq!((f.job_bytes, f.durable_bytes, f.held_bytes), (total, 0, 0));
assert_eq!(f.to_allocate(), total);
let mut need = Need::default();
need.done.insert(0);
need.partial.entry(1).or_default().insert(0, 2 << 20);
let f = SpaceFigures::of(&m, &need, cutoff).unwrap();
assert_eq!(f.durable_bytes, 100 + (2 << 20));
assert_eq!(f.to_allocate(), (3 << 20) + (8 << 20));
}
#[test]
fn a_preallocated_part_is_credited_by_what_the_console_reports_held() {
let cutoff = 1 << 20;
let m = mf(&[5 << 20, 8 << 20]);
let mut need = Need::default();
need.partial.entry(0).or_default().insert(0, 2 << 20);
need.held = 5 << 20;
let f = SpaceFigures::of(&m, &need, cutoff).unwrap();
assert_eq!(f.to_allocate(), 8 << 20);
need.held = 1 << 20;
assert_eq!(
SpaceFigures::of(&m, &need, cutoff).unwrap().to_allocate(),
(3 << 20) + (8 << 20)
);
}
#[test]
fn a_held_figure_larger_than_the_unfinished_large_files_is_clamped_not_believed() {
let cutoff = 1 << 20;
let m = mf(&[100, 5 << 20]);
let need = Need {
held: u64::MAX,
..Need::default()
};
let f = SpaceFigures::of(&m, &need, cutoff).unwrap();
assert_eq!(f.to_allocate(), 100);
let mut bad = Need::default();
bad.partial.entry(1).or_default().insert(0, 9 << 20);
assert!(SpaceFigures::of(&m, &bad, cutoff).is_err());
}
#[test]
fn durable_ranges_beyond_a_file_are_a_protocol_error_not_an_underflow() {
let mut have = RangeSet::new();
have.insert(0, 100);
assert_eq!(large_remaining(300, &have).unwrap(), 200);
assert!(matches!(
large_remaining(50, &have),
Err(SendError::Protocol(_))
));
let mut far = RangeSet::new();
far.insert(1000, 1010);
assert!(matches!(
large_remaining(500, &far),
Err(SendError::Protocol(_))
));
let mut all = RangeSet::new();
all.insert(0, u64::MAX);
assert!(large_remaining(10, &all).is_err());
}
#[test]
fn a_stall_with_bytes_outstanding_is_a_slow_receiver_not_a_dead_one() {
let mut w = Window::new(64 << 20);
assert_eq!(w.outstanding(), 0);
assert_eq!(stall_limit(w.outstanding()), STALL_FATAL);
assert!(w.sent(0, 1, (64 << 20) - 18_260));
assert_eq!(w.available(), 18_260);
assert_eq!(w.outstanding(), (64 << 20) - 18_260);
assert!(stall_limit(w.outstanding()) >= Duration::from_secs(60));
w.received(1); assert!(stall_limit(w.outstanding()) >= Duration::from_secs(60));
w.credit((64 << 20) - 18_260); assert_eq!(w.outstanding(), 0);
assert_eq!(stall_limit(w.outstanding()), STALL_FATAL);
}
#[test]
fn skip_set_is_the_complement_of_the_read_set() {
let dir = std::env::temp_dir().join(format!(
"ava1-skip-{}-{:?}",
std::process::id(),
std::thread::current().id()
));
std::fs::create_dir_all(&dir).unwrap();
let size = gen::LARGE_CUTOFF as u64 + 3 * GROUP;
let m = Manifest {
entries: vec![
Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size,
mtime: 0,
path: "partial".into(),
root: None,
},
Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: 3,
mtime: 0,
path: "done".into(),
root: None,
},
],
};
let mut durable = Need::default();
durable.done.insert(1);
durable.partial.entry(0).or_default().insert(0, 2 * GROUP);
let mut ob = Outboard::open(&dir.join("0.ob"), verify::groups(size)).unwrap();
ob.put(0, &[1; 32]).unwrap();
ob.sync().unwrap();
let skip = skip_set(&m, &durable, Some(&dir));
assert!(skip.done.contains(&1));
assert_eq!(
skip.partial[&0].iter().collect::<Vec<_>>(),
vec![(0, GROUP)]
);
for chunk in [GROUP, 4 * GROUP] {
let plan = pieces(size, &durable.partial[&0], &|g| g == 0, chunk);
for g in 0..verify::groups(size) {
let off = g * GROUP;
let skipped = skip.partial[&0].covers(off, off + (size - off).min(GROUP));
let read = plan
.iter()
.any(|p| p.offset <= off && off < p.offset + p.len);
assert_ne!(skipped, read, "group {g}, chunk {chunk}");
}
}
std::fs::remove_dir_all(dir).unwrap();
}
const G: u64 = crate::verify::GROUP;
#[test]
fn pieces_send_the_missing_groups_and_hash_the_unknown_ones() {
let mut durable = RangeSet::new();
durable.insert(0, 2 * G);
durable.insert(4 * G, 5 * G);
let size = 6 * G + 7;
let p = pieces(size, &durable, &|g| g == 0, 2 * G);
let got: Vec<(u64, u64, bool)> = p.iter().map(|x| (x.offset / G, x.len, x.send)).collect();
assert_eq!(
got,
vec![
(1, G, false),
(2, 2 * G, true),
(4, G, false),
(5, G + 7, true)
]
);
}
#[test]
fn pieces_never_cross_a_chunk_limit_or_a_group_boundary() {
let p = pieces(10 * G + 1, &RangeSet::new(), &|_| false, 3 * G);
assert!(p
.iter()
.all(|x| x.offset % G == 0 && x.len <= 3 * G && x.send));
assert_eq!(p.iter().map(|x| x.len).sum::<u64>(), 10 * G + 1);
}
#[test]
fn every_chunk_but_a_files_last_is_whole_groups() {
let mut rng = SplitMix(11);
for _ in 0..300 {
let size = rng.below(16 * G) + 1;
let mut durable = RangeSet::new();
let mut left = rng.below(6);
while left > 0 {
let s = rng.below(size);
durable.insert(s, (s + rng.below(size - s) + 1).min(size));
left -= 1;
}
let chunk = (rng.below(5) + 1) * G;
let cv_at = rng.below(16);
let p = pieces(size, &durable, &|g| g == cv_at, chunk);
for x in &p {
assert!(x.len <= chunk && x.offset % G == 0, "{x:?} of {size}");
if x.offset + x.len < size {
assert_eq!(x.len % G, 0, "a mid-file piece is not whole groups: {x:?}");
}
}
}
}
#[test]
fn a_lane_death_keeps_the_window_charged_until_the_receiver_accounts() {
let mut w = Window::new(100);
assert!(w.can_send(1, 60, 1000));
w.sent(1, 7, 60);
assert!(!w.can_send(2, 60, 1000), "credit");
w.lane_down(1); assert_eq!(w.refunded.len(), 1, "the dead frame is held");
assert_eq!(w.available(), 40, "a lane death releases no window credit");
assert!(
!w.can_send(2, 60, 1000),
"the dead lane's bytes could be spent twice"
);
w.credit(60); assert_eq!(
w.refunded.len(),
1,
"a Credit closes nothing out: the charge lives in the window"
);
w.received(7); assert_eq!(
w.refunded.len(),
0,
"the late Received closed the frame out"
);
assert_eq!(
w.available(),
100,
"a late Received double-charged the sender"
);
w.received(7); assert_eq!(w.refunded.len(), 0, "a duplicated Received changes nothing");
assert!(w.can_send(2, 60, 1000));
w.sent(2, 8, 60);
assert!(!w.can_send(2, 30, 70), "the lane's in-flight cap");
assert_eq!(w.received(8), Some(60));
}
#[test]
fn a_dead_lane_releases_only_frames_the_writer_never_took() {
let mut w = Window::new(8 << 20);
assert!(w.sent(1, 1, 4 << 20)); assert!(w.sent(1, 2, 4 << 20)); assert_eq!(w.lane_down(1), vec![1, 2]);
assert_eq!(
w.available(),
0,
"a lane death must not release credit itself"
);
assert_eq!(w.release(2), Some(4 << 20));
assert_eq!(w.available(), 4 << 20);
assert_eq!(w.release(2), None, "released twice");
assert_eq!(w.received(1), None);
assert_eq!(w.available(), 4 << 20, "a late Received releases nothing");
w.credit(4 << 20); assert_eq!(w.available(), 8 << 20);
assert_eq!(
w.received(2),
None,
"a late Received for a released frame is nothing"
);
assert_eq!(w.available(), 8 << 20);
}
#[test]
fn lane_churn_never_wedges_the_window_when_the_untaken_frames_are_released() {
for _ in 0..200 {
let mut w = Window::new(8 << 20);
assert!(w.sent(1, 1, 4 << 20)); assert!(w.sent(1, 2, 4 << 20)); w.lane_down(1);
w.release(2);
assert_eq!(w.received(1), None); w.credit(4 << 20); assert_eq!(
w.available(),
8 << 20,
"the window shrank: an untaken frame's charge leaked"
);
}
}
#[test]
fn peer_credit_overflow_saturates_instead_of_wrapping() {
let mut w = Window::new(u64::MAX / 2);
w.credit(u64::MAX);
assert_eq!(w.available(), u64::MAX);
assert!(w.can_send(1, u64::MAX, u64::MAX));
}
#[test]
fn a_lane_skips_a_front_frame_that_does_not_fit_and_takes_one_behind_it() {
let mut s = Sched {
decision: Some(governor::Decision {
lanes: 1,
chunk: 4 << 20,
bundle: 1 << 20,
bottleneck: gen::BN_NETWORK,
mode: Mode::StreamOnly,
prefer: Class::Stream,
sequential: false,
}),
..Default::default()
};
s.chunks.push_back(test_frame(Chunk::TYPE, 15 << 20));
s.chunks.push_back(test_frame(Chunk::TYPE, 4 << 20));
let w = Window::new(8 << 20);
let f = pick_any(&mut s, 1, &w, 64 << 20).unwrap();
assert_eq!(
f.payload,
4 << 20,
"the fitting frame behind the oversized front was taken"
);
assert_eq!(s.chunks.len(), 1, "the oversized front frame stayed queued");
assert_eq!(s.chunks.front().unwrap().payload, 15 << 20);
let tiny = Window::new(3 << 20);
assert!(pick_any(&mut s, 1, &tiny, 64 << 20).is_none());
assert_eq!(s.chunks.len(), 1);
}
#[test]
fn sent_never_takes_more_than_the_credit() {
let mut w = Window::new(50);
assert!(
!w.sent(1, 1, 60),
"a frame larger than the credit is refused"
);
assert_eq!(w.available(), 50);
assert!(w.sent(1, 2, 50));
assert_eq!(w.available(), 0);
}
#[test]
fn the_window_never_underflows_under_mixed_sizes() {
let mut rng = SplitMix(3);
let mut w = Window::new(500);
let mut credited = 0u64;
for _ in 0..20_000 {
match rng.below(6) {
0 => {
let lane = (rng.below(3) + 1) as u16;
let len = rng.below(600) + 1;
if w.can_send(lane, len, 1000) {
assert!(w.sent(lane, rng.below(500) as u32, len));
}
}
1 => {
let _ = w.received(rng.below(500) as u32);
}
2 => {
let lane = (rng.below(3) + 1) as u16;
let _ = w.lane_down(lane);
}
3 => {
let n = rng.below(400);
w.credit(n);
credited += n;
}
_ => {
let _ = w.available();
}
}
assert!(w.available() <= 500 + credited, "credit over-accounted");
}
}
#[test]
fn lane_rates_are_smoothed_and_survive_idle_ticks() {
let mut bytes: HashMap<u16, u64> = HashMap::new();
let mut rate: HashMap<u16, f64> = HashMap::new();
for _ in 0..8 {
bytes.insert(1, 10_000_000);
smooth_rates(&mut bytes, &mut rate, 1.0);
}
assert!(
rate[&1] > 9.9e6,
"the EWMA converges to the true rate: {}",
rate[&1]
);
let before = rate[&1];
smooth_rates(&mut bytes, &mut rate, 1.0);
assert_eq!(rate[&1], before, "an idle tick must not reset the rate");
assert_eq!(
governor::inflight_cap(4 << 20, rate[&1]),
(rate[&1] * 2.0) as u64,
"the cap is two seconds of the lane's rate"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_dropped_job_wakes_a_reader_parked_on_the_budget() {
let sh = test_shared(governor::START_CHUNK, 4);
let held = sh.bytes_budget.clone().acquire_many_owned(4).await.unwrap();
let rt = tokio::runtime::Handle::current();
let sh2 = sh.clone();
let reader = tokio::task::spawn_blocking(move || {
rt.block_on(sh2.bytes_budget.clone().acquire_many_owned(4))
.is_err()
});
tokio::time::sleep(Duration::from_millis(100)).await;
let guard = ReleaseReaders {
stop: Arc::new(AtomicBool::new(false)),
sh: sh.clone(),
source: Arc::new(crate::source::LocalSource::new(std::env::temp_dir())),
seq: None,
};
drop(guard);
let woke = tokio::time::timeout(Duration::from_secs(5), reader)
.await
.expect("the parked reader never woke")
.unwrap();
assert!(
woke,
"the reader got a permit instead of seeing the budget closed"
);
drop(held);
}
#[tokio::test(start_paused = true)]
async fn the_speed_cap_holds_for_the_whole_job_not_for_each_lane() {
const LANES: usize = 4;
const FRAME: u64 = 64 << 10;
let cap = 8u64 << 20;
let pacer = Pacer::new(Some(cap)).unwrap();
let sent = Arc::new(AtomicU64::new(0));
let stop = Arc::new(AtomicBool::new(false));
let t0 = tokio::time::Instant::now();
let lanes: Vec<_> = (0..LANES)
.map(|_| {
let (pacer, sent, stop) = (pacer.clone(), sent.clone(), stop.clone());
tokio::spawn(async move {
while !stop.load(Ordering::Relaxed) {
pacer.pace(FRAME).await;
sent.fetch_add(FRAME, Ordering::Relaxed);
}
})
})
.collect();
tokio::time::sleep(Duration::from_secs(4)).await;
stop.store(true, Ordering::Relaxed);
let rate = sent.load(Ordering::Relaxed) as f64 / t0.elapsed().as_secs_f64();
for l in lanes {
l.await.unwrap();
}
let cap = cap as f64;
assert!(
rate > cap * 0.95 && rate <= cap * 1.01,
"{LANES} lanes sent {rate:.0} B/s against a {cap:.0} B/s cap"
);
}
#[tokio::test(start_paused = true)]
async fn an_idle_stretch_does_not_bank_a_burst_past_the_cap() {
let pacer = Pacer::new(Some(1 << 20)).unwrap();
pacer.pace(1 << 20).await;
tokio::time::sleep(Duration::from_secs(10)).await;
let t = tokio::time::Instant::now();
pacer.pace(1 << 20).await;
assert!(t.elapsed() >= Duration::from_millis(990));
}
#[tokio::test(start_paused = true)]
async fn a_cancelled_wait_gives_its_budget_back() {
let pacer = Pacer::new(Some(1 << 20)).unwrap();
let _ = tokio::time::timeout(Duration::from_millis(10), pacer.pace(10 << 20)).await;
let t = tokio::time::Instant::now();
pacer.pace(1 << 20).await;
let waited = t.elapsed();
assert!(
waited <= Duration::from_millis(1010),
"the cancelled reservation still held the cap: waited {waited:?}"
);
}
#[test]
fn no_cap_or_a_zero_cap_means_no_pacer() {
assert!(Pacer::new(None).is_none());
assert!(Pacer::new(Some(0)).is_none());
}
fn test_shared(chunk: u32, budget_kib: usize) -> Arc<Shared> {
Arc::new(Shared {
sched: Mutex::new(Sched {
floor: 4,
..Default::default()
}),
window: Mutex::new(Window::new(1 << 20)),
wake_tx: watch::channel(0).0,
chunk: AtomicU32::new(chunk),
bundle: AtomicU32::new(governor::START_BUNDLE),
bytes_budget: Arc::new(Semaphore::new(budget_kib)),
stall: Mutex::new(None),
})
}
fn test_frame(ty: u8, payload: u64) -> OutFrame {
OutFrame {
ty,
body: Arc::new(
Chunk {
job_id: [0x78; 16],
file_id: 0,
offset: 0,
data: vec![0x22; payload as usize],
}
.to_bytes()
.unwrap(),
),
class: Class::Stream,
payload,
resend: false,
_budget: Vec::new(),
taken: None,
}
}
#[tokio::test]
async fn a_wake_between_the_check_and_the_wait_is_not_lost() {
let sh = test_shared(governor::START_CHUNK, READ_AHEAD_KIB as usize);
let (armed, fire) = tokio::sync::oneshot::channel();
let waiter = {
let sh = sh.clone();
tokio::spawn(async move {
let mut rx = sh.wake_tx.subscribe();
let mut armed = Some(armed);
loop {
let _ = *rx.borrow_and_update(); let work = {
let s = sh.sched.lock().unwrap();
!s.chunks.is_empty() || !s.bundles.is_empty() || !s.requeue.is_empty()
};
if work {
return;
}
if let Some(a) = armed.take() {
if a.send(()).is_err() {
return;
}
}
if rx.changed().await.is_err() {
return;
}
}
})
};
fire.await.unwrap(); sh.sched
.lock()
.unwrap()
.chunks
.push_back(test_frame(Chunk::TYPE, 64));
sh.wake();
tokio::time::timeout(Duration::from_secs(1), waiter)
.await
.expect("the wake between the check and the wait was lost")
.unwrap();
}
struct MemRead {
data: Vec<u8>,
}
impl crate::source::ReadAt for MemRead {
fn read_at(&mut self, off: u64, buf: &mut [u8]) -> io::Result<usize> {
let n = buf.len().min(self.data.len().saturating_sub(off as usize));
if n > 0 {
buf[..n].copy_from_slice(&self.data[off as usize..off as usize + n]);
}
Ok(n)
}
}
struct MemSource(Vec<u8>);
impl Source for MemSource {
fn open(&self, _rel: &str) -> io::Result<Box<dyn crate::source::ReadAt>> {
Ok(Box::new(MemRead {
data: self.0.clone(),
}))
}
fn list(&self, _rel: &str) -> io::Result<Vec<(String, SourceMeta)>> {
Ok(Vec::new())
}
fn stat(&self, _rel: &str) -> io::Result<SourceMeta> {
Ok(SourceMeta {
size: self.0.len() as u64,
mtime: 0,
mode: 0o644,
is_dir: false,
})
}
}
struct GatedSource {
data: Vec<u8>,
gate: Arc<(Mutex<bool>, std::sync::Condvar)>,
opens: Arc<AtomicUsize>,
}
impl Source for GatedSource {
fn open(&self, _rel: &str) -> io::Result<Box<dyn crate::source::ReadAt>> {
self.opens.fetch_add(1, Ordering::SeqCst);
let (m, cv) = &*self.gate;
let _open = cv.wait_while(m.lock().unwrap(), |open| !*open).unwrap();
Ok(Box::new(MemRead {
data: self.data.clone(),
}))
}
fn list(&self, _rel: &str) -> io::Result<Vec<(String, SourceMeta)>> {
Ok(Vec::new())
}
fn stat(&self, _rel: &str) -> io::Result<SourceMeta> {
Ok(SourceMeta {
size: self.data.len() as u64,
mtime: 0,
mode: 0o644,
is_dir: false,
})
}
}
#[tokio::test(flavor = "multi_thread")]
async fn a_retry_pushed_while_a_reader_is_live_is_read_by_it_and_needs_no_new_set() {
let sh = test_shared(governor::START_CHUNK, READ_AHEAD_KIB as usize);
let m = Arc::new(mf(&[100, 100]));
let gate = Arc::new((Mutex::new(false), std::sync::Condvar::new()));
let opens = Arc::new(AtomicUsize::new(0));
let src: Arc<dyn Source> = Arc::new(GatedSource {
data: vec![7; 100],
gate: gate.clone(),
opens: opens.clone(),
});
let (tx, mut rx) = mpsc::unbounded_channel();
let small = Arc::new(Mutex::new(VecDeque::from([0u32])));
let large = Arc::new(Mutex::new(VecDeque::new()));
let live = LiveReaders::default();
let handles = spawn_readers(
m,
src,
small.clone(),
large,
sh,
&SendOptions::upload(""),
tx,
Arc::new(AtomicBool::new(false)),
&live,
3,
true,
);
tokio::time::timeout(Duration::from_secs(10), async {
while opens.load(Ordering::SeqCst) == 0
|| live.small.load(Ordering::SeqCst) != 1
|| live.large.load(Ordering::SeqCst) != 0
{
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("one reader live, parked in the source");
{
let mut q = small.lock().unwrap();
q.push_back(1);
assert_eq!(live.small.load(Ordering::SeqCst), 1);
}
*gate.0.lock().unwrap() = true;
gate.1.notify_all();
let mut got = Vec::new();
for _ in 0..2 {
match tokio::time::timeout(Duration::from_secs(10), rx.recv()).await {
Ok(Some(Read::Record { file_id, .. })) => got.push(file_id),
_ => panic!("the live reader reads both files"),
}
}
got.sort();
assert_eq!(got, [0, 1]);
for h in handles {
h.await.unwrap();
}
assert_eq!(live.small.load(Ordering::SeqCst), 0, "every reader left");
assert_eq!(opens.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn reader_bytes_stay_under_the_read_ahead_budget() {
let g = G as usize;
let size = 16 * g;
let sh = test_shared(g as u32, g / 1024 + 8); let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "f".into(),
root: None,
}],
});
let src: Arc<dyn Source> = Arc::new(MemSource(vec![0x5a; size]));
let (tx, mut rx) = mpsc::unbounded_channel();
let small = Arc::new(Mutex::new(VecDeque::new()));
let large = Arc::new(Mutex::new(VecDeque::from([(0u32, RangeSet::new())])));
let _handles = spawn_readers(
m,
src,
small,
large,
sh.clone(),
&SendOptions::upload(""),
tx,
Arc::new(AtomicBool::new(false)),
&LiveReaders::default(),
8,
true,
);
let budget = sh.bytes_budget.clone();
let first = tokio::time::timeout(Duration::from_secs(10), rx.recv())
.await
.unwrap()
.unwrap();
match &first {
Read::Chunk { data, .. } => assert!(data.len() <= g),
_ => panic!("expected a chunk"),
}
assert!(
budget.available_permits() < g / 1024 + 1,
"{}",
budget.available_permits()
);
assert!(tokio::time::timeout(Duration::from_millis(200), rx.recv())
.await
.is_err());
drop(first);
let second = tokio::time::timeout(Duration::from_secs(10), rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(second, Read::Chunk { .. }));
}
#[tokio::test]
async fn a_lane_task_stops_promptly_when_the_job_ends() {
let timing = Timing {
ping_every: Duration::from_secs(3600),
dead_after: Duration::from_secs(3600),
handshake: Duration::from_secs(1),
min_frame_rate: crate::link::MIN_FRAME_RATE,
};
let (a, b) = duplex(1 << 20);
let (ar, aw) = split(a);
let (tx, _rx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (_link, outbox) =
crate::link::drive(FrameReader::new(ar), FrameWriter::new(aw), timing, tx);
let mut peer = FrameReader::new(b);
peer.set_max_body(crate::frame::MAX_BODY);
let drain = tokio::spawn(async move {
let _ = peer.recv().await;
});
let router = Arc::new(Router::default());
router.lane_up(1, outbox);
let lane = router.lane(1).unwrap();
let sh = test_shared(governor::START_CHUNK, READ_AHEAD_KIB as usize);
sh.sched
.lock()
.unwrap()
.chunks
.push_back(test_frame(Chunk::TYPE, 1 << 20));
let stop = Arc::new(AtomicBool::new(false));
let h = tokio::spawn(lane_task(lane, sh.clone(), None, stop.clone()));
tokio::time::timeout(Duration::from_secs(5), drain)
.await
.unwrap()
.unwrap();
stop.store(true, Ordering::Relaxed);
sh.wake();
tokio::time::timeout(Duration::from_secs(1), h)
.await
.expect("the lane task did not stop")
.unwrap();
}
#[tokio::test(flavor = "multi_thread")]
async fn a_frame_taken_after_the_sweep_but_before_the_writers_death_keeps_its_charge() {
let timing = Timing {
ping_every: Duration::from_secs(3600),
dead_after: Duration::from_secs(3600),
handshake: Duration::from_secs(1),
min_frame_rate: crate::link::MIN_FRAME_RATE,
};
let (a, b) = duplex(1 << 20);
let (ar, aw) = split(a);
let (br, _bw) = split(b);
let (tx, _rx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (link, outbox) =
crate::link::drive(FrameReader::new(ar), FrameWriter::new(aw), timing, tx);
tokio::spawn(async move {
let mut peer = FrameReader::new(br);
peer.set_max_body(crate::frame::MAX_BODY);
while peer.recv().await.is_ok() {}
});
let sh = Arc::new(Shared {
sched: Mutex::new(Sched {
floor: 4,
..Default::default()
}),
window: Mutex::new(Window::new(1 << 20)),
wake_tx: watch::channel(0).0,
chunk: AtomicU32::new(governor::START_CHUNK),
bundle: AtomicU32::new(governor::START_BUNDLE),
bytes_budget: Arc::new(Semaphore::new(READ_AHEAD_KIB as usize)),
stall: Mutex::new(None),
});
let taken = Arc::new(AtomicBool::new(false));
let mut f = test_frame(Chunk::TYPE, 1024);
let len = f.body.len() as u64;
f.taken = Some(taken.clone());
assert!(sh.window.lock().unwrap().sent(1, 7, len));
sh.sched.lock().unwrap().inflight.insert(7, (1, f));
let tx = ConnTx::new(outbox);
let death = tokio::spawn(lane_death(sh.clone(), 1, tx.clone()));
tokio::time::timeout(Duration::from_secs(10), async {
while !sh.sched.lock().unwrap().inflight.is_empty() {
tokio::time::sleep(Duration::from_millis(1)).await;
}
})
.await
.expect("the first sweep ran within the bound");
tx.send_raw_marked(Chunk::TYPE, 0, 99, vec![0x11; len as usize], taken.clone())
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(10), async {
while !taken.load(Ordering::SeqCst) {
tokio::time::sleep(Duration::from_millis(1)).await;
}
})
.await
.expect("the writer took the frame within the bound");
drop(link);
tokio::time::timeout(Duration::from_secs(10), death)
.await
.expect("the lane death handling finished within the bound")
.expect("the lane death handling did not panic");
assert_eq!(
sh.window.lock().unwrap().available(),
(1 << 20) - len,
"the frame's charge was released although the writer took it: the window was spent twice"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_pick_that_slips_past_the_stop_after_the_lane_death_strands_nothing() {
let timing = Timing {
ping_every: Duration::from_secs(3600),
dead_after: Duration::from_secs(3600),
handshake: Duration::from_secs(1),
min_frame_rate: crate::link::MIN_FRAME_RATE,
};
let (a, _b) = duplex(1 << 20);
let (ar, aw) = split(a);
let (tx, _rx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (link, outbox) =
crate::link::drive(FrameReader::new(ar), FrameWriter::new(aw), timing, tx);
let tx = ConnTx::new(outbox);
let running = Arc::new(AtomicBool::new(false));
tx.send_raw_marked(Chunk::TYPE, 0, 99, vec![0x11; 64], running.clone())
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(10), async {
while !running.load(Ordering::SeqCst) {
tokio::time::sleep(Duration::from_millis(1)).await;
}
})
.await
.expect("the writer started within the bound");
let sh = test_shared(governor::START_CHUNK, READ_AHEAD_KIB as usize);
let mut f = test_frame(Chunk::TYPE, 1024);
f.taken = Some(Arc::new(AtomicBool::new(false)));
assert!(sh.window.lock().unwrap().sent(1, 7, f.body.len() as u64));
sh.sched.lock().unwrap().inflight.insert(7, (1, f));
let stop = AtomicBool::new(false);
stop.store(true, Ordering::Relaxed);
drop(link);
tokio::time::timeout(Duration::from_secs(10), async {
while !tx.writer_dead() {
tokio::time::sleep(Duration::from_millis(1)).await;
}
})
.await
.expect("the writer died within the bound");
lane_death(sh.clone(), 1, tx).await;
{
let s = sh.sched.lock().unwrap();
assert!(
s.inflight.is_empty() && s.requeue.len() == 1,
"the sweep requeued it"
);
}
assert_eq!(sh.window.lock().unwrap().available(), 1 << 20);
assert!(
matches!(pick_for_lane(&sh, 1, &stop), Pick::Stopped),
"a stopped lane picked a frame after its death was handled"
);
{
let s = sh.sched.lock().unwrap();
assert!(
s.inflight.is_empty(),
"a frame was charged to a dead lane after its sweeps: no lane will ever carry it"
);
assert_eq!(s.requeue.len(), 1, "the frame stays queued for a live lane");
}
assert_eq!(sh.window.lock().unwrap().available(), 1 << 20);
assert!(matches!(
pick_for_lane(&sh, 2, &AtomicBool::new(false)),
Pick::Frame((1, _, _, _))
));
}
type LaneHarness = (
tokio::task::JoinHandle<()>,
Arc<Shared>,
tokio::sync::mpsc::UnboundedReceiver<Result<u32, ()>>,
Arc<AtomicBool>,
crate::link::Link,
);
fn lane_harness(credit: u64, chunk: u32) -> LaneHarness {
let timing = Timing {
ping_every: Duration::from_secs(3600),
dead_after: Duration::from_secs(3600),
handshake: Duration::from_secs(1),
min_frame_rate: crate::link::MIN_FRAME_RATE,
};
let (a, b) = duplex(1 << 20);
let (ar, aw) = split(a);
let (tx, _rx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (link, outbox) =
crate::link::drive(FrameReader::new(ar), FrameWriter::new(aw), timing, tx);
let mut peer = FrameReader::new(b);
peer.set_max_body(crate::frame::MAX_BODY);
let (drain_tx, drain_rx) = mpsc::unbounded_channel();
tokio::spawn(async move {
while let Ok(f) = peer.recv().await {
if is_data_type(f.ty) {
let _ = drain_tx.send(Ok(f.channel));
}
}
let _ = drain_tx.send(Err(()));
});
let router = Arc::new(Router::default());
router.lane_up(1, outbox);
let sh = Arc::new(Shared {
sched: Mutex::new(Sched {
floor: 4,
decision: Some(governor::Decision {
lanes: 1,
chunk: 4 << 20,
bundle: 1 << 20,
bottleneck: gen::BN_NETWORK,
mode: Mode::StreamOnly,
prefer: Class::Stream,
sequential: false,
}),
..Default::default()
}),
window: Mutex::new(Window::new(credit)),
wake_tx: watch::channel(0).0,
chunk: AtomicU32::new(chunk),
bundle: AtomicU32::new(governor::START_BUNDLE),
bytes_budget: Arc::new(Semaphore::new(READ_AHEAD_KIB as usize)),
stall: Mutex::new(None),
});
let stop = Arc::new(AtomicBool::new(false));
let h = tokio::spawn(lane_task(
router.lane(1).unwrap(),
sh.clone(),
None,
stop.clone(),
));
(h, sh, drain_rx, stop, link)
}
#[tokio::test]
async fn a_lane_blocked_by_its_inflight_cap_is_backpressure_never_a_stall() {
let (h, sh, mut wire, stop, _link) = lane_harness(16 << 20, governor::START_CHUNK);
assert!(sh.window.lock().unwrap().sent(1, 77, 8 << 20));
sh.sched
.lock()
.unwrap()
.chunks
.push_back(test_frame(Chunk::TYPE, 1 << 20));
sh.wake();
tokio::time::sleep(Duration::from_millis(300)).await;
assert!(
sh.stall.lock().unwrap().is_none(),
"a lane blocked by its in-flight cap marked a fatal stall"
);
let got = wire.try_recv();
assert!(got.is_err(), "the cap-blocked lane sent a frame: {got:?}");
assert_eq!(sh.window.lock().unwrap().received(77), Some(8 << 20));
sh.wake();
tokio::time::timeout(Duration::from_secs(5), wire.recv())
.await
.expect("the lane resumed once the Received arrived")
.expect("a frame arrived")
.expect("the wire carried a frame");
stop.store(true, Ordering::Relaxed);
sh.wake();
tokio::time::timeout(Duration::from_secs(1), h)
.await
.expect("the lane task stopped")
.unwrap();
}
#[tokio::test]
async fn a_frame_the_window_cannot_hold_marks_a_credit_stall() {
let (h, sh, _wire, stop, _link) = lane_harness(512 << 10, governor::START_CHUNK);
sh.sched
.lock()
.unwrap()
.chunks
.push_back(test_frame(Chunk::TYPE, 1 << 20)); sh.wake();
let st = tokio::time::timeout(Duration::from_secs(5), async {
loop {
let st = *sh.stall.lock().unwrap();
if let Some(st) = st {
return st;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the credit stall was recorded within the bound");
assert_eq!(st.grant, 512 << 10);
assert!(st.smallest >= 1 << 20, "smallest: {}", st.smallest);
stop.store(true, Ordering::Relaxed);
sh.wake();
tokio::time::timeout(Duration::from_secs(1), h)
.await
.expect("the lane task stopped")
.unwrap();
}
#[tokio::test]
async fn an_injected_receiver_error_ends_the_job_promptly_with_no_lane_left_sending() {
let timing = Timing {
ping_every: Duration::from_secs(3600),
dead_after: Duration::from_secs(3600),
handshake: Duration::from_secs(5),
min_frame_rate: crate::link::MIN_FRAME_RATE,
};
let job = [0x79u8; 16];
let (ca, cb) = duplex(1 << 20);
let (car, caw) = split(ca);
let (cbr, cbw) = split(cb);
let (ctx, mut crx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (_clink, coutbox) =
crate::link::drive(FrameReader::new(car), FrameWriter::new(caw), timing, ctx);
let (la, lb) = duplex(1 << 20);
let (lar, law) = split(la);
let (lbr, _lbw) = split(lb);
let (ltx, _lrx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (_llink, loutbox) =
crate::link::drive(FrameReader::new(lar), FrameWriter::new(law), timing, ltx);
let router = Arc::new(Router::default());
router.lane_up(1, loutbox);
let mut link = JobLink::new(job, router.clone(), ConnTx::new(coutbox), None);
let r2 = router.clone();
tokio::spawn(async move {
while let Some(f) = crx.recv().await {
if is_data_type(f.ty) {
let _ = r2.route_control(f).await;
}
}
r2.close("the session ended");
});
let lane_seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let ls = lane_seen.clone();
tokio::spawn(async move {
let mut peer = FrameReader::new(cbr);
peer.set_max_body(crate::frame::MAX_BODY);
let mut lane_peer = FrameReader::new(lbr);
lane_peer.set_max_body(crate::frame::MAX_BODY);
let mut peer_w = FrameWriter::new(cbw);
loop {
tokio::select! {
f = peer.recv() => match f {
Ok(f) if f.ty == JobOpen::TYPE => {
let open: JobOpen = f.decode().unwrap();
peer_w.send_msg(0, &JobOpenAck {
job_id: open.job_id,
status: 0,
credit: 64 << 20,
staged: 1,
workers: 4,
message: None,
}).await.unwrap();
}
Ok(f) if f.ty == ManifestEnd::TYPE => {
peer_w.send_msg(0, &JobMap {
job_id: job,
status: 0,
last: 1,
done: vec![],
partial: vec![],
message: None,
held: None,
}).await.unwrap();
}
Ok(_) => {}
Err(_) => return,
},
f = lane_peer.recv() => match f {
Ok(f) if is_data_type(f.ty) => {
let n = ls.fetch_add(1, Ordering::Relaxed);
let _ = peer_w.send_msg(0, &Received { job_id: job, lane: 1, seq: f.channel }).await;
if n == 0 {
peer_w.send(Received::TYPE, 0, &job).await.unwrap();
}
}
Ok(_) => {}
Err(_) => return,
},
}
}
});
let dir = std::env::temp_dir().join(format!("ava1-send-err-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("f"), vec![0x33u8; 64 << 20]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: 64 << 20,
mtime: 1,
path: "f".into(),
root: None,
}],
});
let result = tokio::time::timeout(
Duration::from_secs(15),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await;
let err = result
.expect("the job ended promptly")
.expect_err("the injected error surfaces");
assert!(matches!(err, SendError::Protocol(_)), "{err:?}");
let n = lane_seen.load(Ordering::Relaxed);
tokio::time::sleep(Duration::from_millis(300)).await;
assert_eq!(
lane_seen.load(Ordering::Relaxed),
n,
"a lane task kept sending after the error"
);
let _ = std::fs::remove_dir_all(&dir);
}
struct FakeReceiver {
credit: u64,
credit_on_apply: bool,
done_when_complete: bool,
credit_after_ack: Option<u64>,
malformed_status: bool,
retry_unknown: bool,
malformed_status_open: bool,
wellformed_status_open: bool,
pause_control: Option<Duration>,
control_pipe: usize,
cancel_seen: Arc<AtomicU32>,
no_map: bool,
}
impl Default for FakeReceiver {
fn default() -> Self {
Self {
credit: 0,
credit_on_apply: false,
done_when_complete: false,
credit_after_ack: None,
malformed_status: false,
retry_unknown: false,
malformed_status_open: false,
wellformed_status_open: false,
pause_control: None,
control_pipe: 1 << 20,
cancel_seen: Arc::new(AtomicU32::new(0)),
no_map: false,
}
}
}
fn fake_link(
rcv: FakeReceiver,
opener: Option<Arc<dyn LaneOpener>>,
) -> (
JobLink,
Arc<std::sync::atomic::AtomicUsize>,
Vec<crate::link::Link>,
) {
let timing = Timing {
ping_every: Duration::from_secs(3600),
dead_after: Duration::from_secs(3600),
handshake: Duration::from_secs(5),
min_frame_rate: crate::link::MIN_FRAME_RATE,
};
let job = [0x7d; 16];
let (ca, cb) = duplex(rcv.control_pipe);
let (car, caw) = split(ca);
let (cbr, cbw) = split(cb);
let (ctx, mut crx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (clink, coutbox) =
crate::link::drive(FrameReader::new(car), FrameWriter::new(caw), timing, ctx);
let (la, lb) = duplex(1 << 20);
let (lar, law) = split(la);
let (lbr, _lbw) = split(lb);
let (ltx, _lrx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (llink, loutbox) =
crate::link::drive(FrameReader::new(lar), FrameWriter::new(law), timing, ltx);
let keep: Vec<crate::link::Link> = vec![clink, llink];
let router = Arc::new(Router::default());
router.lane_up(1, loutbox);
let link = JobLink::new(job, router.clone(), ConnTx::new(coutbox), opener);
let r2 = router.clone();
tokio::spawn(async move {
while let Some(f) = crx.recv().await {
if is_data_type(f.ty) {
let _ = r2.route_control(f).await;
}
}
r2.close("the session ended");
});
let lane_seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let ls = lane_seen.clone();
tokio::spawn(async move {
fn pump<R: tokio::io::AsyncRead + Unpin + Send + 'static>(
r: R,
) -> mpsc::Receiver<Result<Frame, ()>> {
let (tx, rx) = mpsc::channel(1);
tokio::spawn(async move {
let mut reader = FrameReader::new(r);
reader.set_max_body(crate::frame::MAX_BODY);
loop {
let f = reader.recv().await.map_err(|_| ());
let end = f.is_err();
if tx.send(f).await.is_err() || end {
return;
}
}
});
rx
}
let mut peer = pump(cbr);
let mut lane_peer = pump(lbr);
let mut peer_w = FrameWriter::new(cbw);
let cancel_seen = rcv.cancel_seen.clone();
let mut control_hold = Box::pin(tokio::time::sleep(Duration::ZERO));
let mut flood = Box::pin(tokio::time::sleep(Duration::from_secs(3600)));
let mut flood_started = false;
let (mut total, mut manifest_bytes, mut manifest_files) = (0u64, 0u64, 0u32);
loop {
tokio::select! {
f = async {
(&mut control_hold).await;
peer.recv().await.unwrap_or(Err(()))
} => match f {
Ok(f) if f.ty == JobOpen::TYPE => {
let open: JobOpen = f.decode().unwrap();
peer_w.send_msg(0, &JobOpenAck {
job_id: open.job_id,
status: 0,
credit: rcv.credit,
staged: 1,
workers: 4,
message: None,
}).await.unwrap();
if let Some(bytes) = rcv.credit_after_ack {
let _ = peer_w.send_msg(0, &Credit { job_id: job, bytes }).await;
}
if rcv.malformed_status_open {
let _ = peer_w.send(Status::TYPE, 0, &job).await;
}
if rcv.wellformed_status_open {
let _ = peer_w.send_msg(0, &Status {
job_id: job,
..Default::default()
}).await;
}
}
Ok(f) if f.ty == ManifestEnd::TYPE => {
let me: ManifestEnd = f.decode().unwrap();
manifest_bytes = me.bytes;
manifest_files = me.files;
if rcv.no_map {
continue;
}
peer_w.send_msg(0, &JobMap {
job_id: job,
status: 0,
last: 1,
done: vec![],
partial: vec![],
message: None,
held: None,
}).await.unwrap();
if let Some(p) = rcv.pause_control {
control_hold.as_mut().reset(tokio::time::Instant::now() + p);
flood.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(1));
flood_started = true;
}
if rcv.retry_unknown {
let _ = peer_w.send_msg(0, &FileRetry {
job_id: job,
file_id: 99,
reason: 1,
}).await;
}
}
Ok(f) if f.ty == gen::JobCancel::TYPE => {
let c: gen::JobCancel = f.decode().unwrap();
cancel_seen.store(c.reason as u32, Ordering::SeqCst);
}
Ok(_) => {}
Err(_) => return,
},
_ = &mut flood, if flood_started => {
let _ = peer_w.send_msg(0, &Status {
job_id: job,
current: Some("s".repeat(32 << 10)),
..Default::default()
}).await;
flood.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(1));
}
f = async { lane_peer.recv().await.unwrap_or(Err(())) } => match f {
Ok(f) if is_data_type(f.ty) => {
let n = ls.fetch_add(1, Ordering::Relaxed);
let payload = f.decode::<Chunk>().map(|c| c.data.len() as u64).unwrap_or(0);
let _ = peer_w.send_msg(0, &Received { job_id: job, lane: 1, seq: f.channel }).await;
if rcv.credit_on_apply {
let _ = peer_w.send_msg(0, &Credit { job_id: job, bytes: f.body.len() as u64 }).await;
}
if rcv.malformed_status && n == 0 {
let _ = peer_w.send(Status::TYPE, 0, &job).await;
}
total += payload;
if rcv.done_when_complete && total >= manifest_bytes {
let _ = peer_w.send_msg(0, &JobDone {
job_id: job,
status: 0,
files: manifest_files,
bytes: manifest_bytes,
message: None, settling: None,}).await;
}
}
Ok(_) => {}
Err(_) => return,
},
}
}
});
(link, lane_seen, keep)
}
struct CountingRead {
data: Arc<Vec<u8>>,
reads: Arc<std::sync::atomic::AtomicUsize>,
active: Arc<std::sync::atomic::AtomicUsize>,
slow: bool,
}
impl crate::source::ReadAt for CountingRead {
fn read_at(&mut self, off: u64, buf: &mut [u8]) -> io::Result<usize> {
if self.slow {
std::thread::sleep(Duration::from_millis(5)); }
let n = buf
.len()
.min(64 * 1024)
.min(self.data.len().saturating_sub(off as usize));
if n > 0 {
buf[..n].copy_from_slice(&self.data[off as usize..off as usize + n]);
}
self.reads.fetch_add(n, Ordering::Relaxed);
Ok(n)
}
}
impl Drop for CountingRead {
fn drop(&mut self) {
self.active.fetch_sub(1, Ordering::Relaxed);
}
}
struct CountingSource {
data: Arc<Vec<u8>>,
reads: Arc<std::sync::atomic::AtomicUsize>,
active: Arc<std::sync::atomic::AtomicUsize>,
slow: bool,
}
impl Source for CountingSource {
fn open(&self, _rel: &str) -> io::Result<Box<dyn crate::source::ReadAt>> {
self.active.fetch_add(1, Ordering::Relaxed);
Ok(Box::new(CountingRead {
data: self.data.clone(),
reads: self.reads.clone(),
active: self.active.clone(),
slow: self.slow,
}))
}
fn list(&self, _rel: &str) -> io::Result<Vec<(String, SourceMeta)>> {
Ok(Vec::new())
}
fn stat(&self, _rel: &str) -> io::Result<SourceMeta> {
Ok(SourceMeta {
size: self.data.len() as u64,
mtime: 0,
mode: 0o644,
is_dir: false,
})
}
}
struct FailingAfter {
data: Vec<u8>,
opens: Arc<std::sync::atomic::AtomicUsize>,
active: Arc<std::sync::atomic::AtomicUsize>,
fail_after: usize,
}
impl Source for FailingAfter {
fn open(&self, _rel: &str) -> io::Result<Box<dyn crate::source::ReadAt>> {
if self.opens.fetch_add(1, Ordering::SeqCst) >= self.fail_after {
return Err(std::io::Error::other("the test source failed"));
}
std::thread::sleep(Duration::from_millis(1)); self.active.fetch_add(1, Ordering::SeqCst);
Ok(Box::new(CountingRead {
data: Arc::new(self.data.clone()),
reads: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
active: self.active.clone(),
slow: false,
}))
}
fn list(&self, _rel: &str) -> io::Result<Vec<(String, SourceMeta)>> {
Ok(Vec::new())
}
fn stat(&self, _rel: &str) -> io::Result<SourceMeta> {
Ok(SourceMeta {
size: self.data.len() as u64,
mtime: 0,
mode: 0o644,
is_dir: false,
})
}
}
#[tokio::test(flavor = "multi_thread")]
async fn a_cancelled_job_stops_the_readers_and_leaves_none_parked() {
let size = 160usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 4 << 20,
credit_on_apply: true,
done_when_complete: false,
credit_after_ack: None,
malformed_status: false,
retry_unknown: false,
..Default::default()
},
None,
);
let reads = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let src: Arc<dyn Source> = Arc::new(CountingSource {
data: Arc::new(vec![0x5a; size]),
reads: reads.clone(),
active: active.clone(),
slow: true,
});
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let cancel = Arc::new(AtomicBool::new(false));
let mut opts = SendOptions::upload("dest");
opts.cancel = cancel.clone();
let job_task = tokio::spawn(async move { send_job(&mut link, m, src, opts).await });
tokio::time::timeout(Duration::from_secs(12), async {
while reads.load(Ordering::Relaxed) < 2 * (4 << 20) {
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("the reader made progress within the bound");
cancel.store(true, Ordering::Relaxed);
let result = tokio::time::timeout(Duration::from_secs(15), job_task)
.await
.expect("the sender task completed within the bound")
.expect("the sender task did not panic");
assert!(matches!(result, Err(SendError::Cancelled)), "{result:?}");
let at = reads.load(Ordering::Relaxed);
tokio::time::sleep(Duration::from_secs(1)).await;
assert_eq!(
reads.load(Ordering::Relaxed),
at,
"a reader kept reading after the job returned"
);
assert_eq!(active.load(Ordering::Relaxed), 0, "a reader thread leaked");
}
struct ParkedSource {
closed: Arc<AtomicBool>,
parked: Arc<AtomicBool>,
}
struct ParkedRead(Arc<AtomicBool>, Arc<AtomicBool>);
impl crate::source::ReadAt for ParkedRead {
fn read_at(&mut self, _off: u64, _buf: &mut [u8]) -> std::io::Result<usize> {
self.1.store(true, Ordering::Relaxed);
while !self.0.load(Ordering::Relaxed) {
std::thread::sleep(Duration::from_millis(10));
}
Err(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"closed",
))
}
}
impl Source for ParkedSource {
fn open(&self, _rel: &str) -> std::io::Result<Box<dyn crate::source::ReadAt>> {
Ok(Box::new(ParkedRead(
self.closed.clone(),
self.parked.clone(),
)))
}
fn list(&self, _rel: &str) -> std::io::Result<Vec<(String, SourceMeta)>> {
Ok(Vec::new())
}
fn stat(&self, _rel: &str) -> std::io::Result<SourceMeta> {
Ok(SourceMeta {
size: 64 << 20,
mtime: 0,
mode: 0o644,
is_dir: false,
})
}
fn close(&self) {
self.closed.store(true, Ordering::Relaxed);
}
}
#[tokio::test(flavor = "multi_thread")]
async fn teardown_closes_the_source_so_a_parked_reader_cannot_hold_the_job() {
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 4 << 20,
credit_on_apply: true,
done_when_complete: false,
credit_after_ack: None,
malformed_status: false,
retry_unknown: false,
..Default::default()
},
None,
);
let (closed, parked) = (
Arc::new(AtomicBool::new(false)),
Arc::new(AtomicBool::new(false)),
);
let src: Arc<dyn Source> = Arc::new(ParkedSource {
closed: closed.clone(),
parked: parked.clone(),
});
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: 64 << 20,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let cancel = Arc::new(AtomicBool::new(false));
let mut opts = SendOptions::upload("dest");
opts.cancel = cancel.clone();
let job_task = tokio::spawn(async move { send_job(&mut link, m, src, opts).await });
tokio::time::timeout(Duration::from_secs(10), async {
while !parked.load(Ordering::Relaxed) {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("a reader parked in the source");
cancel.store(true, Ordering::Relaxed);
let result = tokio::time::timeout(Duration::from_secs(5), job_task)
.await
.expect("teardown waited on a parked reader")
.unwrap();
assert!(matches!(result, Err(SendError::Cancelled)), "{result:?}");
assert!(closed.load(Ordering::Relaxed));
}
struct FailOpenOnSignal(Arc<tokio::sync::Notify>);
impl LaneOpener for FailOpenOnSignal {
fn open(&self) -> BoxFut<'_, Result<u16, crate::Ava1Error>> {
let n = self.0.clone();
Box::pin(async move {
n.notified().await;
Err(crate::Ava1Error::Io(std::io::Error::other(
"the test opener refuses the lane",
)))
})
}
fn close(&self, _id: u16) {}
}
#[tokio::test(flavor = "multi_thread")]
async fn a_lane_open_failure_still_runs_the_teardown_and_joins_the_readers() {
let size = 160usize << 20;
let signal = Arc::new(tokio::sync::Notify::new());
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 4 << 20,
credit_on_apply: false,
done_when_complete: false,
credit_after_ack: None,
malformed_status: false,
retry_unknown: false,
..Default::default()
},
Some(Arc::new(FailOpenOnSignal(signal.clone())) as Arc<dyn LaneOpener>),
);
let reads = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let src: Arc<dyn Source> = Arc::new(CountingSource {
data: Arc::new(vec![0x5b; size]),
reads: reads.clone(),
active: active.clone(),
slow: true,
});
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let job_task =
tokio::spawn(
async move { send_job(&mut link, m, src, SendOptions::upload("dest")).await },
);
tokio::time::timeout(Duration::from_secs(120), async {
loop {
let a = reads.load(Ordering::Relaxed);
tokio::time::sleep(Duration::from_millis(150)).await;
if a > 0 && reads.load(Ordering::Relaxed) == a && active.load(Ordering::Relaxed) > 0
{
break;
}
}
})
.await
.expect("a reader parked on the budget within the bound");
signal.notify_one();
let result = tokio::time::timeout(Duration::from_secs(15), job_task)
.await
.expect("the job ended promptly after the open failure")
.expect("the sender task did not panic");
assert!(
matches!(result, Err(SendError::Disconnected(_))),
"{result:?}"
);
let at = reads.load(Ordering::Relaxed);
tokio::time::sleep(Duration::from_secs(1)).await;
assert_eq!(
reads.load(Ordering::Relaxed),
at,
"a reader kept reading after the job returned"
);
assert_eq!(active.load(Ordering::Relaxed), 0, "a reader thread leaked");
}
struct LaneFactory {
reqs: mpsc::UnboundedSender<tokio::sync::oneshot::Sender<u16>>,
}
impl LaneOpener for LaneFactory {
fn open(&self) -> BoxFut<'_, Result<u16, crate::Ava1Error>> {
let reqs = self.reqs.clone();
Box::pin(async move {
let (tx, rx) = tokio::sync::oneshot::channel();
if reqs.send(tx).is_err() {
return Err(crate::Ava1Error::Io(std::io::Error::other(
"the test's lane factory ended",
)));
}
rx.await
.map_err(|_| crate::Ava1Error::Lost("the test's lane factory ended".into()))
})
}
fn close(&self, _id: u16) {}
}
#[tokio::test(flavor = "multi_thread")]
async fn repeated_lane_kills_with_frames_stuck_in_the_outbox_never_wedge_the_window() {
let size = 256usize << 20;
let credit: u64 = 32 << 20;
let reads_on = Arc::new(AtomicBool::new(true));
let (ack_tx, mut ack_rx) = mpsc::unbounded_channel::<(u32, u64, u32, u64, u64)>();
let timing = Timing {
ping_every: Duration::from_secs(3600),
dead_after: Duration::from_secs(3600),
handshake: Duration::from_secs(5),
min_frame_rate: crate::link::MIN_FRAME_RATE,
};
let job = [0x7eu8; 16];
let (ca, cb) = duplex(1 << 20);
let (car, caw) = split(ca);
let (cbr, cbw) = split(cb);
let (ctx, mut crx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (clink, coutbox) =
crate::link::drive(FrameReader::new(car), FrameWriter::new(caw), timing, ctx);
let router = Arc::new(Router::default());
let (la, lb) = duplex(1 << 20);
let (lar, law) = split(la);
let (lbr, _lbw) = split(lb);
let (ltx, mut lrx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (llink1, loutbox) =
crate::link::drive(FrameReader::new(lar), FrameWriter::new(law), timing, ltx);
let gen1 = router.lane_up(1, loutbox);
let r1 = router.clone();
tokio::spawn(async move {
while lrx.recv().await.is_some() {}
r1.lane_down(1, gen1);
});
{
let reads_on = reads_on.clone();
let ack_tx = ack_tx.clone();
tokio::spawn(async move {
let mut peer = FrameReader::new(lbr);
peer.set_max_body(crate::frame::MAX_BODY);
loop {
while !reads_on.load(Ordering::Relaxed) {
tokio::time::sleep(Duration::from_millis(10)).await;
}
match peer.recv().await {
Ok(f) if is_data_type(f.ty) => {
let c: Chunk = f.decode().unwrap();
let _ = ack_tx.send((
f.channel,
f.body.len() as u64,
c.file_id,
c.offset,
c.data.len() as u64,
));
}
Ok(_) => {} Err(_) => return,
}
}
});
}
tokio::spawn(async move {
let mut peer = FrameReader::new(cbr);
peer.set_max_body(crate::frame::MAX_BODY);
let mut peer_w = FrameWriter::new(cbw);
let (mut manifest_bytes, mut manifest_files) = (0u64, 0u32);
let mut covered = RangeSet::new();
loop {
tokio::select! {
f = peer.recv() => match f {
Ok(f) if f.ty == JobOpen::TYPE => {
let open: JobOpen = f.decode().unwrap();
peer_w.send_msg(0, &JobOpenAck {
job_id: open.job_id,
status: 0,
credit,
staged: 1,
workers: 4,
message: None,
}).await.unwrap();
}
Ok(f) if f.ty == ManifestEnd::TYPE => {
let me: ManifestEnd = f.decode().unwrap();
manifest_bytes = me.bytes;
manifest_files = me.files;
peer_w.send_msg(0, &JobMap {
job_id: job,
status: 0,
last: 1,
done: vec![],
partial: vec![],
message: None,
held: None,
}).await.unwrap();
}
Ok(_) => {}
Err(_) => return,
},
ack = ack_rx.recv() => match ack {
Some((seq, len, _file_id, offset, data_len)) => {
tokio::time::sleep(Duration::from_millis(100)).await;
let _ = peer_w.send_msg(0, &Received { job_id: job, lane: 1, seq }).await;
let _ = peer_w.send_msg(0, &Credit { job_id: job, bytes: len }).await;
covered.insert(offset, offset + data_len);
if manifest_bytes > 0 && covered.covered() >= manifest_bytes {
let _ = peer_w.send_msg(0, &JobDone {
job_id: job,
status: 0,
files: manifest_files,
bytes: manifest_bytes,
message: None, settling: None,}).await;
}
}
None => return,
},
}
}
});
let r2 = router.clone();
tokio::spawn(async move {
while let Some(f) = crx.recv().await {
if is_data_type(f.ty) {
let _ = r2.route_control(f).await;
}
}
r2.close("the session ended");
});
let (req_tx, mut req_rx) = mpsc::unbounded_channel::<tokio::sync::oneshot::Sender<u16>>();
let (new_tx, mut new_rx) = mpsc::unbounded_channel::<(u16, crate::link::Link)>();
let opener = Arc::new(LaneFactory { reqs: req_tx }) as Arc<dyn LaneOpener>;
let mut link = JobLink::new(job, router.clone(), ConnTx::new(coutbox), Some(opener));
let (reads_on_f, ack_tx_f) = (reads_on.clone(), ack_tx.clone());
tokio::spawn(async move {
let mut next_id = 2u16;
while let Some(tx) = req_rx.recv().await {
let id = next_id;
next_id += 1;
let (la, lb) = duplex(1 << 20);
let (lar, law) = split(la);
let (lbr, _lbw) = split(lb);
let (ltx, mut lrx) = mpsc::channel(crate::link::DELIVER_DEPTH);
let (llink, loutbox) =
crate::link::drive(FrameReader::new(lar), FrameWriter::new(law), timing, ltx);
let gen = router.lane_up(id, loutbox);
let r3 = router.clone();
tokio::spawn(async move {
while lrx.recv().await.is_some() {}
r3.lane_down(id, gen);
});
tokio::spawn({
let reads_on = reads_on_f.clone();
let ack_tx = ack_tx_f.clone();
async move {
let mut peer = FrameReader::new(lbr);
peer.set_max_body(crate::frame::MAX_BODY);
loop {
while !reads_on.load(Ordering::Relaxed) {
tokio::time::sleep(Duration::from_millis(10)).await;
}
match peer.recv().await {
Ok(f) if is_data_type(f.ty) => {
let c: Chunk = f.decode().unwrap();
let _ = ack_tx.send((
f.channel,
f.body.len() as u64,
c.file_id,
c.offset,
c.data.len() as u64,
));
}
Ok(_) => {} Err(_) => return,
}
}
}
});
let _ = new_tx.send((id, llink));
let _ = tx.send(id);
}
});
let reads = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let src: Arc<dyn Source> = Arc::new(CountingSource {
data: Arc::new(vec![0x5c; size]),
reads: reads.clone(),
active: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
slow: false,
});
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let job_task =
tokio::spawn(
async move { send_job(&mut link, m, src, SendOptions::upload("dest")).await },
);
tokio::time::timeout(Duration::from_secs(30), async {
let start = Instant::now();
while reads.load(Ordering::Relaxed) < 24 * (1 << 20)
|| start.elapsed() < Duration::from_secs(3)
{
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("the job warmed up within the bound");
reads_on.store(false, Ordering::Relaxed);
tokio::time::sleep(Duration::from_millis(700)).await;
drop(llink1);
let (_, kill2) = tokio::time::timeout(Duration::from_secs(15), new_rx.recv())
.await
.expect("lane 2's Link arrived within the bound")
.expect("the lane factory is alive");
tokio::time::sleep(Duration::from_millis(700)).await;
drop(kill2);
let (_, kill3) = tokio::time::timeout(Duration::from_secs(15), new_rx.recv())
.await
.expect("a replacement lane came up within the bound")
.expect("the lane factory is alive");
tokio::time::sleep(Duration::from_millis(700)).await;
drop(kill3);
reads_on.store(true, Ordering::Relaxed);
let report = tokio::time::timeout(Duration::from_secs(90), job_task)
.await
.expect("the job completed within the bound")
.expect("the sender task did not panic")
.expect("the upload completed");
assert_eq!(report.status, 0, "{:?}", report.message);
assert_eq!(report.bytes, size as u64);
assert!(
report.resent > 0,
"the dead lanes' frames were requeued and re-sent"
);
let _keep_control = clink;
}
struct HungOpener(Arc<AtomicUsize>);
impl LaneOpener for HungOpener {
fn open(&self) -> BoxFut<'_, Result<u16, crate::Ava1Error>> {
self.0.fetch_add(1, Ordering::SeqCst);
Box::pin(std::future::pending())
}
fn close(&self, _id: u16) {}
}
#[tokio::test(flavor = "multi_thread")]
async fn a_hung_lane_open_never_stops_the_loop_draining_received_and_credit() {
let size = 24usize << 20;
let opens = Arc::new(AtomicUsize::new(0));
let (mut link, lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 4 << 20,
credit_on_apply: true,
done_when_complete: true,
..Default::default()
},
Some(Arc::new(HungOpener(opens.clone())) as Arc<dyn LaneOpener>),
);
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let src: Arc<dyn Source> = Arc::new(MemSource(vec![0x6c; size]));
let report = tokio::time::timeout(
Duration::from_secs(60),
send_job(&mut link, m, src, SendOptions::upload("dest")),
)
.await
.expect("the job completed while the lane open hung")
.expect("the upload completed");
assert_eq!(report.status, 0, "{:?}", report.message);
assert_eq!(report.bytes, size as u64);
assert!(
lane_seen.load(Ordering::SeqCst) >= 6,
"every chunk went on lane 1"
);
assert_eq!(
opens.load(Ordering::SeqCst),
1,
"at most one open in flight"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_job_completes_when_the_grant_is_smaller_than_the_chunk() {
let size = 12usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 3 << 20,
credit_on_apply: true,
done_when_complete: true,
credit_after_ack: None,
malformed_status: false,
retry_unknown: false,
..Default::default()
},
None,
);
let dir = std::env::temp_dir().join(format!("ava1-send-grant-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("big"), vec![0x44u8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let report = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await
.expect("the job completed within the bound (a grant below the chunk must not hang it)")
.expect("the upload completed");
assert_eq!(report.status, 0, "{:?}", report.message);
assert_eq!(report.bytes, size as u64);
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_finished_job_leaves_its_telemetry_on_the_progress() {
let size = 3usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 16 << 20,
credit_on_apply: true,
done_when_complete: true,
credit_after_ack: None,
malformed_status: false,
retry_unknown: false,
..Default::default()
},
None,
);
let dir = std::env::temp_dir().join(format!("ava1-send-telem-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("f"), vec![0x11u8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "f".into(),
root: None,
}],
});
let opts = SendOptions::upload("dest");
let progress = opts.progress.clone();
let report = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
opts,
),
)
.await
.expect("bounded")
.expect("the upload completed");
assert_eq!(report.status, 0, "{:?}", report.message);
assert_eq!(progress.attempts.load(Ordering::Relaxed), 1);
let t = progress.telemetry.lock().unwrap().clone();
assert_eq!(
t.shares.history.len() as u32,
t.shares.ticks.min(super::governor::HISTORY_MAX as u32)
);
assert_eq!(progress.unswept_peak.load(Ordering::Relaxed), 0);
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_grant_smaller_than_one_group_fails_loudly_instead_of_hanging() {
let size = 8usize << 20;
let cancel_seen = Arc::new(AtomicU32::new(0));
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 512 << 10,
credit_on_apply: true,
done_when_complete: true,
credit_after_ack: None,
malformed_status: false,
retry_unknown: false,
cancel_seen: cancel_seen.clone(),
..Default::default()
},
None,
);
let dir = std::env::temp_dir().join(format!("ava1-send-tinygrant-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("big"), vec![0x45u8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let err = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await
.expect("the job ended within the bound (it must not hang)")
.expect_err("a window below one group cannot send anything");
assert!(
matches!(&err, SendError::Protocol(msg) if msg.contains("524288") && msg.contains("window")),
"{err:?}"
);
tokio::time::timeout(Duration::from_secs(10), async {
while cancel_seen.load(Ordering::SeqCst) != gen::ERR_STALLED as u32 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the receiver saw the JobCancel (reason ERR_STALLED) within the bound");
let _ = std::fs::remove_dir_all(&dir);
}
async fn quiet_receiver(tag: &str, no_map: bool) {
let cancel_seen = Arc::new(AtomicU32::new(0));
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 16 << 20,
credit_on_apply: true,
done_when_complete: false,
cancel_seen: cancel_seen.clone(),
no_map,
..Default::default()
},
None,
);
let dir = std::env::temp_dir().join(format!("ava1-send-{tag}-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("f"), vec![0x51u8; 4096]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: 4096,
mtime: 1,
path: "f".into(),
root: None,
}],
});
let mut o = SendOptions::upload("dest");
o.receiver_silent_timeout = Some(Duration::from_secs(2));
let err = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
o,
),
)
.await
.expect("the sender waited forever on a receiver that went quiet")
.expect_err("the job never ended on the receiver");
assert!(
matches!(&err, SendError::Protocol(msg) if msg.contains("said nothing")),
"{err:?}"
);
tokio::time::timeout(Duration::from_secs(10), async {
while cancel_seen.load(Ordering::SeqCst) != gen::ERR_STALLED as u32 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the receiver saw the JobCancel (reason ERR_STALLED)");
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_receiver_that_never_sends_the_map_does_not_hold_the_sender_forever() {
quiet_receiver("quiet-map", true).await;
}
#[tokio::test(flavor = "multi_thread")]
async fn a_receiver_that_goes_quiet_after_the_data_does_not_hold_the_sender_forever() {
quiet_receiver("quiet-done", false).await;
}
#[tokio::test(flavor = "multi_thread")]
async fn a_credit_between_the_ack_and_the_map_is_not_discarded() {
let size = 8usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 0,
credit_on_apply: true,
done_when_complete: true,
credit_after_ack: Some(4 << 20),
malformed_status: false,
retry_unknown: false,
..Default::default()
},
None,
);
let dir =
std::env::temp_dir().join(format!("ava1-send-earlycredit-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("big"), vec![0x46u8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let report = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await
.expect("the job completed within the bound")
.expect("the upload completed");
assert_eq!(report.status, 0, "{:?}", report.message);
assert_eq!(report.bytes, size as u64);
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_malformed_status_is_a_protocol_error() {
let size = 4usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 64 << 20,
credit_on_apply: true,
done_when_complete: true,
credit_after_ack: None,
malformed_status: true,
retry_unknown: false,
..Default::default()
},
None,
);
let dir = std::env::temp_dir().join(format!("ava1-send-badstatus-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("big"), vec![0x47u8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let err = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await
.expect("the job ended within the bound")
.expect_err("a malformed Status ends the job");
assert!(matches!(err, SendError::Protocol(_)), "{err:?}");
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_malformed_status_between_the_ack_and_the_map_is_a_protocol_error() {
let size = 4usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 64 << 20,
credit_on_apply: true,
done_when_complete: true,
malformed_status_open: true,
..Default::default()
},
None,
);
let dir =
std::env::temp_dir().join(format!("ava1-send-badstatusopen-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("big"), vec![0x49u8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let err = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await
.expect("the job ended within the bound")
.expect_err("a malformed Status in the open window ends the job");
assert!(matches!(err, SendError::Protocol(_)), "{err:?}");
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_well_formed_status_between_the_ack_and_the_map_is_absorbed() {
let size = 4usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 64 << 20,
credit_on_apply: true,
done_when_complete: true,
wellformed_status_open: true,
..Default::default()
},
None,
);
let dir =
std::env::temp_dir().join(format!("ava1-send-okstatusopen-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("big"), vec![0x4au8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let report = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await
.expect("the job completed within the bound")
.expect("the upload completed");
assert_eq!(report.status, 0, "{:?}", report.message);
assert_eq!(report.bytes, size as u64);
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_reader_failure_sends_the_cancel_and_ends_the_job_without_deadlocking_on_a_full_outbox(
) {
let cancel_seen = Arc::new(AtomicU32::new(0));
let (mut link, lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 64 << 20,
credit_on_apply: true,
done_when_complete: false,
pause_control: Some(Duration::from_millis(1200)),
control_pipe: 4 << 10,
cancel_seen: cancel_seen.clone(),
..Default::default()
},
None,
);
let opens = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let src: Arc<dyn Source> = Arc::new(FailingAfter {
data: vec![0x5d; 256 << 10],
opens: opens.clone(),
active: active.clone(),
fail_after: 150,
});
let m = Arc::new(Manifest {
entries: (0..300)
.map(|i| Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: 256 << 10,
mtime: 1,
path: format!("f{i:03}"),
root: None,
})
.collect(),
});
let job_task =
tokio::spawn(
async move { send_job(&mut link, m, src, SendOptions::upload("dest")).await },
);
let err = tokio::time::timeout(Duration::from_secs(30), job_task)
.await
.expect("the job ended promptly after the reader failure")
.expect("the sender task did not panic")
.expect_err("a reader failure ends the job");
assert!(matches!(err, SendError::Source(_)), "{err:?}");
assert!(
lane_seen.load(Ordering::Relaxed) > 0,
"no frame was sent before the failure"
);
tokio::time::timeout(Duration::from_secs(10), async {
while cancel_seen.load(Ordering::SeqCst) != gen::ERR_IO as u32 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the receiver saw the JobCancel (reason ERR_IO) within the bound");
tokio::time::timeout(Duration::from_secs(10), async {
while active.load(Ordering::SeqCst) != 0 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("a reader thread leaked past the job's end");
}
#[tokio::test(flavor = "multi_thread")]
async fn a_file_retry_for_an_unknown_file_is_a_protocol_error() {
let size = 4usize << 20;
let (mut link, _lane_seen, _keep) = fake_link(
FakeReceiver {
credit: 64 << 20,
credit_on_apply: true,
done_when_complete: true,
credit_after_ack: None,
malformed_status: false,
retry_unknown: true,
..Default::default()
},
None,
);
let dir = std::env::temp_dir().join(format!("ava1-send-badretry-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("big"), vec![0x48u8; size]).unwrap();
let m = Arc::new(Manifest {
entries: vec![Entry {
kind: gen::ENTRY_FILE,
mode: 0o644,
size: size as u64,
mtime: 1,
path: "big".into(),
root: None,
}],
});
let err = tokio::time::timeout(
Duration::from_secs(20),
send_job(
&mut link,
m,
Arc::new(crate::source::LocalSource::new(dir.clone())),
SendOptions::upload("dest"),
),
)
.await
.expect("the job ended within the bound")
.expect_err("a FileRetry for an unknown file ends the job");
assert!(
matches!(&err, SendError::Protocol(msg) if msg.contains("FileRetry")),
"{err:?}"
);
let _ = std::fs::remove_dir_all(&dir);
}
}