diff --git a/datastore/postgres/updatevulnerabilities.go b/datastore/postgres/updatevulnerabilities.go index a9da5ff09..d605ca662 100644 --- a/datastore/postgres/updatevulnerabilities.go +++ b/datastore/postgres/updatevulnerabilities.go @@ -6,7 +6,6 @@ import ( "crypto/md5" "fmt" "log/slog" - "slices" "strconv" "strings" "time" @@ -151,34 +150,43 @@ func (s *MatcherStore) updateVulnerabilities(ctx context.Context, updater string $3, (SELECT id FROM vuln WHERE hash_kind = $1 AND hash = $2)) ON CONFLICT DO NOTHING;` - refreshView = `REFRESH MATERIALIZED VIEW CONCURRENTLY latest_update_operations;` - insertAliasNamespaces = `INSERT INTO alias_namespace (namespace) VALUES (unnest($1::TEXT[])) ON CONFLICT DO NOTHING;` - insertAliases = `INSERT INTO alias (namespace, name) - SELECT ns.id, input.name - FROM - (SELECT unnest($1::TEXT[]) AS space, unnest($2::TEXT[]) AS name) AS input - JOIN - alias_namespace AS ns ON input.space = ns.namespace -ON CONFLICT DO NOTHING;` - insertSelfAlias = `INSERT INTO vulnerability_self (vulnerability, self) VALUES ( - (SELECT id FROM vuln WHERE hash_kind = $1 AND hash = $2), - (SELECT a.id FROM alias AS a JOIN alias_namespace AS ns ON a.namespace = ns.id WHERE ns.namespace = $3 AND a.name = $4) -) ON CONFLICT DO NOTHING;` - insertVulnerabilityAliases = `INSERT INTO vulnerability_alias - SELECT vuln.id, alias.id - FROM - (SELECT id FROM vuln WHERE hash_kind = $1 AND hash = $2) AS vuln, - (SELECT a.id + refreshView = `REFRESH MATERIALIZED VIEW CONCURRENTLY latest_update_operations;` + // bulkLinkAliases links all vulnerability→alias rows in one statement by + // joining the flattened (hash_kind, hash, alias_space, alias_name) arrays + // against the already-populated vuln and alias tables. + bulkLinkAliases = ` + INSERT INTO vulnerability_alias (vulnerability, alias) + SELECT v.id, a.id FROM - (SELECT unnest($3::TEXT[]) AS space, unnest($4::TEXT[]) AS name) AS input + unnest($1::TEXT[], $2::BYTEA[], $3::TEXT[], $4::TEXT[]) + AS input(hash_kind, hash, alias_space, alias_name) JOIN - alias_namespace AS ns ON ns.namespace = input.space + vuln v ON v.hash_kind = input.hash_kind AND v.hash = input.hash JOIN - alias AS a ON a.name = input.name AND a.namespace = ns.id - ) AS alias -ON CONFLICT DO NOTHING` + alias_namespace ns ON ns.namespace = input.alias_space + JOIN + alias a ON a.name = input.alias_name AND a.namespace = ns.id + ON CONFLICT DO NOTHING` + // bulkLinkSelf links all vulnerability→self rows in one statement. + bulkLinkSelf = ` + INSERT INTO vulnerability_self (vulnerability, self) + SELECT v.id, a.id + FROM + unnest($1::TEXT[], $2::BYTEA[], $3::TEXT[], $4::TEXT[]) + AS input(hash_kind, hash, self_space, self_name) + JOIN + vuln v ON v.hash_kind = input.hash_kind AND v.hash = input.hash + JOIN + alias_namespace ns ON ns.namespace = input.self_space + JOIN + alias a ON a.name = input.self_name AND a.namespace = ns.id + ON CONFLICT DO NOTHING` ) + if err := s.insertAliases(ctx, vulnIter); err != nil { + return uuid.Nil, err + } + var uoID uint64 var ref uuid.UUID @@ -270,12 +278,6 @@ ON CONFLICT DO NOTHING` skipCt := 0 vulnCt := 0 start = time.Now() - - // Keep track of the alias namespaces seen in this transaction. - // - // There's a limited set of namespaces, but it's either do this book keeping - // or issue a LOT of pointless inserts. - seenSpace := make(map[unique.Handle[string]]struct{}) var batch pgx.Batch flush := func() (err error) { err = tx.SendBatch(ctx, &batch).Close() @@ -284,6 +286,16 @@ ON CONFLICT DO NOTHING` 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. + var ( + vaHashKinds, vsHashKinds []string + vaHashes, vsHashes [][]byte + vaSpaces, vsSpaces []string + vaNames, vsNames []string + ) + vulnIter(func(vuln *claircore.Vulnerability, iterErr error) bool { if iterErr != nil { err = iterErr @@ -318,53 +330,23 @@ ON CONFLICT DO NOTHING` ) batch.Queue(assoc, hashKind, hash, uoID) - // TODO(hank) Switch the fail-open validity checks to errors. - aliases := append(vuln.Aliases, vuln.Self) - spaces := slices.Collect(func(yield func(string) bool) { - for _, a := range aliases { - if !a.Valid() { - continue - } - if !yield(a.Space.Value()) { - return - } - } - }) - names := slices.Collect(func(yield func(string) bool) { - for _, a := range aliases { - if !a.Valid() { - continue - } - if !yield(a.Name) { - return - } - } - }) - uniq := slices.Collect(func(yield func(string) bool) { - for _, a := range aliases { - if !a.Valid() { - continue - } - s := a.Space - if _, ok := seenSpace[s]; !ok { - if !yield(s.Value()) { - return - } - seenSpace[s] = struct{}{} - } + // 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. + for _, a := range vuln.Aliases { + if !a.Valid() { + continue } - }) - if len(uniq) > 0 { - batch.Queue(insertAliasNamespaces, uniq) - } - // TODO(hank) Remove these conditionals and assume that aliases are - // always present. - if len(names) > 0 { - batch.Queue(insertAliases, spaces, names) - batch.Queue(insertVulnerabilityAliases, hashKind, hash, spaces, names) + vaHashKinds = append(vaHashKinds, hashKind) + vaHashes = append(vaHashes, hash) + vaSpaces = append(vaSpaces, a.Space.Value()) + vaNames = append(vaNames, a.Name) } if vuln.Self.Valid() { - batch.Queue(insertSelfAlias, hashKind, hash, vuln.Self.Space.Value(), vuln.Self.Name) + vsHashKinds = append(vsHashKinds, hashKind) + vsHashes = append(vsHashes, hash) + vsSpaces = append(vsSpaces, vuln.Self.Space.Value()) + vsNames = append(vsNames, vuln.Self.Name) } if ct := batch.Len(); ct < 1000 { @@ -386,6 +368,22 @@ 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()) + // 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) + } + } + updateVulnerabilitiesCounter.WithLabelValues("link_aliases", strconv.FormatBool(delta)).Add(1) + updateVulnerabilitiesDuration.WithLabelValues("link_aliases", strconv.FormatBool(delta)).Observe(time.Since(start).Seconds()) + if err := tx.Commit(ctx); err != nil { return uuid.Nil, fmt.Errorf("failed to commit transaction: %w", err) } @@ -484,3 +482,85 @@ func rangefmt(r *claircore.Range) (kind *string, lower, upper string) { return kind, lower, upper } + +// insertAliases creates all needed alias namespaces and aliases outside of any +// transaction. +// +// Running outside a transaction avoids the deadlock that arises when concurrent +// updaters race to insert the same alias namespaces inside a serialised +// transaction. Each statement here auto-commits immediately, so no row-level +// locks are held between them. +// +// BUG(hank) The alias insertion code does not remove aliases in the case of a +// failure during the update operation. +func (s *MatcherStore) insertAliases(ctx context.Context, vulnIter datastore.VulnerabilityIter) error { + const ( + insertAliasNamespaces = `INSERT INTO alias_namespace (namespace) VALUES (unnest($1::TEXT[])) ON CONFLICT DO NOTHING;` + insertAliases = `INSERT INTO alias (namespace, name) + SELECT ns.id, input.name + FROM + (SELECT unnest($1::TEXT[]) AS space, unnest($2::TEXT[]) AS name) AS input + JOIN + alias_namespace AS ns ON input.space = ns.namespace +ON CONFLICT DO NOTHING;` + ) + + // Collect unique namespaces and aliases in one pass using only the dedup + // maps, then derive the slices the bulk statements need from the maps after + // iteration. This avoids holding the maps and parallel slices in memory at + // the same time. + seenSpace := make(map[unique.Handle[string]]struct{}) + seenAlias := make(map[claircore.Alias]struct{}) + var iterErr error + vulnIter(func(vuln *claircore.Vulnerability, err error) bool { + if err != nil { + iterErr = err + return false + } + for _, a := range vuln.Aliases { + if !a.Valid() { + continue + } + seenSpace[a.Space] = struct{}{} + seenAlias[a] = struct{}{} + } + if vuln.Self.Valid() { + seenSpace[vuln.Self.Space] = struct{}{} + seenAlias[vuln.Self] = struct{}{} + } + return true + }) + if iterErr != nil { + return iterErr + } + if len(seenSpace) == 0 { + return nil + } + + spaces := make([]string, 0, len(seenSpace)) + for h := range seenSpace { + spaces = append(spaces, h.Value()) + } + // aliasSpaces and aliasNames are parallel: aliasSpaces[i] is the namespace + // for aliasNames[i], matching what the unnest query expects. + 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) + } + + conn, err := s.pool.Acquire(ctx) + if err != nil { + return fmt.Errorf("acquiring connection for aliases: %w", err) + } + defer conn.Release() + + if _, err := conn.Exec(ctx, insertAliasNamespaces, spaces); err != nil { + return fmt.Errorf("failed to insert alias namespaces: %w", err) + } + if _, err := conn.Exec(ctx, insertAliases, aliasSpaces, aliasNames); err != nil { + return fmt.Errorf("failed to insert aliases: %w", err) + } + return nil +} diff --git a/datastore/postgres/updatevulnerabilities_test.go b/datastore/postgres/updatevulnerabilities_test.go index 09b23906d..0bcdff6f1 100644 --- a/datastore/postgres/updatevulnerabilities_test.go +++ b/datastore/postgres/updatevulnerabilities_test.go @@ -1,9 +1,15 @@ package postgres import ( + "iter" + "slices" + "strconv" + "sync" "testing" + "unique" "github.com/google/uuid" + "golang.org/x/sync/errgroup" "github.com/quay/claircore" "github.com/quay/claircore/datastore" @@ -253,6 +259,10 @@ func TestGetLatestVulnerabilities(t *testing.T) { Package: &claircore.Package{ Name: "jq", }, + Self: claircore.Alias{ + Space: unique.Make("CVE"), + Name: "000", + }, }, }, }, @@ -265,6 +275,10 @@ func TestGetLatestVulnerabilities(t *testing.T) { Package: &claircore.Package{ Name: "jq-libs", }, + Self: claircore.Alias{ + Space: unique.Make("CVE"), + Name: "456", + }, }, { Updater: updater, @@ -272,6 +286,10 @@ func TestGetLatestVulnerabilities(t *testing.T) { Package: &claircore.Package{ Name: "jq-docs", }, + Self: claircore.Alias{ + Space: unique.Make("CVE"), + Name: "789", + }, }, }, }, @@ -321,3 +339,78 @@ func TestGetLatestVulnerabilities(t *testing.T) { }) } } + +func TestConcurrentUpdateVulnerabilities(t *testing.T) { + integration.NeedDB(t) + ctx := test.Logging(t) + + pool := pgtest.TestMatcherDB(ctx, t) + store := NewMatcherStore(pool) + eg, ctx := errgroup.WithContext(ctx) + + mkVuln := func(tn int) iter.Seq[*claircore.Vulnerability] { + return func(yield func(*claircore.Vulnerability) bool) { + const sz = 500 + for n := range sz { + seq := func(yield func(claircore.Alias) bool) { + for i := range 100 { + if i == n { + continue + } + a := claircore.Alias{ + Space: unique.Make("TEST"), + Name: strconv.Itoa(i), + } + if !yield(a) { + return + } + } + } + s := claircore.Alias{ + Space: unique.Make("TEST"), + Name: strconv.Itoa(n), + } + v := claircore.Vulnerability{ + Updater: t.Name() + strconv.Itoa(tn), + Name: s.String(), + Description: t.Name() + strconv.Itoa(tn) + strconv.Itoa(n), + Package: &claircore.Package{ + Name: "test", + }, + Self: s, + Aliases: slices.Collect(seq), + } + if !yield(&v) { + return + } + } + } + } + start := make(chan struct{}) + var ready sync.WaitGroup + runfunc := func(n int, start chan struct{}) func() error { + fp := driver.Fingerprint(uuid.New().String()) + vulns := slices.Collect(mkVuln(n)) + if n%2 == 0 { + slices.Reverse(vulns) + } + ready.Done() + return func() error { + <-start + _, err := store.UpdateVulnerabilities(ctx, t.Name()+strconv.Itoa(n), fp, vulns) + return err + } + } + + const n = 2 + ready.Add(n) + for n := range n { + eg.Go(runfunc(n, start)) + } + + ready.Wait() + close(start) + if err := eg.Wait(); err != nil { + t.Error(err) + } +}