[Ingest Manager] Add more Fleet concurrency tests #71744

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

@@ -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', () => {
@@ -33,3 +37,186 @@ describe('registerLimitedConcurrencyRoutes', () => {
     expect(mockSetup.http.registerOnPreAuth).toHaveBeenCalledTimes(1);
   });
 });
+
+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();
+
+    // @ts-ignore error re: mockPreAuthToolkit return type
+    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();
+
+    // @ts-ignore error re: mockPreAuthToolkit return type
+    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();
+
+    // @ts-ignore error re: mockPreAuthToolkit return type
+    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();
+
+    // @ts-ignore error re: mockPreAuthToolkit return type
+    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();
+
+      // @ts-ignore error re: mockPreAuthToolkit return type
+      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();
+
+      // @ts-ignore error re: mockPreAuthToolkit return type
+      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);
+  });
+});

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

@@ -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() {
@@ -33,40 +34,56 @@ 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();
+      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),
+    })
+  );
 }