Optimize memory usage in CVE chart cron job (#50385)

Resolves #50266. At production numbers the table looks like this - 20,691 CVEs × 83,000 hosts, ~268M raw (cve, host) rows (software + OS joins combined): ``` ┌─────────────────────────┬─────────────────────────┬───────────────────────┐ │ Shape of host IDs │ Old (map[string][]uint) │ New (roaring bitmaps) │ ├─────────────────────────┼─────────────────────────┼───────────────────────┤ │ Dense (contiguous runs) │ 2,479 MB │ 4.5 MB │ ├─────────────────────────┼─────────────────────────┼───────────────────────┤ │ Sparse (random) │ 2,488 MB │ 282 MB │ └─────────────────────────┴─────────────────────────┴───────────────────────┘ ``` A few things worth noting about how these map to your real data: - The old cost is shape-independent: ~2.5 GB retained just for the result map (268M rows × 8 bytes plus append slack), and the peak during collection is higher still because append doubling leaves garbage behind. That's the number that was blowing up the cron. - The new sparse figure is an overstated worst case. Your 268M rows include duplicates — multiple vulnerable software rows per host for the same CVE (the multi-kernel case) and overlap between the software and OS joins. The old code retained every raw row; the bitmap dedupes on Add, so it's bounded by unique pairs, and real fleets with AUTO_INCREMENT host IDs sit much closer to the dense row than the sparse one. - The new representation also has a hard ceiling the old one doesn't: a roaring bitmap over 83k host IDs maxes out around 16 KB per CVE regardless of contents, so even a pathological dataset caps at ~330 MB for all 20,691 CVEs — versus the old form growing linearly with join rows, unbounded. TL;DR: At scale the change is roughly a 550× reduction in the realistic (dense) case, and at minimum ~9× in the theoretical worst case. - [X] Changes file added for user-visible changes in `changes/`, `orbit/changes/` or `ee/fleetd-chrome/changes`. ## Testing - [X] Added/updated automated tests - [X] QA'd all new/changed functionality manually For unreleased bug fixes in a release candidate, one of: - [X] Confirmed that the fix is not expected to adversely impact load test results - [X] Alerted the release DRI if additional load testing is needed  ## Summary by CodeRabbit - **Performance** - Reduced memory usage for CVE chart data collection. - Improved efficiency when processing large CVE and affected-host datasets. - **Bug Fixes** - Preserved correct CVE filtering, duplicate-host handling, disabled-fleet exclusions, and empty-result behavior. - Added coverage for CVEs sourced from both software and operating-system data.
2026-08-02 14:16:34 -03:00
parent 1ed4b5c89a
commit 001b57cb9d
7 changed files with 181 additions and 43 deletions
@@ -0,0 +1 @@
+* Optimized memory usage of CVE chart cron job.
@@ -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.
@@ -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
@@ -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()
 }
@@ -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)
@@ -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
@@ -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
				`@@ -0,0 +1 @@`
				`* Optimized memory usage of CVE chart cron job.`