diff --git a/cmd/fleet/cron.go b/cmd/fleet/cron.go index 351aff8ffe..1832c2c390 100644 --- a/cmd/fleet/cron.go +++ b/cmd/fleet/cron.go @@ -1504,6 +1504,41 @@ func cronHostVitalsLabelMembership( return nil } +func newBatchActivityCompletionCheckerSchedule( + ctx context.Context, + instanceID string, + ds fleet.Datastore, + logger kitlog.Logger, +) (*schedule.Schedule, error) { + const ( + name = string(fleet.CronBatchActivityCompletionChecker) + interval = 5 * time.Minute + ) + logger = kitlog.With(logger, "cron", name) + s := schedule.New( + ctx, name, instanceID, interval, ds, ds, + schedule.WithLogger(logger), + schedule.WithJob( + "cron_batch_activity_completion_checker", + func(ctx context.Context) error { + return cronBatchActivityCompletionChecker(ctx, ds) + }, + ), + ) + return s, nil +} + +func cronBatchActivityCompletionChecker( + ctx context.Context, + ds fleet.Datastore, +) error { + if err := ds.MarkActivitiesAsCompleted(ctx); err != nil { + return ctxerr.Wrap(ctx, err, "mark batch activities as completed") + } + // TODO -- add an entry in the global feed for each completed activity? + return nil +} + func stringSliceToUintSlice(s []string, logger kitlog.Logger) []uint { result := make([]uint, 0, len(s)) for _, v := range s { diff --git a/cmd/fleet/serve.go b/cmd/fleet/serve.go index 50265995c7..ac51dd9f7b 100644 --- a/cmd/fleet/serve.go +++ b/cmd/fleet/serve.go @@ -1045,6 +1045,13 @@ the way that the Fleet server works. initFatal(err, "failed to register host vitals label membership schedule") } + // Start the service that marks activities as completed. + if err := cronSchedules.StartCronSchedule(func() (fleet.CronSchedule, error) { + return newBatchActivityCompletionCheckerSchedule(ctx, instanceID, ds, logger) + }); err != nil { + initFatal(err, "failed to register batch activity completion checker schedule") + } + level.Info(logger).Log("msg", fmt.Sprintf("started cron schedules: %s", strings.Join(cronSchedules.ScheduleNames(), ", "))) // StartCollectors starts a goroutine per collector, using ctx to cancel. diff --git a/server/datastore/mysql/scripts.go b/server/datastore/mysql/scripts.go index 1bd4f88dff..0be2fdadb1 100644 --- a/server/datastore/mysql/scripts.go +++ b/server/datastore/mysql/scripts.go @@ -2227,3 +2227,47 @@ WHERE return count, nil } + +func (ds *Datastore) markActivitiesAsCompleted(ctx context.Context, tx sqlx.ExtContext) error { + const stmt = ` +UPDATE batch_activities AS ba +JOIN ( + SELECT + ba2.id AS batch_id, + COUNT(*) AS num_targeted, + COUNT(bahr.error) AS num_incompatible, + COUNT(IF(hsr.exit_code = 0, 1, NULL)) AS num_ran, + COUNT(IF(hsr.exit_code > 0, 1, NULL)) AS num_errored, + COUNT(IF(hsr.canceled = 1 AND hsr.exit_code IS NULL, 1, NULL)) AS num_canceled + FROM batch_activity_host_results AS bahr + LEFT JOIN host_script_results AS hsr + ON bahr.host_execution_id = hsr.execution_id + JOIN batch_activities AS ba2 + ON ba2.execution_id = bahr.batch_execution_id + WHERE ba2.status = 'started' + GROUP BY ba2.id + HAVING (num_incompatible + num_ran + num_errored + num_canceled) >= num_targeted +) AS agg + ON agg.batch_id = ba.id +SET + ba.status = 'finished', + ba.finished_at = NOW(), + ba.num_targeted = agg.num_targeted, + ba.num_incompatible = agg.num_incompatible, + ba.num_ran = agg.num_ran, + ba.num_errored = agg.num_errored, + ba.num_canceled = agg.num_canceled, + ba.num_pending = 0 +WHERE ba.status = 'started'; +` + // TODO -- use `RETURNING` to return the IDs of the updated activities? + _, err := tx.ExecContext(ctx, stmt) + if err != nil { + return ctxerr.Wrap(ctx, err, "marking activities as completed") + } + return nil +} + +func (ds *Datastore) MarkActivitiesAsCompleted(ctx context.Context) error { + return ds.markActivitiesAsCompleted(ctx, ds.writer(ctx)) +} diff --git a/server/datastore/mysql/scripts_test.go b/server/datastore/mysql/scripts_test.go index 94a04089bf..89da44e5bd 100644 --- a/server/datastore/mysql/scripts_test.go +++ b/server/datastore/mysql/scripts_test.go @@ -46,6 +46,7 @@ func TestScripts(t *testing.T) { {"BatchExecute", testBatchExecute}, {"BatchExecuteWithStatus", testBatchExecuteWithStatus}, {"BatchScriptSchedule", testBatchScriptSchedule}, + {"TestMarkActivitiesAsCompleted", testMarkActivitiesAsCompleted}, {"DeleteScriptActivatesNextActivity", testDeleteScriptActivatesNextActivity}, {"BatchSetScriptActivatesNextActivity", testBatchSetScriptActivatesNextActivity}, } @@ -2217,6 +2218,100 @@ func testBatchScriptSchedule(t *testing.T, ds *Datastore) { require.Len(t, hostResults, 3) } +func testMarkActivitiesAsCompleted(t *testing.T, ds *Datastore) { + ctx := context.Background() + + user := test.NewUser(t, ds, "user1", "user@example.com", true) + + team1, err := ds.NewTeam(ctx, &fleet.Team{Name: "team1"}) + require.NoError(t, err) + + hostNoScripts := test.NewHost(t, ds, "hostNoScripts", "10.0.0.1", "hostnoscripts", "hostnoscriptsuuid", time.Now()) + hostWindows := test.NewHost(t, ds, "hostWin", "10.0.0.2", "hostWinKey", "hostWinUuid", time.Now(), test.WithPlatform("windows")) + host1 := test.NewHost(t, ds, "host1", "10.0.0.3", "host1key", "host1uuid", time.Now()) + host2 := test.NewHost(t, ds, "host2", "10.0.0.4", "host2key", "host2uuid", time.Now()) + host3 := test.NewHost(t, ds, "host3", "10.0.0.4", "host3key", "host3uuid", time.Now()) + hostTeam1 := test.NewHost(t, ds, "hostTeam1", "10.0.0.5", "hostTeam1key", "hostTeam1uuid", time.Now(), test.WithTeamID(team1.ID)) + + test.SetOrbitEnrollment(t, hostWindows, ds) + test.SetOrbitEnrollment(t, host1, ds) + test.SetOrbitEnrollment(t, host2, ds) + test.SetOrbitEnrollment(t, host3, ds) + test.SetOrbitEnrollment(t, hostTeam1, ds) + + script, err := ds.NewScript(ctx, &fleet.Script{ + Name: "script1.sh", + ScriptContents: "echo hi", + }) + require.NoError(t, err) + + // Actual good execution + execID, err := ds.BatchExecuteScript(ctx, &user.ID, script.ID, []uint{hostNoScripts.ID, hostWindows.ID, host1.ID, host2.ID, host3.ID}) + require.NoError(t, err) + require.NotEmpty(t, execID) + + // Schedule another one + execID2, err := ds.BatchExecuteScript(ctx, &user.ID, script.ID, []uint{hostNoScripts.ID, hostWindows.ID, host1.ID, host2.ID, host3.ID}) + require.NoError(t, err) + require.NotEmpty(t, execID2) + + // Get the upcoming activities for each host + host1Upcoming, err := ds.listUpcomingHostScriptExecutions(ctx, host1.ID, false, false) + require.NoError(t, err) + host2Upcoming, err := ds.listUpcomingHostScriptExecutions(ctx, host2.ID, false, false) + require.NoError(t, err) + host3Upcoming, err := ds.listUpcomingHostScriptExecutions(ctx, host3.ID, false, false) + require.NoError(t, err) + + // Set host 1 to have a successful script result + _, _, err = ds.SetHostScriptExecutionResult(ctx, &fleet.HostScriptResultPayload{ + HostID: host1.ID, + ExecutionID: host1Upcoming[0].ExecutionID, + Output: "foo", + ExitCode: 0, + }) + require.NoError(t, err) + + // Set host 2 to have a failed script result + _, _, err = ds.SetHostScriptExecutionResult(ctx, &fleet.HostScriptResultPayload{ + HostID: host2.ID, + ExecutionID: host2Upcoming[0].ExecutionID, + Output: "bar", + ExitCode: 1, + }) + require.NoError(t, err) + + // Cancel the execution for host 3 + _, err = ds.CancelHostUpcomingActivity(ctx, host3.ID, host3Upcoming[0].ExecutionID) + require.NoError(t, err) + + // Update the batch activity status to "started" + ExecAdhocSQL(t, ds, func(tx sqlx.ExtContext) error { + _, err := tx.ExecContext(ctx, "UPDATE batch_activities SET status='started' WHERE execution_id IN (?,?)", execID, execID2) + return err + }) + + // Mark activities as completed + err = ds.MarkActivitiesAsCompleted(ctx) + require.NoError(t, err) + + // First activity should be marked as finished and updated accordingly. + batchActivity, err := ds.GetBatchActivity(ctx, execID) + require.NoError(t, err) + require.Equal(t, fleet.BatchExecutionFinished, batchActivity.Status) + require.Equal(t, uint(5), *batchActivity.NumTargeted) + require.Equal(t, uint(1), *batchActivity.NumRan) + require.Equal(t, uint(1), *batchActivity.NumErrored) + require.Equal(t, uint(2), *batchActivity.NumIncompatible) + require.Equal(t, uint(1), *batchActivity.NumCanceled) + require.Equal(t, uint(0), *batchActivity.NumPending) + + // Second activity should still be in "started" status. + batchActivity2, err := ds.GetBatchActivity(ctx, execID2) + require.NoError(t, err) + require.Equal(t, fleet.BatchExecutionStarted, batchActivity2.Status) +} + func testDeleteScriptActivatesNextActivity(t *testing.T, ds *Datastore) { ctx := t.Context() u := test.NewUser(t, ds, "Alice", "alice@example.com", true) diff --git a/server/fleet/cron_schedules.go b/server/fleet/cron_schedules.go index b520ae8788..cc8571776f 100644 --- a/server/fleet/cron_schedules.go +++ b/server/fleet/cron_schedules.go @@ -31,10 +31,11 @@ const ( CronMaintainedApps CronScheduleName = "maintained_apps" // CronRefreshVPPAppVersions updates the versions of VPP apps in Fleet to the latest value. Runs // every 1h. - CronRefreshVPPAppVersions CronScheduleName = "refresh_vpp_app_versions" - CronAppleMDMIPhoneIPadReviver CronScheduleName = "apple_mdm_iphone_ipad_reviver" - CronUpcomingActivitiesMaintenance CronScheduleName = "upcoming_activities_maintenance" - CronHostVitalsLabelMembership CronScheduleName = "host_vitals_label_membership" + CronRefreshVPPAppVersions CronScheduleName = "refresh_vpp_app_versions" + CronAppleMDMIPhoneIPadReviver CronScheduleName = "apple_mdm_iphone_ipad_reviver" + CronUpcomingActivitiesMaintenance CronScheduleName = "upcoming_activities_maintenance" + CronHostVitalsLabelMembership CronScheduleName = "host_vitals_label_membership" + CronBatchActivityCompletionChecker CronScheduleName = "batch_activity_completion_checker" ) type CronSchedulesService interface { diff --git a/server/fleet/datastore.go b/server/fleet/datastore.go index ab61790f2e..2c0029d5d6 100644 --- a/server/fleet/datastore.go +++ b/server/fleet/datastore.go @@ -1833,6 +1833,9 @@ type Datastore interface { // CountBatchScriptExecutions returns the number of batch script executions matching the filter. CountBatchScriptExecutions(ctx context.Context, filter BatchExecutionStatusFilter) (int64, error) + // MarkActivitiesAsCompleted updates the status of the specified activities to "completed". + MarkActivitiesAsCompleted(ctx context.Context) error + // GetHostLockWipeStatus gets the lock/unlock and wipe status for the host. GetHostLockWipeStatus(ctx context.Context, host *Host) (*HostLockWipeStatus, error) diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go index af8335e843..46a703a717 100644 --- a/server/mock/datastore_mock.go +++ b/server/mock/datastore_mock.go @@ -1177,6 +1177,8 @@ type ListBatchScriptExecutionsFunc func(ctx context.Context, filter fleet.BatchE type CountBatchScriptExecutionsFunc func(ctx context.Context, filter fleet.BatchExecutionStatusFilter) (int64, error) +type MarkActivitiesAsCompletedFunc func(ctx context.Context) error + type BatchScheduleScriptFunc func(ctx context.Context, userID *uint, scriptID uint, hostIDs []uint, notBefore time.Time) (string, error) type GetBatchActivityFunc func(ctx context.Context, executionID string) (*fleet.BatchActivity, error) @@ -3186,6 +3188,9 @@ type DataStore struct { CountBatchScriptExecutionsFunc CountBatchScriptExecutionsFunc CountBatchScriptExecutionsFuncInvoked bool + MarkActivitiesAsCompletedFunc MarkActivitiesAsCompletedFunc + MarkActivitiesAsCompletedFuncInvoked bool + GetHostLockWipeStatusFunc GetHostLockWipeStatusFunc GetHostLockWipeStatusFuncInvoked bool @@ -7642,6 +7647,13 @@ func (s *DataStore) CountBatchScriptExecutions(ctx context.Context, filter fleet return s.CountBatchScriptExecutionsFunc(ctx, filter) } +func (s *DataStore) MarkActivitiesAsCompleted(ctx context.Context) error { + s.mu.Lock() + s.MarkActivitiesAsCompletedFuncInvoked = true + s.mu.Unlock() + return s.MarkActivitiesAsCompletedFunc(ctx) +} + func (s *DataStore) GetHostLockWipeStatus(ctx context.Context, host *fleet.Host) (*fleet.HostLockWipeStatus, error) { s.mu.Lock() s.GetHostLockWipeStatusFuncInvoked = true