Speedup worker-based device release on ADE enrollment setup (#29892)

This commit is contained in:
Martin Angers
2025-06-16 13:14:25 -04:00
committed by GitHub
parent 4ab8208231
commit fbc8fc031a
9 changed files with 354 additions and 24 deletions
@@ -0,0 +1 @@
* Improved releasing a macOS device during ADE enrollment, by increasing the frequency of checks for readiness.
+8
View File
@@ -998,6 +998,14 @@ func newCleanupsAndAggregationSchedule(
schedule.WithJob("cleanup_host_mdm_apple_profiles", func(ctx context.Context) error {
return ds.CleanupHostMDMAppleProfiles(ctx)
}),
schedule.WithJob("cleanup_worker_jobs", func(ctx context.Context) error {
const (
failedSince = 365 * 24 * time.Hour // keep failed jobs for 1 year
completedSince = 90 * 24 * time.Hour // keep completed (successful) jobs for ~3 months
)
_, err := ds.CleanupWorkerJobs(ctx, failedSince, completedSince)
return err
}),
)
return s, nil
+27
View File
@@ -4,6 +4,7 @@ import (
"context"
"time"
"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/jmoiron/sqlx"
)
@@ -84,3 +85,29 @@ WHERE
return job, nil
}
func (ds *Datastore) CleanupWorkerJobs(ctx context.Context, failedSince, completedSince time.Duration) (int64, error) {
// using not_before instead of created_at/updated_at to be able to use the
// existing index, and the difference between those timestamps will be
// minimal (max 5 retries for failed jobs, with a few hours difference).
const stmt = `
DELETE FROM
jobs
WHERE
(state = ? AND not_before < ?) OR
(state = ? AND not_before < ?)
`
now := time.Now().UTC()
failedBefore := now.Add(-failedSince)
completedBefore := now.Add(-completedSince)
res, err := ds.writer(ctx).ExecContext(ctx, stmt,
fleet.JobStateFailure, failedBefore,
fleet.JobStateSuccess, completedBefore)
if err != nil {
return 0, ctxerr.Wrap(ctx, err, "cleanup worker jobs")
}
n, _ := res.RowsAffected()
return n, nil
}
+88
View File
@@ -2,6 +2,7 @@ package mysql
import (
"context"
"fmt"
"testing"
"time"
@@ -20,6 +21,7 @@ func TestJobs(t *testing.T) {
fn func(t *testing.T, ds *Datastore)
}{
{"QueueAndProcessJobs", testQueueAndProcessJobs},
{"CleanupWorkerJobs", testCleanupWorkerJobs},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
@@ -78,3 +80,89 @@ func testQueueAndProcessJobs(t *testing.T, ds *Datastore) {
require.NotZero(t, jobs[0].NotBefore)
require.False(t, jobs[0].NotBefore.After(time.Now())) // before or equal
}
func testCleanupWorkerJobs(t *testing.T, ds *Datastore) {
ctx := t.Context()
// no job yet
n, err := ds.CleanupWorkerJobs(ctx, time.Second, time.Second)
require.NoError(t, err)
require.EqualValues(t, 0, n)
setJobTimestamp := func(j *fleet.Job, subtract time.Duration) {
j.NotBefore = time.Now().UTC().Add(-subtract)
_, err = ds.UpdateJob(ctx, j.ID, j)
require.NoError(t, err)
}
setJobStatus := func(j *fleet.Job, state fleet.JobState) {
j.State = state
_, err = ds.UpdateJob(ctx, j.ID, j)
require.NoError(t, err)
}
// enqueue a job
j1 := &fleet.Job{Name: "j1", State: fleet.JobStateQueued}
j1, err = ds.NewJob(ctx, j1)
require.NoError(t, err)
setJobTimestamp(j1, 2*time.Second)
// still clears nothing as it is not in a final state
n, err = ds.CleanupWorkerJobs(ctx, time.Second, time.Second)
require.NoError(t, err)
require.EqualValues(t, 0, n)
// job is still returned as a queued job
jobs, err := ds.GetQueuedJobs(ctx, 10, time.Time{})
require.NoError(t, err)
require.Len(t, jobs, 1)
require.Equal(t, j1.ID, jobs[0].ID)
// mark it as done
setJobStatus(j1, fleet.JobStateSuccess)
// does not clear it if the completed duration is not far enough in the past
n, err = ds.CleanupWorkerJobs(ctx, time.Second, time.Minute)
require.NoError(t, err)
require.EqualValues(t, 0, n)
// does clear it if the completed duration is far enough in the past
n, err = ds.CleanupWorkerJobs(ctx, time.Second, time.Second)
require.NoError(t, err)
require.EqualValues(t, 1, n)
jobs, err = ds.GetQueuedJobs(ctx, 10, time.Time{})
require.NoError(t, err)
require.Len(t, jobs, 0)
// enqueue a few more jobs
queuedJobs := make([]*fleet.Job, 0, 4)
for i := 0; i < 4; i++ {
j := &fleet.Job{Name: "j" + fmt.Sprint(i+1), State: fleet.JobStateQueued}
j, err = ds.NewJob(ctx, j)
require.NoError(t, err)
setJobTimestamp(j, time.Duration(i+1)*time.Minute)
queuedJobs = append(queuedJobs, j)
}
// make jobs[1] and jobs[3] failed, jobs[2] successful, jobs[0] queued
setJobStatus(queuedJobs[1], fleet.JobStateFailure) // failed 2m ago
setJobStatus(queuedJobs[3], fleet.JobStateFailure) // failed 4m ago
setJobStatus(queuedJobs[2], fleet.JobStateSuccess) // successful 3m ago
// cleanup failed > 3m, success > 1m, should delete jobs[3] and jobs[2]
n, err = ds.CleanupWorkerJobs(ctx, 3*time.Minute, time.Minute)
require.NoError(t, err)
require.EqualValues(t, 2, n)
// cleanup failed > 1m, success > 10m, should delete jobs[1]
n, err = ds.CleanupWorkerJobs(ctx, 1*time.Minute, 10*time.Minute)
require.NoError(t, err)
require.EqualValues(t, 1, n)
// jobs[0] is still queued
jobs, err = ds.GetQueuedJobs(ctx, 10, time.Time{})
require.NoError(t, err)
require.Len(t, jobs, 1)
require.Equal(t, queuedJobs[0].ID, jobs[0].ID)
}
+4
View File
@@ -1042,6 +1042,10 @@ type Datastore interface {
// UpdateJobs updates an existing job. Call this after processing a job.
UpdateJob(ctx context.Context, id uint, job *Job) (*Job, error)
// CleanupWorkerJobs deletes jobs in a final state that are older than the
// provided durations. It returns the number of jobs deleted and an error.
CleanupWorkerJobs(ctx context.Context, failedSince, completedSince time.Duration) (int64, error)
///////////////////////////////////////////////////////////////////////////////
// Debug
+12
View File
@@ -744,6 +744,8 @@ type GetQueuedJobsFunc func(ctx context.Context, maxNumJobs int, now time.Time)
type UpdateJobFunc func(ctx context.Context, id uint, job *fleet.Job) (*fleet.Job, error)
type CleanupWorkerJobsFunc func(ctx context.Context, failedSince time.Duration, completedSince time.Duration) (int64, error)
type InnoDBStatusFunc func(ctx context.Context) (string, error)
type ProcessListFunc func(ctx context.Context) ([]fleet.MySQLProcess, error)
@@ -2472,6 +2474,9 @@ type DataStore struct {
UpdateJobFunc UpdateJobFunc
UpdateJobFuncInvoked bool
CleanupWorkerJobsFunc CleanupWorkerJobsFunc
CleanupWorkerJobsFuncInvoked bool
InnoDBStatusFunc InnoDBStatusFunc
InnoDBStatusFuncInvoked bool
@@ -5968,6 +5973,13 @@ func (s *DataStore) UpdateJob(ctx context.Context, id uint, job *fleet.Job) (*fl
return s.UpdateJobFunc(ctx, id, job)
}
func (s *DataStore) CleanupWorkerJobs(ctx context.Context, failedSince time.Duration, completedSince time.Duration) (int64, error) {
s.mu.Lock()
s.CleanupWorkerJobsFuncInvoked = true
s.mu.Unlock()
return s.CleanupWorkerJobsFunc(ctx, failedSince, completedSince)
}
func (s *DataStore) InnoDBStatus(ctx context.Context) (string, error) {
s.mu.Lock()
s.InnoDBStatusFuncInvoked = true
+53 -17
View File
@@ -60,10 +60,12 @@ type appleMDMArgs struct {
// associated with the device.
//
// FIXME: Rename this to IdPAccountUUID or something similar.
EnrollReference string `json:"enroll_reference,omitempty"`
EnrollmentCommands []string `json:"enrollment_commands,omitempty"`
Platform string `json:"platform,omitempty"`
UseWorkerDeviceRelease bool `json:"use_worker_device_release,omitempty"`
EnrollReference string `json:"enroll_reference,omitempty"`
EnrollmentCommands []string `json:"enrollment_commands,omitempty"`
Platform string `json:"platform,omitempty"`
UseWorkerDeviceRelease bool `json:"use_worker_device_release,omitempty"`
ReleaseDeviceAttempt int `json:"release_device_attempt,omitempty"` // number of attempts to release the device
ReleaseDeviceStartedAt *time.Time `json:"release_device_started_at,omitempty"` // time when the release device task first started
}
// Run executes the apple_mdm job.
@@ -268,13 +270,12 @@ func (a *AppleMDM) getIdPDisplayName(ctx context.Context, acct *fleet.MDMIdPAcco
return scimUser.DisplayName(), nil
}
// This job is deprecated for macos because releasing devices is now done via
// the orbit endpoint /setup_experience/status that is polled by a swift dialog
// UI window during the setup process (unless there are no setup experience
// items, in which case this worker job is used), and automatically releases
// the device once all pending setup tasks are done. However, it must remain
// implemented for iOS and iPadOS and in case there are such jobs to process
// after a Fleet migration to a new version.
// This job is used only for iDevices or for macos devices that don't use any
// setup experience items (software installs, script exec) - see
// appleMDMArgs.UseWorkerDeviceRelease. Otherwise releasing devices is now done
// via the orbit endpoint /setup_experience/status that is polled by a swift
// dialog UI window during the setup process, and automatically releases the
// device once all pending setup tasks are done.
func (a *AppleMDM) runPostDEPReleaseDevice(ctx context.Context, args appleMDMArgs) error {
// Edge cases:
// - if the device goes offline for a long time, should we go ahead and
@@ -289,20 +290,49 @@ func (a *AppleMDM) runPostDEPReleaseDevice(ctx context.Context, args appleMDMArg
// We opted "yes" to all those, and we want to release after a few minutes,
// not hours, so we'll allow only a couple retries.
const (
maxWaitTime = 15 * time.Minute
minAttempts = 10
maxAttempts = 30
nextAttemptMinDelay = 30 * time.Second
)
args.ReleaseDeviceAttempt++
if args.ReleaseDeviceStartedAt == nil {
now := time.Now().UTC()
args.ReleaseDeviceStartedAt = &now
}
level.Debug(a.Log).Log(
"task", "runPostDEPReleaseDevice",
"msg", fmt.Sprintf("awaiting commands %v and profiles to settle for host %s", args.EnrollmentCommands, args.HostUUID),
"attempt", args.ReleaseDeviceAttempt,
"started_at", args.ReleaseDeviceStartedAt.Format(time.RFC3339),
)
if retryNum, _ := ctx.Value(retryNumberCtxKey).(int); retryNum > 2 {
// give up and release the device
a.Log.Log("info", "releasing device after too many attempts", "host_uuid", args.HostUUID, "retries", retryNum)
// if we've reached the minimum number of attempts and the maximum time to
// wait, we release the device even if some commands or profiles are still
// pending. We also release in case it reached the maximum number of
// attempts, to prevent an issue with clock skew where the wait delay does
// not appear to be reached.
if (args.ReleaseDeviceAttempt >= minAttempts && time.Since(*args.ReleaseDeviceStartedAt) >= maxWaitTime) ||
(args.ReleaseDeviceAttempt >= maxAttempts) {
a.Log.Log("info", "releasing device after too many attempts or too long wait", "host_uuid", args.HostUUID, "attempts", args.ReleaseDeviceAttempt)
if err := a.Commander.DeviceConfigured(ctx, args.HostUUID, uuid.NewString()); err != nil {
return ctxerr.Wrapf(ctx, err, "failed to enqueue DeviceConfigured command after %d retries", retryNum)
return ctxerr.Wrapf(ctx, err, "failed to enqueue DeviceConfigured command after %d attempts", args.ReleaseDeviceAttempt)
}
return nil
}
reenqueueTask := func() error {
// re-enqueue the same job, but now
// ReleaseDeviceAttempt/ReleaseDeviceStartedAt have been incremented/set,
// and run it not before a delay so it doesn't run again until the next
// worker cycle.
_, err := QueueJobWithDelay(ctx, a.Datastore, appleMDMJobName, args, nextAttemptMinDelay)
return err
}
for _, cmdUUID := range args.EnrollmentCommands {
if cmdUUID == "" {
continue
@@ -326,7 +356,10 @@ func (a *AppleMDM) runPostDEPReleaseDevice(ctx context.Context, args appleMDMArg
if !completed {
// DEP enrollment commands are not done being delivered to that device,
// cannot release it now.
return fmt.Errorf("device not ready for release, still awaiting result for command %s, will retry", cmdUUID)
if err := reenqueueTask(); err != nil {
return fmt.Errorf("failed to re-enqueue task: %w", err)
}
return nil
}
level.Debug(a.Log).Log(
"task", "runPostDEPReleaseDevice",
@@ -350,7 +383,10 @@ func (a *AppleMDM) runPostDEPReleaseDevice(ctx context.Context, args appleMDMArg
// if it has any pending profiles, then its profiles are not done being
// delivered (installed or removed).
if prof.Status == nil || *prof.Status == fleet.MDMDeliveryPending {
return fmt.Errorf("device not ready for release, profile %s is still pending, will retry", prof.Identifier)
if err := reenqueueTask(); err != nil {
return fmt.Errorf("failed to re-enqueue task: %w", err)
}
return nil
}
level.Debug(a.Log).Log(
"task", "runPostDEPReleaseDevice",
+161
View File
@@ -3,6 +3,7 @@ package worker
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"os"
@@ -604,6 +605,166 @@ func TestAppleMDM(t *testing.T) {
require.Equal(t, appleMDMJobName, jobs[0].Name)
require.Contains(t, string(*jobs[0].Args), AppleMDMPostDEPReleaseDeviceTask)
})
t.Run("automatic release retries and give up", func(t *testing.T) {
mysql.SetTestABMAssets(t, ds, testOrgName)
defer mysql.TruncateTables(t, ds)
h := createEnrolledHost(t, 1, nil, true)
mdmWorker := &AppleMDM{
Datastore: ds,
Log: nopLog,
Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
}
w := NewWorker(ds, nopLog)
w.Register(mdmWorker)
err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", true)
require.NoError(t, err)
// run the worker, should succeed
err = w.ProcessJobs(ctx)
require.NoError(t, err)
// ensure the job's not_before allows it to be returned if it were to run
// again
time.Sleep(time.Second)
require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))
// the release device job got enqueued, and it will constantly re-enqueue
// itself because the command is never acknowledged
var (
previousID uint
firstStartedAt time.Time
)
for i := 0; i <= 10; i++ {
jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute)) // release job is always added with a delay
require.NoError(t, err)
require.Len(t, jobs, 1)
releaseJob := jobs[0]
require.Equal(t, fleet.JobStateQueued, releaseJob.State)
require.Equal(t, appleMDMJobName, releaseJob.Name)
require.NotEqual(t, previousID, releaseJob.ID)
previousID = releaseJob.ID
var args appleMDMArgs
err = json.Unmarshal([]byte(*releaseJob.Args), &args)
require.NoError(t, err)
require.Equal(t, args.Task, AppleMDMPostDEPReleaseDeviceTask)
require.EqualValues(t, i, args.ReleaseDeviceAttempt)
if i == 0 {
// first time, there is no release device started at
require.Nil(t, args.ReleaseDeviceStartedAt)
} else {
require.NotNil(t, args.ReleaseDeviceStartedAt)
if i == 1 {
firstStartedAt = *args.ReleaseDeviceStartedAt
} else {
require.True(t, firstStartedAt.Equal(*args.ReleaseDeviceStartedAt))
}
}
if i == 10 {
// finally, after 10 attempts, update the release started at to make it
// meet the maximum wait time and actually do the release on the next
// processing.
startedAt := firstStartedAt.Add(-time.Hour)
args.ReleaseDeviceStartedAt = &startedAt
b, err := json.Marshal(args)
require.NoError(t, err)
mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
_, err := q.ExecContext(ctx, `UPDATE jobs SET args = ? WHERE id = ?`, string(b), releaseJob.ID)
return err
})
}
// update the job to make it available to run immediately
releaseJob.NotBefore = time.Now().UTC().Add(-time.Minute)
_, err = ds.UpdateJob(ctx, releaseJob.ID, releaseJob)
require.NoError(t, err)
// run the worker, should succeed and re-enqueue a new job with the same args
err = w.ProcessJobs(ctx)
require.NoError(t, err)
}
// on the last processing, it did end up releasing the device due to the
// limit of attempts and wait delay being reached.
require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "DeviceConfigured"}, getEnqueuedCommandTypes(t))
// job queue is now empty
jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute))
require.NoError(t, err)
require.Len(t, jobs, 0)
})
t.Run("automatic release succeeds after a few attempts", func(t *testing.T) {
mysql.SetTestABMAssets(t, ds, testOrgName)
defer mysql.TruncateTables(t, ds)
h := createEnrolledHost(t, 1, nil, true)
mdmWorker := &AppleMDM{
Datastore: ds,
Log: nopLog,
Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
}
w := NewWorker(ds, nopLog)
w.Register(mdmWorker)
err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", true)
require.NoError(t, err)
// run the worker, should succeed
err = w.ProcessJobs(ctx)
require.NoError(t, err)
// ensure the job's not_before allows it to be returned if it were to run
// again
time.Sleep(time.Second)
require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))
for i := 0; i <= 4; i++ {
jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute)) // release job is always added with a delay
require.NoError(t, err)
require.Len(t, jobs, 1)
releaseJob := jobs[0]
require.Equal(t, fleet.JobStateQueued, releaseJob.State)
require.Equal(t, appleMDMJobName, releaseJob.Name)
if i == 4 {
// after 4 attempts, record a result for the command so it gets released
mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
_, err := q.ExecContext(ctx, `INSERT INTO nano_command_results (id, command_uuid, status, result)
SELECT ?, command_uuid, ?, ? FROM nano_commands`,
h.UUID, "Acknowledged", `<?xml`)
return err
})
}
// update the job to make it available to run immediately
releaseJob.NotBefore = time.Now().UTC().Add(-time.Minute)
_, err = ds.UpdateJob(ctx, releaseJob.ID, releaseJob)
require.NoError(t, err)
// run the worker, should succeed and re-enqueue a new job with the same args
err = w.ProcessJobs(ctx)
require.NoError(t, err)
}
// on the last processing, it did release the device due to all pending
// commands being completed.
require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "DeviceConfigured"}, getEnqueuedCommandTypes(t))
// job queue is now empty
jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute))
require.NoError(t, err)
require.Len(t, jobs, 0)
})
}
func TestGetSignedURL(t *testing.T) {
-7
View File
@@ -12,17 +12,11 @@ import (
"github.com/go-kit/log/level"
)
type ctxKey int
const (
maxRetries = 5
// nvdCVEURL is the base link to a CVE on the NVD website, only the CVE code
// needs to be appended to make it a valid link.
nvdCVEURL = "https://nvd.nist.gov/vuln/detail/"
// context key for the retry number of a job, made available via the context
// to the job processor.
retryNumberCtxKey = ctxKey(0)
)
const (
@@ -209,7 +203,6 @@ func (w *Worker) processJob(ctx context.Context, job *fleet.Job) error {
args = *job.Args
}
ctx = context.WithValue(ctx, retryNumberCtxKey, job.Retries)
return j.Run(ctx, args)
}