feat: git::Repository cancellation support (#1288)

feat: `git::Repository` support cancellation.

To make sure users can cancel git operation via signal, e.g. when the
git operation fail or users no longer want to install.

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

crates/binstalk-downloader/Cargo.toml

@@ -17,6 +17,7 @@ binstalk-types = { version = "0.5.0", path = "../binstalk-types" }
 bytes = "1.4.0"
 bzip2 = "0.4.4"
 compact_str = "0.7.0"
+derive_destructure2 = "0.1"
 flate2 = { version = "1.0.26", default-features = false }
 futures-util = "0.3.28"
 generic-array = "0.14.7"

crates/binstalk-downloader/src/git.rs

@@ -1,6 +1,5 @@
 use std::{fmt, mem, num::NonZeroU32, path::Path, str::FromStr, sync::atomic::AtomicBool};
 
-use compact_str::CompactString;
 use gix::{clone, create, open, remote, Url};
 use thiserror::Error as ThisError;
 use tracing::debug;
@@ -8,6 +7,9 @@ use tracing::debug;
 mod progress_tracing;
 use progress_tracing::TracingProgress;
 
+mod cancellation_token;
+pub use cancellation_token::{GitCancelOnDrop, GitCancellationToken};
+
 pub use gix::url::parse::Error as GitUrlParseError;
 
 #[derive(Debug, ThisError)]
@@ -116,14 +118,21 @@ impl Repository {
     /// async context then you must wrap the call in [`tokio::task::spawn_blocking`].
     ///
     /// WARNING: This function must be called after tokio runtime is initialized.
-    pub fn shallow_clone_bare(url: GitUrl, path: &Path) -> Result<Self, GitError> {
+    pub fn shallow_clone_bare(
+        url: GitUrl,
+        path: &Path,
+        cancellation_token: Option<GitCancellationToken>,
+    ) -> Result<Self, GitError> {
         debug!("Shallow cloning {url} to {}", path.display());
 
         Ok(Self(
             Self::prepare_fetch(url, path, create::Kind::Bare)?
                 .fetch_only(
-                    &mut TracingProgress::new(CompactString::new("Cloning")),
-                    &AtomicBool::new(false),
+                    &mut TracingProgress::new("Cloning bare"),
+                    cancellation_token
+                        .as_ref()
+                        .map(GitCancellationToken::get_atomic)
+                        .unwrap_or(&AtomicBool::new(false)),
                 )?
                 .0
                 .into(),
@@ -134,16 +143,26 @@ impl Repository {
     /// async context then you must wrap the call in [`tokio::task::spawn_blocking`].
     ///
     /// WARNING: This function must be called after tokio runtime is initialized.
-    pub fn shallow_clone(url: GitUrl, path: &Path) -> Result<Self, GitError> {
+    pub fn shallow_clone(
+        url: GitUrl,
+        path: &Path,
+        cancellation_token: Option<GitCancellationToken>,
+    ) -> Result<Self, GitError> {
         debug!("Shallow cloning {url} to {} with worktree", path.display());
 
-        let mut progress = TracingProgress::new(CompactString::new("Cloning"));
+        let mut progress = TracingProgress::new("Cloning with worktree");
 
         Ok(Self(
             Self::prepare_fetch(url, path, create::Kind::WithWorktree)?
                 .fetch_then_checkout(&mut progress, &AtomicBool::new(false))?
                 .0
-                .main_worktree(&mut progress, &AtomicBool::new(false))?
+                .main_worktree(
+                    &mut progress,
+                    cancellation_token
+                        .as_ref()
+                        .map(GitCancellationToken::get_atomic)
+                        .unwrap_or(&AtomicBool::new(false)),
+                )?
                 .0
                 .into(),
         ))

crates/binstalk-downloader/src/git/cancellation_token.rs

@@ -0,0 +1,44 @@
+use std::sync::{
+    atomic::{AtomicBool, Ordering::Relaxed},
+    Arc,
+};
+
+use derive_destructure2::destructure;
+
+/// Token that can be used to cancel git operation.
+#[derive(Clone, Debug, Default)]
+pub struct GitCancellationToken(Arc<AtomicBool>);
+
+impl GitCancellationToken {
+    /// Create a guard that cancel the git operation on drop.
+    #[must_use = "You must assign the guard to a variable, \
+otherwise it is equivalent to `GitCancellationToken::cancel()`"]
+    pub fn cancel_on_drop(self) -> GitCancelOnDrop {
+        GitCancelOnDrop(self)
+    }
+
+    /// Cancel the git operation.
+    pub fn cancel(&self) {
+        self.0.store(true, Relaxed)
+    }
+
+    pub(super) fn get_atomic(&self) -> &AtomicBool {
+        &self.0
+    }
+}
+
+/// Guard used to cancel git operation on drop
+#[derive(Debug, destructure)]
+pub struct GitCancelOnDrop(GitCancellationToken);
+
+impl Drop for GitCancelOnDrop {
+    fn drop(&mut self) {
+        self.0.cancel()
+    }
+}
+impl GitCancelOnDrop {
+    /// Disarm the guard, return the token.
+    pub fn disarm(self) -> GitCancellationToken {
+        self.destructure().0
+    }
+}

crates/binstalk-downloader/src/git/progress_tracing.rs

@@ -27,7 +27,7 @@ const SEP: &str = "::";
 
 impl TracingProgress {
     /// Create a new instanCompactce from `name`.
-    pub fn new(name: CompactString) -> Self {
+    pub fn new(name: &str) -> Self {
         let trigger = Arc::new(AtomicBool::new(true));
         tokio::spawn({
             let mut interval = time::interval(Duration::from_secs_f32(EMIT_LOG_EVERY_S));
@@ -43,7 +43,7 @@ impl TracingProgress {
             }
         });
         Self {
-            name,
+            name: CompactString::new(name),
             id: UNKNOWN,
             max: None,
             step: 0,

crates/binstalk/src/drivers/registry/git_registry.rs

@@ -16,7 +16,7 @@ use crate::{
     errors::BinstallError,
     helpers::{
         cargo_toml::Manifest,
-        git::{GitUrl, Repository},
+        git::{GitCancellationToken, GitUrl, Repository},
         remote::Client,
     },
     manifests::cargo_toml_binstall::Meta,
@@ -30,10 +30,14 @@ struct GitIndex {
 }
 
 impl GitIndex {
-    fn new(url: GitUrl) -> Result<Self, BinstallError> {
+    fn new(url: GitUrl, cancellation_token: GitCancellationToken) -> Result<Self, BinstallError> {
         let tempdir = TempDir::new()?;
 
-        let repo = Repository::shallow_clone_bare(url.clone(), tempdir.as_ref())?;
+        let repo = Repository::shallow_clone_bare(
+            url.clone(),
+            tempdir.as_ref(),
+            Some(cancellation_token),
+        )?;
 
         let config: RegistryConfig = {
             let config = repo
@@ -108,6 +112,10 @@ impl GitRegistry {
         let version_req = version_req.clone();
         let this = self.clone();
 
+        let cancellation_token = GitCancellationToken::default();
+        // Cancel git operation if the future is cancelled (dropped).
+        let cancel_on_drop = cancellation_token.clone().cancel_on_drop();
+
         let (matched_version, dl_url) = spawn_blocking(move || {
             let GitIndex {
                 _tempdir: _,
@@ -116,7 +124,7 @@ impl GitRegistry {
             } = this
                 .0
                 .git_index
-                .get_or_try_init(|| GitIndex::new(this.0.url.clone()))?;
+                .get_or_try_init(|| GitIndex::new(this.0.url.clone(), cancellation_token))?;
 
             let matched_version =
                 Self::find_crate_matched_ver(repo, &crate_name, &crate_prefix, &version_req)?;
@@ -132,6 +140,9 @@ impl GitRegistry {
         })
         .await??;
 
+        // Git operation done, disarm it
+        cancel_on_drop.disarm();
+
         parse_manifest(client, name, dl_url, matched_version).await
     }
 }

crates/binstalk/src/ops/resolve.rs

@@ -373,16 +373,26 @@ impl PackageInfo {
             }
             #[cfg(feature = "git")]
             Some(Git(git_url)) => {
+                use helpers::git::{GitCancellationToken, Repository as GitRepository};
+
                 let git_url = git_url.clone();
                 let name = name.clone();
+                let cancellation_token = GitCancellationToken::default();
+                // Cancel git operation if the future is cancelled (dropped).
+                let cancel_on_drop = cancellation_token.clone().cancel_on_drop();
 
-                spawn_blocking(move || {
+                let ret = spawn_blocking(move || {
                     let dir = TempDir::new()?;
-                    helpers::git::Repository::shallow_clone(git_url, dir.as_ref())?;
+                    GitRepository::shallow_clone(git_url, dir.as_ref(), Some(cancellation_token))?;
 
                     load_manifest_from_workspace(dir.as_ref(), &name).map_err(BinstallError::from)
                 })
-                .await??
+                .await??;
+
+                // Git operation done, disarm it
+                cancel_on_drop.disarm();
+
+                ret
             }
             None => {
                 Box::pin(