Cherry Pick Scope pending host profile rebuilds (#23772) (#23868)

#23772

changes/21338-scope-profile-pending-rebuild

@@ -0,0 +1 @@
+- Speed up adding and removing profiles to large teams by an order of magnitude

cmd/fleet/serve.go

@@ -189,6 +189,7 @@ the way that the Fleet server works.
 			if config.MysqlReadReplica.Address != "" {
 				opts = append(opts, mysql.Replica(&config.MysqlReadReplica))
 			}
+			// NOTE this will disable OTEL/APM interceptor
 			if dev && os.Getenv("FLEET_DEV_ENABLE_SQL_INTERCEPTOR") != "" {
 				opts = append(opts, mysql.WithInterceptor(&devSQLInterceptor{
 					logger: kitlog.With(logger, "component", "sql-interceptor"),

server/datastore/mysql/apple_mdm.go

@@ -1940,12 +1940,16 @@ func (ds *Datastore) bulkDeleteMDMAppleHostsConfigProfilesDB(ctx context.Context
 	return nil
 }
 
+// NOTE If onlyProfileUUIDs is provided (not nil), only profiles with
+// those UUIDs will be update instead of rebuilding all pending
+// profiles for hosts
 func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 	ctx context.Context,
 	tx sqlx.ExtContext,
-	uuids []string,
+	hostUUIDs []string,
+	onlyProfileUUIDs []string,
 ) (updatedDB bool, err error) {
-	if len(uuids) == 0 {
+	if len(hostUUIDs) == 0 {
 		return false, nil
 	}
 
@@ -1972,6 +1976,11 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 	// considers "remove" operations that have NULL status, which it would
 	// update to make its status to NULL).
 
+	profileHostIn := "h.uuid IN (?)"
+	if len(onlyProfileUUIDs) > 0 {
+		profileHostIn = "mae.profile_uuid IN (?) AND " + profileHostIn
+	}
+
 	toInstallStmt := fmt.Sprintf(`
 	SELECT
 		ds.profile_uuid as profile_uuid,
@@ -1990,7 +1999,7 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 		( hmap.profile_uuid IS NULL AND hmap.host_uuid IS NULL ) OR
 		-- profiles in A and B but with operation type "remove"
 		( 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 (?)", "h.uuid IN (?)"))
+`, fmt.Sprintf(appleMDMProfilesDesiredStateQuery, profileHostIn, profileHostIn, profileHostIn, profileHostIn))
 
 	// batches of 10K hosts because h.uuid appears three times in the
 	// query, and the max number of prepared statements is 65K, this was
@@ -2000,19 +2009,35 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 	if ds.testSelectMDMProfilesBatchSize > 0 {
 		selectProfilesBatchSize = ds.testSelectMDMProfilesBatchSize
 	}
-	selectProfilesTotalBatches := int(math.Ceil(float64(len(uuids)) / float64(selectProfilesBatchSize)))
+	selectProfilesTotalBatches := int(math.Ceil(float64(len(hostUUIDs)) / float64(selectProfilesBatchSize)))
 
 	var wantedProfiles []*fleet.MDMAppleProfilePayload
 	for i := 0; i < selectProfilesTotalBatches; i++ {
 		start := i * selectProfilesBatchSize
 		end := start + selectProfilesBatchSize
-		if end > len(uuids) {
-			end = len(uuids)
+		if end > len(hostUUIDs) {
+			end = len(hostUUIDs)
 		}
 
-		batchUUIDs := uuids[start:end]
+		batchUUIDs := hostUUIDs[start:end]
+
+		var stmt string
+		var args []any
+		var err error
+
+		if len(onlyProfileUUIDs) > 0 {
+			stmt, args, err = sqlx.In(
+				toInstallStmt,
+				onlyProfileUUIDs, batchUUIDs,
+				onlyProfileUUIDs, batchUUIDs,
+				onlyProfileUUIDs, batchUUIDs,
+				onlyProfileUUIDs, batchUUIDs,
+				fleet.MDMOperationTypeRemove,
+			)
+		} else {
+			stmt, args, err = sqlx.In(toInstallStmt, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, fleet.MDMOperationTypeRemove)
+		}
 
-		stmt, args, err := sqlx.In(toInstallStmt, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, fleet.MDMOperationTypeRemove)
 		if err != nil {
 			return false, ctxerr.Wrapf(ctx, err, "building statement to select profiles to install, batch %d of %d", i,
 				selectProfilesTotalBatches)
@@ -2030,6 +2055,11 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 	// Exclude macOS only profiles from iPhones/iPads.
 	wantedProfiles = fleet.FilterMacOSOnlyProfilesFromIOSIPadOS(wantedProfiles)
 
+	narrowByProfiles := ""
+	if len(onlyProfileUUIDs) > 0 {
+		narrowByProfiles = "AND hmap.profile_uuid IN (?)"
+	}
+
 	toRemoveStmt := fmt.Sprintf(`
 	SELECT
 		hmap.profile_uuid as profile_uuid,
@@ -2045,7 +2075,7 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 		RIGHT JOIN host_mdm_apple_profiles hmap
 			ON hmap.profile_uuid = ds.profile_uuid AND hmap.host_uuid = ds.host_uuid
 	WHERE
-		hmap.host_uuid IN (?) AND
+		hmap.host_uuid IN (?) %s AND
 		-- profiles that are in B but not in A
 		ds.profile_uuid IS NULL AND ds.host_uuid IS NULL AND
 		-- except "remove" operations in any state
@@ -2059,19 +2089,36 @@ func (ds *Datastore) bulkSetPendingMDMAppleHostProfilesDB(
 			mcpl.apple_profile_uuid = hmap.profile_uuid AND
 			mcpl.label_id IS NULL
 		)
-`, fmt.Sprintf(appleMDMProfilesDesiredStateQuery, "h.uuid IN (?)", "h.uuid IN (?)", "h.uuid IN (?)", "h.uuid IN (?)"))
+`, fmt.Sprintf(appleMDMProfilesDesiredStateQuery, profileHostIn, profileHostIn, profileHostIn, profileHostIn), narrowByProfiles)
 
 	var currentProfiles []*fleet.MDMAppleProfilePayload
 	for i := 0; i < selectProfilesTotalBatches; i++ {
 		start := i * selectProfilesBatchSize
 		end := start + selectProfilesBatchSize
-		if end > len(uuids) {
-			end = len(uuids)
+		if end > len(hostUUIDs) {
+			end = len(hostUUIDs)
 		}
 
-		batchUUIDs := uuids[start:end]
+		batchUUIDs := hostUUIDs[start:end]
+
+		var stmt string
+		var args []any
+		var err error
+
+		if len(onlyProfileUUIDs) > 0 {
+			stmt, args, err = sqlx.In(
+				toRemoveStmt,
+				onlyProfileUUIDs, batchUUIDs,
+				onlyProfileUUIDs, batchUUIDs,
+				onlyProfileUUIDs, batchUUIDs,
+				onlyProfileUUIDs, batchUUIDs,
+				batchUUIDs, onlyProfileUUIDs,
+				fleet.MDMOperationTypeRemove,
+			)
+		} else {
+			stmt, args, err = sqlx.In(toRemoveStmt, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, fleet.MDMOperationTypeRemove)
+		}
 
-		stmt, args, err := sqlx.In(toRemoveStmt, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, fleet.MDMOperationTypeRemove)
 		if err != nil {
 			return false, ctxerr.Wrap(ctx, err, "building profiles to remove statement")
 		}
@@ -2431,7 +2478,7 @@ func generateDesiredStateQuery(entityType string) string {
 		COUNT(*) as ${countEntityLabelsColumn},
 		COUNT(mel.label_id) as count_non_broken_labels,
 		COUNT(lm.label_id) as count_host_labels,
-		-- this helps avoid the case where the host is not a member of a label 
+		-- this helps avoid the case where the host is not a member of a label
 		-- just because it hasn't reported results for that label yet.
 		SUM(CASE WHEN lbl.created_at IS NOT NULL AND h.label_updated_at >= lbl.created_at THEN 1 ELSE 0 END) as count_host_updated_after_labels
 	FROM

server/datastore/mysql/mdm.go

@@ -461,28 +461,36 @@ func (ds *Datastore) bulkSetPendingMDMHostProfilesDB(
 		}
 
 	case len(macProfUUIDs) > 0:
-		// TODO: if a very large number (~65K) of profile UUIDs was provided, could
+		// TODO: if a very large number (~65K/2) of profile UUIDs was provided, could
 		// result in too many placeholders (not an immediate concern).
 		uuidStmt = `
 SELECT DISTINCT h.uuid, h.platform
 FROM hosts h
 JOIN mdm_apple_configuration_profiles macp
 	ON h.team_id = macp.team_id OR (h.team_id IS NULL AND macp.team_id = 0)
+LEFT JOIN host_mdm_apple_profiles hmap
+	ON h.uuid = hmap.host_uuid
 WHERE
-	macp.profile_uuid IN (?) AND (h.platform = 'darwin' OR h.platform = 'ios' OR h.platform = 'ipados')`
-		args = append(args, macProfUUIDs)
+	macp.profile_uuid IN (?) AND (h.platform = 'darwin' OR h.platform = 'ios' OR h.platform = 'ipados')
+OR
+	hmap.profile_uuid IN (?) AND (h.platform = 'darwin' OR h.platform = 'ios' OR h.platform = 'ipados')`
+		args = append(args, macProfUUIDs, macProfUUIDs)
 
 	case len(winProfUUIDs) > 0:
-		// TODO: if a very large number (~65K) of profile IDs was provided, could
+		// TODO: if a very large number (~65K/2) of profile IDs was provided, could
 		// result in too many placeholders (not an immediate concern).
 		uuidStmt = `
 SELECT DISTINCT h.uuid, h.platform
 FROM hosts h
 JOIN mdm_windows_configuration_profiles mawp
 	ON h.team_id = mawp.team_id OR (h.team_id IS NULL AND mawp.team_id = 0)
+LEFT JOIN host_mdm_windows_profiles hmwp
+	ON h.uuid = hmwp.host_uuid
 WHERE
-	mawp.profile_uuid IN (?) AND h.platform = 'windows'`
-		args = append(args, winProfUUIDs)
+	mawp.profile_uuid IN (?) AND h.platform = 'windows'
+OR
+	hmwp.profile_uuid IN (?) AND h.platform = 'windows'`
+		args = append(args, winProfUUIDs, winProfUUIDs)
 
 	}
 
@@ -515,12 +523,12 @@ WHERE
 		}
 	}
 
-	updates.AppleConfigProfile, err = ds.bulkSetPendingMDMAppleHostProfilesDB(ctx, tx, appleHosts)
+	updates.AppleConfigProfile, err = ds.bulkSetPendingMDMAppleHostProfilesDB(ctx, tx, appleHosts, profileUUIDs)
 	if err != nil {
 		return updates, ctxerr.Wrap(ctx, err, "bulk set pending apple host profiles")
 	}
 
-	updates.WindowsConfigProfile, err = ds.bulkSetPendingMDMWindowsHostProfilesDB(ctx, tx, winHosts)
+	updates.WindowsConfigProfile, err = ds.bulkSetPendingMDMWindowsHostProfilesDB(ctx, tx, winHosts, profileUUIDs)
 	if err != nil {
 		return updates, ctxerr.Wrap(ctx, err, "bulk set pending windows host profiles")
 	}
@@ -834,7 +842,7 @@ FROM
 		GROUP BY checksum
 	) cs ON macp.checksum = cs.checksum
 WHERE
-	macp.team_id = ? AND 
+	macp.team_id = ? AND
 	NOT EXISTS (
 		SELECT
 			1
@@ -865,16 +873,16 @@ FROM
 			mdm_apple_configuration_profiles
 		GROUP BY checksum
 	) cs ON macp.checksum = cs.checksum
-	JOIN mdm_configuration_profile_labels mcpl 
+	JOIN mdm_configuration_profile_labels mcpl
 		ON mcpl.apple_profile_uuid = macp.profile_uuid AND mcpl.exclude = 0
-	LEFT OUTER JOIN label_membership lm 
+	LEFT OUTER JOIN label_membership lm
 		ON lm.label_id = mcpl.label_id AND lm.host_id = ?
 WHERE
 	macp.team_id = ?
 GROUP BY
 	identifier
 HAVING
-	count_profile_labels > 0 AND 
+	count_profile_labels > 0 AND
 	count_host_labels = count_profile_labels
 
 UNION
@@ -897,17 +905,17 @@ FROM
 			mdm_apple_configuration_profiles
 		GROUP BY checksum
 	) cs ON macp.checksum = cs.checksum
-	JOIN mdm_configuration_profile_labels mcpl 
+	JOIN mdm_configuration_profile_labels mcpl
 		ON mcpl.apple_profile_uuid = macp.profile_uuid AND mcpl.exclude = 1
-	LEFT OUTER JOIN label_membership lm 
+	LEFT OUTER JOIN label_membership lm
 		ON lm.label_id = mcpl.label_id AND lm.host_id = ?
 WHERE
 	macp.team_id = ?
 GROUP BY
 	identifier
 HAVING
 	-- considers only the profiles with labels, without any broken label, and with the host not in any label
-	count_profile_labels > 0 AND 
+	count_profile_labels > 0 AND
 	count_profile_labels = count_non_broken_labels AND
 	count_host_labels = 0
 `

server/datastore/mysql/microsoft_mdm.go

@@ -1267,7 +1267,7 @@ func (ds *Datastore) ListMDMWindowsProfilesToInstall(ctx context.Context) ([]*fl
 	// be without and use the reader replica?
 	err := ds.withTx(ctx, func(tx sqlx.ExtContext) error {
 		var err error
-		result, err = listMDMWindowsProfilesToInstallDB(ctx, tx, nil)
+		result, err = listMDMWindowsProfilesToInstallDB(ctx, tx, nil, nil)
 		return err
 	})
 	return result, err
@@ -1277,6 +1277,7 @@ func listMDMWindowsProfilesToInstallDB(
 	ctx context.Context,
 	tx sqlx.ExtContext,
 	hostUUIDs []string,
+	onlyProfileUUIDs []string,
 ) ([]*fleet.MDMWindowsProfilePayload, error) {
 	// The query below is a set difference between:
 	//
@@ -1318,14 +1319,29 @@ func listMDMWindowsProfilesToInstallDB(
 
 	hostFilter := "TRUE"
 	if len(hostUUIDs) > 0 {
-		hostFilter = "h.uuid IN (?)"
+		if len(onlyProfileUUIDs) > 0 {
+			hostFilter = "mwcp.profile_uuid IN (?) AND h.uuid IN (?)"
+		} else {
+			hostFilter = "h.uuid IN (?)"
+		}
 	}
 
 	var err error
 	args := []any{fleet.MDMOperationTypeInstall}
 	query = fmt.Sprintf(query, hostFilter, hostFilter, hostFilter, hostFilter)
 	if len(hostUUIDs) > 0 {
-		query, args, err = sqlx.In(query, hostUUIDs, hostUUIDs, hostUUIDs, hostUUIDs, fleet.MDMOperationTypeInstall)
+		if len(onlyProfileUUIDs) > 0 {
+			query, args, err = sqlx.In(
+				query,
+				onlyProfileUUIDs, hostUUIDs,
+				onlyProfileUUIDs, hostUUIDs,
+				onlyProfileUUIDs, hostUUIDs,
+				onlyProfileUUIDs, hostUUIDs,
+				fleet.MDMOperationTypeInstall,
+			)
+		} else {
+			query, args, err = sqlx.In(query, hostUUIDs, hostUUIDs, hostUUIDs, hostUUIDs, fleet.MDMOperationTypeInstall)
+		}
 		if err != nil {
 			return nil, ctxerr.Wrap(ctx, err, "building sqlx.In")
 		}
@@ -1340,7 +1356,7 @@ func (ds *Datastore) ListMDMWindowsProfilesToRemove(ctx context.Context) ([]*fle
 	var result []*fleet.MDMWindowsProfilePayload
 	err := ds.withTx(ctx, func(tx sqlx.ExtContext) error {
 		var err error
-		result, err = listMDMWindowsProfilesToRemoveDB(ctx, tx, nil)
+		result, err = listMDMWindowsProfilesToRemoveDB(ctx, tx, nil, nil)
 		return err
 	})
 
@@ -1351,6 +1367,7 @@ func listMDMWindowsProfilesToRemoveDB(
 	ctx context.Context,
 	tx sqlx.ExtContext,
 	hostUUIDs []string,
+	onlyProfileUUIDs []string,
 ) ([]*fleet.MDMWindowsProfilePayload, error) {
 	// The query below is a set difference between:
 	//
@@ -1374,7 +1391,11 @@ func listMDMWindowsProfilesToRemoveDB(
 
 	hostFilter := "TRUE"
 	if len(hostUUIDs) > 0 {
-		hostFilter = "hmwp.host_uuid IN (?)"
+		if len(onlyProfileUUIDs) > 0 {
+			hostFilter = "hmwp.profile_uuid IN (?) AND hmwp.host_uuid IN (?)"
+		} else {
+			hostFilter = "hmwp.host_uuid IN (?)"
+		}
 	}
 
 	query := fmt.Sprintf(`
@@ -1408,7 +1429,11 @@ func listMDMWindowsProfilesToRemoveDB(
 	var err error
 	var args []any
 	if len(hostUUIDs) > 0 {
-		query, args, err = sqlx.In(query, hostUUIDs)
+		if len(onlyProfileUUIDs) > 0 {
+			query, args, err = sqlx.In(query, onlyProfileUUIDs, hostUUIDs)
+		} else {
+			query, args, err = sqlx.In(query, hostUUIDs)
+		}
 		if err != nil {
 			return nil, err
 		}
@@ -1917,18 +1942,19 @@ ON DUPLICATE KEY UPDATE
 func (ds *Datastore) bulkSetPendingMDMWindowsHostProfilesDB(
 	ctx context.Context,
 	tx sqlx.ExtContext,
-	uuids []string,
+	hostUUIDs []string,
+	onlyProfileUUIDs []string,
 ) (updatedDB bool, err error) {
-	if len(uuids) == 0 {
+	if len(hostUUIDs) == 0 {
 		return false, nil
 	}
 
-	profilesToInstall, err := listMDMWindowsProfilesToInstallDB(ctx, tx, uuids)
+	profilesToInstall, err := listMDMWindowsProfilesToInstallDB(ctx, tx, hostUUIDs, onlyProfileUUIDs)
 	if err != nil {
 		return false, ctxerr.Wrap(ctx, err, "list profiles to install")
 	}
 
-	profilesToRemove, err := listMDMWindowsProfilesToRemoveDB(ctx, tx, uuids)
+	profilesToRemove, err := listMDMWindowsProfilesToRemoveDB(ctx, tx, hostUUIDs, onlyProfileUUIDs)
 	if err != nil {
 		return false, ctxerr.Wrap(ctx, err, "list profiles to remove")
 	}

server/service/apple_mdm.go

@@ -847,7 +847,7 @@ func (svc *Service) DeleteMDMAppleConfigProfile(ctx context.Context, profileUUID
 	}
 
 	// cannot use the profile ID as it is now deleted
-	if _, err := svc.ds.BulkSetPendingMDMHostProfiles(ctx, nil, []uint{teamID}, nil, nil); err != nil {
+	if _, err := svc.ds.BulkSetPendingMDMHostProfiles(ctx, nil, nil, []string{profileUUID}, nil); err != nil {
 		return ctxerr.Wrap(ctx, err, "bulk set pending host profiles")
 	}