[ResponseOps] Allow tasks to indicate they should disable themselves after run (elastic#230998)

resolves elastic#228306
resolves elastic#227894

Adds the capability of a task to disable itself after running, by
returning an indication in the task response. We then use this with
alerting to allow a rule's task to disable itself if it ran but found
the rule itself was disabled.

Specifically:
- A rule throwing an error with a new reason
`RuleExecutionStatusErrorReasons.Disabled` will cause the alerting task
runner to return to task manager with a new flag `shouldDisableTask:
true`.
- During the post processing of the task, if the task indicated it
should disable itself via the new flag, it will log a warning that it's
going to do that, and then update the task with `enabled: false`.

Task store has been updated as well, to use `overwrite: true` when doing
`bulkSchedule`, as this code currently assumes that - at this point -
the task doc does not exist, so the bulk update fails with a conflict if
one does.

---------

Co-authored-by: kibanamachine <[email protected]>
Co-authored-by: Elastic Machine <[email protected]>

Author: 3 people

x-pack/platform/plugins/shared/alerting/server/task_runner/task_runner.test.ts

@@ -7,6 +7,7 @@
 
 import sinon from 'sinon';
 import { usageCountersServiceMock } from '@kbn/usage-collection-plugin/server/usage_counters/usage_counters_service.mock';
+import type { SavedObject } from '@kbn/core/server';
 import type {
   RuleExecutorOptions,
   RuleTypeParams,
@@ -21,6 +22,7 @@ import {
   RuleExecutionStatusWarningReasons,
   DEFAULT_FLAPPING_SETTINGS,
   DEFAULT_QUERY_DELAY_SETTINGS,
+  type RawRule,
 } from '../types';
 import type { ConcreteTaskInstance } from '@kbn/task-manager-plugin/server';
 import {
@@ -3425,6 +3427,37 @@ describe('Task Runner', () => {
     });
   });
 
+  test('should return shouldDisableTask when task is enabled but rule is not', async () => {
+    const taskRunner = new TaskRunner({
+      ruleType,
+      internalSavedObjectsRepository,
+      taskInstance: { ...mockedTaskInstance },
+      context: taskRunnerFactoryInitializerParams,
+      inMemoryMetrics,
+    });
+    expect(AlertingEventLogger).toHaveBeenCalled();
+    const mockedRuleTypeSavedObjectDisabled = {
+      ...(mockedRuleTypeSavedObject as Rule),
+      enabled: false,
+    };
+    const mockedRawRuleSODisabled: SavedObject<RawRule> = {
+      ...mockedRawRuleSO,
+    };
+    mockedRawRuleSODisabled.attributes.enabled = false;
+
+    mockGetRuleFromRaw.mockReturnValue(mockedRuleTypeSavedObjectDisabled);
+    encryptedSavedObjectsClient.getDecryptedAsInternalUser.mockResolvedValue(mockedRawRuleSO);
+
+    const result = await taskRunner.run();
+
+    expect(ruleType.executor).not.toHaveBeenCalled();
+
+    expect(result.shouldDisableTask).toEqual(true);
+    expect(result.taskRunError?.message).toBe(
+      'Rule failed to execute because rule ran after it was disabled.'
+    );
+  });
+
   function testAlertingEventLogCalls({
     ruleContext = alertingEventLoggerInitializer,
     ruleTypeDef = ruleType,

x-pack/platform/plugins/shared/alerting/server/task_runner/task_runner.ts

@@ -687,6 +687,7 @@ export class TaskRunner<
 
     let runRuleResult: Result<RunRuleResult, Error>;
     let schedule: Result<IntervalSchedule, Error>;
+    let shouldDisableTask = false;
     try {
       const validatedRuleData = await this.prepareToRun();
 
@@ -705,6 +706,7 @@ export class TaskRunner<
 
       runRuleResult = asErr(err);
       schedule = asErr(err);
+      shouldDisableTask = err.reason === RuleExecutionStatusErrorReasons.Disabled;
     }
 
     await withAlertingSpan('alerting:process-run-results-and-update-rule', () =>
@@ -736,6 +738,8 @@ export class TaskRunner<
         ruleTypeId: this.ruleType.id,
         ruleId,
       }),
+      // added this way so we don't add shouldDisableTask: false explicitly
+      ...(shouldDisableTask ? { shouldDisableTask } : {}),
     };
   }
 

x-pack/platform/plugins/shared/alerting/server/task_runner/types.ts

@@ -61,6 +61,7 @@ export interface RuleTaskRunResult {
   schedule: IntervalSchedule | undefined;
   taskRunError?: DecoratedError;
   shouldDeleteTask?: boolean;
+  shouldDisableTask?: boolean;
 }
 
 export const getDeleteRuleTaskRunResult = (): RuleTaskRunResult => ({

x-pack/platform/plugins/shared/task_manager/server/task.ts

@@ -78,6 +78,7 @@ export type SuccessfulRunResult = {
   taskRunError?: DecoratedError;
   shouldValidate?: boolean;
   shouldDeleteTask?: boolean;
+  shouldDisableTask?: boolean;
 } & (
   | // ensure a SuccessfulRunResult can either specify a new `runAt` or a new `schedule`, but not both
   {

x-pack/platform/plugins/shared/task_manager/server/task_running/task_runner.test.ts

@@ -2722,6 +2722,43 @@ describe('TaskManagerRunner', () => {
         tags: ['task:end', 'foo', 'bar'],
       });
     });
+
+    test('Disables task if shouldDisableTask: true is returned', async () => {
+      const { runner, store, logger } = await readyToRunStageSetup({
+        instance: {
+          schedule: {
+            interval: `1m`,
+          },
+        },
+        definitions: {
+          bar: {
+            title: 'Bar!',
+            createTaskRunner: () => ({
+              async run() {
+                return {
+                  state: {},
+                  shouldDisableTask: true,
+                };
+              },
+            }),
+          },
+        },
+      });
+      await runner.run();
+      expect(store.partialUpdate).toHaveBeenCalledWith(
+        expect.objectContaining({
+          enabled: false,
+          id: 'foo',
+          status: 'idle',
+        }),
+        expect.anything()
+      );
+      expect(logger.warn).toBeCalledTimes(1);
+      expect(logger.warn).toBeCalledWith(
+        'Disabling task bar:foo as it indicated it should disable itself',
+        { tags: ['bar'] }
+      );
+    });
   });
 
   describe('isAdHocTaskAndOutOfAttempts', () => {

x-pack/platform/plugins/shared/task_manager/server/task_running/task_runner.ts

@@ -664,6 +664,7 @@ export class TaskManagerRunner implements TaskRunner {
     result: Result<SuccessfulRunResult, FailedRunResult>
   ): Promise<TaskRunResult> {
     const hasTaskRunFailed = isOk(result);
+    let shouldTaskBeDisabled = false;
     const fieldUpdates: Partial<ConcreteTaskInstance> & Pick<ConcreteTaskInstance, 'status'> = flow(
       // if running the task has failed ,try to correct by scheduling a retry in the near future
       mapErr(this.rescheduleFailedRun),
@@ -675,12 +676,18 @@ export class TaskManagerRunner implements TaskRunner {
           state,
           attempts = 0,
           shouldDeleteTask,
+          shouldDisableTask,
         }: SuccessfulRunResult & { attempts: number }) => {
           if (shouldDeleteTask) {
             // set the status to failed so task will get deleted
             return asOk({ status: TaskStatus.ShouldDelete });
           }
 
+          if (shouldDisableTask) {
+            shouldTaskBeDisabled = true;
+            return asOk({ status: TaskStatus.Idle });
+          }
+
           const updatedTaskSchedule = reschedule ?? this.instance.task.schedule;
           return asOk({
             runAt:
@@ -744,13 +751,22 @@ export class TaskManagerRunner implements TaskRunner {
         }
       } else {
         shouldUpdateTask = true;
+
+        if (shouldTaskBeDisabled) {
+          const label = `${this.taskType}:${this.instance.task.id}`;
+          this.logger.warn(`Disabling task ${label} as it indicated it should disable itself`, {
+            tags: [this.taskType],
+          });
+        }
+
         partialTask = {
           ...partialTask,
           ...fieldUpdates,
           // reset fields that track the lifecycle of the concluded `task run`
           startedAt: null,
           retryAt: null,
           ownerId: null,
+          ...(shouldTaskBeDisabled ? { enabled: false } : {}),
         };
       }
 

x-pack/platform/plugins/shared/task_manager/server/task_store.test.ts

@@ -2464,7 +2464,10 @@ describe('TaskStore', () => {
             },
           },
         ],
-        { refresh: false }
+        {
+          overwrite: true,
+          refresh: false,
+        }
       );
 
       expect(result).toEqual([
@@ -2913,7 +2916,10 @@ describe('TaskStore', () => {
             },
           },
         ],
-        { refresh: false }
+        {
+          overwrite: true,
+          refresh: false,
+        }
       );
     });
   });

x-pack/platform/plugins/shared/task_manager/server/task_store.ts

@@ -417,6 +417,7 @@ export class TaskStore {
     try {
       savedObjects = await soClient.bulkCreate<SerializedConcreteTaskInstance>(objects, {
         refresh: false,
+        overwrite: true,
       });
       this.adHocTaskCounter.increment(
         taskInstances.filter((task) => {

x-pack/platform/test/alerting_api_integration/common/lib/task_manager_utils.ts

@@ -162,4 +162,16 @@ export class TaskManagerUtils {
       }
     });
   }
+
+  async setTaskEnabled(id: string, enabled: boolean) {
+    await this.es.update({
+      id: `task:${id}`,
+      index: '.kibana_task_manager',
+      doc: {
+        task: {
+          enabled,
+        },
+      },
+    });
+  }
 }

x-pack/platform/test/alerting_api_integration/spaces_only/tests/alerting/group1/disable.ts

@@ -7,6 +7,7 @@
 
 import expect from '@kbn/expect';
 import { RULE_SAVED_OBJECT_TYPE } from '@kbn/alerting-plugin/server';
+import type { RawRule } from '@kbn/alerting-plugin/server/types';
 import { ES_TEST_INDEX_NAME } from '@kbn/alerting-api-integration-helpers';
 import { ALERT_STATUS } from '@kbn/rule-data-utils';
 import { Spaces } from '../../../scenarios';
@@ -19,6 +20,7 @@ import {
   getTestRuleData,
   ObjectRemover,
   getEventLog,
+  TaskManagerUtils,
 } from '../../../../common/lib';
 import { validateEvent } from './event_log';
 
@@ -29,6 +31,7 @@ export default function createDisableRuleTests({ getService }: FtrProviderContex
   const supertestWithoutAuth = getService('supertestWithoutAuth');
   const retry = getService('retry');
   const supertest = getService('supertest');
+  const taskManagerUtils = new TaskManagerUtils(es, retry);
 
   describe('disable', function () {
     this.tags('skipFIPS');
@@ -55,6 +58,14 @@ export default function createDisableRuleTests({ getService }: FtrProviderContex
       return scheduledTask._source!;
     }
 
+    async function getRule(id: string): Promise<RawRule> {
+      const ruleDoc = await es.get<{ alert: RawRule }>({
+        id: `alert:${id}`,
+        index: '.kibana_alerting_cases',
+      });
+      return ruleDoc._source!.alert;
+    }
+
     it('should handle disable rule request appropriately', async () => {
       const { body: createdRule } = await supertestWithoutAuth
         .post(`${getUrlPrefix(Spaces.space1.id)}/api/alerting/rule`)
@@ -265,5 +276,46 @@ export default function createDisableRuleTests({ getService }: FtrProviderContex
         id: createdRule.id,
       });
     });
+
+    it('should disable task if task is run but rule is disabled', async () => {
+      const { body: createdRule } = await supertestWithoutAuth
+        .post(`${getUrlPrefix(Spaces.space1.id)}/api/alerting/rule`)
+        .set('kbn-xsrf', 'foo')
+        .send(
+          getTestRuleData({
+            enabled: true,
+            schedule: {
+              interval: '3s',
+            },
+          })
+        )
+        .expect(200);
+      objectRemover.add(Spaces.space1.id, createdRule.id, 'rule', 'alerting');
+
+      // wait for rule to run once
+      await retry.try(async () => {
+        const rule = await getRule(createdRule.id);
+        expect(rule?.monitoring?.run?.last_run?.timestamp).not.to.be(undefined);
+      });
+
+      // disable rule, and implicitly task
+      await ruleUtils.disable(createdRule.id);
+
+      // wait for the task to be disabled
+      await waitForDisabledTask(createdRule.scheduled_task_id);
+
+      // manually enable task
+      await taskManagerUtils.setTaskEnabled(createdRule.scheduled_task_id, true);
+
+      // wait for the task to be disabled
+      await waitForDisabledTask(createdRule.scheduled_task_id);
+    });
+
+    async function waitForDisabledTask(taskId: string) {
+      await retry.try(async () => {
+        const taskDoc = await getScheduledTask(taskId);
+        expect(taskDoc.task.enabled).to.be(false);
+      });
+    }
   });
 }