Skip to content

Commit

Permalink
[Security Solution] Deduplicate requests to ML (#153244)
Browse files Browse the repository at this point in the history
**Resolves: #134826

## Summary

We've recently had a couple of SDHs related to the slow performance of
Security Solution. The performance issues were related to slow responses
from ML API. In this PR, I rewrite data-fetching hooks to React Query to
leverage request deduplication and caching. That significantly reduces
the number of outgoing HTTP requests to ML routes on page load.

**Before**
9 HTTP requests to ML endpoints on initial page load, 5 of which are
duplicates.

![Screenshot 2023-03-20 at 12 09
55](https://user-images.githubusercontent.com/1938181/226322831-b3f594c9-a08a-4c8a-acc9-247df8d4a028.png)

**After**
4 HTTP requests to ML endpoints on initial page load, no duplicates.

![Screenshot 2023-03-20 at 11 59
33](https://user-images.githubusercontent.com/1938181/226322847-a69129a2-afcd-408d-911c-a7b767dc4fdb.png)
  • Loading branch information
xcrzx authored Mar 23, 2023
1 parent 6a16932 commit 9669cec
Show file tree
Hide file tree
Showing 11 changed files with 227 additions and 119 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -5,30 +5,27 @@
* 2.0.
*/

import type { HttpSetup } from '@kbn/core/public';
import type { MlSummaryJob } from '@kbn/ml-plugin/public';
import { KibanaServices } from '../../../lib/kibana';

export interface GetJobsSummaryArgs {
http: HttpSetup;
jobIds?: string[];
signal: AbortSignal;
signal?: AbortSignal;
}

/**
* Fetches a summary of all ML jobs currently installed
*
* @param http HTTP Service
* @param jobIds Array of job IDs to filter against
* @param signal to cancel request
*
* @throws An error if response is not OK
*/
export const getJobsSummary = async ({
http,
jobIds,
signal,
}: GetJobsSummaryArgs): Promise<MlSummaryJob[]> =>
http.fetch<MlSummaryJob[]>('/api/ml/jobs/jobs_summary', {
KibanaServices.get().http.fetch<MlSummaryJob[]>('/api/ml/jobs/jobs_summary', {
method: 'POST',
body: JSON.stringify({ jobIds: jobIds ?? [] }),
asSystemRequest: true,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
/*
* 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 type { MlSummaryJob } from '@kbn/ml-plugin/public';
import type { UseQueryOptions } from '@tanstack/react-query';
import { useQuery, useQueryClient } from '@tanstack/react-query';
import { useCallback } from 'react';
import type { GetJobsSummaryArgs } from '../api/get_jobs_summary';
import { getJobsSummary } from '../api/get_jobs_summary';

const ONE_MINUTE = 60000;
export const GET_JOBS_SUMMARY_QUERY_KEY = ['POST', '/api/ml/jobs/jobs_summary'];

export const useFetchJobsSummaryQuery = (
queryArgs: Omit<GetJobsSummaryArgs, 'signal'>,
options?: UseQueryOptions<MlSummaryJob[]>
) => {
return useQuery<MlSummaryJob[]>(
[GET_JOBS_SUMMARY_QUERY_KEY, queryArgs],
async ({ signal }) => getJobsSummary({ signal, ...queryArgs }),
{
refetchIntervalInBackground: false,
staleTime: ONE_MINUTE * 5,
retry: false,
...options,
}
);
};

export const useInvalidateFetchJobsSummaryQuery = () => {
const queryClient = useQueryClient();

return useCallback(() => {
queryClient.invalidateQueries(GET_JOBS_SUMMARY_QUERY_KEY, {
refetchType: 'active',
});
}, [queryClient]);
};

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import { useAppToastsMock } from '../../../hooks/use_app_toasts.mock';
import { mockJobsSummaryResponse } from '../../ml_popover/api.mock';
import { getJobsSummary } from '../api/get_jobs_summary';
import { useInstalledSecurityJobs } from './use_installed_security_jobs';
import { TestProviders } from '../../../mock';

jest.mock('../../../../../common/machine_learning/has_ml_user_permissions');
jest.mock('../../../../../common/machine_learning/has_ml_license');
Expand All @@ -37,7 +38,9 @@ describe('useInstalledSecurityJobs', () => {
});

it('returns jobs and permissions', async () => {
const { result, waitForNextUpdate } = renderHook(() => useInstalledSecurityJobs());
const { result, waitForNextUpdate } = renderHook(() => useInstalledSecurityJobs(), {
wrapper: TestProviders,
});
await waitForNextUpdate();

expect(result.current.jobs).toHaveLength(3);
Expand Down Expand Up @@ -68,7 +71,9 @@ describe('useInstalledSecurityJobs', () => {
});

it('filters out non-security jobs', async () => {
const { result, waitForNextUpdate } = renderHook(() => useInstalledSecurityJobs());
const { result, waitForNextUpdate } = renderHook(() => useInstalledSecurityJobs(), {
wrapper: TestProviders,
});
await waitForNextUpdate();

expect(result.current.jobs.length).toBeGreaterThan(0);
Expand All @@ -77,7 +82,9 @@ describe('useInstalledSecurityJobs', () => {

it('renders a toast error if the ML call fails', async () => {
(getJobsSummary as jest.Mock).mockRejectedValue('whoops');
const { waitForNextUpdate } = renderHook(() => useInstalledSecurityJobs());
const { waitForNextUpdate } = renderHook(() => useInstalledSecurityJobs(), {
wrapper: TestProviders,
});
await waitForNextUpdate();

expect(appToastsMock.addError).toHaveBeenCalledWith('whoops', {
Expand All @@ -93,7 +100,9 @@ describe('useInstalledSecurityJobs', () => {
});

it('returns empty jobs and false predicates', () => {
const { result } = renderHook(() => useInstalledSecurityJobs());
const { result } = renderHook(() => useInstalledSecurityJobs(), {
wrapper: TestProviders,
});

expect(result.current.jobs).toEqual([]);
expect(result.current.isMlUser).toEqual(false);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,15 @@
* 2.0.
*/

import { useEffect, useMemo, useState } from 'react';

import type { MlSummaryJob } from '@kbn/ml-plugin/public';
import { hasMlUserPermissions } from '../../../../../common/machine_learning/has_ml_user_permissions';
import { useMemo } from 'react';
import { hasMlLicense } from '../../../../../common/machine_learning/has_ml_license';
import { hasMlUserPermissions } from '../../../../../common/machine_learning/has_ml_user_permissions';
import { isSecurityJob } from '../../../../../common/machine_learning/is_security_job';
import { useAppToasts } from '../../../hooks/use_app_toasts';
import { useHttp } from '../../../lib/kibana';
import { useMlCapabilities } from './use_ml_capabilities';
import * as i18n from '../translations';
import { useGetJobsSummary } from './use_get_jobs_summary';
import { useFetchJobsSummaryQuery } from './use_fetch_jobs_summary_query';
import { useMlCapabilities } from './use_ml_capabilities';

export interface UseInstalledSecurityJobsReturn {
loading: boolean;
Expand All @@ -35,35 +33,24 @@ export interface UseInstalledSecurityJobsReturn {
*
*/
export const useInstalledSecurityJobs = (): UseInstalledSecurityJobsReturn => {
const [jobs, setJobs] = useState<MlSummaryJob[]>([]);
const { addError } = useAppToasts();
const mlCapabilities = useMlCapabilities();
const http = useHttp();
const { error, loading, result, start } = useGetJobsSummary();

const isMlUser = hasMlUserPermissions(mlCapabilities);
const isLicensed = hasMlLicense(mlCapabilities);

useEffect(() => {
if (isMlUser && isLicensed) {
start({ http });
}
}, [http, isMlUser, isLicensed, start]);

useEffect(() => {
if (result) {
const securityJobs = result.filter(isSecurityJob);
setJobs(securityJobs);
const { isFetching, data: jobs = [] } = useFetchJobsSummaryQuery(
{},
{
enabled: isMlUser && isLicensed,
onError: (error) => {
addError(error, { title: i18n.SIEM_JOB_FETCH_FAILURE });
},
}
}, [result]);
);

useEffect(() => {
if (error) {
addError(error, { title: i18n.SIEM_JOB_FETCH_FAILURE });
}
}, [addError, error]);
const securityJobs = jobs.filter(isSecurityJob);

return { isLicensed, isMlUser, jobs, loading };
return { isLicensed, isMlUser, jobs: securityJobs, loading: isFetching };
};

export const useInstalledSecurityJobsIds = () => {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
/*
* 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 type { UseQueryOptions } from '@tanstack/react-query';
import { useQuery, useQueryClient } from '@tanstack/react-query';
import { useCallback } from 'react';
import { getModules } from '../api';
import type { GetModulesProps, Module } from '../types';

const ONE_MINUTE = 60000;
export const GET_MODULES_QUERY_KEY = ['GET', '/api/ml/modules/get_module/:moduleId'];

export const useFetchModulesQuery = (
queryArgs: Omit<GetModulesProps, 'signal'>,
options?: UseQueryOptions<Module[]>
) => {
return useQuery<Module[]>(
[GET_MODULES_QUERY_KEY, queryArgs],
async ({ signal }) => getModules({ signal, ...queryArgs }),
{
refetchIntervalInBackground: false,
staleTime: ONE_MINUTE * 5,
retry: false,
...options,
}
);
};

export const useInvalidateFetchModulesQuery = () => {
const queryClient = useQueryClient();

return useCallback(() => {
queryClient.invalidateQueries(GET_MODULES_QUERY_KEY, { refetchType: 'active' });
}, [queryClient]);
};
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
/*
* 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 type { UseQueryOptions } from '@tanstack/react-query';
import { useQuery, useQueryClient } from '@tanstack/react-query';
import { useCallback } from 'react';
import { checkRecognizer } from '../api';
import type { CheckRecognizerProps, RecognizerModule } from '../types';

const ONE_MINUTE = 60000;
export const GET_RECOGNIZER_QUERY_KEY = ['GET', '/api/ml/modules/recognize/:indexPatterns'];

export const useFetchRecognizerQuery = (
queryArgs: Omit<CheckRecognizerProps, 'signal'>,
options?: UseQueryOptions<RecognizerModule[]>
) => {
return useQuery<RecognizerModule[]>(
[GET_RECOGNIZER_QUERY_KEY, queryArgs],
async ({ signal }) => checkRecognizer({ signal, ...queryArgs }),
{
refetchIntervalInBackground: false,
staleTime: ONE_MINUTE * 5,
retry: false,
...options,
}
);
};

export const useInvalidateFetchRecognizerQuery = () => {
const queryClient = useQueryClient();

return useCallback(() => {
queryClient.invalidateQueries(GET_RECOGNIZER_QUERY_KEY, { refetchType: 'active' });
}, [queryClient]);
};
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import {
checkRecognizerSuccess,
} from '../api.mock';
import { useSecurityJobs } from './use_security_jobs';
import { TestProviders } from '../../../mock';

jest.mock('../../../../../common/machine_learning/has_ml_admin_permissions');
jest.mock('../../../../../common/machine_learning/has_ml_license');
Expand Down Expand Up @@ -71,15 +72,19 @@ describe('useSecurityJobs', () => {
bucketSpanSeconds: 900,
};

const { result, waitForNextUpdate } = renderHook(() => useSecurityJobs());
const { result, waitForNextUpdate } = renderHook(() => useSecurityJobs(), {
wrapper: TestProviders,
});
await waitForNextUpdate();

expect(result.current.jobs).toHaveLength(6);
expect(result.current.jobs).toEqual(expect.arrayContaining([expectedSecurityJob]));
});

it('returns those permissions', async () => {
const { result, waitForNextUpdate } = renderHook(() => useSecurityJobs());
const { result, waitForNextUpdate } = renderHook(() => useSecurityJobs(), {
wrapper: TestProviders,
});
await waitForNextUpdate();

expect(result.current.isMlAdmin).toEqual(true);
Expand All @@ -88,7 +93,9 @@ describe('useSecurityJobs', () => {

it('renders a toast error if an ML call fails', async () => {
(getModules as jest.Mock).mockRejectedValue('whoops');
const { waitForNextUpdate } = renderHook(() => useSecurityJobs());
const { waitForNextUpdate } = renderHook(() => useSecurityJobs(), {
wrapper: TestProviders,
});
await waitForNextUpdate();

expect(appToastsMock.addError).toHaveBeenCalledWith('whoops', {
Expand All @@ -104,7 +111,9 @@ describe('useSecurityJobs', () => {
});

it('returns empty jobs and false predicates', () => {
const { result } = renderHook(() => useSecurityJobs());
const { result } = renderHook(() => useSecurityJobs(), {
wrapper: TestProviders,
});

expect(result.current.jobs).toEqual([]);
expect(result.current.isMlAdmin).toEqual(false);
Expand Down
Loading

0 comments on commit 9669cec

Please sign in to comment.