Proof-of-concept: Fleet agent concurrency

elastic · Jul 1, 2020 · d4f3755 · d4f3755
1 parent ad5ccfd
commit d4f3755
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 6 deletions.
diff --git a/src/core/server/http/http_server.ts b/src/core/server/http/http_server.ts
@@ -282,7 +282,7 @@ export class HttpServer {
       this.log.warn(`registerOnPreAuth called after stop`);
     }
 
-    this.server.ext('onRequest', adoptToHapiOnPreAuthFormat(fn, this.log));
+    this.server.ext('onPreAuth', adoptToHapiOnPreAuthFormat(fn, this.log));
   }
 
   private registerOnPreResponse(fn: OnPreResponseHandler) {

diff --git a/x-pack/plugins/ingest_manager/common/constants/agent.ts b/x-pack/plugins/ingest_manager/common/constants/agent.ts
@@ -16,3 +16,5 @@ export const AGENT_POLLING_THRESHOLD_MS = 30000;
 export const AGENT_POLLING_INTERVAL = 1000;
 export const AGENT_UPDATE_LAST_CHECKIN_INTERVAL_MS = 30000;
 export const AGENT_UPDATE_ACTIONS_INTERVAL_MS = 5000;
+
+export const AGENT_ROUTE_TAG = 'fleet:agent-route';
diff --git a/x-pack/plugins/ingest_manager/server/constants/index.ts b/x-pack/plugins/ingest_manager/server/constants/index.ts
@@ -26,6 +26,8 @@ export {
   SETUP_API_ROUTE,
   SETTINGS_API_ROUTES,
   APP_API_ROUTES,
+  // Route Tags
+  AGENT_ROUTE_TAG,
   // Saved object types
   AGENT_SAVED_OBJECT_TYPE,
   AGENT_EVENT_SAVED_OBJECT_TYPE,

diff --git a/x-pack/plugins/ingest_manager/server/plugin.ts b/x-pack/plugins/ingest_manager/server/plugin.ts
@@ -13,6 +13,11 @@ import {
   PluginInitializerContext,
   SavedObjectsServiceStart,
   HttpServiceSetup,
+  KibanaRequest,
+  LifecycleResponseFactory,
+  OnPreAuthToolkit,
+  OnPreResponseToolkit,
+  OnPreResponseInfo,
 } from 'kibana/server';
 import { LicensingPluginSetup, ILicense } from '../../licensing/server';
 import {
@@ -45,7 +50,7 @@ import {
   registerSettingsRoutes,
   registerAppRoutes,
 } from './routes';
-import { IngestManagerConfigType, NewDatasource } from '../common';
+import { IngestManagerConfigType, NewDatasource, AGENT_ROUTE_TAG } from '../common';
 import {
   appContextService,
   licenseService,
@@ -152,6 +157,35 @@ export class IngestManagerPlugin
   }
 
   public async setup(core: CoreSetup, deps: IngestManagerSetupDeps) {
+    const maxConcurrentRequests = 1;
+    let concurrentRequests = 0;
+    const isAgentRequest = (request: KibanaRequest) => {
+      const tags = request.route.options.tags;
+      return tags.includes(AGENT_ROUTE_TAG);
+    };
+    core.http.registerOnPreAuth(
+      (request: KibanaRequest, response: LifecycleResponseFactory, toolkit: OnPreAuthToolkit) => {
+        if (!isAgentRequest(request)) {
+          return toolkit.next();
+        }
+
+        if (concurrentRequests >= maxConcurrentRequests) {
+          return response.customError({ body: 'Too Many Agents', statusCode: 429 });
+        }
+
+        concurrentRequests += 1;
+        return toolkit.next();
+      }
+    );
+    core.http.registerOnPreResponse(
+      (request: KibanaRequest, preResponse: OnPreResponseInfo, toolkit: OnPreResponseToolkit) => {
+        if (isAgentRequest(request) && preResponse.statusCode !== 429) {
+          concurrentRequests -= 1;
+        }
+
+        return toolkit.next();
+      }
+    );
     this.httpSetup = core.http;
     this.licensing$ = deps.licensing.license$;
     if (deps.security) {

diff --git a/x-pack/plugins/ingest_manager/server/routes/agent/index.ts b/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, AGENT_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: [AGENT_ROUTE_TAG] },
     },
     postAgentCheckinHandler
   );
@@ -95,7 +95,7 @@ export const registerRoutes = (router: IRouter) => {
     {
       path: AGENT_API_ROUTES.ENROLL_PATTERN,
       validate: PostAgentEnrollRequestSchema,
-      options: { tags: [] },
+      options: { tags: [AGENT_ROUTE_TAG] },
     },
     postAgentEnrollHandler
   );
@@ -105,7 +105,7 @@ export const registerRoutes = (router: IRouter) => {
     {
       path: AGENT_API_ROUTES.ACKS_PATTERN,
       validate: PostAgentAcksRequestSchema,
-      options: { tags: [] },
+      options: { tags: [AGENT_ROUTE_TAG] },
     },
     postAgentAcksHandlerBuilder({
       acknowledgeAgentActions: AgentService.acknowledgeAgentActions,