[Search] Search Sessions Monitoring Task

x-pack/plugins/alerts/server/alerts_client/alerts_client.ts

@@ -43,7 +43,6 @@ import {
 import { EncryptedSavedObjectsClient } from '../../../encrypted_saved_objects/server';
 import { TaskManagerStartContract } from '../../../task_manager/server';
 import { taskInstanceToAlertTaskInstance } from '../task_runner/alert_task_instance';
-import { deleteTaskIfItExists } from '../lib/delete_task_if_it_exists';
 import { RegistryAlertType, UntypedNormalizedAlertType } from '../alert_type_registry';
 import { AlertsAuthorization, WriteOperations, ReadOperations } from '../authorization';
 import { IEventLogClient } from '../../../../plugins/event_log/server';
@@ -602,7 +601,7 @@ export class AlertsClient {
     const removeResult = await this.unsecuredSavedObjectsClient.delete('alert', id);
 
     await Promise.all([
-      taskIdToRemove ? deleteTaskIfItExists(this.taskManager, taskIdToRemove) : null,
+      taskIdToRemove ? this.taskManager.deleteTaskIfItExists(taskIdToRemove) : null,
       apiKeyToInvalidate
         ? markApiKeyForInvalidation(
             { apiKey: apiKeyToInvalidate },
@@ -1060,7 +1059,7 @@ export class AlertsClient {
 
       await Promise.all([
         attributes.scheduledTaskId
-          ? deleteTaskIfItExists(this.taskManager, attributes.scheduledTaskId)
+          ? this.taskManager.deleteTaskIfItExists(attributes.scheduledTaskId)
           : null,
         apiKeyToInvalidate
           ? await markApiKeyForInvalidation(

x-pack/plugins/data_enhanced/common/index.ts

@@ -13,7 +13,6 @@ export {
   IAsyncSearchOptions,
   pollSearch,
   SearchSessionSavedObjectAttributes,
-  SearchSessionFindOptions,
   SearchSessionStatus,
   SearchSessionRequestInfo,
 } from './search';

x-pack/plugins/data_enhanced/common/search/session/types.ts

@@ -13,18 +13,55 @@ export interface SearchSessionSavedObjectAttributes {
    * App that created the session. e.g 'discover'
    */
   appId: string;
+  /**
+   * Creation time of the session
+   */
   created: string;
+  /**
+   * Expiration time of the session. Expiration itself is managed by Elasticsearch.
+   */
   expires: string;
+  /**
+   * status
+   */
   status: string;
+  /**
+   * urlGeneratorId
+   */
   urlGeneratorId: string;
+  /**
+   * The application state that was used to create the session.
+   * Should be used, for example, to re-load an expired search session.
+   */
   initialState: Record<string, unknown>;
+  /**
+   * Application state that should be used to restore the session.
+   * For example, relative dates are conveted to absolute ones.
+   */
   restoreState: Record<string, unknown>;
+  /**
+   * Mapping of search request hashes to their corresponsing info (async search id, etc.)
+   */
   idMapping: Record<string, SearchSessionRequestInfo>;
 }
 
 export interface SearchSessionRequestInfo {
-  id: string; // ID of the async search request
-  strategy: string; // Search strategy used to submit the search request
+  /**
+   * ID of the async search request
+   */
+  id: string;
+  /**
+   * Search strategy used to submit the search request
+   */
+  strategy: string;
+  /**
+   * status
+   */
+  status: string;
+  /**
+   * An optional error. Set if status is set to error.
+   */
+  error?: string;
 }
 
 export interface SearchSessionFindOptions {

x-pack/plugins/data_enhanced/kibana.json

@@ -8,7 +8,8 @@
   "requiredPlugins": [
     "bfetch",
     "data",
-    "features"
+    "features",
+    "taskManager"
   ],
   "optionalPlugins": ["kibanaUtils", "usageCollection"],
   "server": true,

x-pack/plugins/data_enhanced/server/plugin.ts

@@ -5,6 +5,7 @@
  */
 
 import { CoreSetup, CoreStart, Logger, Plugin, PluginInitializerContext } from 'kibana/server';
+import { TaskManagerSetupContract, TaskManagerStartContract } from '../../task_manager/server';
 import {
   PluginSetup as DataPluginSetup,
   PluginStart as DataPluginStart,
@@ -24,9 +25,15 @@ import { getUiSettings } from './ui_settings';
 interface SetupDependencies {
   data: DataPluginSetup;
   usageCollection?: UsageCollectionSetup;
+  taskManager: TaskManagerSetupContract;
+}
+export interface StartDependencies {
+  data: DataPluginStart;
+  taskManager: TaskManagerStartContract;
 }
 
-export class EnhancedDataServerPlugin implements Plugin<void, void, SetupDependencies> {
+export class EnhancedDataServerPlugin
+  implements Plugin<void, void, SetupDependencies, StartDependencies> {
   private readonly logger: Logger;
   private sessionService!: SearchSessionService;
 
@@ -65,10 +72,17 @@ export class EnhancedDataServerPlugin implements Plugin<void, void, SetupDepende
 
     const router = core.http.createRouter();
     registerSessionRoutes(router);
+
+    this.sessionService.setup(core, {
+      taskManager: deps.taskManager,
+    });
   }
 
-  public start(core: CoreStart) {
-    this.sessionService.start(core, this.initializerContext.config.create());
+  public start(core: CoreStart, { taskManager }: StartDependencies) {
+    this.sessionService.start(core, {
+      taskManager,
+      config$: this.initializerContext.config.create(),
+    });
   }
 
   public stop() {

x-pack/plugins/data_enhanced/server/search/session/check_running_sessions.test.ts

@@ -0,0 +1,191 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { checkRunningSessions } from './check_running_sessions';
+import { SearchSessionStatus, SearchSessionSavedObjectAttributes } from '../../../common';
+import { savedObjectsClientMock } from '../../../../../../src/core/server/mocks';
+import type { SavedObjectsClientContract } from 'kibana/server';
+import { SearchStatus } from './types';
+
+describe('getSearchStatus', () => {
+  let mockClient: any;
+  let savedObjectsClient: jest.Mocked<SavedObjectsClientContract>;
+  const mockLogger: any = {
+    debug: jest.fn(),
+    warn: jest.fn(),
+    error: jest.fn(),
+  };
+
+  beforeEach(() => {
+    savedObjectsClient = savedObjectsClientMock.create();
+    mockClient = {
+      asyncSearch: {
+        status: jest.fn(),
+      },
+    };
+  });
+
+  test('does nothing if there are no open sessions', async () => {
+    savedObjectsClient.bulkUpdate = jest.fn();
+    savedObjectsClient.find.mockResolvedValue({
+      saved_objects: [],
+      total: 0,
+    } as any);
+
+    await checkRunningSessions(savedObjectsClient, mockClient, mockLogger);
+
+    expect(savedObjectsClient.bulkUpdate).not.toBeCalled();
+  });
+
+  test('does nothing if there are no searchIds in the saved object', async () => {
+    savedObjectsClient.bulkUpdate = jest.fn();
+    savedObjectsClient.find.mockResolvedValue({
+      saved_objects: [
+        {
+          attributes: {
+            idMapping: {},
+          },
+        },
+      ],
+      total: 1,
+    } as any);
+
+    await checkRunningSessions(savedObjectsClient, mockClient, mockLogger);
+
+    expect(savedObjectsClient.bulkUpdate).not.toBeCalled();
+  });
+
+  test('does nothing if the search is still running', async () => {
+    savedObjectsClient.bulkUpdate = jest.fn();
+    const so = {
+      attributes: {
+        idMapping: {
+          'search-hash': {
+            id: 'search-id',
+            strategy: 'cool',
+            status: SearchStatus.IN_PROGRESS,
+          },
+        },
+      },
+    };
+    savedObjectsClient.find.mockResolvedValue({
+      saved_objects: [so],
+      total: 1,
+    } as any);
+
+    mockClient.asyncSearch.status.mockResolvedValue({
+      body: {
+        is_partial: true,
+        is_running: true,
+      },
+    });
+
+    await checkRunningSessions(savedObjectsClient, mockClient, mockLogger);
+
+    expect(savedObjectsClient.bulkUpdate).not.toBeCalled();
+  });
+
+  test("doesn't re-check completed or errored searches", async () => {
+    savedObjectsClient.bulkUpdate = jest.fn();
+    const so = {
+      attributes: {
+        idMapping: {
+          'search-hash': {
+            id: 'search-id',
+            strategy: 'cool',
+            status: SearchStatus.COMPLETE,
+          },
+          'another-search-hash': {
+            id: 'search-id',
+            strategy: 'cool',
+            status: SearchStatus.ERROR,
+          },
+        },
+      },
+    };
+    savedObjectsClient.find.mockResolvedValue({
+      saved_objects: [so],
+      total: 1,
+    } as any);
+
+    await checkRunningSessions(savedObjectsClient, mockClient, mockLogger);
+
+    expect(mockClient.asyncSearch.status).not.toBeCalled();
+  });
+
+  test('updates to complete if the search is done', async () => {
+    savedObjectsClient.bulkUpdate = jest.fn();
+    const so = {
+      attributes: {
+        idMapping: {
+          'search-hash': {
+            id: 'search-id',
+            strategy: 'cool',
+            status: SearchStatus.IN_PROGRESS,
+          },
+        },
+      },
+    };
+    savedObjectsClient.find.mockResolvedValue({
+      saved_objects: [so],
+      total: 1,
+    } as any);
+
+    mockClient.asyncSearch.status.mockResolvedValue({
+      body: {
+        is_partial: false,
+        is_running: false,
+        completion_status: 200,
+      },
+    });
+
+    await checkRunningSessions(savedObjectsClient, mockClient, mockLogger);
+
+    expect(mockClient.asyncSearch.status).toBeCalledWith({ id: 'search-id' });
+    const [updateInput] = savedObjectsClient.bulkUpdate.mock.calls[0];
+    const updatedAttributes = updateInput[0].attributes as SearchSessionSavedObjectAttributes;
+    expect(updatedAttributes.status).toBe(SearchSessionStatus.COMPLETE);
+    expect(updatedAttributes.idMapping['search-hash'].status).toBe(SearchStatus.COMPLETE);
+    expect(updatedAttributes.idMapping['search-hash'].error).toBeUndefined();
+  });
+
+  test('updates to error if the search is errored', async () => {
+    savedObjectsClient.bulkUpdate = jest.fn();
+    const so = {
+      attributes: {
+        idMapping: {
+          'search-hash': {
+            id: 'search-id',
+            strategy: 'cool',
+            status: SearchStatus.IN_PROGRESS,
+          },
+        },
+      },
+    };
+    savedObjectsClient.find.mockResolvedValue({
+      saved_objects: [so],
+      total: 1,
+    } as any);
+
+    mockClient.asyncSearch.status.mockResolvedValue({
+      body: {
+        is_partial: false,
+        is_running: false,
+        completion_status: 500,
+      },
+    });
+
+    await checkRunningSessions(savedObjectsClient, mockClient, mockLogger);
+    const [updateInput] = savedObjectsClient.bulkUpdate.mock.calls[0];
+
+    const updatedAttributes = updateInput[0].attributes as SearchSessionSavedObjectAttributes;
+    expect(updatedAttributes.status).toBe(SearchSessionStatus.ERROR);
+    expect(updatedAttributes.idMapping['search-hash'].status).toBe(SearchStatus.ERROR);
+    expect(updatedAttributes.idMapping['search-hash'].error).toBe(
+      'Search completed with a 500 status'
+    );
+  });
+});