[Discover] Cancel long running requests in Discover alert

x-pack/plugins/alerting/common/rule_task_instance.ts

@@ -18,6 +18,7 @@ export const ruleStateSchema = t.partial({
 
 const ruleExecutionMetricsSchema = t.partial({
   numSearches: t.number,
+  numSearchSourceSearches: t.number,
   totalSearchDurationMs: t.number,
   esSearchDurationMs: t.number,
 });

x-pack/plugins/alerting/server/lib/merge_search_metrics.ts

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { RuleExecutionMetrics } from '../types';
+import type { ScopedClusterClientMetrics } from './wrap_scoped_cluster_client';
+import type { SearchSourceClientMetrics } from './wrap_search_source_fetch';
+
+export const mergeSearchMetrics = (
+  scopedClusterClientMetrics: ScopedClusterClientMetrics,
+  searchSourceMetrics: SearchSourceClientMetrics
+): RuleExecutionMetrics => {
+  return {
+    numSearches: scopedClusterClientMetrics.numSearches,
+    numSearchSourceSearches: searchSourceMetrics.numSearchSourceSearches,
+    totalSearchDurationMs:
+      scopedClusterClientMetrics.totalSearchDurationMs + searchSourceMetrics.totalSearchDurationMs,
+    esSearchDurationMs:
+      scopedClusterClientMetrics.esSearchDurationMs + searchSourceMetrics.esSearchDurationMs,
+  };
+};

x-pack/plugins/alerting/server/lib/types.ts

@@ -7,6 +7,8 @@
 
 import * as t from 'io-ts';
 import { either } from 'fp-ts/lib/Either';
+import { Rule } from '../types';
+
 // represents a Date from an ISO string
 export const DateFromString = new t.Type<Date, string, unknown>(
   'DateFromString',
@@ -24,3 +26,10 @@ export const DateFromString = new t.Type<Date, string, unknown>(
     ),
   (valueToEncode) => valueToEncode.toISOString()
 );
+
+export type RuleInfo = Pick<Rule, 'name' | 'alertTypeId' | 'id'> & { spaceId: string };
+
+export interface LogSearchMetricsOpts {
+  esSearchDuration: number;
+  totalSearchDuration: number;
+}

x-pack/plugins/alerting/server/lib/wrap_scoped_cluster_client.ts

@@ -21,10 +21,14 @@ import type {
   AggregationsAggregate,
 } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { IScopedClusterClient, ElasticsearchClient, Logger } from '@kbn/core/server';
-import { RuleExecutionMetrics } from '../types';
-import { Rule } from '../types';
+import { LogSearchMetricsOpts, RuleInfo } from './types';
+
+export type LogSearchMetricsFn = (metrics: LogSearchMetricsOpts) => void;
+
+type WrapEsClientOpts = Omit<WrapScopedClusterClientOpts, 'scopedClusterClient'> & {
+  esClient: ElasticsearchClient;
+};
 
-type RuleInfo = Pick<Rule, 'name' | 'alertTypeId' | 'id'> & { spaceId: string };
 interface WrapScopedClusterClientFactoryOpts {
   scopedClusterClient: IScopedClusterClient;
   rule: RuleInfo;
@@ -36,15 +40,11 @@ type WrapScopedClusterClientOpts = WrapScopedClusterClientFactoryOpts & {
   logMetricsFn: LogSearchMetricsFn;
 };
 
-type WrapEsClientOpts = Omit<WrapScopedClusterClientOpts, 'scopedClusterClient'> & {
-  esClient: ElasticsearchClient;
-};
-
-interface LogSearchMetricsOpts {
-  esSearchDuration: number;
-  totalSearchDuration: number;
+export interface ScopedClusterClientMetrics {
+  esSearchDurationMs: number;
+  totalSearchDurationMs: number;
+  numSearches: number;
 }
-type LogSearchMetricsFn = (metrics: LogSearchMetricsOpts) => void;
 
 export function createWrappedScopedClusterClientFactory(opts: WrapScopedClusterClientFactoryOpts) {
   let numSearches: number = 0;
@@ -61,7 +61,7 @@ export function createWrappedScopedClusterClientFactory(opts: WrapScopedClusterC
 
   return {
     client: () => wrappedClient,
-    getMetrics: (): RuleExecutionMetrics => {
+    getMetrics: (): ScopedClusterClientMetrics => {
       return {
         esSearchDurationMs,
         totalSearchDurationMs,

x-pack/plugins/alerting/server/lib/wrap_search_source_fetch.test.ts

@@ -0,0 +1,108 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { loggingSystemMock } from '@kbn/core/server/mocks';
+import { searchSourceInstanceMock } from '@kbn/data-plugin/common/search/search_source/mocks';
+import { wrapSearchSourceFetch } from './wrap_search_source_fetch';
+
+const logger = loggingSystemMock.create().get();
+
+const rule = {
+  name: 'test-rule',
+  alertTypeId: '.test-rule-type',
+  id: 'abcdefg',
+  spaceId: 'my-space',
+};
+
+describe('wrapSearchSourceFetch', () => {
+  beforeAll(() => {
+    jest.useFakeTimers();
+  });
+
+  afterAll(() => {
+    jest.useRealTimers();
+  });
+
+  afterEach(() => {
+    jest.resetAllMocks();
+  });
+
+  test('searches properly', async () => {
+    const abortController = new AbortController();
+
+    const { fetch: wrappedFetch } = wrapSearchSourceFetch({ logger, rule, abortController });
+    wrappedFetch(searchSourceInstanceMock);
+
+    expect(searchSourceInstanceMock.fetch).toHaveBeenCalledWith({
+      abortSignal: abortController.signal,
+    });
+  });
+
+  test('re-throws error when search throws error', async () => {
+    const abortController = new AbortController();
+
+    (searchSourceInstanceMock.fetch as jest.Mock).mockRejectedValueOnce(
+      new Error('something went wrong!')
+    );
+    const { fetch: wrappedFetch } = wrapSearchSourceFetch({ logger, rule, abortController });
+
+    await expect(
+      wrappedFetch.bind({}, searchSourceInstanceMock)
+    ).rejects.toThrowErrorMatchingInlineSnapshot(`"something went wrong!"`);
+  });
+
+  test('throws error when search throws abort error', async () => {
+    const abortController = new AbortController();
+    abortController.abort();
+
+    (searchSourceInstanceMock.fetch as jest.Mock).mockRejectedValueOnce(
+      new Error('Request has been aborted by the user')
+    );
+
+    const { fetch: wrappedFetch } = wrapSearchSourceFetch({ logger, rule, abortController });
+
+    await expect(
+      wrappedFetch.bind({}, searchSourceInstanceMock)
+    ).rejects.toThrowErrorMatchingInlineSnapshot(
+      `"Search has been aborted due to cancelled execution"`
+    );
+  });
+
+  test('keeps track of number of queries', async () => {
+    const abortController = new AbortController();
+    abortController.abort();
+
+    (searchSourceInstanceMock.fetch as jest.Mock).mockResolvedValue({ took: 333 });
+
+    const { fetch: wrappedFetch, getMetrics } = wrapSearchSourceFetch({
+      logger,
+      rule,
+      abortController,
+    });
+
+    await wrappedFetch(searchSourceInstanceMock);
+    await wrappedFetch(searchSourceInstanceMock);
+    await wrappedFetch(searchSourceInstanceMock);
+
+    expect(searchSourceInstanceMock.fetch).toHaveBeenCalledTimes(3);
+
+    const stats = getMetrics();
+    expect(stats.numSearchSourceSearches).toEqual(3);
+    expect(stats.esSearchDurationMs).toEqual(999);
+
+    expect(logger.debug).toHaveBeenCalledWith(
+      'executing query for rule .test-rule-type:abcdefg in space my-space'
+    );
+  });
+});

x-pack/plugins/alerting/server/lib/wrap_search_source_fetch.ts

@@ -0,0 +1,62 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { Logger } from '@kbn/core/server';
+import { ISearchSource } from '@kbn/data-plugin/common';
+import { LogSearchMetricsOpts, RuleInfo } from './types';
+
+export interface SearchSourceClientMetrics {
+  esSearchDurationMs: number;
+  totalSearchDurationMs: number;
+  numSearchSourceSearches: number;
+}
+
+export function wrapSearchSourceFetch({
+  logger,
+  rule,
+  abortController,
+}: {
+  logger: Logger;
+  rule: RuleInfo;
+  abortController: AbortController;
+}) {
+  let numSearchSourceSearches: number = 0;
+  let esSearchDurationMs: number = 0;
+  let totalSearchDurationMs: number = 0;
+
+  function logMetrics(metrics: LogSearchMetricsOpts) {
+    numSearchSourceSearches++;
+    esSearchDurationMs += metrics.esSearchDuration;
+    totalSearchDurationMs += metrics.totalSearchDuration;
+  }
+
+  return {
+    fetch: async (searchSource: ISearchSource) => {
+      try {
+        const start = Date.now();
+        logger.debug(
+          `executing query for rule ${rule.alertTypeId}:${rule.id} in space ${rule.spaceId}`
+        );
+        const result = await searchSource.fetch({ abortSignal: abortController.signal });
+
+        const durationMs = Date.now() - start;
+        logMetrics({ esSearchDuration: result.took ?? 0, totalSearchDuration: durationMs });
+        return result;
+      } catch (e) {
+        if (abortController.signal.aborted) {
+          throw new Error('Search has been aborted due to cancelled execution');
+        }
+        throw e;
+      }
+    },
+    getMetrics: (): SearchSourceClientMetrics => ({
+      esSearchDurationMs,
+      totalSearchDurationMs,
+      numSearchSourceSearches,
+    }),
+  };
+}

x-pack/plugins/alerting/server/mocks.ts

@@ -9,9 +9,11 @@ import {
   elasticsearchServiceMock,
   savedObjectsClientMock,
   uiSettingsServiceMock,
-  httpServerMock,
 } from '@kbn/core/server/mocks';
-import { dataPluginMock } from '@kbn/data-plugin/server/mocks';
+import {
+  searchSourceCommonMock,
+  searchSourceInstanceMock,
+} from '@kbn/data-plugin/common/search/search_source/mocks';
 import { rulesClientMock } from './rules_client.mock';
 import { PluginSetupContract, PluginStartContract } from './plugin';
 import { Alert, AlertFactoryDoneUtils } from './alert';
@@ -113,11 +115,10 @@ const createRuleExecutorServicesMock = <
     shouldWriteAlerts: () => true,
     shouldStopExecution: () => true,
     search: createAbortableSearchServiceMock(),
-    searchSourceClient: Promise.resolve(
-      dataPluginMock
-        .createStartContract()
-        .search.searchSource.asScoped(httpServerMock.createKibanaRequest())
-    ),
+    searchSourceUtils: {
+      searchSourceClient: searchSourceCommonMock,
+      wrappedFetch: () => searchSourceInstanceMock.fetch(),
+    },
   };
 };
 export type RuleExecutorServicesMock = ReturnType<typeof createRuleExecutorServicesMock>;

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

@@ -240,7 +240,7 @@ describe('Task Runner', () => {
     expect(logger.debug).nthCalledWith(1, 'executing rule test:1 at 1970-01-01T00:00:00.000Z');
     expect(logger.debug).nthCalledWith(
       2,
-      'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"esSearchDurationMs":33,"totalSearchDurationMs":23423},"numberOfTriggeredActions":0,"numberOfGeneratedActions":0,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"ok"}'
+      'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"numSearchSourceSearches":0,"totalSearchDurationMs":23423,"esSearchDurationMs":33},"numberOfTriggeredActions":0,"numberOfGeneratedActions":0,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"ok"}'
     );
 
     const eventLogger = taskRunnerFactoryInitializerParams.eventLogger;
@@ -320,7 +320,7 @@ describe('Task Runner', () => {
       );
       expect(logger.debug).nthCalledWith(
         3,
-        'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"esSearchDurationMs":33,"totalSearchDurationMs":23423},"numberOfTriggeredActions":1,"numberOfGeneratedActions":1,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
+        'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"numSearchSourceSearches":0,"totalSearchDurationMs":23423,"esSearchDurationMs":33},"numberOfTriggeredActions":1,"numberOfGeneratedActions":1,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
       );
 
       const eventLogger = customTaskRunnerFactoryInitializerParams.eventLogger;
@@ -428,7 +428,7 @@ describe('Task Runner', () => {
     );
     expect(logger.debug).nthCalledWith(
       4,
-      'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"esSearchDurationMs":33,"totalSearchDurationMs":23423},"numberOfTriggeredActions":0,"numberOfGeneratedActions":0,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
+      'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"numSearchSourceSearches":0,"totalSearchDurationMs":23423,"esSearchDurationMs":33},"numberOfTriggeredActions":0,"numberOfGeneratedActions":0,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
     );
 
     const eventLogger = taskRunnerFactoryInitializerParams.eventLogger;
@@ -593,7 +593,7 @@ describe('Task Runner', () => {
       );
       expect(logger.debug).nthCalledWith(
         4,
-        'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"esSearchDurationMs":33,"totalSearchDurationMs":23423},"numberOfTriggeredActions":1,"numberOfGeneratedActions":1,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
+        'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"numSearchSourceSearches":0,"totalSearchDurationMs":23423,"esSearchDurationMs":33},"numberOfTriggeredActions":1,"numberOfGeneratedActions":1,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
       );
       expect(mockUsageCounter.incrementCounter).not.toHaveBeenCalled();
     }
@@ -1121,7 +1121,7 @@ describe('Task Runner', () => {
       );
       expect(logger.debug).nthCalledWith(
         4,
-        'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"esSearchDurationMs":33,"totalSearchDurationMs":23423},"numberOfTriggeredActions":2,"numberOfGeneratedActions":2,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
+        'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"numSearchSourceSearches":0,"totalSearchDurationMs":23423,"esSearchDurationMs":33},"numberOfTriggeredActions":2,"numberOfGeneratedActions":2,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
       );
 
       const eventLogger = customTaskRunnerFactoryInitializerParams.eventLogger;
@@ -1262,7 +1262,7 @@ describe('Task Runner', () => {
       );
       expect(logger.debug).nthCalledWith(
         4,
-        `ruleExecutionStatus for test:${alertId}: {"metrics":{"numSearches":3,"esSearchDurationMs":33,"totalSearchDurationMs":23423},"numberOfTriggeredActions":2,"numberOfGeneratedActions":2,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}`
+        'ruleExecutionStatus for test:e558aaad-fd81-46d2-96fc-3bd8fc3dc03f: {"metrics":{"numSearches":3,"numSearchSourceSearches":0,"totalSearchDurationMs":23423,"esSearchDurationMs":33},"numberOfTriggeredActions":2,"numberOfGeneratedActions":2,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"active"}'
       );
 
       const eventLogger = customTaskRunnerFactoryInitializerParams.eventLogger;
@@ -2489,7 +2489,7 @@ describe('Task Runner', () => {
     expect(logger.debug).nthCalledWith(1, 'executing rule test:1 at 1970-01-01T00:00:00.000Z');
     expect(logger.debug).nthCalledWith(
       2,
-      'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"esSearchDurationMs":33,"totalSearchDurationMs":23423},"numberOfTriggeredActions":0,"numberOfGeneratedActions":0,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"ok"}'
+      'ruleExecutionStatus for test:1: {"metrics":{"numSearches":3,"numSearchSourceSearches":0,"totalSearchDurationMs":23423,"esSearchDurationMs":33},"numberOfTriggeredActions":0,"numberOfGeneratedActions":0,"lastExecutionDate":"1970-01-01T00:00:00.000Z","status":"ok"}'
     );
 
     const eventLogger = taskRunnerFactoryInitializerParams.eventLogger;