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

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<FetchError>,
+    },
 }

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<SegmentAttemptStats> 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<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?;
 
+    // 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());
     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<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,
     index: SegmentIndex,
     range: ByteRange,
     client: reqwest::Client,
-    parts_dir: PathBuf,
-) -> Result<SegmentStats, FetchError> {
-    let started = Instant::now();
-
-    let part_path = part_file(&parts_dir, index);
+    parts_dir: &Path,
+) -> Result<u64, FetchError> {
+    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<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
@@ -199,6 +496,11 @@ async fn collect_part_files(parts_dir: &Path) -> Result<Vec<PathBuf>, 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()))
 }

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;

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<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);
+        }
+    }
+}