Skip to content

datastore/postgres: fix iterator exhaustion in alias insertion #89

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jvdm wants to merge 1 commit into
base: faster-alias-insert
Choose a base branch
from
Open

datastore/postgres: fix iterator exhaustion in alias insertion #89

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 45 additions & 85 deletions datastore/postgres/updatevulnerabilities.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,12 +181,19 @@ func (s *MatcherStore) updateVulnerabilities(ctx context.Context, updater string
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 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;`
)

if err := s.insertAliases(ctx, vulnIter); err != nil {
return uuid.Nil, err
}

var uoID uint64
var ref uuid.UUID

Expand Down Expand Up @@ -295,6 +302,8 @@ func (s *MatcherStore) updateVulnerabilities(ctx context.Context, updater string
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 {
Expand Down Expand Up @@ -337,12 +346,16 @@ func (s *MatcherStore) updateVulnerabilities(ctx context.Context, updater string
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)
}
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())
Expand All @@ -368,6 +381,34 @@ func (s *MatcherStore) updateVulnerabilities(ctx context.Context, updater string
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()
Expand Down Expand Up @@ -483,84 +524,3 @@ 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
}
64 changes: 64 additions & 0 deletions datastore/postgres/updatevulnerabilities_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -414,3 +414,67 @@ func TestConcurrentUpdateVulnerabilities(t *testing.T) {
t.Error(err)
}
}

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)
}