libvuln: offline import performance improvements

libvuln/jsonblob/jsonblob.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"iter"
 	"os"
 	"sort"
 	"sync"
@@ -45,29 +46,43 @@ type Store struct {
 }
 
 // Load reads in all the records serialized in the provided [io.Reader].
+//
+// Deprecated: This just calls [NewLoader].
+//
+//go:fix inline
 func Load(ctx context.Context, r io.Reader) (*Loader, error) {
+	return NewLoader(ctx, r)
+}
+
+// NewLoader returns a loader configured to read the records serialized in the
+// provided [io.Reader].
+func NewLoader(ctx context.Context, r io.Reader) (*Loader, error) {
 	l := Loader{
 		dec: json.NewDecoder(r),
 		cur: uuid.Nil,
 	}
 	return &l, nil
 }
 
-// Loader is an iterator that returns a series of [Entry].
+// Loader is an iterator over serialized records.
 //
-// Users should call [*Loader.Next] until it reports false, then check for
-// errors via [*Loader.Err].
+// Users should consume the [Loader.All] iterator, then check for errors via
+// [Loader.Err].
 type Loader struct {
-	err error
-	e   *Entry
+	err     error
+	dec     *json.Decoder
+	started bool
 
-	dec  *json.Decoder
+	// Used in the deprecated [Loader.Next]+[Loader.Entry] flow.
+	e    *Entry
 	next *Entry
 	de   diskEntry
 	cur  uuid.UUID
 }
 
 // Next reports whether there's an [Entry] to be processed.
+//
+// Deprecated: Use the [Loader.All] iterator.
 func (l *Loader) Next() bool {
 	if l.err != nil {
 		return false
@@ -113,10 +128,78 @@ func (l *Loader) Next() bool {
 }
 
 // Entry returns the latest loaded [Entry].
+//
+// Deprecated: Use the [Loader.All] iterator.
 func (l *Loader) Entry() *Entry {
 	return l.e
 }
 
+// All returns an iterator-of-iterators yielding all entries in the Loader.
+//
+// The inner iterator returns zero or one of the two values. If both pointers
+// are nil, the caller must check [Loader.Err].
+func (l *Loader) All() iter.Seq2[*Entry, iter.Seq2[*claircore.Vulnerability, *driver.EnrichmentRecord]] {
+	// These are shared across the two iterators:
+	cur := uuid.Nil
+	var de diskEntry
+
+	// The inner iterator decodes every entry after the first one, stopping
+	// iteration and passing control to the outer iterator when the Ref
+	// changes.
+	inner := func(yield func(*claircore.Vulnerability, *driver.EnrichmentRecord) bool) {
+		for ; l.err == nil && cur == de.Ref; l.err = l.dec.Decode(&de) {
+			var v *claircore.Vulnerability
+			var e *driver.EnrichmentRecord
+			switch de.Kind {
+			case driver.VulnerabilityKind:
+				v = getVulnerability()
+				l.err = json.Unmarshal(de.Vuln.buf, v)
+			case driver.EnrichmentKind:
+				e = getEnrichment()
+				l.err = json.Unmarshal(de.Enrichment.buf, e)
+			}
+			if l.err != nil {
+				yield(nil, nil)
+				return
+			}
+			if !yield(v, e) {
+				return
+			}
+		}
+		if l.err != nil && !errors.Is(l.err, io.EOF) {
+			yield(nil, nil)
+		}
+	}
+	// Outer reads the first entry and handles when the Ref changes.
+	outer := func(yield func(*Entry, iter.Seq2[*claircore.Vulnerability, *driver.EnrichmentRecord]) bool) {
+		if l.started {
+			l.err = errors.Join(l.err, fmt.Errorf("attempted to re-use a Loader"))
+			return
+		}
+		l.started = true
+
+		for l.err = l.dec.Decode(&de); l.err == nil; {
+			ent := &Entry{
+				CommonEntry: CommonEntry{
+					Updater:     de.Updater,
+					Fingerprint: de.Fingerprint,
+					Date:        de.Date,
+					Kind:        de.Kind,
+				},
+			}
+			cur = de.Ref
+			if !yield(ent, inner) {
+				return
+			}
+			for ; l.err == nil && cur == de.Ref; l.err = l.dec.Decode(&de) {
+				// Skip entries if the inner iterator was not fully consumed.
+			}
+		}
+	}
+
+	return outer
+}
+
 // Err is the latest encountered error.
 func (l *Loader) Err() error {
 	// Don't report EOF as an error.
@@ -148,10 +231,10 @@ func (s *Store) Store(w io.Writer) error {
 			shim := newBufShim(f)
 			defer shim.Close()
 			for range ct {
+				e.Kind = k
 				dent := diskEntry{
 					CommonEntry: e,
 					Ref:         id,
-					Kind:        k,
 				}
 				switch k {
 				case driver.EnrichmentKind:
@@ -237,6 +320,7 @@ type CommonEntry struct {
 	Updater     string
 	Fingerprint driver.Fingerprint
 	Date        time.Time
+	Kind        driver.UpdateKind
 }
 
 // DiskEntry is a single vulnerability or enrichment. It's made from unpacking an
@@ -249,7 +333,6 @@ type diskEntry struct {
 	Ref        uuid.UUID
 	Vuln       *bufShim `json:",omitempty"`
 	Enrichment *bufShim `json:",omitempty"`
-	Kind       driver.UpdateKind
 }
 
 // Entries returns a map containing all the Entries stored by calls to
@@ -477,6 +560,7 @@ func getBuf() []byte {
 	}
 	return b
 }
+
 func putBuf(b []byte) {
 	bufPool.Put(b)
 }

libvuln/jsonblob/pool.go

@@ -0,0 +1,48 @@
+package jsonblob
+
+import (
+	"sync"
+
+	"github.com/quay/claircore"
+	"github.com/quay/claircore/libvuln/driver"
+)
+
+var (
+	vulnerability sync.Pool
+	enrichment    sync.Pool
+)
+
+func getVulnerability() *claircore.Vulnerability {
+	if v := vulnerability.Get(); v != nil {
+		return v.(*claircore.Vulnerability)
+	}
+	return new(claircore.Vulnerability)
+}
+
+// ReturnVulnerability can be used by callers to return
+// [claircore.Vulnerability] objects from [Loader.All] iterators to a common
+// pool.
+//
+// This may take some pressure off the garbage collector.
+func ReturnVulnerability(v *claircore.Vulnerability) {
+	// Reset the fields.
+	*v = claircore.Vulnerability{}
+	vulnerability.Put(v)
+}
+
+func getEnrichment() *driver.EnrichmentRecord {
+	if v := enrichment.Get(); v != nil {
+		return v.(*driver.EnrichmentRecord)
+	}
+	return new(driver.EnrichmentRecord)
+}
+
+// ReturnEnrichment can be used by callers to return [driver.EnrichmentRecord]
+// objects from [Loader.All] iterators to a common pool.
+//
+// This may take some pressure off the garbage collector.
+func ReturnEnrichment(e *driver.EnrichmentRecord) {
+	// Reset the fields.
+	*e = driver.EnrichmentRecord{}
+	enrichment.Put(e)
+}

libvuln/offlineimport_test.go

@@ -0,0 +1,120 @@
+package libvuln
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/klauspost/compress/zstd"
+
+	"github.com/quay/claircore/datastore/postgres"
+	"github.com/quay/claircore/libvuln/driver"
+	"github.com/quay/claircore/libvuln/jsonblob"
+	"github.com/quay/claircore/test"
+	"github.com/quay/claircore/test/integration"
+	pgtest "github.com/quay/claircore/test/postgres"
+)
+
+var importFile *string
+
+func init() {
+	flag.Func(`load-file`, "run the integration test reading from `FILE` (must be zstd compressed)", func(v string) error {
+		importFile = &v
+		return nil
+	})
+}
+
+func TestMain(m *testing.M) {
+	var c int
+	defer func() { os.Exit(c) }()
+	defer integration.DBSetup()()
+	c = m.Run()
+}
+
+// TestLiveOfflineImport is meant to be used for profiling and testing the
+// system on a "real" export as produced by `clairctl`.
+func TestOfflineImport(t *testing.T) {
+	if importFile == nil {
+		t.Skip(`needed flag "-load-file" not provided`)
+	}
+	integration.NeedDB(t)
+
+	t.Run("Old", testOneOfflineImport(oldOfflineImport))
+	t.Run("New", testOneOfflineImport(OfflineImport))
+}
+
+func testOneOfflineImport(inner func(context.Context, *pgxpool.Pool, io.Reader) error) func(*testing.T) {
+	return func(t *testing.T) {
+		ctx := test.Logging(t)
+
+		f, err := os.Open(*importFile)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer f.Close()
+		zr, err := zstd.NewReader(f)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer zr.Close()
+
+		pool := pgtest.TestMatcherDB(ctx, t)
+		if err := inner(ctx, pool, zr); err != nil {
+			t.Error(err)
+		}
+	}
+}
+
+// OldOfflineImport is a copy of the previous implementation of [OfflineImport],
+// kept here for the above test.
+func oldOfflineImport(ctx context.Context, pool *pgxpool.Pool, in io.Reader) error {
+	s := postgres.NewMatcherStore(pool)
+	l, err := jsonblob.NewLoader(ctx, in)
+	if err != nil {
+		return err
+	}
+
+	ops, err := s.GetUpdateOperations(ctx, driver.VulnerabilityKind)
+	if err != nil {
+		return err
+	}
+
+Update:
+	for l.Next() {
+		e := l.Entry()
+		log := slog.With("updater", e.Updater)
+		for _, op := range ops[e.Updater] {
+			// This only helps if updaters don't keep something that
+			// changes in the fingerprint.
+			if op.Fingerprint == e.Fingerprint {
+				log.InfoContext(ctx, "fingerprint match, skipping")
+				continue Update
+			}
+		}
+		var ref uuid.UUID
+		if e.Enrichment != nil {
+			if ref, err = s.UpdateEnrichments(ctx, e.Updater, e.Fingerprint, e.Enrichment); err != nil {
+				return fmt.Errorf("updating enrichements: %w", err)
+			}
+		}
+		if e.Vuln != nil {
+			if ref, err = s.UpdateVulnerabilities(ctx, e.Updater, e.Fingerprint, e.Vuln); err != nil {
+				return fmt.Errorf("updating vulnerabilities: %w", err)
+			}
+		}
+		log.InfoContext(ctx, "update imported",
+			"ref", ref,
+			"vuln_count", len(e.Vuln),
+			"enrichment_count", len(e.Enrichment))
+	}
+	if err := l.Err(); err != nil {
+		return err
+	}
+	return nil
+}