Fix concurrency bugs

Author: huangsam

pkg/dns/search.go

@@ -9,7 +9,6 @@ import (
 )
 
 func SearchByProbe(name string, max int) ([]model.DNSResult, error) {
-	var wg sync.WaitGroup
 	domains := []string{"com", "org", "net", "io", "tech", "ai", "me", "shop"}
 	domainChan := make(chan string)
 
@@ -21,8 +20,10 @@ func SearchByProbe(name string, max int) ([]model.DNSResult, error) {
 	}()
 
 	result := []model.DNSResult{}
+	resultCount := 0
 	errorCount := 0
 	var mu sync.Mutex
+	var wg sync.WaitGroup
 
 	for i := 0; i < 4; i++ {
 		wg.Add(1)
@@ -35,11 +36,14 @@ func SearchByProbe(name string, max int) ([]model.DNSResult, error) {
 					errorCount++
 				}
 				mu.Lock()
-				result = append(result, model.DNSResult{FQDN: fullDomain, IPList: ips})
+				if resultCount < max {
+					result = append(result, model.DNSResult{FQDN: fullDomain, IPList: ips})
+					resultCount++
+				}
+				mu.Unlock()
 				if len(result) >= max {
 					return
 				}
-				mu.Unlock()
 			}
 		}()
 	}

pkg/pypi/search.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/huangsam/namigo/internal/model"
@@ -35,30 +36,45 @@ func SearchByAPI(name string, max int) ([]model.PyPIPackageResult, error) {
 	}()
 
 	result := []model.PyPIPackageResult{}
-	count := 0
-	for pkg := range taskChan {
-		bd, err := util.RESTAPIQuery(client, detail(pkg))
-		if err != nil {
-			continue
-		}
-		var detailRes PypiDetailResponse
-		if err := json.Unmarshal(bd, &detailRes); err != nil {
-			continue
-		}
-		description := detailRes.Info.Summary
-		if len(description) == 0 {
-			description = model.NoDescription
-		}
-		author := detailRes.Info.Author
-		if len(author) == 0 {
-			author = model.NoAuthor
-		}
-		result = append(result, model.PyPIPackageResult{Name: pkg, Description: description, Author: author})
-		count++
-		if count >= max {
-			break
-		}
+	resultCount := 0
+	var mu sync.Mutex
+	var wg sync.WaitGroup
+
+	for i := 0; i < 4; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for pkg := range taskChan {
+				bd, err := util.RESTAPIQuery(client, detail(pkg))
+				if err != nil {
+					continue
+				}
+				var detailRes PypiDetailResponse
+				if err := json.Unmarshal(bd, &detailRes); err != nil {
+					continue
+				}
+				description := detailRes.Info.Summary
+				if len(description) == 0 {
+					description = model.NoDescription
+				}
+				author := detailRes.Info.Author
+				if len(author) == 0 {
+					author = model.NoAuthor
+				}
+				mu.Lock()
+				if resultCount < max {
+					result = append(result, model.PyPIPackageResult{Name: pkg, Description: description, Author: author})
+					resultCount++
+				}
+				mu.Unlock()
+				if resultCount >= max {
+					break
+				}
+			}
+		}()
 	}
 
+	wg.Wait()
+
 	return result, nil
 }