Skip to content

Commit

Permalink
[Ingest Manager] Fix limited concurrency helper (elastic#73976)
Browse files Browse the repository at this point in the history 
# Conflicts:
#	x-pack/plugins/ingest_manager/server/routes/limited_concurrency.test.ts
#	x-pack/plugins/ingest_manager/server/routes/limited_concurrency.ts
  • Loading branch information
@nchaulet
nchaulet committed Aug 1, 2020
1 parent ec583ce commit e4d7dde
Show file tree
Hide file tree
Showing 2 changed files with 218 additions and 19 deletions.
188 changes: 186 additions & 2 deletions x-pack/plugins/ingest_manager/server/routes/limited_concurrency.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,12 @@
* you may not use this file except in compliance with the Elastic License.
*/

import { coreMock } from 'src/core/server/mocks';
import { registerLimitedConcurrencyRoutes } from './limited_concurrency';
import { coreMock, httpServerMock, httpServiceMock } from 'src/core/server/mocks';
import {
createLimitedPreAuthHandler,
isLimitedRoute,
registerLimitedConcurrencyRoutes,
} from './limited_concurrency';
import { IngestManagerConfigType } from '../index';

describe('registerLimitedConcurrencyRoutes', () => {
Expand Down Expand Up @@ -33,3 +37,183 @@ describe('registerLimitedConcurrencyRoutes', () => {
expect(mockSetup.http.registerOnPreAuth).toHaveBeenCalledTimes(1);
});
});

// assertions for calls to .decrease are commented out because it's called on the
// "req.events.completed$ observable (which) will never emit from a mocked request in a jest unit test environment"
// https://github.com/elastic/kibana/pull/72338#issuecomment-661908791
describe('preAuthHandler', () => {
test(`ignores routes when !isMatch`, async () => {
const mockMaxCounter = {
increase: jest.fn(),
decrease: jest.fn(),
lessThanMax: jest.fn(),
};
const preAuthHandler = createLimitedPreAuthHandler({
isMatch: jest.fn().mockImplementation(() => false),
maxCounter: mockMaxCounter,
});

const mockRequest = httpServerMock.createKibanaRequest({
path: '/no/match',
});
const mockResponse = httpServerMock.createResponseFactory();
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();

await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);

expect(mockMaxCounter.increase).not.toHaveBeenCalled();
expect(mockMaxCounter.decrease).not.toHaveBeenCalled();
expect(mockMaxCounter.lessThanMax).not.toHaveBeenCalled();
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
});

test(`ignores routes which don't have the correct tag`, async () => {
const mockMaxCounter = {
increase: jest.fn(),
decrease: jest.fn(),
lessThanMax: jest.fn(),
};
const preAuthHandler = createLimitedPreAuthHandler({
isMatch: isLimitedRoute,
maxCounter: mockMaxCounter,
});

const mockRequest = httpServerMock.createKibanaRequest({
path: '/no/match',
});
const mockResponse = httpServerMock.createResponseFactory();
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();

await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);

expect(mockMaxCounter.increase).not.toHaveBeenCalled();
expect(mockMaxCounter.decrease).not.toHaveBeenCalled();
expect(mockMaxCounter.lessThanMax).not.toHaveBeenCalled();
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
});

test(`processes routes which have the correct tag`, async () => {
const mockMaxCounter = {
increase: jest.fn(),
decrease: jest.fn(),
lessThanMax: jest.fn().mockImplementation(() => true),
};
const preAuthHandler = createLimitedPreAuthHandler({
isMatch: isLimitedRoute,
maxCounter: mockMaxCounter,
});

const mockRequest = httpServerMock.createKibanaRequest({
path: '/should/match',
routeTags: ['ingest:limited-concurrency'],
});
const mockResponse = httpServerMock.createResponseFactory();
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();

await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);

// will call lessThanMax because isMatch succeeds
expect(mockMaxCounter.lessThanMax).toHaveBeenCalledTimes(1);
// will not error because lessThanMax is true
expect(mockResponse.customError).not.toHaveBeenCalled();
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
});

test(`updates the counter when isMatch & lessThanMax`, async () => {
const mockMaxCounter = {
increase: jest.fn(),
decrease: jest.fn(),
lessThanMax: jest.fn().mockImplementation(() => true),
};
const preAuthHandler = createLimitedPreAuthHandler({
isMatch: jest.fn().mockImplementation(() => true),
maxCounter: mockMaxCounter,
});

const mockRequest = httpServerMock.createKibanaRequest();
const mockResponse = httpServerMock.createResponseFactory();
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();

await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);

expect(mockMaxCounter.increase).toHaveBeenCalled();
// expect(mockMaxCounter.decrease).toHaveBeenCalled();
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
});

test(`lessThanMax ? next : error`, async () => {
const mockMaxCounter = {
increase: jest.fn(),
decrease: jest.fn(),
lessThanMax: jest
.fn()
// call 1
.mockImplementationOnce(() => true)
// calls 2, 3, 4
.mockImplementationOnce(() => false)
.mockImplementationOnce(() => false)
.mockImplementationOnce(() => false)
// calls 5+
.mockImplementationOnce(() => true)
.mockImplementation(() => true),
};

const preAuthHandler = createLimitedPreAuthHandler({
isMatch: isLimitedRoute,
maxCounter: mockMaxCounter,
});

function makeRequestExpectNext() {
const request = httpServerMock.createKibanaRequest({
path: '/should/match/',
routeTags: ['ingest:limited-concurrency'],
});
const response = httpServerMock.createResponseFactory();
const toolkit = httpServiceMock.createOnPreAuthToolkit();

preAuthHandler(request, response, toolkit);
expect(toolkit.next).toHaveBeenCalledTimes(1);
expect(response.customError).not.toHaveBeenCalled();
}

function makeRequestExpectError() {
const request = httpServerMock.createKibanaRequest({
path: '/should/match/',
routeTags: ['ingest:limited-concurrency'],
});
const response = httpServerMock.createResponseFactory();
const toolkit = httpServiceMock.createOnPreAuthToolkit();

preAuthHandler(request, response, toolkit);
expect(toolkit.next).not.toHaveBeenCalled();
expect(response.customError).toHaveBeenCalledTimes(1);
expect(response.customError).toHaveBeenCalledWith({
statusCode: 429,
body: 'Too Many Requests',
});
}

// request 1 succeeds
makeRequestExpectNext();
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(1);
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(1);

// requests 2, 3, 4 fail
makeRequestExpectError();
makeRequestExpectError();
makeRequestExpectError();

// requests 5+ succeed
makeRequestExpectNext();
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(2);
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(2);

makeRequestExpectNext();
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(3);
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(3);

makeRequestExpectNext();
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(4);
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(4);
});
});
49 changes: 32 additions & 17 deletions x-pack/plugins/ingest_manager/server/routes/limited_concurrency.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ import {
} from 'kibana/server';
import { LIMITED_CONCURRENCY_ROUTE_TAG } from '../../common';
import { IngestManagerConfigType } from '../index';
class MaxCounter {

export class MaxCounter {
constructor(private readonly max: number = 1) {}
private counter = 0;
valueOf() {
Expand All @@ -33,40 +34,54 @@ class MaxCounter {
}
}

function shouldHandleRequest(request: KibanaRequest) {
export type IMaxCounter = Pick<MaxCounter, 'increase' | 'decrease' | 'lessThanMax'>;

export function isLimitedRoute(request: KibanaRequest) {
const tags = request.route.options.tags;
return tags.includes(LIMITED_CONCURRENCY_ROUTE_TAG);
return !!tags.includes(LIMITED_CONCURRENCY_ROUTE_TAG);
}

export function registerLimitedConcurrencyRoutes(core: CoreSetup, config: IngestManagerConfigType) {
const max = config.fleet.maxConcurrentConnections;
if (!max) return;

const counter = new MaxCounter(max);
core.http.registerOnPreAuth(function preAuthHandler(
export function createLimitedPreAuthHandler({
isMatch,
maxCounter,
}: {
isMatch: (request: KibanaRequest) => boolean;
maxCounter: IMaxCounter;
}) {
return function preAuthHandler(
request: KibanaRequest,
response: LifecycleResponseFactory,
toolkit: OnPreAuthToolkit
) {
if (!shouldHandleRequest(request)) {
if (!isMatch(request)) {
return toolkit.next();
}

if (!counter.lessThanMax()) {
if (!maxCounter.lessThanMax()) {
return response.customError({
body: 'Too Many Requests',
statusCode: 429,
});
}

counter.increase();
maxCounter.increase();

// requests.events.aborted$ has a bug (but has test which explicitly verifies) where it's fired even when the request completes
// https://github.com/elastic/kibana/pull/70495#issuecomment-656288766
request.events.aborted$.toPromise().then(() => {
counter.decrease();
request.events.completed$.toPromise().then(() => {
maxCounter.decrease();
});

return toolkit.next();
});
};
}

export function registerLimitedConcurrencyRoutes(core: CoreSetup, config: IngestManagerConfigType) {
const max = config.fleet.maxConcurrentConnections;
if (!max) return;

core.http.registerOnPreAuth(
createLimitedPreAuthHandler({
isMatch: isLimitedRoute,
maxCounter: new MaxCounter(max),
})
);
}

0 comments on commit e4d7dde

Please sign in to comment.