From 4e4c21f8e1c819dbb5164a30b34d67083186a177 Mon Sep 17 00:00:00 2001 From: Luc Varoqui Date: Mon, 5 Oct 2026 18:34:27 +0200 Subject: [PATCH 1/3] feat(upload): upload profile archives as concurrent multipart uploads S3 rejects single uploads above 5 GiB, and a single connection to S3 only reaches about 20-25 MiB/s on GitHub-hosted runners, so large profile archives were slow or impossible to upload. Every archive is now sent as an S3 multipart upload, replacing the single-request upload: about two parts per concurrent upload, each between 16 MiB and 256 MiB, so a small archive is a single part. The md5 of every part is computed in the same pass as the archive md5 and sent in the upload metadata (version 12) as `profileMultipart`. The API answers with `multipartUploadUrls`: the parts are uploaded 8 at a time (overridable with `CODSPEED_UPLOAD_CONCURRENCY`), each with its own retries, then the upload is completed with the part ETags in order. This applies to both on-disk and in-memory (gzip) archives. This requires an upload endpoint that accepts metadata version 12 and answers with `multipartUploadUrls`. Walltime profile folders above 5 GiB are no longer gzipped on disk to fit in a single request, and the runner no longer caps the archive size itself: the upload endpoint rejects archives above its limit, with the reason shown in the runner output, so the limit can change without a runner release. Archives are now hashed while streaming on the blocking thread pool, instead of being read whole into memory on the async runtime. Closes COD-3700 Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + Cargo.toml | 1 + src/run_environment/provider.rs | 3 +- src/upload/interfaces.rs | 51 +- src/upload/mod.rs | 1 + src/upload/profile_archive.rs | 271 +++++++- src/upload/s3.rs | 152 ++++ ...ata__tests__get_local_metadata_hash-2.snap | 9 +- ..._metadata__tests__get_metadata_hash-2.snap | 11 +- src/upload/upload_metadata.rs | 20 +- src/upload/uploader.rs | 651 +++++++++++++----- 11 files changed, 955 insertions(+), 216 deletions(-) create mode 100644 src/upload/s3.rs diff --git a/Cargo.lock b/Cargo.lock index 6a23121e5..acae15118 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -708,6 +708,7 @@ dependencies = [ "axoupdater", "base64", "bincode", + "bytes", "caps", "clap", "console", diff --git a/Cargo.toml b/Cargo.toml index 9df21458c..4587aa90d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -69,6 +69,7 @@ uuid = { version = "1.23.1", features = ["v4"] } which = "8.0.2" crc32fast = "1.5.0" samply = { path = "crates/samply-codspeed/samply" } +bytes = "1" # Memory profiling (memtrack) and the capability handling around it are Linux-only. [target.'cfg(target_os = "linux")'.dependencies] diff --git a/src/run_environment/provider.rs b/src/run_environment/provider.rs index 8858e16fa..7b887355e 100644 --- a/src/run_environment/provider.rs +++ b/src/run_environment/provider.rs @@ -136,8 +136,9 @@ pub trait RunEnvironmentProvider { tokenless: api_client.token().is_none(), repository_provider: self.get_repository_provider(), run_environment_metadata, - profile_md5: profile_archive.hash.clone(), + profile_md5: profile_archive.md5.clone(), profile_encoding: profile_archive.content.encoding(), + profile_multipart: profile_archive.multipart.clone(), commit_hash, allow_empty: config.allow_empty, runner: Runner { diff --git a/src/upload/interfaces.rs b/src/upload/interfaces.rs index fbf6abad7..376b3242d 100644 --- a/src/upload/interfaces.rs +++ b/src/upload/interfaces.rs @@ -5,7 +5,7 @@ use crate::instruments::InstrumentName; use crate::run_environment::{RepositoryProvider, RunEnvironment, RunEnvironmentMetadata, RunPart}; use crate::system::SystemInfo; -pub const LATEST_UPLOAD_METADATA_VERSION: u32 = 11; +pub const LATEST_UPLOAD_METADATA_VERSION: u32 = 12; #[derive(Serialize, Debug)] #[serde(rename_all = "camelCase")] @@ -15,6 +15,7 @@ pub struct UploadMetadata { pub tokenless: bool, pub profile_md5: String, pub profile_encoding: Option, + pub profile_multipart: ProfileMultipart, pub runner: Runner, pub run_environment: RunEnvironment, pub run_part: Option, @@ -24,6 +25,18 @@ pub struct UploadMetadata { pub run_environment_metadata: RunEnvironmentMetadata, } +/// Layout of a profile archive uploaded as an S3 multipart upload, in consecutive +/// `part_size` chunks, the last one holding the remainder. S3 requires parts of 5 MiB +/// to 5 GiB (the last one excepted from the minimum), and at most 10,000 of them. +#[derive(Serialize, Debug, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct ProfileMultipart { + pub size: u64, + pub part_size: u64, + /// Base64 md5 of each part, in upload order + pub part_md5s: Vec, +} + #[derive(Serialize, Debug)] #[serde(rename_all = "camelCase")] pub struct Runner { @@ -46,12 +59,46 @@ pub struct Runner { #[serde(rename_all = "camelCase")] pub struct UploadData { pub status: String, - pub upload_url: String, + pub multipart_upload_urls: MultipartUploadUrls, pub run_id: String, } +#[derive(Deserialize, Serialize, Debug)] +#[serde(rename_all = "camelCase")] +pub struct MultipartUploadUrls { + /// Presigned S3 `UploadPart` URLs, one per part, in upload order + pub part_urls: Vec, + /// Presigned S3 `CompleteMultipartUpload` URL + pub complete_url: String, +} + #[derive(Deserialize, Debug)] #[serde(rename_all = "camelCase")] pub struct UploadError { pub error: String, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_multipart_upload_response() { + let upload_data: UploadData = serde_json::from_str( + r#"{ + "status": "success", + "runId": "run-id", + "multipartUploadUrls": { + "partUrls": ["https://part/1", "https://part/2"], + "completeUrl": "https://complete" + } + }"#, + ) + .unwrap(); + + assert_eq!(upload_data.run_id, "run-id"); + let urls = upload_data.multipart_upload_urls; + assert_eq!(urls.part_urls, ["https://part/1", "https://part/2"]); + assert_eq!(urls.complete_url, "https://complete"); + } +} diff --git a/src/upload/mod.rs b/src/upload/mod.rs index cb4e855ff..aafc1d5cb 100644 --- a/src/upload/mod.rs +++ b/src/upload/mod.rs @@ -3,6 +3,7 @@ mod interfaces; pub mod poll_results; mod profile_archive; mod run_index_state; +mod s3; mod upload_metadata; mod uploader; diff --git a/src/upload/profile_archive.rs b/src/upload/profile_archive.rs index 5484b96cc..8adaa8185 100644 --- a/src/upload/profile_archive.rs +++ b/src/upload/profile_archive.rs @@ -1,59 +1,167 @@ use base64::{Engine, engine::general_purpose}; use crate::prelude::*; +use crate::upload::interfaces::ProfileMultipart; +use bytes::Bytes; +use std::io::{Read, SeekFrom}; use std::path::PathBuf; +use std::sync::LazyLock; +use tokio::io::{AsyncReadExt, AsyncSeekExt}; + +/// Default number of multipart upload parts sent at the same time +/// +/// Measured on `ubuntu:latest` and `macro runner gen 2` with a sweep from 1, 2, 4, 8, 16 +/// After 8 no more performance improvement was observed. +const DEFAULT_CONCURRENT_PART_UPLOADS: usize = 8; + +/// Overrides [`DEFAULT_CONCURRENT_PART_UPLOADS`] +const CONCURRENT_PART_UPLOADS_ENV: &str = "CODSPEED_UPLOAD_CONCURRENCY"; + +static CONCURRENT_PART_UPLOADS: LazyLock = LazyLock::new(|| { + let Ok(value) = std::env::var(CONCURRENT_PART_UPLOADS_ENV) else { + return DEFAULT_CONCURRENT_PART_UPLOADS; + }; + match value.parse::() { + Ok(concurrency) if concurrency > 0 => concurrency, + _ => { + warn!( + "Ignoring invalid {CONCURRENT_PART_UPLOADS_ENV} value {value:?}, using {DEFAULT_CONCURRENT_PART_UPLOADS}" + ); + DEFAULT_CONCURRENT_PART_UPLOADS + } + } +}); + +pub(super) fn concurrent_part_uploads() -> usize { + *CONCURRENT_PART_UPLOADS +} + +const MULTIPART_MIN_PART_SIZE_BYTES: u64 = 16 * 1024 * 1024; // 16 MiB +/// Bounds how much a failed part has to re-send +const MULTIPART_MAX_PART_SIZE_BYTES: u64 = 256 * 1024 * 1024; // 256 MiB +/// More parts than concurrent uploads, so that the upload slots freed by fast parts +/// pick up remaining work instead of idling while the slowest part finishes +const PARTS_PER_CONCURRENT_UPLOAD: u64 = 2; +const HASH_READ_BUFFER_SIZE: usize = 8 * 1024 * 1024; // 8 MiB #[derive(Debug)] pub struct ProfileArchive { - pub hash: String, + pub md5: String, pub content: ProfileArchiveContent, + pub multipart: ProfileMultipart, } #[derive(Debug)] pub enum ProfileArchiveContent { - CompressedInMemory { data: Vec }, + CompressedInMemory { data: Bytes }, UncompressedOnDisk { path: PathBuf }, - CompressedOnDisk { path: PathBuf }, } -impl ProfileArchive { - pub fn new_compressed_in_memory(data: Vec) -> Self { - let hash = general_purpose::STANDARD.encode(md5::compute(&data).0); - ProfileArchive { - hash, - content: ProfileArchiveContent::CompressedInMemory { data }, +fn encode_md5(digest: md5::Digest) -> String { + general_purpose::STANDARD.encode(digest.0) +} + +/// Read the content once to compute the md5 of all of it and of each consecutive +/// `part_size` chunk of it. +/// +/// CPU heavy + potentially reading from blocking IO, so it is recommended to +/// run this on a blocking thread pool. +fn compute_md5s_from_reader( + mut reader: impl Read, + part_size: u64, +) -> Result<(String, Vec)> { + let mut buffer = vec![0u8; HASH_READ_BUFFER_SIZE]; + let mut whole_context = md5::Context::new(); + let mut part_context = md5::Context::new(); + let mut part_length = 0u64; + let mut part_md5s = Vec::new(); + + loop { + let read = reader.read(&mut buffer)?; + if read == 0 { + break; + } + let mut chunk = &buffer[..read]; + whole_context.consume(chunk); + + while !chunk.is_empty() { + let taken = (part_size - part_length).min(chunk.len() as u64) as usize; + part_context.consume(&chunk[..taken]); + part_length += taken as u64; + chunk = &chunk[taken..]; + if part_length == part_size { + let finished_part = std::mem::replace(&mut part_context, md5::Context::new()); + part_md5s.push(encode_md5(finished_part.finalize())); + part_length = 0; + } } } + if part_length > 0 { + part_md5s.push(encode_md5(part_context.finalize())); + } - pub fn new_uncompressed_on_disk(path: PathBuf) -> Result { - let metadata = std::fs::metadata(&path)?; - if !metadata.is_file() { - return Err(anyhow!("The provided path is not a file")); + Ok((encode_md5(whole_context.finalize()), part_md5s)) +} + +/// [`compute_md5s_from_reader`] over the content, run on the blocking thread pool as +/// hashing the whole archive would otherwise stall the async runtime. +async fn compute_md5s( + content: &ProfileArchiveContent, + part_size: u64, +) -> Result<(String, Vec)> { + match content { + ProfileArchiveContent::CompressedInMemory { data } => { + let data = data.clone(); + tokio::task::spawn_blocking(move || compute_md5s_from_reader(&data[..], part_size)) + .await? } - let mut file = std::fs::File::open(&path)?; - let mut buffer = Vec::new(); - use std::io::Read; - file.read_to_end(&mut buffer)?; - let hash = general_purpose::STANDARD.encode(md5::compute(&buffer).0); - Ok(ProfileArchive { - hash, - content: ProfileArchiveContent::UncompressedOnDisk { path }, - }) + ProfileArchiveContent::UncompressedOnDisk { path } => { + let path = path.clone(); + tokio::task::spawn_blocking(move || { + compute_md5s_from_reader(std::fs::File::open(path)?, part_size) + }) + .await? + } + } +} + +/// Split the archive in about `PARTS_PER_CONCURRENT_UPLOAD` parts per concurrent upload, +/// within the part size bounds. An archive no larger than the minimum part size is a +/// single part. +fn choose_multipart_part_size(size: u64, concurrent_uploads: usize) -> u64 { + let target_part_count = concurrent_uploads as u64 * PARTS_PER_CONCURRENT_UPLOAD; + size.div_ceil(target_part_count) + .clamp(MULTIPART_MIN_PART_SIZE_BYTES, MULTIPART_MAX_PART_SIZE_BYTES) +} + +impl ProfileArchive { + pub async fn new_compressed_in_memory(data: Vec) -> Result { + Self::new(ProfileArchiveContent::CompressedInMemory { data: data.into() }).await } - pub fn new_compressed_on_disk(path: PathBuf) -> Result { - let metadata = std::fs::metadata(&path)?; + pub async fn new_uncompressed_on_disk(path: PathBuf) -> Result { + let metadata = tokio::fs::metadata(&path).await?; if !metadata.is_file() { return Err(anyhow!("The provided path is not a file")); } - let mut file = std::fs::File::open(&path)?; - let mut buffer = Vec::new(); - use std::io::Read; - file.read_to_end(&mut buffer)?; - let hash = general_purpose::STANDARD.encode(md5::compute(&buffer).0); + Self::new(ProfileArchiveContent::UncompressedOnDisk { path }).await + } + + async fn new(content: ProfileArchiveContent) -> Result { + let size = content.size().await?; + let part_size = choose_multipart_part_size(size, concurrent_part_uploads()); + + let (md5, part_md5s) = compute_md5s(&content, part_size).await?; + let multipart = ProfileMultipart { + size, + part_size, + part_md5s, + }; + Ok(ProfileArchive { - hash, - content: ProfileArchiveContent::CompressedOnDisk { path }, + md5, + content, + multipart, }) } } @@ -62,8 +170,7 @@ impl ProfileArchiveContent { pub async fn size(&self) -> Result { match &self { ProfileArchiveContent::CompressedInMemory { data } => Ok(data.len() as u64), - ProfileArchiveContent::UncompressedOnDisk { path } - | ProfileArchiveContent::CompressedOnDisk { path } => { + ProfileArchiveContent::UncompressedOnDisk { path } => { let metadata = tokio::fs::metadata(path).await?; Ok(metadata.len()) } @@ -73,20 +180,108 @@ impl ProfileArchiveContent { pub fn encoding(&self) -> Option { match self { ProfileArchiveContent::CompressedInMemory { .. } => Some("gzip".to_string()), - ProfileArchiveContent::CompressedOnDisk { .. } => Some("gzip".to_string()), - _ => None, + ProfileArchiveContent::UncompressedOnDisk { .. } => None, + } + } + + /// Request body holding `length` bytes of the content starting at `offset`. + /// On-disk content is streamed rather than loaded in memory. + pub async fn body(&self, offset: u64, length: u64) -> Result { + match self { + ProfileArchiveContent::CompressedInMemory { data } => { + let start = offset as usize; + Ok(data.slice(start..start + length as usize).into()) + } + ProfileArchiveContent::UncompressedOnDisk { path } => { + let mut file = tokio::fs::File::open(path) + .await + .context(format!("Failed to open file at path: {}", path.display()))?; + file.seek(SeekFrom::Start(offset)).await?; + let stream = tokio_util::io::ReaderStream::new(file.take(length)); + Ok(reqwest::Body::wrap_stream(stream)) + } } } } impl Drop for ProfileArchiveContent { fn drop(&mut self) { - if let ProfileArchiveContent::UncompressedOnDisk { path } - | ProfileArchiveContent::CompressedOnDisk { path } = self - { + if let ProfileArchiveContent::UncompressedOnDisk { path } = self { if path.exists() { let _ = std::fs::remove_file(path); } } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn write_temp_file(content: &[u8]) -> PathBuf { + let path = tempfile::NamedTempFile::new() + .unwrap() + .into_temp_path() + .keep() + .unwrap(); + std::fs::write(&path, content).unwrap(); + path + } + + #[test] + fn computes_whole_and_part_md5s_in_one_pass() { + // Not a multiple of the part size, and spanning several read buffers + let content: Vec = (0..HASH_READ_BUFFER_SIZE * 2 + 123) + .map(|i| (i % 251) as u8) + .collect(); + let part_size = (HASH_READ_BUFFER_SIZE / 3) as u64; + + let (hash, part_md5s) = compute_md5s_from_reader(&content[..], part_size).unwrap(); + + assert_eq!(hash, encode_md5(md5::compute(&content))); + let expected_part_md5s: Vec = content + .chunks(part_size as usize) + .map(|part| encode_md5(md5::compute(part))) + .collect(); + assert_eq!(part_md5s, expected_part_md5s); + } + + #[tokio::test] + async fn small_archive_is_a_single_part() { + let path = write_temp_file(b"profile-archive"); + + let archive = ProfileArchive::new_uncompressed_on_disk(path) + .await + .unwrap(); + + let md5 = encode_md5(md5::compute(b"profile-archive")); + assert_eq!(archive.md5, md5); + assert_eq!( + archive.multipart, + ProfileMultipart { + size: b"profile-archive".len() as u64, + part_size: MULTIPART_MIN_PART_SIZE_BYTES, + part_md5s: vec![md5], + } + ); + } + + #[test] + fn archive_is_split_in_two_parts_per_concurrent_upload() { + const MIB: u64 = 1024 * 1024; + assert_eq!(choose_multipart_part_size(1024 * MIB, 8), 64 * MIB); + } + + #[test] + fn part_size_is_clamped() { + const MIB: u64 = 1024 * 1024; + assert_eq!( + choose_multipart_part_size(64 * MIB, 8), + MULTIPART_MIN_PART_SIZE_BYTES + ); + assert_eq!( + choose_multipart_part_size(15 * 1024 * MIB, 8), + MULTIPART_MAX_PART_SIZE_BYTES + ); + } +} diff --git a/src/upload/s3.rs b/src/upload/s3.rs new file mode 100644 index 000000000..2c621a5aa --- /dev/null +++ b/src/upload/s3.rs @@ -0,0 +1,152 @@ +//! Client side of S3 uploads over presigned URLs. +//! +//! An archive is uploaded as a multipart upload created beforehand: one +//! [UploadPart](https://docs.aws.amazon.com/AmazonS3/latest/API/API_UploadPart.html) +//! request per part, then a +//! [CompleteMultipartUpload](https://docs.aws.amazon.com/AmazonS3/latest/API/API_CompleteMultipartUpload.html) +//! request. This module handles the S3-specific parts of these requests. +//! +//! The part URLs are signed with the `Content-MD5` and `Content-Length` headers, so S3 +//! rejects a part that does not send them with the signed values. The content type and +//! encoding of the archive are set when the multipart upload is created. + +use crate::prelude::*; +use console::style; + +/// `UploadPart` request uploading one part of a multipart upload. +pub(super) fn upload_part( + client: &reqwest::Client, + url: &str, + content_md5: &str, +) -> reqwest::RequestBuilder { + client.put(url).header("Content-MD5", content_md5) +} + +/// ETag S3 assigned to an uploaded part, only obtainable from a part upload response +#[derive(Debug)] +pub(super) struct PartETag(String); + +/// Read the ETag S3 assigned to an uploaded part, needed to complete the upload. +pub(super) fn part_etag(response: &reqwest::Response) -> Result { + Ok(PartETag( + response + .headers() + .get(reqwest::header::ETAG) + .context("Missing ETag in the part upload response")? + .to_str()? + .to_owned(), + )) +} + +/// `CompleteMultipartUpload` request assembling the uploaded parts into the final +/// object. `etags` are in part order. +pub(super) fn complete_upload( + client: &reqwest::Client, + url: &str, + etags: &[PartETag], +) -> reqwest::RequestBuilder { + client + .post(url) + .header("Content-Type", "application/xml") + .body(build_complete_body(etags)) +} + +/// Error reported by S3 in the XML body of a response +#[derive(Debug)] +pub(super) struct S3Error { + code: Option, + body: String, +} + +impl S3Error { + fn from_body(body: String) -> Option { + if !body.contains("") { + return None; + } + let code = body + .split_once("") + .and_then(|(_, rest)| rest.split_once("")) + .map(|(code, _)| code.to_owned()); + Some(S3Error { code, body }) + } + + /// Whether S3 documents the error as worth retrying, as opposed to errors such as + /// `InvalidPart` that the same request would hit again. + /// + pub(super) fn is_transient(&self) -> bool { + matches!( + self.code.as_deref(), + Some("InternalError" | "ServiceUnavailable" | "SlowDown" | "RequestTimeout") + ) + } +} + +impl std::fmt::Display for S3Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "Failed to complete the performance report upload: {}\n -> {} {}", + self.code.as_deref().unwrap_or("unknown S3 error"), + style("Reason:").bold(), + style(&self.body).red() + ) + } +} + +impl std::error::Error for S3Error {} + +/// Check the response of a successful `CompleteMultipartUpload` request. S3 sends the +/// `200 OK` status as soon as it starts assembling the parts, and reports a failure +/// happening afterwards in the body. +/// +pub(super) async fn check_complete_upload_response( + response: reqwest::Response, +) -> reqwest::Result> { + let body = response.text().await?; + Ok(match S3Error::from_body(body) { + Some(error) => Err(error), + None => Ok(()), + }) +} + +fn build_complete_body(etags: &[PartETag]) -> String { + let parts: String = etags + .iter() + .enumerate() + .map(|(index, PartETag(etag))| { + format!( + "{}{}", + index + 1, + etag + ) + }) + .collect(); + format!("{parts}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn success_body_is_not_an_error() { + let body = + "\"etag\""; + assert!(S3Error::from_body(body.to_owned()).is_none()); + } + + #[test] + fn internal_error_is_transient() { + let error = S3Error::from_body( + "InternalErrorPlease try again.".into(), + ) + .unwrap(); + assert!(error.is_transient()); + } + + #[test] + fn invalid_part_is_permanent() { + let error = S3Error::from_body("InvalidPart".into()).unwrap(); + assert!(!error.is_transient()); + } +} diff --git a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap index 1679aa850..832d7553c 100644 --- a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap +++ b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap @@ -4,10 +4,17 @@ expression: upload_metadata --- { "repositoryProvider": "PROJECT", - "version": 11, + "version": 12, "tokenless": false, "profileMd5": "tfC4VxYiYdJcTWpHpv4Ouw==", "profileEncoding": "gzip", + "profileMultipart": { + "size": 20480, + "partSize": 16777216, + "partMd5s": [ + "tfC4VxYiYdJcTWpHpv4Ouw==" + ] + }, "runner": { "name": "codspeed-runner", "version": "4.11.1", diff --git a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap index fd695b09f..ea985f875 100644 --- a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap +++ b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap @@ -4,10 +4,19 @@ expression: upload_metadata --- { "repositoryProvider": "GITHUB", - "version": 11, + "version": 12, "tokenless": true, "profileMd5": "jp/k05RKuqP3ERQuIIvx4Q==", "profileEncoding": "gzip", + "profileMultipart": { + "size": 629145600, + "partSize": 268435456, + "partMd5s": [ + "eEKfdGLWNqhNnJIvSVWZxQ==", + "vKfHJAI2Hyo68jWwUc6H9A==", + "bQY//BUrm3jASNTyWj/3Aw==" + ] + }, "runner": { "name": "codspeed-runner", "version": "2.1.0", diff --git a/src/upload/upload_metadata.rs b/src/upload/upload_metadata.rs index b90e75070..fc95c4c2d 100644 --- a/src/upload/upload_metadata.rs +++ b/src/upload/upload_metadata.rs @@ -22,7 +22,7 @@ mod tests { Sender, }; use crate::system::SystemInfo; - use crate::upload::{LATEST_UPLOAD_METADATA_VERSION, Runner, UploadMetadata}; + use crate::upload::{LATEST_UPLOAD_METADATA_VERSION, ProfileMultipart, Runner, UploadMetadata}; #[test] fn test_get_metadata_hash() { @@ -32,6 +32,15 @@ mod tests { tokenless: true, profile_md5: "jp/k05RKuqP3ERQuIIvx4Q==".into(), profile_encoding: Some("gzip".into()), + profile_multipart: ProfileMultipart { + size: 629145600, + part_size: 268435456, + part_md5s: vec![ + "eEKfdGLWNqhNnJIvSVWZxQ==".into(), + "vKfHJAI2Hyo68jWwUc6H9A==".into(), + "bQY//BUrm3jASNTyWj/3Aw==".into(), + ], + }, runner: Runner { name: "codspeed-runner".into(), version: "2.1.0".into(), @@ -73,7 +82,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"b6e221583869b0a49498d71538432a5396b08fec7783081dcdf23c9e037a9365" + @"307ed5a2218e04b1befe21d355e19a2a3adb7a617b19d1b6e8406c23355a2664" ); assert_json_snapshot!(upload_metadata); } @@ -86,6 +95,11 @@ mod tests { tokenless: false, profile_md5: "tfC4VxYiYdJcTWpHpv4Ouw==".into(), profile_encoding: Some("gzip".into()), + profile_multipart: ProfileMultipart { + size: 20480, + part_size: 16777216, + part_md5s: vec!["tfC4VxYiYdJcTWpHpv4Ouw==".into()], + }, runner: Runner { name: "codspeed-runner".into(), version: "4.11.1".into(), @@ -144,7 +158,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"5c960260ea5b5ceaafa20ea220e566743e20b65106d0fd672a844a63bba3835a" + @"dcbb07dc8bf0c97cc88566ef78acec873f452420ce004807f3cdd6218544a458" ); assert_json_snapshot!(upload_metadata); } diff --git a/src/upload/uploader.rs b/src/upload/uploader.rs index 1afa4715a..b018b4e1e 100644 --- a/src/upload/uploader.rs +++ b/src/upload/uploader.rs @@ -10,53 +10,48 @@ use crate::{ }; use async_compression::tokio::write::GzipEncoder; use console::style; +use futures::{StreamExt, TryStreamExt}; use reqwest::StatusCode; use reqwest_retry::{ DefaultRetryableStrategy, RetryDecision, RetryPolicy, Retryable, RetryableStrategy, }; use serde_json::Value; use std::collections::BTreeMap; -use std::time::SystemTime; +use std::time::{Duration, Instant, SystemTime}; use tokio::fs::File; use tokio::io::AsyncWriteExt; use tokio_tar::Builder; -use super::interfaces::{UploadData, UploadMetadata}; -use super::profile_archive::ProfileArchive; - -fn bytes_to_mib(bytes: u64) -> u64 { - bytes / (1024 * 1024) -} - -/// Maximum allowed profile archive size in bytes before upload is rejected -const MAX_UPLOAD_PROFILE_SIZE_BYTES: u64 = 1024 * 1024 * 1024 * 5; // 5 GiB - -/// Calculate the total size of a directory in bytes -async fn calculate_folder_size(path: &std::path::Path) -> Result { - let mut total_size = 0u64; - let mut dirs_to_process = vec![path.to_path_buf()]; - - while let Some(current_dir) = dirs_to_process.pop() { - let mut entries = tokio::fs::read_dir(¤t_dir).await?; - - while let Some(entry) = entries.next_entry().await? { - let metadata = entry.metadata().await?; - if metadata.is_file() { - total_size += metadata.len(); - } else if metadata.is_dir() { - dirs_to_process.push(entry.path()); - } - } +use super::interfaces::{MultipartUploadUrls, ProfileMultipart, UploadData, UploadMetadata}; +use super::profile_archive::{ProfileArchive, concurrent_part_uploads}; +use super::s3; + +fn human_bytes(bytes: u64) -> String { + const KB: u64 = 1024; + const MB: u64 = KB * 1024; + const GB: u64 = MB * 1024; + + if bytes >= GB { + format!("{:.2} GB", bytes as f64 / GB as f64) + } else if bytes >= MB { + format!("{:.2} MB", bytes as f64 / MB as f64) + } else if bytes >= KB { + format!("{:.2} KB", bytes as f64 / KB as f64) + } else { + format!("{bytes}") } +} - Ok(total_size) +fn human_bytes_per_second(bytes: u64, elapsed: std::time::Duration) -> String { + let bytes_per_second = bytes as f64 / elapsed.as_secs_f64(); + format!("{}/s", human_bytes(bytes_per_second as u64)) } /// Create a profile archive from the profile folder and return its md5 hash encoded in base64 /// /// For Valgrind, we create a gzip-compressed tar archive of the entire profile folder. -/// For WallTime, we check the folder size and create either a compressed or uncompressed tar archive -/// based on the [`MAX_UPLOAD_PROFILE_SIZE_BYTES`] threshold. +/// For WallTime and Memory, we create an uncompressed tar archive on disk: their +/// profiles are already compressed, so gzip would barely shrink them. async fn create_profile_archive( profile_folder: &std::path::Path, executor_name: ExecutorName, @@ -71,13 +66,10 @@ async fn create_profile_archive( let mut gzip_encoder = tar.into_inner().await?; gzip_encoder.shutdown().await?; let data = gzip_encoder.into_inner(); - ProfileArchive::new_compressed_in_memory(data) + ProfileArchive::new_compressed_in_memory(data).await? } ExecutorName::Memory | ExecutorName::WallTime => { - // Check folder size to decide on compression - let folder_size_bytes = calculate_folder_size(profile_folder).await?; - let should_compress = folder_size_bytes >= MAX_UPLOAD_PROFILE_SIZE_BYTES; - + debug!("Creating uncompressed tar archive on disk"); let temp_file = tempfile::NamedTempFile::new()?; let temp_path = temp_file.path().to_path_buf(); @@ -87,32 +79,11 @@ async fn create_profile_archive( // Persist the temporary file to prevent deletion when temp_file goes out of scope let persistent_path = temp_file.into_temp_path().keep()?; - if should_compress { - debug!( - "Profile folder size ({} MiB) exceeds threshold ({} MiB), creating compressed tar.gz archive on disk", - bytes_to_mib(folder_size_bytes), - bytes_to_mib(MAX_UPLOAD_PROFILE_SIZE_BYTES) - ); - let enc = GzipEncoder::new(file); - let mut tar = Builder::new(enc); - tar.append_dir_all(".", profile_folder).await?; - let mut gzip_encoder = tar.into_inner().await?; - gzip_encoder.shutdown().await?; - gzip_encoder.into_inner().sync_all().await?; - - ProfileArchive::new_compressed_on_disk(persistent_path)? - } else { - debug!( - "Profile folder size ({} MiB) is below threshold ({} MiB), creating uncompressed tar archive on disk", - bytes_to_mib(folder_size_bytes), - bytes_to_mib(MAX_UPLOAD_PROFILE_SIZE_BYTES) - ); - let mut tar = Builder::new(file); - tar.append_dir_all(".", profile_folder).await?; - tar.into_inner().await?.sync_all().await?; + let mut tar = Builder::new(file); + tar.append_dir_all(".", profile_folder).await?; + tar.into_inner().await?.sync_all().await?; - ProfileArchive::new_uncompressed_on_disk(persistent_path)? - } + ProfileArchive::new_uncompressed_on_disk(persistent_path).await? } }; @@ -123,10 +94,6 @@ async fn create_profile_archive( time_start.elapsed() ); - if archive_size > MAX_UPLOAD_PROFILE_SIZE_BYTES { - bail!("Profile archive exceeds the maximum allowed size"); - } - Ok(profile_archive) } @@ -195,110 +162,217 @@ async fn retrieve_upload_data( } } -/// The retry middleware can't replay a consumed stream, so we rebuild the body from -/// disk on each attempt. Response-level errors (4xx/5xx) are left for the caller. -async fn send_streamed_with_retry( - upload_data: &UploadData, - path: &std::path::Path, - archive_size: u64, - archive_hash: &str, - encoding: Option, -) -> Result { +/// A byte range of the archive content, sent as a request body. +struct ContentRange<'a> { + content: &'a ProfileArchiveContent, + offset: u64, + length: u64, +} + +impl ContentRange<'_> { + /// Attach this range as the body of `request`. The body is rebuilt on every call, + /// since a streamed body is consumed by the request that sends it. + async fn attach(&self, request: reqwest::RequestBuilder) -> Result { + let body = self.content.body(self.offset, self.length).await?; + Ok(request.header("Content-Length", self.length).body(body)) + } +} + +/// Failure of one attempt of an upload request +enum AttemptError { + /// Worth sending the request again + Transient(Error), + Permanent(Error), +} + +impl From for AttemptError { + fn from(error: Error) -> Self { + AttemptError::Permanent(error) + } +} + +/// Run `attempt` until it succeeds, retrying transient failures with the +/// [`upload_backoff`] policy. +async fn with_upload_retry(mut attempt: impl FnMut() -> Fut) -> Result +where + Fut: Future>, +{ 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); - } - - 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?); + let error = match attempt().await { + Ok(value) => return Ok(value), + Err(AttemptError::Permanent(error)) => return Err(error), + Err(AttemptError::Transient(error)) => error, + }; + let RetryDecision::Retry { execute_after } = policy.should_retry(start, n_past_retries) + else { + return Err(error); + }; + let wait = execute_after + .duration_since(SystemTime::now()) + .unwrap_or_default(); + debug!("Upload attempt failed (transient), retrying in {wait:?}: {error}"); + tokio::time::sleep(wait).await; + n_past_retries += 1; } } -async fn upload_profile_archive( - upload_data: &UploadData, - profile_archive: ProfileArchive, -) -> Result<()> { - let archive_size = profile_archive.content.size().await?; - let archive_hash = profile_archive.hash; - - let response = match &profile_archive.content { - content @ ProfileArchiveContent::CompressedInMemory { data } => { - // Use regular client with retry middleware for compressed data - let mut request = REQUEST_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(data.clone()).send().await? - } - content @ ProfileArchiveContent::UncompressedOnDisk { path } - | content @ ProfileArchiveContent::CompressedOnDisk { path } => { - send_streamed_with_retry( - upload_data, - path, - archive_size, - &archive_hash, - content.encoding(), +/// Send an upload request, failing on a non-success status. Connection errors and +/// statuses such as 5xx or 429 are reported as transient. +async fn send_upload_request( + request: reqwest::RequestBuilder, +) -> std::result::Result { + /// Bounds reading the body of a failed upload response, only used in the error message + const ERROR_BODY_READ_TIMEOUT: Duration = Duration::from_secs(10); + + let result = request + .send() + .await + .map_err(reqwest_middleware::Error::Reqwest); + let is_transient = matches!( + DefaultRetryableStrategy.handle(&result), + Some(Retryable::Transient) + ); + let error = match result { + Ok(response) if response.status().is_success() => return Ok(response), + Ok(response) => { + let status = response.status(); + // A stalled error body must not keep a retryable failure from being retried + let error_text = tokio::time::timeout(ERROR_BODY_READ_TIMEOUT, response.text()) + .await + .ok() + .and_then(Result::ok) + .unwrap_or_default(); + anyhow!( + "Failed to upload performance report: {}\n -> {} {}", + status, + style("Reason:").bold(), + style(error_text).red() ) - .await? } + Err(error) => error.into(), }; + Err(if is_transient { + AttemptError::Transient(error) + } else { + AttemptError::Permanent(error) + }) +} - if !response.status().is_success() { - let status = response.status(); - let error_text = response.text().await?; +/// Upload the parts concurrently, then assemble them with the S3 complete request. +/// Parts of a failed upload are cleaned up by the bucket lifecycle rules. +async fn upload_multipart_profile_archive( + multipart_upload_urls: &MultipartUploadUrls, + multipart: &ProfileMultipart, + content: &ProfileArchiveContent, +) -> Result<()> { + debug!("Starting multipart upload for profile archive"); + debug!( + "Multipart upload details: part_count={}, part_size={}, total_size={}", + multipart.part_md5s.len(), + multipart.part_size, + multipart.size + ); + + let part_count = multipart.part_md5s.len(); + if multipart_upload_urls.part_urls.len() != part_count { bail!( - "Failed to upload performance report: {}\n -> {} {}", - status, - style("Reason:").bold(), - style(error_text).red() + "Received {} part upload URLs for {} parts", + multipart_upload_urls.part_urls.len(), + part_count ); } - Ok(()) + let concurrency = concurrent_part_uploads(); + let upload_start = Instant::now(); + // Unordered, so that a part finishing frees its slot even while an earlier part is still + // uploading + let mut indexed_etags: Vec<_> = futures::stream::iter( + multipart_upload_urls + .part_urls + .iter() + .zip(&multipart.part_md5s) + .enumerate(), + ) + .map(|(index, (part_url, part_md5))| async move { + let offset = index as u64 * multipart.part_size; + let range = ContentRange { + content, + offset, + length: multipart.part_size.min(multipart.size - offset), + }; + debug!( + "Uploading part {}/{} ({} bytes)", + index + 1, + part_count, + range.length + ); + let part_start = Instant::now(); + let etag = with_upload_retry(|| async { + let request = s3::upload_part(&STREAMING_CLIENT, part_url, part_md5); + let response = send_upload_request(range.attach(request).await?).await?; + Ok(s3::part_etag(&response)?) + }) + .await?; + let part_elapsed = part_start.elapsed(); + debug!( + "Uploaded part {}/{} in {:.1?} ({})", + index + 1, + part_count, + part_elapsed, + human_bytes_per_second(range.length, part_elapsed) + ); + Ok::<_, anyhow::Error>((index, etag)) + }) + .buffer_unordered(concurrency) + .try_collect() + .await?; + + indexed_etags.sort_unstable_by_key(|(index, _)| *index); + let etags: Vec<_> = indexed_etags.into_iter().map(|(_, etag)| etag).collect(); + + let upload_elapsed = upload_start.elapsed(); + info!( + "Uploaded {} part{} ({}) in {:.1?} with {} concurrent uploads ({})", + part_count, + if part_count > 1 { "s" } else { "" }, + human_bytes(multipart.size), + upload_elapsed, + concurrency, + human_bytes_per_second(multipart.size, upload_elapsed) + ); + + with_upload_retry(|| async { + let request = s3::complete_upload( + &STREAMING_CLIENT, + &multipart_upload_urls.complete_url, + &etags, + ); + let response = send_upload_request(request).await?; + match s3::check_complete_upload_response(response).await { + Ok(Ok(())) => Ok(()), + Ok(Err(error)) if error.is_transient() => Err(AttemptError::Transient(error.into())), + Ok(Err(error)) => Err(AttemptError::Permanent(error.into())), + // The connection stays open while S3 assembles the parts, and can drop + Err(error) => Err(AttemptError::Transient(error.into())), + } + }) + .await +} + +async fn upload_profile_archive( + upload_data: &UploadData, + profile_archive: ProfileArchive, +) -> Result<()> { + upload_multipart_profile_archive( + &upload_data.multipart_upload_urls, + &profile_archive.multipart, + &profile_archive.content, + ) + .await } #[derive(Clone)] @@ -442,7 +516,7 @@ mod tests { const EXPECTED_ATTEMPTS: usize = crate::request_client::UPLOAD_RETRY_COUNT as usize + 1; /// Answers `503` to each of the next `max_conns` connections, then exits. Returns - /// the URL, a counter of connections received, and the server's join handle. + /// the base URL, a counter of connections received, and the server's join handle. fn spawn_mock_returning_503( max_conns: usize, ) -> ( @@ -456,7 +530,7 @@ mod tests { 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 url = format!("http://{}", listener.local_addr().unwrap()); let hits = Arc::new(AtomicUsize::new(0)); let hits_loop = hits.clone(); @@ -479,16 +553,8 @@ mod tests { (url, hits, handle) } - fn upload_data_for(url: String) -> UploadData { - UploadData { - status: "success".to_string(), - upload_url: url, - run_id: "test-run".to_string(), - } - } - - /// On-disk archives stream through `send_streamed_with_retry`, which retries - /// transient failures itself since `STREAMING_CLIENT` has no retry middleware. + /// `send_with_retry` retries transient failures itself, since `STREAMING_CLIENT` + /// has no retry middleware. #[tokio::test] async fn streamed_upload_is_retried() { use std::sync::atomic::Ordering; @@ -501,9 +567,11 @@ mod tests { .keep() .unwrap(); std::fs::write(&path, b"profile-archive").unwrap(); - let archive = ProfileArchive::new_uncompressed_on_disk(path).unwrap(); + let archive = ProfileArchive::new_uncompressed_on_disk(path) + .await + .unwrap(); - let result = upload_profile_archive(&upload_data_for(url), archive).await; + let result = upload_profile_archive(&multipart_upload_data_for(&url, 1), archive).await; server.join().unwrap(); assert!( @@ -517,17 +585,17 @@ mod tests { ); } - /// In-memory archives go through `REQUEST_CLIENT`, whose retry middleware handles - /// transient failures. #[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 archive = ProfileArchive::new_compressed_in_memory(b"profile-archive".to_vec()) + .await + .unwrap(); - let result = upload_profile_archive(&upload_data_for(url), archive).await; + let result = upload_profile_archive(&multipart_upload_data_for(&url, 1), archive).await; server.join().unwrap(); assert!(result.is_err(), "a 503 should surface as an error"); @@ -537,4 +605,247 @@ mod tests { "in-memory upload should be attempted 1 + UPLOAD_RETRY_COUNT times" ); } + + struct RecordedRequest { + method: String, + path: String, + headers: BTreeMap, + body: Vec, + } + + /// Serves the next `max_conns` connections, answering each request with the + /// response `respond` builds for it, and returns every request it received. + fn spawn_recording_mock( + max_conns: usize, + respond: impl Fn(&RecordedRequest) -> String + Send + 'static, + ) -> (String, std::thread::JoinHandle>) { + use std::io::{BufRead, BufReader, Read, Write}; + use std::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + + let handle = std::thread::spawn(move || { + let mut requests = Vec::new(); + for stream in listener.incoming().take(max_conns) { + let mut stream = stream.unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + + let mut request_line = String::new(); + reader.read_line(&mut request_line).unwrap(); + let mut parts = request_line.split_whitespace(); + let method = parts.next().unwrap().to_string(); + let path = parts.next().unwrap().to_string(); + + let mut headers = BTreeMap::new(); + loop { + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + let line = line.trim_end(); + if line.is_empty() { + break; + } + let (name, value) = line.split_once(':').unwrap(); + headers.insert(name.to_lowercase(), value.trim().to_string()); + } + let content_length: usize = headers + .get("content-length") + .map_or(0, |value| value.parse().unwrap()); + let mut body = vec![0u8; content_length]; + reader.read_exact(&mut body).unwrap(); + + let request = RecordedRequest { + method, + path, + headers, + body, + }; + stream.write_all(respond(&request).as_bytes()).unwrap(); + requests.push(request); + } + requests + }); + + (base_url, handle) + } + + fn ok_response(extra_headers: &str, body: &str) -> String { + format!( + "HTTP/1.1 200 OK\r\n{extra_headers}Content-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + } + + fn multipart_archive(content: &[u8], part_size: u64, in_memory: bool) -> ProfileArchive { + let encode = |data: &[u8]| { + use base64::Engine; + base64::engine::general_purpose::STANDARD.encode(md5::compute(data).0) + }; + let archive_content = if in_memory { + ProfileArchiveContent::CompressedInMemory { + data: content.to_vec().into(), + } + } else { + let path = tempfile::NamedTempFile::new() + .unwrap() + .into_temp_path() + .keep() + .unwrap(); + std::fs::write(&path, content).unwrap(); + ProfileArchiveContent::UncompressedOnDisk { path } + }; + ProfileArchive { + md5: encode(content), + content: archive_content, + multipart: ProfileMultipart { + size: content.len() as u64, + part_size, + part_md5s: content.chunks(part_size as usize).map(encode).collect(), + }, + } + } + + fn multipart_upload_data_for(base_url: &str, part_count: usize) -> UploadData { + UploadData { + status: "success".to_string(), + multipart_upload_urls: MultipartUploadUrls { + part_urls: (1..=part_count) + .map(|part| format!("{base_url}/part/{part}")) + .collect(), + complete_url: format!("{base_url}/complete"), + }, + run_id: "test-run".to_string(), + } + } + + #[tokio::test] + async fn multipart_upload_sends_each_part_then_completes() { + assert_multipart_upload_sends_each_part_then_completes(false).await; + } + + #[tokio::test] + async fn multipart_upload_of_in_memory_archive_sends_each_part_then_completes() { + assert_multipart_upload_sends_each_part_then_completes(true).await; + } + + async fn assert_multipart_upload_sends_each_part_then_completes(in_memory: bool) { + let content = b"0123456789"; + let archive = multipart_archive(content, 4, in_memory); + let part_md5s = archive.multipart.part_md5s.clone(); + + let (base_url, server) = spawn_recording_mock(4, |request| { + if request.method == "PUT" { + let part = request.path.trim_start_matches("/part/"); + ok_response(&format!("ETag: \"etag-{part}\"\r\n"), "") + } else { + ok_response("", "") + } + }); + + upload_profile_archive(&multipart_upload_data_for(&base_url, 3), archive) + .await + .unwrap(); + let requests = server.join().unwrap(); + + let expected_parts: [&[u8]; 3] = [b"0123", b"4567", b"89"]; + // Parts are uploaded concurrently, so they can reach the server in any order + for (index, expected_body) in expected_parts.iter().enumerate() { + let path = format!("/part/{}", index + 1); + let request = requests[..3] + .iter() + .find(|request| request.path == path) + .unwrap(); + assert_eq!(request.method, "PUT"); + assert_eq!(request.body, *expected_body); + assert_eq!(request.headers["content-md5"], part_md5s[index]); + } + + let complete = &requests[3]; + assert_eq!(complete.method, "POST"); + assert_eq!(complete.path, "/complete"); + assert_eq!( + String::from_utf8(complete.body.clone()).unwrap(), + "\ + 1\"etag-1\"\ + 2\"etag-2\"\ + 3\"etag-3\"\ + " + ); + } + + #[tokio::test] + async fn multipart_upload_retries_completion_error_body() { + let archive = multipart_archive(b"0123", 4, false); + + let (base_url, server) = spawn_recording_mock(1 + EXPECTED_ATTEMPTS, |request| { + if request.method == "PUT" { + ok_response("ETag: \"etag-1\"\r\n", "") + } else { + ok_response("", "InternalError") + } + }); + + let result = + upload_profile_archive(&multipart_upload_data_for(&base_url, 1), archive).await; + let requests = server.join().unwrap(); + + let error = result.unwrap_err().to_string(); + assert!( + error.contains("Failed to complete the performance report upload"), + "unexpected error: {error}" + ); + let completion_attempts = requests + .iter() + .filter(|request| request.path == "/complete") + .count(); + assert_eq!(completion_attempts, EXPECTED_ATTEMPTS); + } + + #[tokio::test] + async fn multipart_upload_does_not_retry_permanent_completion_error() { + let archive = multipart_archive(b"0123", 4, false); + + let (base_url, server) = spawn_recording_mock(2, |request| { + if request.method == "PUT" { + ok_response("ETag: \"etag-1\"\r\n", "") + } else { + ok_response("", "InvalidPart") + } + }); + + let result = + upload_profile_archive(&multipart_upload_data_for(&base_url, 1), archive).await; + let requests = server.join().unwrap(); + + let error = result.unwrap_err().to_string(); + assert!(error.contains("InvalidPart"), "unexpected error: {error}"); + let completion_attempts = requests + .iter() + .filter(|request| request.path == "/complete") + .count(); + assert_eq!(completion_attempts, 1); + } + + #[tokio::test] + async fn multipart_upload_completes_after_a_retried_error_body() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let archive = multipart_archive(b"0123", 4, false); + + let completion_attempts = AtomicUsize::new(0); + let (base_url, server) = spawn_recording_mock(3, move |request| { + if request.method == "PUT" { + ok_response("ETag: \"etag-1\"\r\n", "") + } else if completion_attempts.fetch_add(1, Ordering::SeqCst) == 0 { + ok_response("", "InternalError") + } else { + ok_response("", "") + } + }); + + upload_profile_archive(&multipart_upload_data_for(&base_url, 1), archive) + .await + .unwrap(); + server.join().unwrap(); + } } From 6363b74cd1533e6d08ae795f461e4817a342ae19 Mon Sep 17 00:00:00 2001 From: Luc Varoqui Date: Thu, 8 Oct 2026 14:22:36 +0200 Subject: [PATCH 2/3] feat(upload): check multipart uploads against CRC64NVME checksums Replace the md5 of each part and of the archive with CRC64NVME checksums, and describe the profile archive with a single `profileArchive` field in the upload metadata (version 12): its encoding, size and CRC64NVME, and the size and CRC64NVME of its parts. It replaces `profileEncoding`, `profileMd5` and `profileMultipart`. Each part is hashed once, and the part CRCs are combined into the archive's. S3 can check a CRC64NVME on the whole archive of a multipart upload, which md5 does not support, so a completion that leaves out a part or assembles another archive is rejected, and not only a corrupted part. It is also much cheaper to compute: about 0.5s for a 6 GiB archive on a GitHub-hosted runner, against 24s for the part and archive md5s. The upload endpoint now answers with `multipartUpload`, replacing `multipartUploadUrls`: a presigned request per part and one to complete the upload, each with the headers to send as is, as they are part of the signature. The runner sends them without knowing which ones S3 checks, so the endpoint can change them without a runner release. `crc-fast` is pinned to 1.9, the last release supporting Rust 1.88, which requires `crc` 3.3. Refs COD-3700 Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 30 ++-- Cargo.toml | 2 +- src/run_environment/provider.rs | 4 +- src/upload/interfaces.rs | 69 +++++--- src/upload/profile_archive.rs | 109 +++++++------ src/upload/s3.rs | 30 ++-- ...ata__tests__get_local_metadata_hash-2.snap | 10 +- ..._metadata__tests__get_metadata_hash-2.snap | 14 +- src/upload/upload_metadata.rs | 28 ++-- src/upload/uploader.rs | 148 +++++++++--------- 10 files changed, 253 insertions(+), 191 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index acae15118..1fdf3b3b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -712,6 +712,7 @@ dependencies = [ "caps", "clap", "console", + "crc-fast", "crc32fast", "debugid", "escargot", @@ -728,7 +729,6 @@ dependencies = [ "libc", "linux-perf-data 0.12.0 (git+https://github.com/mstange/linux-perf-data.git?rev=da5bce4b9fb724e84b1eea0cb6ab9c8a291bc676)", "log", - "md5", "memmap2", "memtrack", "nestify", @@ -890,9 +890,9 @@ dependencies = [ [[package]] name = "crc" -version = "3.4.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +checksum = "9710d3b3739c2e349eb44fe848ad0b7c8cb1e42bd87ee49371df2f7acaf3e675" dependencies = [ "crc-catalog", ] @@ -903,6 +903,18 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" +[[package]] +name = "crc-fast" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fd92aca2c6001b1bf5ba0ff84ee74ec8501b52bbef0cac80bf25a6c1d87a83d" +dependencies = [ + "crc", + "digest", + "rustversion", + "spin", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -2369,12 +2381,6 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4facc753ae494aeb6e3c22f839b158aebd4f9270f55cd3c79906c45476c47ab4" -[[package]] -name = "md5" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae960838283323069879657ca3de837e9f7bbb4c7bf6ea7f1b290d5e9476d2e0" - [[package]] name = "memchr" version = "2.8.0" @@ -4331,6 +4337,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spin" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" + [[package]] name = "srcsrv" version = "0.2.3" diff --git a/Cargo.toml b/Cargo.toml index 4587aa90d..ab83135d8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,7 +29,7 @@ sha256 = "1.6" tokio = { version = "1", features = ["macros", "rt"] } tokio-tar = { package = "astral-tokio-tar", version = "0.6.2" } tokio-util = "0.7.18" -md5 = "0.8" +crc-fast = { version = "1.9.0", default-features = false, features = ["std"] } base64 = "0.22.1" async-compression = { version = "0.4.42", features = ["tokio", "gzip"] } schemars = "1.2.1" diff --git a/src/run_environment/provider.rs b/src/run_environment/provider.rs index 7b887355e..7af90a483 100644 --- a/src/run_environment/provider.rs +++ b/src/run_environment/provider.rs @@ -136,9 +136,7 @@ pub trait RunEnvironmentProvider { tokenless: api_client.token().is_none(), repository_provider: self.get_repository_provider(), run_environment_metadata, - profile_md5: profile_archive.md5.clone(), - profile_encoding: profile_archive.content.encoding(), - profile_multipart: profile_archive.multipart.clone(), + profile_archive: profile_archive.metadata.clone(), commit_hash, allow_empty: config.allow_empty, runner: Runner { diff --git a/src/upload/interfaces.rs b/src/upload/interfaces.rs index 376b3242d..d5893b340 100644 --- a/src/upload/interfaces.rs +++ b/src/upload/interfaces.rs @@ -1,3 +1,5 @@ +use std::collections::BTreeMap; + use serde::{Deserialize, Serialize}; use crate::executor::ExecutorName; @@ -13,9 +15,7 @@ pub struct UploadMetadata { pub repository_provider: RepositoryProvider, pub version: Option, pub tokenless: bool, - pub profile_md5: String, - pub profile_encoding: Option, - pub profile_multipart: ProfileMultipart, + pub profile_archive: ProfileMetadata, pub runner: Runner, pub run_environment: RunEnvironment, pub run_part: Option, @@ -25,16 +25,20 @@ pub struct UploadMetadata { pub run_environment_metadata: RunEnvironmentMetadata, } -/// Layout of a profile archive uploaded as an S3 multipart upload, in consecutive -/// `part_size` chunks, the last one holding the remainder. S3 requires parts of 5 MiB -/// to 5 GiB (the last one excepted from the minimum), and at most 10,000 of them. +/// Profile archive, uploaded as an S3 multipart upload in consecutive `part_size` +/// chunks, the last one holding the remainder. S3 requires parts of 5 MiB to 5 GiB +/// (the last one excepted from the minimum), and at most 10,000 of them. #[derive(Serialize, Debug, Clone, PartialEq)] #[serde(rename_all = "camelCase")] -pub struct ProfileMultipart { +pub struct ProfileMetadata { + /// `Content-Encoding` of the archive, such as `gzip` + pub encoding: Option, pub size: u64, + /// Base64 big-endian CRC64NVME of the whole archive + pub crc64nvme: String, pub part_size: u64, - /// Base64 md5 of each part, in upload order - pub part_md5s: Vec, + /// Base64 big-endian CRC64NVME of each part, in upload order + pub part_crc64nvmes: Vec, } #[derive(Serialize, Debug)] @@ -59,17 +63,25 @@ pub struct Runner { #[serde(rename_all = "camelCase")] pub struct UploadData { pub status: String, - pub multipart_upload_urls: MultipartUploadUrls, + pub multipart_upload: MultipartUpload, pub run_id: String, } #[derive(Deserialize, Serialize, Debug)] #[serde(rename_all = "camelCase")] -pub struct MultipartUploadUrls { - /// Presigned S3 `UploadPart` URLs, one per part, in upload order - pub part_urls: Vec, - /// Presigned S3 `CompleteMultipartUpload` URL - pub complete_url: String, +pub struct MultipartUpload { + /// Presigned S3 `UploadPart` requests, one per part, in upload order + pub parts: Vec, + /// Presigned S3 `CompleteMultipartUpload` request + pub complete: PresignedRequest, +} + +/// Request presigned by the API, to send to `url` with `headers` as is: they are +/// part of the signature +#[derive(Deserialize, Serialize, Debug)] +pub struct PresignedRequest { + pub url: String, + pub headers: BTreeMap, } #[derive(Deserialize, Debug)] @@ -88,17 +100,32 @@ mod tests { r#"{ "status": "success", "runId": "run-id", - "multipartUploadUrls": { - "partUrls": ["https://part/1", "https://part/2"], - "completeUrl": "https://complete" + "multipartUpload": { + "parts": [ + { "url": "https://part/1", "headers": { "x-amz-checksum-crc64nvme": "nq48tdaL2no=" } }, + { "url": "https://part/2", "headers": { "x-amz-checksum-crc64nvme": "XMXoclwBfLo=" } } + ], + "complete": { + "url": "https://complete", + "headers": { "x-amz-checksum-type": "FULL_OBJECT" } + } } }"#, ) .unwrap(); assert_eq!(upload_data.run_id, "run-id"); - let urls = upload_data.multipart_upload_urls; - assert_eq!(urls.part_urls, ["https://part/1", "https://part/2"]); - assert_eq!(urls.complete_url, "https://complete"); + let upload = upload_data.multipart_upload; + let part_urls: Vec<_> = upload.parts.iter().map(|part| part.url.as_str()).collect(); + assert_eq!(part_urls, ["https://part/1", "https://part/2"]); + assert_eq!( + upload.parts[1].headers, + BTreeMap::from([("x-amz-checksum-crc64nvme".into(), "XMXoclwBfLo=".into())]) + ); + assert_eq!(upload.complete.url, "https://complete"); + assert_eq!( + upload.complete.headers, + BTreeMap::from([("x-amz-checksum-type".into(), "FULL_OBJECT".into())]) + ); } } diff --git a/src/upload/profile_archive.rs b/src/upload/profile_archive.rs index 8adaa8185..d8b7f2e5c 100644 --- a/src/upload/profile_archive.rs +++ b/src/upload/profile_archive.rs @@ -1,7 +1,8 @@ use base64::{Engine, engine::general_purpose}; +use crc_fast::{CrcAlgorithm, Digest}; use crate::prelude::*; -use crate::upload::interfaces::ProfileMultipart; +use crate::upload::interfaces::ProfileMetadata; use bytes::Bytes; use std::io::{Read, SeekFrom}; use std::path::PathBuf; @@ -46,9 +47,8 @@ const HASH_READ_BUFFER_SIZE: usize = 8 * 1024 * 1024; // 8 MiB #[derive(Debug)] pub struct ProfileArchive { - pub md5: String, pub content: ProfileArchiveContent, - pub multipart: ProfileMultipart, + pub metadata: ProfileMetadata, } #[derive(Debug)] @@ -57,24 +57,24 @@ pub enum ProfileArchiveContent { UncompressedOnDisk { path: PathBuf }, } -fn encode_md5(digest: md5::Digest) -> String { - general_purpose::STANDARD.encode(digest.0) +/// Base64 of the big-endian CRC64NVME, the encoding S3 expects +fn encode_crc64nvme(crc: u64) -> String { + general_purpose::STANDARD.encode(crc.to_be_bytes()) } -/// Read the content once to compute the md5 of all of it and of each consecutive -/// `part_size` chunk of it. +/// Read the content once to compute the CRC64NVME of each consecutive `part_size` +/// chunk of it, and of all of it by combining the part CRCs. /// /// CPU heavy + potentially reading from blocking IO, so it is recommended to /// run this on a blocking thread pool. -fn compute_md5s_from_reader( +fn compute_crc64nvmes_from_reader( mut reader: impl Read, part_size: u64, ) -> Result<(String, Vec)> { let mut buffer = vec![0u8; HASH_READ_BUFFER_SIZE]; - let mut whole_context = md5::Context::new(); - let mut part_context = md5::Context::new(); - let mut part_length = 0u64; - let mut part_md5s = Vec::new(); + let mut whole_digest = Digest::new(CrcAlgorithm::Crc64Nvme); + let mut part_digest = Digest::new(CrcAlgorithm::Crc64Nvme); + let mut part_crc64nvmes = Vec::new(); loop { let read = reader.read(&mut buffer)?; @@ -82,43 +82,42 @@ fn compute_md5s_from_reader( break; } let mut chunk = &buffer[..read]; - whole_context.consume(chunk); - while !chunk.is_empty() { - let taken = (part_size - part_length).min(chunk.len() as u64) as usize; - part_context.consume(&chunk[..taken]); - part_length += taken as u64; + let taken = (part_size - part_digest.get_amount()).min(chunk.len() as u64) as usize; + part_digest.update(&chunk[..taken]); chunk = &chunk[taken..]; - if part_length == part_size { - let finished_part = std::mem::replace(&mut part_context, md5::Context::new()); - part_md5s.push(encode_md5(finished_part.finalize())); - part_length = 0; + if part_digest.get_amount() == part_size { + whole_digest.combine(&part_digest); + part_crc64nvmes.push(encode_crc64nvme(part_digest.finalize_reset())); } } } - if part_length > 0 { - part_md5s.push(encode_md5(part_context.finalize())); + if part_digest.get_amount() > 0 { + whole_digest.combine(&part_digest); + part_crc64nvmes.push(encode_crc64nvme(part_digest.finalize())); } - Ok((encode_md5(whole_context.finalize()), part_md5s)) + Ok((encode_crc64nvme(whole_digest.finalize()), part_crc64nvmes)) } -/// [`compute_md5s_from_reader`] over the content, run on the blocking thread pool as -/// hashing the whole archive would otherwise stall the async runtime. -async fn compute_md5s( +/// [`compute_crc64nvmes_from_reader`] over the content, run on the blocking thread pool +/// as hashing the whole archive would otherwise stall the async runtime. +async fn compute_crc64nvmes( content: &ProfileArchiveContent, part_size: u64, ) -> Result<(String, Vec)> { match content { ProfileArchiveContent::CompressedInMemory { data } => { let data = data.clone(); - tokio::task::spawn_blocking(move || compute_md5s_from_reader(&data[..], part_size)) - .await? + tokio::task::spawn_blocking(move || { + compute_crc64nvmes_from_reader(&data[..], part_size) + }) + .await? } ProfileArchiveContent::UncompressedOnDisk { path } => { let path = path.clone(); tokio::task::spawn_blocking(move || { - compute_md5s_from_reader(std::fs::File::open(path)?, part_size) + compute_crc64nvmes_from_reader(std::fs::File::open(path)?, part_size) }) .await? } @@ -151,18 +150,16 @@ impl ProfileArchive { let size = content.size().await?; let part_size = choose_multipart_part_size(size, concurrent_part_uploads()); - let (md5, part_md5s) = compute_md5s(&content, part_size).await?; - let multipart = ProfileMultipart { + let (crc64nvme, part_crc64nvmes) = compute_crc64nvmes(&content, part_size).await?; + let metadata = ProfileMetadata { + encoding: content.encoding(), size, + crc64nvme, part_size, - part_md5s, + part_crc64nvmes, }; - Ok(ProfileArchive { - md5, - content, - multipart, - }) + Ok(ProfileArchive { content, metadata }) } } @@ -228,22 +225,31 @@ mod tests { path } + fn crc64nvme(data: &[u8]) -> String { + encode_crc64nvme(crc_fast::checksum(CrcAlgorithm::Crc64Nvme, data)) + } + #[test] - fn computes_whole_and_part_md5s_in_one_pass() { + fn computes_whole_and_part_crc64nvmes_in_one_pass() { // Not a multiple of the part size, and spanning several read buffers let content: Vec = (0..HASH_READ_BUFFER_SIZE * 2 + 123) .map(|i| (i % 251) as u8) .collect(); let part_size = (HASH_READ_BUFFER_SIZE / 3) as u64; - let (hash, part_md5s) = compute_md5s_from_reader(&content[..], part_size).unwrap(); + let (whole, part_crc64nvmes) = + compute_crc64nvmes_from_reader(&content[..], part_size).unwrap(); - assert_eq!(hash, encode_md5(md5::compute(&content))); - let expected_part_md5s: Vec = content - .chunks(part_size as usize) - .map(|part| encode_md5(md5::compute(part))) - .collect(); - assert_eq!(part_md5s, expected_part_md5s); + assert_eq!(whole, crc64nvme(&content)); + let expected_part_crc64nvmes: Vec = + content.chunks(part_size as usize).map(crc64nvme).collect(); + assert_eq!(part_crc64nvmes, expected_part_crc64nvmes); + } + + #[test] + fn crc64nvme_is_base64_of_the_big_endian_crc() { + // Check value of the CRC-64/NVME catalogue entry: 0xAE8B14860A799888 + assert_eq!(crc64nvme(b"123456789"), "rosUhgp5mIg="); } #[tokio::test] @@ -254,14 +260,15 @@ mod tests { .await .unwrap(); - let md5 = encode_md5(md5::compute(b"profile-archive")); - assert_eq!(archive.md5, md5); + let crc = crc64nvme(b"profile-archive"); assert_eq!( - archive.multipart, - ProfileMultipart { + archive.metadata, + ProfileMetadata { + encoding: None, size: b"profile-archive".len() as u64, + crc64nvme: crc.clone(), part_size: MULTIPART_MIN_PART_SIZE_BYTES, - part_md5s: vec![md5], + part_crc64nvmes: vec![crc], } ); } diff --git a/src/upload/s3.rs b/src/upload/s3.rs index 2c621a5aa..5f38a3895 100644 --- a/src/upload/s3.rs +++ b/src/upload/s3.rs @@ -6,20 +6,33 @@ //! [CompleteMultipartUpload](https://docs.aws.amazon.com/AmazonS3/latest/API/API_CompleteMultipartUpload.html) //! request. This module handles the S3-specific parts of these requests. //! -//! The part URLs are signed with the `Content-MD5` and `Content-Length` headers, so S3 -//! rejects a part that does not send them with the signed values. The content type and -//! encoding of the archive are set when the multipart upload is created. +//! Each request is presigned with headers, such as the checksums S3 checks the upload +//! against, which have to be sent with their signed values or S3 rejects the request. +//! The content type and encoding of the archive are set when the multipart upload is +//! created. use crate::prelude::*; +use crate::upload::interfaces::PresignedRequest; use console::style; +/// Request to the presigned URL, with its signed headers +fn presigned( + client: &reqwest::Client, + method: reqwest::Method, + request: &PresignedRequest, +) -> reqwest::RequestBuilder { + request.headers.iter().fold( + client.request(method, &request.url), + |builder, (name, value)| builder.header(name, value), + ) +} + /// `UploadPart` request uploading one part of a multipart upload. pub(super) fn upload_part( client: &reqwest::Client, - url: &str, - content_md5: &str, + request: &PresignedRequest, ) -> reqwest::RequestBuilder { - client.put(url).header("Content-MD5", content_md5) + presigned(client, reqwest::Method::PUT, request) } /// ETag S3 assigned to an uploaded part, only obtainable from a part upload response @@ -42,11 +55,10 @@ pub(super) fn part_etag(response: &reqwest::Response) -> Result { /// object. `etags` are in part order. pub(super) fn complete_upload( client: &reqwest::Client, - url: &str, + request: &PresignedRequest, etags: &[PartETag], ) -> reqwest::RequestBuilder { - client - .post(url) + presigned(client, reqwest::Method::POST, request) .header("Content-Type", "application/xml") .body(build_complete_body(etags)) } diff --git a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap index 832d7553c..89db98cbf 100644 --- a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap +++ b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap @@ -6,13 +6,13 @@ expression: upload_metadata "repositoryProvider": "PROJECT", "version": 12, "tokenless": false, - "profileMd5": "tfC4VxYiYdJcTWpHpv4Ouw==", - "profileEncoding": "gzip", - "profileMultipart": { + "profileArchive": { + "encoding": "gzip", "size": 20480, + "crc64nvme": "XMXoclwBfLo=", "partSize": 16777216, - "partMd5s": [ - "tfC4VxYiYdJcTWpHpv4Ouw==" + "partCrc64nvmes": [ + "XMXoclwBfLo=" ] }, "runner": { diff --git a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap index ea985f875..f218237d9 100644 --- a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap +++ b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap @@ -6,15 +6,15 @@ expression: upload_metadata "repositoryProvider": "GITHUB", "version": 12, "tokenless": true, - "profileMd5": "jp/k05RKuqP3ERQuIIvx4Q==", - "profileEncoding": "gzip", - "profileMultipart": { + "profileArchive": { + "encoding": "gzip", "size": 629145600, + "crc64nvme": "SVWSLdzczJc=", "partSize": 268435456, - "partMd5s": [ - "eEKfdGLWNqhNnJIvSVWZxQ==", - "vKfHJAI2Hyo68jWwUc6H9A==", - "bQY//BUrm3jASNTyWj/3Aw==" + "partCrc64nvmes": [ + "w7Dr86r5+2I=", + "puujmsfPTns=", + "Tt/QAZdhKk4=" ] }, "runner": { diff --git a/src/upload/upload_metadata.rs b/src/upload/upload_metadata.rs index fc95c4c2d..272786f0b 100644 --- a/src/upload/upload_metadata.rs +++ b/src/upload/upload_metadata.rs @@ -22,7 +22,7 @@ mod tests { Sender, }; use crate::system::SystemInfo; - use crate::upload::{LATEST_UPLOAD_METADATA_VERSION, ProfileMultipart, Runner, UploadMetadata}; + use crate::upload::{LATEST_UPLOAD_METADATA_VERSION, ProfileMetadata, Runner, UploadMetadata}; #[test] fn test_get_metadata_hash() { @@ -30,15 +30,15 @@ mod tests { repository_provider: RepositoryProvider::GitHub, version: Some(LATEST_UPLOAD_METADATA_VERSION), tokenless: true, - profile_md5: "jp/k05RKuqP3ERQuIIvx4Q==".into(), - profile_encoding: Some("gzip".into()), - profile_multipart: ProfileMultipart { + profile_archive: ProfileMetadata { + encoding: Some("gzip".into()), size: 629145600, + crc64nvme: "SVWSLdzczJc=".into(), part_size: 268435456, - part_md5s: vec![ - "eEKfdGLWNqhNnJIvSVWZxQ==".into(), - "vKfHJAI2Hyo68jWwUc6H9A==".into(), - "bQY//BUrm3jASNTyWj/3Aw==".into(), + part_crc64nvmes: vec![ + "w7Dr86r5+2I=".into(), + "puujmsfPTns=".into(), + "Tt/QAZdhKk4=".into(), ], }, runner: Runner { @@ -82,7 +82,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"307ed5a2218e04b1befe21d355e19a2a3adb7a617b19d1b6e8406c23355a2664" + @"07b92482167eaabedd8111906c7d96763d0da5ea951a5be04492413508206261" ); assert_json_snapshot!(upload_metadata); } @@ -93,12 +93,12 @@ mod tests { repository_provider: RepositoryProvider::Project, version: Some(LATEST_UPLOAD_METADATA_VERSION), tokenless: false, - profile_md5: "tfC4VxYiYdJcTWpHpv4Ouw==".into(), - profile_encoding: Some("gzip".into()), - profile_multipart: ProfileMultipart { + profile_archive: ProfileMetadata { + encoding: Some("gzip".into()), size: 20480, + crc64nvme: "XMXoclwBfLo=".into(), part_size: 16777216, - part_md5s: vec!["tfC4VxYiYdJcTWpHpv4Ouw==".into()], + part_crc64nvmes: vec!["XMXoclwBfLo=".into()], }, runner: Runner { name: "codspeed-runner".into(), @@ -158,7 +158,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"dcbb07dc8bf0c97cc88566ef78acec873f452420ce004807f3cdd6218544a458" + @"5fd485b7ad81097b43268e651dfeec697feb7985f3a8e6e1e5eea39d213262cf" ); assert_json_snapshot!(upload_metadata); } diff --git a/src/upload/uploader.rs b/src/upload/uploader.rs index b018b4e1e..fcf72b712 100644 --- a/src/upload/uploader.rs +++ b/src/upload/uploader.rs @@ -22,7 +22,7 @@ use tokio::fs::File; use tokio::io::AsyncWriteExt; use tokio_tar::Builder; -use super::interfaces::{MultipartUploadUrls, ProfileMultipart, UploadData, UploadMetadata}; +use super::interfaces::{MultipartUpload, ProfileMetadata, UploadData, UploadMetadata}; use super::profile_archive::{ProfileArchive, concurrent_part_uploads}; use super::s3; @@ -47,7 +47,7 @@ fn human_bytes_per_second(bytes: u64, elapsed: std::time::Duration) -> String { format!("{}/s", human_bytes(bytes_per_second as u64)) } -/// Create a profile archive from the profile folder and return its md5 hash encoded in base64 +/// Create a profile archive from the profile folder, along with its description /// /// For Valgrind, we create a gzip-compressed tar archive of the entire profile folder. /// For WallTime and Memory, we create an uncompressed tar archive on disk: their @@ -265,23 +265,21 @@ async fn send_upload_request( /// Upload the parts concurrently, then assemble them with the S3 complete request. /// Parts of a failed upload are cleaned up by the bucket lifecycle rules. async fn upload_multipart_profile_archive( - multipart_upload_urls: &MultipartUploadUrls, - multipart: &ProfileMultipart, + multipart_upload: &MultipartUpload, + profile: &ProfileMetadata, content: &ProfileArchiveContent, ) -> Result<()> { debug!("Starting multipart upload for profile archive"); + let part_count = profile.part_crc64nvmes.len(); debug!( "Multipart upload details: part_count={}, part_size={}, total_size={}", - multipart.part_md5s.len(), - multipart.part_size, - multipart.size + part_count, profile.part_size, profile.size ); - let part_count = multipart.part_md5s.len(); - if multipart_upload_urls.part_urls.len() != part_count { + if multipart_upload.parts.len() != part_count { bail!( - "Received {} part upload URLs for {} parts", - multipart_upload_urls.part_urls.len(), + "Received {} part upload requests for {} parts", + multipart_upload.parts.len(), part_count ); } @@ -290,46 +288,41 @@ async fn upload_multipart_profile_archive( let upload_start = Instant::now(); // Unordered, so that a part finishing frees its slot even while an earlier part is still // uploading - let mut indexed_etags: Vec<_> = futures::stream::iter( - multipart_upload_urls - .part_urls - .iter() - .zip(&multipart.part_md5s) - .enumerate(), - ) - .map(|(index, (part_url, part_md5))| async move { - let offset = index as u64 * multipart.part_size; - let range = ContentRange { - content, - offset, - length: multipart.part_size.min(multipart.size - offset), - }; - debug!( - "Uploading part {}/{} ({} bytes)", - index + 1, - part_count, - range.length - ); - let part_start = Instant::now(); - let etag = with_upload_retry(|| async { - let request = s3::upload_part(&STREAMING_CLIENT, part_url, part_md5); - let response = send_upload_request(range.attach(request).await?).await?; - Ok(s3::part_etag(&response)?) + let part_requests = multipart_upload.parts.iter().enumerate(); + let mut indexed_etags: Vec<_> = futures::stream::iter(part_requests) + .map(|(index, part_request)| async move { + let offset = index as u64 * profile.part_size; + let range = ContentRange { + content, + offset, + length: profile.part_size.min(profile.size - offset), + }; + debug!( + "Uploading part {}/{} ({} bytes)", + index + 1, + part_count, + range.length + ); + let part_start = Instant::now(); + let etag = with_upload_retry(|| async { + let request = s3::upload_part(&STREAMING_CLIENT, part_request); + let response = send_upload_request(range.attach(request).await?).await?; + Ok(s3::part_etag(&response)?) + }) + .await?; + let part_elapsed = part_start.elapsed(); + debug!( + "Uploaded part {}/{} in {:.1?} ({})", + index + 1, + part_count, + part_elapsed, + human_bytes_per_second(range.length, part_elapsed) + ); + Ok::<_, anyhow::Error>((index, etag)) }) + .buffer_unordered(concurrency) + .try_collect() .await?; - let part_elapsed = part_start.elapsed(); - debug!( - "Uploaded part {}/{} in {:.1?} ({})", - index + 1, - part_count, - part_elapsed, - human_bytes_per_second(range.length, part_elapsed) - ); - Ok::<_, anyhow::Error>((index, etag)) - }) - .buffer_unordered(concurrency) - .try_collect() - .await?; indexed_etags.sort_unstable_by_key(|(index, _)| *index); let etags: Vec<_> = indexed_etags.into_iter().map(|(_, etag)| etag).collect(); @@ -339,18 +332,14 @@ async fn upload_multipart_profile_archive( "Uploaded {} part{} ({}) in {:.1?} with {} concurrent uploads ({})", part_count, if part_count > 1 { "s" } else { "" }, - human_bytes(multipart.size), + human_bytes(profile.size), upload_elapsed, concurrency, - human_bytes_per_second(multipart.size, upload_elapsed) + human_bytes_per_second(profile.size, upload_elapsed) ); with_upload_retry(|| async { - let request = s3::complete_upload( - &STREAMING_CLIENT, - &multipart_upload_urls.complete_url, - &etags, - ); + let request = s3::complete_upload(&STREAMING_CLIENT, &multipart_upload.complete, &etags); let response = send_upload_request(request).await?; match s3::check_complete_upload_response(response).await { Ok(Ok(())) => Ok(()), @@ -368,8 +357,8 @@ async fn upload_profile_archive( profile_archive: ProfileArchive, ) -> Result<()> { upload_multipart_profile_archive( - &upload_data.multipart_upload_urls, - &profile_archive.multipart, + &upload_data.multipart_upload, + &profile_archive.metadata, &profile_archive.content, ) .await @@ -438,6 +427,7 @@ mod tests { use url::Url; use super::*; + use crate::upload::interfaces::PresignedRequest; use std::path::PathBuf; // TODO: remove the ignore when implementing network mocking @@ -676,11 +666,10 @@ mod tests { ) } + /// Archive split in `part_size` parts. The checksums are placeholders: the uploader + /// sends the headers the API presigned, not checksums of its own. fn multipart_archive(content: &[u8], part_size: u64, in_memory: bool) -> ProfileArchive { - let encode = |data: &[u8]| { - use base64::Engine; - base64::engine::general_purpose::STANDARD.encode(md5::compute(data).0) - }; + let part_count = content.len().div_ceil(part_size as usize); let archive_content = if in_memory { ProfileArchiveContent::CompressedInMemory { data: content.to_vec().into(), @@ -694,25 +683,39 @@ mod tests { std::fs::write(&path, content).unwrap(); ProfileArchiveContent::UncompressedOnDisk { path } }; + let encoding = archive_content.encoding(); ProfileArchive { - md5: encode(content), content: archive_content, - multipart: ProfileMultipart { + metadata: ProfileMetadata { + encoding, size: content.len() as u64, + crc64nvme: "crc".to_string(), part_size, - part_md5s: content.chunks(part_size as usize).map(encode).collect(), + part_crc64nvmes: (1..=part_count).map(|part| format!("crc-{part}")).collect(), }, } } fn multipart_upload_data_for(base_url: &str, part_count: usize) -> UploadData { + let presigned = |path: String, header: (&str, String)| PresignedRequest { + url: format!("{base_url}{path}"), + headers: BTreeMap::from([(header.0.to_string(), header.1)]), + }; UploadData { status: "success".to_string(), - multipart_upload_urls: MultipartUploadUrls { - part_urls: (1..=part_count) - .map(|part| format!("{base_url}/part/{part}")) + multipart_upload: MultipartUpload { + parts: (1..=part_count) + .map(|part| { + presigned( + format!("/part/{part}"), + ("x-amz-checksum-crc64nvme", format!("signed-crc-{part}")), + ) + }) .collect(), - complete_url: format!("{base_url}/complete"), + complete: presigned( + "/complete".to_string(), + ("x-amz-mp-object-size", "signed-size".to_string()), + ), }, run_id: "test-run".to_string(), } @@ -731,7 +734,6 @@ mod tests { async fn assert_multipart_upload_sends_each_part_then_completes(in_memory: bool) { let content = b"0123456789"; let archive = multipart_archive(content, 4, in_memory); - let part_md5s = archive.multipart.part_md5s.clone(); let (base_url, server) = spawn_recording_mock(4, |request| { if request.method == "PUT" { @@ -757,12 +759,16 @@ mod tests { .unwrap(); assert_eq!(request.method, "PUT"); assert_eq!(request.body, *expected_body); - assert_eq!(request.headers["content-md5"], part_md5s[index]); + assert_eq!( + request.headers["x-amz-checksum-crc64nvme"], + format!("signed-crc-{}", index + 1) + ); } let complete = &requests[3]; assert_eq!(complete.method, "POST"); assert_eq!(complete.path, "/complete"); + assert_eq!(complete.headers["x-amz-mp-object-size"], "signed-size"); assert_eq!( String::from_utf8(complete.body.clone()).unwrap(), "\ From acfbe6a4670dea3cd86aa0ac2f471c3acd520ff1 Mon Sep 17 00:00:00 2001 From: Luc Varoqui Date: Fri, 9 Oct 2026 11:17:30 +0200 Subject: [PATCH 3/3] review changes (will squash merge) --- src/run_environment/provider.rs | 2 +- src/upload/interfaces.rs | 8 +- src/upload/profile_archive.rs | 47 +++++----- src/upload/s3.rs | 27 ------ ...ata__tests__get_local_metadata_hash-2.snap | 2 +- ..._metadata__tests__get_metadata_hash-2.snap | 2 +- src/upload/upload_metadata.rs | 12 +-- src/upload/uploader.rs | 89 ++++++------------- 8 files changed, 69 insertions(+), 120 deletions(-) diff --git a/src/run_environment/provider.rs b/src/run_environment/provider.rs index 7af90a483..31b3328a9 100644 --- a/src/run_environment/provider.rs +++ b/src/run_environment/provider.rs @@ -136,7 +136,7 @@ pub trait RunEnvironmentProvider { tokenless: api_client.token().is_none(), repository_provider: self.get_repository_provider(), run_environment_metadata, - profile_archive: profile_archive.metadata.clone(), + profile_archive_metadata: profile_archive.metadata.clone(), commit_hash, allow_empty: config.allow_empty, runner: Runner { diff --git a/src/upload/interfaces.rs b/src/upload/interfaces.rs index d5893b340..c03f72620 100644 --- a/src/upload/interfaces.rs +++ b/src/upload/interfaces.rs @@ -15,7 +15,7 @@ pub struct UploadMetadata { pub repository_provider: RepositoryProvider, pub version: Option, pub tokenless: bool, - pub profile_archive: ProfileMetadata, + pub profile_archive_metadata: ProfileArchiveMetadata, pub runner: Runner, pub run_environment: RunEnvironment, pub run_part: Option, @@ -25,12 +25,12 @@ pub struct UploadMetadata { pub run_environment_metadata: RunEnvironmentMetadata, } -/// Profile archive, uploaded as an S3 multipart upload in consecutive `part_size` -/// chunks, the last one holding the remainder. S3 requires parts of 5 MiB to 5 GiB +/// Metadata of the profile archive, uploaded as an S3 multipart upload in consecutive +/// `part_size` chunks, the last one holding the remainder. S3 requires parts of 5 MiB to 5 GiB /// (the last one excepted from the minimum), and at most 10,000 of them. #[derive(Serialize, Debug, Clone, PartialEq)] #[serde(rename_all = "camelCase")] -pub struct ProfileMetadata { +pub struct ProfileArchiveMetadata { /// `Content-Encoding` of the archive, such as `gzip` pub encoding: Option, pub size: u64, diff --git a/src/upload/profile_archive.rs b/src/upload/profile_archive.rs index d8b7f2e5c..9308810f7 100644 --- a/src/upload/profile_archive.rs +++ b/src/upload/profile_archive.rs @@ -2,7 +2,7 @@ use base64::{Engine, engine::general_purpose}; use crc_fast::{CrcAlgorithm, Digest}; use crate::prelude::*; -use crate::upload::interfaces::ProfileMetadata; +use crate::upload::interfaces::ProfileArchiveMetadata; use bytes::Bytes; use std::io::{Read, SeekFrom}; use std::path::PathBuf; @@ -18,7 +18,7 @@ const DEFAULT_CONCURRENT_PART_UPLOADS: usize = 8; /// Overrides [`DEFAULT_CONCURRENT_PART_UPLOADS`] const CONCURRENT_PART_UPLOADS_ENV: &str = "CODSPEED_UPLOAD_CONCURRENCY"; -static CONCURRENT_PART_UPLOADS: LazyLock = LazyLock::new(|| { +pub static CONCURRENT_PART_UPLOADS: LazyLock = LazyLock::new(|| { let Ok(value) = std::env::var(CONCURRENT_PART_UPLOADS_ENV) else { return DEFAULT_CONCURRENT_PART_UPLOADS; }; @@ -33,22 +33,21 @@ static CONCURRENT_PART_UPLOADS: LazyLock = LazyLock::new(|| { } }); -pub(super) fn concurrent_part_uploads() -> usize { - *CONCURRENT_PART_UPLOADS -} - const MULTIPART_MIN_PART_SIZE_BYTES: u64 = 16 * 1024 * 1024; // 16 MiB /// Bounds how much a failed part has to re-send const MULTIPART_MAX_PART_SIZE_BYTES: u64 = 256 * 1024 * 1024; // 256 MiB -/// More parts than concurrent uploads, so that the upload slots freed by fast parts -/// pick up remaining work instead of idling while the slowest part finishes -const PARTS_PER_CONCURRENT_UPLOAD: u64 = 2; +/// Target number of parts per concurrent upload. Queuing several parts per upload slot +/// lets a slot freed by a fast part pick up a remaining one instead of idling while the +/// slowest part finishes. +/// +/// See [`choose_multipart_part_size`] for how the target is used to determine the actual part size. +const TARGET_PARTS_PER_CONCURRENT_UPLOAD: u64 = 2; const HASH_READ_BUFFER_SIZE: usize = 8 * 1024 * 1024; // 8 MiB #[derive(Debug)] pub struct ProfileArchive { pub content: ProfileArchiveContent, - pub metadata: ProfileMetadata, + pub metadata: ProfileArchiveMetadata, } #[derive(Debug)] @@ -100,8 +99,7 @@ fn compute_crc64nvmes_from_reader( Ok((encode_crc64nvme(whole_digest.finalize()), part_crc64nvmes)) } -/// [`compute_crc64nvmes_from_reader`] over the content, run on the blocking thread pool -/// as hashing the whole archive would otherwise stall the async runtime. +/// [`compute_crc64nvmes_from_reader`] over the content, on the blocking thread pool async fn compute_crc64nvmes( content: &ProfileArchiveContent, part_size: u64, @@ -124,11 +122,13 @@ async fn compute_crc64nvmes( } } -/// Split the archive in about `PARTS_PER_CONCURRENT_UPLOAD` parts per concurrent upload, -/// within the part size bounds. An archive no larger than the minimum part size is a -/// single part. +/// Size parts to split the archive in [`TARGET_PARTS_PER_CONCURRENT_UPLOAD`] parts per +/// concurrent upload, clamped to the part size bounds. The part count therefore drifts +/// from the target at both ends: a larger archive is split in more parts of +/// [`MULTIPART_MAX_PART_SIZE_BYTES`], and a smaller one in fewer parts of +/// [`MULTIPART_MIN_PART_SIZE_BYTES`], down to a single part. fn choose_multipart_part_size(size: u64, concurrent_uploads: usize) -> u64 { - let target_part_count = concurrent_uploads as u64 * PARTS_PER_CONCURRENT_UPLOAD; + let target_part_count = concurrent_uploads as u64 * TARGET_PARTS_PER_CONCURRENT_UPLOAD; size.div_ceil(target_part_count) .clamp(MULTIPART_MIN_PART_SIZE_BYTES, MULTIPART_MAX_PART_SIZE_BYTES) } @@ -148,10 +148,10 @@ impl ProfileArchive { async fn new(content: ProfileArchiveContent) -> Result { let size = content.size().await?; - let part_size = choose_multipart_part_size(size, concurrent_part_uploads()); + let part_size = choose_multipart_part_size(size, *CONCURRENT_PART_UPLOADS); let (crc64nvme, part_crc64nvmes) = compute_crc64nvmes(&content, part_size).await?; - let metadata = ProfileMetadata { + let metadata = ProfileArchiveMetadata { encoding: content.encoding(), size, crc64nvme, @@ -263,7 +263,7 @@ mod tests { let crc = crc64nvme(b"profile-archive"); assert_eq!( archive.metadata, - ProfileMetadata { + ProfileArchiveMetadata { encoding: None, size: b"profile-archive".len() as u64, crc64nvme: crc.clone(), @@ -274,9 +274,14 @@ mod tests { } #[test] - fn archive_is_split_in_two_parts_per_concurrent_upload() { + fn part_count_hits_the_target_within_the_part_size_bounds() { const MIB: u64 = 1024 * 1024; - assert_eq!(choose_multipart_part_size(1024 * MIB, 8), 64 * MIB); + let size = 1024 * MIB; + let part_size = choose_multipart_part_size(size, 8); + assert_eq!( + size.div_ceil(part_size), + 8 * TARGET_PARTS_PER_CONCURRENT_UPLOAD + ); } #[test] diff --git a/src/upload/s3.rs b/src/upload/s3.rs index 5f38a3895..16c028bcf 100644 --- a/src/upload/s3.rs +++ b/src/upload/s3.rs @@ -135,30 +135,3 @@ fn build_complete_body(etags: &[PartETag]) -> String { .collect(); format!("{parts}") } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn success_body_is_not_an_error() { - let body = - "\"etag\""; - assert!(S3Error::from_body(body.to_owned()).is_none()); - } - - #[test] - fn internal_error_is_transient() { - let error = S3Error::from_body( - "InternalErrorPlease try again.".into(), - ) - .unwrap(); - assert!(error.is_transient()); - } - - #[test] - fn invalid_part_is_permanent() { - let error = S3Error::from_body("InvalidPart".into()).unwrap(); - assert!(!error.is_transient()); - } -} diff --git a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap index 89db98cbf..dd8ee43b0 100644 --- a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap +++ b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_local_metadata_hash-2.snap @@ -6,7 +6,7 @@ expression: upload_metadata "repositoryProvider": "PROJECT", "version": 12, "tokenless": false, - "profileArchive": { + "profileArchiveMetadata": { "encoding": "gzip", "size": 20480, "crc64nvme": "XMXoclwBfLo=", diff --git a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap index f218237d9..eb03ece8f 100644 --- a/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap +++ b/src/upload/snapshots/codspeed_runner__upload__upload_metadata__tests__get_metadata_hash-2.snap @@ -6,7 +6,7 @@ expression: upload_metadata "repositoryProvider": "GITHUB", "version": 12, "tokenless": true, - "profileArchive": { + "profileArchiveMetadata": { "encoding": "gzip", "size": 629145600, "crc64nvme": "SVWSLdzczJc=", diff --git a/src/upload/upload_metadata.rs b/src/upload/upload_metadata.rs index 272786f0b..5d3c8945b 100644 --- a/src/upload/upload_metadata.rs +++ b/src/upload/upload_metadata.rs @@ -22,7 +22,9 @@ mod tests { Sender, }; use crate::system::SystemInfo; - use crate::upload::{LATEST_UPLOAD_METADATA_VERSION, ProfileMetadata, Runner, UploadMetadata}; + use crate::upload::{ + LATEST_UPLOAD_METADATA_VERSION, ProfileArchiveMetadata, Runner, UploadMetadata, + }; #[test] fn test_get_metadata_hash() { @@ -30,7 +32,7 @@ mod tests { repository_provider: RepositoryProvider::GitHub, version: Some(LATEST_UPLOAD_METADATA_VERSION), tokenless: true, - profile_archive: ProfileMetadata { + profile_archive_metadata: ProfileArchiveMetadata { encoding: Some("gzip".into()), size: 629145600, crc64nvme: "SVWSLdzczJc=".into(), @@ -82,7 +84,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"07b92482167eaabedd8111906c7d96763d0da5ea951a5be04492413508206261" + @"75a9b81695bbea8934d88e221dd3dd98392170904d4775a84145767117555a60" ); assert_json_snapshot!(upload_metadata); } @@ -93,7 +95,7 @@ mod tests { repository_provider: RepositoryProvider::Project, version: Some(LATEST_UPLOAD_METADATA_VERSION), tokenless: false, - profile_archive: ProfileMetadata { + profile_archive_metadata: ProfileArchiveMetadata { encoding: Some("gzip".into()), size: 20480, crc64nvme: "XMXoclwBfLo=".into(), @@ -158,7 +160,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"5fd485b7ad81097b43268e651dfeec697feb7985f3a8e6e1e5eea39d213262cf" + @"c3b2c0af52a3416543c2f9696aa4c12013027d6e190386d5b0babf9fd895bbc7" ); assert_json_snapshot!(upload_metadata); } diff --git a/src/upload/uploader.rs b/src/upload/uploader.rs index fcf72b712..1f32729cf 100644 --- a/src/upload/uploader.rs +++ b/src/upload/uploader.rs @@ -22,8 +22,8 @@ use tokio::fs::File; use tokio::io::AsyncWriteExt; use tokio_tar::Builder; -use super::interfaces::{MultipartUpload, ProfileMetadata, UploadData, UploadMetadata}; -use super::profile_archive::{ProfileArchive, concurrent_part_uploads}; +use super::interfaces::{UploadData, UploadMetadata}; +use super::profile_archive::{CONCURRENT_PART_UPLOADS, ProfileArchive}; use super::s3; fn human_bytes(bytes: u64) -> String { @@ -47,7 +47,7 @@ fn human_bytes_per_second(bytes: u64, elapsed: std::time::Duration) -> String { format!("{}/s", human_bytes(bytes_per_second as u64)) } -/// Create a profile archive from the profile folder, along with its description +/// Create a profile archive from the profile folder /// /// For Valgrind, we create a gzip-compressed tar archive of the entire profile folder. /// For WallTime and Memory, we create an uncompressed tar archive on disk: their @@ -262,20 +262,20 @@ async fn send_upload_request( }) } -/// Upload the parts concurrently, then assemble them with the S3 complete request. -/// Parts of a failed upload are cleaned up by the bucket lifecycle rules. -async fn upload_multipart_profile_archive( - multipart_upload: &MultipartUpload, - profile: &ProfileMetadata, - content: &ProfileArchiveContent, +async fn upload_profile_archive( + upload_data: &UploadData, + profile_archive: ProfileArchive, ) -> Result<()> { - debug!("Starting multipart upload for profile archive"); - let part_count = profile.part_crc64nvmes.len(); + let multipart_upload = &upload_data.multipart_upload; + let metadata = &profile_archive.metadata; + let content = &profile_archive.content; + let part_count = metadata.part_crc64nvmes.len(); + let concurrency = *CONCURRENT_PART_UPLOADS; + debug!( - "Multipart upload details: part_count={}, part_size={}, total_size={}", - part_count, profile.part_size, profile.size + "Starting multipart upload for profile archive: part_count={}, part_size={}, total_size={}", + part_count, metadata.part_size, metadata.size ); - if multipart_upload.parts.len() != part_count { bail!( "Received {} part upload requests for {} parts", @@ -284,18 +284,15 @@ async fn upload_multipart_profile_archive( ); } - let concurrency = concurrent_part_uploads(); let upload_start = Instant::now(); - // Unordered, so that a part finishing frees its slot even while an earlier part is still - // uploading let part_requests = multipart_upload.parts.iter().enumerate(); let mut indexed_etags: Vec<_> = futures::stream::iter(part_requests) .map(|(index, part_request)| async move { - let offset = index as u64 * profile.part_size; + let offset = index as u64 * metadata.part_size; let range = ContentRange { content, offset, - length: profile.part_size.min(profile.size - offset), + length: metadata.part_size.min(metadata.size - offset), }; debug!( "Uploading part {}/{} ({} bytes)", @@ -320,10 +317,13 @@ async fn upload_multipart_profile_archive( ); Ok::<_, anyhow::Error>((index, etag)) }) + // Unordered, so that fast connections pick up remaining parts without waiting + // for slower ones .buffer_unordered(concurrency) .try_collect() .await?; + // ETags need to be sorted in the complete upload request indexed_etags.sort_unstable_by_key(|(index, _)| *index); let etags: Vec<_> = indexed_etags.into_iter().map(|(_, etag)| etag).collect(); @@ -332,12 +332,11 @@ async fn upload_multipart_profile_archive( "Uploaded {} part{} ({}) in {:.1?} with {} concurrent uploads ({})", part_count, if part_count > 1 { "s" } else { "" }, - human_bytes(profile.size), + human_bytes(metadata.size), upload_elapsed, concurrency, - human_bytes_per_second(profile.size, upload_elapsed) + human_bytes_per_second(metadata.size, upload_elapsed) ); - with_upload_retry(|| async { let request = s3::complete_upload(&STREAMING_CLIENT, &multipart_upload.complete, &etags); let response = send_upload_request(request).await?; @@ -352,18 +351,6 @@ async fn upload_multipart_profile_archive( .await } -async fn upload_profile_archive( - upload_data: &UploadData, - profile_archive: ProfileArchive, -) -> Result<()> { - upload_multipart_profile_archive( - &upload_data.multipart_upload, - &profile_archive.metadata, - &profile_archive.content, - ) - .await -} - #[derive(Clone)] pub struct UploadResult { pub run_id: String, @@ -422,7 +409,10 @@ pub async fn upload( #[cfg(test)] mod tests { - use crate::api_client::CodSpeedAPIClient; + use crate::{ + api_client::CodSpeedAPIClient, + upload::{MultipartUpload, ProfileArchiveMetadata}, + }; use temp_env::async_with_vars; use url::Url; @@ -543,10 +533,10 @@ mod tests { (url, hits, handle) } - /// `send_with_retry` retries transient failures itself, since `STREAMING_CLIENT` + /// `with_upload_retry` retries transient failures itself, since `STREAMING_CLIENT` /// has no retry middleware. #[tokio::test] - async fn streamed_upload_is_retried() { + async fn part_upload_is_retried() { use std::sync::atomic::Ordering; let (url, hits, server) = spawn_mock_returning_503(EXPECTED_ATTEMPTS); @@ -571,28 +561,7 @@ mod tests { assert_eq!( hits.load(Ordering::SeqCst), EXPECTED_ATTEMPTS, - "streamed upload should be attempted 1 + UPLOAD_RETRY_COUNT times" - ); - } - - #[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()) - .await - .unwrap(); - - let result = upload_profile_archive(&multipart_upload_data_for(&url, 1), 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" + "part upload should be attempted 1 + UPLOAD_RETRY_COUNT times" ); } @@ -686,7 +655,7 @@ mod tests { let encoding = archive_content.encoding(); ProfileArchive { content: archive_content, - metadata: ProfileMetadata { + metadata: ProfileArchiveMetadata { encoding, size: content.len() as u64, crc64nvme: "crc".to_string(),