refactor(upload): refactor profile-archive

Author: adriencaccia

src/run/run_environment/local/provider.rs

@@ -8,7 +8,7 @@ use crate::run::config::RepositoryOverride;
 use crate::run::helpers::{GitRemote, parse_git_remote};
 use crate::run::run_environment::{RunEnvironment, RunPart};
 use crate::run::runner::ExecutorName;
-use crate::run::uploader::{Runner, UploadMetadata};
+use crate::run::uploader::{ProfileArchive, Runner, UploadMetadata};
 use crate::run::{
     config::Config,
     helpers::find_repository_root,
@@ -147,8 +147,7 @@ impl RunEnvironmentProvider for LocalProvider {
         &self,
         config: &Config,
         system_info: &SystemInfo,
-        archive_hash: &str,
-        content_encoding: Option<String>,
+        profile_archive: &ProfileArchive,
         executor_name: ExecutorName,
     ) -> Result<UploadMetadata> {
         let run_environment_metadata = self.get_run_environment_metadata()?;
@@ -159,8 +158,8 @@ impl RunEnvironmentProvider for LocalProvider {
             repository_provider: self.get_repository_provider(),
             commit_hash: run_environment_metadata.ref_.clone(),
             run_environment_metadata,
-            profile_md5: archive_hash.into(),
-            content_encoding,
+            profile_md5: profile_archive.hash.clone(),
+            profile_encoding: profile_archive.content.encoding(),
             runner: Runner {
                 name: "codspeed-runner".into(),
                 version: crate::VERSION.into(),

src/run/run_environment/provider.rs

@@ -5,7 +5,7 @@ use crate::prelude::*;
 use crate::run::check_system::SystemInfo;
 use crate::run::config::Config;
 use crate::run::runner::ExecutorName;
-use crate::run::uploader::{Runner, UploadMetadata};
+use crate::run::uploader::{ProfileArchive, Runner, UploadMetadata};
 
 use super::interfaces::{RepositoryProvider, RunEnvironment, RunEnvironmentMetadata, RunPart};
 
@@ -66,8 +66,7 @@ pub trait RunEnvironmentProvider {
         &self,
         config: &Config,
         system_info: &SystemInfo,
-        archive_hash: &str,
-        content_encoding: Option<String>,
+        profile_archive: &ProfileArchive,
         executor_name: ExecutorName,
     ) -> Result<UploadMetadata> {
         let run_environment_metadata = self.get_run_environment_metadata()?;
@@ -79,8 +78,8 @@ pub trait RunEnvironmentProvider {
             tokenless: config.token.is_none(),
             repository_provider: self.get_repository_provider(),
             run_environment_metadata,
-            profile_md5: archive_hash.into(),
-            content_encoding,
+            profile_md5: profile_archive.hash.clone(),
+            profile_encoding: profile_archive.content.encoding(),
             commit_hash,
             runner: Runner {
                 name: "codspeed-runner".into(),

src/run/uploader/interfaces.rs

@@ -14,7 +14,7 @@ pub struct UploadMetadata {
     pub version: Option<u32>,
     pub tokenless: bool,
     pub profile_md5: String,
-    pub content_encoding: Option<String>,
+    pub profile_encoding: Option<String>,
     pub runner: Runner,
     pub run_environment: RunEnvironment,
     pub run_part: Option<RunPart>,

src/run/uploader/mod.rs

@@ -1,6 +1,8 @@
 mod interfaces;
+mod profile_archive;
 mod upload;
 mod upload_metadata;
 
 pub use interfaces::*;
+pub use profile_archive::ProfileArchive;
 pub use upload::upload;

src/run/uploader/profile_archive.rs

@@ -0,0 +1,71 @@
+use base64::{Engine, engine::general_purpose};
+
+use crate::prelude::*;
+use std::path::PathBuf;
+
+#[derive(Debug)]
+pub struct ProfileArchive {
+    pub hash: String,
+    pub content: ProfileArchiveContent,
+}
+
+#[derive(Debug)]
+pub enum ProfileArchiveContent {
+    CompressedInMemory { data: Vec<u8> },
+    UncompressedOnDisk { path: PathBuf },
+}
+
+impl ProfileArchive {
+    pub fn new_compressed_in_memory(data: Vec<u8>) -> Self {
+        let hash = general_purpose::STANDARD.encode(md5::compute(&data).0);
+        ProfileArchive {
+            hash,
+            content: ProfileArchiveContent::CompressedInMemory { data },
+        }
+    }
+
+    pub fn new_uncompressed_on_disk(path: PathBuf) -> Result<Self> {
+        let metadata = std::fs::metadata(&path)?;
+        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::UncompressedOnDisk { path },
+        })
+    }
+}
+
+impl ProfileArchiveContent {
+    pub async fn size(&self) -> Result<u64> {
+        match &self {
+            ProfileArchiveContent::CompressedInMemory { data } => Ok(data.len() as u64),
+            ProfileArchiveContent::UncompressedOnDisk { path } => {
+                let metadata = tokio::fs::metadata(path).await?;
+                Ok(metadata.len())
+            }
+        }
+    }
+
+    pub fn encoding(&self) -> Option<String> {
+        match self {
+            ProfileArchiveContent::CompressedInMemory { .. } => Some("gzip".to_string()),
+            _ => None,
+        }
+    }
+}
+
+impl Drop for ProfileArchiveContent {
+    fn drop(&mut self) {
+        if let ProfileArchiveContent::UncompressedOnDisk { path } = self {
+            if path.exists() {
+                let _ = std::fs::remove_file(path);
+            }
+        }
+    }
+}

src/run/uploader/upload.rs

@@ -2,59 +2,22 @@ use crate::run::{
     check_system::SystemInfo,
     config::Config,
     run_environment::{RunEnvironment, RunEnvironmentProvider},
-    runner::ExecutorName,
-    runner::RunData,
-    uploader::UploadError,
+    runner::{ExecutorName, RunData},
+    uploader::{UploadError, profile_archive::ProfileArchiveContent},
 };
 use crate::{
     prelude::*,
     request_client::{REQUEST_CLIENT, STREAMING_CLIENT},
 };
 use async_compression::tokio::write::GzipEncoder;
-use base64::{Engine as _, engine::general_purpose};
 use console::style;
 use reqwest::StatusCode;
-use std::path::PathBuf;
 use tokio::fs::File;
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::io::AsyncWriteExt;
 use tokio_tar::Builder;
 
 use super::interfaces::{UploadData, UploadMetadata};
-
-#[derive(Debug)]
-enum ProfileArchive {
-    CompressedInMemory(Vec<u8>),
-    UncompressedOnDisk(PathBuf),
-}
-
-impl ProfileArchive {
-    async fn size(&self) -> Result<u64> {
-        match self {
-            ProfileArchive::CompressedInMemory(data) => Ok(data.len() as u64),
-            ProfileArchive::UncompressedOnDisk(path) => {
-                let metadata = tokio::fs::metadata(path).await?;
-                Ok(metadata.len())
-            }
-        }
-    }
-
-    fn to_content_encoding(&self) -> Option<String> {
-        match self {
-            ProfileArchive::CompressedInMemory(_) => Some("gzip".to_string()),
-            ProfileArchive::UncompressedOnDisk(_) => None,
-        }
-    }
-}
-
-impl Drop for ProfileArchive {
-    fn drop(&mut self) {
-        if let ProfileArchive::UncompressedOnDisk(path) = self {
-            if path.exists() {
-                let _ = std::fs::remove_file(path);
-            }
-        }
-    }
-}
+use super::profile_archive::ProfileArchive;
 
 /// Create a profile archive from the profile folder and return its md5 hash encoded in base64
 ///
@@ -63,7 +26,7 @@ impl Drop for ProfileArchive {
 async fn create_profile_archive(
     run_data: &RunData,
     executor_name: ExecutorName,
-) -> Result<(ProfileArchive, String)> {
+) -> Result<ProfileArchive> {
     let time_start = std::time::Instant::now();
     let profile_archive = match executor_name {
         ExecutorName::Valgrind => {
@@ -74,7 +37,8 @@ async fn create_profile_archive(
                 .await?;
             let mut gzip_encoder = tar.into_inner().await?;
             gzip_encoder.shutdown().await?;
-            ProfileArchive::CompressedInMemory(gzip_encoder.into_inner())
+            let data = gzip_encoder.into_inner();
+            ProfileArchive::new_compressed_in_memory(data)
         }
         ExecutorName::WallTime => {
             debug!("Creating uncompressed tar archive for WallTime on disk");
@@ -92,44 +56,18 @@ 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()?;
-            ProfileArchive::UncompressedOnDisk(persistent_path)
-        }
-    };
 
-    let (archive_digest, archive_size) = match &profile_archive {
-        ProfileArchive::CompressedInMemory(data) => {
-            let digest = md5::compute(data.as_slice());
-            (digest, data.len() as u64)
-        }
-        ProfileArchive::UncompressedOnDisk(path) => {
-            let mut file = File::open(path).await.context(format!(
-                "Failed to open uncompressed file at path: {}",
-                path.display()
-            ))?;
-            let mut hasher = md5::Context::new();
-            let mut buffer = [0; 8192];
-            let mut total_size = 0u64;
-
-            loop {
-                let bytes_read = file.read(&mut buffer).await?;
-                if bytes_read == 0 {
-                    break;
-                }
-                hasher.consume(&buffer[..bytes_read]);
-                total_size += bytes_read as u64;
-            }
-            (hasher.compute(), total_size)
+            ProfileArchive::new_uncompressed_on_disk(persistent_path)?
         }
     };
 
-    let archive_hash = general_purpose::STANDARD.encode(archive_digest.0);
     debug!(
         "Created archive ({} bytes) in {:.2?}",
-        archive_size,
+        profile_archive.content.size().await?,
         time_start.elapsed()
     );
 
-    Ok((profile_archive, archive_hash))
+    Ok(profile_archive)
 }
 
 async fn retrieve_upload_data(
@@ -187,23 +125,26 @@ async fn retrieve_upload_data(
 async fn upload_profile_archive(
     upload_data: &UploadData,
     profile_archive: ProfileArchive,
-    archive_hash: &String,
 ) -> Result<()> {
-    let archive_size = profile_archive.size().await?;
+    let archive_size = profile_archive.content.size().await?;
+    let archive_hash = profile_archive.hash;
 
-    let response = match &profile_archive {
-        ProfileArchive::CompressedInMemory(data) => {
+    let response = match &profile_archive.content {
+        ProfileArchiveContent::CompressedInMemory { data } => {
             // Use regular client with retry middleware for compressed data
-            let request = REQUEST_CLIENT
+            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)
-                .header("Content-Encoding", "gzip");
+                .header("Content-MD5", archive_hash);
+
+            if let Some(encoding) = profile_archive.content.encoding() {
+                request = request.header("Content-Encoding", encoding);
+            }
 
             request.body(data.clone()).send().await?
         }
-        ProfileArchive::UncompressedOnDisk(path) => {
+        ProfileArchiveContent::UncompressedOnDisk { path } => {
             // Use streaming client without retry middleware for file streams
             let file = File::open(path)
                 .await
@@ -248,21 +189,15 @@ pub async fn upload(
     run_data: &RunData,
     executor_name: ExecutorName,
 ) -> Result<UploadResult> {
-    let (profile_archive, archive_hash) =
-        create_profile_archive(run_data, executor_name.clone()).await?;
+    let profile_archive = create_profile_archive(run_data, executor_name.clone()).await?;
 
     debug!(
         "Run Environment provider detected: {:?}",
         provider.get_run_environment()
     );
 
-    let upload_metadata = provider.get_upload_metadata(
-        config,
-        system_info,
-        &archive_hash,
-        profile_archive.to_content_encoding(),
-        executor_name,
-    )?;
+    let upload_metadata =
+        provider.get_upload_metadata(config, system_info, &profile_archive, executor_name)?;
     debug!("Upload metadata: {upload_metadata:#?}");
     info!(
         "Linked repository: {}\n",
@@ -283,8 +218,11 @@ pub async fn upload(
     debug!("runId: {}", upload_data.run_id);
 
     info!("Uploading performance data...");
-    debug!("Uploading {} bytes...", profile_archive.size().await?);
-    upload_profile_archive(&upload_data, profile_archive, &archive_hash).await?;
+    debug!(
+        "Uploading {} bytes...",
+        profile_archive.content.size().await?
+    );
+    upload_profile_archive(&upload_data, profile_archive).await?;
     info!("Performance data uploaded");
 
     Ok(UploadResult {

src/run/uploader/upload_metadata.rs

@@ -33,7 +33,7 @@ mod tests {
             version: Some(7),
             tokenless: true,
             profile_md5: "jp/k05RKuqP3ERQuIIvx4Q==".into(),
-            content_encoding: Some("gzip".into()),
+            profile_encoding: Some("gzip".into()),
             runner: Runner {
                 name: "codspeed-runner".into(),
                 version: "2.1.0".into(),
@@ -77,7 +77,7 @@ mod tests {
             hash,
             // Caution: when changing this value, we need to ensure that
             // the related backend snapshot remains the same
-            @"f827f6a834c26d39900c0a9e2dddfaaf22956494c8db911fc06fef72878b0c70"
+            @"7275243b4457a8fa70215c084210bea7518a3994b863e4437fa33c5ae2bd219e"
         );
         assert_json_snapshot!(upload_metadata);
     }