Optimize ReparseWindowsHostCertificates migration to fix DB overload during upgrade

changes/50228-fix-reparse-certs-migration

@@ -0,0 +1 @@
+* Improved performance of the Windows host certificates re-parse migration to avoid sustained high DB load during upgrade.

server/datastore/mysql/migrations/tables/20260723181402_ReparseWindowsHostCertificates.go

@@ -32,41 +32,50 @@
 	return total, err
 }
 
-// softDeleteWindowsHostCertsForReparse walks the osquery-origin Windows host certificates in id-keyed batches and
-// soft-deletes each one, calling increment per row so progress is reported.
+// softDeleteWindowsHostCertsForReparse collects Windows host IDs once upfront,
+// then soft-deletes their certificate rows in batches using only the
+// host_certificates primary key. This avoids re-joining the hosts table on
+// every batch, which was causing ~14 minutes of sustained full DB load at
+// 100K hosts (#50228).
 func softDeleteWindowsHostCertsForReparse(tx *sql.Tx, increment incrementCountFn) error {
 	txx := sqlx.Tx{Tx: tx, Mapper: reflectx.NewMapperFunc("db", sqlx.NameMapper)}
 
-	const batchSize = 1000
-	var lastID uint64
-	for {
-		var ids []uint64
-		if err := txx.Select(&ids, `
-			SELECT hc.id
-			FROM host_certificates hc
-			JOIN hosts h ON h.id = hc.host_id
-			WHERE h.platform = 'windows' AND hc.origin = 'osquery' AND hc.deleted_at IS NULL AND hc.id > ?
-			ORDER BY hc.id
-			LIMIT ?`, lastID, batchSize); err != nil {
-			return fmt.Errorf("selecting windows host certs batch after id %d: %w", lastID, err)
-		}
-		if len(ids) == 0 {
-			return nil
+	// Step 1: collect all Windows host IDs in one query (bounded by host count, not cert count).
+	var windowsHostIDs []uint64
+	if err := txx.Select(&windowsHostIDs, `SELECT id FROM hosts WHERE platform = 'windows'`); err != nil {
+		return fmt.Errorf("selecting windows host ids: %w", err)
+	}
+	if len(windowsHostIDs) == 0 {
+		return nil
+	}
+
+	// Step 2: soft-delete certs in batches of host IDs. Each batch UPDATE
+	// hits the host_certificates index on (host_id) without joining hosts.
+	const hostBatchSize = 500
+	for i := 0; i < len(windowsHostIDs); i += hostBatchSize {
+		end := i + hostBatchSize
+		if end > len(windowsHostIDs) {
+			end = len(windowsHostIDs)
 		}
+		batch := windowsHostIDs[i:end]
 
-		query, args, err := sqlx.In(`UPDATE host_certificates SET deleted_at = NOW(6) WHERE id IN (?)`, ids)
+		query, args, err := sqlx.In(`
+			UPDATE host_certificates
+			SET deleted_at = NOW(6)
+			WHERE host_id IN (?) AND origin = 'osquery' AND deleted_at IS NULL`, batch)
 		if err != nil {
-			return fmt.Errorf("building soft-delete query: %w", err)
+			return fmt.Errorf("building soft-delete query for host batch: %w", err)
 		}
-		if _, err := txx.Exec(query, args...); err != nil {
-			return fmt.Errorf("soft-deleting windows host certs batch after id %d: %w", lastID, err)
+		result, err := txx.Exec(query, args...)
+		if err != nil {
+			return fmt.Errorf("soft-deleting certs for host batch starting at index %d: %w", i, err)
 		}
-
-		for range ids {
+		affected, _ := result.RowsAffected()
+		for range affected {
 			increment()
 		}
-		lastID = ids[len(ids)-1]
 	}
+	return nil
 }
 
 func Down_20260723181402(tx *sql.Tx) error {

server/datastore/mysql/migrations/tables/20260723181402_ReparseWindowsHostCertificates_bench_test.go

@@ -0,0 +1,106 @@
+package tables
+
+import (
+	"crypto/sha256"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestUp_20260723181402_Performance inserts a realistic number of hosts and
+// certificates, then runs the migration and reports how long it takes. This
+// validates the fix for #50228 where the original migration took ~14 minutes
+// at 100K hosts due to repeated JOIN on the hosts table per batch.
+//
+// The test uses 1,000 Windows hosts x 20 certs each = 20,000 cert rows.
+// This is enough to exercise the batching logic (hostBatchSize=500) while
+// staying fast enough for CI (~seconds, not minutes).
+func TestUp_20260723181402_Performance(t *testing.T) {
+	db := applyUpToPrev(t)
+
+	const (
+		numWindowsHosts = 1000
+		numMacHosts     = 200
+		certsPerHost    = 20
+	)
+
+	// Insert Windows hosts.
+	for i := 0; i < numWindowsHosts; i++ {
+		uid := fmt.Sprintf("win-%d", i)
+		execNoErr(t, db, `INSERT INTO hosts (osquery_host_id, node_key, hostname, uuid, platform) VALUES (?, ?, ?, ?, 'windows')`,
+			uid, uid, uid, uid)
+	}
+
+	// Insert macOS hosts (should be untouched by migration).
+	for i := 0; i < numMacHosts; i++ {
+		uid := fmt.Sprintf("mac-%d", i)
+		execNoErr(t, db, `INSERT INTO hosts (osquery_host_id, node_key, hostname, uuid, platform) VALUES (?, ?, ?, ?, 'darwin')`,
+			uid, uid, uid, uid)
+	}
+
+	// Collect host IDs.
+	var windowsIDs []uint
+	require.NoError(t, db.Select(&windowsIDs, `SELECT id FROM hosts WHERE platform = 'windows'`))
+	var macIDs []uint
+	require.NoError(t, db.Select(&macIDs, `SELECT id FROM hosts WHERE platform = 'darwin'`))
+
+	// Insert certs for all hosts.
+	insertCerts := func(hostIDs []uint, prefix string) {
+		for _, hid := range hostIDs {
+			for j := 0; j < certsPerHost; j++ {
+				fullHash := sha256.Sum256([]byte(fmt.Sprintf("%s-%d-%d", prefix, hid, j)))
+				sha1Bytes := fullHash[:20] // sha1_sum column is 20 bytes
+				serial := fmt.Sprintf("%s-%d-%d", prefix, hid, j)
+				execNoErr(t, db, `
+					INSERT INTO host_certificates (
+						host_id, not_valid_after, not_valid_before, certificate_authority,
+						common_name, key_algorithm, key_strength, key_usage,
+						serial, signing_algorithm,
+						subject_country, subject_org, subject_org_unit, subject_common_name,
+						issuer_country, issuer_org, issuer_org_unit, issuer_common_name,
+						sha1_sum, origin, deleted_at
+					) VALUES (?, '2027-01-01', '2026-01-01', 0, 'cn', 'rsa', 2048, 'digitalSignature',
+						?, 'sha256WithRSAEncryption', '', '', '', '', '', '', '', '',
+						?, 'osquery', NULL)`,
+					hid, serial, sha1Bytes)
+			}
+		}
+	}
+	insertCerts(windowsIDs, "win")
+	insertCerts(macIDs, "mac")
+
+	// Verify row counts before migration.
+	var totalCerts int
+	require.NoError(t, db.Get(&totalCerts, `SELECT COUNT(*) FROM host_certificates WHERE deleted_at IS NULL`))
+	t.Logf("Total live certs before migration: %d", totalCerts)
+	require.Equal(t, (numWindowsHosts+numMacHosts)*certsPerHost, totalCerts)
+
+	// Run the migration and time it.
+	start := time.Now()
+	applyNext(t, db)
+	elapsed := time.Since(start)
+	t.Logf("Migration completed in %s for %d Windows hosts x %d certs = %d cert rows",
+		elapsed, numWindowsHosts, certsPerHost, numWindowsHosts*certsPerHost)
+
+	// Verify: all Windows osquery certs are soft-deleted.
+	var windowsLive int
+	require.NoError(t, db.Get(&windowsLive, `
+		SELECT COUNT(*) FROM host_certificates hc
+		JOIN hosts h ON h.id = hc.host_id
+		WHERE h.platform = 'windows' AND hc.origin = 'osquery' AND hc.deleted_at IS NULL`))
+	require.Equal(t, 0, windowsLive, "all Windows osquery certs should be soft-deleted")
+
+	// Verify: all macOS certs are untouched.
+	var macLive int
+	require.NoError(t, db.Get(&macLive, `
+		SELECT COUNT(*) FROM host_certificates hc
+		JOIN hosts h ON h.id = hc.host_id
+		WHERE h.platform = 'darwin' AND hc.origin = 'osquery' AND hc.deleted_at IS NULL`))
+	require.Equal(t, numMacHosts*certsPerHost, macLive, "macOS certs should be untouched")
+
+	// At 1,000 hosts x 20 certs, the migration should complete in under 30 seconds.
+	// The original implementation took proportionally ~14 minutes at 100K hosts.
+	require.Less(t, elapsed, 30*time.Second, "migration took too long, possible performance regression")
+}