diff --git a/server/crates/sn-fetch/src/error.rs b/server/crates/sn-fetch/src/error.rs index 8a9b2db..387694d 100644 --- a/server/crates/sn-fetch/src/error.rs +++ b/server/crates/sn-fetch/src/error.rs @@ -47,4 +47,12 @@ pub enum FetchError { #[error("invalid sha256 hex string")] InvalidDigest, + + #[error("segment {index}: all {attempts} attempt(s) failed; last error: {last_error}")] + SegmentAttemptsExhausted { + index: SegmentIndex, + attempts: u32, + #[source] + last_error: Box, + }, } diff --git a/server/crates/sn-fetch/src/exec.rs b/server/crates/sn-fetch/src/exec.rs index 711726d..f24f2e7 100644 --- a/server/crates/sn-fetch/src/exec.rs +++ b/server/crates/sn-fetch/src/exec.rs @@ -1,7 +1,18 @@ -//! Execution: download planned segments in parallel and reassemble the -//! target file. +//! Execution: download planned segments in parallel (with per-segment +//! 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::sync::{Arc, Mutex, MutexGuard}; use std::time::{Duration, Instant}; use sha2::{Digest, Sha256}; @@ -10,6 +21,7 @@ use tokio::io::AsyncWriteExt; use crate::error::FetchError; use crate::exit::ExitRegistry; use crate::plan::FetchPlan; +use crate::progress::{ProgressEvent, ProgressSink}; use crate::types::{ByteRange, ExitId, FileSize, SegmentIndex, Sha256Digest, TargetUrl}; /// Statistics for one finished segment. @@ -20,6 +32,50 @@ pub struct SegmentStats { 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 for SegmentStats { + fn from(stats: SegmentAttemptStats) -> Self { + SegmentStats { + index: stats.index, + bytes: stats.bytes, + duration: stats.duration, + } + } +} + /// Result of a complete fetch. #[derive(Debug, Clone)] pub struct FetchReport { @@ -41,19 +97,63 @@ pub async fn run_segments( exits: &ExitRegistry, parts_dir: &Path, ) -> Result, 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, FetchError> { 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 = 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()); for seg in &plan.segments { - let client = exits.get(seg.exit)?.client().clone(); - handles.push(tokio::spawn(download_segment( - plan.url.clone(), - seg.exit, - seg.index, - seg.range, - client, - parts_dir.to_path_buf(), - ))); + let task = SegmentTask { + url: plan.url.clone(), + index: seg.index, + range: seg.range, + planned_exit: seg.exit, + clients: Arc::clone(&clients), + selector: Arc::clone(&selector), + 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()); @@ -67,18 +167,158 @@ pub async fn run_segments( Ok(stats) } -/// Download one byte range and write it to `parts_dir/seg-NNNNNN.part`. -async fn download_segment( +/// Everything one segment task needs, bundled to keep the task +/// signature small. +struct SegmentTask { url: TargetUrl, + index: SegmentIndex, + range: ByteRange, + planned_exit: ExitId, + clients: Arc>, + selector: Arc, + 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 { + 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, index: SegmentIndex, range: ByteRange, client: reqwest::Client, - parts_dir: PathBuf, -) -> Result { - let started = Instant::now(); - - let part_path = part_file(&parts_dir, index); + parts_dir: &Path, +) -> Result { + let part_path = part_file(parts_dir, index); let mut response = client .get(url.inner().clone()) .header(reqwest::header::RANGE, range.http_value()) @@ -106,7 +346,7 @@ async fn download_segment( file.flush().await?; if written != range.len() { - tokio::fs::remove_file(&part_path).await.ok(); + let _ = tokio::fs::remove_file(&part_path).await; return Err(FetchError::SegmentLengthMismatch { index, got: written, @@ -114,11 +354,68 @@ async fn download_segment( }); } - Ok(SegmentStats { - index, - bytes: written, - duration: started.elapsed(), - }) + Ok(written) +} + +/// 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, + last_failed: Mutex>, +} + +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 = None; + for &id in &self.ids { + if id == avoid { + continue; + } + best = Some(match best { + None => id, + Some(current) => { + let prefer_id = match (guard.get(¤t), 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> { + match self.last_failed.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } + } } /// Reassemble part files (in index order) into the output file while @@ -199,6 +496,11 @@ async fn collect_part_files(parts_dir: &Path) -> Result, FetchError 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 { parts_dir.join(format!("seg-{:06}.part", index.value())) } diff --git a/server/crates/sn-fetch/src/lib.rs b/server/crates/sn-fetch/src/lib.rs index 346eeb9..1429523 100644 --- a/server/crates/sn-fetch/src/lib.rs +++ b/server/crates/sn-fetch/src/lib.rs @@ -18,5 +18,6 @@ pub mod exec; pub mod exit; pub mod plan; pub mod probe; +pub mod progress; pub mod runner; pub mod types; diff --git a/server/crates/sn-fetch/src/progress.rs b/server/crates/sn-fetch/src/progress.rs new file mode 100644 index 0000000..f759108 --- /dev/null +++ b/server/crates/sn-fetch/src/progress.rs @@ -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, +} + +impl ProgressSink { + /// Wrap any event callback into a sink. + pub fn new(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, +} + +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); + } + } +} diff --git a/server/crates/sn-fetch/tests/common/mod.rs b/server/crates/sn-fetch/tests/common/mod.rs index 57efbc0..fcbea8e 100644 --- a/server/crates/sn-fetch/tests/common/mod.rs +++ b/server/crates/sn-fetch/tests/common/mod.rs @@ -1,5 +1,7 @@ //! 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)] // Tests are allowed to unwrap/expect; production code denies them via @@ -8,8 +10,36 @@ use std::io::{BufRead, BufReader, Write}; use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; 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 addr: SocketAddr, @@ -21,14 +51,25 @@ impl TestServer { /// server behaves like a range-hostile origin: no `Accept-Ranges` /// header, and `Range` requests are answered with the full body. pub fn start(data: Vec, ranges: bool) -> Self { + Self::start_with(data, ranges, ServerConfig::default()) + } + + /// Start with a custom [`ServerConfig`] (failure simulation). + pub fn start_with(data: Vec, ranges: bool, config: ServerConfig) -> Self { 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 addr = listener.local_addr().expect("addr in test"); + let addr = listener.local_addr().expect("local addr in test"); let handle = std::thread::spawn(move || { 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 || { - 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 request_line = String::new(); 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 len = data.len() as u64; - let method = request_line .split_whitespace() .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(); match (ranges, range) { (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; write!( stream, @@ -111,6 +161,28 @@ fn handle_conn(stream: TcpStream, data: &[u8], ranges: bool) -> std::io::Result< 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 /// (open-ended, malformed) is treated as absent. fn parse_range(value: &str) -> Option> { @@ -124,3 +196,35 @@ fn parse_range(value: &str) -> Option> { 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 + } +} diff --git a/server/crates/sn-fetch/tests/retry.rs b/server/crates/sn-fetch/tests/retry.rs new file mode 100644 index 0000000..fa83a97 --- /dev/null +++ b/server/crates/sn-fetch/tests/retry.rs @@ -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>>; + +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 { + 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 = (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, + ®istry, + 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, + ®istry, + 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 = 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, ®istry, 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 = (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, + ®istry, + 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 = 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, + ®istry, + 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" + ); +}