feat: Exporter 12x faster hopefully (#5691)

Made a few changes to vastly increase the speed of the exporter.

- Increased each channel buffer up to 100 from 0.
- 5x ed the worker count from 600 to 3000
- Added a fast object list GCS library (mostly copied over from
vulnfeeds), and split up the initial object listing loop at various
places.
- This is mostly guesswork at the moment, ideally we should do a review
of where to roughly evenly split the work every now and then. But good
enough for now.
  - This does mean the exporter will slow down near the end.
- Refactored message receiving to use the `for` syntax. This should be a
slight behavior change where we won't cancel based on the context, but
only when the channel is closed. (Or when the context is cancelled on
the sending side, just never on the receiving side)
- Fixed README typos
- Use logger.ErrorContext instead of logger.Error
-

Author: another-rex

deployment/clouddeploy/gke-workers/environments/oss-vdb-test/exporter.yaml

@@ -15,7 +15,9 @@ spec:
             args:
               # TODO(michaelkedar): single source of truth w/ terraform config
               - "--upload-to-gcs=true"
-              - "--workers=600"
+              - "--workers=3000"
               - "--bucket=osv-test-vulnerabilities"
               - "--osv-vulns-bucket=osv-test-vulnerabilities"
+              - "--breakdown-prefixes=ALSA-,BIT-,CGA-0,CGA-a,CGA-m,CVE-,DLA-,DSA-,GHSA-,GO-,MAL-,MGASA-,OSV-,PYSEC-,RLSA-,RUSTSEC-,SUSE-,UBUNTU-"
+
 

deployment/clouddeploy/gke-workers/environments/oss-vdb/exporter.yaml

@@ -14,7 +14,9 @@ spec:
               value: oss-vdb
             args:
               - "--upload-to-gcs=true"
-              - "--workers=600"
+              - "--workers=3000"
               - "--bucket=osv-vulnerabilities"
               - "--osv-vulns-bucket=osv-vulnerabilities"
+              - "--breakdown-prefixes=ALSA-,BIT-,CGA-0,CGA-a,CGA-m,CVE-,DLA-,DSA-,GHSA-,GO-,MAL-,MGASA-,OSV-,PYSEC-,RLSA-,RUSTSEC-,SUSE-,UBUNTU-"
+
 

deployment/clouddeploy/gke-workers/environments/private/exporter.yaml

@@ -14,7 +14,7 @@ spec:
               value: oss-vdb-test
             args:
               - "--upload-to-gcs=true"
-              - "--workers=600"
+              - "--workers=3000"
               - "--bucket=osv-test-vulnerabilities-private"
               - "--osv-vulns-bucket=osv-test-vulnerabilities-private"
 

go/cmd/exporter/README.md

@@ -17,7 +17,7 @@ To run the exporter locally, run the exporter from within the `go/cmd/exporter`
 
 ```sh
 # Example
-go run . -osv_vulns_bucket osv-test-vulnerabilities -uploadToGCS=false -bucket /tmp/osv-export
+go run . -osv-vulns-bucket osv-test-vulnerabilities -upload-to-gcs=false -bucket /tmp/osv-export
 ```
 
 This will write the exported files to the `/tmp/osv-export` directory.

go/cmd/exporter/downloader.go

@@ -5,6 +5,7 @@ import (
 	"log/slog"
 	"sync"
 
+	"github.com/google/osv.dev/go/internal/osvutil"
 	"github.com/google/osv.dev/go/logger"
 	"github.com/google/osv.dev/go/osv/clients"
 	"github.com/ossf/osv-schema/bindings/go/osvschema"
@@ -16,29 +17,24 @@ import (
 // result to outCh.
 func downloader(ctx context.Context, client clients.CloudStorage, inCh <-chan string, outCh chan<- *osvschema.Vulnerability, wg *sync.WaitGroup) {
 	defer wg.Done()
-	for {
-		var path string
-		var ok bool
-
-		// Wait to receive an object path, or be cancelled.
-		select {
-		case path, ok = <-inCh:
-			if !ok {
-				return // Channel closed.
-			}
-		case <-ctx.Done():
-			return
-		}
-
+	for path := range inCh {
 		// Process object.
 		data, err := client.ReadObject(ctx, path)
 		if err != nil {
-			logger.Error("failed to read vulnerability", slog.String("obj", path), slog.Any("err", err))
+			if osvutil.IsContextError(err) {
+				return
+			}
+			logger.ErrorContext(ctx, "failed to read vulnerability", slog.String("obj", path), slog.Any("err", err))
+
 			continue
 		}
 		vuln := &osvschema.Vulnerability{}
 		if err := proto.Unmarshal(data, vuln); err != nil {
-			logger.Error("failed to unmarshal vulnerability", slog.String("obj", path), slog.Any("err", err))
+			if osvutil.IsContextError(err) {
+				return
+			}
+			logger.ErrorContext(ctx, "failed to unmarshal vulnerability", slog.String("obj", path), slog.Any("err", err))
+
 			continue
 		}
 

go/cmd/exporter/exporter.go

@@ -7,9 +7,10 @@ import (
 	"flag"
 	"log/slog"
 	"os"
-	"path/filepath"
+	"os/signal"
 	"strings"
 	"sync"
+	"syscall"
 
 	"cloud.google.com/go/storage"
 	"github.com/google/osv.dev/go/logger"
@@ -27,21 +28,26 @@ func main() {
 	logger.InitGlobalLogger()
 	defer logger.Close()
 
-	ctx, span := otel.Tracer("exporter").Start(context.Background(), "exporter")
+	ctx, stopSignal := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stopSignal()
+
+	ctx, span := otel.Tracer("exporter").Start(ctx, "exporter")
 	defer span.End()
 
 	outBucketName := flag.String("bucket", "osv-test-vulnerabilities", "Output bucket or directory name. If -local is true, this is a local path; otherwise, it's a GCS bucket name.")
 	vulnBucketName := flag.String("osv-vulns-bucket", os.Getenv("OSV_VULNERABILITIES_BUCKET"), "GCS bucket to read vulnerability protobufs from. Can also be set with the OSV_VULNERABILITIES_BUCKET environment variable.")
 	uploadToGCS := flag.Bool("upload-to-gcs", false, "If false, writes the output to a local directory specified by -bucket instead of a GCS bucket.")
-	numWorkers := flag.Int("workers", 200, "The total number of concurrent workers to use for downloading from GCS and writing the output.")
+	numWorkers := flag.Int("workers", 1000, "The total number of concurrent workers to use for downloading from GCS and writing the output.")
+	breakdownPrefixesStr := flag.String("breakdown-prefixes", "", "Comma-separated list of prefix breakdowns for parallel GCS object listing. Defaults to A-Z if empty.")
 
 	flag.Parse()
 
 	logger.InfoContext(ctx, "exporter starting",
 		slog.String("bucket", *outBucketName),
 		slog.String("osv-vulns-bucket", *vulnBucketName),
 		slog.Bool("upload-to-gcs", *uploadToGCS),
-		slog.Int("workers", *numWorkers))
+		slog.Int("workers", *numWorkers),
+		slog.String("breakdown-prefixes", *breakdownPrefixesStr))
 
 	if *vulnBucketName == "" {
 		logger.FatalContext(ctx, "OSV_VULNERABILITIES_BUCKET must be set")
@@ -76,9 +82,25 @@ func main() {
 	// 4. The `ecosystemWorker`s and the `allEcosystemWorker` process the vulnerabilities and
 	//    generate the final files, sending the data to be written to `routerToWriteCh`.
 	// 5. A pool of `writer` workers receive the file data and write it to the output.
-	gcsPathToDownloaderCh := make(chan string)
-	downloaderToRouterCh := make(chan *osvschema.Vulnerability)
-	routerToWriteCh := make(chan writeMsg)
+	gcsPathToDownloaderCh := make(chan string, 100)
+	downloaderToRouterCh := make(chan *osvschema.Vulnerability, 100)
+	routerToWriteCh := make(chan writeMsg, 100)
+
+	var breakdownPrefixes []string
+	if *breakdownPrefixesStr != "" {
+		for _, p := range strings.Split(*breakdownPrefixesStr, ",") {
+			p = strings.TrimSpace(p)
+			if p != "" {
+				breakdownPrefixes = append(breakdownPrefixes, p)
+			}
+		}
+	}
+	// Fall back to simple A-Z prefixes
+	if len(breakdownPrefixes) == 0 {
+		for ch := 'A'; ch <= 'Z'; ch++ {
+			breakdownPrefixes = append(breakdownPrefixes, string(ch))
+		}
+	}
 
 	var downloaderWg sync.WaitGroup
 	for range *numWorkers / 2 {
@@ -95,21 +117,13 @@ func main() {
 	routerWg.Add(1)
 	go ecosystemRouter(ctx, downloaderToRouterCh, routerToWriteCh, &routerWg)
 
-	prevPrefix := ""
 MainLoop:
-	for obj, err := range vulnClient.Objects(ctx, gcsProtoPrefix) {
+	for objName, err := range vulnClient.ObjectsFast(ctx, gcsProtoPrefix, breakdownPrefixes) {
 		if err != nil {
 			logger.FatalContext(ctx, "failed to list objects", slog.Any("err", err))
 		}
-		// Only log when we see a new ID prefix (i.e. roughly once per data source)
-		prefix := filepath.Base(obj.Name)
-		prefix, _, _ = strings.Cut(prefix, "-")
-		if prefix != prevPrefix {
-			logger.InfoContext(ctx, "iterating vulnerabilities", slog.String("now_at", obj.Name))
-			prevPrefix = prefix
-		}
 		select {
-		case gcsPathToDownloaderCh <- obj.Name:
+		case gcsPathToDownloaderCh <- objName:
 		case <-ctx.Done():
 			break MainLoop
 		}

go/cmd/exporter/worker.go

@@ -38,7 +38,7 @@ type ecosystemWorker struct {
 
 // newEcosystemWorker creates and starts a new ecosystemWorker.
 func newEcosystemWorker(ctx context.Context, ecosystem string, outCh chan<- writeMsg, wg *sync.WaitGroup) *ecosystemWorker {
-	ch := make(chan *osvschema.Vulnerability)
+	ch := make(chan *osvschema.Vulnerability, 100)
 	worker := &ecosystemWorker{
 		ecosystem: ecosystem,
 		inCh:      ch,
@@ -65,24 +65,11 @@ func (w *ecosystemWorker) run(ctx context.Context, outCh chan<- writeMsg, wg *sy
 	defer span.End()
 
 	logger.InfoContext(ctx, "new ecosystem worker started", slog.String("ecosystem", w.ecosystem))
-	var allVulns []vulnData
-	var csvData [][]string
+	// 500 size is around the minimum each ecosystem would need, most ecosystems are much bigger.
+	allVulns := make([]vulnData, 0, 500)
+	csvData := make([][]string, 0, 500)
 	var vanirVulns []vulnData
-WorkLoop:
-	for {
-		var v *osvschema.Vulnerability
-		var ok bool
-
-		// Wait to receive a vulnerability, or be cancelled.
-		select {
-		case <-ctx.Done():
-			logger.WarnContext(ctx, "ecosystem worker cancelled", slog.String("ecosystem", w.ecosystem), slog.Any("err", ctx.Err()))
-			return
-		case v, ok = <-w.inCh:
-			if !ok {
-				break WorkLoop
-			}
-		}
+	for v := range w.inCh {
 		// Process vulnerability.
 		b, err := marshalToJSON(v)
 		if err != nil {
@@ -114,6 +101,10 @@ WorkLoop:
 		}
 	}
 
+	if ctx.Err() != nil {
+		return
+	}
+
 	logger.InfoContext(ctx, "All vulnerabilities processed", slog.String("ecosystem", w.ecosystem))
 	writeModifiedIDCSV(ctx, filepath.Join(w.ecosystem, modifiedCSVFilename), csvData, outCh)
 	writeZIP(ctx, filepath.Join(w.ecosystem, allZipFilename), allVulns, outCh)
@@ -143,7 +134,7 @@ type allEcosystemWorker struct {
 
 // newAllEcosystemWorker creates and starts a new allEcosystemWorker.
 func newAllEcosystemWorker(ctx context.Context, outCh chan<- writeMsg, wg *sync.WaitGroup) *allEcosystemWorker {
-	ch := make(chan vulnAndEcos)
+	ch := make(chan vulnAndEcos, 100)
 	worker := &allEcosystemWorker{
 		inCh: ch,
 	}
@@ -161,32 +152,31 @@ func (w *allEcosystemWorker) run(ctx context.Context, outCh chan<- writeMsg, wg
 	defer span.End()
 
 	logger.InfoContext(ctx, "all-ecosystem worker started")
-	var allVulns []vulnData
-	var csvData [][]string
+	// We have currently about 1.8 million entries, so start at 100k
+	allVulns := make([]vulnData, 0, 100000)
+	csvData := make([][]string, 0, 100000)
 	ecosystems := make(map[string]struct{})
-WorkLoop:
-	for {
-		select {
-		case <-ctx.Done():
-			logger.WarnContext(ctx, "all-ecosystem worker cancelled", slog.Any("err", ctx.Err()))
-			return
-		case v, ok := <-w.inCh:
-			if !ok {
-				break WorkLoop
-			}
-			b, err := marshalToJSON(v.Vulnerability)
-			if err != nil {
-				logger.ErrorContext(ctx, "failed to marshal vulnerability to json", slog.String("id", v.GetId()), slog.Any("err", err))
-				continue
-			}
-			modified := v.GetModified().AsTime()
-			allVulns = append(allVulns, vulnData{id: v.GetId(), modified: modified, data: b})
-			for _, e := range v.ecosystems {
-				ecosystems[e] = struct{}{}
-				csvData = append(csvData, []string{modified.Format(time.RFC3339Nano), e + "/" + v.GetId()})
+	for v := range w.inCh {
+		b, err := marshalToJSON(v.Vulnerability)
+		if err != nil {
+			logger.ErrorContext(ctx, "failed to marshal vulnerability to json", slog.String("id", v.GetId()), slog.Any("err", err))
+			continue
+		}
+		modified := v.GetModified().AsTime()
+		allVulns = append(allVulns, vulnData{id: v.GetId(), modified: modified, data: b})
+		for _, e := range v.ecosystems {
+			ecosystems[e] = struct{}{}
+			csvData = append(csvData, []string{modified.Format(time.RFC3339Nano), e + "/" + v.GetId()})
+			if len(csvData)%10000 == 0 {
+				logger.InfoContext(ctx, "processed N vulnerabilities", slog.Int("n", len(csvData)))
 			}
 		}
 	}
+
+	if ctx.Err() != nil {
+		return
+	}
+
 	writeModifiedIDCSV(ctx, modifiedCSVFilename, csvData, outCh)
 	writeZIP(ctx, allZipFilename, allVulns, outCh)
 	ecos := slices.Collect(maps.Keys(ecosystems))

go/cmd/exporter/writer.go

@@ -27,46 +27,37 @@ type writeMsg struct {
 // bucket or a local directory.
 func writer(ctx context.Context, cancel context.CancelFunc, inCh <-chan writeMsg, client clients.CloudStorage, pathPrefix string, wg *sync.WaitGroup) {
 	defer wg.Done()
-	for {
-		select {
-		case msg, ok := <-inCh:
-			if !ok {
-				// Channel closed.
-				return
+	for msg := range inCh {
+		path := filepath.Join(pathPrefix, msg.path)
+		if client != nil {
+			// Skip the upload if the object already has the same content.
+			if gcsContentUnchanged(ctx, client, path, msg.data) {
+				continue
 			}
-			path := filepath.Join(pathPrefix, msg.path)
-			if client != nil {
-				// Skip the upload if the object already has the same content.
-				if gcsContentUnchanged(ctx, client, path, msg.data) {
-					break
-				}
-				err := client.WriteObject(ctx, path, msg.data, &clients.WriteOptions{
-					ContentType: msg.mimeType,
-				})
-				if err != nil {
-					logger.Error("failed to write file", slog.String("path", path), slog.Any("err", err))
-					cancel()
+			err := client.WriteObject(ctx, path, msg.data, &clients.WriteOptions{
+				ContentType: msg.mimeType,
+			})
+			if err != nil {
+				logger.Error("failed to write file", slog.String("path", path), slog.Any("err", err))
+				cancel()
 
-					break
-				}
-			} else {
-				// Write locally.
-				dir := filepath.Dir(path)
-				if err := os.MkdirAll(dir, 0755); err != nil {
-					logger.Error("failed to create directories", slog.String("dir", dir), slog.Any("err", err))
-					cancel()
+				break
+			}
+		} else {
+			// Write locally.
+			dir := filepath.Dir(path)
+			if err := os.MkdirAll(dir, 0755); err != nil {
+				logger.Error("failed to create directories", slog.String("dir", dir), slog.Any("err", err))
+				cancel()
 
-					break
-				}
-				if err := os.WriteFile(path, msg.data, 0600); err != nil {
-					logger.Error("failed to write file", slog.String("path", path), slog.Any("err", err))
-					cancel()
+				break
+			}
+			if err := os.WriteFile(path, msg.data, 0600); err != nil {
+				logger.Error("failed to write file", slog.String("path", path), slog.Any("err", err))
+				cancel()
 
-					break
-				}
+				break
 			}
-		case <-ctx.Done():
-			return
 		}
 	}
 }