adjustments to support 60k MDM hosts after load testing (#21247)

for https://github.com/fleetdm/fleet/issues/20007

cmd/osquery-perf/agent.go

@@ -2509,7 +2509,7 @@ func main() {
 		// osquery-perf will send log requests with results only if there are scheduled queries configured AND it's their time to run.
 		logInterval         = flag.Duration("logger_tls_period", 10*time.Second, "Interval for scheduled queries log requests")
 		queryInterval       = flag.Duration("query_interval", 10*time.Second, "Interval for distributed query requests")
-		mdmCheckInInterval  = flag.Duration("mdm_check_in_interval", 10*time.Second, "Interval for performing MDM check-ins (applies to both macOS and Windows)")
+		mdmCheckInInterval  = flag.Duration("mdm_check_in_interval", 1*time.Minute, "Interval for performing MDM check-ins (applies to both macOS and Windows)")
 		onlyAlreadyEnrolled = flag.Bool("only_already_enrolled", false, "Only start agents that are already enrolled")
 		nodeKeyFile         = flag.String("node_key_file", "", "File with node keys to use")
 

infrastructure/loadtesting/terraform/readme.md

@@ -97,8 +97,8 @@ There are a few main places of interest to monitor the load and resource usage:
 
 You can deploy new code changes to an environment the following way:
 
-1. Push the code changes to the `BRANCH_NAME` and wait for the [Docker publish](https://github.com/fleetdm/fleet/actions/workflows/goreleaser-snapshot-fleet.yaml) action to complete.
-2. Find the docker image ID corresponding to your branch:
+1. Push the code changes to the `BRANCH_NAME`, trigger a manual run of the [Docker publish](https://github.com/fleetdm/fleet/actions/workflows/goreleaser-snapshot-fleet.yaml) workflow (make sure to select the branch) and wait for it to complete.
+2. Find the docker image IDs corresponding to your branch:
 ```sh
 docker images | grep 'BRANCH_NAME' | awk '{print $3}'
 ```

server/datastore/mysql/apple_mdm.go

@@ -12,6 +12,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"os"
 	"strings"
 	"time"
@@ -1712,18 +1713,34 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 		( hmap.host_uuid IS NOT NULL AND ( hmap.operation_type = ? OR hmap.operation_type IS NULL ) )
 `, fmt.Sprintf(appleMDMProfilesDesiredStateQuery, "h.uuid IN (?)", "h.uuid IN (?)", "h.uuid IN (?)"))
 
-	// TODO: if a very large number (~65K) of host uuids was matched (via
-	// uuids, teams or profile IDs), could result in too many placeholders (not
-	// an immediate concern).
-	stmt, args, err := sqlx.In(toInstallStmt, uuids, uuids, uuids, fleet.MDMOperationTypeRemove)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "building profiles to install statement")
+	selectProfilesBatchSize := 10_000
+	if ds.testSelectMDMProfilesBatchSize > 0 {
+		selectProfilesBatchSize = ds.testSelectMDMProfilesBatchSize
 	}
+	selectProfilesTotalBatches := int(math.Ceil(float64(len(uuids)) / float64(selectProfilesBatchSize)))
 
 	var wantedProfiles []*fleet.MDMAppleProfilePayload
-	err = sqlx.SelectContext(ctx, tx, &wantedProfiles, stmt, args...)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "bulk set pending profile status execute")
+	for i := 0; i < selectProfilesTotalBatches; i++ {
+		start := i * selectProfilesBatchSize
+		end := start + selectProfilesBatchSize
+		if end > len(uuids) {
+			end = len(uuids)
+		}
+
+		batchUUIDs := uuids[start:end]
+
+		stmt, args, err := sqlx.In(toInstallStmt, batchUUIDs, batchUUIDs, batchUUIDs, fleet.MDMOperationTypeRemove)
+		if err != nil {
+			return ctxerr.Wrapf(ctx, err, "building statement to select profiles to install, batch %d of %d", i, selectProfilesTotalBatches)
+		}
+
+		var partialResult []*fleet.MDMAppleProfilePayload
+		err = sqlx.SelectContext(ctx, tx, &partialResult, stmt, args...)
+		if err != nil {
+			return ctxerr.Wrapf(ctx, err, "selecting profiles to install, batch %d of %d", i, selectProfilesTotalBatches)
+		}
+
+		wantedProfiles = append(wantedProfiles, partialResult...)
 	}
 
 	// Exclude macOS only profiles from iPhones/iPads.
@@ -1760,17 +1777,27 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 		)
 `, fmt.Sprintf(appleMDMProfilesDesiredStateQuery, "h.uuid IN (?)", "h.uuid IN (?)", "h.uuid IN (?)"))
 
-	// TODO: if a very large number (~65K) of host uuids was matched (via
-	// uuids, teams or profile IDs), could result in too many placeholders (not
-	// an immediate concern). Note that uuids are provided twice.
-	stmt, args, err = sqlx.In(toRemoveStmt, uuids, uuids, uuids, uuids, fleet.MDMOperationTypeRemove)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "building profiles to remove statement")
-	}
 	var currentProfiles []*fleet.MDMAppleProfilePayload
-	err = sqlx.SelectContext(ctx, tx, &currentProfiles, stmt, args...)
-	if err != nil {
-		return ctxerr.Wrap(ctx, err, "fetching profiles to remove")
+	for i := 0; i < selectProfilesTotalBatches; i++ {
+		start := i * selectProfilesBatchSize
+		end := start + selectProfilesBatchSize
+		if end > len(uuids) {
+			end = len(uuids)
+		}
+
+		batchUUIDs := uuids[start:end]
+
+		stmt, args, err := sqlx.In(toRemoveStmt, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, fleet.MDMOperationTypeRemove)
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "building profiles to remove statement")
+		}
+		var partialResult []*fleet.MDMAppleProfilePayload
+		err = sqlx.SelectContext(ctx, tx, &partialResult, stmt, args...)
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "fetching profiles to remove")
+		}
+
+		currentProfiles = append(currentProfiles, partialResult...)
 	}
 
 	if len(wantedProfiles) == 0 && len(currentProfiles) == 0 {
@@ -1780,6 +1807,9 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 	// delete all host profiles to start from a clean slate, new entries will be added next
 	// TODO(roberto): is this really necessary? this was pre-existing
 	// behavior but I think it can be refactored. For now leaving it as-is.
+	//
+	// TODO part II(roberto): we found this call to be a major bottleneck during load testing
+	// https://github.com/fleetdm/fleet/issues/21338
 	if err := ds.bulkDeleteMDMAppleHostsConfigProfilesDB(ctx, tx, wantedProfiles); err != nil {
 		return ctxerr.Wrap(ctx, err, "bulk delete all profiles")
 	}

server/datastore/mysql/apple_mdm_test.go

@@ -1330,16 +1330,20 @@ func teamConfigProfileForTest(t *testing.T, name, identifier, uuid string, teamI
 }
 
 func testMDMAppleProfileManagementBatch2(t *testing.T, ds *Datastore) {
+	ds.testSelectMDMProfilesBatchSize = 2
 	ds.testUpsertMDMDesiredProfilesBatchSize = 2
 	t.Cleanup(func() {
+		ds.testSelectMDMProfilesBatchSize = 0
 		ds.testUpsertMDMDesiredProfilesBatchSize = 0
 	})
 	testMDMAppleProfileManagement(t, ds)
 }
 
 func testMDMAppleProfileManagementBatch3(t *testing.T, ds *Datastore) {
+	ds.testSelectMDMProfilesBatchSize = 3
 	ds.testUpsertMDMDesiredProfilesBatchSize = 3
 	t.Cleanup(func() {
+		ds.testSelectMDMProfilesBatchSize = 0
 		ds.testUpsertMDMDesiredProfilesBatchSize = 0
 	})
 	testMDMAppleProfileManagement(t, ds)

server/datastore/mysql/migrations/tables/20240826111228_AdjustHostIndexes.go

@@ -0,0 +1,41 @@
+package tables
+
+import (
+	"database/sql"
+	"fmt"
+)
+
+func init() {
+	MigrationClient.AddMigration(Up_20240826111228, Down_20240826111228)
+}
+
+func Up_20240826111228(tx *sql.Tx) error {
+	_, err := tx.Exec(`
+	    ALTER TABLE hosts
+	    DROP INDEX host_ip_mac_search
+	`)
+	if err != nil {
+		return fmt.Errorf("dropping host_ip_mac_search index: %w", err)
+	}
+
+	_, err = tx.Exec(`
+	    ALTER TABLE hosts
+	    DROP INDEX hosts_search
+	`)
+	if err != nil {
+		return fmt.Errorf("dropping hosts_search index: %w", err)
+	}
+
+	_, err = tx.Exec(`
+	    ALTER TABLE hosts
+	    ADD INDEX idx_hosts_uuid (uuid);
+	`)
+	if err != nil {
+		return fmt.Errorf("adding hosts_uuid index: %w", err)
+	}
+	return nil
+}
+
+func Down_20240826111228(tx *sql.Tx) error {
+	return nil
+}

server/datastore/mysql/mysql.go

@@ -82,6 +82,8 @@ type Datastore struct {
 	testDeleteMDMProfilesBatchSize int
 	// for tests, set to override the default batch size.
 	testUpsertMDMDesiredProfilesBatchSize int
+	// for tests set to override the default batch size.
+	testSelectMDMProfilesBatchSize int
 
 	// set this in tests to simulate an error at various stages in the
 	// batchSetMDMAppleProfilesDB execution: if the string starts with "insert", it