[Fleet] Run agent policy schema in batches during fleet setup + add xpack.fleet.setup.agentPolicySchemaUpgradeBatchSize config

x-pack/plugins/fleet/common/types/index.ts

@@ -36,6 +36,9 @@ export interface FleetConfigType {
   packageVerification?: {
     gpgKeyPath?: string;
   };
+  setup?: {
+    agentPolicySchemaUpgradeBatchSize?: number;
+  };
   developer?: {
     disableRegistryVersionCheck?: boolean;
     bundledPackageLocation?: string;

x-pack/plugins/fleet/scripts/create_agents/create_agents.ts

@@ -25,6 +25,8 @@ const printUsage = () =>
     [--kibana]: full url of kibana instance to create agents and policy in e.g http://localhost:5601/mybase, defaults to http://localhost:5601
     [--username]: username for kibana, defaults to elastic
     [--password]: password for kibana, defaults to changeme
+    [--batches]: run the script in batches, defaults to 1 e.g if count is 50 and batches is 10, 500 agents will be created and 10 agent policies
+    [--concurrentBatches]: how many batches to run concurrently, defaults to 10
 `);
 
 const DEFAULT_KIBANA_URL = 'http://localhost:5601';
@@ -49,6 +51,8 @@ const {
   agentVersion: agentVersionArg,
   username: kbnUsername = DEFAULT_KIBANA_USERNAME,
   password: kbnPassword = DEFAULT_KIBANA_PASSWORD,
+  batches: batchesArg,
+  concurrentBatches: concurrentBatchesArg = 10,
   // ignore yargs positional args, we only care about named args
   _,
   $0,
@@ -59,6 +63,8 @@ const statusesArg = (statusArg as string).split(',') as AgentStatus[];
 const inactivityTimeout = inactivityTimeoutArg
   ? Number(inactivityTimeoutArg).valueOf()
   : DEFAULT_UNENROLL_TIMEOUT;
+const batches = inactivityTimeoutArg ? Number(batchesArg).valueOf() : 1;
+const concurrentBatches = concurrentBatchesArg ? Number(concurrentBatchesArg).valueOf() : 10;
 const count = countArg ? Number(countArg).valueOf() : DEFAULT_AGENT_COUNT;
 const kbnAuth = 'Basic ' + Buffer.from(kbnUsername + ':' + kbnPassword).toString('base64');
 
@@ -258,7 +264,7 @@ async function createAgentPolicy(id: string) {
       name: id,
       namespace: 'default',
       description: '',
-      monitoring_enabled: ['logs'],
+      monitoring_enabled: ['logs', 'metrics'],
       inactivity_timeout: inactivityTimeout,
     }),
     headers: {
@@ -314,23 +320,40 @@ export async function run() {
     logger.info(`Deleted ${deleteRes.deleted} agents, took ${deleteRes.took}ms`);
   }
 
-  logger.info('Creating agent policy');
-
-  const agentPolicyId = 'script-create-agent-' + uuidv4();
-  const agentPolicy = await createAgentPolicy(agentPolicyId);
-  logger.info(`Created agent policy ${agentPolicy.item.id}`);
-
   logger.info('Creating fleet superuser');
   const { role, user } = await createSuperUser();
   logger.info(`Role "${ES_SUPERUSER}" ${role.role.created ? 'created' : 'already exists'}`);
   logger.info(`User "${ES_SUPERUSER}" ${user.created ? 'created' : 'already exists'}`);
 
-  logger.info('Creating agent documents');
-  const statusMap = statusesArg.reduce((acc, status) => ({ ...acc, [status]: count }), {});
-  logStatusMap(statusMap);
-  const agents = createAgentsWithStatuses(statusMap, agentPolicyId, agentVersion);
-  const createRes = await createAgentDocsBulk(agents);
-  logger.info(
-    `Created ${createRes.items.length} agent docs, took ${createRes.took}, errors: ${createRes.errors}`
-  );
+  let batchesRemaining = batches;
+  let totalAgents = 0;
+  while (batchesRemaining > 0) {
+    const currentBatchSize = Math.min(concurrentBatches, batchesRemaining);
+    if (batches > 1) {
+      logger.info(`Running ${currentBatchSize} batches. ${batchesRemaining} batches remaining`);
+    }
+
+    await Promise.all(
+      Array(currentBatchSize)
+        .fill(0)
+        .map(async (__, i) => {
+          const agentPolicyId = 'script-create-agent-' + uuidv4();
+          const agentPolicy = await createAgentPolicy(agentPolicyId);
+          logger.info(`Created agent policy ${agentPolicy.item.id}`);
+
+          const statusMap = statusesArg.reduce((acc, status) => ({ ...acc, [status]: count }), {});
+          logStatusMap(statusMap);
+          const agents = createAgentsWithStatuses(statusMap, agentPolicyId, agentVersion);
+          const createRes = await createAgentDocsBulk(agents);
+          logger.info(
+            `Batch complete, created ${createRes.items.length} agent docs, took ${createRes.took}, errors: ${createRes.errors}`
+          );
+          totalAgents += createRes.items.length;
+        })
+    );
+
+    batchesRemaining -= currentBatchSize;
+  }
+
+  logger.info(`All batches complete. Created ${totalAgents} agents in total. Goodbye!`);
 }

x-pack/plugins/fleet/server/config.ts

@@ -120,6 +120,11 @@ export const config: PluginConfigDescriptor = {
     fleetServerHosts: PreconfiguredFleetServerHostsSchema,
     proxies: PreconfiguredFleetProxiesSchema,
     agentIdVerificationEnabled: schema.boolean({ defaultValue: true }),
+    setup: schema.maybe(
+      schema.object({
+        agentPolicySchemaUpgradeBatchSize: schema.maybe(schema.number()),
+      })
+    ),
     developer: schema.object({
       disableRegistryVersionCheck: schema.boolean({ defaultValue: false }),
       allowAgentUpgradeSourceUri: schema.boolean({ defaultValue: false }),

x-pack/plugins/fleet/server/services/agent_policy.ts

@@ -315,15 +315,22 @@ class AgentPolicyService {
     soClient: SavedObjectsClientContract,
     options: ListWithKuery & {
       withPackagePolicies?: boolean;
+      fields?: string[];
     }
-  ): Promise<{ items: AgentPolicy[]; total: number; page: number; perPage: number }> {
+  ): Promise<{
+    items: AgentPolicy[];
+    total: number;
+    page: number;
+    perPage: number;
+  }> {
     const {
       page = 1,
       perPage = 20,
       sortField = 'updated_at',
       sortOrder = 'desc',
       kuery,
       withPackagePolicies = false,
+      fields,
     } = options;
 
     const baseFindParams = {
@@ -332,6 +339,7 @@ class AgentPolicyService {
       sortOrder,
       page,
       perPage,
+      ...(fields ? { fields } : {}),
     };
     const filter = kuery ? normalizeKuery(SAVED_OBJECT_TYPE, kuery) : undefined;
     let agentPoliciesSO;

x-pack/plugins/fleet/server/services/setup/upgrade_agent_policy_schema_version.ts

@@ -10,14 +10,16 @@ import type { SavedObjectsClientContract } from '@kbn/core/server';
 import {
   AGENT_POLICY_SAVED_OBJECT_TYPE,
   FLEET_AGENT_POLICIES_SCHEMA_VERSION,
-  SO_SEARCH_LIMIT,
 } from '../../constants';
 import { agentPolicyService } from '../agent_policy';
+import { appContextService } from '../app_context';
 
-function getOutdatedAgentPoliciesBatch(soClient: SavedObjectsClientContract) {
+const DEFAULT_BATCH_SIZE = 100;
+function getOutdatedAgentPoliciesBatch(soClient: SavedObjectsClientContract, batchSize: number) {
   return agentPolicyService.list(soClient, {
-    perPage: SO_SEARCH_LIMIT,
+    perPage: batchSize,
     kuery: `NOT ${AGENT_POLICY_SAVED_OBJECT_TYPE}.schema_version:${FLEET_AGENT_POLICIES_SCHEMA_VERSION}`,
+    fields: ['id'], // we only need the ID of the agent policy
   });
 }
 
@@ -26,13 +28,23 @@ function getOutdatedAgentPoliciesBatch(soClient: SavedObjectsClientContract) {
 // deploy outdated policies to .fleet-policies index
 // bump oudated SOs schema_version
 export async function upgradeAgentPolicySchemaVersion(soClient: SavedObjectsClientContract) {
-  let outdatedAgentPolicies = await getOutdatedAgentPoliciesBatch(soClient);
+  const config = appContextService.getConfig();
+  const logger = appContextService.getLogger();
 
+  const batchSize = config?.setup?.agentPolicySchemaUpgradeBatchSize ?? DEFAULT_BATCH_SIZE;
+  let outdatedAgentPolicies = await getOutdatedAgentPoliciesBatch(soClient, batchSize);
+  logger.debug(`Found ${outdatedAgentPolicies.total} outdated agent policies`);
   while (outdatedAgentPolicies.total > 0) {
+    const start = Date.now();
     const outdatedAgentPolicyIds = outdatedAgentPolicies.items.map(
       (outdatedAgentPolicy) => outdatedAgentPolicy.id
     );
     await agentPolicyService.deployPolicies(soClient, outdatedAgentPolicyIds);
-    outdatedAgentPolicies = await getOutdatedAgentPoliciesBatch(soClient);
+    outdatedAgentPolicies = await getOutdatedAgentPoliciesBatch(soClient, batchSize);
+    logger.debug(
+      `Upgraded ${outdatedAgentPolicyIds.length} agent policies in ${Date.now() - start}ms, ${
+        outdatedAgentPolicies.total
+      } remaining`
+    );
   }
 }