fix: calculate push page size based on monitor size per batch (#993)

* retry on payload

* add paylaod based chunking

* add paylaod based chunking

* Fix unit tests

* conflicts resolution 101

* update tests

* PR feedback

* PR feedback

* PR feedback

* show only for 413

* keep size as mb and kb

---------

Co-authored-by: emilioalvap <[email protected]>
Co-authored-by: vigneshshanmugam <[email protected]>

Author: 3 people

__tests__/push/utils.test.ts

@@ -23,7 +23,7 @@
  *
  */
 
-import { normalizeMonitorName } from '../../src/push/utils';
+import { getSizedBatches, normalizeMonitorName } from '../../src/push/utils';
 
 describe('Push Utils', () => {
   it("normalize monitor's name", () => {
@@ -38,4 +38,42 @@ describe('Push Utils', () => {
     expect(normalizeMonitorName('x.y.z')).toBe('x_y_z');
     expect(normalizeMonitorName('foo:bar:blah')).toBe('foo_bar_blah');
   });
+
+  describe('getSizedChunks', () => {
+    it('should return empty array when input is empty', () => {
+      expect(getSizedBatches([], new Map(), 100, 100)).toEqual([]);
+    });
+
+    it('should return chunks of input array based on maxChunkSizeKB', () => {
+      const input = [
+        { id: '1' },
+        { id: '2' },
+        { id: '3' },
+        { id: '4' },
+        { id: '5' },
+        { id: '6' },
+      ];
+      // map of id and sizes
+      const sizes = new Map([
+        ['1', 10 * 1024],
+        ['2', 20 * 1024],
+        ['3', 30 * 1024],
+        ['4', 40 * 1024],
+        ['5', 50 * 1024],
+        ['6', 60 * 1024],
+      ]);
+      expect(getSizedBatches(input, sizes, 1000, 100)).toEqual([input]);
+      expect(getSizedBatches(input, sizes, 100, 3)).toEqual([
+        [input[0], input[1], input[2]],
+        [input[3], input[4]],
+        [input[5]],
+      ]);
+      expect(getSizedBatches(input, sizes, 100, 2)).toEqual([
+        [input[0], input[1]],
+        [input[2], input[3]],
+        [input[4]],
+        [input[5]],
+      ]);
+    });
+  });
 });

src/push/index.ts

@@ -50,10 +50,12 @@ import {
   bulkGetMonitors,
   bulkPutMonitors,
   createMonitorsLegacy,
-  CHUNK_SIZE,
+  BATCH_SIZE,
+  MAX_PAYLOAD_SIZE_KB,
 } from './kibana_api';
 import {
-  getChunks,
+  getBatches,
+  getSizedBatches,
   isBulkAPISupported,
   isLightweightMonitorSupported,
   logDiff,
@@ -102,7 +104,10 @@ export async function push(monitors: Monitor[], options: PushOptions) {
     for (const value of sizes.values()) {
       totalSize += value;
     }
-    progress('total size of the monitors payload is ' + printBytes(totalSize));
+    progress(
+      `total size of the ${sizes.size} monitors payload is ` +
+        printBytes(totalSize)
+    );
   }
 
   if (options.dryRun) {
@@ -115,11 +120,19 @@ export async function push(monitors: Monitor[], options: PushOptions) {
   const updatedMonitors = new Set<string>([...changedIDs, ...newIDs]);
   if (updatedMonitors.size > 0) {
     const updatedMonSchemas = schemas.filter(s => updatedMonitors.has(s.id));
-    const chunks = getChunks(updatedMonSchemas, CHUNK_SIZE);
-    for (const chunk of chunks) {
+    const batches = getSizedBatches(
+      updatedMonSchemas,
+      sizes,
+      MAX_PAYLOAD_SIZE_KB,
+      BATCH_SIZE
+    );
+    if (batches.length > 1) {
+      progress(`Monitors will be pushed as ${batches.length} batches.`);
+    }
+    for (const batch of batches) {
       await liveProgress(
-        bulkPutMonitors(options, chunk),
-        `creating or updating ${chunk.length} monitors`
+        bulkPutMonitors(options, batch),
+        `creating or updating ${batch.length} monitors`
       );
     }
   }
@@ -136,11 +149,11 @@ export async function push(monitors: Monitor[], options: PushOptions) {
         options.yes
       );
     }
-    const chunks = getChunks(Array.from(removedIDs), CHUNK_SIZE);
-    for (const chunk of chunks) {
+    const batches = getBatches(Array.from(removedIDs), BATCH_SIZE);
+    for (const batch of batches) {
       await liveProgress(
-        bulkDeleteMonitors(options, chunk),
-        `deleting ${chunk.length} monitors`
+        bulkDeleteMonitors(options, batch),
+        `deleting ${batch.length} monitors`
       );
     }
   }
@@ -317,14 +330,18 @@ export async function pushLegacy(monitors: Monitor[], options: PushOptions) {
   }
 
   let schemas: MonitorSchema[] = [];
+  let sizes: Map<string, number> = new Map();
   if (monitors.length > 0) {
     progress(`preparing ${monitors.length} monitors`);
-    ({ schemas } = await buildMonitorSchema(monitors, false));
-    const chunks = getChunks(schemas, 10);
-    for (const chunk of chunks) {
+    ({ schemas, sizes } = await buildMonitorSchema(monitors, false));
+    const batches = getSizedBatches(schemas, sizes, MAX_PAYLOAD_SIZE_KB, 10);
+    if (batches.length > 1) {
+      progress(`Monitors will be pushed as ${batches.length} batches.`);
+    }
+    for (const batch of batches) {
       await liveProgress(
-        createMonitorsLegacy({ schemas: chunk, keepStale: true, options }),
-        `creating or updating ${chunk.length} monitors`
+        createMonitorsLegacy({ schemas: batch, keepStale: true, options }),
+        `creating or updating ${batch.length} monitors`
       );
     }
   } else {

src/push/kibana_api.ts

@@ -36,8 +36,9 @@ import {
 } from './request';
 import { generateURL } from './utils';
 
-// Default chunk size for bulk put / delete
-export const CHUNK_SIZE = parseInt(process.env.CHUNK_SIZE) || 100;
+// Default batch size for bulk put / delete
+export const BATCH_SIZE = parseInt(process.env.CHUNK_SIZE) || 250;
+export const MAX_PAYLOAD_SIZE_KB = 50 * 1000; // 50 MB in KB
 
 export type PutResponse = {
   createdMonitors: string[];
@@ -49,12 +50,15 @@ export async function bulkPutMonitors(
   options: PushOptions,
   schemas: MonitorSchema[]
 ) {
+  const url = generateURL(options, 'bulk_update') + '/_bulk_update';
+
   const resp = await sendReqAndHandleError<PutResponse>({
-    url: generateURL(options, 'bulk_update') + '/_bulk_update',
+    url,
     method: 'PUT',
     auth: options.auth,
     body: JSON.stringify({ monitors: schemas }),
   });
+
   const { failedMonitors } = resp;
   if (failedMonitors && failedMonitors.length > 0) {
     throw formatFailedMonitors(failedMonitors);

src/push/monitor.ts

@@ -50,6 +50,7 @@ export type MonitorSchema = Omit<MonitorConfig, 'locations'> & {
   content?: string;
   filter?: Monitor['filter'];
   hash?: string;
+  size?: number;
 };
 
 // Abbreviated monitor info, as often returned by the API,
@@ -300,7 +301,7 @@ export function buildMonitorFromYaml(
   });
 
   /**
-   * Params support is only available for lighweight monitors
+   * Params support is only available for lightweight monitors
    * post 8.7.2 stack
    */
   if (isParamOptionSupported(options.kibanaVersion)) {

src/push/request.ts

@@ -58,8 +58,14 @@ export async function sendReqAndHandleError<T>(
   options: APIRequestOptions
 ): Promise<T> {
   const { statusCode, body } = await sendRequest(options);
+
   return (
-    await handleError(statusCode, options.url, body)
+    await handleError(
+      statusCode,
+      options.url,
+      body,
+      statusCode === 413 ? `${options.body?.length} bytes sent` : ''
+    )
   ).json() as Promise<T>;
 }
 
@@ -74,7 +80,8 @@ type APIError = {
 export async function handleError(
   statusCode: number,
   url: string,
-  body: Dispatcher.ResponseData['body']
+  body: Dispatcher.ResponseData['body'],
+  extraMessage?: string
 ): Promise<Dispatcher.ResponseData['body']> {
   if (statusCode === 404) {
     throw formatNotFoundError(url, await body.text());
@@ -84,9 +91,19 @@ export async function handleError(
       const resp = await body.text();
       parsed = JSON.parse(resp) as APIError;
     } catch (e) {
-      throw formatAPIError(statusCode, 'unexpected error', e.message);
+      throw formatAPIError(
+        statusCode,
+        'unexpected error',
+        e.message,
+        extraMessage
+      );
     }
-    throw formatAPIError(statusCode, parsed.error, parsed.message);
+    throw formatAPIError(
+      statusCode,
+      parsed.error,
+      parsed.message,
+      extraMessage
+    );
   }
 
   return body;
@@ -111,16 +128,18 @@ export function formatNotFoundError(url: string, message: string) {
 }
 
 export function formatAPIError(
-  statuCode: number,
+  statusCode: number,
   error: string,
-  message: string
+  message: string,
+  extraMessage = ''
 ) {
   let outer = bold(`${symbols['failed']} Error\n`);
   let inner = bold(
-    `${symbols['failed']} monitor creation failed - ${statuCode}:${error}\n`
+    `${symbols['failed']} monitor creation failed - ${statusCode}:${error}\n`
   );
   inner += indent(message, '    ');
   outer += indent(inner);
+  outer += extraMessage ? indent(extraMessage, '    ') : '';
   return red(outer);
 }
 

src/push/utils.ts

@@ -90,12 +90,44 @@ export function printBytes(bytes: number) {
   return `${bytes.toFixed(1)} ${BYTE_UNITS[exponent]}`;
 }
 
-export function getChunks<T>(arr: Array<T>, size: number): Array<T[]> {
-  const chunks = [];
+export function getBatches(arr: string[], size: number): string[][] {
+  const batches = [];
   for (let i = 0; i < arr.length; i += size) {
-    chunks.push(arr.slice(i, i + size));
+    batches.push(arr.slice(i, i + size));
   }
-  return chunks;
+  return batches;
+}
+
+export function getSizedBatches<T extends { id?: string }>(
+  arr: Array<T>,
+  sizes: Map<string, number>,
+  maxBatchSizeKB: number,
+  maxBatchItems: number
+): Array<T[]> {
+  const batches: Array<T[]> = [];
+  let currentBatch: T[] = [];
+  let currentSize = 0;
+
+  for (const item of arr) {
+    const sizeKB = item.id ? Math.round(sizes.get(item.id) / 1000) : 1;
+    // If adding the current item would exceed limits, create a new chunk
+    if (
+      currentBatch.length >= maxBatchItems ||
+      currentSize + sizeKB > maxBatchSizeKB
+    ) {
+      batches.push(currentBatch);
+      currentBatch = [];
+      currentSize = 0;
+    }
+    currentBatch.push(item);
+    currentSize += sizeKB;
+  }
+
+  // Push the last chunk if it contains any items
+  if (currentBatch.length > 0) {
+    batches.push(currentBatch);
+  }
+  return batches;
 }
 
 type Operation =