jsonblob: add an iterator interface

This has the goal of reducing peak memory usage.

Signed-off-by: Hank Donnay <hdonnay@redhat.com>
Change-Id: I387d9289aad128724e57b9cebd213a3c6a6a6964

Author: hdonnay

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 one of the [Loader.Vulnerabilities] or
+// [Loader.Enrichments] iterators, 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,75 @@ 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 exactly one of the two values.
+func (l *Loader) All() iter.Seq2[*Entry, iter.Seq2[*claircore.Vulnerability, *driver.EnrichmentRecord]] {
+	if l.started {
+		l.err = errors.Join(l.err, fmt.Errorf("attempted to re-use a Loader"))
+		return noopIterator
+	}
+
+	// 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(v, e) {
+				return
+			}
+		}
+	}
+	// Outer reads the first entry and handles when the Ref changes.
+	outer := func(yield func(*Entry, iter.Seq2[*claircore.Vulnerability, *driver.EnrichmentRecord]) bool) {
+		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
+}
+
+// NoopIterator is an iterator that does nothing.
+func noopIterator(_ func(*Entry, iter.Seq2[*claircore.Vulnerability, *driver.EnrichmentRecord]) bool) {
+}
+
 // Err is the latest encountered error.
 func (l *Loader) Err() error {
 	// Don't report EOF as an error.
@@ -148,10 +228,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 +317,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 +330,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 +557,7 @@ func getBuf() []byte {
 	}
 	return b
 }
+
 func putBuf(b []byte) {
 	bufPool.Put(b)
 }

libvuln/jsonblob/pool.go

@@ -0,0 +1,44 @@
+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) {
+	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) {
+	enrichment.Put(e)
+}