Perform startup/init code eagerly in entry::install_crates (#880)

The startup/init code in `entry::install_crates` performs a lot of blocking operations, from testing if dir exists to reading from files and there is no `.await` in there until after the startup.

There are also a few cases where `block_in_place` should be called (e.g. loading manifests, loading TLS certificates) but is missing.

Most of the `Args` passed to `entry::install_crates` are actually consumed before the first `.await` point, so performing startup/init code eagerly would make the generated future much smaller, reduce codegen and also makes it easier to optimize.

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>

crates/bin/src/bin_util.rs

@@ -35,29 +35,33 @@ impl Termination for MainExit {
 }
 
 impl MainExit {
-    pub fn new(result: Result<Result<()>, BinstallError>, done: Duration) -> Self {
-        result.map_or_else(MainExit::Error, |res| {
-            res.map(|()| MainExit::Success(Some(done)))
-                .unwrap_or_else(|err| {
-                    err.downcast::<BinstallError>()
-                        .map(MainExit::Error)
-                        .unwrap_or_else(MainExit::Report)
-                })
-        })
+    pub fn new(res: Result<()>, done: Duration) -> Self {
+        res.map(|()| MainExit::Success(Some(done)))
+            .unwrap_or_else(|err| {
+                err.downcast::<BinstallError>()
+                    .map(MainExit::Error)
+                    .unwrap_or_else(MainExit::Report)
+            })
     }
 }
 
 /// This function would start a tokio multithreading runtime,
-/// spawn a new task on it that runs `f`, then `block_on` it.
+/// spawn a new task on it that runs `f()`, then `block_on` it.
 ///
 /// It will cancel the future if user requested cancellation
 /// via signal.
-pub fn run_tokio_main<F, T>(f: F) -> Result<T, BinstallError>
+pub fn run_tokio_main<Func, Fut>(f: Func) -> Result<()>
 where
-    F: Future<Output = T> + Send + 'static,
-    T: Send + 'static,
+    Func: FnOnce() -> Result<Option<Fut>>,
+    Fut: Future<Output = Result<()>> + Send + 'static,
 {
-    let rt = Runtime::new()?;
-    let handle = AutoAbortJoinHandle::new(rt.spawn(f));
-    rt.block_on(cancel_on_user_sig_term(handle))
+    let rt = Runtime::new().map_err(BinstallError::from)?;
+    let _guard = rt.enter();
+
+    if let Some(fut) = f()? {
+        let handle = AutoAbortJoinHandle::new(rt.spawn(fut));
+        rt.block_on(cancel_on_user_sig_term(handle))?
+    } else {
+        Ok(())
+    }
 }

crates/bin/src/entry.rs

@@ -1,5 +1,6 @@
 use std::{
     env, fs,
+    future::Future,
     path::{Path, PathBuf},
     sync::Arc,
     time::Duration,
@@ -35,7 +36,10 @@ use crate::{
     ui::confirm,
 };
 
-pub async fn install_crates(args: Args, jobserver_client: LazyJobserverClient) -> Result<()> {
+pub fn install_crates(
+    args: Args,
+    jobserver_client: LazyJobserverClient,
+) -> Result<Option<impl Future<Output = Result<()>>>> {
     // Compute Resolvers
     let mut cargo_install_fallback = false;
 
@@ -63,7 +67,7 @@ pub async fn install_crates(args: Args, jobserver_client: LazyJobserverClient) -
 
     if crate_names.peek().is_none() {
         debug!("Nothing to do");
-        return Ok(());
+        return Ok(None);
     }
 
     // Launch target detection
@@ -131,53 +135,55 @@ pub async fn install_crates(args: Args, jobserver_client: LazyJobserverClient) -
         })
         .collect();
 
-    // Collect results
-    let mut resolution_fetchs = Vec::new();
-    let mut resolution_sources = Vec::new();
+    Ok(Some(async move {
+        // Collect results
+        let mut resolution_fetchs = Vec::new();
+        let mut resolution_sources = Vec::new();
 
-    for task in tasks {
-        match task.await?? {
-            Resolution::AlreadyUpToDate => {}
-            Resolution::Fetch(fetch) => {
-                fetch.print(&binstall_opts);
-                resolution_fetchs.push(fetch)
-            }
-            Resolution::InstallFromSource(source) => {
-                source.print();
-                resolution_sources.push(source)
+        for task in tasks {
+            match task.await?? {
+                Resolution::AlreadyUpToDate => {}
+                Resolution::Fetch(fetch) => {
+                    fetch.print(&binstall_opts);
+                    resolution_fetchs.push(fetch)
+                }
+                Resolution::InstallFromSource(source) => {
+                    source.print();
+                    resolution_sources.push(source)
+                }
             }
         }
-    }
 
-    if resolution_fetchs.is_empty() && resolution_sources.is_empty() {
-        debug!("Nothing to do");
-        return Ok(());
-    }
+        if resolution_fetchs.is_empty() && resolution_sources.is_empty() {
+            debug!("Nothing to do");
+            return Ok(());
+        }
 
-    // Confirm
-    if !dry_run && !no_confirm {
-        confirm().await?;
-    }
+        // Confirm
+        if !dry_run && !no_confirm {
+            confirm().await?;
+        }
 
-    do_install_fetches(
-        resolution_fetchs,
-        manifests,
-        &binstall_opts,
-        dry_run,
-        temp_dir,
-        no_cleanup,
-    )?;
+        do_install_fetches(
+            resolution_fetchs,
+            manifests,
+            &binstall_opts,
+            dry_run,
+            temp_dir,
+            no_cleanup,
+        )?;
 
-    let tasks: Vec<_> = resolution_sources
-        .into_iter()
-        .map(|source| AutoAbortJoinHandle::spawn(source.install(binstall_opts.clone())))
-        .collect();
+        let tasks: Vec<_> = resolution_sources
+            .into_iter()
+            .map(|source| AutoAbortJoinHandle::spawn(source.install(binstall_opts.clone())))
+            .collect();
 
-    for task in tasks {
-        task.await??;
-    }
+        for task in tasks {
+            task.await??;
+        }
 
-    Ok(())
+        Ok(())
+    }))
 }
 
 fn do_read_root_cert(path: &Path) -> Result<Option<Certificate>, BinstallError> {
@@ -228,43 +234,41 @@ fn compute_paths_and_load_manifests(
     roots: Option<PathBuf>,
     install_path: Option<PathBuf>,
 ) -> Result<(PathBuf, Option<Manifests>, tempfile::TempDir)> {
-    block_in_place(|| {
-        // Compute cargo_roots
-        let cargo_roots = install_path::get_cargo_roots_path(roots).ok_or_else(|| {
-            error!("No viable cargo roots path found of specified, try `--roots`");
-            miette!("No cargo roots path found or specified")
-        })?;
+    // Compute cargo_roots
+    let cargo_roots = install_path::get_cargo_roots_path(roots).ok_or_else(|| {
+        error!("No viable cargo roots path found of specified, try `--roots`");
+        miette!("No cargo roots path found or specified")
+    })?;
 
-        // Compute install directory
-        let (install_path, custom_install_path) =
-            install_path::get_install_path(install_path, Some(&cargo_roots));
-        let install_path = install_path.ok_or_else(|| {
-            error!("No viable install path found of specified, try `--install-path`");
-            miette!("No install path found or specified")
-        })?;
-        fs::create_dir_all(&install_path).map_err(BinstallError::Io)?;
-        debug!("Using install path: {}", install_path.display());
+    // Compute install directory
+    let (install_path, custom_install_path) =
+        install_path::get_install_path(install_path, Some(&cargo_roots));
+    let install_path = install_path.ok_or_else(|| {
+        error!("No viable install path found of specified, try `--install-path`");
+        miette!("No install path found or specified")
+    })?;
+    fs::create_dir_all(&install_path).map_err(BinstallError::Io)?;
+    debug!("Using install path: {}", install_path.display());
 
-        // Load manifests
-        let manifests = if !custom_install_path {
-            Some(Manifests::open_exclusive(&cargo_roots)?)
-        } else {
-            None
-        };
+    // Load manifests
+    let manifests = if !custom_install_path {
+        Some(Manifests::open_exclusive(&cargo_roots)?)
+    } else {
+        None
+    };
 
-        // Create a temporary directory for downloads etc.
-        //
-        // Put all binaries to a temporary directory under `dst` first, catching
-        // some failure modes (e.g., out of space) before touching the existing
-        // binaries. This directory will get cleaned up via RAII.
-        let temp_dir = tempfile::Builder::new()
-            .prefix("cargo-binstall")
-            .tempdir_in(&install_path)
-            .map_err(BinstallError::from)
-            .wrap_err("Creating a temporary directory failed.")?;
+    // Create a temporary directory for downloads etc.
+    //
+    // Put all binaries to a temporary directory under `dst` first, catching
+    // some failure modes (e.g., out of space) before touching the existing
+    // binaries. This directory will get cleaned up via RAII.
+    let temp_dir = tempfile::Builder::new()
+        .prefix("cargo-binstall")
+        .tempdir_in(&install_path)
+        .map_err(BinstallError::from)
+        .wrap_err("Creating a temporary directory failed.")?;
 
-        Ok((install_path, manifests, temp_dir))
-    })
+    Ok((install_path, manifests, temp_dir))
 }
 
 /// Return vec of (crate_name, current_version)

crates/bin/src/main.rs

@@ -32,7 +32,7 @@ fn main() -> MainExit {
 
         let start = Instant::now();
 
-        let result = run_tokio_main(entry::install_crates(args, jobserver_client));
+        let result = run_tokio_main(|| entry::install_crates(args, jobserver_client));
 
         let done = start.elapsed();
         debug!("run time: {done:?}");