[ML] ML client and shared services optimizations

x-pack/plugins/apm/server/lib/anomaly_detection/get_ml_jobs_with_apm_group.ts

@@ -5,7 +5,8 @@
  * 2.0.
  */
 
-import { MlPluginSetup } from '@kbn/ml-plugin/server';
+import { DATAFEED_STATE, JOB_STATE } from '@kbn/ml-plugin/common';
+import { MlAnomalyDetectors } from '@kbn/ml-plugin/server';
 import { ApmMlJob } from '../../../common/anomaly_detection/apm_ml_job';
 import { Environment } from '../../../common/environment_rt';
 import { withApmSpan } from '../../utils/with_apm_span';
@@ -23,11 +24,11 @@ function catch404(e: any) {
 }
 
 export function getMlJobsWithAPMGroup(
-  anomalyDetectors: ReturnType<MlPluginSetup['anomalyDetectorsProvider']>
+  anomalyDetectors: MlAnomalyDetectors
 ): Promise<ApmMlJob[]> {
   return withApmSpan('get_ml_jobs_with_apm_group', async () => {
     try {
-      const [jobs, allJobStats, allDatafeedStats] = await Promise.all([
+      const [jobs, jobStats, datafeedStats] = await Promise.all([
         anomalyDetectors
           .jobs(APM_ML_JOB_GROUP)
           .then((response) => response.jobs),
@@ -41,26 +42,27 @@ export function getMlJobsWithAPMGroup(
           .catch(catch404),
       ]);
 
-      return jobs.map((job): ApmMlJob => {
-        const jobStats = allJobStats.find(
-          (stats) => stats.job_id === job.job_id
-        );
+      const datafeedStateMap = Object.fromEntries(
+        datafeedStats.map((d) => [d.datafeed_id, d.state as DATAFEED_STATE])
+      );
 
-        const datafeedStats = allDatafeedStats.find(
-          (stats) => stats.datafeed_id === job.datafeed_config?.datafeed_id
-        );
+      const jobStateMap = Object.fromEntries(
+        jobStats.map((j) => [j.job_id, j.state as JOB_STATE])
+      );
 
+      return jobs.map((job): ApmMlJob => {
+        const jobId = job.job_id;
+        const datafeedId = job.datafeed_config?.datafeed_id;
         return {
+          jobId,
+          jobState: jobStateMap[jobId],
+          datafeedId,
+          datafeedState: datafeedId ? datafeedStateMap[datafeedId] : undefined,
+          version: Number(job?.custom_settings?.job_tags?.apm_ml_version ?? 1),
           environment: String(
-            job.custom_settings?.job_tags?.environment
+            job?.custom_settings?.job_tags?.environment
           ) as Environment,
-          jobId: job.job_id,
-          jobState: jobStats?.state as ApmMlJob['jobState'],
-          version: Number(job.custom_settings?.job_tags?.apm_ml_version ?? 1),
-          datafeedId: datafeedStats?.datafeed_id,
-          datafeedState: datafeedStats?.state as ApmMlJob['datafeedState'],
-          // @ts-expect-error bucket_span is of type `estypes.Duration`
-          bucketSpan: job.analysis_config?.bucket_span,
+          bucketSpan: job?.analysis_config.bucket_span as string,
         };
       });
     } catch (e) {

x-pack/plugins/apm/server/lib/helpers/get_ml_client.ts

@@ -5,13 +5,18 @@
  * 2.0.
  */
 
+import {
+  MlAnomalyDetectors,
+  MlMlSystem,
+  MlModules,
+} from '@kbn/ml-plugin/server';
 import { isActivePlatinumLicense } from '../../../common/license_check';
 import { APMRouteHandlerResources } from '../../routes/typings';
 
 export interface MlClient {
-  mlSystem: any;
-  anomalyDetectors: any;
-  modules: any;
+  mlSystem: MlMlSystem;
+  anomalyDetectors: MlAnomalyDetectors;
+  modules: MlModules;
 }
 
 export async function getMlClient({

x-pack/plugins/apm/server/routes/service_map/get_service_anomalies.ts

@@ -9,7 +9,7 @@ import Boom from '@hapi/boom';
 import { sortBy, uniqBy } from 'lodash';
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import type { ESSearchResponse } from '@kbn/es-types';
-import { MlPluginSetup } from '@kbn/ml-plugin/server';
+import type { MlAnomalyDetectors } from '@kbn/ml-plugin/server';
 import { rangeQuery } from '@kbn/observability-plugin/server';
 import { getSeverity, ML_ERRORS } from '../../../common/anomaly_detection';
 import { ENVIRONMENT_ALL } from '../../../common/environment_filter_values';
@@ -154,7 +154,7 @@ export async function getServiceAnomalies({
 }
 
 export async function getMLJobs(
-  anomalyDetectors: ReturnType<MlPluginSetup['anomalyDetectorsProvider']>,
+  anomalyDetectors: MlAnomalyDetectors,
   environment?: string
 ) {
   const jobs = await getMlJobsWithAPMGroup(anomalyDetectors);
@@ -173,7 +173,7 @@ export async function getMLJobs(
 }
 
 export async function getMLJobIds(
-  anomalyDetectors: ReturnType<MlPluginSetup['anomalyDetectorsProvider']>,
+  anomalyDetectors: MlAnomalyDetectors,
   environment?: string
 ) {
   const mlJobs = await getMLJobs(anomalyDetectors, environment);

x-pack/plugins/ml/server/lib/capabilities/check_capabilities.ts

@@ -6,6 +6,7 @@
  */
 
 import { KibanaRequest } from '@kbn/core/server';
+import { once } from 'lodash';
 import type { MlClient } from '../ml_client';
 import { mlLog } from '../log';
 import {
@@ -60,24 +61,28 @@ function disableAdminPrivileges(capabilities: MlCapabilities) {
 
 export type HasMlCapabilities = (capabilities: MlCapabilitiesKey[]) => Promise<void>;
 
-export function hasMlCapabilitiesProvider(resolveMlCapabilities: ResolveMlCapabilities) {
-  return (request: KibanaRequest): HasMlCapabilities => {
-    let mlCapabilities: MlCapabilities | null = null;
-    return async (capabilities: MlCapabilitiesKey[]) => {
-      try {
-        mlCapabilities = await resolveMlCapabilities(request);
-      } catch (e) {
-        mlLog.error(e);
-        throw new UnknownMLCapabilitiesError(`Unable to perform ML capabilities check ${e}`);
-      }
+export function hasMlCapabilitiesProvider(
+  resolveMlCapabilities: ResolveMlCapabilities,
+  request: KibanaRequest
+) {
+  let mlCapabilities: MlCapabilities | null = null;
 
-      if (mlCapabilities === null) {
-        throw new MLPrivilegesUninitialized('ML capabilities have not been initialized');
-      }
+  const resolveMlCapabilitiesOnce = once(resolveMlCapabilities);
 
-      if (capabilities.every((c) => mlCapabilities![c] === true) === false) {
-        throw new InsufficientMLCapabilities('Insufficient privileges to access feature');
-      }
-    };
+  return async (capabilities: MlCapabilitiesKey[]) => {
+    try {
+      mlCapabilities = await resolveMlCapabilitiesOnce(request);
+    } catch (e) {
+      mlLog.error(e);
+      throw new UnknownMLCapabilitiesError(`Unable to perform ML capabilities check ${e}`);
+    }
+
+    if (mlCapabilities === null) {
+      throw new MLPrivilegesUninitialized('ML capabilities have not been initialized');
+    }
+
+    if (capabilities.every((c) => mlCapabilities![c] === true) === false) {
+      throw new InsufficientMLCapabilities('Insufficient privileges to access feature');
+    }
   };
 }

x-pack/plugins/ml/server/saved_objects/service.ts

@@ -6,6 +6,7 @@
  */
 
 import RE2 from 're2';
+import { memoize } from 'lodash';
 import {
   KibanaRequest,
   SavedObjectsClientContract,
@@ -57,6 +58,11 @@ export function mlSavedObjectServiceFactory(
   client: IScopedClusterClient,
   isMlReady: () => Promise<void>
 ) {
+  const _savedObjectsClientFindMemo = memoize(
+    async <T>(options: SavedObjectsFindOptions) => savedObjectsClient.find<T>(options),
+    (options: SavedObjectsFindOptions) => JSON.stringify(options)
+  );
+
   async function _getJobObjects(
     jobType?: JobType,
     jobId?: string,
@@ -87,8 +93,7 @@ export function mlSavedObjectServiceFactory(
       filter,
     };
 
-    const jobs = await savedObjectsClient.find<JobObject>(options);
-
+    const jobs = await _savedObjectsClientFindMemo<JobObject>(options);
     return jobs.saved_objects;
   }
 
@@ -231,7 +236,8 @@ export function mlSavedObjectServiceFactory(
       filter,
     };
 
-    return (await internalSavedObjectsClient.find<JobObject>(options)).saved_objects;
+    const jobs = await _savedObjectsClientFindMemo<JobObject>(options);
+    return jobs.saved_objects;
   }
 
   async function addDatafeed(datafeedId: string, jobId: string) {
@@ -462,8 +468,7 @@ export function mlSavedObjectServiceFactory(
       filter,
     };
 
-    const models = await savedObjectsClient.find<TrainedModelObject>(options);
-
+    const models = await _savedObjectsClientFindMemo<TrainedModelObject>(options);
     return models.saved_objects;
   }
 
@@ -677,7 +682,7 @@ export function mlSavedObjectServiceFactory(
         searchFields,
         filter,
       };
-      return savedObjectsClient.find<TrainedModelObject>(options);
+      return _savedObjectsClientFindMemo<TrainedModelObject>(options);
     });
 
     const finedResult = await Promise.all(searches);

x-pack/plugins/ml/server/shared_services/providers/alerting_service.ts

@@ -16,9 +16,10 @@ export function getAlertingServiceProvider(getGuards: GetGuards) {
       savedObjectsClient: SavedObjectsClientContract,
       request: KibanaRequest
     ) {
+      const guards = getGuards(request, savedObjectsClient);
       return {
         preview: async (...args: Parameters<MlAlertingService['preview']>) => {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(({ mlClient, scopedClient, getFieldsFormatRegistry, getDataViewsService }) =>
@@ -33,7 +34,7 @@ export function getAlertingServiceProvider(getGuards: GetGuards) {
         execute: async (
           ...args: Parameters<MlAlertingService['execute']>
         ): ReturnType<MlAlertingService['execute']> => {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(({ mlClient, scopedClient, getFieldsFormatRegistry, getDataViewsService }) =>

x-pack/plugins/ml/server/shared_services/providers/anomaly_detectors.ts

@@ -32,9 +32,10 @@ export function getAnomalyDetectorsProvider(getGuards: GetGuards): AnomalyDetect
       request: KibanaRequest,
       savedObjectsClient: SavedObjectsClientContract
     ) {
+      const guards = getGuards(request, savedObjectsClient);
       return {
         async jobs(jobId?: string) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(async ({ mlClient }) => {
@@ -45,7 +46,7 @@ export function getAnomalyDetectorsProvider(getGuards: GetGuards): AnomalyDetect
             });
         },
         async jobStats(jobId?: string) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(async ({ mlClient }) => {
@@ -56,7 +57,7 @@ export function getAnomalyDetectorsProvider(getGuards: GetGuards): AnomalyDetect
             });
         },
         async datafeeds(datafeedId?: string) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetDatafeeds'])
             .ok(async ({ mlClient }) => {
@@ -67,7 +68,7 @@ export function getAnomalyDetectorsProvider(getGuards: GetGuards): AnomalyDetect
             });
         },
         async datafeedStats(datafeedId?: string) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetDatafeeds'])
             .ok(async ({ mlClient }) => {

x-pack/plugins/ml/server/shared_services/providers/index.ts

@@ -0,0 +1,20 @@
+/*
+ * 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.
+ */
+
+export type { MlSystemProvider } from './system';
+export type { JobServiceProvider } from './job_service';
+export type { ModuleSetupPayload, ModulesProvider } from './modules';
+export type { ResultsServiceProvider } from './results_service';
+export type { TrainedModelsProvider } from './trained_models';
+export type { AnomalyDetectorsProvider } from './anomaly_detectors';
+
+export { getMlSystemProvider } from './system';
+export { getJobServiceProvider } from './job_service';
+export { getModulesProvider } from './modules';
+export { getResultsServiceProvider } from './results_service';
+export { getTrainedModelsProvider } from './trained_models';
+export { getAnomalyDetectorsProvider } from './anomaly_detectors';

x-pack/plugins/ml/server/shared_services/providers/job_service.ts

@@ -25,9 +25,10 @@ export interface JobServiceProvider {
 export function getJobServiceProvider(getGuards: GetGuards): JobServiceProvider {
   return {
     jobServiceProvider(request: KibanaRequest, savedObjectsClient: SavedObjectsClientContract) {
+      const guards = getGuards(request, savedObjectsClient);
       return {
         jobsSummary: async (...args) => {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(({ scopedClient, mlClient }) => {
@@ -36,7 +37,7 @@ export function getJobServiceProvider(getGuards: GetGuards): JobServiceProvider
             });
         },
         forceStartDatafeeds: async (...args) => {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canStartStopDatafeed'])
             .ok(({ scopedClient, mlClient }) => {
@@ -45,7 +46,7 @@ export function getJobServiceProvider(getGuards: GetGuards): JobServiceProvider
             });
         },
         stopDatafeeds: async (...args) => {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canStartStopDatafeed'])
             .ok(({ scopedClient, mlClient }) => {

x-pack/plugins/ml/server/shared_services/providers/modules.ts

@@ -33,9 +33,10 @@ export function getModulesProvider(
 ): ModulesProvider {
   return {
     modulesProvider(request: KibanaRequest, savedObjectsClient: SavedObjectsClientContract) {
+      const guards = getGuards(request, savedObjectsClient);
       return {
         async recognize(...args) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(async ({ scopedClient, mlClient, mlSavedObjectService, getDataViewsService }) => {
@@ -52,7 +53,7 @@ export function getModulesProvider(
             });
         },
         async getModule(moduleId: string) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(async ({ scopedClient, mlClient, mlSavedObjectService, getDataViewsService }) => {
@@ -69,7 +70,7 @@ export function getModulesProvider(
             });
         },
         async listModules() {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(async ({ scopedClient, mlClient, mlSavedObjectService, getDataViewsService }) => {
@@ -86,7 +87,7 @@ export function getModulesProvider(
             });
         },
         async setup(payload: ModuleSetupPayload) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canCreateJob'])
             .ok(async ({ scopedClient, mlClient, mlSavedObjectService, getDataViewsService }) => {

x-pack/plugins/ml/server/shared_services/providers/results_service.ts

@@ -23,9 +23,10 @@ export interface ResultsServiceProvider {
 export function getResultsServiceProvider(getGuards: GetGuards): ResultsServiceProvider {
   return {
     resultsServiceProvider(request: KibanaRequest, savedObjectsClient: SavedObjectsClientContract) {
+      const guards = getGuards(request, savedObjectsClient);
       return {
         async getAnomaliesTableData(...args) {
-          return await getGuards(request, savedObjectsClient)
+          return await guards
             .isFullLicense()
             .hasMlCapabilities(['canGetJobs'])
             .ok(async ({ mlClient }) => {