Improved Windows MDM performance - #43912
Conversation
|
@coderabbitai full review |
✅ Actions performedFull review triggered. |
WalkthroughThis pull request refactors Windows MDM pending command retrieval and device response handling to reduce database query overhead during device check-ins. The changes modify the datastore layer to accept enrollment identifiers instead of device identifiers, thread the already-resolved enrolled device through the service layer call chain to avoid redundant lookups, add deterministic ordering to enrollment queries, and introduce a fast-path existence check before executing the full pending-command query. Interface signatures, mock implementations, and test cases are updated accordingly. Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
server/service/microsoft_mdm.go (1)
1722-1760:⚠️ Potential issue | 🟠 MajorDon’t persist command responses before authentication succeeds.
For
RequestAuthStateChallengeandRequestAuthStateUnauthorized,saveResponsecan still insert command results, update profile state, and dequeue commands. Keep these states challenge-only, or persist only a raw audit response that cannot mutate command state.Proposed fix
responseCmds = append(responseCmds, ackMsg) - err = saveResponse([]string{}) - if err != nil { - return nil, err - } return responseCmds, nil }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@server/service/microsoft_mdm.go` around lines 1722 - 1760, The current flow for RequestAuthStateChallenge/RequestAuthStateUnauthorized builds the ackMsg (NewSyncMLCmdStatus) and then calls saveResponse([]string{}) which may persist command results and mutate device/profile state; change this so that for requestAuthState == RequestAuthStateChallenge || RequestAuthStateUnauthorized you do NOT call saveResponse or you call a non-mutating audit-only persistence path instead (e.g., add a new saveAuditResponse or a saveResponse flag) so only the nonce is stored via svc.keyValueStore.Set and the ackMsg is returned without dequeuing or marking commands done; update places referencing saveResponse, responseCmds, and ackMsg to use the non-persisting behavior.server/datastore/mysql/microsoft_mdm.go (1)
3379-3397:⚠️ Potential issue | 🟡 MinorAvoid leaving a trailing
ORwhen skipped command IDs appear at the end.If
failedCommandIdsends with"", the loop skips the final element and never trims the previous" OR ", producing invalid SQL beforeORDER BY.Proposed fix
args := []any{deviceID} - for idx, commandId := range failedCommandIds { + clauses := make([]string, 0, len(failedCommandIds)) + for _, commandId := range failedCommandIds { if commandId == "" { continue } - stmt += " wmc.raw_command LIKE ? OR " + clauses = append(clauses, "wmc.raw_command LIKE ?") args = append(args, "%<CmdID>"+commandId+"</CmdID>%") - if idx == len(failedCommandIds)-1 { - stmt = strings.TrimSuffix(stmt, " OR ") - } } if len(args) == 1 { // No valid command IDs were provided, return empty result to avoid returning all commands for the device. return []*fleet.MDMWindowsCommand{}, nil } - stmt += fmt.Sprintf(" ORDER BY created_at DESC LIMIT %d", len(failedCommandIds)) + stmt += " " + strings.Join(clauses, " OR ") + stmt += fmt.Sprintf(" ORDER BY created_at DESC LIMIT %d", len(clauses))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@server/datastore/mysql/microsoft_mdm.go` around lines 3379 - 3397, The SQL builder loop can leave a trailing " OR " when the last failedCommandIds entries are empty because the TrimSuffix is only applied inside the loop; fix by removing the per-iteration TrimSuffix and instead trim the trailing " OR " once after the loop (or build the WHERE fragments into a slice and strings.Join with " OR "). Ensure you append args with "%<CmdID>"+commandId+"</CmdID>%" only for non-empty commandId, keep the initial deviceID in args, and update the LIMIT to use the count of added command IDs (len(args)-1) rather than len(failedCommandIds); references: stmt, args, failedCommandIds, deviceID.
🧹 Nitpick comments (1)
server/datastore/mysql/microsoft_mdm_test.go (1)
1753-1756: Minor: the "non-existent" assertion is a bit weaker than before.Passing
0always misses because auto-increment enrollment ids start at 1, so this validates the cheap "no rows" path but can't catch a bug where the query accidentally drops theenrollment_idfilter and returns any queued row. Consider also asserting the empty result for a plausibly-valid-but-unused id (e.g.dID + 1000) so a regression that widens the filter would be caught.Suggested tweak
- // non-existent enrollment - cmds, err = ds.MDMWindowsGetPendingCommands(ctx, 0) - require.NoError(t, err) - require.Empty(t, cmds) + // non-existent enrollment (sentinel id 0 and an unused id well past any real row) + for _, missing := range []uint{0, dID + 1000} { + cmds, err = ds.MDMWindowsGetPendingCommands(ctx, missing) + require.NoError(t, err) + require.Empty(t, cmds) + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@server/datastore/mysql/microsoft_mdm_test.go` around lines 1753 - 1756, The test currently calls ds.MDMWindowsGetPendingCommands(ctx, 0) which only verifies the "no rows" path; add a second assertion using a plausibly-valid-but-unused enrollment id (e.g. use the existing dID variable and call ds.MDMWindowsGetPendingCommands(ctx, dID+1000)) to ensure the query still filters by enrollment_id (call require.NoError and require.Empty for that call). This targets the MDMWindowsGetPendingCommands code path and will catch regressions that drop or widen the enrollment_id filter.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@server/datastore/mysql/microsoft_mdm.go`:
- Around line 3379-3397: The SQL builder loop can leave a trailing " OR " when
the last failedCommandIds entries are empty because the TrimSuffix is only
applied inside the loop; fix by removing the per-iteration TrimSuffix and
instead trim the trailing " OR " once after the loop (or build the WHERE
fragments into a slice and strings.Join with " OR "). Ensure you append args
with "%<CmdID>"+commandId+"</CmdID>%" only for non-empty commandId, keep the
initial deviceID in args, and update the LIMIT to use the count of added command
IDs (len(args)-1) rather than len(failedCommandIds); references: stmt, args,
failedCommandIds, deviceID.
In `@server/service/microsoft_mdm.go`:
- Around line 1722-1760: The current flow for
RequestAuthStateChallenge/RequestAuthStateUnauthorized builds the ackMsg
(NewSyncMLCmdStatus) and then calls saveResponse([]string{}) which may persist
command results and mutate device/profile state; change this so that for
requestAuthState == RequestAuthStateChallenge || RequestAuthStateUnauthorized
you do NOT call saveResponse or you call a non-mutating audit-only persistence
path instead (e.g., add a new saveAuditResponse or a saveResponse flag) so only
the nonce is stored via svc.keyValueStore.Set and the ackMsg is returned without
dequeuing or marking commands done; update places referencing saveResponse,
responseCmds, and ackMsg to use the non-persisting behavior.
---
Nitpick comments:
In `@server/datastore/mysql/microsoft_mdm_test.go`:
- Around line 1753-1756: The test currently calls
ds.MDMWindowsGetPendingCommands(ctx, 0) which only verifies the "no rows" path;
add a second assertion using a plausibly-valid-but-unused enrollment id (e.g.
use the existing dID variable and call ds.MDMWindowsGetPendingCommands(ctx,
dID+1000)) to ensure the query still filters by enrollment_id (call
require.NoError and require.Empty for that call). This targets the
MDMWindowsGetPendingCommands code path and will catch regressions that drop or
widen the enrollment_id filter.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 45b37307-9e01-4529-9181-5293d07d785f
📒 Files selected for processing (8)
changes/43875-windows-mdm-pending-commands-floodserver/datastore/mysql/mdm_test.goserver/datastore/mysql/microsoft_mdm.goserver/datastore/mysql/microsoft_mdm_test.goserver/fleet/datastore.goserver/mock/datastore_mock.goserver/service/microsoft_mdm.goserver/service/microsoft_mdm_test.go
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #43912 +/- ##
=======================================
Coverage 65.08% 65.09%
=======================================
Files 2603 2603
Lines 253269 253273 +4
Branches 9356 9356
=======================================
+ Hits 164846 164862 +16
+ Misses 75656 75647 -9
+ Partials 12767 12764 -3
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review to trigger a review and subscribe this PR to future pushes, or @claude review once for a one-time review.
Tip: disable this comment in your organization's Code Review settings.
There was a problem hiding this comment.
Pull request overview
Optimizes the hot Windows MDM management check-in path (POST /api/mdm/microsoft/management) by reducing redundant datastore lookups and avoiding the expensive pending-commands query when the per-enrollment queue is empty.
Changes:
- Add a fast-path probe in
MDMWindowsGetPendingCommandsto short-circuit when no commands are queued, and query pending commands byenrollment_id(dropping the enrollments join). - Resolve the Windows enrollment once during request trust evaluation and thread the resulting
enrolledDevicethrough downstream service logic to avoid repeated lookups. - Update mocks/tests and datastore interfaces accordingly; add a release note entry.
Reviewed changes
Copilot reviewed 7 out of 8 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| server/service/microsoft_mdm.go | Threads enrolledDevice from trust check through management response flow; switches pending-commands lookup to enrollment ID. |
| server/datastore/mysql/microsoft_mdm.go | Adds EXISTS probe fast-path and changes pending-commands query to use enrollment_id; makes enrollment selection deterministic with created_at, id ordering; updates save-response signature to accept the already-resolved enrollment. |
| server/fleet/datastore.go | Updates datastore interface signatures for Windows pending commands and save response. |
| server/mock/datastore_mock.go | Updates mock signatures to match datastore interface changes. |
| server/datastore/mysql/microsoft_mdm_test.go | Updates MySQL datastore tests to use enrollment IDs and to reflect per-enrollment pending command semantics. |
| server/datastore/mysql/mdm_test.go | Updates Windows MDM command test to pass full enrollment into MDMWindowsSaveResponse. |
| server/service/microsoft_mdm_test.go | Updates service test mocking to validate enrollment ID is threaded through pending-commands lookup. |
| changes/43875-windows-mdm-pending-commands-flood | Adds release note for Windows MDM performance improvement. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // isTrustedRequest checks if the incoming request was sent from an MDM-enrolled | ||
| // device. It returns the matched enrollment (when the device was found), the | ||
| // auth state, and an error only when the request is malformed or otherwise | ||
| // cannot be processed. Expected non-trusted outcomes (for example | ||
| // RequestAuthStateChallenge or RequestAuthStateUnauthorized) are reported via | ||
| // the returned auth state and may return a nil error. The returned enrolled | ||
| // device may be nil when the state is RequestAuthStateUntrusted. | ||
| func (svc *Service) isTrustedRequest(ctx context.Context, reqSyncML *fleet.SyncML, reqCerts []*x509.Certificate) (*fleet.MDMWindowsEnrolledDevice, requestAuthState, error) { | ||
| if reqSyncML == nil { | ||
| return RequestAuthStateUntrusted, fleet.NewInvalidArgumentError("syncml req message", "message is not present") | ||
| return nil, RequestAuthStateUntrusted, fleet.NewInvalidArgumentError("syncml req message", "message is not present") | ||
| } | ||
|
|
||
| // Checking if calling request is coming from an already MDM enrolled device | ||
| deviceID, err := reqSyncML.GetSource() | ||
| if err != nil || deviceID == "" { | ||
| return RequestAuthStateUntrusted, fmt.Errorf("invalid SyncML message %w", err) | ||
| return nil, RequestAuthStateUntrusted, fmt.Errorf("invalid SyncML message %w", err) | ||
| } | ||
|
|
||
| enrolledDevice, err := svc.ds.MDMWindowsGetEnrolledDeviceWithDeviceID(ctx, deviceID) | ||
| if err != nil || enrolledDevice == nil { | ||
| return RequestAuthStateUntrusted, errors.New("device was not MDM enrolled") | ||
| return nil, RequestAuthStateUntrusted, errors.New("device was not MDM enrolled") | ||
| } |
There was a problem hiding this comment.
The doc comment for isTrustedRequest says it returns an error only for malformed/unprocessable requests and that expected non-trusted outcomes are communicated via the auth state, but the implementation still returns a non-nil error when the device isn't enrolled (RequestAuthStateUntrusted). Please update the comment to match the actual error semantics (or adjust the function/caller behavior so that 'not enrolled' is treated as a non-error auth state).
Related issue: Resolves #43875
POST /api/mdm/microsoft/managementis the hot endpoint for any Windows-MDM-enrolledfleet. Every enrolled host hits it twice per check-in interval. At 40k hosts that's a
four-figure sustained queries-per-second rate on the database reader pool, dominated
by one expensive query plus a handful of redundant
MDMWindowsGetEnrolledDeviceWithDeviceIDlookups on the same row.
This PR cuts that load by:
overwhelming common case). Replaces a 3-table join plus anti-join with a cheap
primary-key probe.
isTrustedRequestand threading itthrough to every downstream consumer instead of re-fetching it three times.
No behavior change to the protocol, no schema change. Also filed a related issue: #43897
Checklist for submitter
If some of the following don't apply, delete the relevant line.
changes/,orbit/changes/oree/fleetd-chrome/changes.Testing
For unreleased bug fixes in a release candidate, one of:
Summary by CodeRabbit