postgres: alias insertion deadlock

datastore/postgres/updatevulnerabilities.go

@@ -6,7 +6,6 @@ import (
 	"crypto/md5"
 	"fmt"
 	"log/slog"
-	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -151,32 +150,48 @@ 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;`
+		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
+			unnest($1::TEXT[], $2::BYTEA[], $3::TEXT[], $4::TEXT[])
+				AS input(hash_kind, hash, alias_space, alias_name)
+		JOIN
+			vuln v ON v.hash_kind = input.hash_kind AND v.hash = input.hash
+		JOIN
+			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`
+		// insertAliasNamespaces creates all needed namespace rows outside any
+		// transaction so concurrent updaters do not deadlock.
 		insertAliasNamespaces = `INSERT INTO alias_namespace (namespace) VALUES (unnest($1::TEXT[])) ON CONFLICT DO NOTHING;`
-		insertAliases         = `INSERT INTO alias (namespace, name)
+		// insertAliases creates all needed alias rows outside any transaction.
+		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
-		FROM
-			(SELECT unnest($3::TEXT[]) AS space, unnest($4::TEXT[]) AS name) AS input
-		JOIN
-			alias_namespace AS ns ON ns.namespace = input.space
-		JOIN
-			alias AS a ON a.name = input.name AND a.namespace = ns.id
-		) AS alias
-ON CONFLICT DO NOTHING`
 	)
 
 	var uoID uint64
@@ -270,12 +285,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 +293,18 @@ 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
+	)
+	seenSpace := make(map[unique.Handle[string]]struct{})
+	seenAlias := make(map[claircore.Alias]struct{})
+
 	vulnIter(func(vuln *claircore.Vulnerability, iterErr error) bool {
 		if iterErr != nil {
 			err = iterErr
@@ -307,7 +328,8 @@ ON CONFLICT DO NOTHING`
 		hashKind, hash := md5Vuln(vuln)
 		vKind, _, _ := rangefmt(vuln.Range)
 
-		batch.Queue(insert,
+		batch.Queue(
+			insert,
 			hashKind, hash,
 			vuln.Name, vuln.Updater, vuln.Description, vuln.Issued, vuln.Links, vuln.Severity, vuln.NormalizedSeverity,
 			pkg.Name, pkg.Version, pkg.Module, pkg.Arch, pkg.Kind,
@@ -318,53 +340,27 @@ 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)
+			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)
 		}
 		if vuln.Self.Valid() {
-			batch.Queue(insertSelfAlias, hashKind, hash, vuln.Self.Space.Value(), vuln.Self.Name)
+			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)
 		}
 
 		if ct := batch.Len(); ct < 1000 {
@@ -386,6 +382,50 @@ 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)
+		}
+
+		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)
+		}
+	}
+	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)
 	}

datastore/postgres/updatevulnerabilities_test.go

@@ -1,7 +1,9 @@
 package postgres
 
 import (
+	"sync"
 	"testing"
+	"unique"
 
 	"github.com/google/uuid"
 
@@ -253,6 +255,10 @@ func TestGetLatestVulnerabilities(t *testing.T) {
 						Package: &claircore.Package{
 							Name: "jq",
 						},
+						Self: claircore.Alias{
+							Space: unique.Make("CVE"),
+							Name:  "000",
+						},
 					},
 				},
 			},
@@ -265,13 +271,21 @@ func TestGetLatestVulnerabilities(t *testing.T) {
 						Package: &claircore.Package{
 							Name: "jq-libs",
 						},
+						Self: claircore.Alias{
+							Space: unique.Make("CVE"),
+							Name:  "456",
+						},
 					},
 					{
 						Updater: updater,
 						Name:    "CVE-789",
 						Package: &claircore.Package{
 							Name: "jq-docs",
 						},
+						Self: claircore.Alias{
+							Space: unique.Make("CVE"),
+							Name:  "789",
+						},
 					},
 				},
 			},
@@ -321,3 +335,67 @@ func TestGetLatestVulnerabilities(t *testing.T) {
 		})
 	}
 }
+
+func TestUpdateVulnerabilitiesIterSinglePass(t *testing.T) {
+	integration.NeedDB(t)
+	ctx := test.Logging(t)
+
+	pool := pgtest.TestMatcherDB(ctx, t)
+	store := NewMatcherStore(pool)
+
+	vulns := []*claircore.Vulnerability{
+		{
+			Updater: t.Name(),
+			Name:    "CVE-2024-0001",
+			Package: &claircore.Package{Name: "test-pkg"},
+			Self:    claircore.Alias{Space: unique.Make("CVE"), Name: "CVE-2024-0001"},
+			Aliases: []claircore.Alias{
+				{Space: unique.Make("GHSA"), Name: "GHSA-xxxx-yyyy-zzzz"},
+			},
+		},
+		{
+			Updater: t.Name(),
+			Name:    "CVE-2024-0002",
+			Package: &claircore.Package{Name: "test-pkg-2"},
+			Self:    claircore.Alias{Space: unique.Make("CVE"), Name: "CVE-2024-0002"},
+			Aliases: []claircore.Alias{
+				{Space: unique.Make("GHSA"), Name: "GHSA-aaaa-bbbb-cccc"},
+			},
+		},
+	}
+
+	// Single-pass iterator: yields data only on the first call, mimicking
+	// jsonblob.RecordIter which streams from a compressed file.
+	var once sync.Once
+	singlePass := datastore.VulnerabilityIter(func(yield func(*claircore.Vulnerability, error) bool) {
+		once.Do(func() {
+			for _, v := range vulns {
+				if !yield(v, nil) {
+					return
+				}
+			}
+		})
+	})
+
+	_, err := store.UpdateVulnerabilitiesIter(ctx, t.Name(), driver.Fingerprint(uuid.New().String()), singlePass)
+	if err != nil {
+		t.Fatalf("UpdateVulnerabilitiesIter: %v", err)
+	}
+
+	var vulnCount int
+	if err := pool.QueryRow(ctx, `SELECT count(*) FROM vuln WHERE updater = $1`, t.Name()).Scan(&vulnCount); err != nil {
+		t.Fatalf("counting vulns: %v", err)
+	}
+	if vulnCount != len(vulns) {
+		t.Fatalf("vuln table has %d rows, want %d; single-pass iterator was likely exhausted before the main insertion loop", vulnCount, len(vulns))
+	}
+
+	var aliasCount int
+	if err := pool.QueryRow(ctx, `SELECT count(*) FROM alias`).Scan(&aliasCount); err != nil {
+		t.Fatalf("counting aliases: %v", err)
+	}
+	if aliasCount == 0 {
+		t.Fatal("alias table has 0 rows")
+	}
+	t.Logf("vuln rows: %d, alias rows: %d", vulnCount, aliasCount)
+}