Skip to content

Commit 2ba5b78

Browse files
John Schulzelasticmachine
John Schulz
and
elasticmachine
authored
[Ingest Manager] Add more Fleet concurrency tests #71744 (#72338) (#73213)
* Refactor to make more testable. Add more tests. * Remove ts-ignores. Add comment re: testing limitation Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com> Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
1 parent 33ab341 commit 2ba5b78

File tree

2 files changed

+217
-16
lines changed

2 files changed

+217
-16
lines changed

‎x-pack/plugins/ingest_manager/server/routes/limited_concurrency.test.ts‎

Lines changed: 186 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,12 @@
44
* you may not use this file except in compliance with the Elastic License.
55
*/
66

7-
import { coreMock } from 'src/core/server/mocks';
8-
import { registerLimitedConcurrencyRoutes } from './limited_concurrency';
7+
import { coreMock, httpServerMock, httpServiceMock } from 'src/core/server/mocks';
8+
import {
9+
createLimitedPreAuthHandler,
10+
isLimitedRoute,
11+
registerLimitedConcurrencyRoutes,
12+
} from './limited_concurrency';
913
import { IngestManagerConfigType } from '../index';
1014

1115
describe('registerLimitedConcurrencyRoutes', () => {
@@ -33,3 +37,183 @@ describe('registerLimitedConcurrencyRoutes', () => {
3337
expect(mockSetup.http.registerOnPreAuth).toHaveBeenCalledTimes(1);
3438
});
3539
});
40+
41+
// assertions for calls to .decrease are commented out because it's called on the
42+
// "req.events.aborted$ observable (which) will never emit from a mocked request in a jest unit test environment"
43+
// https://github.com/elastic/kibana/pull/72338#issuecomment-661908791
44+
describe('preAuthHandler', () => {
45+
test(`ignores routes when !isMatch`, async () => {
46+
const mockMaxCounter = {
47+
increase: jest.fn(),
48+
decrease: jest.fn(),
49+
lessThanMax: jest.fn(),
50+
};
51+
const preAuthHandler = createLimitedPreAuthHandler({
52+
isMatch: jest.fn().mockImplementation(() => false),
53+
maxCounter: mockMaxCounter,
54+
});
55+
56+
const mockRequest = httpServerMock.createKibanaRequest({
57+
path: '/no/match',
58+
});
59+
const mockResponse = httpServerMock.createResponseFactory();
60+
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
61+
62+
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
63+
64+
expect(mockMaxCounter.increase).not.toHaveBeenCalled();
65+
expect(mockMaxCounter.decrease).not.toHaveBeenCalled();
66+
expect(mockMaxCounter.lessThanMax).not.toHaveBeenCalled();
67+
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
68+
});
69+
70+
test(`ignores routes which don't have the correct tag`, async () => {
71+
const mockMaxCounter = {
72+
increase: jest.fn(),
73+
decrease: jest.fn(),
74+
lessThanMax: jest.fn(),
75+
};
76+
const preAuthHandler = createLimitedPreAuthHandler({
77+
isMatch: isLimitedRoute,
78+
maxCounter: mockMaxCounter,
79+
});
80+
81+
const mockRequest = httpServerMock.createKibanaRequest({
82+
path: '/no/match',
83+
});
84+
const mockResponse = httpServerMock.createResponseFactory();
85+
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
86+
87+
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
88+
89+
expect(mockMaxCounter.increase).not.toHaveBeenCalled();
90+
expect(mockMaxCounter.decrease).not.toHaveBeenCalled();
91+
expect(mockMaxCounter.lessThanMax).not.toHaveBeenCalled();
92+
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
93+
});
94+
95+
test(`processes routes which have the correct tag`, async () => {
96+
const mockMaxCounter = {
97+
increase: jest.fn(),
98+
decrease: jest.fn(),
99+
lessThanMax: jest.fn().mockImplementation(() => true),
100+
};
101+
const preAuthHandler = createLimitedPreAuthHandler({
102+
isMatch: isLimitedRoute,
103+
maxCounter: mockMaxCounter,
104+
});
105+
106+
const mockRequest = httpServerMock.createKibanaRequest({
107+
path: '/should/match',
108+
routeTags: ['ingest:limited-concurrency'],
109+
});
110+
const mockResponse = httpServerMock.createResponseFactory();
111+
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
112+
113+
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
114+
115+
// will call lessThanMax because isMatch succeeds
116+
expect(mockMaxCounter.lessThanMax).toHaveBeenCalledTimes(1);
117+
// will not error because lessThanMax is true
118+
expect(mockResponse.customError).not.toHaveBeenCalled();
119+
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
120+
});
121+
122+
test(`updates the counter when isMatch & lessThanMax`, async () => {
123+
const mockMaxCounter = {
124+
increase: jest.fn(),
125+
decrease: jest.fn(),
126+
lessThanMax: jest.fn().mockImplementation(() => true),
127+
};
128+
const preAuthHandler = createLimitedPreAuthHandler({
129+
isMatch: jest.fn().mockImplementation(() => true),
130+
maxCounter: mockMaxCounter,
131+
});
132+
133+
const mockRequest = httpServerMock.createKibanaRequest();
134+
const mockResponse = httpServerMock.createResponseFactory();
135+
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
136+
137+
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
138+
139+
expect(mockMaxCounter.increase).toHaveBeenCalled();
140+
// expect(mockMaxCounter.decrease).toHaveBeenCalled();
141+
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
142+
});
143+
144+
test(`lessThanMax ? next : error`, async () => {
145+
const mockMaxCounter = {
146+
increase: jest.fn(),
147+
decrease: jest.fn(),
148+
lessThanMax: jest
149+
.fn()
150+
// call 1
151+
.mockImplementationOnce(() => true)
152+
// calls 2, 3, 4
153+
.mockImplementationOnce(() => false)
154+
.mockImplementationOnce(() => false)
155+
.mockImplementationOnce(() => false)
156+
// calls 5+
157+
.mockImplementationOnce(() => true)
158+
.mockImplementation(() => true),
159+
};
160+
161+
const preAuthHandler = createLimitedPreAuthHandler({
162+
isMatch: isLimitedRoute,
163+
maxCounter: mockMaxCounter,
164+
});
165+
166+
function makeRequestExpectNext() {
167+
const request = httpServerMock.createKibanaRequest({
168+
path: '/should/match/',
169+
routeTags: ['ingest:limited-concurrency'],
170+
});
171+
const response = httpServerMock.createResponseFactory();
172+
const toolkit = httpServiceMock.createOnPreAuthToolkit();
173+
174+
preAuthHandler(request, response, toolkit);
175+
expect(toolkit.next).toHaveBeenCalledTimes(1);
176+
expect(response.customError).not.toHaveBeenCalled();
177+
}
178+
179+
function makeRequestExpectError() {
180+
const request = httpServerMock.createKibanaRequest({
181+
path: '/should/match/',
182+
routeTags: ['ingest:limited-concurrency'],
183+
});
184+
const response = httpServerMock.createResponseFactory();
185+
const toolkit = httpServiceMock.createOnPreAuthToolkit();
186+
187+
preAuthHandler(request, response, toolkit);
188+
expect(toolkit.next).not.toHaveBeenCalled();
189+
expect(response.customError).toHaveBeenCalledTimes(1);
190+
expect(response.customError).toHaveBeenCalledWith({
191+
statusCode: 429,
192+
body: 'Too Many Requests',
193+
});
194+
}
195+
196+
// request 1 succeeds
197+
makeRequestExpectNext();
198+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(1);
199+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(1);
200+
201+
// requests 2, 3, 4 fail
202+
makeRequestExpectError();
203+
makeRequestExpectError();
204+
makeRequestExpectError();
205+
206+
// requests 5+ succeed
207+
makeRequestExpectNext();
208+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(2);
209+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(2);
210+
211+
makeRequestExpectNext();
212+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(3);
213+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(3);
214+
215+
makeRequestExpectNext();
216+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(4);
217+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(4);
218+
});
219+
});

‎x-pack/plugins/ingest_manager/server/routes/limited_concurrency.ts‎

Lines changed: 31 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@ import {
1212
} from 'kibana/server';
1313
import { LIMITED_CONCURRENCY_ROUTE_TAG } from '../../common';
1414
import { IngestManagerConfigType } from '../index';
15-
class MaxCounter {
15+
16+
export class MaxCounter {
1617
constructor(private readonly max: number = 1) {}
1718
private counter = 0;
1819
valueOf() {
@@ -33,40 +34,56 @@ class MaxCounter {
3334
}
3435
}
3536

36-
function shouldHandleRequest(request: KibanaRequest) {
37+
export type IMaxCounter = Pick<MaxCounter, 'increase' | 'decrease' | 'lessThanMax'>;
38+
39+
export function isLimitedRoute(request: KibanaRequest) {
3740
const tags = request.route.options.tags;
38-
return tags.includes(LIMITED_CONCURRENCY_ROUTE_TAG);
41+
return !!tags.includes(LIMITED_CONCURRENCY_ROUTE_TAG);
3942
}
4043

41-
export function registerLimitedConcurrencyRoutes(core: CoreSetup, config: IngestManagerConfigType) {
42-
const max = config.fleet.maxConcurrentConnections;
43-
if (!max) return;
44-
45-
const counter = new MaxCounter(max);
46-
core.http.registerOnPreAuth(function preAuthHandler(
44+
export function createLimitedPreAuthHandler({
45+
isMatch,
46+
maxCounter,
47+
}: {
48+
isMatch: (request: KibanaRequest) => boolean;
49+
maxCounter: IMaxCounter;
50+
}) {
51+
return function preAuthHandler(
4752
request: KibanaRequest,
4853
response: LifecycleResponseFactory,
4954
toolkit: OnPreAuthToolkit
5055
) {
51-
if (!shouldHandleRequest(request)) {
56+
if (!isMatch(request)) {
5257
return toolkit.next();
5358
}
5459

55-
if (!counter.lessThanMax()) {
60+
if (!maxCounter.lessThanMax()) {
5661
return response.customError({
5762
body: 'Too Many Requests',
5863
statusCode: 429,
5964
});
6065
}
6166

62-
counter.increase();
67+
maxCounter.increase();
6368

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

7075
return toolkit.next();
71-
});
76+
};
77+
}
78+
79+
export function registerLimitedConcurrencyRoutes(core: CoreSetup, config: IngestManagerConfigType) {
80+
const max = config.fleet.maxConcurrentConnections;
81+
if (!max) return;
82+
83+
core.http.registerOnPreAuth(
84+
createLimitedPreAuthHandler({
85+
isMatch: isLimitedRoute,
86+
maxCounter: new MaxCounter(max),
87+
})
88+
);
7289
}

0 commit comments

Comments
 (0)