Merge pull request #83 from github/threads

nickrolfe · web-flow · commit e2453820575d · 2021-01-07T17:14:41.000Z
Parallelize extraction
diff --git a/.github/workflows/dataset_measure.yml b/.github/workflows/dataset_measure.yml
@@ -27,6 +27,7 @@ jobs:
           unzip -q codeql-linux64.zip
         env:
           GITHUB_TOKEN: ${{ github.token }}
+          CODEQL_THREADS: 4 # TODO: remove this once it's set by the CLI
       - uses: actions/cache@v2
         with:
           path: |
@@ -46,6 +47,7 @@ jobs:
         run: |
           codeql/codeql database create \
             --search-path "${{ github.workspace }}" \
+            --threads 4 \
             --language ruby --source-root "${{ github.workspace }}/repo" \
             "${{ runner.temp }}/database"
       - name: Measure database
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/extractor/Cargo.toml b/extractor/Cargo.toml
@@ -14,3 +14,5 @@ tree-sitter-ruby = { git = "https://github.com/tree-sitter/tree-sitter-ruby.git"
 clap = "2.33"
 tracing = "0.1"
 tracing-subscriber = { version = "0.2", features = ["env-filter"] }
+rayon = "1.5.0"
+num_cpus = "1.13.0"
diff --git a/extractor/src/extractor.rs b/extractor/src/extractor.rs
@@ -148,55 +148,40 @@ impl TrapWriter {
     }
 }
 
-pub struct Extractor {
-    pub parser: Parser,
-    pub schema: NodeTypeMap,
-}
-
-pub fn create(language: Language, schema: NodeTypeMap) -> Extractor {
-    let mut parser = Parser::new();
-    parser.set_language(language).unwrap();
+/// Extracts the source file at `path`, which is assumed to be canonicalized.
+pub fn extract(language: Language, schema: &NodeTypeMap, path: &Path) -> std::io::Result<Program> {
+    let span = span!(
+        Level::TRACE,
+        "extract",
+        file = %path.display()
+    );
 
-    Extractor { parser, schema }
-}
+    let _enter = span.enter();
 
-impl Extractor {
-    /// Extracts the source file at `path`, which is assumed to be canonicalized.
-    pub fn extract<'a>(&'a mut self, path: &Path) -> std::io::Result<Program> {
-        let span = span!(
-            Level::TRACE,
-            "extract",
-            file = %path.display()
-        );
+    info!("extracting: {}", path.display());
 
-        let _enter = span.enter();
-
-        info!("extracting: {}", path.display());
-
-        let source = std::fs::read(&path)?;
-        let tree = &self
-            .parser
-            .parse(&source, None)
-            .expect("Failed to parse file");
-        let mut trap_writer = new_trap_writer();
-        trap_writer.comment(format!("Auto-generated TRAP file for {}", path.display()));
-        let file_label = &trap_writer.populate_file(path);
-        let mut visitor = Visitor {
-            source: &source,
-            trap_writer: trap_writer,
-            // TODO: should we handle path strings that are not valid UTF8 better?
-            path: format!("{}", path.display()),
-            file_label: *file_label,
-            token_counter: 0,
-            toplevel_child_counter: 0,
-            stack: Vec::new(),
-            schema: &self.schema,
-        };
-        traverse(&tree, &mut visitor);
-
-        &self.parser.reset();
-        Ok(Program(visitor.trap_writer.trap_output))
-    }
+    let mut parser = Parser::new();
+    parser.set_language(language).unwrap();
+    let source = std::fs::read(&path)?;
+    let tree = parser.parse(&source, None).expect("Failed to parse file");
+    let mut trap_writer = new_trap_writer();
+    trap_writer.comment(format!("Auto-generated TRAP file for {}", path.display()));
+    let file_label = &trap_writer.populate_file(path);
+    let mut visitor = Visitor {
+        source: &source,
+        trap_writer: trap_writer,
+        // TODO: should we handle path strings that are not valid UTF8 better?
+        path: format!("{}", path.display()),
+        file_label: *file_label,
+        token_counter: 0,
+        toplevel_child_counter: 0,
+        stack: Vec::new(),
+        schema,
+    };
+    traverse(&tree, &mut visitor);
+
+    parser.reset();
+    Ok(Program(visitor.trap_writer.trap_output))
 }
 
 /// Normalizes the path according the common CodeQL specification. Assumes that
diff --git a/extractor/src/main.rs b/extractor/src/main.rs
@@ -1,7 +1,10 @@
 mod extractor;
 
+extern crate num_cpus;
+
 use clap;
 use flate2::write::GzEncoder;
+use rayon::prelude::*;
 use std::fs;
 use std::io::{BufRead, BufWriter, Write};
 use std::path::{Path, PathBuf};
@@ -42,6 +45,39 @@ impl TrapCompression {
     }
 }
 
+/**
+ * Gets the number of threads the extractor should use, by reading the
+ * CODEQL_THREADS environment variable and using it as described in the
+ * extractor spec:
+ *
+ * "If the number is positive, it indicates the number of threads that should
+ * be used. If the number is negative or zero, it should be added to the number
+ * of cores available on the machine to determine how many threads to use
+ * (minimum of 1). If unspecified, should be considered as set to 1."
+ */
+fn num_codeql_threads() -> usize {
+    match std::env::var("CODEQL_THREADS") {
+        // Use 1 thread if the environment variable isn't set.
+        Err(_) => 1,
+
+        Ok(num) => match num.parse::<i32>() {
+            Ok(num) if num <= 0 => {
+                let reduction = -num as usize;
+                num_cpus::get() - reduction
+            }
+            Ok(num) => num as usize,
+
+            Err(_) => {
+                tracing::error!(
+                    "Unable to parse CODEQL_THREADS value '{}'; defaulting to 1 thread.",
+                    &num
+                );
+                1
+            }
+        },
+    }
+}
+
 fn main() -> std::io::Result<()> {
     tracing_subscriber::fmt()
         .with_target(false)
@@ -50,6 +86,21 @@ fn main() -> std::io::Result<()> {
         .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
         .init();
 
+    let num_threads = num_codeql_threads();
+    tracing::info!(
+        "Using {} {}",
+        num_threads,
+        if num_threads == 1 {
+            "thread"
+        } else {
+            "threads"
+        }
+    );
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(num_threads)
+        .build_global()
+        .unwrap();
+
     let matches = clap::App::new("Ruby extractor")
         .version("1.0")
         .author("GitHub")
@@ -76,28 +127,28 @@ fn main() -> std::io::Result<()> {
 
     let language = tree_sitter_ruby::language();
     let schema = node_types::read_node_types_str(tree_sitter_ruby::NODE_TYPES)?;
-    let mut extractor = extractor::create(language, schema);
-    for line in std::io::BufReader::new(file_list).lines() {
-        let path = PathBuf::from(line?).canonicalize()?;
+    let lines: std::io::Result<Vec<String>> = std::io::BufReader::new(file_list).lines().collect();
+    let lines = lines?;
+    lines.par_iter().try_for_each(|line| {
+        let path = PathBuf::from(line).canonicalize()?;
         let trap_file = path_for(&trap_dir, &path, trap_compression.extension());
         let src_archive_file = path_for(&src_archive_dir, &path, "");
-        let trap = extractor.extract(&path)?;
+        let trap = extractor::extract(language, &schema, &path)?;
         std::fs::create_dir_all(&src_archive_file.parent().unwrap())?;
         std::fs::copy(&path, &src_archive_file)?;
         std::fs::create_dir_all(&trap_file.parent().unwrap())?;
         let trap_file = std::fs::File::create(&trap_file)?;
         let mut trap_file = BufWriter::new(trap_file);
         match trap_compression {
             TrapCompression::None => {
-                write!(trap_file, "{}", trap)?;
+                write!(trap_file, "{}", trap)
             }
             TrapCompression::Gzip => {
                 let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
-                write!(compressed_writer, "{}", trap)?;
+                write!(compressed_writer, "{}", trap)
             }
         }
-    }
-    return Ok(());
+    })
 }
 
 fn path_for(dir: &Path, path: &Path, ext: &str) -> PathBuf {
