diff --git a/datastore/postgres/updatevulnerabilities.go b/datastore/postgres/updatevulnerabilities.go index 3009e25fd..ae1ddc728 100644 --- a/datastore/postgres/updatevulnerabilities.go +++ b/datastore/postgres/updatevulnerabilities.go @@ -196,7 +196,11 @@ ON CONFLICT DO NOTHING;` start := time.Now() - tx, err := s.pool.Begin(ctx) + // The isolation level must be pinned to "read committed" (rather than + // inheriting default_transaction_isolation) because linkFlush's + // INSERT..SELECT statements need to see alias rows committed outside this + // transaction after it began. + tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) if err != nil { return uuid.Nil, fmt.Errorf("unable to start transaction: %w", err) } @@ -285,23 +289,66 @@ ON CONFLICT DO NOTHING;` var batch pgx.Batch flush := func() (err error) { err = tx.SendBatch(ctx, &batch).Close() - clear(batch.QueuedQueries) - batch.QueuedQueries = batch.QueuedQueries[:0] + resetSlices(&batch.QueuedQueries) return err } - // Flattened parallel arrays for the bulk alias-linking statements run after - // all vuln inserts are done. Each entry in va* corresponds to one - // (vuln, alias) pair; each entry in vs* to one (vuln, self) pair. + // Columns for the bulk alias-linking statements, flushed alongside the + // vulnerability insert batch to bound memory usage. Each entry in va + // corresponds to one (vuln, alias) pair; each entry in vs to one + // (vuln, self) pair. + var va, vs linkCols + // Namespaces and aliases not yet inserted, i.e., first seen in the + // current chunk. var ( - vaHashKinds, vsHashKinds []string - vaHashes, vsHashes [][]byte - vaSpaces, vsSpaces []string - vaNames, vsNames []string + newSpaces []string + newAliasSpaces, newAliasNames []string ) seenSpace := make(map[unique.Handle[string]]struct{}) seenAlias := make(map[claircore.Alias]struct{}) + // linkDur accumulates time spent in linkFlush so the link_aliases metric + // covers the per-chunk flushes and insert_batch excludes them. + var linkDur time.Duration + + // linkFlush inserts the alias namespaces and aliases first seen in the + // current chunk (outside the transaction to avoid deadlocks between + // concurrent updaters), then links the current chunk's (vuln, alias) and + // (vuln, self) pairs inside the transaction. The vulnerability rows of + // the chunk must have been flushed to the transaction already. + linkFlush := func() error { + defer func(t time.Time) { linkDur += time.Since(t) }(time.Now()) + if len(newSpaces) > 0 { + if _, err := s.pool.Exec(ctx, insertAliasNamespaces, newSpaces); err != nil { + return fmt.Errorf("failed to insert alias namespaces: %w", err) + } + } + if len(newAliasNames) > 0 { + if _, err := s.pool.Exec(ctx, insertAliases, newAliasSpaces, newAliasNames); err != nil { + return fmt.Errorf("failed to insert aliases: %w", err) + } + } + // Enforce using uncached plans because with more than a few chunks, + // the generic plan (built for unnest's default row estimate) is + // chosen and is far slower for the actual array sizes. + if len(va.hashKinds) > 0 { + if _, err := tx.Exec(ctx, bulkLinkAliases, pgx.QueryExecModeExec, va.hashKinds, va.hashes, va.spaces, va.names); err != nil { + return fmt.Errorf("failed to bulk link vulnerability aliases: %w", err) + } + } + if len(vs.hashKinds) > 0 { + if _, err := tx.Exec(ctx, bulkLinkSelf, pgx.QueryExecModeExec, vs.hashKinds, vs.hashes, vs.spaces, vs.names); err != nil { + return fmt.Errorf("failed to bulk link vulnerability self aliases: %w", err) + } + } + clear(seenSpace) + clear(seenAlias) + resetSlices(&newSpaces, &newAliasSpaces, &newAliasNames) + va.reset() + vs.reset() + return nil + } + vulnIter(func(vuln *claircore.Vulnerability, iterErr error) bool { if iterErr != nil { err = iterErr @@ -337,27 +384,28 @@ ON CONFLICT DO NOTHING;` ) batch.Queue(assoc, hashKind, hash, uoID) - // Accumulate alias links for the bulk statements below. The hash is - // repeated once per alias so the unnest join can match each row to its - // vuln. + // Accumulate alias links for the bulk statements below. + stage := func(a claircore.Alias) { + if _, ok := seenSpace[a.Space]; !ok { + seenSpace[a.Space] = struct{}{} + newSpaces = append(newSpaces, a.Space.Value()) + } + if _, ok := seenAlias[a]; !ok { + seenAlias[a] = struct{}{} + newAliasSpaces = append(newAliasSpaces, a.Space.Value()) + newAliasNames = append(newAliasNames, a.Name) + } + } for _, a := range vuln.Aliases { if !a.Valid() { continue } - seenSpace[a.Space] = struct{}{} - seenAlias[a] = struct{}{} - vaHashKinds = append(vaHashKinds, hashKind) - vaHashes = append(vaHashes, hash) - vaSpaces = append(vaSpaces, a.Space.Value()) - vaNames = append(vaNames, a.Name) + stage(a) + va.add(hashKind, hash, a) } if vuln.Self.Valid() { - seenSpace[vuln.Self.Space] = struct{}{} - seenAlias[vuln.Self] = struct{}{} - vsHashKinds = append(vsHashKinds, hashKind) - vsHashes = append(vsHashes, hash) - vsSpaces = append(vsSpaces, vuln.Self.Space.Value()) - vsNames = append(vsNames, vuln.Self.Name) + stage(vuln.Self) + vs.add(hashKind, hash, vuln.Self) } if ct := batch.Len(); ct < 1000 { @@ -367,6 +415,9 @@ ON CONFLICT DO NOTHING;` err = fmt.Errorf("failed batching: %w", err) return false } + if err = linkFlush(); err != nil { + return false + } return true }) if err != nil { @@ -377,51 +428,14 @@ ON CONFLICT DO NOTHING;` } updateVulnerabilitiesCounter.WithLabelValues("insert_batch", strconv.FormatBool(delta)).Add(1) - updateVulnerabilitiesDuration.WithLabelValues("insert_batch", strconv.FormatBool(delta)).Observe(time.Since(start).Seconds()) - - // Insert alias namespaces and aliases outside the transaction to avoid - // deadlocks when concurrent updaters race to insert the same namespaces. - if len(seenSpace) > 0 { - spaces := make([]string, 0, len(seenSpace)) - for h := range seenSpace { - spaces = append(spaces, h.Value()) - } - aliasSpaces := make([]string, 0, len(seenAlias)) - aliasNames := make([]string, 0, len(seenAlias)) - for a := range seenAlias { - aliasSpaces = append(aliasSpaces, a.Space.Value()) - aliasNames = append(aliasNames, a.Name) - } + updateVulnerabilitiesDuration.WithLabelValues("insert_batch", strconv.FormatBool(delta)).Observe((time.Since(start) - linkDur).Seconds()) - conn, err := s.pool.Acquire(ctx) - if err != nil { - return uuid.Nil, fmt.Errorf("acquiring connection for aliases: %w", err) - } - defer conn.Release() - - if _, err := conn.Exec(ctx, insertAliasNamespaces, spaces); err != nil { - return uuid.Nil, fmt.Errorf("failed to insert alias namespaces: %w", err) - } - if _, err := conn.Exec(ctx, insertAliases, aliasSpaces, aliasNames); err != nil { - return uuid.Nil, fmt.Errorf("failed to insert aliases: %w", err) - } - } - - // Bulk-link aliases and self references. Two single statements replace the - // former per-vulnerability hash-lookup subqueries queued in the batch above. - start = time.Now() - if len(vaHashKinds) > 0 { - if _, err := tx.Exec(ctx, bulkLinkAliases, vaHashKinds, vaHashes, vaSpaces, vaNames); err != nil { - return uuid.Nil, fmt.Errorf("failed to bulk link vulnerability aliases: %w", err) - } - } - if len(vsHashKinds) > 0 { - if _, err := tx.Exec(ctx, bulkLinkSelf, vsHashKinds, vsHashes, vsSpaces, vsNames); err != nil { - return uuid.Nil, fmt.Errorf("failed to bulk link vulnerability self aliases: %w", err) - } + // Link any remainder not covered by the periodic flushes above. + if err := linkFlush(); err != nil { + return uuid.Nil, err } updateVulnerabilitiesCounter.WithLabelValues("link_aliases", strconv.FormatBool(delta)).Add(1) - updateVulnerabilitiesDuration.WithLabelValues("link_aliases", strconv.FormatBool(delta)).Observe(time.Since(start).Seconds()) + updateVulnerabilitiesDuration.WithLabelValues("link_aliases", strconv.FormatBool(delta)).Observe(linkDur.Seconds()) if err := tx.Commit(ctx); err != nil { return uuid.Nil, fmt.Errorf("failed to commit transaction: %w", err) @@ -436,6 +450,38 @@ ON CONFLICT DO NOTHING;` return ref, nil } +// LinkCols is the set of flattened parallel arrays passed to one of the bulk +// alias-linking statements: one entry per (vuln, alias) pair. The hash is +// repeated once per alias so the unnest join can match each row to its vuln. +type linkCols struct { + hashKinds []string + hashes [][]byte + spaces []string + names []string +} + +func (l *linkCols) add(hashKind string, hash []byte, a claircore.Alias) { + l.hashKinds = append(l.hashKinds, hashKind) + l.hashes = append(l.hashes, hash) + l.spaces = append(l.spaces, a.Space.Value()) + l.names = append(l.names, a.Name) +} + +func (l *linkCols) reset() { + resetSlices(&l.hashKinds, &l.spaces, &l.names) + resetSlices(&l.hashes) +} + +// ResetSlices truncates the passed slices, keeping their backing arrays for +// reuse but clearing the elements so anything they reference can be +// collected. +func resetSlices[T any](slices ...*[]T) { + for _, s := range slices { + clear(*s) + *s = (*s)[:0] + } +} + // SkipVulnerability reports if the provided [claircore.Vulnerability] should // not be uploaded to the database. func skipVulnerability(v *claircore.Vulnerability) bool { diff --git a/datastore/postgres/updatevulnerabilities_benchmark_test.go b/datastore/postgres/updatevulnerabilities_benchmark_test.go new file mode 100644 index 000000000..29e263dc8 --- /dev/null +++ b/datastore/postgres/updatevulnerabilities_benchmark_test.go @@ -0,0 +1,74 @@ +package postgres + +import ( + "runtime" + "runtime/metrics" + "strconv" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/quay/claircore/libvuln/driver" + "github.com/quay/claircore/test" + "github.com/quay/claircore/test/integration" + pgtest "github.com/quay/claircore/test/postgres" +) + +// liveHeapBytes reports bytes occupied by live (and not-yet-collected) heap +// objects without stopping the world. +func liveHeapBytes(s []metrics.Sample) uint64 { + metrics.Read(s) + return s[0].Value.Uint64() +} + +func Benchmark_UpdateVulnerabilities(b *testing.B) { + integration.NeedDB(b) + // Consider using `-benchtime 1x` when running the 50000, 75000, and 100000. + for _, sz := range []int{100, 500, 1200, 50000, 75000, 100000} { + b.Run(strconv.Itoa(sz)+" vulnerabilities", func(b *testing.B) { + ctx := test.Logging(b) + pool := pgtest.TestMatcherDB(ctx, b) + store := NewMatcherStore(pool) + vulns := genAliasVulns(b.Name(), sz) + + // Sample live heap during the run: B/op reports cumulative + // allocation and misses how long allocations stay reachable, + // which is what the chunked alias flush is meant to bound. + sample := []metrics.Sample{{Name: "/memory/classes/heap/objects:bytes"}} + runtime.GC() + base := liveHeapBytes(sample) + var peak atomic.Uint64 + done := make(chan struct{}) + defer close(done) + go func() { + tick := time.NewTicker(25 * time.Millisecond) + defer tick.Stop() + sample := []metrics.Sample{{Name: "/memory/classes/heap/objects:bytes"}} + for { + select { + case <-done: + return + case <-tick.C: + if h := liveHeapBytes(sample); h > peak.Load() { + peak.Store(h) + } + } + } + }() + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + if _, err := store.UpdateVulnerabilities(ctx, b.Name(), driver.Fingerprint(uuid.New().String()), vulns); err != nil { + b.Fatalf("UpdateVulnerabilities: %v", err) + } + } + b.StopTimer() + if p := peak.Load(); p > base { + b.ReportMetric(float64(p-base), "peak-heap-B") + } + }) + } +} diff --git a/datastore/postgres/updatevulnerabilities_test.go b/datastore/postgres/updatevulnerabilities_test.go index 505e21614..c066401bb 100644 --- a/datastore/postgres/updatevulnerabilities_test.go +++ b/datastore/postgres/updatevulnerabilities_test.go @@ -1,6 +1,7 @@ package postgres import ( + "fmt" "sync" "testing" "unique" @@ -399,3 +400,80 @@ func TestUpdateVulnerabilitiesIterSinglePass(t *testing.T) { } t.Logf("vuln rows: %d, alias rows: %d", vulnCount, aliasCount) } + +// SharedAliasName is attached to every vulnerability generated by +// genAliasVulns, so it recurs in every flush chunk. +const sharedAliasName = "GHSA-shared-0000-0000" + +func genAliasVulns(updater string, n int) []*claircore.Vulnerability { + shared := claircore.Alias{Space: unique.Make("GHSA"), Name: sharedAliasName} + vulns := make([]*claircore.Vulnerability, n) + for i := range vulns { + name := fmt.Sprintf("CVE-2024-%04d", i) + vulns[i] = &claircore.Vulnerability{ + Updater: updater, + Name: name, + Package: &claircore.Package{Name: "test-pkg"}, + Self: claircore.Alias{Space: unique.Make("CVE"), Name: name}, + Aliases: []claircore.Alias{ + shared, + {Space: unique.Make("GHSA"), Name: "GHSA-" + name}, + }, + } + } + return vulns +} + +func TestUpdateVulnerabilitiesChunked(t *testing.T) { + integration.NeedDB(t) + ctx := test.Logging(t) + + pool := pgtest.TestMatcherDB(ctx, t) + store := NewMatcherStore(pool) + + // The insert batch flushes every 500 vulnerabilities (two queued queries + // per vulnerability, flushed at 1000), so 1200 crosses two chunk + // boundaries and leaves a remainder for the final flush. + const vulnCt = 1200 + vulns := genAliasVulns(t.Name(), vulnCt) + + if _, err := store.UpdateVulnerabilities(ctx, t.Name(), driver.Fingerprint(uuid.New().String()), vulns); err != nil { + t.Fatalf("UpdateVulnerabilities: %v", err) + } + + checks := []struct { + desc string + query string + want int + }{ + {"vuln rows", `SELECT count(*) FROM vuln WHERE updater = $1`, vulnCt}, + {"self links", `SELECT count(*) FROM vulnerability_self s JOIN vuln v ON s.vulnerability = v.id WHERE v.updater = $1`, vulnCt}, + // One shared alias plus one unique alias per vulnerability. + {"alias links", `SELECT count(*) FROM vulnerability_alias a JOIN vuln v ON a.vulnerability = v.id WHERE v.updater = $1`, 2 * vulnCt}, + } + for _, c := range checks { + var got int + if err := pool.QueryRow(ctx, c.query, t.Name()).Scan(&got); err != nil { + t.Fatalf("counting %s: %v", c.desc, err) + } + if got != c.want { + t.Errorf("%s: got %d, want %d", c.desc, got, c.want) + } + } + + // The shared alias row is created during the first chunk; vulnerabilities + // in later chunks must still link to it. + var sharedCt int + err := pool.QueryRow(ctx, ` + SELECT count(*) + FROM vulnerability_alias va + JOIN vuln v ON va.vulnerability = v.id + JOIN alias a ON va.alias = a.id + WHERE v.updater = $1 AND a.name = $2`, t.Name(), sharedAliasName).Scan(&sharedCt) + if err != nil { + t.Fatalf("counting shared alias links: %v", err) + } + if sharedCt != vulnCt { + t.Errorf("shared alias links: got %d, want %d", sharedCt, vulnCt) + } +}