sn-fetch: per-segment retry with exit reassignment and progress sinks

This commit is contained in:
2026-09-24 15:58:11 +02:00
parent f4e614acce
commit 810529d502
6 changed files with 1044 additions and 30 deletions
+8
View File
@@ -47,4 +47,12 @@ pub enum FetchError {
#[error("invalid sha256 hex string")] #[error("invalid sha256 hex string")]
InvalidDigest, InvalidDigest,
#[error("segment {index}: all {attempts} attempt(s) failed; last error: {last_error}")]
SegmentAttemptsExhausted {
index: SegmentIndex,
attempts: u32,
#[source]
last_error: Box<FetchError>,
},
} }
+326 -24
View File
@@ -1,7 +1,18 @@
//! Execution: download planned segments in parallel and reassemble the //! Execution: download planned segments in parallel (with per-segment
//! target file. //! retry and exit reassignment) and reassemble the target file.
//!
//! API note: [`SegmentStats`] keeps its original shape and `run_segments`
//! keeps its signature, so existing callers (runner.rs constructs
//! `SegmentStats` literals directly) stay source-compatible. The
//! retry-aware path [`run_segments_with_policy`] returns
//! [`SegmentAttemptStats`] instead, which additionally records the exit
//! that served the successful attempt and how many attempts (retries
//! included) the segment needed; `run_segments` converts via `From`.
use std::collections::BTreeMap;
use std::num::NonZeroU32;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
@@ -10,6 +21,7 @@ use tokio::io::AsyncWriteExt;
use crate::error::FetchError; use crate::error::FetchError;
use crate::exit::ExitRegistry; use crate::exit::ExitRegistry;
use crate::plan::FetchPlan; use crate::plan::FetchPlan;
use crate::progress::{ProgressEvent, ProgressSink};
use crate::types::{ByteRange, ExitId, FileSize, SegmentIndex, Sha256Digest, TargetUrl}; use crate::types::{ByteRange, ExitId, FileSize, SegmentIndex, Sha256Digest, TargetUrl};
/// Statistics for one finished segment. /// Statistics for one finished segment.
@@ -20,6 +32,50 @@ pub struct SegmentStats {
pub duration: Duration, pub duration: Duration,
} }
/// Retry behaviour for segment downloads.
#[derive(Debug, Clone, Copy)]
pub struct RetryPolicy {
/// Total tries per segment; attempt 1 is the initial try.
pub max_attempts: NonZeroU32,
}
impl RetryPolicy {
#[must_use]
pub const fn new(max_attempts: NonZeroU32) -> Self {
Self { max_attempts }
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_attempts: NonZeroU32::new(3).unwrap_or(NonZeroU32::MIN),
}
}
}
/// Statistics for one finished segment, including retry information.
/// `duration` spans all attempts (retries included); `exit` is the exit
/// that served the successful attempt.
#[derive(Debug, Clone)]
pub struct SegmentAttemptStats {
pub index: SegmentIndex,
pub exit: ExitId,
pub bytes: u64,
pub duration: Duration,
pub attempts: u32,
}
impl From<SegmentAttemptStats> for SegmentStats {
fn from(stats: SegmentAttemptStats) -> Self {
SegmentStats {
index: stats.index,
bytes: stats.bytes,
duration: stats.duration,
}
}
}
/// Result of a complete fetch. /// Result of a complete fetch.
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct FetchReport { pub struct FetchReport {
@@ -41,19 +97,63 @@ pub async fn run_segments(
exits: &ExitRegistry, exits: &ExitRegistry,
parts_dir: &Path, parts_dir: &Path,
) -> Result<Vec<SegmentStats>, FetchError> { ) -> Result<Vec<SegmentStats>, FetchError> {
let stats = run_segments_with_policy(
plan,
exits,
parts_dir,
&RetryPolicy::default(),
&ProgressSink::noop(),
)
.await?;
Ok(stats.into_iter().map(SegmentStats::from).collect())
}
/// Download every planned segment in parallel, retrying failed segments
/// according to `policy` — preferably on a different exit — and report
/// progress to `progress`. Each segment is written to its own part file
/// inside `parts_dir` (named `seg-NNNNNN.part`); the part file of a
/// failed attempt is removed before the next attempt.
///
/// Failures abort the whole fetch: a segment that exhausts all attempts
/// fails the fetch with [`FetchError::SegmentAttemptsExhausted`], a
/// non-retryable failure aborts immediately. A partially assembled file
/// is never presented as success.
pub async fn run_segments_with_policy(
plan: &FetchPlan,
exits: &ExitRegistry,
parts_dir: &Path,
policy: &RetryPolicy,
progress: &ProgressSink,
) -> Result<Vec<SegmentAttemptStats>, FetchError> {
tokio::fs::create_dir_all(parts_dir).await?; tokio::fs::create_dir_all(parts_dir).await?;
// Validate that every planned exit exists, then snapshot the
// clients of all registered exits: tasks must fetch clients by id
// without borrowing the registry ('static spawn requirement).
for seg in &plan.segments {
exits.get(seg.exit)?;
}
let mut clients: BTreeMap<ExitId, reqwest::Client> = BTreeMap::new();
for id in exits.ids() {
clients.insert(id, exits.get(id)?.client().clone());
}
let clients = Arc::new(clients);
let selector = Arc::new(ExitSelector::new(exits));
let mut handles = Vec::with_capacity(plan.segments.len()); let mut handles = Vec::with_capacity(plan.segments.len());
for seg in &plan.segments { for seg in &plan.segments {
let client = exits.get(seg.exit)?.client().clone(); let task = SegmentTask {
handles.push(tokio::spawn(download_segment( url: plan.url.clone(),
plan.url.clone(), index: seg.index,
seg.exit, range: seg.range,
seg.index, planned_exit: seg.exit,
seg.range, clients: Arc::clone(&clients),
client, selector: Arc::clone(&selector),
parts_dir.to_path_buf(), parts_dir: parts_dir.to_path_buf(),
))); policy: *policy,
progress: progress.clone(),
};
handles.push(tokio::spawn(download_with_retry(task)));
} }
let mut stats = Vec::with_capacity(handles.len()); let mut stats = Vec::with_capacity(handles.len());
@@ -67,18 +167,158 @@ pub async fn run_segments(
Ok(stats) Ok(stats)
} }
/// Download one byte range and write it to `parts_dir/seg-NNNNNN.part`. /// Everything one segment task needs, bundled to keep the task
async fn download_segment( /// signature small.
struct SegmentTask {
url: TargetUrl, url: TargetUrl,
index: SegmentIndex,
range: ByteRange,
planned_exit: ExitId,
clients: Arc<BTreeMap<ExitId, reqwest::Client>>,
selector: Arc<ExitSelector>,
parts_dir: PathBuf,
policy: RetryPolicy,
progress: ProgressSink,
}
/// Download one segment with up to `policy.max_attempts` attempts,
/// reassigning to a different exit after a retryable failure (while the
/// registry offers more than one exit).
async fn download_with_retry(task: SegmentTask) -> Result<SegmentAttemptStats, FetchError> {
let SegmentTask {
url,
index,
range,
planned_exit,
clients,
selector,
parts_dir,
policy,
progress,
} = task;
let max_attempts = policy.max_attempts.get();
let started = Instant::now();
let mut exit = planned_exit;
let mut attempt: u32 = 0;
loop {
attempt += 1;
let client = clients
.get(&exit)
.ok_or(FetchError::UnknownExit(exit))?
.clone();
progress.on_event(&ProgressEvent::SegmentStarted {
index,
exit,
attempt,
});
tracing::debug!(segment = index.value(), exit = %exit, attempt, "segment attempt started");
match download_attempt(&url, exit, index, range, client, &parts_dir).await {
Ok(bytes) => {
let duration = started.elapsed();
progress.on_event(&ProgressEvent::SegmentFinished {
index,
exit,
bytes,
duration,
});
tracing::info!(
segment = index.value(),
exit = %exit,
attempts = attempt,
bytes,
"segment finished"
);
return Ok(SegmentAttemptStats {
index,
exit,
bytes,
duration,
attempts: attempt,
});
}
Err(err) => {
let retryable = is_retryable(&err);
let err_text = err.to_string();
progress.on_event(&ProgressEvent::SegmentFailed {
index,
exit,
attempt,
error: err_text,
});
remove_part_file(&parts_dir, index).await;
if !retryable {
tracing::warn!(
segment = index.value(),
exit = %exit,
attempt,
"segment failed with non-retryable error: {err}"
);
return Err(err);
}
selector.mark_failed(exit);
if attempt >= max_attempts {
tracing::warn!(
segment = index.value(),
attempts = attempt,
"segment failed after all attempts were exhausted: {err}"
);
return Err(FetchError::SegmentAttemptsExhausted {
index,
attempts: attempt,
last_error: Box::new(err),
});
}
tracing::warn!(
segment = index.value(),
exit = %exit,
attempt,
"segment attempt failed; retrying: {err}"
);
exit = selector.pick_retry(exit);
}
}
}
}
/// Retryable: transport-level failures (network, timeouts), 5xx
/// responses and length mismatches. Non-retryable: 4xx range rejections
/// (e.g. 416), corrupt parts, plan or digest errors — retrying those on
/// another exit cannot succeed.
fn is_retryable(error: &FetchError) -> bool {
match error {
FetchError::Http(err) => {
!err.is_status() || err.status().is_some_and(|status| status.is_server_error())
}
FetchError::Io(_) | FetchError::SegmentLengthMismatch { .. } => true,
FetchError::RangeRejected { status, .. } => *status >= 500,
FetchError::InvalidUrl(_)
| FetchError::ZeroFileSize
| FetchError::TotalLengthMismatch { .. }
| FetchError::HashMismatch { .. }
| FetchError::InvalidPlan(_)
| FetchError::UnknownExit(_)
| FetchError::CorruptPart(_)
| FetchError::InvalidDigest
| FetchError::SegmentAttemptsExhausted { .. } => false,
}
}
/// Download one byte range and write it to `parts_dir/seg-NNNNNN.part`.
/// Returns the number of bytes written. A failed attempt can leave a
/// partial part file behind; the retry loop removes it before the next
/// attempt.
async fn download_attempt(
url: &TargetUrl,
exit: ExitId, exit: ExitId,
index: SegmentIndex, index: SegmentIndex,
range: ByteRange, range: ByteRange,
client: reqwest::Client, client: reqwest::Client,
parts_dir: PathBuf, parts_dir: &Path,
) -> Result<SegmentStats, FetchError> { ) -> Result<u64, FetchError> {
let started = Instant::now(); let part_path = part_file(parts_dir, index);
let part_path = part_file(&parts_dir, index);
let mut response = client let mut response = client
.get(url.inner().clone()) .get(url.inner().clone())
.header(reqwest::header::RANGE, range.http_value()) .header(reqwest::header::RANGE, range.http_value())
@@ -106,7 +346,7 @@ async fn download_segment(
file.flush().await?; file.flush().await?;
if written != range.len() { if written != range.len() {
tokio::fs::remove_file(&part_path).await.ok(); let _ = tokio::fs::remove_file(&part_path).await;
return Err(FetchError::SegmentLengthMismatch { return Err(FetchError::SegmentLengthMismatch {
index, index,
got: written, got: written,
@@ -114,11 +354,68 @@ async fn download_segment(
}); });
} }
Ok(SegmentStats { Ok(written)
index, }
bytes: written,
duration: started.elapsed(), /// Chooses exits for retry attempts across all segment tasks of one
}) /// run. Preference: the exit that failed least recently (exits that
/// never failed first, ties broken by ascending id). The exit an
/// attempt just failed on is avoided while at least one other exit
/// exists; with a single-exit registry, retries stay on that exit.
#[derive(Debug)]
struct ExitSelector {
ids: Vec<ExitId>,
last_failed: Mutex<BTreeMap<ExitId, Instant>>,
}
impl ExitSelector {
fn new(exits: &ExitRegistry) -> Self {
Self {
ids: exits.ids(),
last_failed: Mutex::new(BTreeMap::new()),
}
}
fn mark_failed(&self, id: ExitId) {
let mut guard = self.lock();
guard.insert(id, Instant::now());
}
fn pick_retry(&self, avoid: ExitId) -> ExitId {
let guard = self.lock();
let mut best: Option<ExitId> = None;
for &id in &self.ids {
if id == avoid {
continue;
}
best = Some(match best {
None => id,
Some(current) => {
let prefer_id = match (guard.get(&current), guard.get(&id)) {
// tie or current never failed: keep the lower id
(None, None | Some(_)) => false,
(Some(_), None) => true, // id never failed: better
(Some(current_failed), Some(id_failed)) => *id_failed < *current_failed,
};
if prefer_id {
id
} else {
current
}
}
});
}
best.unwrap_or(avoid) // single-exit registry: retry on the same exit
}
/// std mutexes can be poisoned by a panicking holder; since the
/// map is only advisory (exit preference), keep using its contents.
fn lock(&self) -> MutexGuard<'_, BTreeMap<ExitId, Instant>> {
match self.last_failed.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}
} }
/// Reassemble part files (in index order) into the output file while /// Reassemble part files (in index order) into the output file while
@@ -199,6 +496,11 @@ async fn collect_part_files(parts_dir: &Path) -> Result<Vec<PathBuf>, FetchError
Ok(out) Ok(out)
} }
/// Best-effort removal of a failed attempt's part file.
async fn remove_part_file(parts_dir: &Path, index: SegmentIndex) {
let _ = tokio::fs::remove_file(part_file(parts_dir, index)).await;
}
fn part_file(parts_dir: &Path, index: SegmentIndex) -> PathBuf { fn part_file(parts_dir: &Path, index: SegmentIndex) -> PathBuf {
parts_dir.join(format!("seg-{:06}.part", index.value())) parts_dir.join(format!("seg-{:06}.part", index.value()))
} }
+1
View File
@@ -18,5 +18,6 @@ pub mod exec;
pub mod exit; pub mod exit;
pub mod plan; pub mod plan;
pub mod probe; pub mod probe;
pub mod progress;
pub mod runner; pub mod runner;
pub mod types; pub mod types;
+171
View File
@@ -0,0 +1,171 @@
//! Progress reporting for segmented downloads.
//!
//! [`ProgressSink`] is a cloneable callback handle rather than a trait:
//! spawned segment tasks need an owned, `'static` reporter, and a
//! callback handle can be cloned cheaply from a plain `&ProgressSink`
//! parameter without extra plumbing. [`ProgressEvent`] carries one
//! event per segment attempt; [`BytesTracker`] aggregates the bytes
//! finished across all segments.
#![allow(clippy::module_name_repetitions)] // `Progress*` types read naturally inside `mod progress`
use std::fmt;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use crate::types::{ExitId, SegmentIndex};
/// One event describing the state of a segment download attempt.
#[derive(Debug, Clone)]
pub enum ProgressEvent {
/// A segment attempt started (attempt 1 is the initial try).
SegmentStarted {
index: SegmentIndex,
exit: ExitId,
attempt: u32,
},
/// A segment finished successfully. `duration` spans the whole
/// segment, all attempts (retries included).
SegmentFinished {
index: SegmentIndex,
exit: ExitId,
bytes: u64,
duration: Duration,
},
/// A segment attempt failed. `error` is the formatted failure
/// reason; the typed error surfaces through the fetch's result.
SegmentFailed {
index: SegmentIndex,
exit: ExitId,
attempt: u32,
error: String,
},
}
impl fmt::Display for ProgressEvent {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::SegmentStarted {
index,
exit,
attempt,
} => {
write!(f, "segment {index} on {exit}: attempt {attempt} started")
}
Self::SegmentFinished {
index,
exit,
bytes,
duration,
} => {
write!(
f,
"segment {index} on {exit}: {bytes} bytes in {duration:.2?}"
)
}
Self::SegmentFailed {
index,
exit,
attempt,
error,
} => {
write!(
f,
"segment {index} on {exit}: attempt {attempt} failed: {error}"
)
}
}
}
}
/// Cloneable handle to a progress callback; `Send + Sync`, so it can be
/// shared across tokio tasks. All clones call the same callback.
#[derive(Clone)]
pub struct ProgressSink {
on_event: Arc<dyn Fn(&ProgressEvent) + Send + Sync>,
}
impl ProgressSink {
/// Wrap any event callback into a sink.
pub fn new<F>(on_event: F) -> Self
where
F: Fn(&ProgressEvent) + Send + Sync + 'static,
{
Self {
on_event: Arc::new(on_event),
}
}
/// A sink that discards every event.
#[must_use]
pub fn noop() -> Self {
Self::new(|_| ())
}
/// A sink that logs every event through `tracing`: started attempts
/// at debug, finished segments at info, failures at warn.
#[must_use]
pub fn tracing_log() -> Self {
Self::new(|event| match event {
ProgressEvent::SegmentStarted { .. } => tracing::debug!("{event}"),
ProgressEvent::SegmentFinished { .. } => tracing::info!("{event}"),
ProgressEvent::SegmentFailed { .. } => tracing::warn!("{event}"),
})
}
/// Combine two sinks: both receive every event.
#[must_use]
pub fn fanout(first: &Self, second: &Self) -> Self {
let first = first.clone();
let second = second.clone();
Self::new(move |event| {
first.on_event(event);
second.on_event(event);
})
}
/// Deliver one event to the underlying callback.
pub fn on_event(&self, event: &ProgressEvent) {
(self.on_event)(event);
}
}
impl fmt::Debug for ProgressSink {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ProgressSink").finish_non_exhaustive()
}
}
/// Counts the bytes of successfully finished segments across the whole
/// fetch. Cheap to clone: every clone observes the same total.
#[derive(Debug, Default, Clone)]
pub struct BytesTracker {
bytes: Arc<AtomicU64>,
}
impl BytesTracker {
#[must_use]
pub fn new() -> Self {
Self::default()
}
/// Total bytes reported as finished so far.
#[must_use]
pub fn total(&self) -> u64 {
self.bytes.load(Ordering::Relaxed)
}
/// A sink that feeds this tracker; the tracker stays queryable.
#[must_use]
pub fn sink(&self) -> ProgressSink {
let tracker = self.clone();
ProgressSink::new(move |event| tracker.record(event))
}
fn record(&self, event: &ProgressEvent) {
if let ProgressEvent::SegmentFinished { bytes, .. } = event {
self.bytes.fetch_add(*bytes, Ordering::Relaxed);
}
}
}
+110 -6
View File
@@ -1,5 +1,7 @@
//! A minimal HTTP server for integration tests. Supports HEAD, GET with //! A minimal HTTP server for integration tests. Supports HEAD, GET with
//! optional range handling (configurable), `Connection: close` semantics. //! optional range handling (configurable), simulated failures for range
//! GETs (flaky/exhausted/non-retryable modes), and a broken exit whose
//! client fails at transport level, for retry tests.
#![forbid(unsafe_code)] #![forbid(unsafe_code)]
// Tests are allowed to unwrap/expect; production code denies them via // Tests are allowed to unwrap/expect; production code denies them via
@@ -8,8 +10,36 @@
use std::io::{BufRead, BufReader, Write}; use std::io::{BufRead, BufReader, Write};
use std::net::{SocketAddr, TcpListener, TcpStream}; use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc; use std::sync::Arc;
use std::thread::JoinHandle; use std::thread::JoinHandle;
use std::time::Duration;
use sn_fetch::exit::HttpExit;
use sn_fetch::types::ExitId;
/// When the test origin should fail `Range` GETs (retry simulation).
/// Only used by the retry test binary; allow dead_code for the others.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[allow(dead_code)]
pub enum FailMode {
#[default]
Never,
/// Fail the first `n` range GETs with HTTP 500, then serve normally.
FirstN(u32),
/// Fail every range GET with HTTP 500.
Always,
/// Answer every range GET with the given status (e.g. 416 to
/// simulate a non-retryable range rejection).
AlwaysStatus(u16),
}
/// Knobs for the test origin.
#[derive(Debug, Clone, Copy, Default)]
#[allow(dead_code)]
pub struct ServerConfig {
pub fail_mode: FailMode,
}
pub struct TestServer { pub struct TestServer {
pub addr: SocketAddr, pub addr: SocketAddr,
@@ -21,14 +51,25 @@ impl TestServer {
/// server behaves like a range-hostile origin: no `Accept-Ranges` /// server behaves like a range-hostile origin: no `Accept-Ranges`
/// header, and `Range` requests are answered with the full body. /// header, and `Range` requests are answered with the full body.
pub fn start(data: Vec<u8>, ranges: bool) -> Self { pub fn start(data: Vec<u8>, ranges: bool) -> Self {
Self::start_with(data, ranges, ServerConfig::default())
}
/// Start with a custom [`ServerConfig`] (failure simulation).
pub fn start_with(data: Vec<u8>, ranges: bool, config: ServerConfig) -> Self {
let data = Arc::new(data); let data = Arc::new(data);
let config = Arc::new(config);
let range_get_count = Arc::new(AtomicU32::new(0));
let listener = TcpListener::bind("127.0.0.1:0").expect("bind in test"); let listener = TcpListener::bind("127.0.0.1:0").expect("bind in test");
let addr = listener.local_addr().expect("addr in test"); let addr = listener.local_addr().expect("local addr in test");
let handle = std::thread::spawn(move || { let handle = std::thread::spawn(move || {
for stream in listener.incoming().flatten() { for stream in listener.incoming().flatten() {
let data = Arc::clone(&data); let (data, config, range_get_count) = (
Arc::clone(&data),
Arc::clone(&config),
Arc::clone(&range_get_count),
);
std::thread::spawn(move || { std::thread::spawn(move || {
let _ = handle_conn(stream, &data, ranges); let _ = handle_conn(stream, &data, ranges, &config, &range_get_count);
}); });
} }
}); });
@@ -43,7 +84,13 @@ impl TestServer {
} }
} }
fn handle_conn(stream: TcpStream, data: &[u8], ranges: bool) -> std::io::Result<()> { fn handle_conn(
stream: TcpStream,
data: &[u8],
ranges: bool,
config: &ServerConfig,
range_get_count: &AtomicU32,
) -> std::io::Result<()> {
let mut reader = BufReader::new(stream.try_clone()?); let mut reader = BufReader::new(stream.try_clone()?);
let mut request_line = String::new(); let mut request_line = String::new();
reader.read_line(&mut request_line)?; reader.read_line(&mut request_line)?;
@@ -64,7 +111,6 @@ fn handle_conn(stream: TcpStream, data: &[u8], ranges: bool) -> std::io::Result<
let mut stream = stream; let mut stream = stream;
let len = data.len() as u64; let len = data.len() as u64;
let method = request_line let method = request_line
.split_whitespace() .split_whitespace()
.next() .next()
@@ -83,6 +129,10 @@ fn handle_conn(stream: TcpStream, data: &[u8], ranges: bool) -> std::io::Result<
let range = range_header.as_deref().and_then(parse_range).flatten(); let range = range_header.as_deref().and_then(parse_range).flatten();
match (ranges, range) { match (ranges, range) {
(true, Some((start, end))) => { (true, Some((start, end))) => {
if should_fail_range_get(config, range_get_count) {
write_fail_response(&mut stream, config.fail_mode)?;
return Ok(());
}
let slice_len = end - start + 1; let slice_len = end - start + 1;
write!( write!(
stream, stream,
@@ -111,6 +161,28 @@ fn handle_conn(stream: TcpStream, data: &[u8], ranges: bool) -> std::io::Result<
Ok(()) Ok(())
} }
/// Decide (and count) whether this range GET should fail. HEAD and
/// non-range GETs are never counted nor failed.
fn should_fail_range_get(config: &ServerConfig, range_get_count: &AtomicU32) -> bool {
match config.fail_mode {
FailMode::Never => false,
FailMode::Always | FailMode::AlwaysStatus(_) => true,
FailMode::FirstN(limit) => range_get_count.fetch_add(1, Ordering::Relaxed) < limit,
}
}
fn write_fail_response(stream: &mut TcpStream, mode: FailMode) -> std::io::Result<()> {
let status = match mode {
FailMode::AlwaysStatus(status) => status,
_ => 500,
};
write!(
stream,
"HTTP/1.1 {status} Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
)?;
stream.flush()
}
/// Parse `bytes=a-b` into inclusive `(start, end)`. Anything else /// Parse `bytes=a-b` into inclusive `(start, end)`. Anything else
/// (open-ended, malformed) is treated as absent. /// (open-ended, malformed) is treated as absent.
fn parse_range(value: &str) -> Option<Option<(u64, u64)>> { fn parse_range(value: &str) -> Option<Option<(u64, u64)>> {
@@ -124,3 +196,35 @@ fn parse_range(value: &str) -> Option<Option<(u64, u64)>> {
Some(None) Some(None)
} }
} }
/// An exit whose client sends every request through a proxy that
/// refuses connections — every request through this exit fails at
/// transport level (retryable), as needed for exit-reassignment tests.
#[allow(dead_code)]
pub struct BrokenProxyExit {
id: ExitId,
client: reqwest::Client,
}
impl BrokenProxyExit {
#[allow(dead_code)]
pub fn new(id: ExitId) -> Self {
let proxy = reqwest::Proxy::http("http://127.0.0.1:1").expect("static proxy url parses");
let client = reqwest::Client::builder()
.proxy(proxy)
.connect_timeout(Duration::from_secs(10))
.build()
.expect("client builds in test");
Self { id, client }
}
}
impl HttpExit for BrokenProxyExit {
fn id(&self) -> ExitId {
self.id
}
fn client(&self) -> &reqwest::Client {
&self.client
}
}
+428
View File
@@ -0,0 +1,428 @@
//! Retry behaviour against a local HTTP server: a flaky first attempt
//! (retried, reassigned to another exit), a consistently failing exit
//! (reassigned), retry exhaustion (whole fetch fails, no output file)
//! and a non-retryable status (immediate abort).
#![forbid(unsafe_code)]
// Tests are allowed to unwrap/expect; production code denies them via
// workspace lints (server/Cargo.toml).
#![allow(clippy::unwrap_used, clippy::expect_used)]
mod common;
use std::num::{NonZeroU32, NonZeroU64, NonZeroU8, NonZeroUsize};
use std::path::Path;
use std::sync::{Arc, Mutex};
use sha2::{Digest, Sha256};
use tempfile::TempDir;
use sn_fetch::error::FetchError;
use sn_fetch::exec::{self, RetryPolicy};
use sn_fetch::exit::{ExitRegistry, LocalExit};
use sn_fetch::plan::{ExitSpecs, FetchPlan};
use sn_fetch::progress::{BytesTracker, ProgressEvent, ProgressSink};
use sn_fetch::types::{
Egress, ExitId, ExitWeight, FileSize, MinSegmentBytes, SegmentsPerExit, TargetUrl,
};
use common::{BrokenProxyExit, FailMode, ServerConfig, TestServer};
fn spec(id: u8) -> (ExitId, ExitWeight, Egress) {
(
ExitId::new(NonZeroU8::new(id).expect("test ids are 1..=255")),
ExitWeight::new(NonZeroU32::new(1).expect("1 is non-zero")),
Egress::Direct,
)
}
fn exit_id(id: u8) -> ExitId {
ExitId::new(NonZeroU8::new(id).expect("test ids are 1..=255"))
}
fn local_registry(specs: &ExitSpecs) -> ExitRegistry {
let mut registry = ExitRegistry::empty();
for (id, _, _) in specs {
registry.register(Arc::new(LocalExit::new(*id).expect("local exit")));
}
registry
}
/// Single-exit plan + registry, for exhaustion tests.
fn single_exit_setup(
url: &TargetUrl,
file_size: u64,
min_segment: u64,
) -> (FetchPlan, ExitRegistry) {
let specs: ExitSpecs = vec![spec(1)];
let plan = FetchPlan::build(
url.clone(),
FileSize::try_new(file_size).expect("non-zero in test"),
&specs,
SegmentsPerExit::new(NonZeroUsize::MIN),
MinSegmentBytes::new(NonZeroU64::new(min_segment).expect("non-zero in test")),
)
.expect("plan builds");
let registry = local_registry(&specs);
(plan, registry)
}
type RecordedEvents = Arc<Mutex<Vec<ProgressEvent>>>;
fn recording_sink() -> (ProgressSink, RecordedEvents) {
let events: RecordedEvents = Arc::new(Mutex::new(Vec::new()));
let sink = ProgressSink::new({
let events = Arc::clone(&events);
move |event| {
events
.lock()
.expect("mutex poisoned in test")
.push(event.clone())
}
});
(sink, events)
}
fn failed_events(events: &RecordedEvents) -> Vec<(u32, u32, String)> {
events
.lock()
.expect("mutex poisoned in test")
.iter()
.filter_map(|event| match event {
ProgressEvent::SegmentFailed {
index,
exit: _,
attempt,
error,
} => Some((index.value(), *attempt, error.clone())),
_ => None,
})
.collect()
}
/// Names of leftover part files in `parts_dir`.
fn part_files(dir: &Path) -> Vec<String> {
std::fs::read_dir(dir)
.expect("parts dir readable in test")
.flatten()
.filter(|entry| entry.file_name().to_string_lossy().ends_with(".part"))
.map(|entry| entry.file_name().to_string_lossy().into_owned())
.collect()
}
#[tokio::test]
async fn flaky_segment_retried_on_other_exit() {
let data: Vec<u8> = (0..(2 * 1024 * 1024)).map(|i| (i % 253) as u8).collect();
let server = TestServer::start_with(
data.clone(),
true,
ServerConfig {
fail_mode: FailMode::FirstN(1),
},
);
let url = TargetUrl::parse(&server.url("/file.bin")).expect("valid url");
let specs: ExitSpecs = vec![spec(1), spec(2)];
let plan = FetchPlan::build(
url,
FileSize::try_new(data.len() as u64).expect("non-zero in test"),
&specs,
SegmentsPerExit::new(NonZeroUsize::new(2).expect("2 is non-zero")),
MinSegmentBytes::new(NonZeroU64::new(1024).expect("1024 is non-zero")),
)
.expect("plan builds");
let registry = local_registry(&specs);
let parts = TempDir::new().expect("tempdir in test");
let (progress, events) = recording_sink();
let tracker = BytesTracker::new();
let combined = ProgressSink::fanout(&progress, &tracker.sink());
let stats = exec::run_segments_with_policy(
&plan,
&registry,
parts.path(),
&RetryPolicy::default(),
&combined,
)
.await
.expect("fetch succeeds despite one flaky attempt");
assert_eq!(stats.len(), plan.segments.len());
let total: u64 = stats.iter().map(|s| s.bytes).sum();
assert_eq!(total, data.len() as u64);
assert_eq!(
tracker.total(),
data.len() as u64,
"tracker must see all finished bytes"
);
let retried: Vec<_> = stats.iter().filter(|s| s.attempts > 1).collect();
assert_eq!(retried.len(), 1, "exactly one segment needed a retry");
let retried = retried[0];
let planned_exit = plan
.segments
.iter()
.find(|seg| seg.index == retried.index)
.expect("planned segment")
.exit;
assert_ne!(retried.exit, planned_exit, "retry must use another exit");
assert_eq!(retried.attempts, 2);
let started: Vec<(u8, u32)> = events
.lock()
.expect("mutex poisoned in test")
.iter()
.filter_map(|event| match event {
ProgressEvent::SegmentStarted {
index,
exit,
attempt,
} if index == &retried.index => Some((exit.value(), *attempt)),
_ => None,
})
.collect();
assert_eq!(
started,
vec![(planned_exit.value(), 1), (retried.exit.value(), 2)],
"the retried segment must start once per attempt, on different exits"
);
let failed = failed_events(&events);
assert_eq!(
failed,
vec![(retried.index.value(), 1, failed[0].2.clone())],
"exactly one failed attempt, recorded on attempt 1"
);
assert!(
failed[0].2.contains("500"),
"failure reason should mention the 500 status: {}",
failed[0].2
);
}
#[tokio::test]
async fn exhausted_retries_fail_without_output() {
let server = TestServer::start_with(
vec![7u8; 32 * 1024],
true,
ServerConfig {
fail_mode: FailMode::Always,
},
);
let url = TargetUrl::parse(&server.url("/file.bin")).expect("valid url");
let (plan, registry) = single_exit_setup(&url, 32 * 1024, 32 * 1024);
assert_eq!(plan.segments.len(), 1);
let parts = TempDir::new().expect("tempdir in test");
let (progress, events) = recording_sink();
let tracker = BytesTracker::new();
let combined = ProgressSink::fanout(&progress, &tracker.sink());
let result = exec::run_segments_with_policy(
&plan,
&registry,
parts.path(),
&RetryPolicy::default(),
&combined,
)
.await;
let err = result.expect_err("fetch must fail when every attempt fails");
assert!(
matches!(
err,
FetchError::SegmentAttemptsExhausted { attempts: 3, .. }
),
"unexpected error: {err}"
);
assert_eq!(
tracker.total(),
0,
"nothing finished, tracker stays at zero"
);
let failed: Vec<u32> = failed_events(&events)
.into_iter()
.map(|(_, attempt, error)| {
assert!(error.contains("500"), "failure reason: {error}");
attempt
})
.collect();
assert_eq!(failed, vec![1, 2, 3], "default policy means three attempts");
assert!(
part_files(parts.path()).is_empty(),
"failed part files must be cleaned up between attempts"
);
let output = parts.path().join("output.bin");
assert!(
!output.exists(),
"no output file may be created for a failed fetch"
);
}
#[tokio::test]
async fn single_attempt_policy_fails_fast() {
let server = TestServer::start_with(
vec![5u8; 16 * 1024],
true,
ServerConfig {
fail_mode: FailMode::Always,
},
);
let url = TargetUrl::parse(&server.url("/file.bin")).expect("valid url");
let (plan, registry) = single_exit_setup(&url, 16 * 1024, 16 * 1024);
let parts = TempDir::new().expect("tempdir in test");
let (progress, events) = recording_sink();
let policy = RetryPolicy::new(NonZeroU32::new(1).expect("1 is non-zero"));
let result =
exec::run_segments_with_policy(&plan, &registry, parts.path(), &policy, &progress).await;
let err = result.expect_err("fetch must fail after the single attempt");
assert!(
matches!(
err,
FetchError::SegmentAttemptsExhausted { attempts: 1, .. }
),
"unexpected error: {err}"
);
assert_eq!(failed_events(&events).len(), 1);
assert!(part_files(parts.path()).is_empty());
}
#[tokio::test]
async fn failing_exit_reassigned_to_working_exit() {
let data: Vec<u8> = (0..(1024 * 1024)).map(|i| (i % 241) as u8).collect();
let expected_digest = sn_fetch::types::Sha256Digest::from_bytes(Sha256::digest(&data).into());
let server = TestServer::start(data.clone(), true);
let url = TargetUrl::parse(&server.url("/file.bin")).expect("valid url");
let specs: ExitSpecs = vec![spec(1), spec(2)];
let plan = FetchPlan::build(
url,
FileSize::try_new(data.len() as u64).expect("non-zero in test"),
&specs,
SegmentsPerExit::new(NonZeroUsize::MIN),
MinSegmentBytes::new(NonZeroU64::new(1024).expect("1024 is non-zero")),
)
.expect("plan builds");
let mut registry = ExitRegistry::empty();
registry.register(Arc::new(BrokenProxyExit::new(exit_id(1))));
registry.register(Arc::new(LocalExit::new(exit_id(2)).expect("local exit")));
let parts = TempDir::new().expect("tempdir in test");
let (progress, events) = recording_sink();
let stats = exec::run_segments_with_policy(
&plan,
&registry,
parts.path(),
&RetryPolicy::default(),
&progress,
)
.await
.expect("fetch succeeds by reassigning off the broken exit");
for stat in &stats {
let planned = plan
.segments
.iter()
.find(|seg| seg.index == stat.index)
.expect("planned segment");
if planned.exit.value() == 1 {
assert_eq!(
stat.attempts, 2,
"segment planned on the broken exit needs exactly one retry"
);
assert_eq!(stat.exit.value(), 2, "retry must land on the working exit");
} else {
assert_eq!(stat.attempts, 1);
assert_eq!(stat.exit.value(), 2);
}
}
let started_exits: Vec<u8> = events
.lock()
.expect("mutex poisoned in test")
.iter()
.filter_map(|event| match event {
ProgressEvent::SegmentStarted { index, exit, .. } if index.value() == 0 => {
Some(exit.value())
}
_ => None,
})
.collect();
assert_eq!(
started_exits,
vec![1, 2],
"segment 0 must be tried on the broken exit first, then reassigned"
);
let output = parts.path().join("out.bin");
let digest = exec::assemble(
parts.path(),
&output,
FileSize::try_new(data.len() as u64).expect("non-zero in test"),
Some(&expected_digest),
)
.await
.expect("assembly verifies");
assert_eq!(digest, expected_digest);
}
#[tokio::test]
async fn non_retryable_status_aborts_immediately() {
let server = TestServer::start_with(
vec![3u8; 32 * 1024],
true,
ServerConfig {
fail_mode: FailMode::AlwaysStatus(416),
},
);
let url = TargetUrl::parse(&server.url("/file.bin")).expect("valid url");
let (plan, registry) = single_exit_setup(&url, 32 * 1024, 32 * 1024);
let parts = TempDir::new().expect("tempdir in test");
let (progress, events) = recording_sink();
let result = exec::run_segments_with_policy(
&plan,
&registry,
parts.path(),
&RetryPolicy::default(),
&progress,
)
.await;
let err = result.expect_err("416 must abort the fetch");
assert!(
matches!(
err,
FetchError::Http(ref http_err)
if http_err.status() == Some(reqwest::StatusCode::RANGE_NOT_SATISFIABLE)
),
"the raw 416 must surface unchanged, got: {err}"
);
// Non-retryable: one started attempt, one failure event, no retry.
let guard = events.lock().expect("mutex poisoned in test");
let started = guard
.iter()
.filter(|event| matches!(event, ProgressEvent::SegmentStarted { .. }))
.count();
assert_eq!(started, 1, "non-retryable failure must not be retried");
let failed = guard
.iter()
.filter(|event| matches!(event, ProgressEvent::SegmentFailed { .. }))
.count();
assert_eq!(failed, 1);
drop(guard);
assert!(
part_files(parts.path()).is_empty(),
"failed part files must be cleaned up"
);
}