Atomic vulnerability count calculations

changes/35043-missing-vuln-counts

@@ -0,0 +1,2 @@
+- fixed issue where vulnerabilities would occasionally show as missing
+- added vulnerability seeding and performance testing tools

server/datastore/mysql/vulnerabilities.go

@@ -540,48 +540,27 @@ func getVulnHostCountQuery(scope CountScope) string {
 }
 
 func (ds *Datastore) UpdateVulnerabilityHostCounts(ctx context.Context, maxRoutines int) error {
-	// set all counts to 0 to later identify rows to delete
-	_, err := ds.writer(ctx).ExecContext(ctx, "UPDATE vulnerability_host_counts SET host_count = 0")
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "initializing vulnerability host counts")
-	}
-
 	globalHostCounts, err := ds.batchFetchVulnerabilityCounts(ctx, GlobalCount, maxRoutines)
 	if err != nil {
 		return ctxerr.Wrap(ctx, err, "fetching global vulnerability host counts")
 	}
 
-	err = ds.batchInsertHostCounts(ctx, globalHostCounts)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "inserting global vulnerability host counts")
-	}
-
 	teamHostCounts, err := ds.batchFetchVulnerabilityCounts(ctx, TeamCount, maxRoutines)
 	if err != nil {
 		return ctxerr.Wrap(ctx, err, "fetching team vulnerability host counts")
 	}
 
-	err = ds.batchInsertHostCounts(ctx, teamHostCounts)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "inserting team vulnerability host counts")
-	}
-
 	noTeamHostCounts, err := ds.batchFetchVulnerabilityCounts(ctx, NoTeamCount, maxRoutines)
 	if err != nil {
 		return ctxerr.Wrap(ctx, err, "fetching no team vulnerability host counts")
 	}
 
-	err = ds.batchInsertHostCounts(ctx, noTeamHostCounts)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "inserting team vulnerability host counts")
-	}
-
-	err = ds.cleanupVulnerabilityHostCounts(ctx)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "cleaning up vulnerability host counts")
-	}
+	allCounts := make([]hostCount, 0, len(globalHostCounts)+len(teamHostCounts)+len(noTeamHostCounts))
+	allCounts = append(allCounts, globalHostCounts...)
+	allCounts = append(allCounts, teamHostCounts...)
+	allCounts = append(allCounts, noTeamHostCounts...)
 
-	return nil
+	return ds.atomicTableSwapVulnerabilityCounts(ctx, allCounts)
 }
 
 type hostCount struct {
@@ -591,46 +570,101 @@ type hostCount struct {
 	GlobalStats bool   `db:"global_stats"`
 }
 
-func (ds *Datastore) cleanupVulnerabilityHostCounts(ctx context.Context) error {
-	_, err := ds.writer(ctx).ExecContext(ctx, "DELETE FROM vulnerability_host_counts WHERE host_count = 0")
+const vulnerabilityHostCountsSwapTableSchema = `
+	CREATE TABLE IF NOT EXISTS vulnerability_host_counts_swap (
+		cve varchar(255) NOT NULL,
+		team_id int unsigned NOT NULL DEFAULT 0,
+		host_count int unsigned NOT NULL DEFAULT 0,
+		global_stats tinyint(1) NOT NULL DEFAULT 0,
+		created_at timestamp DEFAULT CURRENT_TIMESTAMP,
+		updated_at timestamp DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
+		UNIQUE KEY cve_team_id_global_stats (cve, team_id, global_stats)
+	) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+`
+
+// atomicTableSwapVulnerabilityCounts implements atomic table swap pattern
+// 1. Populate swap table with new data
+// 2. Atomically rename tables to swap them
+// 3. Clean up old table
+func (ds *Datastore) atomicTableSwapVulnerabilityCounts(ctx context.Context, counts []hostCount) error {
+	err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
+		// Create/recreate the swap table fresh
+		_, err := tx.ExecContext(ctx, "DROP TABLE IF EXISTS vulnerability_host_counts_swap")
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "dropping existing swap table")
+		}
+
+		_, err = tx.ExecContext(ctx, vulnerabilityHostCountsSwapTableSchema)
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "creating swap table")
+		}
+
+		if len(counts) > 0 {
+			err = ds.insertHostCountsIntoTable(ctx, tx, counts, "vulnerability_host_counts_swap")
+			if err != nil {
+				return ctxerr.Wrap(ctx, err, "populating swap table")
+			}
+		}
+
+		return nil
+	})
 	if err != nil {
-		return fmt.Errorf("deleting zero host count entries: %w", err)
+		return err
 	}
 
-	return nil
+	// Atomic table swap using RENAME TABLE
+	return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
+		_, err := tx.ExecContext(ctx, `
+			RENAME TABLE
+				vulnerability_host_counts TO vulnerability_host_counts_old,
+				vulnerability_host_counts_swap TO vulnerability_host_counts
+		`)
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "atomic table swap")
+		}
+
+		// Clean up old table (drop it)
+		_, err = tx.ExecContext(ctx, "DROP TABLE vulnerability_host_counts_old")
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "dropping old table")
+		}
+
+		// Recreate empty swap table for next run
+		_, err = tx.ExecContext(ctx, vulnerabilityHostCountsSwapTableSchema)
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "recreating swap table")
+		}
+
+		return nil
+	})
 }
 
-func (ds *Datastore) batchInsertHostCounts(ctx context.Context, counts []hostCount) error {
+// insertHostCountsIntoTable inserts counts into specified table
+func (ds *Datastore) insertHostCountsIntoTable(ctx context.Context, tx sqlx.ExtContext, counts []hostCount, tableName string) error {
 	if len(counts) == 0 {
 		return nil
 	}
 
-	insertStmt := "INSERT INTO vulnerability_host_counts (team_id, cve, host_count, global_stats) VALUES "
-	var insertArgs []interface{}
+	insertStmt := fmt.Sprintf("INSERT INTO %s (team_id, cve, host_count, global_stats) VALUES ", tableName)
 
+	// Use smaller chunks to avoid parameter limits
 	chunkSize := 100
 	for i := 0; i < len(counts); i += chunkSize {
-		end := i + chunkSize
-		if end > len(counts) {
-			end = len(counts)
-		}
+		end := min(i+chunkSize, len(counts))
+
+		valueStrings := make([]string, 0, end-i)
+		chunkArgs := make([]interface{}, 0, (end-i)*4)
 
-		valueStrings := make([]string, 0, chunkSize)
 		for _, count := range counts[i:end] {
 			valueStrings = append(valueStrings, "(?, ?, ?, ?)")
-			insertArgs = append(insertArgs, count.TeamID, count.CVE, count.HostCount, count.GlobalStats)
+			chunkArgs = append(chunkArgs, count.TeamID, count.CVE, count.HostCount, count.GlobalStats)
 		}
 
-		insertStmt += strings.Join(valueStrings, ", ")
-		insertStmt += " ON DUPLICATE KEY UPDATE host_count = VALUES(host_count);"
-
-		_, err := ds.writer(ctx).ExecContext(ctx, insertStmt, insertArgs...)
+		fullStmt := insertStmt + strings.Join(valueStrings, ", ")
+		_, err := tx.ExecContext(ctx, fullStmt, chunkArgs...)
 		if err != nil {
-			return fmt.Errorf("inserting host counts: %w", err)
+			return fmt.Errorf("inserting host counts chunk %d-%d into %s: %w", i, end-1, tableName, err)
 		}
-
-		insertStmt = "INSERT INTO vulnerability_host_counts (team_id, cve, host_count, global_stats) VALUES "
-		insertArgs = nil
 	}
 
 	return nil

tools/software/vulnerabilities/performance_test/README.md

@@ -0,0 +1,135 @@
+# Vulnerability Performance Testing Tools
+
+This directory contains tools for testing the performance of Fleet's vulnerability-related datastore methods.
+
+## Tools
+
+### Seeder (`seeder/volume_vuln_seeder.go`)
+
+Seeds the database with test data for performance testing.
+
+**Usage:**
+
+```bash
+go run seeder/volume_vuln_seeder.go [options]
+```
+
+**Options:**
+
+- `-hosts=N` - Number of hosts to create (default: 100)
+- `-teams=N` - Number of teams to create (default: 5)
+- `-cves=N` - Total number of unique CVEs in the system (default: 500)
+- `-software=N` - Total number of unique software packages (default: 500)
+- `-help` - Show help information
+- `-verbose` - Enable verbose timing output for each step
+
+**Example:**
+
+```bash
+go run seeder/volume_vuln_seeder.go -hosts=1000 -teams=10 -cves=2000 -software=4000
+```
+
+### Performance Tester (`tester/performance_tester.go`)
+
+Benchmarks any Fleet datastore method with statistical analysis.
+
+**Usage:**
+
+```bash
+go run tester/performance_tester.go [options]
+```
+
+**Options:**
+
+- `-funcs=NAME[,NAME2,...]` - Comma-separated list of test functions (default: "UpdateVulnerabilityHostCounts")
+- `-iterations=N` - Number of iterations per test (default: 5)
+- `-verbose` - Show timing for each iteration
+- `-details` - Show detailed statistics including percentiles
+- `-list` - List available test functions
+- `-help` - Show help information
+
+**Available Test Functions:**
+
+- `UpdateVulnerabilityHostCounts` - Test vulnerability host count updates
+
+### Adding New Test Functions
+
+To add support for additional datastore methods, edit the `testFunctions` map in `tester/performance_tester.go`:
+
+```go
+var testFunctions = map[string]TestFunction{
+    // Existing functions...
+
+    // Add new function
+    "CountHosts": func(ctx context.Context, ds *mysql.Datastore) error {
+        _, err := ds.CountHosts(ctx, fleet.TeamFilter{User: &fleet.User{}}, fleet.HostListOptions{})
+        return err
+    },
+
+    // Add function with parameters
+    "ListHosts:100": func(ctx context.Context, ds *mysql.Datastore) error {
+        _, err := ds.ListHosts(ctx, fleet.TeamFilter{User: &fleet.User{}}, fleet.HostListOptions{
+            ListOptions: fleet.ListOptions{Page: 0, PerPage: 100},
+        })
+        return err
+    },
+}
+```
+
+Each function should:
+
+1. Accept `context.Context` and `*mysql.Datastore` as parameters
+2. Return only an `error`
+3. Handle any return values from the datastore method (discard non-error returns)
+4. Use meaningful parameter values for realistic testing
+
+**Examples:**
+
+```bash
+# Test single function with details
+go run tester/performance_tester.go -funcs=UpdateVulnerabilityHostCounts -iterations=10 -details
+
+# Test different batch sizes
+go run tester/performance_tester.go -funcs=UpdateVulnerabilityHostCounts:5,UpdateVulnerabilityHostCounts:20 -iterations=5
+
+# Verbose output
+go run tester/performance_tester.go -funcs=UpdateVulnerabilityHostCounts -verbose
+```
+
+## Performance Analysis
+
+The tools provide comprehensive performance metrics:
+
+- **Total time** - Sum of all successful iterations
+- **Average time** - Mean execution time
+- **Min/Max time** - Fastest and slowest iterations
+- **Success rate** - Percentage of successful vs failed iterations
+- **Percentiles** - P50, P90, P99 response times (with `-details`)
+
+## Typical Workflow
+
+1. **Seed test data:**
+
+   ```bash
+   go run seeder/volume_vuln_seeder.go -hosts=1000 -teams=10 -cves=2000 -software=4000
+   ```
+
+2. **Test baseline performance:**
+
+   ```bash
+   go run tester/performance_tester.go -funcs=UpdateVulnerabilityHostCounts -iterations=10 -details
+   ```
+
+3. **Make code changes to optimize**
+
+4. **Test optimized performance:**
+
+   ```bash
+   go run tester/performance_tester.go -funcs=UpdateVulnerabilityHostCounts -iterations=10 -details
+   ```
+
+5. **Compare results**
+
+## Notes
+
+- The seeder is not idempotent - run `make db-reset` to reset the database before reseeding.