diff --git a/changes/50266-mem-optimization-cve-chart b/changes/50266-mem-optimization-cve-chart new file mode 100644 index 0000000000..ebb9c6cf0a --- /dev/null +++ b/changes/50266-mem-optimization-cve-chart @@ -0,0 +1 @@ +* Optimized memory usage of CVE chart cron job. diff --git a/server/chart/api/chart.go b/server/chart/api/chart.go index d1c431e4a6..b39cb83b54 100644 --- a/server/chart/api/chart.go +++ b/server/chart/api/chart.go @@ -80,12 +80,14 @@ type DatasetStore interface { // Used by datasets like uptime. FindOnlineHostIDs(ctx context.Context, now time.Time, disabledFleetIDs []uint) ([]uint, error) - // AffectedHostIDsByCVE returns host IDs grouped by CVE, scoped to the given - // cves set. nil or empty cves returns an empty map — callers must pass the - // CVE set they want to collect for. Unresolved-only is implicit in the - // underlying joins: a host's software/OS row transitions when it upgrades - // past the vulnerable version, so the join naturally stops matching. - AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string][]uint, error) + // AffectedHostIDsByCVE returns a bitmap of affected host IDs per CVE, + // scoped to the given cves set. nil or empty cves returns an empty map — + // callers must pass the CVE set they want to collect for. Unresolved-only + // is implicit in the underlying joins: a host's software/OS row transitions + // when it upgrades past the vulnerable version, so the join naturally + // stops matching. Bitmaps are returned in op form, ready to pass to + // RecordBucketData. + AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string]*roaring.Bitmap, error) // CollectibleCVEs returns every CVE ID, at all severities, on the curated // set of tracked software unioned with all operating-system vulnerabilities. diff --git a/server/chart/datasets.go b/server/chart/datasets.go index 11aca2d15e..a06f5b55ef 100644 --- a/server/chart/datasets.go +++ b/server/chart/datasets.go @@ -43,20 +43,17 @@ func (c *CVEDataset) Collect(ctx context.Context, store api.DatasetStore, now ti // Collect CVEs at all severities on the curated set of tracked software and // OS vulnerabilities. Display-time narrowing (critical-only this round, // plus user filters) happens at read time via ResolveCVEChartEntities. - // TODO: implement bitmap compression so we can track more CVEs. tracked, err := store.CollectibleCVEs(ctx) if err != nil { return err } - hostIDsByCVE, err := store.AffectedHostIDsByCVE(ctx, disabledFleetIDs, tracked) + // The store sets bits while streaming the vulnerability joins, so peak + // memory here is one bitmap per CVE — never the raw (CVE, host) pairs. + bitmaps, err := store.AffectedHostIDsByCVE(ctx, disabledFleetIDs, tracked) if err != nil { return err } - bitmaps := make(map[string]*roaring.Bitmap, len(hostIDsByCVE)) - for cve, hostIDs := range hostIDsByCVE { - bitmaps[cve] = NewBitmap(hostIDs) - } bucketStart := now.UTC().Truncate(time.Hour) // Always call RecordBucketData, even when bitmaps is empty: snapshot // semantics use an empty input to close any open rows for entities no diff --git a/server/chart/internal/mysql/charts.go b/server/chart/internal/mysql/charts.go index 824008ea91..de1724d811 100644 --- a/server/chart/internal/mysql/charts.go +++ b/server/chart/internal/mysql/charts.go @@ -5,10 +5,12 @@ import ( "context" "fmt" "log/slog" + "math" "slices" "strings" "time" + "github.com/RoaringBitmap/roaring" "github.com/fleetdm/fleet/v4/server/chart/api" "github.com/fleetdm/fleet/v4/server/chart/internal/types" "github.com/fleetdm/fleet/v4/server/contexts/ctxdb" @@ -156,8 +158,9 @@ func (ds *Datastore) FindOnlineHostIDs(ctx context.Context, now time.Time, disab } // The matcher list exists as a performance optimization that bounds which CVEs -// the chart collects. -// TODO: implement bitmap compression so we can collect more CVE data. +// the chart collects. Collection RAM is no longer the constraint (bits are set +// into roaring bitmaps while streaming — see streamCVEHostPairs); the list now +// bounds the join size and the host_scd_data row count per bucket. // cveSoftwareMatcher filters `software` rows by a MySQL LIKE pattern and an // optional source allowlist. Empty Sources means any source. Category groups @@ -203,15 +206,17 @@ var trackedCVESoftwareMatchers = []cveSoftwareMatcher{ {api.CVECategoryOS, "kernel-%", []string{"rpm_packages"}}, } -// AffectedHostIDsByCVE returns host IDs grouped by CVE, scoped to the given -// cves set. It streams two joins (software-level and OS-level vulnerabilities) -// and merges the results into a single map. Duplicates across sources are -// harmless — the downstream HostIDsToBlob setBit is idempotent. +// AffectedHostIDsByCVE returns a bitmap of affected host IDs per CVE, scoped +// to the given cves set. It streams two joins (software-level and OS-level +// vulnerabilities) and merges the results into a single map, setting bits +// while scanning so the raw (cve, host_id) rows — millions on a large fleet — +// are never materialized. Duplicates across sources are harmless — Bitmap.Add +// is idempotent. // // nil or empty cves returns an empty map without running any query. -// TODO: support `nil` meaning "all CVEs" once bitmap compression is implemented. -func (ds *Datastore) AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string][]uint, error) { - result := make(map[string][]uint) +// TODO: support `nil` meaning "all CVEs". +func (ds *Datastore) AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string]*roaring.Bitmap, error) { + result := make(map[string]*roaring.Bitmap) if len(cves) == 0 { return result, nil } @@ -255,6 +260,12 @@ func (ds *Datastore) AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs return nil, ctxerr.Wrap(ctx, err, "stream OS CVE host pairs") } + // Compact container representations now that all bits are set, so the + // retained op form stays small until RecordBucketData serializes it. + for _, rb := range result { + rb.RunOptimize() + } + return result, nil } @@ -450,14 +461,20 @@ func streamCVEStrings(ctx context.Context, q sqlx.QueryerContext, query string, return rows.Err() } -// streamCVEHostPairs runs a query yielding (cve, host_id) pairs and appends -// host IDs into out under each CVE key. Streams rather than materializing the -// join result, since on a large fleet the (cve, host_id) row count can reach -// millions. +// streamCVEHostPairs runs a query yielding (cve, host_id) pairs and sets each +// host's bit in out's bitmap for that CVE, allocating the bitmap on first +// sight of the CVE. Setting bits while scanning keeps peak memory at one +// bitmap per CVE (KBs even at 50k hosts) instead of retaining every raw +// (cve, host_id) pair, whose row count can reach many millions on a large +// fleet. Duplicate pairs — several matching software rows on one host, or +// overlap between the software and OS queries — are no-op Adds. +// +// Host IDs of 0 or above MaxUint32 are skipped, mirroring chart.NewBitmap — +// Fleet host IDs are AUTO_INCREMENT starting at 1. // // args are expanded via sqlx.In for slice arguments (e.g. team IDs) and // rebinds to the driver dialect. -func (ds *Datastore) streamCVEHostPairs(ctx context.Context, query string, args []any, out map[string][]uint) error { +func (ds *Datastore) streamCVEHostPairs(ctx context.Context, query string, args []any, out map[string]*roaring.Bitmap) error { if len(args) > 0 { expanded, expandedArgs, err := sqlx.In(query, args...) if err != nil { @@ -481,7 +498,15 @@ func (ds *Datastore) streamCVEHostPairs(ctx context.Context, query string, args if err := rows.Scan(&cve, &hostID); err != nil { return err } - out[cve] = append(out[cve], hostID) + if hostID == 0 || hostID > math.MaxUint32 { + continue + } + rb, ok := out[cve] + if !ok { + rb = roaring.New() + out[cve] = rb + } + rb.Add(uint32(hostID)) } return rows.Err() } diff --git a/server/chart/internal/mysql/charts_test.go b/server/chart/internal/mysql/charts_test.go index 36238319a2..1d12eafaae 100644 --- a/server/chart/internal/mysql/charts_test.go +++ b/server/chart/internal/mysql/charts_test.go @@ -1,6 +1,7 @@ package mysql import ( + "crypto/sha256" "testing" "time" @@ -361,6 +362,118 @@ func testFindOnlineMobileDisabledEnrollment(t *testing.T, tdb *testutils.TestDB, assert.ElementsMatch(t, []uint{ids[0]}, got) } +// seedHostVulnSoftware inserts one software row, links it to the given CVEs, +// and installs it on the host. Unlike seedSoftware (cve_filter_test.go), which +// only attributes a CVE to software, this attributes CVEs to a specific host +// via host_software. +func seedHostVulnSoftware(t *testing.T, tdb *testutils.TestDB, hostID uint, name, source string, cves ...string) { + t.Helper() + ctx := t.Context() + + // checksum is binary(16) UNIQUE NOT NULL; derive it from the row's own + // identifying inputs so each seeded row is unique. + sum := sha256.Sum256([]byte(name + "\x00" + source + "\x00" + itoa(hostID))) + res, err := tdb.DB.ExecContext(ctx, + `INSERT INTO software (name, version, source, checksum) VALUES (?, '1.0', ?, ?)`, + name, source, sum[:16]) + require.NoError(t, err) + swID, err := res.LastInsertId() + require.NoError(t, err) + + for _, cve := range cves { + _, err = tdb.DB.ExecContext(ctx, + `INSERT INTO software_cve (software_id, cve) VALUES (?, ?)`, swID, cve) + require.NoError(t, err) + } + _, err = tdb.DB.ExecContext(ctx, + `INSERT INTO host_software (host_id, software_id) VALUES (?, ?)`, hostID, swID) + require.NoError(t, err) +} + +// seedHostOS creates an operating_systems row with the given id (satisfying +// host_operating_system's FK) and links the host to it. Attribute CVEs to the +// OS afterwards via seedOSVuln. +func seedHostOS(t *testing.T, tdb *testutils.TestDB, hostID, osID uint) { + t.Helper() + ctx := t.Context() + _, err := tdb.DB.ExecContext(ctx, + `INSERT INTO operating_systems (id, name, version, arch, kernel_version, platform) + VALUES (?, ?, '1.0', 'x86_64', '1.0', 'linux')`, + osID, "os-"+itoa(osID)) + require.NoError(t, err) + _, err = tdb.DB.ExecContext(ctx, + `INSERT INTO host_operating_system (host_id, os_id) VALUES (?, ?)`, hostID, osID) + require.NoError(t, err) +} + +// u32 narrows a seeded host id for bitmap-content comparison. +func u32(id uint) uint32 { + return uint32(id) //nolint:gosec // G115: AUTO_INCREMENT primary key fits in uint32 +} + +// TestAffectedHostIDsByCVE covers the CVE collector's host-set query: rows are +// grouped per CVE as bitmaps, duplicate (cve, host) rows (several vulnerable +// software rows on one host) collapse to a single bit, software- and OS-level +// sources merge, the cves argument scopes the result, and hosts in disabled +// fleets are excluded. +func TestAffectedHostIDsByCVE(t *testing.T) { + tdb := testutils.SetupTestDB(t, "chart_mysql") + defer tdb.TruncateTables(t) + ds := NewDatastore(tdb.Conns(), tdb.Logger) + ctx := t.Context() + + now := time.Now().UTC().Truncate(time.Second) + ids := seedHosts(t, tdb, []hostSeed{ + {teamID: 0, seenTime: now}, // 0: no team + {teamID: 1, seenTime: now}, // 1 + {teamID: 2, seenTime: now}, // 2 + }) + + // Host 0 carries two distinct software rows both vulnerable to CVE-A (the + // multiple-installed-kernels shape) — the duplicate (CVE-A, host 0) rows + // must collapse to a single bit. + seedHostVulnSoftware(t, tdb, ids[0], "linux-image-6.1", "deb_packages", "CVE-A") + seedHostVulnSoftware(t, tdb, ids[0], "linux-image-6.5", "deb_packages", "CVE-A") + seedHostVulnSoftware(t, tdb, ids[1], "Google Chrome", "apps", "CVE-A", "CVE-B") + // Host 2 gets CVE-B via its OS (merging with host 1's software-side CVE-B) + // and CVE-C via software; CVE-C is never requested so it must not appear. + seedHostOS(t, tdb, ids[2], 1) + seedOSVuln(t, tdb, 1, "CVE-B") + seedHostVulnSoftware(t, tdb, ids[2], "Firefox", "apps", "CVE-C") + + t.Run("GroupsDedupesAndMergesSources", func(t *testing.T) { + got, err := ds.AffectedHostIDsByCVE(ctx, nil, []string{"CVE-A", "CVE-B"}) + require.NoError(t, err) + require.Len(t, got, 2) + assert.Equal(t, []uint32{u32(ids[0]), u32(ids[1])}, got["CVE-A"].ToArray(), + "host 0's two vulnerable software rows must produce one bit") + assert.Equal(t, []uint32{u32(ids[1]), u32(ids[2])}, got["CVE-B"].ToArray(), + "software-side and OS-side hosts must merge under one CVE") + }) + + t.Run("DisabledFleetsExcluded", func(t *testing.T) { + got, err := ds.AffectedHostIDsByCVE(ctx, []uint{1}, []string{"CVE-A", "CVE-B"}) + require.NoError(t, err) + require.Len(t, got, 2) + assert.Equal(t, []uint32{u32(ids[0])}, got["CVE-A"].ToArray(), + "disabled-fleet host dropped; NULL-team host retained") + assert.Equal(t, []uint32{u32(ids[2])}, got["CVE-B"].ToArray()) + }) + + t.Run("FullyExcludedCVELeavesNoKey", func(t *testing.T) { + got, err := ds.AffectedHostIDsByCVE(ctx, []uint{1, 2}, []string{"CVE-B"}) + require.NoError(t, err) + assert.Empty(t, got, "a CVE whose only affected hosts are excluded must not appear") + }) + + t.Run("EmptyCVEsShortCircuits", func(t *testing.T) { + got, err := ds.AffectedHostIDsByCVE(ctx, nil, nil) + require.NoError(t, err) + require.NotNil(t, got) + assert.Empty(t, got) + }) +} + func testFindOnlineMobileDisabledFleet(t *testing.T, tdb *testutils.TestDB, ds *Datastore) { ctx := t.Context() now := time.Now().UTC().Truncate(time.Second) diff --git a/server/chart/internal/service/service_test.go b/server/chart/internal/service/service_test.go index aeccecf218..43fdfecdd6 100644 --- a/server/chart/internal/service/service_test.go +++ b/server/chart/internal/service/service_test.go @@ -60,7 +60,7 @@ type mockDatastore struct { getSCDDataFunc func(ctx context.Context, dataset string, startDate, endDate time.Time, bucketSize time.Duration, strategy api.SampleStrategy, filterMask *roaring.Bitmap, entityIDs []string) ([]api.DataPoint, error) getHostIDsForFilterFunc func(ctx context.Context, hostFilter *types.HostFilter) ([]uint, error) findOnlineHostIDsFn func(ctx context.Context, now time.Time, disabledFleetIDs []uint) ([]uint, error) - affectedHostIDsByCVEFn func(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string][]uint, error) + affectedHostIDsByCVEFn func(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string]*roaring.Bitmap, error) collectibleCVEsFn func(ctx context.Context) ([]string, error) resolveCVEEntitiesFn func(ctx context.Context, filter types.CVEChartFilter) ([]string, error) recordBucketDataFn func(ctx context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error @@ -77,7 +77,7 @@ func (m *mockDatastore) FindOnlineHostIDs(ctx context.Context, now time.Time, di return nil, nil } -func (m *mockDatastore) AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string][]uint, error) { +func (m *mockDatastore) AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string]*roaring.Bitmap, error) { if m.affectedHostIDsByCVEFn != nil { return m.affectedHostIDsByCVEFn(ctx, disabledFleetIDs, cves) } @@ -647,11 +647,11 @@ func TestCollectDatasetsCVE(t *testing.T) { return wantTracked, nil } var gotCVEs []string - ds.affectedHostIDsByCVEFn = func(_ context.Context, _ []uint, cves []string) (map[string][]uint, error) { + ds.affectedHostIDsByCVEFn = func(_ context.Context, _ []uint, cves []string) (map[string]*roaring.Bitmap, error) { gotCVEs = cves - return map[string][]uint{ - "CVE-2024-0001": {1, 2, 3}, - "CVE-2024-0002": {2, 4}, + return map[string]*roaring.Bitmap{ + "CVE-2024-0001": roaring.BitmapOf(1, 2, 3), + "CVE-2024-0002": roaring.BitmapOf(2, 4), }, nil } ds.recordBucketDataFn = func(_ context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error { @@ -684,9 +684,9 @@ func TestCollectDatasetsCVEEmptyTracked(t *testing.T) { ds.collectibleCVEsFn = func(_ context.Context) ([]string, error) { return []string{}, nil } - ds.affectedHostIDsByCVEFn = func(_ context.Context, _ []uint, cves []string) (map[string][]uint, error) { + ds.affectedHostIDsByCVEFn = func(_ context.Context, _ []uint, cves []string) (map[string]*roaring.Bitmap, error) { assert.Empty(t, cves, "empty tracked set must propagate as empty cves filter") - return map[string][]uint{}, nil + return map[string]*roaring.Bitmap{}, nil } var gotBitmaps map[string]*roaring.Bitmap ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error { @@ -752,9 +752,9 @@ func TestCollectDatasetsForwardsScope(t *testing.T) { return []string{"CVE-1"}, nil } var gotDisabled []uint - ds.affectedHostIDsByCVEFn = func(_ context.Context, disabled []uint, _ []string) (map[string][]uint, error) { + ds.affectedHostIDsByCVEFn = func(_ context.Context, disabled []uint, _ []string) (map[string]*roaring.Bitmap, error) { gotDisabled = disabled - return map[string][]uint{"CVE-1": {1}}, nil + return map[string]*roaring.Bitmap{"CVE-1": roaring.BitmapOf(1)}, nil } ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ map[string]*roaring.Bitmap) error { return nil diff --git a/server/chart/internal/types/chart.go b/server/chart/internal/types/chart.go index aaba158768..5c017cf9e3 100644 --- a/server/chart/internal/types/chart.go +++ b/server/chart/internal/types/chart.go @@ -59,12 +59,12 @@ type Datastore interface { // like uptime. FindOnlineHostIDs(ctx context.Context, now time.Time, disabledFleetIDs []uint) ([]uint, error) - // AffectedHostIDsByCVE returns host IDs grouped by CVE, scoped to the given - // cves set. nil or empty cves returns an empty map. Unresolved-only is - // implicit in the underlying joins: a host's software/OS row transitions - // when it upgrades past the vulnerable version, so the join naturally - // stops matching. - AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string][]uint, error) + // AffectedHostIDsByCVE returns a bitmap of affected host IDs per CVE, + // scoped to the given cves set. nil or empty cves returns an empty map. + // Unresolved-only is implicit in the underlying joins: a host's software/OS + // row transitions when it upgrades past the vulnerable version, so the join + // naturally stops matching. + AffectedHostIDsByCVE(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string]*roaring.Bitmap, error) // CollectibleCVEs returns every CVE ID, at all severities, on the curated // set of tracked software (trackedCVESoftwareMatchers) unioned with all