feat: concurrent parsing for Gazelle (#630)

* feat: concurrent parsing for Gazelle

Signed-off-by: Thulio Ferraz Assis <[email protected]>

* rm: unused code

Signed-off-by: Thulio Ferraz Assis <[email protected]>

* doc: why 2 workers

Author: f0rmiga

gazelle/BUILD.bazel

@@ -55,7 +55,6 @@ go_test(
         ":gazelle_python_binary",
         ":parse",
         ":std_modules",
-        #"@python_interpreter//:bazel_install/bin/python3",
     ] + glob(["testdata/**"]),
     deps = [
         "@bazel_gazelle//testtools:go_default_library",

gazelle/fix.go

@@ -10,4 +10,4 @@ import (
 // that delete or rename rules should not be performed.
 func (py *Python) Fix(c *config.Config, f *rule.File) {
 	// TODO(f0rmiga): implement.
-}
+}

gazelle/generate.go

@@ -191,7 +191,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 
 	var pyLibrary *rule.Rule
 	if !pyLibraryFilenames.Empty() {
-		deps, err := parser.parseAll(pyLibraryFilenames)
+		deps, err := parser.parse(pyLibraryFilenames)
 		if err != nil {
 			log.Fatalf("ERROR: %v\n", err)
 		}
@@ -228,7 +228,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 	}
 
 	if hasPyBinary {
-		deps, err := parser.parse(pyBinaryEntrypointFilename)
+		deps, err := parser.parseSingle(pyBinaryEntrypointFilename)
 		if err != nil {
 			log.Fatalf("ERROR: %v\n", err)
 		}
@@ -275,7 +275,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 			// the file exists on disk.
 			pyTestFilenames.Add(pyTestEntrypointFilename)
 		}
-		deps, err := parser.parseAll(pyTestFilenames)
+		deps, err := parser.parse(pyTestFilenames)
 		if err != nil {
 			log.Fatalf("ERROR: %v\n", err)
 		}

gazelle/kinds.go

@@ -85,4 +85,4 @@ var pyLoads = []rule.LoadInfo{
 			pyTestKind,
 		},
 	},
-}
+}

gazelle/language.go

@@ -15,4 +15,4 @@ type Python struct {
 // interface. This is the entrypoint for the extension initialization.
 func NewLanguage() language.Language {
 	return &Python{}
-}
+}

gazelle/parse.py

@@ -1,15 +1,17 @@
 # parse.py is a long-living program that communicates over STDIN and STDOUT.
-# STDIN receives filepaths, one per line. For each parsed file, it outputs to
-# STDOUT the modules parsed out of the import statements.
+# STDIN receives parse requests, one per line. It outputs the parsed modules and
+# comments from all the files from each request.
 
 import ast
+import concurrent.futures
 import json
+import os
 import sys
 from io import BytesIO
 from tokenize import COMMENT, tokenize
 
 
-def parse_import_statements(content):
+def parse_import_statements(content, filepath):
     modules = list()
     tree = ast.parse(content)
     for node in ast.walk(tree):
@@ -18,12 +20,14 @@ def parse_import_statements(content):
                 module = {
                     "name": subnode.name,
                     "lineno": node.lineno,
+                    "filepath": filepath,
                 }
                 modules.append(module)
         elif isinstance(node, ast.ImportFrom) and node.level == 0:
             module = {
                 "name": node.module,
                 "lineno": node.lineno,
+                "filepath": filepath,
             }
             modules.append(module)
     return modules
@@ -38,25 +42,45 @@ def parse_comments(content):
     return comments
 
 
-def parse(stdout, filepath):
-    with open(filepath, "r") as file:
+def parse(repo_root, rel_package_path, filename):
+    rel_filepath = os.path.join(rel_package_path, filename)
+    abs_filepath = os.path.join(repo_root, rel_filepath)
+    with open(abs_filepath, "r") as file:
         content = file.read()
-        modules = parse_import_statements(content)
-        comments = parse_comments(content)
+       # From simple benchmarks, 2 workers gave the best performance here.
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+            modules_future = executor.submit(parse_import_statements, content, rel_filepath)
+            comments_future = executor.submit(parse_comments, content)
+        modules = modules_future.result()
+        comments = comments_future.result()
         output = {
             "modules": modules,
             "comments": comments,
         }
-        print(json.dumps(output), end="", file=stdout)
-        stdout.flush()
-        stdout.buffer.write(bytes([0]))
-        stdout.flush()
+        return output
 
 
 def main(stdin, stdout):
-    for filepath in stdin:
-        filepath = filepath.rstrip()
-        parse(stdout, filepath)
+    with concurrent.futures.ProcessPoolExecutor() as executor:
+        for parse_request in stdin:
+            parse_request = json.loads(parse_request)
+            repo_root = parse_request["repo_root"]
+            rel_package_path = parse_request["rel_package_path"]
+            filenames = parse_request["filenames"]
+            outputs = list()
+            if len(filenames) == 1:
+                outputs.append(parse(repo_root, rel_package_path, filenames[0]))
+            else:
+                futures = [
+                    executor.submit(parse, repo_root, rel_package_path, filename)
+                    for filename in filenames
+                    if filename != ""
+                ]
+                for future in concurrent.futures.as_completed(futures):
+                    outputs.append(future.result())
+            print(json.dumps(outputs), end="", file=stdout, flush=True)
+            stdout.buffer.write(bytes([0]))
+            stdout.flush()
 
 
 if __name__ == "__main__":

gazelle/parser.go

@@ -9,7 +9,6 @@ import (
 	"log"
 	"os"
 	"os/exec"
-	"path/filepath"
 	"strings"
 	"sync"
 	"time"
@@ -91,64 +90,61 @@ func newPython3Parser(
 	}
 }
 
-// parseAll parses all provided Python files by consecutively calling p.parse.
-func (p *python3Parser) parseAll(pyFilepaths *treeset.Set) (*treeset.Set, error) {
-	allModules := treeset.NewWith(moduleComparator)
-	it := pyFilepaths.Iterator()
-	for it.Next() {
-		modules, err := p.parse(it.Value().(string))
-		if err != nil {
-			return nil, err
-		}
-		modulesIt := modules.Iterator()
-		for modulesIt.Next() {
-			allModules.Add(modulesIt.Value())
-		}
-	}
-	return allModules, nil
+// parseSingle parses a single Python file and returns the extracted modules
+// from the import statements as well as the parsed comments.
+func (p *python3Parser) parseSingle(pyFilename string) (*treeset.Set, error) {
+	pyFilenames := treeset.NewWith(godsutils.StringComparator)
+	pyFilenames.Add(pyFilename)
+	return p.parse(pyFilenames)
 }
 
-// parse parses a Python file and returns the extracted modules from the import
-// statements. An error is raised if communicating with the long-lived Python
-// parser over stdin and stdout fails.
-func (p *python3Parser) parse(pyFilepath string) (*treeset.Set, error) {
+// parse parses multiple Python files and returns the extracted modules from
+// the import statements as well as the parsed comments.
+func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, error) {
 	parserMutex.Lock()
 	defer parserMutex.Unlock()
 
 	modules := treeset.NewWith(moduleComparator)
 
-	relFilepath := filepath.Join(p.relPackagePath, pyFilepath)
-	absFilepath := filepath.Join(p.repoRoot, relFilepath)
-	fmt.Fprintln(parserStdin, absFilepath)
+	req := map[string]interface{}{
+		"repo_root":        p.repoRoot,
+		"rel_package_path": p.relPackagePath,
+		"filenames":        pyFilenames.Values(),
+	}
+	encoder := json.NewEncoder(parserStdin)
+	if err := encoder.Encode(&req); err != nil {
+		return nil, fmt.Errorf("failed to parse: %w", err)
+	}
+
 	reader := bufio.NewReader(parserStdout)
 	data, err := reader.ReadBytes(0)
 	if err != nil {
-		return nil, fmt.Errorf("failed to parse %s: %w", pyFilepath, err)
+		return nil, fmt.Errorf("failed to parse: %w", err)
 	}
 	data = data[:len(data)-1]
-	var res parserResponse
-	if err := json.Unmarshal(data, &res); err != nil {
-		return nil, fmt.Errorf("failed to parse %s: %w", pyFilepath, err)
+	var allRes []parserResponse
+	if err := json.Unmarshal(data, &allRes); err != nil {
+		return nil, fmt.Errorf("failed to parse: %w", err)
 	}
 
-	annotations := annotationsFromComments(res.Comments)
-
-	for _, m := range res.Modules {
-		// Check for ignored dependencies set via an annotation to the Python
-		// module.
-		if annotations.ignores(m.Name) {
-			continue
-		}
+	for _, res := range allRes {
+		annotations := annotationsFromComments(res.Comments)
 
-		// Check for ignored dependencies set via a Gazelle directive in a BUILD
-		// file.
-		if p.ignoresDependency(m.Name) {
-			continue
-		}
+		for _, m := range res.Modules {
+			// Check for ignored dependencies set via an annotation to the Python
+			// module.
+			if annotations.ignores(m.Name) {
+				continue
+			}
 
-		m.Filepath = relFilepath
+			// Check for ignored dependencies set via a Gazelle directive in a BUILD
+			// file.
+			if p.ignoresDependency(m.Name) {
+				continue
+			}
 
-		modules.Add(m)
+			modules.Add(m)
+		}
 	}
 
 	return modules, nil
@@ -173,17 +169,7 @@ type module struct {
 	// The line number where the import happened.
 	LineNumber uint32 `json:"lineno"`
 	// The path to the module file relative to the Bazel workspace root.
-	Filepath string
-}
-
-// path returns the replaced dots with the os-specific path separator.
-func (m *module) path() string {
-	return filepath.Join(strings.Split(m.Name, ".")...)
-}
-
-// bazelPath returns the replaced dots with forward slashes.
-func (m *module) bazelPath() string {
-	return strings.ReplaceAll(m.Name, ".", "/")
+	Filepath string `json:"filepath"`
 }
 
 // moduleComparator compares modules by name.

gazelle/python_test.go

@@ -208,4 +208,4 @@ type testYAML struct {
 		Stdout   string `json:"stdout"`
 		Stderr   string `json:"stderr"`
 	} `json:"expect"`
-}
+}

gazelle/std_modules.go

@@ -95,4 +95,4 @@ func isStdModule(m module) (bool, error) {
 		return true, nil
 	}
 	return false, nil
-}
+}

gazelle/target.go

@@ -133,4 +133,4 @@ func (t *targetBuilder) build() *rule.Rule {
 	}
 	r.SetPrivateAttr(resolvedDepsKey, t.resolvedDeps)
 	return r
-}
+}