perf: do api queries in parallel using errgroup
#347
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -12,6 +12,7 @@ import ( | |
| "path" | ||
|
|
||
| "github.com/g-rath/osv-detector/internal" | ||
| "golang.org/x/sync/errgroup" | ||
| ) | ||
|
|
||
| func (db APIDB) buildAPIPayload(pkg internal.PackageDetails) apiQuery { | ||
|
|
@@ -171,15 +172,38 @@ func findOrDefault(vulns Vulnerabilities, def OSV) OSV { | |
| func (db APIDB) Check(pkgs []internal.PackageDetails) ([]Vulnerabilities, error) { | ||
| batches := batchPkgs(pkgs, db.BatchSize) | ||
|
|
||
| var eg errgroup.Group | ||
|
|
||
| // use a sensible upper limit so it's not possible to have inf. operations going | ||
| // even though it's very unlikely there will be more than a couple of batches | ||
| eg.SetLimit(100) | ||
|
|
||
| batchResults := make([][][]ObjectWithID, len(batches)) | ||
|
|
||
| for i, batch := range batches { | ||
| eg.Go(func() error { | ||
| results, err := db.checkBatch(batch) | ||
|
|
||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| batchResults[i] = results | ||
|
Comment on lines
+184
to
+191
There was a problem hiding this comment. Race condition: closure captures loop variable incorrectly The closure passed to This will cause multiple goroutines to process the same batch and write to the same index in Solution: Capture the loop variables before the closure: for i, batch := range batches {
i, batch := i, batch // capture loop variables
eg.Go(func() error {
results, err := db.checkBatch(batch)
if err != nil {
return err
}
batchResults[i] = results
return nil
})
} |
||
|
|
||
| return nil | ||
| }) | ||
| } | ||
|
|
||
| err := eg.Wait() | ||
|
|
||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| vulnerabilities := make([]Vulnerabilities, 0, len(pkgs)) | ||
|
|
||
| // todo: pretty sure some of these loops and slices can be merged and simplified | ||
| for _, results := range batchResults { | ||
| for _, withIDs := range results { | ||
| vulns := make(Vulnerabilities, 0, len(withIDs)) | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -553,11 +553,15 @@ func TestAPIDB_Check_Batches(t *testing.T) { | |
| mux.HandleFunc("/querybatch", func(w http.ResponseWriter, r *http.Request) { | ||
| requestCount++ | ||
|
There was a problem hiding this comment. Race condition in test: concurrent writes to With parallel batch execution, multiple goroutines may increment Solution: Use var requestCount atomic.Int32
mux.HandleFunc("/querybatch", func(w http.ResponseWriter, r *http.Request) {
count := requestCount.Add(1)
if count > 2 {
t.Errorf("unexpected number of requests (%d)", count)
}
// ... rest of the handlerAnd at line 611: if requestCount.Load() != 2 {
Owner
Author
There was a problem hiding this comment. I'm fine with this as it should be very low risk (its not been an issue so far) |
||
|
|
||
| if requestCount > 2 { | ||
| t.Errorf("unexpected number of requests (%d)", requestCount) | ||
| } | ||
|
|
||
| var expectedPayload []apiQuery | ||
| var batchResponse []objectsWithIDs | ||
|
|
||
| // strictly speaking not the best of checks, but it should be good enough | ||
| if r.ContentLength > 100 { | ||
| expectedPayload = []apiQuery{ | ||
| { | ||
| Version: "1.0.0", | ||
|
|
@@ -569,16 +573,14 @@ func TestAPIDB_Check_Batches(t *testing.T) { | |
| }, | ||
| } | ||
| batchResponse = []objectsWithIDs{{}, {}} | ||
| } else if r.ContentLength > 50 { | ||
| expectedPayload = []apiQuery{ | ||
| { | ||
| Version: "2.3.1", | ||
| Package: apiPackage{Name: "their-package", Ecosystem: lockfile.NpmEcosystem}, | ||
| }, | ||
| } | ||
| batchResponse = []objectsWithIDs{{}} | ||
| } | ||
|
|
||
| expectRequestPayload(t, r, expectedPayload) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,65 +2,32 @@ package database | |
|
|
||
| import ( | ||
| "sort" | ||
|
|
||
| "golang.org/x/sync/errgroup" | ||
| ) | ||
|
|
||
| func (db APIDB) FetchAll(ids []string) Vulnerabilities { | ||
| var eg errgroup.Group | ||
|
|
||
| eg.SetLimit(200) | ||
|
|
||
| osvs := make(Vulnerabilities, len(ids)) | ||
|
|
||
| for i, id := range ids { | ||
| eg.Go(func() error { | ||
| // if we error, still report the vulnerability as hopefully the ID should be | ||
| // enough to manually look up the details - in future we should ideally warn | ||
| // the user too, but for now we just silently eat the error | ||
| osv, _ := db.Fetch(id) | ||
|
|
||
| osvs[i] = osv | ||
|
Comment on lines
+17
to
+23
There was a problem hiding this comment. Race condition: closure captures loop variable incorrectly The closure passed to This will result in multiple goroutines writing to the same index in the Solution: Capture the loop variables by passing them as parameters to an immediately invoked function, or assign them to local variables before the closure: for i, id := range ids {
i, id := i, id // capture loop variables
eg.Go(func() error {
osv, _ := db.Fetch(id)
osvs[i] = osv
return nil
})
}
Owner
Author
There was a problem hiding this comment. I'm pretty sure this is incorrect since Go 1.22 |
||
|
|
||
| return nil | ||
| }) | ||
| } | ||
|
|
||
| // errors are handled within the go routines | ||
| _ = eg.Wait() | ||
|
|
||
| sort.Slice(osvs, func(i, j int) bool { | ||
| return osvs[i].ID < osvs[j].ID | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Missing concurrency limit for errgroup
Unlike
FetchAll()which sets a limit of 200 concurrent operations witheg.SetLimit(200)(line 12 in api-fetch-all.go), this errgroup has no limit. This could potentially spawn a large number of goroutines if there are many batches, which may overwhelm system resources or the API server.Suggestion: Consider adding a concurrency limit consistent with the existing pattern: