fix(upload): retry streamed uploads on transient failures

Streamed (on-disk) uploads went through STREAMING_CLIENT, which has no
retry middleware because a consumed stream body can't be replayed. Add a
manual retry loop that rebuilds the stream from disk on each attempt and
reuses reqwest_retry's backoff policy and transient classifier, matching
the in-memory path's behavior.

Extract the shared backoff into request_client::upload_backoff() so both
the retry middleware and the streamed loop use one policy; under cfg(test)
its intervals shrink to milliseconds to keep the retry tests fast. Add
tests asserting both the streamed and in-memory paths retry transient
503s 1 + UPLOAD_RETRY_COUNT times.

Author: not-matthias

src/request_client.rs

@@ -4,15 +4,27 @@ use reqwest::ClientBuilder;
 use reqwest_middleware::{ClientBuilder as ClientWithMiddlewareBuilder, ClientWithMiddleware};
 use reqwest_retry::{RetryTransientMiddleware, policies::ExponentialBackoff};
 
-const UPLOAD_RETRY_COUNT: u32 = 3;
+pub const UPLOAD_RETRY_COUNT: u32 = 3;
 const OIDC_RETRY_COUNT: u32 = 10;
 const USER_AGENT: &str = "codspeed-runner";
 
+/// Shared backoff policy for upload retries, used both by the retry middleware on
+/// [`REQUEST_CLIENT`] and by the manual stream-retry loop in the uploader. Under
+/// `cfg(test)` the intervals are shrunk to milliseconds so retry tests don't sleep
+/// through the real exponential backoff (1s, 2s, 4s).
+pub fn upload_backoff() -> ExponentialBackoff {
+    let builder = ExponentialBackoff::builder();
+    #[cfg(test)]
+    let builder = builder.retry_bounds(
+        std::time::Duration::from_millis(1),
+        std::time::Duration::from_millis(5),
+    );
+    builder.build_with_max_retries(UPLOAD_RETRY_COUNT)
+}
+
 pub static REQUEST_CLIENT: LazyLock<ClientWithMiddleware> = LazyLock::new(|| {
     ClientWithMiddlewareBuilder::new(ClientBuilder::new().user_agent(USER_AGENT).build().unwrap())
-        .with(RetryTransientMiddleware::new_with_policy(
-            ExponentialBackoff::builder().build_with_max_retries(UPLOAD_RETRY_COUNT),
-        ))
+        .with(RetryTransientMiddleware::new_with_policy(upload_backoff()))
         .build()
 });
 

src/upload/uploader.rs

@@ -6,13 +6,17 @@ use crate::run_environment::RunEnvironment;
 use crate::upload::{UploadError, profile_archive::ProfileArchiveContent};
 use crate::{
     prelude::*,
-    request_client::{REQUEST_CLIENT, STREAMING_CLIENT},
+    request_client::{REQUEST_CLIENT, STREAMING_CLIENT, upload_backoff},
 };
 use async_compression::tokio::write::GzipEncoder;
 use console::style;
 use reqwest::StatusCode;
+use reqwest_retry::{
+    DefaultRetryableStrategy, RetryDecision, RetryPolicy, Retryable, RetryableStrategy,
+};
 use serde_json::Value;
 use std::collections::BTreeMap;
+use std::time::SystemTime;
 use tokio::fs::File;
 use tokio::io::AsyncWriteExt;
 use tokio_tar::Builder;
@@ -182,6 +186,66 @@ async fn retrieve_upload_data(
     }
 }
 
+/// Streams a file-backed body with retries. The retry middleware can't replay a
+/// stream body (a consumed stream isn't cloneable), so we rebuild it from disk on
+/// each attempt and reuse reqwest_retry's backoff policy + transient classification
+/// to match the behavior of the middleware-backed in-memory path.
+async fn send_streamed_with_retry(
+    upload_data: &UploadData,
+    path: &std::path::Path,
+    archive_size: u64,
+    archive_hash: &str,
+    encoding: Option<String>,
+) -> Result<reqwest::Response> {
+    let policy = upload_backoff();
+    let start = SystemTime::now();
+    let mut n_past_retries = 0;
+
+    loop {
+        let file = File::open(path)
+            .await
+            .context(format!("Failed to open file at path: {}", path.display()))?;
+        let stream = tokio_util::io::ReaderStream::new(file);
+        let body = reqwest::Body::wrap_stream(stream);
+
+        let mut request = STREAMING_CLIENT
+            .put(upload_data.upload_url.clone())
+            .header("Content-Type", "application/x-tar")
+            .header("Content-Length", archive_size)
+            .header("Content-MD5", archive_hash);
+        if let Some(encoding) = &encoding {
+            request = request.header("Content-Encoding", encoding);
+        }
+
+        // Wrap into reqwest_middleware::Error so the default classifier can be reused.
+        let result = request
+            .body(body)
+            .send()
+            .await
+            .map_err(reqwest_middleware::Error::Reqwest);
+
+        let is_transient = matches!(
+            DefaultRetryableStrategy.handle(&result),
+            Some(Retryable::Transient)
+        );
+        if is_transient {
+            if let RetryDecision::Retry { execute_after } =
+                policy.should_retry(start, n_past_retries)
+            {
+                let wait = execute_after
+                    .duration_since(SystemTime::now())
+                    .unwrap_or_default();
+                debug!("Streamed upload attempt failed (transient), retrying in {wait:?}");
+                tokio::time::sleep(wait).await;
+                n_past_retries += 1;
+                continue;
+            }
+        }
+
+        return Ok(result?);
+    }
+}
+
 async fn upload_profile_archive(
     upload_data: &UploadData,
     profile_archive: ProfileArchive,
@@ -206,24 +270,14 @@ async fn upload_profile_archive(
         }
         content @ ProfileArchiveContent::UncompressedOnDisk { path }
         | content @ ProfileArchiveContent::CompressedOnDisk { path } => {
-            // Use streaming client without retry middleware for file streams
-            let file = File::open(path)
-                .await
-                .context(format!("Failed to open file at path: {}", path.display()))?;
-            let stream = tokio_util::io::ReaderStream::new(file);
-            let body = reqwest::Body::wrap_stream(stream);
-
-            let mut request = STREAMING_CLIENT
-                .put(upload_data.upload_url.clone())
-                .header("Content-Type", "application/x-tar")
-                .header("Content-Length", archive_size)
-                .header("Content-MD5", archive_hash);
-
-            if let Some(encoding) = content.encoding() {
-                request = request.header("Content-Encoding", encoding);
-            }
-
-            request.body(body).send().await?
+            send_streamed_with_retry(
+                upload_data,
+                path,
+                archive_size,
+                &archive_hash,
+                content.encoding(),
+            )
+            .await?
         }
     };
 
@@ -378,4 +432,111 @@ mod tests {
         )
         .await;
     }
+
+    /// Number of upload attempts a transient failure should produce: the initial
+    /// attempt plus [`UPLOAD_RETRY_COUNT`] retries.
+    const EXPECTED_ATTEMPTS: usize = crate::request_client::UPLOAD_RETRY_COUNT as usize + 1;
+
+    /// Spawns a local TCP server that answers `503` (a retryable status) to each of
+    /// the next `max_conns` connections, then exits. Returns the URL, a counter of
+    /// connections received (= upload attempts), and the server's join handle so the
+    /// caller can wait for clean shutdown instead of leaking the thread.
+    fn spawn_mock_returning_503(
+        max_conns: usize,
+    ) -> (
+        String,
+        std::sync::Arc<std::sync::atomic::AtomicUsize>,
+        std::thread::JoinHandle<()>,
+    ) {
+        use std::io::{Read, Write};
+        use std::net::TcpListener;
+        use std::sync::Arc;
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let url = format!("http://{}/upload", listener.local_addr().unwrap());
+        let hits = Arc::new(AtomicUsize::new(0));
+
+        let hits_loop = hits.clone();
+        let handle = std::thread::spawn(move || {
+            for stream in listener.incoming().take(max_conns) {
+                let Ok(mut stream) = stream else { continue };
+                hits_loop.fetch_add(1, Ordering::SeqCst);
+                let mut buf = [0u8; 2048];
+                let _ = stream.read(&mut buf);
+                let body = "transient";
+                let resp = format!(
+                    "HTTP/1.1 503 Service Unavailable\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
+                    body.len(),
+                    body
+                );
+                let _ = stream.write_all(resp.as_bytes());
+            }
+        });
+
+        (url, hits, handle)
+    }
+
+    fn upload_data_for(url: String) -> UploadData {
+        UploadData {
+            status: "success".to_string(),
+            upload_url: url,
+            run_id: "test-run".to_string(),
+        }
+    }
+
+    /// WallTime/Memory path: an on-disk archive is streamed via `STREAMING_CLIENT`,
+    /// which has no retry middleware, so `send_streamed_with_retry` rebuilds the
+    /// stream on each attempt. A transient `503` must be retried `UPLOAD_RETRY_COUNT`
+    /// times.
+    #[tokio::test]
+    async fn streamed_upload_is_retried() {
+        use std::sync::atomic::Ordering;
+
+        let (url, hits, server) = spawn_mock_returning_503(EXPECTED_ATTEMPTS);
+
+        let path = tempfile::NamedTempFile::new()
+            .unwrap()
+            .into_temp_path()
+            .keep()
+            .unwrap();
+        std::fs::write(&path, b"profile-archive").unwrap();
+        let archive = ProfileArchive::new_uncompressed_on_disk(path).unwrap();
+
+        let result = upload_profile_archive(&upload_data_for(url), archive).await;
+        server.join().unwrap();
+
+        assert!(
+            result.is_err(),
+            "a 503 should surface as an error after retries"
+        );
+        assert_eq!(
+            hits.load(Ordering::SeqCst),
+            EXPECTED_ATTEMPTS,
+            "streamed upload should be attempted 1 + UPLOAD_RETRY_COUNT times"
+        );
+    }
+
+    /// Valgrind path: an in-memory archive goes through `REQUEST_CLIENT`, which retries
+    /// transient failures `UPLOAD_RETRY_COUNT` times. Confirms the same `503` mock is
+    /// genuinely retryable, so the single attempt above is due to the client, not the
+    /// status code.
+    #[tokio::test]
+    async fn in_memory_upload_is_retried() {
+        use std::sync::atomic::Ordering;
+
+        let (url, hits, server) = spawn_mock_returning_503(EXPECTED_ATTEMPTS);
+
+        let archive = ProfileArchive::new_compressed_in_memory(b"profile-archive".to_vec());
+
+        let result = upload_profile_archive(&upload_data_for(url), archive).await;
+        server.join().unwrap();
+
+        assert!(result.is_err(), "a 503 should surface as an error");
+        assert_eq!(
+            hits.load(Ordering::SeqCst),
+            EXPECTED_ATTEMPTS,
+            "in-memory upload should be attempted 1 + UPLOAD_RETRY_COUNT times"
+        );
+    }
 }