improved multithreading performance. scan methods are now interruptable

Author: Gal Ben David

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pyrepscan"
-version = "0.8.0"
+version = "0.9.0"
 authors = ["Gal Ben David <[email protected]>"]
 edition = "2018"
 description = "A Git Repository Secrets Scanner written in Rust"
@@ -36,10 +36,11 @@ crate-type = ["cdylib"]
 
 [dependencies]
 regex = "1"
-rayon = "1.5"
 chrono = "0.4"
-parking_lot = "0.11"
 num_cpus = "1"
+parking_lot = "0.11"
+crossbeam = "0.8"
+crossbeam-utils = "0.8"
 
 [dependencies.git2]
 version = "0.13"

pyproject.toml

@@ -14,7 +14,7 @@ strip = true
 
 [tool.poetry]
 name = "pyrepscan"
-version = "0.8.0"
+version = "0.9.0"
 authors = ["Gal Ben David <[email protected]>"]
 description = "A Git Repository Secrets Scanner written in Rust"
 readme = "README.md"

src/git_repository_scanner.rs

@@ -1,20 +1,25 @@
 use crate::rules_manager;
 
 use chrono::prelude::*;
+use crossbeam_utils::thread as crossbeam_thread;
+use crossbeam::queue::SegQueue;
 use git2::{Oid, Repository, Delta};
-use git2::Error;
 use parking_lot::Mutex;
-use rayon::prelude::*;
+use pyo3::prelude::*;
 use std::collections::HashMap;
 use std::path::Path;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::thread;
+use std::time;
 
 fn scan_commit_oid(
+    should_stop: &AtomicBool,
     git_repo: &Repository,
     oid: &Oid,
     rules_manager: &rules_manager::RulesManager,
     output_matches: Arc<Mutex<Vec<HashMap<&str, String>>>>,
-) -> Result<(), Error> {
+) -> Result<(), git2::Error> {
     let commit = git_repo.find_commit(*oid)?;
 
     let commit_parent_count = commit.parent_count();
@@ -34,6 +39,10 @@ fn scan_commit_oid(
     };
 
     for delta in commit_diff.deltas() {
+        if should_stop.load(Ordering::Relaxed) {
+            break;
+        }
+
         match delta.status() {
             Delta::Added | Delta::Modified => {},
             _ => continue,
@@ -119,25 +128,22 @@ fn scan_commit_oid(
 pub fn get_file_content(
     repository_path: &str,
     file_oid: &str,
-) -> Result<Vec<u8>, Error> {
+) -> Result<Vec<u8>, git2::Error> {
     let git_repo = Repository::open(repository_path)?;
     let oid = Oid::from_str(file_oid)?;
     let blob = git_repo.find_blob(oid)?;
 
     Ok(blob.content().to_vec())
 }
 
-pub fn scan_repository(
+fn get_oids(
     repository_path: &str,
     branch_glob_pattern: &str,
     from_timestamp: i64,
-    rules_manager: &rules_manager::RulesManager,
-    output_matches: Arc<Mutex<Vec<HashMap<&str, String>>>>,
-) -> Result<(), Error> {
+) -> Result<Vec<Oid>, git2::Error>{
     let git_repo = Repository::open(repository_path)?;
 
     let mut revwalk = git_repo.revwalk()?;
-
     revwalk.push_head()?;
     revwalk.set_sorting(git2::Sort::TIME)?;
     revwalk.push_glob(branch_glob_pattern)?;
@@ -153,22 +159,79 @@ pub fn scan_repository(
         }
     }
 
-    let chunk_size = (oids.len() as f64 / (num_cpus::get() * 5) as f64).ceil() as usize;
-    if !oids.is_empty() {
-        oids.par_chunks(chunk_size).for_each(
-            |oids| {
-                let git_repo = Repository::open(repository_path).unwrap();
-                for oid in oids {
-                    scan_commit_oid(
-                        &git_repo,
-                        oid,
-                        rules_manager,
-                        output_matches.clone()
-                    ).unwrap_or(());
-                }
-            },
-        );
+    Ok(oids)
+}
+
+pub fn scan_repository(
+    py: &Python,
+    repository_path: &str,
+    branch_glob_pattern: &str,
+    from_timestamp: i64,
+    rules_manager: &rules_manager::RulesManager,
+    output_matches: Arc<Mutex<Vec<HashMap<&str, String>>>>,
+) -> Result<(), PyErr> {
+    let oids_queue = Arc::new(SegQueue::new());
+    match get_oids(
+        repository_path,
+        branch_glob_pattern,
+        from_timestamp
+    ) {
+        Ok(oids) => {
+            for oid in oids {
+                oids_queue.push(oid);
+            }
+        },
+        Err(error) => {
+            return Err(pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))
+        },
     }
+    py.check_signals()?;
+
+    let mut py_signal_error: PyResult<()> = Ok(());
+
+    crossbeam_thread::scope(
+        |scope| {
+            let should_stop  = Arc::new(AtomicBool::new(false));
+
+            for _ in 0..num_cpus::get() {
+                let output_matches = output_matches.clone();
+                let oids_queue = oids_queue.clone();
+                let should_stop = should_stop.clone();
+                scope.spawn(
+                    move |_| {
+                        if let Ok(git_repo) = Repository::open(repository_path) {
+                            while !should_stop.load(Ordering::Relaxed) {
+                                if let Some(oid) = oids_queue.pop() {
+                                    scan_commit_oid(
+                                        &should_stop,
+                                        &git_repo,
+                                        &oid,
+                                        rules_manager,
+                                        output_matches.clone(),
+                                    ).unwrap_or(());
+                                } else {
+                                    break;
+                                }
+                            }
+                        };
+                    }
+                );
+            }
+
+            while !oids_queue.is_empty() {
+                py_signal_error = py.check_signals();
+                if py_signal_error.is_err() {
+                    should_stop.store(true, Ordering::Relaxed);
+
+                    break;
+                }
+
+                thread::sleep(time::Duration::from_millis(100));
+            }
+        }
+    ).unwrap_or(());
+
+    py_signal_error?;
 
     Ok(())
 }

src/lib.rs

@@ -1,12 +1,12 @@
 mod git_repository_scanner;
 mod rules_manager;
 
-use git2::Repository;
 use parking_lot::Mutex;
 use pyo3::exceptions;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use std::collections::HashMap;
+use std::path::PathBuf;
 use std::sync::Arc;
 
 /// GitRepositoryScanner class
@@ -201,17 +201,16 @@ impl GitRepositoryScanner {
         from_timestamp: Option<i64>,
     ) -> PyResult<Py<PyAny>> {
         let matches = Arc::new(Mutex::new(Vec::<HashMap<&str, String>>::with_capacity(10000)));
-
-        if let Err(error) = git_repository_scanner::scan_repository(
+        match git_repository_scanner::scan_repository(
+            &py,
             repository_path,
             branch_glob_pattern.unwrap_or("*"),
             from_timestamp.unwrap_or(0),
             &self.rules_manager,
             matches.clone(),
         ) {
-            Err(exceptions::PyRuntimeError::new_err(error.to_string()))
-        } else {
-            Ok(matches.lock().to_object(py))
+            Ok(_) => Ok(matches.lock().to_object(py)),
+            Err(error) => Err(error),
         }
     }
 
@@ -242,22 +241,26 @@ impl GitRepositoryScanner {
         branch_glob_pattern: Option<&str>,
         from_timestamp: Option<i64>,
     ) -> PyResult<Py<PyAny>> {
-        let matches = Arc::new(Mutex::new(Vec::<HashMap<&str, String>>::with_capacity(10000)));
+        let mut repository_full_path = PathBuf::from(repository_path);
+        repository_full_path.push(url.split('/').last().unwrap_or(""));
 
-        if let Err(error) = Repository::clone(url, repository_path) {
+        let mut builder = git2::build::RepoBuilder::new();
+        builder.bare(true);
+        if let Err(error) = builder.clone(url, repository_full_path.as_path()) {
             return Err(exceptions::PyRuntimeError::new_err(error.to_string()));
         };
 
-        if let Err(error) = git_repository_scanner::scan_repository(
+        let matches = Arc::new(Mutex::new(Vec::<HashMap<&str, String>>::with_capacity(10000)));
+        match git_repository_scanner::scan_repository(
+            &py,
             repository_path,
             branch_glob_pattern.unwrap_or("*"),
             from_timestamp.unwrap_or(0),
             &self.rules_manager,
             matches.clone(),
         ) {
-            Err(exceptions::PyRuntimeError::new_err(error.to_string()))
-        } else {
-            Ok(matches.lock().to_object(py))
+            Ok(_) => Ok(matches.lock().to_object(py)),
+            Err(error) => Err(error),
         }
     }
 }