diff --git a/Cargo.lock b/Cargo.lock index 6a23121e5..1fdf3b3b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -708,9 +708,11 @@ dependencies = [ "axoupdater", "base64", "bincode", + "bytes", "caps", "clap", "console", + "crc-fast", "crc32fast", "debugid", "escargot", @@ -727,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", @@ -889,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", ] @@ -902,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" @@ -2368,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" @@ -4330,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 9df21458c..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" @@ -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..31b3328a9 100644 --- a/src/run_environment/provider.rs +++ b/src/run_environment/provider.rs @@ -136,8 +136,7 @@ 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_encoding: profile_archive.content.encoding(), + 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 fbf6abad7..c03f72620 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; @@ -5,7 +7,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")] @@ -13,8 +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_archive_metadata: ProfileArchiveMetadata, pub runner: Runner, pub run_environment: RunEnvironment, pub run_part: Option, @@ -24,6 +25,22 @@ pub struct UploadMetadata { pub run_environment_metadata: RunEnvironmentMetadata, } +/// 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 ProfileArchiveMetadata { + /// `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 big-endian CRC64NVME of each part, in upload order + pub part_crc64nvmes: Vec, +} + #[derive(Serialize, Debug)] #[serde(rename_all = "camelCase")] pub struct Runner { @@ -46,12 +63,69 @@ pub struct Runner { #[serde(rename_all = "camelCase")] pub struct UploadData { pub status: String, - pub upload_url: String, + pub multipart_upload: MultipartUpload, pub run_id: String, } +#[derive(Deserialize, Serialize, Debug)] +#[serde(rename_all = "camelCase")] +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)] #[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", + "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 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/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..9308810f7 100644 --- a/src/upload/profile_archive.rs +++ b/src/upload/profile_archive.rs @@ -1,60 +1,165 @@ use base64::{Engine, engine::general_purpose}; +use crc_fast::{CrcAlgorithm, Digest}; use crate::prelude::*; +use crate::upload::interfaces::ProfileArchiveMetadata; +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"; + +pub 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 + } + } +}); + +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 +/// 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 hash: String, pub content: ProfileArchiveContent, + pub metadata: ProfileArchiveMetadata, } #[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 }, +/// 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 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_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_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)?; + if read == 0 { + break; + } + let mut chunk = &buffer[..read]; + while !chunk.is_empty() { + 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_digest.get_amount() == part_size { + whole_digest.combine(&part_digest); + part_crc64nvmes.push(encode_crc64nvme(part_digest.finalize_reset())); + } } } + if part_digest.get_amount() > 0 { + whole_digest.combine(&part_digest); + part_crc64nvmes.push(encode_crc64nvme(part_digest.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_crc64nvme(whole_digest.finalize()), part_crc64nvmes)) +} + +/// [`compute_crc64nvmes_from_reader`] over the content, on the blocking thread pool +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_crc64nvmes_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 }, - }) - } - - pub fn new_compressed_on_disk(path: PathBuf) -> Result { - let metadata = std::fs::metadata(&path)?; + ProfileArchiveContent::UncompressedOnDisk { path } => { + let path = path.clone(); + tokio::task::spawn_blocking(move || { + compute_crc64nvmes_from_reader(std::fs::File::open(path)?, part_size) + }) + .await? + } + } +} + +/// 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 * TARGET_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 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); - Ok(ProfileArchive { - hash, - content: ProfileArchiveContent::CompressedOnDisk { path }, - }) + 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 (crc64nvme, part_crc64nvmes) = compute_crc64nvmes(&content, part_size).await?; + let metadata = ProfileArchiveMetadata { + encoding: content.encoding(), + size, + crc64nvme, + part_size, + part_crc64nvmes, + }; + + Ok(ProfileArchive { content, metadata }) } } @@ -62,8 +167,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 +177,123 @@ 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 + } + + fn crc64nvme(data: &[u8]) -> String { + encode_crc64nvme(crc_fast::checksum(CrcAlgorithm::Crc64Nvme, data)) + } + + #[test] + 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 (whole, part_crc64nvmes) = + compute_crc64nvmes_from_reader(&content[..], part_size).unwrap(); + + 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] + 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 crc = crc64nvme(b"profile-archive"); + assert_eq!( + archive.metadata, + ProfileArchiveMetadata { + encoding: None, + size: b"profile-archive".len() as u64, + crc64nvme: crc.clone(), + part_size: MULTIPART_MIN_PART_SIZE_BYTES, + part_crc64nvmes: vec![crc], + } + ); + } + + #[test] + fn part_count_hits_the_target_within_the_part_size_bounds() { + const MIB: u64 = 1024 * 1024; + 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] + 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..16c028bcf --- /dev/null +++ b/src/upload/s3.rs @@ -0,0 +1,137 @@ +//! 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. +//! +//! 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, + request: &PresignedRequest, +) -> reqwest::RequestBuilder { + presigned(client, reqwest::Method::PUT, request) +} + +/// 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, + request: &PresignedRequest, + etags: &[PartETag], +) -> reqwest::RequestBuilder { + presigned(client, reqwest::Method::POST, request) + .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}") +} 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..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 @@ -4,10 +4,17 @@ expression: upload_metadata --- { "repositoryProvider": "PROJECT", - "version": 11, + "version": 12, "tokenless": false, - "profileMd5": "tfC4VxYiYdJcTWpHpv4Ouw==", - "profileEncoding": "gzip", + "profileArchiveMetadata": { + "encoding": "gzip", + "size": 20480, + "crc64nvme": "XMXoclwBfLo=", + "partSize": 16777216, + "partCrc64nvmes": [ + "XMXoclwBfLo=" + ] + }, "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..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 @@ -4,10 +4,19 @@ expression: upload_metadata --- { "repositoryProvider": "GITHUB", - "version": 11, + "version": 12, "tokenless": true, - "profileMd5": "jp/k05RKuqP3ERQuIIvx4Q==", - "profileEncoding": "gzip", + "profileArchiveMetadata": { + "encoding": "gzip", + "size": 629145600, + "crc64nvme": "SVWSLdzczJc=", + "partSize": 268435456, + "partCrc64nvmes": [ + "w7Dr86r5+2I=", + "puujmsfPTns=", + "Tt/QAZdhKk4=" + ] + }, "runner": { "name": "codspeed-runner", "version": "2.1.0", diff --git a/src/upload/upload_metadata.rs b/src/upload/upload_metadata.rs index b90e75070..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, Runner, UploadMetadata}; + use crate::upload::{ + LATEST_UPLOAD_METADATA_VERSION, ProfileArchiveMetadata, Runner, UploadMetadata, + }; #[test] fn test_get_metadata_hash() { @@ -30,8 +32,17 @@ 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_archive_metadata: ProfileArchiveMetadata { + encoding: Some("gzip".into()), + size: 629145600, + crc64nvme: "SVWSLdzczJc=".into(), + part_size: 268435456, + part_crc64nvmes: vec![ + "w7Dr86r5+2I=".into(), + "puujmsfPTns=".into(), + "Tt/QAZdhKk4=".into(), + ], + }, runner: Runner { name: "codspeed-runner".into(), version: "2.1.0".into(), @@ -73,7 +84,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"b6e221583869b0a49498d71538432a5396b08fec7783081dcdf23c9e037a9365" + @"75a9b81695bbea8934d88e221dd3dd98392170904d4775a84145767117555a60" ); assert_json_snapshot!(upload_metadata); } @@ -84,8 +95,13 @@ mod tests { repository_provider: RepositoryProvider::Project, version: Some(LATEST_UPLOAD_METADATA_VERSION), tokenless: false, - profile_md5: "tfC4VxYiYdJcTWpHpv4Ouw==".into(), - profile_encoding: Some("gzip".into()), + profile_archive_metadata: ProfileArchiveMetadata { + encoding: Some("gzip".into()), + size: 20480, + crc64nvme: "XMXoclwBfLo=".into(), + part_size: 16777216, + part_crc64nvmes: vec!["XMXoclwBfLo=".into()], + }, runner: Runner { name: "codspeed-runner".into(), version: "4.11.1".into(), @@ -144,7 +160,7 @@ mod tests { hash, // Caution: when changing this value, we need to ensure that // the related backend snapshot remains the same - @"5c960260ea5b5ceaafa20ea220e566743e20b65106d0fd672a844a63bba3835a" + @"c3b2c0af52a3416543c2f9696aa4c12013027d6e190386d5b0babf9fd895bbc7" ); assert_json_snapshot!(upload_metadata); } diff --git a/src/upload/uploader.rs b/src/upload/uploader.rs index 1afa4715a..16ef3bd0f 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::profile_archive::{CONCURRENT_PART_UPLOADS, ProfileArchive}; +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 +/// Create a profile archive from the profile folder /// /// 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,193 @@ 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 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; + } +} - 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; - } +/// 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() + ) } - - return Ok(result?); - } + Err(error) => error.into(), + }; + Err(if is_transient { + AttemptError::Transient(error) + } else { + AttemptError::Permanent(error) + }) } 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); - } + 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; - 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(), - ) - .await? - } - }; - - if !response.status().is_success() { - let status = response.status(); - let error_text = response.text().await?; + debug!( + "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!( - "Failed to upload performance report: {}\n -> {} {}", - status, - style("Reason:").bold(), - style(error_text).red() + "Received {} part upload requests for {} parts", + multipart_upload.parts.len(), + part_count ); } - Ok(()) + let upload_start = Instant::now(); + 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 * metadata.part_size; + let range = ContentRange { + content, + offset, + length: metadata.part_size.min(metadata.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)) + }) + // 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(); + + let upload_elapsed = upload_start.elapsed(); + debug!( + "Uploaded {} part{plural} ({}) in {:.1?} with {} concurrent upload{plural} ({})", + part_count, + human_bytes(metadata.size), + upload_elapsed, + concurrency.min(part_count), + human_bytes_per_second(metadata.size, upload_elapsed), + plural = if part_count > 1 { "s" } else { "" }, + ); + with_upload_retry(|| async { + 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(()), + 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 } #[derive(Clone)] @@ -359,11 +409,15 @@ 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; use super::*; + use crate::upload::interfaces::PresignedRequest; use std::path::PathBuf; // TODO: remove the ignore when implementing network mocking @@ -442,7 +496,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 +510,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,18 +533,10 @@ 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. + /// `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); @@ -501,9 +547,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!( @@ -513,28 +561,266 @@ mod tests { assert_eq!( hits.load(Ordering::SeqCst), EXPECTED_ATTEMPTS, - "streamed upload should be attempted 1 + UPLOAD_RETRY_COUNT times" + "part upload should be attempted 1 + UPLOAD_RETRY_COUNT times" ); } - /// In-memory archives go through `REQUEST_CLIENT`, whose retry middleware handles - /// transient failures. + 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() + ) + } + + /// 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 part_count = content.len().div_ceil(part_size as usize); + 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 } + }; + let encoding = archive_content.encoding(); + ProfileArchive { + content: archive_content, + metadata: ProfileArchiveMetadata { + encoding, + size: content.len() as u64, + crc64nvme: "crc".to_string(), + part_size, + 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: MultipartUpload { + parts: (1..=part_count) + .map(|part| { + presigned( + format!("/part/{part}"), + ("x-amz-checksum-crc64nvme", format!("signed-crc-{part}")), + ) + }) + .collect(), + complete: presigned( + "/complete".to_string(), + ("x-amz-mp-object-size", "signed-size".to_string()), + ), + }, + run_id: "test-run".to_string(), + } + } + #[tokio::test] - async fn in_memory_upload_is_retried() { - use std::sync::atomic::Ordering; + async fn multipart_upload_sends_each_part_then_completes() { + assert_multipart_upload_sends_each_part_then_completes(false).await; + } - let (url, hits, server) = spawn_mock_returning_503(EXPECTED_ATTEMPTS); + #[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; + } - let archive = ProfileArchive::new_compressed_in_memory(b"profile-archive".to_vec()); + 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 result = upload_profile_archive(&upload_data_for(url), archive).await; - server.join().unwrap(); + 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("", "") + } + }); - assert!(result.is_err(), "a 503 should surface as an error"); + 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["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!( - hits.load(Ordering::SeqCst), - EXPECTED_ATTEMPTS, - "in-memory upload should be attempted 1 + UPLOAD_RETRY_COUNT times" + 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(); } }