add thread-count and chunk-size computation; interrupt capability (#301)

Byron · Byron · commit 8cbe85d13589 · 2022-03-08T16:07:06.000+08:00
diff --git a/git-features/src/parallel/reduce.rs b/git-features/src/parallel/reduce.rs
@@ -10,7 +10,7 @@ mod stepped {
         receive_result: std::sync::mpsc::Receiver<Reduce::Input>,
         /// `join()` will be called on these guards to assure every thread tries to send through a closed channel. When
         /// that happens, they break out of their loops.
-        _threads: Vec<std::thread::JoinHandle<()>>,
+        threads: Vec<std::thread::JoinHandle<()>>,
         /// The reducer is called only in the thread using the iterator, dropping it has no side effects.
         reducer: Option<Reduce>,
     }
@@ -21,7 +21,7 @@ mod stepped {
             drop(std::mem::replace(&mut self.receive_result, sink));
 
             let mut last_err = None;
-            for handle in std::mem::take(&mut self._threads) {
+            for handle in std::mem::take(&mut self.threads) {
                 if let Err(err) = handle.join() {
                     last_err = Some(err);
                 };
@@ -82,7 +82,7 @@ mod stepped {
                 receive_result
             };
             Stepwise {
-                _threads: threads,
+                threads,
                 receive_result,
                 reducer: Some(reducer),
             }
diff --git a/git-worktree/src/index/checkout.rs b/git-worktree/src/index/checkout.rs
@@ -164,6 +164,10 @@ pub struct Outcome {
 pub struct Options {
     /// capabilities of the file system
     pub fs: crate::fs::Capabilities,
+    /// If set, don't use more than this amount of threads.
+    /// Otherwise, usually use as many threads as there are logical cores.
+    /// A value of 0 is interpreted as no-limit
+    pub thread_limit: Option<usize>,
     /// If true, we assume no file to exist in the target directory, and want exclusive access to it.
     /// This should be enabled when cloning to avoid checks for freshness of files. This also enables
     /// detection of collisions based on whether or not exclusive file creation succeeds or fails.
@@ -194,6 +198,7 @@ impl Default for Options {
     fn default() -> Self {
         Options {
             fs: Default::default(),
+            thread_limit: None,
             destination_is_initially_empty: false,
             keep_going: false,
             trust_ctime: true,
diff --git a/git-worktree/src/index/mod.rs b/git-worktree/src/index/mod.rs
@@ -1,6 +1,8 @@
 use bstr::BStr;
+use git_features::interrupt;
 use git_features::progress::Progress;
 use git_hash::oid;
+use std::sync::atomic::AtomicBool;
 
 use crate::index::checkout::{ErrorRecord, PathCache};
 use crate::{index, os};
@@ -14,6 +16,7 @@ pub fn checkout<Find, E>(
     find: Find,
     files: &mut impl Progress,
     bytes: &mut impl Progress,
+    should_interrupt: &AtomicBool,
     options: checkout::Options,
 ) -> Result<checkout::Outcome, checkout::Error<E>>
 where
@@ -35,8 +38,14 @@ where
         find,
         options,
     };
+    let (chunk_size, _, num_threads) = git_features::parallel::optimize_chunk_size_and_thread_limit(
+        100,
+        index.entries().len().into(),
+        options.thread_limit,
+        None,
+    );
 
-    for (entry, entry_path) in index.entries_mut_with_paths() {
+    for (entry, entry_path) in interrupt::Iter::new(index.entries_mut_with_paths(), should_interrupt) {
         // TODO: write test for that
         if entry.flags.contains(git_index::entry::Flags::SKIP_WORKTREE) {
             ctx.files.inc();
diff --git a/git-worktree/tests/index/checkout.rs b/git-worktree/tests/index/checkout.rs
@@ -109,6 +109,7 @@ use git_object::bstr::ByteSlice;
 use git_odb::FindExt;
 use git_worktree::{fs::Capabilities, index, index::checkout::Collision};
 use std::io::ErrorKind::AlreadyExists;
+use std::sync::atomic::AtomicBool;
 use tempfile::TempDir;
 
 use crate::fixture_path;
@@ -468,6 +469,7 @@ fn checkout_index_in_tmp_dir_opts(
         },
         &mut progress::Discard,
         &mut progress::Discard,
+        &AtomicBool::default(),
         opts,
     )?;
     Ok((source_tree, destination, index, outcome))
diff --git a/gitoxide-core/src/index/checkout.rs b/gitoxide-core/src/index/checkout.rs
@@ -4,13 +4,15 @@ use anyhow::bail;
 use git::{odb::FindExt, worktree::index::checkout, Progress};
 use git_repository as git;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicBool;
 
 pub fn checkout_exclusive(
     index_path: impl AsRef<Path>,
     dest_directory: impl AsRef<Path>,
     repo: Option<PathBuf>,
     mut err: impl std::io::Write,
     mut progress: impl Progress,
+    should_interrupt: &AtomicBool,
     index::checkout_exclusive::Options {
         index: Options { object_hash, .. },
         empty_files,
@@ -84,6 +86,7 @@ pub fn checkout_exclusive(
             },
             &mut files,
             &mut bytes,
+            should_interrupt,
             opts,
         ),
         None => git::worktree::index::checkout(
@@ -95,6 +98,7 @@ pub fn checkout_exclusive(
             },
             &mut files,
             &mut bytes,
+            should_interrupt,
             opts,
         ),
     }?;
diff --git a/src/plumbing/main.rs b/src/plumbing/main.rs
@@ -97,6 +97,7 @@ pub fn main() -> Result<()> {
                         repository,
                         err,
                         progress,
+                        &should_interrupt,
                         core::index::checkout_exclusive::Options {
                             index: core::index::Options { object_hash, format },
                             empty_files,
