datastore: gc enrichments

Signed-off-by: RTann <rtannenb@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED

Author: RTann

datastore/postgres/gc.go

@@ -15,6 +15,8 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/quay/zlog"
 	"golang.org/x/sync/semaphore"
+
+	"github.com/quay/claircore/libvuln/driver"
 )
 
 var (
@@ -54,6 +56,8 @@ const (
 // If a full GC is required run this method until the returned int64 value
 // is 0.
 func (s *MatcherStore) GC(ctx context.Context, keep int) (int64, error) {
+	ctx = zlog.ContextWithValues(ctx, "component", "datastore/postgres/GC")
+
 	// obtain update operations which need deletin'
 	ops, totalOps, err := eligibleUpdateOpts(ctx, s.pool, keep)
 	if err != nil {
@@ -81,20 +85,32 @@ func (s *MatcherStore) GC(ctx context.Context, keep int) (int64, error) {
 	cpus := int64(runtime.GOMAXPROCS(0))
 	sem := semaphore.NewWeighted(cpus)
 
-	errC := make(chan error, len(updaters))
+	errC := make(chan error, len(updaters[driver.VulnerabilityKind])+len(updaters[driver.EnrichmentKind]))
 
-	for _, updater := range updaters {
-		err = sem.Acquire(ctx, 1)
-		if err != nil {
-			break
+	for kind, us := range updaters {
+		var cleanup cleanupFunc
+		switch kind {
+		case driver.VulnerabilityKind:
+			cleanup = vulnCleanup
+		case driver.EnrichmentKind:
+			cleanup = enrichmentCleanup
+		default:
+			zlog.Error(ctx).Str("kind", string(kind)).Msg("unknown updater kind; skipping cleanup")
+			continue
 		}
-		go func(u string) {
-			defer sem.Release(1)
-			err := vulnCleanup(ctx, s.pool, u)
+		for _, u := range us {
+			err = sem.Acquire(ctx, 1)
 			if err != nil {
-				errC <- err
+				break
 			}
-		}(updater)
+			go func(cleanup cleanupFunc, u string) {
+				defer sem.Release(1)
+				err := cleanup(ctx, s.pool, u)
+				if err != nil {
+					errC <- err
+				}
+			}(cleanup, u)
+		}
 	}
 
 	// unconditionally wait for all in-flight go routines to return.
@@ -116,11 +132,11 @@ func (s *MatcherStore) GC(ctx context.Context, keep int) (int64, error) {
 
 // distinctUpdaters returns all updaters which have registered an update
 // operation.
-func distinctUpdaters(ctx context.Context, pool *pgxpool.Pool) ([]string, error) {
+func distinctUpdaters(ctx context.Context, pool *pgxpool.Pool) (map[driver.UpdateKind][]string, error) {
 	const (
 		// will always contain at least two update operations
 		selectUpdaters = `
-SELECT DISTINCT(updater) FROM update_operation;
+SELECT DISTINCT(updater), kind FROM update_operation;
 `
 	)
 	rows, err := pool.Query(ctx, selectUpdaters)
@@ -129,17 +145,20 @@ SELECT DISTINCT(updater) FROM update_operation;
 	}
 	defer rows.Close()
 
-	var updaters []string
+	updaters := make(map[driver.UpdateKind][]string)
 	for rows.Next() {
-		var updater string
-		err := rows.Scan(&updater)
+		var (
+			updater string
+			kind    driver.UpdateKind
+		)
+		err := rows.Scan(&updater, &kind)
 		switch err {
 		case nil:
 			// hop out
 		default:
 			return nil, fmt.Errorf("error scanning updater: %v", err)
 		}
-		updaters = append(updaters, updater)
+		updaters[kind] = append(updaters[kind], updater)
 	}
 	if rows.Err() != nil {
 		return nil, rows.Err()
@@ -198,6 +217,8 @@ WHERE array_length(ordered_ops.refs, 1) > $2;
 	return m, int64(len(m)), nil
 }
 
+type cleanupFunc func(context.Context, *pgxpool.Pool, string) error
+
 func vulnCleanup(ctx context.Context, pool *pgxpool.Pool, updater string) error {
 	const (
 		deleteOrphanedVulns = `
@@ -214,7 +235,7 @@ AND v1.id = v2.id;
 	start := time.Now()
 	ctx = zlog.ContextWithValues(ctx, "updater", updater)
 	zlog.Debug(ctx).
-		Msg("starting clean up")
+		Msg("starting vuln clean up")
 	res, err := pool.Exec(ctx, deleteOrphanedVulns, updater)
 	if err != nil {
 		gcCounter.WithLabelValues("deleteVulns", "false").Inc()
@@ -226,3 +247,32 @@ AND v1.id = v2.id;
 
 	return nil
 }
+
+func enrichmentCleanup(ctx context.Context, pool *pgxpool.Pool, updater string) error {
+	const (
+		deleteOrphanedEnrichments = `
+DELETE FROM enrichment e1 USING
+	enrichment e2
+	LEFT JOIN uo_enrich uen
+		ON e2.id = uen.enrich
+	WHERE uen.enrich IS NULL
+	AND e2.updater = $1
+	AND e1.id = e2.id;
+`
+	)
+
+	start := time.Now()
+	ctx = zlog.ContextWithValues(ctx, "updater", updater)
+	zlog.Debug(ctx).
+		Msg("starting enrichment clean up")
+	res, err := pool.Exec(ctx, deleteOrphanedEnrichments, updater)
+	if err != nil {
+		gcCounter.WithLabelValues("deleteEnrichments", "false").Inc()
+		return fmt.Errorf("failed while exec'ing enrichment delete: %w", err)
+	}
+	zlog.Debug(ctx).Int64("rows affected", res.RowsAffected()).Msg("enrichments deleted")
+	gcCounter.WithLabelValues("deleteEnrichments", "true").Inc()
+	gcDuration.WithLabelValues("deleteEnrichments").Observe(time.Since(start).Seconds())
+
+	return nil
+}

datastore/postgres/gc_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/hex"
+	"encoding/json"
 	"io"
 	"net/http"
 	"testing"
@@ -19,6 +20,8 @@ import (
 	pgtest "github.com/quay/claircore/test/postgres"
 )
 
+var _ driver.Updater = (*updaterMock)(nil)
+
 type updaterMock struct {
 	_name  func() string
 	_fetch func(_ context.Context, _ driver.Fingerprint) (io.ReadCloser, driver.Fingerprint, error)
@@ -37,6 +40,30 @@ func (u *updaterMock) Parse(ctx context.Context, contents io.ReadCloser) ([]*cla
 	return u._parse(ctx, contents)
 }
 
+var (
+	_ driver.Updater           = (*enricherMock)(nil)
+	_ driver.EnrichmentUpdater = (*enricherMock)(nil)
+)
+
+type enricherMock struct {
+	driver.NoopUpdater
+	_name  func() string
+	_fetch func(_ context.Context, _ driver.Fingerprint) (io.ReadCloser, driver.Fingerprint, error)
+	_parse func(ctx context.Context, contents io.ReadCloser) ([]driver.EnrichmentRecord, error)
+}
+
+func (e enricherMock) Name() string {
+	return e._name()
+}
+
+func (e enricherMock) FetchEnrichment(ctx context.Context, fingerprint driver.Fingerprint) (io.ReadCloser, driver.Fingerprint, error) {
+	return e._fetch(ctx, fingerprint)
+}
+
+func (e enricherMock) ParseEnrichment(ctx context.Context, contents io.ReadCloser) ([]driver.EnrichmentRecord, error) {
+	return e._parse(ctx, contents)
+}
+
 // TestGC confirms the garbage collection of
 // vulnerabilities works correctly.
 func TestGC(t *testing.T) {
@@ -60,6 +87,23 @@ func TestGC(t *testing.T) {
 		},
 	}
 
+	// mock returns exactly one random enrichment each time its Parse method is called.
+	// each update operation will be associated with a single enrichment.
+	mockEnrich := &enricherMock{
+		_name: func() string { return "MockEnrichmentUpdater" },
+		_fetch: func(_ context.Context, _ driver.Fingerprint) (io.ReadCloser, driver.Fingerprint, error) {
+			return nil, "", nil
+		},
+		_parse: func(ctx context.Context, contents io.ReadCloser) ([]driver.EnrichmentRecord, error) {
+			return []driver.EnrichmentRecord{
+				{
+					Tags:       []string{randString(t)},
+					Enrichment: json.RawMessage("{}"),
+				},
+			}, nil
+		},
+	}
+
 	// these tests maintain a one:one relationship between
 	// update operations and a linked vulnerability for simplicty.
 	// in other words, each update operation inserts one vuln and
@@ -116,14 +160,22 @@ func TestGC(t *testing.T) {
 				locks,
 				http.DefaultClient, // Used on purpose -- shouldn't actually get called by anything.
 				updates.WithEnabled([]string{}),
+				updates.WithFactories(map[string]driver.UpdaterSetFactory{
+					"MockEnrichmentUpdater": func() driver.UpdaterSetFactory {
+						set := driver.NewUpdaterSet()
+						_ = set.Add(mockEnrich)
+						return driver.StaticSet(set)
+					}(),
+				}),
 				updates.WithOutOfTree([]driver.Updater{mock}),
 			)
 			if err != nil {
 				t.Fatalf("failed creating update manager: %v", err)
 			}
 
+			t.Logf("update Opts: %d", tt.updateOps)
 			// run updater n times to create n update operations
-			for i := 0; i < tt.updateOps; i++ {
+			for range tt.updateOps {
 				err := mgr.Run(ctx)
 				if err != nil {
 					t.Fatalf("manager failed to run: %v", err)
@@ -138,9 +190,16 @@ func TestGC(t *testing.T) {
 			if len(ops["MockUpdater"]) != tt.updateOps {
 				t.Fatalf("%s got: %v want: %v", tt.name, len(ops["MockUpdater"]), tt.updateOps)
 			}
+			ops, err = store.GetUpdateOperations(ctx, driver.EnrichmentKind)
+			if err != nil {
+				t.Fatalf("failed obtaining enrichment update ops: %v", err)
+			}
+			if len(ops["MockEnrichmentUpdater"]) != tt.updateOps {
+				t.Fatalf("%s got: %v want: %v", tt.name, len(ops["MockEnrichmentUpdater"]), tt.updateOps)
+			}
 
 			// run gc
-			expectedNotDone := Max(tt.updateOps-tt.keep-GCThrottle, 0)
+			expectedNotDone := max(2*(tt.updateOps-tt.keep)-GCThrottle, 0)
 			notDone, err := store.GC(ctx, tt.keep)
 			switch {
 			case err != nil:
@@ -153,14 +212,20 @@ func TestGC(t *testing.T) {
 			if tt.updateOps < tt.keep {
 				wantKeep = tt.updateOps
 			}
-			ops, err = store.GetUpdateOperations(ctx, driver.VulnerabilityKind)
+			expectedRemaining := 2*wantKeep + expectedNotDone
+
+			updaterOps, err := store.GetUpdateOperations(ctx, driver.VulnerabilityKind)
 			if err != nil {
 				t.Fatalf("failed obtaining update ops: %v", err)
 			}
-			t.Logf("ops %v", ops)
-			expectedRemaining := wantKeep + expectedNotDone
-			if len(ops["MockUpdater"]) != expectedRemaining {
-				t.Fatalf("%s got: %v want: %v", tt.name, len(ops["MockUpdater"]), expectedRemaining)
+			t.Logf("ops %v", updaterOps)
+			enricherOps, err := store.GetUpdateOperations(ctx, driver.EnrichmentKind)
+			if err != nil {
+				t.Fatalf("failed obtaining enrichment update ops: %v", err)
+			}
+			t.Logf("ops %v", enricherOps)
+			if len(updaterOps["MockUpdater"])+len(enricherOps["MockEnrichmentUpdater"]) != expectedRemaining {
+				t.Fatalf("%s got: %v want: %v", tt.name, len(updaterOps["MockUpdater"])+len(enricherOps["MockEnrichmentUpdater"]), expectedRemaining)
 			}
 		})
 	}
@@ -174,10 +239,3 @@ func randString(t *testing.T) string {
 	}
 	return hex.EncodeToString(buf)
 }
-
-func Max(x, y int) int {
-	if x < y {
-		return y
-	}
-	return x
-}

datastore/postgres/migrations/matcher/15-enrichment-indexes.sql

@@ -0,0 +1,3 @@
+CREATE INDEX IF NOT EXISTS uo_enrich_enrich_idx ON uo_enrich (enrich);
+CREATE INDEX IF NOT EXISTS uo_enrich_uo_idx ON uo_enrich (uo);
+CREATE INDEX IF NOT EXISTS enrichment_updater_idx ON enrichment (updater);

datastore/postgres/migrations/migrations.go

@@ -120,4 +120,8 @@ var MatcherMigrations = []migrate.Migration{
 		ID: 14,
 		Up: runFile("matcher/14-delete-rhcc-vulns.sql"),
 	},
+	{
+		ID: 15,
+		Up: runFile("matcher/15-enrichment-indexes.sql"),
+	},
 }