ref(snapshots): Use async I/O for snapshot image processing

Cargo.toml

@@ -69,7 +69,7 @@ sha2 = "0.10.9"
 sourcemap = { version = "9.3.0", features = ["ram_bundle"] }
 symbolic = { version = "12.13.3", features = ["debuginfo-serde", "il2cpp"] }
 thiserror = "1.0.38"
-tokio = { version = "1.47", features = ["rt"] }
+tokio = { version = "1.47", features = ["rt", "fs", "io-util"] }
 url = "2.3.1"
 uuid = { version = "1.3.0", features = ["v4", "serde"] }
 walkdir = "2.3.2"

src/commands/build/snapshots.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
-use std::fs;
 use std::path::{Path, PathBuf};
 use std::str::FromStr as _;
+use std::sync::Arc;
 
 use anyhow::{Context as _, Result};
 use clap::{Arg, ArgMatches, Command};
@@ -11,6 +11,8 @@ use log::{debug, info, warn};
 use objectstore_client::{ClientBuilder, ExpirationPolicy, Usecase};
 use secrecy::ExposeSecret as _;
 use sha2::{Digest as _, Sha256};
+use tokio::io::AsyncReadExt as _;
+use tokio::sync::Mutex;
 use walkdir::WalkDir;
 
 use crate::api::{Api, CreateSnapshotResponse, ImageMetadata, SnapshotsManifest};
@@ -209,11 +211,24 @@ fn validate_image_sizes(images: &[ImageInfo]) -> Result<()> {
     Ok(())
 }
 
-fn compute_sha256_hash(data: &[u8]) -> String {
+async fn compute_sha256_hash(path: &Path) -> Result<String> {
+    let mut file = tokio::fs::File::open(path)
+        .await
+        .with_context(|| format!("Failed to open image for hashing: {}", path.display()))?;
     let mut hasher = Sha256::new();
-    hasher.update(data);
+    let mut buffer = [0u8; 8192];
+    loop {
+        let bytes_read = file
+            .read(&mut buffer)
+            .await
+            .with_context(|| format!("Failed to read image for hashing: {}", path.display()))?;
+        if bytes_read == 0 {
+            break;
+        }
+        hasher.update(&buffer[..bytes_read]);
+    }
     let result = hasher.finalize();
-    format!("{result:x}")
+    Ok(format!("{result:x}"))
 }
 
 fn is_hidden(root: &Path, path: &Path) -> bool {
@@ -254,74 +269,95 @@ fn upload_images(
         .build()?;
 
     let mut scope = Usecase::new("preprod").scope();
-    for (key, value) in &options.objectstore.scopes {
-        scope = scope.push(key, value);
+    let (mut org_id, mut project_id): (Option<String>, Option<String>) = (None, None);
+    for (key, value) in options.objectstore.scopes.into_iter() {
+        scope = scope.push(&key, value.clone());
+        if key == "org" {
+            org_id = Some(value);
+        } else if key == "project" {
+            project_id = Some(value);
+        }
     }
+    let Some(org_id) = org_id else {
+        anyhow::bail!("Missing org in UploadOptions scope");
+    };
+    let Some(project_id) = project_id else {
+        anyhow::bail!("Missing project in UploadOptions scope");
+    };
+
     let session = scope.session(&client)?;
 
     let runtime = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()
         .context("Failed to create tokio runtime")?;
 
-    let mut many_builder = session.many();
-    let mut manifest_entries = HashMap::new();
     let image_count = images.len();
+    let manifest_entries = Arc::new(Mutex::new(HashMap::new()));
 
-    for image in images {
-        debug!("Processing image: {}", image.path.display());
+    runtime.block_on(async {
+        let mut many_builder = session.many();
 
-        let contents = fs::read(&image.path)
-            .with_context(|| format!("Failed to read image: {}", image.path.display()))?;
-        let hash = compute_sha256_hash(&contents);
+        for image in images {
+            debug!("Processing image: {}", image.path.display());
 
-        info!("Queueing {} as {hash}", image.relative_path.display());
-
-        many_builder = many_builder.push(
-            session
-                .put(contents)
-                .key(&hash)
-                .expiration_policy(expiration),
-        );
+            let hash = compute_sha256_hash(&image.path).await?;
+            let file = tokio::fs::File::open(&image.path).await.with_context(|| {
+                format!("Failed to open image for upload: {}", image.path.display())
+            })?;
 
-        let image_file_name = image
-            .relative_path
-            .file_name()
-            .unwrap_or_default()
-            .to_string_lossy()
-            .into_owned();
-        manifest_entries.insert(
-            hash,
-            ImageMetadata {
-                image_file_name,
-                width: image.width,
-                height: image.height,
-            },
-        );
-    }
+            let key = format!("{org_id}/{project_id}/{hash}");
+            info!("Queueing {} as {key}", image.relative_path.display());
 
-    let result = runtime.block_on(async { many_builder.send().error_for_failures().await });
+            many_builder = many_builder.push(
+                session
+                    .put_file(file)
+                    .key(&key)
+                    .expiration_policy(expiration),
+            );
 
-    match result {
-        Ok(()) => {
-            println!(
-                "{} Uploaded {} image {}",
-                style(">").dim(),
-                style(image_count).yellow(),
-                if image_count == 1 { "file" } else { "files" }
+            let image_file_name = image
+                .relative_path
+                .file_name()
+                .unwrap_or_default()
+                .to_string_lossy()
+                .into_owned();
+            manifest_entries.lock().await.insert(
+                hash,
+                ImageMetadata {
+                    image_file_name,
+                    width: image.width,
+                    height: image.height,
+                },
             );
-            Ok(manifest_entries)
         }
-        Err(errors) => {
-            eprintln!("There were errors uploading images:");
-            let mut error_count = 0;
-            for error in errors {
-                eprintln!("  {}", style(error).red());
-                error_count += 1;
+
+        let result = many_builder.send().error_for_failures().await;
+        match result {
+            Ok(()) => {
+                println!(
+                    "{} Uploaded {} image {}",
+                    style(">").dim(),
+                    style(image_count).yellow(),
+                    if image_count == 1 { "file" } else { "files" }
+                );
+                Ok(())
+            }
+            Err(errors) => {
+                eprintln!("There were errors uploading images:");
+                let mut error_count = 0;
+                for error in errors {
+                    eprintln!("  {}", style(error).red());
+                    error_count += 1;
+                }
+                anyhow::bail!("Failed to upload {error_count} out of {image_count} images")
             }
-            anyhow::bail!("Failed to upload {error_count} out of {image_count} images")
         }
-    }
+    })?;
+
+    Ok(Arc::try_unwrap(manifest_entries)
+        .expect("all references should be dropped after runtime completes")
+        .into_inner())
 }
 
 #[cfg(test)]