Merge pull request #10064 from eth3lbert/async-tarball

Migrate `crates_io_tarball` to async/await

Author: Turbo87

Cargo.lock

@@ -8,23 +8,29 @@ edition = "2021"
 workspace = true
 
 [features]
-builder = []
+builder = ["dep:flate2", "dep:tar"]
 
 [dependencies]
 cargo-manifest = "=0.16.1"
-flate2 = "=1.0.35"
+flate2 = { version = "=1.0.35", optional = true }
 serde = { version = "=1.0.215", features = ["derive"] }
 serde_json = "=1.0.133"
-tar = "=0.4.43"
+tar = { version = "=0.4.43", optional = true }
 thiserror = "=2.0.3"
 tracing = "=0.1.40"
+tokio = { version = "=1.41.1", features = ["io-util", "macros", "rt-multi-thread"] }
+async-compression = { version = "=0.4.18", default-features = false, features = ["gzip", "tokio"] }
+krata-tokio-tar = "=0.4.2"
+futures-util = "=0.3.31"
 
 [dev-dependencies]
 anyhow = "=1.0.93"
 claims = "=0.8.0"
 clap = { version = "=4.5.21", features = ["derive", "unicode", "wrap_help"] }
+flate2 = { version = "=1.0.35" }
 indicatif = { version = "=0.17.9", features = ["rayon"] }
 insta = "=1.41.1"
 rayon = "=1.10.0"
+tar = { version = "=0.4.43" }
 tracing-subscriber = { version = "=0.3.18", features = ["env-filter"] }
 walkdir = "=2.5.0"

crates/crates_io_tarball/Cargo.toml

@@ -1,10 +1,11 @@
 use anyhow::anyhow;
 use clap::Parser;
 use crates_io_tarball::process_tarball;
-use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle};
+use futures_util::{stream, StreamExt};
+use indicatif::{ParallelProgressIterator, ProgressBar, ProgressIterator, ProgressStyle};
 use rayon::prelude::*;
-use std::fs::File;
 use std::path::{Path, PathBuf};
+use tokio::fs::File;
 use tracing::{debug, info, warn};
 use tracing_subscriber::filter::LevelFilter;
 use tracing_subscriber::EnvFilter;
@@ -17,7 +18,8 @@ pub struct Options {
     path: PathBuf,
 }
 
-fn main() -> anyhow::Result<()> {
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
     setup_tracing();
 
     let options = Options::parse();
@@ -51,27 +53,30 @@ fn main() -> anyhow::Result<()> {
         ProgressStyle::with_template("{bar:60} ({pos}/{len}, ETA {eta}) {wide_msg}").unwrap(),
     );
 
-    paths
-        .par_iter()
-        .progress_with(pb.clone())
-        .for_each(|path| process_path(path, &pb));
+    stream::iter(paths.iter().progress_with(pb.clone()))
+        .for_each_concurrent(None, |path| {
+            let pb = pb.clone();
+            async move { process_path(path, &pb).await }
+        })
+        .await;
 
     Ok(())
 }
 
-fn process_path(path: &Path, pb: &ProgressBar) {
-    let file =
-        File::open(path).map_err(|error| pb.suspend(|| warn!(%error, "Failed to read crate file")));
+async fn process_path(path: &Path, pb: &ProgressBar) {
+    let file = File::open(path)
+        .await
+        .map_err(|error| pb.suspend(|| warn!(%error, "Failed to read crate file")));
 
-    let Ok(file) = file else {
+    let Ok(mut file) = file else {
         return;
     };
 
     let path_no_ext = path.with_extension("");
     let pkg_name = path_no_ext.file_name().unwrap().to_string_lossy();
     pb.set_message(format!("{pkg_name}"));
 
-    let result = process_tarball(&pkg_name, &file, u64::MAX);
+    let result = process_tarball(&pkg_name, &mut file, u64::MAX).await;
     pb.suspend(|| match result {
         Ok(result) => debug!(%pkg_name, path = %path.display(), ?result),
         Err(error) => warn!(%pkg_name, path = %path.display(), %error, "Failed to process tarball"),

crates/crates_io_tarball/examples/check_all_crates.rs

@@ -1,8 +1,8 @@
 use anyhow::{anyhow, Context};
 use clap::Parser;
 use crates_io_tarball::process_tarball;
-use std::fs::File;
 use std::path::PathBuf;
+use tokio::fs::File;
 use tracing_subscriber::filter::LevelFilter;
 use tracing_subscriber::EnvFilter;
 
@@ -15,7 +15,8 @@ pub struct Options {
     path: PathBuf,
 }
 
-fn main() -> anyhow::Result<()> {
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
     setup_tracing();
 
     let options = Options::parse();
@@ -25,13 +26,14 @@ fn main() -> anyhow::Result<()> {
         return Err(anyhow!("`{}` not found or not a file", path.display()));
     }
 
-    let file = File::open(&path).context("Failed to read tarball")?;
+    let mut file = File::open(&path).await.context("Failed to read tarball")?;
 
     let path_no_ext = path.with_extension("");
     let pkg_name = path_no_ext.file_name().unwrap().to_string_lossy();
 
-    let result =
-        process_tarball(&pkg_name, &file, u64::MAX).context("Failed to process tarball")?;
+    let result = process_tarball(&pkg_name, &mut file, u64::MAX)
+        .await
+        .context("Failed to process tarball")?;
 
     println!("{result:#?}");