[Fleet] Improving bulk actions for more than 10k agents

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

@@ -119,6 +119,7 @@ export const AGENT_API_ROUTES = {
   UPGRADE_PATTERN: `${API_ROOT}/agents/{agentId}/upgrade`,
   BULK_UPGRADE_PATTERN: `${API_ROOT}/agents/bulk_upgrade`,
   CURRENT_UPGRADES_PATTERN: `${API_ROOT}/agents/current_upgrades`,
+  INTERNAL_LIST_PATTERN: `${INTERNAL_ROOT}/agents`, // for testing
 };
 
 export const ENROLLMENT_API_KEY_ROUTES = {

x-pack/plugins/fleet/common/types/models/agent.ts

@@ -87,6 +87,7 @@ export interface Agent extends AgentBase {
   access_api_key?: string;
   status?: AgentStatus;
   packages: string[];
+  sort?: Array<number | string | null>;
 }
 
 export interface AgentSOAttributes extends AgentBase {

x-pack/plugins/fleet/server/routes/agent/handlers.ts

@@ -54,6 +54,32 @@ export const getAgentHandler: RequestHandler<
   }
 };
 
+export const getAllAgentsHandler: RequestHandler<
+  undefined,
+  TypeOf<typeof GetAgentsRequestSchema.query>
+> = async (context, request, response) => {
+  const coreContext = await context.core;
+  const esClient = coreContext.elasticsearch.client.asInternalUser;
+
+  try {
+    const agents = await AgentService.getAgents(esClient, {
+      perPage: request.query.perPage,
+      kuery: request.query.kuery!,
+    });
+
+    const body: GetAgentsResponse = {
+      items: agents.slice(0, 10),
+      total: agents.length,
+      totalInactive: 0,
+      page: 1,
+      perPage: 0,
+    };
+    return response.ok({ body });
+  } catch (error) {
+    return defaultIngestErrorHandler({ error, response });
+  }
+};
+
 export const deleteAgentHandler: RequestHandler<
   TypeOf<typeof DeleteAgentRequestSchema.params>
 > = async (context, request, response) => {

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

@@ -35,6 +35,7 @@ import {
   putAgentsReassignHandler,
   postBulkAgentsReassignHandler,
   getAgentDataHandler,
+  getAllAgentsHandler,
 } from './handlers';
 import {
   postNewAgentActionHandlerBuilder,
@@ -93,6 +94,18 @@ export const registerAPIRoutes = (router: FleetAuthzRouter, config: FleetConfigT
     getAgentsHandler
   );
 
+  // For testing
+  router.get(
+    {
+      path: AGENT_API_ROUTES.INTERNAL_LIST_PATTERN,
+      validate: GetAgentsRequestSchema,
+      fleetAuthz: {
+        fleet: { all: true },
+      },
+    },
+    getAllAgentsHandler
+  );
+
   // Agent actions
   router.post(
     {

x-pack/plugins/fleet/server/services/agents/crud.ts

@@ -7,6 +7,7 @@
 
 import Boom from '@hapi/boom';
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import type { SortResults } from '@elastic/elasticsearch/lib/api/types';
 import type { SavedObjectsClientContract, ElasticsearchClient } from '@kbn/core/server';
 
 import type { KueryNode } from '@kbn/es-query';
@@ -68,6 +69,7 @@ export type GetAgentsOptions =
   | {
       kuery: string;
       showInactive?: boolean;
+      perPage?: number;
     };
 
 export async function getAgents(esClient: ElasticsearchClient, options: GetAgentsOptions) {
@@ -79,6 +81,7 @@ export async function getAgents(esClient: ElasticsearchClient, options: GetAgent
       await getAllAgentsByKuery(esClient, {
         kuery: options.kuery,
         showInactive: options.showInactive ?? false,
+        perPage: options.perPage,
       })
     ).agents;
   } else {
@@ -90,6 +93,108 @@ export async function getAgents(esClient: ElasticsearchClient, options: GetAgent
   return agents;
 }
 
+export async function getAgentsByKueryPit(
+  esClient: ElasticsearchClient,
+  options: ListWithKuery & {
+    showInactive: boolean;
+    pitId: string;
+    searchAfter?: SortResults;
+  }
+): Promise<{
+  agents: Agent[];
+  total: number;
+  page: number;
+  perPage: number;
+}> {
+  const {
+    page = 1,
+    perPage = 20,
+    sortField = 'enrolled_at',
+    sortOrder = 'desc',
+    kuery,
+    showInactive = false,
+    showUpgradeable,
+    searchAfter,
+  } = options;
+  const { pitId } = options;
+  const filters = [];
+
+  if (kuery && kuery !== '') {
+    filters.push(kuery);
+  }
+
+  if (showInactive === false) {
+    filters.push(ACTIVE_AGENT_CONDITION);
+  }
+
+  const kueryNode = _joinFilters(filters);
+  const body = kueryNode ? { query: toElasticsearchQuery(kueryNode) } : {};
+
+  const queryAgents = async (from: number, size: number) => {
+    return esClient.search<FleetServerAgent, {}>({
+      from,
+      size,
+      track_total_hits: true,
+      rest_total_hits_as_int: true,
+      body: {
+        ...body,
+        sort: [{ [sortField]: { order: sortOrder } }],
+      },
+      pit: {
+        id: pitId,
+        keep_alive: '1m',
+      },
+      ...(searchAfter ? { search_after: searchAfter, from: 0 } : {}),
+    });
+  };
+
+  const res = await queryAgents((page - 1) * perPage, perPage);
+
+  let agents = res.hits.hits.map(searchHitToAgent);
+  let total = res.hits.total as number;
+
+  // filtering for a range on the version string will not work,
+  // nor does filtering on a flattened field (local_metadata), so filter here
+  if (showUpgradeable) {
+    // fixing a bug where upgradeable filter was not returning right results https://github.com/elastic/kibana/issues/117329
+    // query all agents, then filter upgradeable, and return the requested page and correct total
+    // if there are more than SO_SEARCH_LIMIT agents, the logic falls back to same as before
+    if (total < SO_SEARCH_LIMIT) {
+      const response = await queryAgents(0, SO_SEARCH_LIMIT);
+      agents = response.hits.hits
+        .map(searchHitToAgent)
+        .filter((agent) => isAgentUpgradeable(agent, appContextService.getKibanaVersion()));
+      total = agents.length;
+      const start = (page - 1) * perPage;
+      agents = agents.slice(start, start + perPage);
+    } else {
+      agents = agents.filter((agent) =>
+        isAgentUpgradeable(agent, appContextService.getKibanaVersion())
+      );
+    }
+  }
+
+  return {
+    agents,
+    total,
+    page,
+    perPage,
+  };
+}
+
+export async function openAgentsPointInTime(esClient: ElasticsearchClient): Promise<string> {
+  const pitKeepAlive = '10m';
+  const pitRes = await esClient.openPointInTime({
+    index: AGENTS_INDEX,
+    keep_alive: pitKeepAlive,
+  });
+  return pitRes.id;
+}
+
+export async function closeAgentsPointInTime(esClient: ElasticsearchClient, pitId: string) {
+  await esClient.closePointInTime({ id: pitId });
+}
+
 export async function getAgentsByKuery(
   esClient: ElasticsearchClient,
   options: ListWithKuery & {
@@ -168,19 +273,92 @@ export async function getAgentsByKuery(
   };
 }
 
+export async function processAgentsInBatches(
+  esClient: ElasticsearchClient,
+  options: Omit<ListWithKuery, 'page'> & {
+    showInactive: boolean;
+  },
+  processAgents: (
+    agents: Agent[],
+    includeSuccess: boolean
+  ) => Promise<{ items: BulkActionResult[] }>
+): Promise<{ items: BulkActionResult[] }> {
+  const pitId = await openAgentsPointInTime(esClient);
+
+  const perPage = options.perPage ?? SO_SEARCH_LIMIT;
+
+  const res = await getAgentsByKueryPit(esClient, {
+    ...options,
+    page: 1,
+    perPage,
+    pitId,
+  });
+
+  let currentAgents = res.agents;
+  // include successful agents if total agents does not exceed 10k
+  const skipSuccess = res.total > SO_SEARCH_LIMIT;
+
+  let results = await processAgents(currentAgents, skipSuccess);
+  let allAgentsProcessed = currentAgents.length;
+
+  while (allAgentsProcessed < res.total) {
+    const lastAgent = currentAgents[currentAgents.length - 1];
+    const nextPage = await getAgentsByKueryPit(esClient, {
+      ...options,
+      page: 1,
+      perPage,
+      pitId,
+      searchAfter: lastAgent.sort!,
+    });
+    currentAgents = nextPage.agents;
+    const currentResults = await processAgents(currentAgents, skipSuccess);
+    results = { items: results.items.concat(currentResults.items) };
+    allAgentsProcessed += currentAgents.length;
+  }
+
+  await closeAgentsPointInTime(esClient, pitId);
+
+  return results;
+}
+
 export async function getAllAgentsByKuery(
   esClient: ElasticsearchClient,
-  options: Omit<ListWithKuery, 'page' | 'perPage'> & {
+  options: Omit<ListWithKuery, 'page'> & {
     showInactive: boolean;
   }
 ): Promise<{
   agents: Agent[];
   total: number;
 }> {
-  const res = await getAgentsByKuery(esClient, { ...options, page: 1, perPage: SO_SEARCH_LIMIT });
+  const pitId = await openAgentsPointInTime(esClient);
+
+  const perPage = options.perPage ?? SO_SEARCH_LIMIT;
+
+  const res = await getAgentsByKueryPit(esClient, {
+    ...options,
+    page: 1,
+    perPage,
+    pitId,
+  });
+
+  let allAgents = res.agents;
+
+  while (allAgents.length < res.total) {
+    const lastAgent = allAgents[allAgents.length - 1];
+    const nextPage = await getAgentsByKueryPit(esClient, {
+      ...options,
+      page: 1,
+      perPage,
+      pitId,
+      searchAfter: lastAgent.sort!,
+    });
+    allAgents = allAgents.concat(nextPage.agents);
+  }
+
+  await closeAgentsPointInTime(esClient, pitId);
 
   return {
-    agents: res.agents,
+    agents: allAgents,
     total: res.total,
   };
 }

x-pack/plugins/fleet/server/services/agents/helpers.ts

@@ -6,7 +6,7 @@
  */
 
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
-
+import type { SortResults } from '@elastic/elasticsearch/lib/api/types';
 import type { SearchHit } from '@kbn/core/types/elasticsearch';
 
 import type { Agent, AgentSOAttributes, FleetServerAgent } from '../../types';
@@ -17,7 +17,7 @@ type FleetServerAgentESResponse =
   | estypes.SearchResponse<FleetServerAgent>['hits']['hits'][0]
   | SearchHit<FleetServerAgent>;
 
-export function searchHitToAgent(hit: FleetServerAgentESResponse): Agent {
+export function searchHitToAgent(hit: FleetServerAgentESResponse & { sort?: SortResults }): Agent {
   // @ts-expect-error @elastic/elasticsearch MultiGetHit._source is optional
   const agent: Agent = {
     id: hit._id,
@@ -26,6 +26,7 @@ export function searchHitToAgent(hit: FleetServerAgentESResponse): Agent {
     access_api_key: undefined,
     status: undefined,
     packages: hit._source?.packages ?? [],
+    sort: hit.sort,
   };
 
   agent.status = getAgentStatus(agent);