Use updated onPreAuth from Platform

Author: John Schulz

x-pack/plugins/ingest_manager/common/constants/routes.ts

@@ -11,6 +11,8 @@ export const PACKAGE_CONFIG_API_ROOT = `${API_ROOT}/package_configs`;
 export const AGENT_CONFIG_API_ROOT = `${API_ROOT}/agent_configs`;
 export const FLEET_API_ROOT = `${API_ROOT}/fleet`;
 
+export const LIMITED_CONCURRENCY_ROUTE_TAG = 'ingest:limited-concurrency';
+
 // EPM API routes
 const EPM_PACKAGES_MANY = `${EPM_API_ROOT}/packages`;
 const EPM_PACKAGES_ONE = `${EPM_PACKAGES_MANY}/{pkgkey}`;

x-pack/plugins/ingest_manager/server/constants/index.ts

@@ -15,6 +15,7 @@ export {
   AGENT_UPDATE_ACTIONS_INTERVAL_MS,
   INDEX_PATTERN_PLACEHOLDER_SUFFIX,
   // Routes
+  LIMITED_CONCURRENCY_ROUTE_TAG,
   PLUGIN_ID,
   EPM_API_ROUTES,
   DATA_STREAM_API_ROUTES,

x-pack/plugins/ingest_manager/server/plugin.ts

@@ -34,6 +34,7 @@ import {
 } from './constants';
 import { registerSavedObjects, registerEncryptedSavedObjects } from './saved_objects';
 import {
+  preAuthHandler,
   registerEPMRoutes,
   registerPackageConfigRoutes,
   registerDataStreamRoutes,
@@ -231,6 +232,9 @@ export class IngestManagerPlugin
             );
           }
         } else {
+          // we currently only use this global interceptor if fleet is enabled
+          // since it would run this func on *every* req (other plugins, CSS, etc)
+          this.httpSetup.registerOnPreAuth(preAuthHandler);
           registerAgentRoutes(router);
           registerEnrollmentApiKeyRoutes(router);
           registerInstallScriptRoutes({

x-pack/plugins/ingest_manager/server/routes/agent/index.ts

@@ -10,7 +10,7 @@
  */
 
 import { IRouter } from 'src/core/server';
-import { PLUGIN_ID, AGENT_API_ROUTES } from '../../constants';
+import { PLUGIN_ID, AGENT_API_ROUTES, LIMITED_CONCURRENCY_ROUTE_TAG } from '../../constants';
 import {
   GetAgentsRequestSchema,
   GetOneAgentRequestSchema,
@@ -85,7 +85,7 @@ export const registerRoutes = (router: IRouter) => {
     {
       path: AGENT_API_ROUTES.CHECKIN_PATTERN,
       validate: PostAgentCheckinRequestSchema,
-      options: { tags: [] },
+      options: { tags: [LIMITED_CONCURRENCY_ROUTE_TAG] },
     },
     postAgentCheckinHandler
   );
@@ -95,7 +95,7 @@ export const registerRoutes = (router: IRouter) => {
     {
       path: AGENT_API_ROUTES.ENROLL_PATTERN,
       validate: PostAgentEnrollRequestSchema,
-      options: { tags: [] },
+      options: { tags: [LIMITED_CONCURRENCY_ROUTE_TAG] },
     },
     postAgentEnrollHandler
   );
@@ -105,7 +105,7 @@ export const registerRoutes = (router: IRouter) => {
     {
       path: AGENT_API_ROUTES.ACKS_PATTERN,
       validate: PostAgentAcksRequestSchema,
-      options: { tags: [] },
+      options: { tags: [LIMITED_CONCURRENCY_ROUTE_TAG] },
     },
     postAgentAcksHandlerBuilder({
       acknowledgeAgentActions: AgentService.acknowledgeAgentActions,

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

@@ -0,0 +1,63 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { KibanaRequest, LifecycleResponseFactory, OnPreAuthToolkit } from 'kibana/server';
+import { LIMITED_CONCURRENCY_ROUTE_TAG } from '../../common';
+
+class MaxCounter {
+  constructor(private readonly max: number = 1) {}
+  private counter = 0;
+  valueOf() {
+    return this.counter;
+  }
+  increase() {
+    if (this.counter < this.max) {
+      this.counter += 1;
+    }
+  }
+  decrease() {
+    this.counter += 1;
+  }
+  lessThanMax() {
+    return this.counter < this.max;
+  }
+}
+
+function shouldHandleRequest(request: KibanaRequest) {
+  const tags = request.route.options.tags;
+  return tags.includes(LIMITED_CONCURRENCY_ROUTE_TAG);
+}
+
+const LIMITED_CONCURRENCY_MAX_REQUESTS = 250;
+const counter = new MaxCounter(LIMITED_CONCURRENCY_MAX_REQUESTS);
+
+export function preAuthHandler(
+  request: KibanaRequest,
+  response: LifecycleResponseFactory,
+  toolkit: OnPreAuthToolkit
+) {
+  if (!shouldHandleRequest(request)) {
+    return toolkit.next();
+  }
+
+  if (!counter.lessThanMax()) {
+    return response.customError({
+      body: 'Too Many Agents',
+      statusCode: 503,
+      headers: {
+        'Retry-After': '30',
+      },
+    });
+  }
+
+  counter.increase();
+
+  // requests.events.aborted$ has a bug where it's fired even when the request completes...
+  // we can take advantage of this bug just for load testing...
+  request.events.aborted$.toPromise().then(() => counter.decrease());
+
+  return toolkit.next();
+}

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

@@ -14,3 +14,4 @@ export { registerRoutes as registerInstallScriptRoutes } from './install_script'
 export { registerRoutes as registerOutputRoutes } from './output';
 export { registerRoutes as registerSettingsRoutes } from './settings';
 export { registerRoutes as registerAppRoutes } from './app';
+export { preAuthHandler } from './global_interceptors';