Add cron schedule to mark activities as finished (#31737)

for #31555 

# Details

This PR adds a new cron schedule "batch_activity_completion_checker"
that runs every 5 minutes and checks whether any batch activities marked
as "started" have completed their runs. In general this is done by
determining whether the sum of the "ran", "incompatible", "errored" and
"canceled" hosts equals the number of "targeted" hosts for the activity.
How that is computed will vary by batch activity type (currently we just
have batch scripts).

When an activity is marked as finished, we cache the final tally of host
statuses (ran, incompatible, errored, canceled) on the record. This is
important so that future queries on activity records don't have to do
the expensive query to compute the host counts on activities where those
counts will never change.

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

- [X] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)

## Testing

- [X] Added/updated automated tests
- [X] Where appropriate, [automated tests simulate multiple hosts and
test for host
isolation](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/reference/patterns-backend.md#unit-testing)
(updates to one hosts's records do not affect another)

- [X] QA'd all new/changed functionality manually
Started a new batch script run using the update run modal (see
https://github.com/fleetdm/fleet/pull/31604) and then triggered the new
job using `fleetctl trigger --name batch_activity_completion_checker`,
and verified that the `batch_activities` record status was `finished`
and the expected fields were populated.



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Introduced an automated process that regularly marks completed batch
activities, ensuring more accurate and up-to-date activity statuses.
* **Bug Fixes**
* Improved reliability in updating the status of batch activities when
all targeted hosts have finished their tasks.
* **Tests**
* Added comprehensive tests to verify correct marking of completed batch
activities.
* **Chores**
* Enhanced internal scheduling and datastore interfaces to support the
new completion-checking process.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

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 {

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.

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))
+}

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)

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 {

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)
 

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