[Discover] Cancel long running requests in Discover alert

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

@@ -21,6 +21,7 @@ import type {
   AggregationsAggregate,
 } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { IScopedClusterClient, ElasticsearchClient, Logger } from 'src/core/server';
+import { ISearchSource } from 'src/plugins/data/common';
 import { RuleExecutionMetrics } from '../types';
 import { Rule } from '../types';
 
@@ -46,7 +47,7 @@ interface LogSearchMetricsOpts {
 }
 type LogSearchMetricsFn = (metrics: LogSearchMetricsOpts) => void;
 
-export function createWrappedScopedClusterClientFactory(opts: WrapScopedClusterClientFactoryOpts) {
+export function createWrappedSearchClientsFactory(opts: WrapScopedClusterClientFactoryOpts) {
   let numSearches: number = 0;
   let esSearchDurationMs: number = 0;
   let totalSearchDurationMs: number = 0;
@@ -57,10 +58,12 @@ export function createWrappedScopedClusterClientFactory(opts: WrapScopedClusterC
     totalSearchDurationMs += metrics.totalSearchDuration;
   }
 
-  const wrappedClient = wrapScopedClusterClient({ ...opts, logMetricsFn: logMetrics });
+  const wrappedScopedClusterClient = wrapScopedClusterClient({ ...opts, logMetricsFn: logMetrics });
+  const wrappedSearchSourceFetch = wrapSearchSourceFetch({ ...opts, logMetricsFn: logMetrics });
 
   return {
-    client: () => wrappedClient,
+    scopedClusterClient: () => wrappedScopedClusterClient,
+    searchSourceFetch: () => wrappedSearchSourceFetch,
     getMetrics: (): RuleExecutionMetrics => {
       return {
         esSearchDurationMs,
@@ -71,6 +74,27 @@ export function createWrappedScopedClusterClientFactory(opts: WrapScopedClusterC
   };
 }
 
+export function wrapSearchSourceFetch(opts: WrapScopedClusterClientOpts) {
+  return async (searchSource: ISearchSource) => {
+    try {
+      const start = Date.now();
+      opts.logger.debug(
+        `executing query for rule ${opts.rule.alertTypeId}:${opts.rule.id} in space ${opts.rule.spaceId}`
+      );
+      const result = await searchSource.fetch({ abortSignal: opts.abortController.signal });
+
+      const durationMs = Date.now() - start;
+      opts.logMetricsFn({ esSearchDuration: result.took ?? 0, totalSearchDuration: durationMs });
+      return result;
+    } catch (e) {
+      if (opts.abortController.signal.aborted) {
+        throw new Error('Search has been aborted due to cancelled execution');
+      }
+      throw e;
+    }
+  };
+}
+
 function wrapScopedClusterClient(opts: WrapScopedClusterClientOpts): IScopedClusterClient {
   const { scopedClusterClient, ...rest } = opts;
   return {

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

@@ -16,7 +16,6 @@ import { ConcreteTaskInstance, throwUnrecoverableError } from '../../../task_man
 import { createExecutionHandler, ExecutionHandler } from './create_execution_handler';
 import { Alert, createAlertFactory } from '../alert';
 import {
-  createWrappedScopedClusterClientFactory,
   ElasticsearchError,
   ErrorWithReason,
   executionStatusFromError,
@@ -74,6 +73,7 @@ import {
   ScheduleActionsForRecoveredAlertsParams,
   TrackAlertDurationsParams,
 } from './types';
+import { createWrappedSearchClientsFactory } from '../lib/wrap_scoped_cluster_client';
 
 const FALLBACK_RETRY_INTERVAL = '5m';
 const CONNECTIVITY_RETRY_INTERVAL = '5m';
@@ -356,7 +356,7 @@ export class TaskRunner<
     const ruleLabel = `${this.ruleType.id}:${ruleId}: '${name}'`;
 
     const scopedClusterClient = this.context.elasticsearch.client.asScoped(fakeRequest);
-    const wrappedScopedClusterClient = createWrappedScopedClusterClientFactory({
+    const wrappedScopedClusterClient = createWrappedSearchClientsFactory({
       scopedClusterClient,
       rule: {
         name: rule.name,
@@ -367,6 +367,11 @@ export class TaskRunner<
       logger: this.logger,
       abortController: this.searchAbortController,
     });
+    const searchSourceClient = await this.context.data.search.searchSource.asScoped(fakeRequest);
+    const searchSourceUtils = {
+      searchSourceClient,
+      wrappedFetch: wrappedScopedClusterClient.searchSourceFetch(),
+    };
 
     let updatedRuleTypeState: void | Record<string, unknown>;
     try {
@@ -389,9 +394,9 @@ export class TaskRunner<
           executionId: this.executionId,
           services: {
             savedObjectsClient,
+            searchSourceUtils,
             uiSettingsClient: this.context.uiSettings.asScopedToClient(savedObjectsClient),
-            scopedClusterClient: wrappedScopedClusterClient.client(),
-            searchSourceClient: this.context.data.search.searchSource.asScoped(fakeRequest),
+            scopedClusterClient: wrappedScopedClusterClient.scopedClusterClient(),
             alertFactory: createAlertFactory<
               InstanceState,
               InstanceContext,

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

@@ -12,6 +12,10 @@ import type {
   IUiSettingsClient,
 } from 'src/core/server';
 import type { PublicMethodsOf } from '@kbn/utility-types';
+import {
+  AggregationsAggregate,
+  SearchResponse,
+} from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { AlertFactoryDoneUtils, PublicAlert } from './alert';
 import { RuleTypeRegistry as OrigruleTypeRegistry } from './rule_type_registry';
 import { PluginSetupContract, PluginStartContract } from './plugin';
@@ -42,7 +46,7 @@ import {
   MappedParams,
 } from '../common';
 import { LicenseType } from '../../licensing/server';
-import { ISearchStartSearchSource } from '../../../../src/plugins/data/common';
+import { ISearchSource, ISearchStartSearchSource } from '../../../../src/plugins/data/common';
 import { RuleTypeConfig } from './config';
 export type WithoutQueryAndParams<T> = Pick<T, Exclude<keyof T, 'query' | 'params'>>;
 export type SpaceIdToNamespaceFunction = (spaceId?: string) => string | undefined;
@@ -74,7 +78,12 @@ export interface RuleExecutorServices<
   InstanceContext extends AlertInstanceContext = AlertInstanceContext,
   ActionGroupIds extends string = never
 > {
-  searchSourceClient: Promise<ISearchStartSearchSource>;
+  searchSourceUtils: {
+    searchSourceClient: ISearchStartSearchSource;
+    wrappedFetch: (
+      searchSource: ISearchSource
+    ) => Promise<SearchResponse<unknown, Record<string, AggregationsAggregate>>>;
+  };
   savedObjectsClient: SavedObjectsClientContract;
   uiSettingsClient: IUiSettingsClient;
   scopedClusterClient: IScopedClusterClient;

x-pack/plugins/stack_alerts/server/alert_types/es_query/executor.ts

@@ -24,7 +24,7 @@ export async function executor(
 ) {
   const esQueryAlert = isEsQueryAlert(options);
   const { alertId, name, services, params, state } = options;
-  const { alertFactory, scopedClusterClient, searchSourceClient } = services;
+  const { alertFactory, scopedClusterClient, searchSourceUtils } = services;
   const currentTimestamp = new Date().toISOString();
   const publicBaseUrl = core.http.basePath.publicBaseUrl ?? '';
 
@@ -51,10 +51,7 @@ export async function executor(
         alertId,
         params as OnlySearchSourceAlertParams,
         latestTimestamp,
-        {
-          searchSourceClient,
-          logger,
-        }
+        { searchSourceUtils, logger }
       );
 
   // apply the alert condition

x-pack/plugins/stack_alerts/server/alert_types/es_query/lib/fetch_search_source_query.ts

@@ -6,6 +6,10 @@
  */
 import { buildRangeFilter, Filter } from '@kbn/es-query';
 import { Logger } from 'kibana/server';
+import {
+  AggregationsAggregate,
+  SearchResponse,
+} from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { OnlySearchSourceAlertParams } from '../types';
 import {
   getTime,
@@ -14,18 +18,26 @@ import {
   SortDirection,
 } from '../../../../../../../src/plugins/data/common';
 
+interface SearchSourceUtils {
+  searchSourceClient: ISearchStartSearchSource;
+  wrappedFetch: (
+    searchSource: ISearchSource
+  ) => Promise<SearchResponse<unknown, Record<string, AggregationsAggregate>>>;
+}
+
 export async function fetchSearchSourceQuery(
   alertId: string,
   params: OnlySearchSourceAlertParams,
   latestTimestamp: string | undefined,
   services: {
     logger: Logger;
-    searchSourceClient: Promise<ISearchStartSearchSource>;
+    searchSourceUtils: SearchSourceUtils;
   }
 ) {
-  const { logger, searchSourceClient } = services;
-  const client = await searchSourceClient;
-  const initialSearchSource = await client.create(params.searchConfiguration);
+  const { logger, searchSourceUtils } = services;
+  const { searchSourceClient, wrappedFetch } = searchSourceUtils;
+
+  const initialSearchSource = await searchSourceClient.create(params.searchConfiguration);
 
   const { searchSource, dateStart, dateEnd } = updateSearchSource(
     initialSearchSource,
@@ -39,7 +51,7 @@ export async function fetchSearchSourceQuery(
     )}`
   );
 
-  const searchResult = await searchSource.fetch();
+  const searchResult = await wrappedFetch(searchSource);
 
   return {
     numMatches: Number(searchResult.hits.total),