perf(exporter): retain pre-compressed deflate payloads in memory

go/cmd/exporter/downloader.go

@@ -7,7 +7,6 @@ import (
 	"hash/crc32"
 	"io"
 	"log/slog"
-	"os"
 	"path/filepath"
 	"strings"
 	"sync"
@@ -28,9 +27,9 @@ var flateWriterPool = sync.Pool{
 
 // downloadThenProcessor is a worker that receives GCS object handles from inCh, downloads
 // the raw protobuf data, unmarshals it into a Vulnerability, marshals it to compact
-// JSON, saves pre-compressed Deflate data to scratch disk, queues individual JSON uploads,
-// and sends the metadata to routerCh.
-func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, client clients.CloudStorage, scratchDir string, inCh <-chan string, routerCh chan<- processedVuln, writeCh chan<- writeMsg, wg *sync.WaitGroup) {
+// JSON, pre-compresses Deflate data in memory, queues individual JSON uploads,
+// and sends the metadata and compressed payload to routerCh.
+func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, client clients.CloudStorage, inCh <-chan string, routerCh chan<- processedVuln, writeCh chan<- writeMsg, wg *sync.WaitGroup) {
 	defer wg.Done()
 	for path := range inCh {
 		// Process object.
@@ -83,18 +82,6 @@ func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, clien
 		compressedBytes := compBuf.Bytes()
 		crc := crc32.ChecksumIEEE(b)
 
-		// Cache pre-compressed Deflate payload to local scratch disk.
-		localPath := filepath.Join(scratchDir, vuln.GetId()+".deflate")
-		//nolint:gosec // G703: Staging temporary file in scratch directory
-		if err := os.WriteFile(localPath, compressedBytes, 0600); err != nil {
-			logger.ErrorContext(ctx, "failed to write cached vulnerability to disk", slog.String("id", vuln.GetId()), slog.Any("err", err))
-			// Cancel the exporter context if writing to the scratch disk fails (e.g. disk full)
-			// to fail fast rather than producing incomplete archives later.
-			cancel()
-
-			return
-		}
-
 		hasVanir := false
 		// Check for Vanir signatures
 		for _, aff := range vuln.GetAffected() {
@@ -136,11 +123,12 @@ func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, clien
 		select {
 		case routerCh <- processedVuln{
 			meta: vulnMeta{
-				id:         vuln.GetId(),
-				modified:   vuln.GetModified().AsTime(),
-				crc32:      crc,
-				uncompSize: uint64(len(b)),
-				compSize:   uint64(len(compressedBytes)),
+				id:          vuln.GetId(),
+				modified:    vuln.GetModified().AsTime(),
+				crc32:       crc,
+				uncompSize:  uint64(len(b)),
+				compSize:    uint64(len(compressedBytes)),
+				deflateData: compressedBytes,
 			},
 			ecosystems: ecoNames,
 			hasVanir:   hasVanir,

go/cmd/exporter/exporter.go

@@ -45,7 +45,7 @@ func main() {
 	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.")
 	scratchDirFlag := flag.String("scratch-dir", defaultScratchDir, "Directory to stage temporary JSON and zip files.")
-	cleanUpScratchDir := flag.Bool("cleanup-scratch-dir", false, "Whether to delete the temporary scratch directory on exit. Defaults to false.")
+	cleanUpScratchDir := flag.Bool("cleanup-scratch-dir", true, "Whether to delete the temporary scratch directory on exit. Defaults to true.")
 
 	flag.Parse()
 
@@ -114,7 +114,7 @@ func main() {
 	var processorWg sync.WaitGroup
 	for range *numWorkers / 2 {
 		processorWg.Add(1)
-		go downloadThenProcessor(ctx, cancel, vulnClient, scratchDir, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg)
+		go downloadThenProcessor(ctx, cancel, vulnClient, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg)
 	}
 
 	var writerWg sync.WaitGroup
@@ -190,7 +190,7 @@ func ecosystemRouter(ctx context.Context, inCh <-chan processedVuln, outCh chan<
 	workers := make(map[string]*ecosystemWorker)
 	var workersWg sync.WaitGroup
 	vulnCounter := 0
-	var vanirVulnIDs []string
+	var vanirVulns []vanirVuln
 
 	allEcosystemWorker := newAllEcosystemWorker(ctx, scratchDir, outCh, &workersWg)
 
@@ -209,7 +209,7 @@ RouterLoop:
 		vulnCounter++
 
 		if vuln.hasVanir {
-			vanirVulnIDs = append(vanirVulnIDs, vuln.meta.id)
+			vanirVulns = append(vanirVulns, vanirVuln{id: vuln.meta.id, deflateData: vuln.meta.deflateData})
 		}
 
 		for _, eco := range vuln.ecosystems {
@@ -237,8 +237,8 @@ RouterLoop:
 	allEcosystemWorker.Finish()
 	workersWg.Wait()
 
-	if len(vanirVulnIDs) > 0 && ctx.Err() == nil {
-		writeVanir(ctx, vanirVulnIDs, outCh, scratchDir)
+	if len(vanirVulns) > 0 && ctx.Err() == nil {
+		writeVanir(ctx, vanirVulns, outCh)
 	}
 
 	if ctx.Err() == nil {

go/cmd/exporter/exporter_test.go

@@ -137,7 +137,7 @@ func TestExporterPipeline_EndToEnd(t *testing.T) {
 
 	var processorWg sync.WaitGroup
 	processorWg.Add(1)
-	go downloadThenProcessor(ctx, cancel, vulnStorage, scratchDir, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg)
+	go downloadThenProcessor(ctx, cancel, vulnStorage, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg)
 
 	var writerWg sync.WaitGroup
 	writerWg.Add(1)

go/cmd/exporter/worker.go

@@ -32,13 +32,20 @@ const (
 	ecosystemsFilename  = "ecosystems.txt"
 )
 
-// vulnMeta holds the ID, modified time, CRC32, and pre-compression sizes for a vulnerability.
+// vulnMeta holds the ID, modified time, CRC32, pre-compression sizes, and deflate payload for a vulnerability.
 type vulnMeta struct {
-	id         string
-	modified   time.Time
-	crc32      uint32
-	uncompSize uint64
-	compSize   uint64
+	id          string
+	modified    time.Time
+	crc32       uint32
+	uncompSize  uint64
+	compSize    uint64
+	deflateData []byte
+}
+
+// vanirVuln holds the ID and deflate payload for a vulnerability with Vanir signatures.
+type vanirVuln struct {
+	id          string
+	deflateData []byte
 }
 
 // csvEntry holds the modified time and the relative entry path.
@@ -234,7 +241,7 @@ func writeModifiedIDCSV(ctx context.Context, path string, csvData []csvEntry, ou
 	write(ctx, path, buf.Bytes(), "text/csv", outCh)
 }
 
-// writeZIP constructs and writes a zip file by streaming pre-compressed local files via CreateRaw.
+// writeZIP constructs and writes a zip file by streaming pre-compressed in-memory payloads via CreateRaw.
 func writeZIP(ctx context.Context, path string, allVulns []vulnMeta, outCh chan<- writeMsg, scratchDir string) {
 	logger.InfoContext(ctx, "constructing zip file", slog.String("path", path))
 	slices.SortFunc(allVulns, func(a, b vulnMeta) int {
@@ -262,16 +269,9 @@ func writeZIP(ctx context.Context, path string, allVulns []vulnMeta, outCh chan<
 			logger.ErrorContext(ctx, "failed to create raw vuln in zip file", slog.String("id", vuln.id), slog.Any("err", err))
 			continue
 		}
-		localPath := filepath.Join(scratchDir, vuln.id+".deflate")
-		f, err := os.Open(localPath)
-		if err != nil {
-			logger.ErrorContext(ctx, "failed to open local vuln deflate file", slog.String("path", localPath), slog.Any("err", err))
-			continue
-		}
-		if _, err := io.Copy(w, f); err != nil {
+		if _, err := w.Write(vuln.deflateData); err != nil {
 			logger.ErrorContext(ctx, "failed to write vuln deflate data to zip file", slog.String("id", vuln.id), slog.Any("err", err))
 		}
-		f.Close()
 	}
 	if err := wr.Close(); err != nil {
 		logger.ErrorContext(ctx, "failed to close zip writer", slog.String("path", path), slog.Any("err", err))
@@ -283,26 +283,20 @@ func writeZIP(ctx context.Context, path string, allVulns []vulnMeta, outCh chan<
 }
 
 // writeVanir constructs and writes the osv_git.json file containing vulnerabilities with Vanir signatures
-// by reading the cached JSON files from disk and marshaling the combined JSON array in memory.
-func writeVanir(ctx context.Context, vanirVulnIDs []string, outCh chan<- writeMsg, scratchDir string) {
-	logger.InfoContext(ctx, "constructing vanir file", slog.Int("count", len(vanirVulnIDs)))
-	slices.Sort(vanirVulnIDs)
-
-	vulns := make([]json.RawMessage, 0, len(vanirVulnIDs))
-	for _, id := range vanirVulnIDs {
-		localPath := filepath.Join(scratchDir, id+".deflate")
-		//nolint:gosec // G703: Reading temporary local file in scratch directory
-		f, err := os.Open(localPath)
-		if err != nil {
-			logger.ErrorContext(ctx, "failed to open local vuln file for vanir", slog.String("id", id), slog.Any("err", err))
-			continue
-		}
-		fr := flate.NewReader(f)
+// by decompressing the in-memory deflate payloads and marshaling the combined JSON array in memory.
+func writeVanir(ctx context.Context, vanirVulns []vanirVuln, outCh chan<- writeMsg) {
+	logger.InfoContext(ctx, "constructing vanir file", slog.Int("count", len(vanirVulns)))
+	slices.SortFunc(vanirVulns, func(a, b vanirVuln) int {
+		return cmp.Compare(a.id, b.id)
+	})
+
+	vulns := make([]json.RawMessage, 0, len(vanirVulns))
+	for _, v := range vanirVulns {
+		fr := flate.NewReader(bytes.NewReader(v.deflateData))
 		data, err := io.ReadAll(fr)
 		_ = fr.Close()
-		_ = f.Close()
 		if err != nil {
-			logger.ErrorContext(ctx, "failed to decompress local vuln file for vanir", slog.String("id", id), slog.Any("err", err))
+			logger.ErrorContext(ctx, "failed to decompress vuln file for vanir", slog.String("id", v.id), slog.Any("err", err))
 			continue
 		}
 		vulns = append(vulns, data)