Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add batching logic when we pull windows profiles to install or remove #26964

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

add batching logic when we pull windows profiles to install or remove #26964

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
81774ec
add batching logic when we pull windows profiles to install or remove
jahzielv Mar 7, 2025
a382201
add a test
jahzielv Mar 7, 2025
1b548d9
make sure the queries run at least once to support cases where all th…
jahzielv Mar 7, 2025
46d856d
chore: changes file
jahzielv Mar 7, 2025
eaf6f67
Add comment about being repro-able at lower host count
iansltx Mar 8, 2025
ee76f32
Split "all hosts" and "specific hosts" handling for Windows MDM profi…
iansltx Mar 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions changes/26918-placeholders-windows
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
- Fixed an issue with assigning Windows MDM profiles to large numbers (> 65k) of hosts by batching
the relevant database queries.
14 changes: 14 additions & 0 deletions server/datastore/mysql/mdm_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ func TestMDMShared(t *testing.T) {
{"TestIsHostConnectedToFleetMDM", testIsHostConnectedToFleetMDM},
{"TestAreHostsConnectedToFleetMDM", testAreHostsConnectedToFleetMDM},
{"TestBulkSetPendingMDMHostProfilesExcludeAny", testBulkSetPendingMDMHostProfilesExcludeAny},
{"TestBulkSetPendingMDMHostProfilesLotsOfHosts", testBulkSetPendingMDMWindowsHostProfilesLotsOfHosts},
}

for _, c := range cases {
Expand Down Expand Up @@ -7614,3 +7615,16 @@ func testBulkSetPendingMDMHostProfilesExcludeAny(t *testing.T, ds *Datastore) {
winHost2: {},
})
}

func testBulkSetPendingMDMWindowsHostProfilesLotsOfHosts(t *testing.T, ds *Datastore) {
ctx := context.Background()

var hostUUIDs []string
// The bug this test was built to reproduce is visible down to ~16400 hosts; keeping this at 66k for scale testing
for range 66000 {
hostUUIDs = append(hostUUIDs, uuid.NewString())
}

_, err := ds.bulkSetPendingMDMWindowsHostProfilesDB(ctx, ds.writer(ctx), hostUUIDs, nil)
require.NoError(t, err)
}
134 changes: 96 additions & 38 deletions server/datastore/mysql/microsoft_mdm.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"encoding/xml"
"errors"
"fmt"
"math"
"strings"

"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
Expand Down Expand Up @@ -1247,13 +1248,13 @@ 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, nil)
result, err = ds.listMDMWindowsProfilesToInstallDB(ctx, tx, nil, nil)
return err
})
return result, err
}

func listMDMWindowsProfilesToInstallDB(
func (ds *Datastore) listMDMWindowsProfilesToInstallDB(
ctx context.Context,
tx sqlx.ExtContext,
hostUUIDs []string,
Expand Down Expand Up @@ -1282,7 +1283,7 @@ func listMDMWindowsProfilesToInstallDB(
// where one of the labels does not exist anymore, will not be considered for
// installation.

query := fmt.Sprintf(`
toInstallQuery := fmt.Sprintf(`
SELECT
ds.profile_uuid,
ds.host_uuid,
Expand All @@ -1306,44 +1307,73 @@ func listMDMWindowsProfilesToInstallDB(
}
}

var err error
args := []any{fleet.MDMOperationTypeInstall}
query = fmt.Sprintf(query, hostFilter, hostFilter, hostFilter, hostFilter)
if len(hostUUIDs) > 0 {
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)
toInstallQuery = fmt.Sprintf(toInstallQuery, hostFilter, hostFilter, hostFilter, hostFilter)

// use a 10k host batch size to match what we do on the macOS side.
selectProfilesBatchSize := 10_000
if ds.testSelectMDMProfilesBatchSize > 0 {
selectProfilesBatchSize = ds.testSelectMDMProfilesBatchSize
}
selectProfilesTotalBatches := int(math.Ceil(float64(len(hostUUIDs)) / float64(selectProfilesBatchSize)))

if len(hostUUIDs) == 0 {
// Make sure we run the query at least once
selectProfilesTotalBatches = 1
}

var wantedProfiles []*fleet.MDMWindowsProfilePayload
for i := range selectProfilesTotalBatches {
start := i * selectProfilesBatchSize
end := min(start+selectProfilesBatchSize, len(hostUUIDs))

batchUUIDs := hostUUIDs[start:end]

var err error
args := []any{fleet.MDMOperationTypeInstall}
stmt := toInstallQuery
if len(hostUUIDs) > 0 {
if len(onlyProfileUUIDs) > 0 {
stmt, args, err = sqlx.In(
toInstallQuery,
onlyProfileUUIDs, batchUUIDs,
onlyProfileUUIDs, batchUUIDs,
onlyProfileUUIDs, batchUUIDs,
onlyProfileUUIDs, batchUUIDs,
fleet.MDMOperationTypeInstall,
)
} else {
stmt, args, err = sqlx.In(toInstallQuery, batchUUIDs, batchUUIDs, batchUUIDs, batchUUIDs, fleet.MDMOperationTypeInstall)
}
if err != nil {
return nil, ctxerr.Wrapf(ctx, err, "building sqlx.In for list MDM windows profiles to install, batch %d of %d", i, selectProfilesTotalBatches)
}
}

var partialResult []*fleet.MDMWindowsProfilePayload
err = sqlx.SelectContext(ctx, tx, &partialResult, stmt, args...)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "building sqlx.In")
return nil, ctxerr.Wrapf(ctx, err, "selecting windows MDM profiles to install, batch %d of %d", i, selectProfilesTotalBatches)
}

wantedProfiles = append(wantedProfiles, partialResult...)

}

var profiles []*fleet.MDMWindowsProfilePayload
err = sqlx.SelectContext(ctx, tx, &profiles, query, args...)
return profiles, err
return wantedProfiles, nil
}

func (ds *Datastore) ListMDMWindowsProfilesToRemove(ctx context.Context) ([]*fleet.MDMWindowsProfilePayload, error) {
var result []*fleet.MDMWindowsProfilePayload
err := ds.withTx(ctx, func(tx sqlx.ExtContext) error {
var err error
result, err = listMDMWindowsProfilesToRemoveDB(ctx, tx, nil, nil)
result, err = ds.listMDMWindowsProfilesToRemoveDB(ctx, tx, nil, nil)
return err
})

return result, err
}

func listMDMWindowsProfilesToRemoveDB(
func (ds *Datastore) listMDMWindowsProfilesToRemoveDB(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same feedback as for the install case. Feels even more like we should do this refactor given that doing so doesn't require us to modify the datastore interface.

ctx context.Context,
tx sqlx.ExtContext,
hostUUIDs []string,
Expand Down Expand Up @@ -1378,7 +1408,7 @@ func listMDMWindowsProfilesToRemoveDB(
}
}

query := fmt.Sprintf(`
toRemoveQuery := fmt.Sprintf(`
SELECT
hmwp.profile_uuid,
hmwp.host_uuid,
Expand Down Expand Up @@ -1406,22 +1436,50 @@ func listMDMWindowsProfilesToRemoveDB(
(%s)
`, fmt.Sprintf(windowsMDMProfilesDesiredStateQuery, "TRUE", "TRUE", "TRUE", "TRUE"), hostFilter)

var err error
var args []any
if len(hostUUIDs) > 0 {
if len(onlyProfileUUIDs) > 0 {
query, args, err = sqlx.In(query, onlyProfileUUIDs, hostUUIDs)
} else {
query, args, err = sqlx.In(query, hostUUIDs)
// use a 10k host batch size to match what we do on the macOS side.
selectProfilesBatchSize := 10_000
if ds.testSelectMDMProfilesBatchSize > 0 {
selectProfilesBatchSize = ds.testSelectMDMProfilesBatchSize
}
selectProfilesTotalBatches := int(math.Ceil(float64(len(hostUUIDs)) / float64(selectProfilesBatchSize)))

if len(hostUUIDs) == 0 {
// Make sure we run the query at least once
selectProfilesTotalBatches = 1
}

var wantedProfiles []*fleet.MDMWindowsProfilePayload

for i := range selectProfilesTotalBatches {
start := i * selectProfilesBatchSize
end := min(start+selectProfilesBatchSize, len(hostUUIDs))

batchUUIDs := hostUUIDs[start:end]

var err error
var args []any
stmt := toRemoveQuery
if len(hostUUIDs) > 0 {
if len(onlyProfileUUIDs) > 0 {
stmt, args, err = sqlx.In(toRemoveQuery, onlyProfileUUIDs, batchUUIDs)
} else {
stmt, args, err = sqlx.In(toRemoveQuery, batchUUIDs)
}
if err != nil {
return nil, ctxerr.Wrapf(ctx, err, "building sqlx.In for list MDM windows profiles to remove, batch %d of %d", i, selectProfilesTotalBatches)
}
}

var partialResult []*fleet.MDMWindowsProfilePayload
err = sqlx.SelectContext(ctx, tx, &partialResult, stmt, args...)
if err != nil {
return nil, err
return nil, ctxerr.Wrapf(ctx, err, "selecting windows MDM profiles to remove, batch %d of %d", i, selectProfilesTotalBatches)
}

wantedProfiles = append(wantedProfiles, partialResult...)
}

var profiles []*fleet.MDMWindowsProfilePayload
err = sqlx.SelectContext(ctx, tx, &profiles, query, args...)
return profiles, err
return wantedProfiles, nil
}

func (ds *Datastore) BulkUpsertMDMWindowsHostProfiles(ctx context.Context, payload []*fleet.MDMWindowsBulkUpsertHostProfilePayload) error {
Expand Down Expand Up @@ -1930,12 +1988,12 @@ func (ds *Datastore) bulkSetPendingMDMWindowsHostProfilesDB(
return false, nil
}

profilesToInstall, err := listMDMWindowsProfilesToInstallDB(ctx, tx, hostUUIDs, onlyProfileUUIDs)
profilesToInstall, err := ds.listMDMWindowsProfilesToInstallDB(ctx, tx, hostUUIDs, onlyProfileUUIDs)
if err != nil {
return false, ctxerr.Wrap(ctx, err, "list profiles to install")
}

profilesToRemove, err := listMDMWindowsProfilesToRemoveDB(ctx, tx, hostUUIDs, onlyProfileUUIDs)
profilesToRemove, err := ds.listMDMWindowsProfilesToRemoveDB(ctx, tx, hostUUIDs, onlyProfileUUIDs)
if err != nil {
return false, ctxerr.Wrap(ctx, err, "list profiles to remove")
}
Expand Down
Loading