perf: optimize combine-to-osv bucket listing

osv.dev.code-workspace

@@ -0,0 +1,17 @@
+{
+	"folders": [
+		{
+			"path": "."
+		},
+		{
+			"path": "go"
+		},
+		{
+			"path": "vulnfeeds"
+		},
+		{
+			"path": "go/cmd/tools/reimport-tui"
+		}
+	],
+	"settings": {}
+}

vulnfeeds/cmd/combine-to-osv/main.go

@@ -4,7 +4,6 @@ package main
 import (
 	"cmp"
 	"context"
-	"errors"
 	"flag"
 	"fmt"
 	"io"
@@ -32,10 +31,10 @@ import (
 )
 
 const (
-	defaultOSVOutputPath  = "osv-output"
-	defaultCVE5Path       = "cve5"
-	defaultNVDOSVPath     = "nvd"
-	defaultListFileMaxAge = 2 * 24 * time.Hour
+	defaultOSVOutputPath = "osv-output"
+	defaultCVE5Path      = "cve5"
+	defaultNVDOSVPath    = "nvd"
+	parallelStartYear    = 2018
 )
 
 // CVEWorkItem represents a unit of work for a single CVE.
@@ -60,78 +59,26 @@ func cveIDFromPath(p string) models.CVEID {
 	return ""
 }
 
-// readListFile reads the list file (usually files.txt) from GCS.
-// The list file is expected to contain newline-separated GCS object names
-// relative to the bucket root (e.g. "prefix/CVE-2023-0001.json").
-// It returns a slice of full GCS paths (e.g. "gs://bucket/prefix/CVE-2023-0001.json").
-func readListFile(ctx context.Context, listObj *storage.ObjectHandle, bucketName string) ([]string, error) {
-	rc, err := listObj.NewReader(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("failed to open reader: %w", err)
-	}
-	defer rc.Close()
-
-	content, err := io.ReadAll(rc)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read content: %w", err)
-	}
-
-	lines := strings.Split(string(content), "\n")
-	fullPaths := make([]string, 0, len(lines))
-	for _, line := range lines {
-		line = strings.TrimSpace(line)
-		if line == "" {
-			continue
-		}
-		fullPaths = append(fullPaths, fmt.Sprintf("gs://%s/%s", bucketName, line))
-	}
-
-	return fullPaths, nil
-}
-
-func listObjects(ctx context.Context, client *storage.Client, pathStr string, maxAge time.Duration) ([]string, error) {
+func listObjects(ctx context.Context, client *storage.Client, pathStr string) ([]string, error) {
 	if strings.HasPrefix(pathStr, "gs://") {
 		trimmed := strings.TrimPrefix(pathStr, "gs://")
 		bucketName, prefix, _ := strings.Cut(trimmed, "/")
 		bucket := client.Bucket(bucketName)
 
-		listObjectName := gcs.FileListFilename
-		if prefix != "" {
-			listObjectName = strings.TrimSuffix(prefix, "/") + "/" + gcs.FileListFilename
-		}
-
-		listObj := bucket.Object(listObjectName)
-		attrs, err := listObj.Attrs(ctx)
-		useListFile := false
-		if err == nil {
-			// Check age
-			age := time.Since(attrs.Updated)
-			if age <= maxAge {
-				useListFile = true
-			} else {
-				logger.Error("List file is too old (generators might be failing)", slog.String("bucket", bucketName), slog.String("prefix", prefix), slog.Duration("age", age))
-			}
-		} else if !errors.Is(err, storage.ErrObjectNotExist) {
-			logger.Warn("Failed to get list file attributes", slog.String("bucket", bucketName), slog.String("prefix", prefix), slog.Any("err", err))
-		}
-
-		if useListFile {
-			logger.Info("Using list file", slog.String("bucket", bucketName), slog.String("prefix", prefix))
-			fullPaths, err := readListFile(ctx, listObj, bucketName)
-			if err == nil {
-				return fullPaths, nil
-			}
-			logger.Error("Failed to read list file, falling back to bucket listing", slog.String("bucket", bucketName), slog.String("prefix", prefix), slog.Any("err", err))
+		currentYear := time.Now().Year()
+		var breakdowns []string
+		for year := parallelStartYear; year <= currentYear; year++ {
+			breakdowns = append(breakdowns, fmt.Sprintf("CVE-%d-", year))
 		}
 
-		logger.Info("Falling back to bucket listing", slog.String("bucket", bucketName), slog.String("prefix", prefix))
-		objs, err := gcs.ListBucketObjects(ctx, bucket, prefix)
+		objs, err := gcs.ListObjectsFast(ctx, bucket, prefix, breakdowns)
 		if err != nil {
 			return nil, err
 		}
+
 		var fullPaths []string
 		for _, obj := range objs {
-			if strings.HasSuffix(obj, "/") || !strings.HasSuffix(obj, ".json") || strings.HasSuffix(obj, ".metrics.json") {
+			if !strings.HasSuffix(obj, ".json") || strings.HasSuffix(obj, ".metrics.json") {
 				continue
 			}
 			fullPaths = append(fullPaths, fmt.Sprintf("gs://%s/%s", bucketName, obj))
@@ -256,7 +203,6 @@ func main() {
 	uploadToGCS := flag.Bool("upload-to-gcs", false, "If true, upload to GCS bucket instead of writing to local disk.")
 	numWorkers := flag.Int("workers", 64, "Number of workers to process records")
 	syncDeletions := flag.Bool("sync-deletions", false, "If false, do not delete files in bucket that are not local")
-	listFileMaxAge := flag.Duration("list-file-max-age", defaultListFileMaxAge, "The maximum age of the list file before falling back to listing the bucket.")
 	flag.Parse()
 
 	err := os.MkdirAll(*osvOutputPath, 0755)
@@ -281,7 +227,7 @@ func main() {
 	go func() {
 		defer listObjWg.Done()
 		logger.Info("Listing CVE5 objects", slog.String("path", *cve5Path))
-		cve5Files, cve5ListErr = listObjects(ctx, client, *cve5Path, *listFileMaxAge)
+		cve5Files, cve5ListErr = listObjects(ctx, client, *cve5Path)
 		if cve5ListErr != nil {
 			logger.Fatal("Failed to list CVE5 objects", slog.Any("err", cve5ListErr))
 		}
@@ -290,7 +236,7 @@ func main() {
 	go func() {
 		defer listObjWg.Done()
 		logger.Info("Listing NVD objects", slog.String("path", *nvdPath))
-		nvdFiles, nvdListErr = listObjects(ctx, client, *nvdPath, *listFileMaxAge)
+		nvdFiles, nvdListErr = listObjects(ctx, client, *nvdPath)
 		if nvdListErr != nil {
 			logger.Fatal("Failed to list NVD objects", slog.Any("err", nvdListErr))
 		}

vulnfeeds/cmd/converters/cve/cve5/bulk-converter/main.go

@@ -8,7 +8,6 @@ import (
 	"flag"
 	"log/slog"
 	"os"
-	"path"
 	"path/filepath"
 	"slices"
 	"strconv"
@@ -23,12 +22,10 @@ import (
 	"github.com/google/osv/vulnfeeds/utility/logger"
 )
 
-const defaultStartYear = "2022"
-
 var (
 	repoDir        = flag.String("cve5-repo", "cvelistV5", "CVEListV5 directory path")
 	localOutputDir = flag.String("out-dir", "cve5", "Path to output results.")
-	startYear      = flag.String("start-year", defaultStartYear, "The first in scope year to process.")
+	startYear      = flag.String("start-year", "2022", "The first in scope year to process.")
 	workers        = flag.Int("workers", 10, "The number of concurrent workers to use for processing CVEs.")
 	gcsWorkers     = flag.Int("gcs-workers", 30, "The number of concurrent workers to use for GCS uploads.")
 	cnaDenyList    = flag.String("cna-denylist", "", "A comma-separated list of CNAs to skip. If not provided, defaults to cna_denylist.txt.")
@@ -46,20 +43,6 @@ func main() {
 	logger.InitGlobalLogger()
 	defer logger.Close()
 
-	startYearInt, err := strconv.Atoi(*startYear)
-	defaultStartYearInt, _ := strconv.Atoi(defaultStartYear)
-	isPartial := false
-	if err != nil {
-		logger.Error("Invalid start-year, assuming partial", slog.String("start-year", *startYear))
-		isPartial = true
-	} else if startYearInt > defaultStartYearInt {
-		isPartial = true
-	}
-
-	if isPartial {
-		logger.Info("Partial run detected (start-year > " + defaultStartYear + "), will skip files.txt upload")
-	}
-
 	logger.Info("Commencing CVE to OSV conversion run")
 	if err := os.MkdirAll(*localOutputDir, 0755); err != nil {
 		logger.Fatal("Failed to create local output directory", slog.Any("err", err))
@@ -87,29 +70,20 @@ func main() {
 		if err != nil {
 			logger.Fatal("Failed to initialize GCS upload pool", slog.Any("err", err))
 		}
+		defer gcsHelper.CloseAndWait()
 		logger.Info("GCS Upload Pool initialized", slog.String("bucket", *outputBucket))
 	}
 
-	outputFilesChan := make(chan string)
-	var collectorWg sync.WaitGroup
-	var outputFiles []string
-	collectorWg.Add(1)
-	go func() {
-		defer collectorWg.Done()
-		for f := range outputFilesChan {
-			outputFiles = append(outputFiles, f)
-		}
-	}()
-
 	// Start the worker pool.
 	for range *workers {
 		wg.Add(1)
-		go worker(&wg, jobs, gcsHelper, *localOutputDir, cnaList, *rejectFailed, outputFilesChan)
+		go worker(&wg, jobs, gcsHelper, *localOutputDir, cnaList, *rejectFailed)
 	}
 
 	// Discover files and send them to the workers.
 	logger.Info("Starting conversion of CVEs...")
 	currentYear := time.Now().Year()
+	startYearInt, _ := strconv.Atoi(*startYear)
 
 	for year := startYearInt; year <= currentYear; year++ {
 		year := strconv.Itoa(year)
@@ -138,37 +112,22 @@ func main() {
 
 	close(jobs)
 	wg.Wait()
-	close(outputFilesChan)
-	collectorWg.Wait()
-
-	if *uploadToGCS && gcsHelper != nil {
-		gcsHelper.CloseAndWait()
-		if !isPartial {
-			logger.Info("Uploading output file list to GCS...")
-			if err := gcs.UploadFileList(ctx, *outputBucket, *gcsPrefix, outputFiles); err != nil {
-				logger.Error("Failed to upload output file list", slog.Any("err", err))
-			}
-		} else {
-			logger.Info("Skipping files.txt upload due to partial run")
-		}
-	}
-
 	logger.Info("CVE5 Conversion run complete")
 }
 
 // worker is a function that processes CVE files from the jobs channel.
-func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDir string, cnas []string, rejectFailed bool, outputFilesChan chan<- string) {
+func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDir string, cnas []string, rejectFailed bool) {
 	defer wg.Done()
-	for cvePath := range jobs {
-		data, err := os.ReadFile(cvePath)
+	for path := range jobs {
+		data, err := os.ReadFile(path)
 		if err != nil {
-			logger.Info("Failed to read file", slog.String("path", cvePath), slog.Any("err", err))
+			logger.Info("Failed to read file", slog.String("path", path), slog.Any("err", err))
 			continue
 		}
 
 		var cve models.CVE5
 		if err := json.Unmarshal(data, &cve); err != nil {
-			logger.Info("Failed to unmarshal JSON", slog.String("path", cvePath), slog.Any("err", err))
+			logger.Info("Failed to unmarshal JSON", slog.String("path", path), slog.Any("err", err))
 			continue
 		}
 
@@ -180,9 +139,9 @@ func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDi
 
 		sourceLink := ""
 		baseDirCVEList := "cves/" // The base folder for the CVEListV5 repository.
-		idx := strings.Index(cvePath, baseDirCVEList)
+		idx := strings.Index(path, baseDirCVEList)
 		if idx != -1 {
-			relPath := cvePath[idx:]
+			relPath := path[idx:]
 			sourceLink = "https://github.com/CVEProject/cvelistV5/tree/main/" + relPath
 		}
 
@@ -194,8 +153,6 @@ func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDi
 				logger.Info("Queueing OSV record for "+string(cveID), slog.String("cve", string(cveID)))
 				if err := writer.UploadVulnIfChangedAsync(gcsHelper, *gcsPrefix, vuln.Vulnerability); err != nil {
 					logger.Error("Failed to queue vulnerability upload", slog.String("cve", string(cveID)), slog.Any("err", err))
-				} else {
-					outputFilesChan <- path.Join(*gcsPrefix, string(cveID)+".json")
 				}
 
 				if err := writer.UploadMetricsToGCSAsync(gcsHelper, *gcsPrefix, cveID, metrics); err != nil {