perf: optimize combine-to-osv bucket listing and avoid partial files.txt upload

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

@@ -4,6 +4,7 @@ package main
 import (
 	"cmp"
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
@@ -16,6 +17,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"cloud.google.com/go/storage"
 	"github.com/google/osv/vulnfeeds/conversion"
@@ -30,9 +32,10 @@ import (
 )
 
 const (
-	defaultOSVOutputPath = "osv-output"
-	defaultCVE5Path      = "cve5"
-	defaultNVDOSVPath    = "nvd"
+	defaultOSVOutputPath  = "osv-output"
+	defaultCVE5Path       = "cve5"
+	defaultNVDOSVPath     = "nvd"
+	defaultListFileMaxAge = 2 * 24 * time.Hour
 )
 
 // CVEWorkItem represents a unit of work for a single CVE.
@@ -57,11 +60,71 @@ func cveIDFromPath(p string) models.CVEID {
 	return ""
 }
 
-func listObjects(ctx context.Context, client *storage.Client, pathStr string) ([]string, error) {
+// 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) {
 	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))
+		}
+
+		logger.Info("Falling back to bucket listing", slog.String("bucket", bucketName), slog.String("prefix", prefix))
 		objs, err := gcs.ListBucketObjects(ctx, bucket, prefix)
 		if err != nil {
 			return nil, err
@@ -193,6 +256,7 @@ 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)
@@ -217,7 +281,7 @@ func main() {
 	go func() {
 		defer listObjWg.Done()
 		logger.Info("Listing CVE5 objects", slog.String("path", *cve5Path))
-		cve5Files, cve5ListErr = listObjects(ctx, client, *cve5Path)
+		cve5Files, cve5ListErr = listObjects(ctx, client, *cve5Path, *listFileMaxAge)
 		if cve5ListErr != nil {
 			logger.Fatal("Failed to list CVE5 objects", slog.Any("err", cve5ListErr))
 		}
@@ -226,7 +290,7 @@ func main() {
 	go func() {
 		defer listObjWg.Done()
 		logger.Info("Listing NVD objects", slog.String("path", *nvdPath))
-		nvdFiles, nvdListErr = listObjects(ctx, client, *nvdPath)
+		nvdFiles, nvdListErr = listObjects(ctx, client, *nvdPath, *listFileMaxAge)
 		if nvdListErr != nil {
 			logger.Fatal("Failed to list NVD objects", slog.Any("err", nvdListErr))
 		}

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

@@ -8,6 +8,7 @@ import (
 	"flag"
 	"log/slog"
 	"os"
+	"path"
 	"path/filepath"
 	"slices"
 	"strconv"
@@ -22,10 +23,12 @@ 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", "2022", "The first in scope year to process.")
+	startYear      = flag.String("start-year", defaultStartYear, "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.")
@@ -43,6 +46,20 @@ 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))
@@ -70,20 +87,29 @@ 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)
+		go worker(&wg, jobs, gcsHelper, *localOutputDir, cnaList, *rejectFailed, outputFilesChan)
 	}
 
 	// 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)
@@ -112,22 +138,37 @@ 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) {
+func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDir string, cnas []string, rejectFailed bool, outputFilesChan chan<- string) {
 	defer wg.Done()
-	for path := range jobs {
-		data, err := os.ReadFile(path)
+	for cvePath := range jobs {
+		data, err := os.ReadFile(cvePath)
 		if err != nil {
-			logger.Info("Failed to read file", slog.String("path", path), slog.Any("err", err))
+			logger.Info("Failed to read file", slog.String("path", cvePath), 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", path), slog.Any("err", err))
+			logger.Info("Failed to unmarshal JSON", slog.String("path", cvePath), slog.Any("err", err))
 			continue
 		}
 
@@ -139,9 +180,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(path, baseDirCVEList)
+		idx := strings.Index(cvePath, baseDirCVEList)
 		if idx != -1 {
-			relPath := path[idx:]
+			relPath := cvePath[idx:]
 			sourceLink = "https://github.com/CVEProject/cvelistV5/tree/main/" + relPath
 		}
 
@@ -153,6 +194,8 @@ 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 {

vulnfeeds/cmd/converters/cve/nvd-cve-osv/main.go

@@ -8,6 +8,7 @@ import (
 	"log/slog"
 	"net/http"
 	"os"
+	"path"
 	"path/filepath"
 	"regexp"
 	"runtime/pprof"
@@ -26,6 +27,8 @@ import (
 	"github.com/google/osv/vulnfeeds/vulns"
 )
 
+const defaultStartYear = 2016
+
 var (
 	jsonPath            = flag.String("nvd-json", "", "Path to NVD CVE JSON to examine.")
 	jsonDir             = flag.String("nvd-json-dir", "", "Path to directory containing NVD CVE JSON files to examine.")
@@ -39,7 +42,7 @@ var (
 	uploadToGCS         = flag.Bool("upload-to-gcs", false, "If true, upload to GCS bucket instead of writing to local disk.")
 	outputBucket        = flag.String("output-bucket", "osv-test-cve-osv-conversion", "The GCS bucket to write to.")
 	gcsPrefix           = flag.String("gcs-prefix", "nvd-osv", "The prefix within the GCS bucket.")
-	startYear           = flag.Int("start-year", 2016, "The first in scope year to process. If 0, process all years.")
+	startYear           = flag.Int("start-year", defaultStartYear, "The first in scope year to process. If 0, process all years.")
 )
 
 var (
@@ -71,6 +74,11 @@ func main() {
 	logger.InitGlobalLogger()
 	defer logger.Close()
 
+	isPartial := *startYear > defaultStartYear
+	if isPartial {
+		logger.Info("Partial run detected (start-year > " + strconv.Itoa(defaultStartYear) + "), will skip files.txt upload")
+	}
+
 	var files []string
 	if *jsonDir != "" {
 		matches, err := filepath.Glob(filepath.Join(*jsonDir, "nvdcve-2.0-*.json"))
@@ -130,16 +138,26 @@ 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)
+		}
+	}()
+
 	jobs := make(chan models.NVDCVE, *workers)
 	var wg sync.WaitGroup
 
 	for range *workers {
 		wg.Add(1)
-		go worker(&wg, jobs, gcsHelper, *outDir, vpRepoCache, repoTagsCache)
+		go worker(&wg, jobs, gcsHelper, *outDir, vpRepoCache, repoTagsCache, outputFilesChan)
 	}
 
 	for _, file := range files {
@@ -164,11 +182,21 @@ func main() {
 
 	close(jobs)
 	wg.Wait()
+	close(outputFilesChan)
+	collectorWg.Wait()
 
 	timesBlocked := int64(0)
 	if gcsHelper != nil {
 		timesBlocked = gcsHelper.GetTimesBlocked()
 		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("Conversion Stats",
 		slog.Uint64("total_processed", totalConversionsCount.Load()),
@@ -199,7 +227,7 @@ func processCVE(cve models.NVDCVE, vpRepoCache *c.VPRepoCache, repoTagsCache git
 	return nvd.CVEToOSV(cve, repos, vpRepoCache, repoTagsCache, metrics)
 }
 
-func worker(wg *sync.WaitGroup, jobs <-chan models.NVDCVE, gcsHelper *gcs.Helper, outDir string, vpRepoCache *c.VPRepoCache, repoTagsCache git.RepoTagsCache) {
+func worker(wg *sync.WaitGroup, jobs <-chan models.NVDCVE, gcsHelper *gcs.Helper, outDir string, vpRepoCache *c.VPRepoCache, repoTagsCache git.RepoTagsCache, outputFilesChan chan<- string) {
 	defer wg.Done()
 	for cve := range jobs {
 		vuln, metrics, outcome := processCVE(cve, vpRepoCache, repoTagsCache)
@@ -245,6 +273,8 @@ func worker(wg *sync.WaitGroup, jobs <-chan models.NVDCVE, gcsHelper *gcs.Helper
 			if vuln != nil {
 				if err := writer.UploadVulnIfChangedAsync(gcsHelper, *gcsPrefix, vuln.Vulnerability); err != nil {
 					logger.Error("Failed to queue vulnerability upload", slog.String("cve", vuln.Id), slog.Any("err", err))
+				} else {
+					outputFilesChan <- path.Join(*gcsPrefix, vuln.Id+".json")
 				}
 			}
 			if *outputMetrics && metrics != nil {