diff --git a/Cargo.lock b/Cargo.lock index 5b29eca5..058ca20d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -290,6 +290,7 @@ dependencies = [ "bytes", "bzip2", "compact_str", + "derive_destructure2", "flate2", "futures-util", "generic-array", @@ -756,6 +757,17 @@ version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7684a49fb1af197853ef7b2ee694bc1f5b4179556f1e5710e1760c5db6f5e929" +[[package]] +name = "derive_destructure2" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35cb7e5875e1028a73e551747d6d0118f25c3d6dbba2dadf97cc0f4d0c53f2f5" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "detect-targets" version = "0.1.10" diff --git a/crates/binstalk-downloader/Cargo.toml b/crates/binstalk-downloader/Cargo.toml index bf54fcdd..8036aaba 100644 --- a/crates/binstalk-downloader/Cargo.toml +++ b/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" diff --git a/crates/binstalk-downloader/src/git.rs b/crates/binstalk-downloader/src/git.rs index 2f9ba0a8..8ff64e39 100644 --- a/crates/binstalk-downloader/src/git.rs +++ b/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 { + pub fn shallow_clone_bare( + url: GitUrl, + path: &Path, + cancellation_token: Option, + ) -> Result { 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 { + pub fn shallow_clone( + url: GitUrl, + path: &Path, + cancellation_token: Option, + ) -> Result { 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(), )) diff --git a/crates/binstalk-downloader/src/git/cancellation_token.rs b/crates/binstalk-downloader/src/git/cancellation_token.rs new file mode 100644 index 00000000..97e65e52 --- /dev/null +++ b/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); + +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 + } +} diff --git a/crates/binstalk-downloader/src/git/progress_tracing.rs b/crates/binstalk-downloader/src/git/progress_tracing.rs index 67ebe562..b88af701 100644 --- a/crates/binstalk-downloader/src/git/progress_tracing.rs +++ b/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, diff --git a/crates/binstalk/src/drivers/registry/git_registry.rs b/crates/binstalk/src/drivers/registry/git_registry.rs index 3f98a3ea..c72b504a 100644 --- a/crates/binstalk/src/drivers/registry/git_registry.rs +++ b/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 { + fn new(url: GitUrl, cancellation_token: GitCancellationToken) -> Result { 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 } } diff --git a/crates/binstalk/src/ops/resolve.rs b/crates/binstalk/src/ops/resolve.rs index f9f031cb..fdf3414a 100644 --- a/crates/binstalk/src/ops/resolve.rs +++ b/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(