From b8a9774923308c415d3b44d707e1bdba3056f833 Mon Sep 17 00:00:00 2001
From: hunshcn <hunsh.cn@gmail.com>
Date: Tue, 14 May 2024 20:36:32 +0800
Subject: [PATCH] use errgroup

---
 gazelle/MODULE.bazel       |  1 +
 gazelle/go.mod             |  1 +
 gazelle/go.sum             |  2 ++
 gazelle/python/BUILD.bazel |  1 +
 gazelle/python/parser.go   | 50 ++++++++++++++------------------------
 5 files changed, 23 insertions(+), 32 deletions(-)

diff --git a/gazelle/MODULE.bazel b/gazelle/MODULE.bazel
index 582e1c1728..aba6508a59 100644
--- a/gazelle/MODULE.bazel
+++ b/gazelle/MODULE.bazel
@@ -20,6 +20,7 @@ use_repo(
     "com_github_smacker_go_tree_sitter",
     "com_github_stretchr_testify",
     "in_gopkg_yaml_v2",
+    "org_golang_x_sync",
 )
 
 non_module_deps = use_extension("//:def.bzl", "non_module_deps")
diff --git a/gazelle/go.mod b/gazelle/go.mod
index 824f65663c..4b65e71d67 100644
--- a/gazelle/go.mod
+++ b/gazelle/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/ghodss/yaml v1.0.0
 	github.com/smacker/go-tree-sitter v0.0.0-20240422154435-0628b34cbf9c
 	github.com/stretchr/testify v1.9.0
+	golang.org/x/sync v0.2.0
 	gopkg.in/yaml.v2 v2.4.0
 )
 
diff --git a/gazelle/go.sum b/gazelle/go.sum
index a0ee3b6c85..46e0127e8f 100644
--- a/gazelle/go.sum
+++ b/gazelle/go.sum
@@ -68,6 +68,8 @@ golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAG
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.2.0 h1:PUR+T4wwASmuSTYdKjYHI5TD22Wy5ogLU5qZCOLxBrI=
+golang.org/x/sync v0.2.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
diff --git a/gazelle/python/BUILD.bazel b/gazelle/python/BUILD.bazel
index d32217e558..75402ca943 100644
--- a/gazelle/python/BUILD.bazel
+++ b/gazelle/python/BUILD.bazel
@@ -43,6 +43,7 @@ go_library(
         "@com_github_emirpasic_gods//utils",
         "@com_github_smacker_go_tree_sitter//:go-tree-sitter",
         "@com_github_smacker_go_tree_sitter//python",
+        "@org_golang_x_sync//errgroup",
     ],
 )
 
diff --git a/gazelle/python/parser.go b/gazelle/python/parser.go
index e4b77ed13a..1b2a90dddf 100644
--- a/gazelle/python/parser.go
+++ b/gazelle/python/parser.go
@@ -18,12 +18,11 @@ import (
 	"context"
 	_ "embed"
 	"fmt"
-	"path/filepath"
 	"strings"
-	"sync"
 
 	"github.com/emirpasic/gods/sets/treeset"
 	godsutils "github.com/emirpasic/gods/utils"
+	"golang.org/x/sync/errgroup"
 )
 
 // python3Parser implements a parser for Python files that extracts the modules
@@ -64,43 +63,30 @@ func (p *python3Parser) parseSingle(pyFilename string) (*treeset.Set, map[string
 func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[string]*treeset.Set, *annotations, error) {
 	modules := treeset.NewWith(moduleComparator)
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	var wg sync.WaitGroup
+	g, ctx := errgroup.WithContext(context.Background())
 	ch := make(chan struct{}, 6) // Limit the number of concurrent parses.
 	chRes := make(chan *ParserOutput, len(pyFilenames.Values()))
-	chErr := make(chan error, 1)
 	for _, v := range pyFilenames.Values() {
-		wg.Add(1)
-		go func(filename string) {
-			defer func() {
-				<-ch
-				wg.Done()
-			}()
-			if err := ctx.Err(); err != nil {
-				return
-			}
-			ch <- struct{}{}
-			res, err := NewFileParser().ParseFile(ctx, p.repoRoot, p.relPackagePath, filename)
-			if err != nil {
-				cancel()
-				select {
-				case chErr <- fmt.Errorf("failed to parse %s: %w", filepath.Join(p.relPackagePath, filename), err):
-				default:
+		ch <- struct{}{}
+		g.Go(func(filename string) func() error {
+			return func() error {
+				defer func() {
+					<-ch
+				}()
+				res, err := NewFileParser().ParseFile(ctx, p.repoRoot, p.relPackagePath, filename)
+				if err != nil {
+					return err
 				}
-				return
+				chRes <- res
+				return nil
 			}
-			chRes <- res
-		}(v.(string))
+		}(v.(string)))
 	}
-	wg.Wait()
-	close(chErr)
-	close(chRes)
-	if len(chErr) > 0 {
-		return nil, nil, nil, <-chErr
+	if err := g.Wait(); err != nil {
+		return nil, nil, nil, err
 	}
-
+	close(ch)
+	close(chRes)
 	mainModules := make(map[string]*treeset.Set, len(chRes))
 	allAnnotations := new(annotations)
 	allAnnotations.ignore = make(map[string]struct{})