Evenly distribute bulk-enabled alerting rules (#174656)

resolves: #171980 This is a follow-on issue of: #172742 Above issue randomises `runAt` of the bulk enabled rules. And creates new tasks (by using `scheduleTask` for each one of them) if they don't have any. But, as we create the tasks already enabled `bulkEnable` method of the TM skips them. This PR replaces scheduleTask with bulkSchedule and creates task as disabled, so `bulkEnable` can pick them up. ## To verify: Add Security Solutions' [prebuilt detection rules](http://localhost:5601/app/security/rules/management?sourcerer=(default:(id:security-solution-default,selectedPatterns:!()))&timeline=(activeTab:query,graphEventId:%27%27,isOpen:!f)) and bulk enable them after installing all.
elastic · Jan 17, 2024 · e3fed0c · e3fed0c
1 parent bcfdb56
commit e3fed0c
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 109 deletions.
diff --git a/x-pack/plugins/alerting/server/rules_client/methods/bulk_enable.ts b/x-pack/plugins/alerting/server/rules_client/methods/bulk_enable.ts
@@ -11,7 +11,8 @@ import { SavedObjectsBulkUpdateObject, SavedObjectsFindResult } from '@kbn/core/
 import { withSpan } from '@kbn/apm-utils';
 import { Logger } from '@kbn/core/server';
 import { TaskManagerStartContract, TaskStatus } from '@kbn/task-manager-plugin/server';
-import { RawRule, IntervalSchedule } from '../../types';
+import { TaskInstanceWithDeprecatedFields } from '@kbn/task-manager-plugin/server/task';
+import { RawRule } from '../../types';
 import { convertRuleIdsToKueryNode } from '../../lib';
 import { ruleAuditEvent, RuleAuditAction } from '../common/audit_events';
 import {
@@ -24,7 +25,6 @@ import {
   getAuthorizationFilter,
   checkAuthorizationAndGetTotal,
   getAlertFromRaw,
-  scheduleTask,
   updateMeta,
   createNewAPIKeySet,
   migrateLegacyActions,
@@ -128,6 +128,7 @@ const bulkEnableRulesWithOCC = async (
 
   const rulesFinderRules: Array<SavedObjectsFindResult<RawRule>> = [];
   const rulesToEnable: Array<SavedObjectsBulkUpdateObject<RawRule>> = [];
+  const tasksToSchedule: TaskInstanceWithDeprecatedFields[] = [];
   const errors: BulkOperationError[] = [];
   const ruleNameToRuleIdMapping: Record<string, string> = {};
   const username = await context.getUserName();
@@ -208,31 +209,37 @@ const bulkEnableRulesWithOCC = async (
               error: null,
               warning: null,
             },
+            scheduledTaskId: rule.id,
           });
 
           const shouldScheduleTask = await getShouldScheduleTask(
             context,
             rule.attributes.scheduledTaskId
           );
 
-          let scheduledTaskId;
           if (shouldScheduleTask) {
-            const scheduledTask = await scheduleTask(context, {
+            tasksToSchedule.push({
               id: rule.id,
-              consumer: rule.attributes.consumer,
-              ruleTypeId: rule.attributes.alertTypeId,
-              schedule: rule.attributes.schedule as IntervalSchedule,
-              throwOnConflict: false,
+              taskType: `alerting:${rule.attributes.alertTypeId}`,
+              schedule: rule.attributes.schedule,
+              params: {
+                alertId: rule.id,
+                spaceId: context.spaceId,
+                consumer: rule.attributes.consumer,
+              },
+              state: {
+                previousStartedAt: null,
+                alertTypeState: {},
+                alertInstances: {},
+              },
+              scope: ['alerting'],
+              enabled: false, // we create the task as disabled, taskManager.bulkEnable will enable them by randomising their schedule datetime
             });
-            scheduledTaskId = scheduledTask.id;
           }
 
           rulesToEnable.push({
             ...rule,
-            attributes: {
-              ...updatedAttributes,
-              ...(scheduledTaskId ? { scheduledTaskId } : undefined),
-            },
+            attributes: updatedAttributes,
             ...(migratedActions.hasLegacyActions
               ? { references: migratedActions.resultedReferences }
               : {}),
@@ -264,6 +271,12 @@ const bulkEnableRulesWithOCC = async (
     }
   );
 
+  if (tasksToSchedule.length > 0) {
+    await withSpan({ name: 'taskManager.bulkSchedule', type: 'tasks' }, () =>
+      context.taskManager.bulkSchedule(tasksToSchedule)
+    );
+  }
+
   const result = await withSpan(
     { name: 'unsecuredSavedObjectsClient.bulkCreate', type: 'rules' },
     () =>

diff --git a/x-pack/plugins/alerting/server/rules_client/tests/bulk_enable.test.ts b/x-pack/plugins/alerting/server/rules_client/tests/bulk_enable.test.ts
@@ -182,6 +182,11 @@ describe('bulkEnableRules', () => {
       { overwrite: true }
     );
 
+    expect(taskManager.bulkSchedule).not.toHaveBeenCalled();
+
+    expect(taskManager.bulkEnable).toHaveBeenCalledTimes(1);
+    expect(taskManager.bulkEnable).toHaveBeenCalledWith(['id1', 'id2']);
+
     expect(result).toStrictEqual({
       errors: [],
       rules: [returnedRule1, returnedRule2],
@@ -523,45 +528,34 @@ describe('bulkEnableRules', () => {
     test('should schedule task when scheduledTaskId is defined but task with that ID does not', async () => {
       // One rule gets the task successfully, one rule doesn't so only one task should be scheduled
       taskManager.get.mockRejectedValueOnce(new Error('Failed to get task!'));
-      taskManager.schedule.mockResolvedValueOnce({
-        id: 'id1',
-        taskType: 'alerting:fakeType',
-        scheduledAt: new Date(),
-        attempts: 1,
-        status: TaskStatus.Idle,
-        runAt: new Date(),
-        startedAt: null,
-        retryAt: null,
-        state: {},
-        params: {},
-        ownerId: null,
-      });
       unsecuredSavedObjectsClient.bulkCreate.mockResolvedValue({
         saved_objects: [enabledRule1, enabledRule2],
       });
 
       const result = await rulesClient.bulkEnableRules({ ids: ['id1', 'id2'] });
 
-      expect(taskManager.schedule).toHaveBeenCalledTimes(1);
-      expect(taskManager.schedule).toHaveBeenCalledWith({
-        id: 'id1',
-        taskType: `alerting:fakeType`,
-        params: {
-          alertId: 'id1',
-          spaceId: 'default',
-          consumer: 'fakeConsumer',
-        },
-        schedule: {
-          interval: '5m',
-        },
-        enabled: true,
-        state: {
-          alertInstances: {},
-          alertTypeState: {},
-          previousStartedAt: null,
+      expect(taskManager.bulkSchedule).toHaveBeenCalledTimes(1);
+      expect(taskManager.bulkSchedule).toHaveBeenCalledWith([
+        {
+          id: 'id1',
+          taskType: `alerting:fakeType`,
+          params: {
+            alertId: 'id1',
+            spaceId: 'default',
+            consumer: 'fakeConsumer',
+          },
+          schedule: {
+            interval: '5m',
+          },
+          enabled: false,
+          state: {
+            alertInstances: {},
+            alertTypeState: {},
+            previousStartedAt: null,
+          },
+          scope: ['alerting'],
         },
-        scope: ['alerting'],
-      });
+      ]);
 
       expect(unsecuredSavedObjectsClient.bulkCreate).toHaveBeenCalledTimes(1);
       expect(unsecuredSavedObjectsClient.bulkCreate).toHaveBeenCalledWith(
@@ -607,44 +601,33 @@ describe('bulkEnableRules', () => {
             };
           },
         });
-      taskManager.schedule.mockResolvedValueOnce({
-        id: 'id1',
-        taskType: 'alerting:fakeType',
-        scheduledAt: new Date(),
-        attempts: 1,
-        status: TaskStatus.Idle,
-        runAt: new Date(),
-        startedAt: null,
-        retryAt: null,
-        state: {},
-        params: {},
-        ownerId: null,
-      });
       unsecuredSavedObjectsClient.bulkCreate.mockResolvedValue({
         saved_objects: [enabledRule1, enabledRule2],
       });
       const result = await rulesClient.bulkEnableRules({ ids: ['id1', 'id2'] });
 
-      expect(taskManager.schedule).toHaveBeenCalledTimes(1);
-      expect(taskManager.schedule).toHaveBeenCalledWith({
-        id: 'id1',
-        taskType: `alerting:fakeType`,
-        params: {
-          alertId: 'id1',
-          spaceId: 'default',
-          consumer: 'fakeConsumer',
-        },
-        schedule: {
-          interval: '5m',
-        },
-        enabled: true,
-        state: {
-          alertInstances: {},
-          alertTypeState: {},
-          previousStartedAt: null,
+      expect(taskManager.bulkSchedule).toHaveBeenCalledTimes(1);
+      expect(taskManager.bulkSchedule).toHaveBeenCalledWith([
+        {
+          id: 'id1',
+          taskType: `alerting:fakeType`,
+          params: {
+            alertId: 'id1',
+            spaceId: 'default',
+            consumer: 'fakeConsumer',
+          },
+          schedule: {
+            interval: '5m',
+          },
+          enabled: false,
+          state: {
+            alertInstances: {},
+            alertTypeState: {},
+            previousStartedAt: null,
+          },
+          scope: ['alerting'],
         },
-        scope: ['alerting'],
-      });
+      ]);
 
       expect(unsecuredSavedObjectsClient.bulkCreate).toHaveBeenCalledTimes(1);
       expect(unsecuredSavedObjectsClient.bulkCreate).toHaveBeenCalledWith(
@@ -690,19 +673,6 @@ describe('bulkEnableRules', () => {
         ownerId: null,
         enabled: false,
       });
-      taskManager.schedule.mockResolvedValueOnce({
-        id: 'id1',
-        taskType: 'alerting:fakeType',
-        scheduledAt: new Date(),
-        attempts: 1,
-        status: TaskStatus.Idle,
-        runAt: new Date(),
-        startedAt: null,
-        retryAt: null,
-        state: {},
-        params: {},
-        ownerId: null,
-      });
       unsecuredSavedObjectsClient.bulkCreate.mockResolvedValue({
         saved_objects: [enabledRule1, enabledRule2],
       });
@@ -711,26 +681,28 @@ describe('bulkEnableRules', () => {
 
       expect(taskManager.removeIfExists).toHaveBeenCalledTimes(1);
       expect(taskManager.removeIfExists).toHaveBeenCalledWith('id1');
-      expect(taskManager.schedule).toHaveBeenCalledTimes(1);
-      expect(taskManager.schedule).toHaveBeenCalledWith({
-        id: 'id1',
-        taskType: `alerting:fakeType`,
-        params: {
-          alertId: 'id1',
-          spaceId: 'default',
-          consumer: 'fakeConsumer',
-        },
-        schedule: {
-          interval: '5m',
-        },
-        enabled: true,
-        state: {
-          alertInstances: {},
-          alertTypeState: {},
-          previousStartedAt: null,
+      expect(taskManager.bulkSchedule).toHaveBeenCalledTimes(1);
+      expect(taskManager.bulkSchedule).toHaveBeenCalledWith([
+        {
+          id: 'id1',
+          taskType: `alerting:fakeType`,
+          params: {
+            alertId: 'id1',
+            spaceId: 'default',
+            consumer: 'fakeConsumer',
+          },
+          schedule: {
+            interval: '5m',
+          },
+          enabled: false,
+          state: {
+            alertInstances: {},
+            alertTypeState: {},
+            previousStartedAt: null,
+          },
+          scope: ['alerting'],
         },
-        scope: ['alerting'],
-      });
+      ]);
 
       expect(unsecuredSavedObjectsClient.bulkCreate).toHaveBeenCalledTimes(1);
       expect(unsecuredSavedObjectsClient.bulkCreate).toHaveBeenCalledWith(