Migrations: dynamically adjust batchSize when reading

packages/core/saved-objects/core-saved-objects-base-server-internal/src/saved_objects_config.ts

@@ -9,13 +9,18 @@
 import { valid } from 'semver';
 import { schema, TypeOf } from '@kbn/config-schema';
 import type { ServiceConfigDescriptor } from '@kbn/core-base-server-internal';
+import buffer from 'buffer';
 
 const migrationSchema = schema.object({
   algorithm: schema.oneOf([schema.literal('v2'), schema.literal('zdt')], {
     defaultValue: 'v2',
   }),
   batchSize: schema.number({ defaultValue: 1_000 }),
   maxBatchSizeBytes: schema.byteSize({ defaultValue: '100mb' }), // 100mb is the default http.max_content_length Elasticsearch config value
+  maxReadBatchSizeBytes: schema.byteSize({
+    defaultValue: buffer.constants.MAX_STRING_LENGTH,
+    max: buffer.constants.MAX_STRING_LENGTH,
+  }),
   discardUnknownObjects: schema.maybe(
     schema.string({
       validate: (value: string) =>

packages/core/saved-objects/core-saved-objects-migration-server-internal/index.ts

@@ -47,6 +47,7 @@ export type {
   ReindexResponse,
   UpdateByQueryResponse,
   UpdateAndPickupMappingsResponse,
+  EsResponseTooLargeError,
 } from './src/actions';
 export {
   isClusterShardLimitExceeded,

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/actions/catch_retryable_es_client_errors.test.ts

@@ -72,7 +72,7 @@ describe('catchRetryableEsClientErrors', () => {
       });
     });
     it('ResponseError with retryable status code', async () => {
-      const statusCodes = [503, 401, 403, 408, 410];
+      const statusCodes = [503, 401, 403, 408, 410, 429];
       return Promise.all(
         statusCodes.map(async (status) => {
           const error = new esErrors.ResponseError(

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/actions/catch_retryable_es_client_errors.ts

@@ -15,6 +15,7 @@ const retryResponseStatuses = [
   403, // AuthenticationException
   408, // RequestTimeout
   410, // Gone
+  429, // TooManyRequests -> ES circuit breaker
 ];
 
 export interface RetryableEsClientError {

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/actions/index.ts

@@ -146,6 +146,11 @@ export interface RequestEntityTooLargeException {
   type: 'request_entity_too_large_exception';
 }
 
+export interface EsResponseTooLargeError {
+  type: 'es_response_too_large';
+  contentLength: number;
+}
+
 /** @internal */
 export interface AcknowledgeResponse {
   acknowledged: boolean;
@@ -168,6 +173,7 @@ export interface ActionErrorTypeMap {
   index_not_green_timeout: IndexNotGreenTimeout;
   index_not_yellow_timeout: IndexNotYellowTimeout;
   cluster_shard_limit_exceeded: ClusterShardLimitExceeded;
+  es_response_too_large: EsResponseTooLargeError;
 }
 
 /**

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/actions/read_with_pit.test.ts

@@ -32,23 +32,54 @@ describe('readWithPit', () => {
       pitId: 'pitId',
       query: { match_all: {} },
       batchSize: 10_000,
+      maxResponseSizeBytes: 100_000,
     })();
 
     expect(client.search).toHaveBeenCalledTimes(1);
-    expect(client.search).toHaveBeenCalledWith({
-      allow_partial_search_results: false,
-      pit: {
-        id: 'pitId',
-        keep_alive: '10m',
-      },
-      query: {
-        match_all: {},
+    expect(client.search).toHaveBeenCalledWith(
+      {
+        allow_partial_search_results: false,
+        pit: {
+          id: 'pitId',
+          keep_alive: '10m',
+        },
+        query: {
+          match_all: {},
+        },
+        search_after: undefined,
+        seq_no_primary_term: undefined,
+        size: 10000,
+        sort: '_shard_doc:asc',
+        track_total_hits: true,
       },
-      search_after: undefined,
-      seq_no_primary_term: undefined,
-      size: 10000,
-      sort: '_shard_doc:asc',
-      track_total_hits: true,
+      { maxResponseSize: 100_000 }
+    );
+  });
+
+  it('returns left es_response_too_large when client throws RequestAbortedError', async () => {
+    // Create a mock client that rejects all methods with a RequestAbortedError
+    // response.
+    const retryableError = new EsErrors.RequestAbortedError(
+      'The content length (536870889) is bigger than the maximum allow string (536870888)'
+    );
+    const client = elasticsearchClientMock.createInternalClient(
+      elasticsearchClientMock.createErrorTransportRequestPromise(retryableError)
+    );
+
+    const task = readWithPit({
+      client,
+      pitId: 'pitId',
+      query: { match_all: {} },
+      batchSize: 10_000,
+    });
+    try {
+      await task();
+    } catch (e) {
+      /** ignore */
+    }
+    await expect(task()).resolves.toEqual({
+      _tag: 'Left',
+      left: { contentLength: 536870889, type: 'es_response_too_large' },
     });
   });
 

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/actions/read_with_pit.ts

@@ -9,13 +9,15 @@
 import * as Either from 'fp-ts/lib/Either';
 import * as TaskEither from 'fp-ts/lib/TaskEither';
 import * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import { errors as EsErrors } from '@elastic/elasticsearch';
 import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
 import type { SavedObjectsRawDoc } from '@kbn/core-saved-objects-server';
 import {
   catchRetryableEsClientErrors,
   type RetryableEsClientError,
 } from './catch_retryable_es_client_errors';
 import { DEFAULT_PIT_KEEP_ALIVE } from './open_pit';
+import { EsResponseTooLargeError } from '.';
 
 /** @internal */
 export interface ReadWithPit {
@@ -32,6 +34,7 @@ export interface ReadWithPitParams {
   batchSize: number;
   searchAfter?: number[];
   seqNoPrimaryTerm?: boolean;
+  maxResponseSizeBytes?: number;
 }
 
 /*
@@ -45,32 +48,39 @@ export const readWithPit =
     batchSize,
     searchAfter,
     seqNoPrimaryTerm,
-  }: ReadWithPitParams): TaskEither.TaskEither<RetryableEsClientError, ReadWithPit> =>
+    maxResponseSizeBytes,
+  }: ReadWithPitParams): TaskEither.TaskEither<
+    RetryableEsClientError | EsResponseTooLargeError,
+    ReadWithPit
+  > =>
   () => {
     return client
-      .search<SavedObjectsRawDoc>({
-        seq_no_primary_term: seqNoPrimaryTerm,
-        // Fail if the index being searched doesn't exist or is closed
-        // allow_no_indices: false,
-        // By default ES returns a 200 with partial results if there are shard
-        // request timeouts or shard failures which can lead to data loss for
-        // migrations
-        allow_partial_search_results: false,
-        // Sort fields are required to use searchAfter so we sort by the
-        // natural order of the index which is the most efficient option
-        // as order is not important for the migration
-        sort: '_shard_doc:asc',
-        pit: { id: pitId, keep_alive: DEFAULT_PIT_KEEP_ALIVE },
-        size: batchSize,
-        search_after: searchAfter,
-        /**
-         * We want to know how many documents we need to process so we can log the progress.
-         * But we also want to increase the performance of these requests,
-         * so we ask ES to report the total count only on the first request (when searchAfter does not exist)
-         */
-        track_total_hits: typeof searchAfter === 'undefined',
-        query,
-      })
+      .search<SavedObjectsRawDoc>(
+        {
+          seq_no_primary_term: seqNoPrimaryTerm,
+          // Fail if the index being searched doesn't exist or is closed
+          // allow_no_indices: false,
+          // By default ES returns a 200 with partial results if there are shard
+          // request timeouts or shard failures which can lead to data loss for
+          // migrations
+          allow_partial_search_results: false,
+          // Sort fields are required to use searchAfter so we sort by the
+          // natural order of the index which is the most efficient option
+          // as order is not important for the migration
+          sort: '_shard_doc:asc',
+          pit: { id: pitId, keep_alive: DEFAULT_PIT_KEEP_ALIVE },
+          size: batchSize,
+          search_after: searchAfter,
+          /**
+           * We want to know how many documents we need to process so we can log the progress.
+           * But we also want to increase the performance of these requests,
+           * so we ask ES to report the total count only on the first request (when searchAfter does not exist)
+           */
+          track_total_hits: typeof searchAfter === 'undefined',
+          query,
+        },
+        { maxResponseSize: maxResponseSizeBytes }
+      )
       .then((body) => {
         const totalHits =
           typeof body.hits.total === 'number'
@@ -93,5 +103,22 @@ export const readWithPit =
           totalHits,
         });
       })
+      .catch((e) => {
+        if (
+          e instanceof EsErrors.RequestAbortedError &&
+          e.message.match(/The content length \(\d+\) is bigger than the maximum/) != null
+        ) {
+          return Either.left({
+            type: 'es_response_too_large' as const,
+            contentLength: Number.parseInt(
+              e.message.match(/The content length \((\d+)\) is bigger than the maximum/)?.[1] ??
+                '-1',
+              10
+            ),
+          });
+        } else {
+          throw e;
+        }
+      })
       .catch(catchRetryableEsClientErrors);
   };

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/initial_state.test.ts

@@ -27,6 +27,7 @@ const migrationsConfig = {
   retryAttempts: 15,
   batchSize: 1000,
   maxBatchSizeBytes: ByteSizeValue.parse('100mb'),
+  maxReadBatchSizeBytes: ByteSizeValue.parse('500mb'),
 } as unknown as SavedObjectsMigrationConfigType;
 
 const createInitialStateCommonParams = {
@@ -217,7 +218,9 @@ describe('createInitialState', () => {
         "knownTypes": Array [],
         "legacyIndex": ".kibana_task_manager",
         "logs": Array [],
+        "maxBatchSize": 1000,
         "maxBatchSizeBytes": 104857600,
+        "maxReadBatchSizeBytes": 524288000,
         "migrationDocLinks": Object {
           "clusterShardLimitExceeded": "https://www.elastic.co/guide/en/kibana/test-branch/resolve-migrations-failures.html#cluster-shard-limit-exceeded",
           "repeatedTimeoutRequests": "https://www.elastic.co/guide/en/kibana/test-branch/resolve-migrations-failures.html#_repeated_time_out_requests_that_eventually_fail",

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/initial_state.ts

@@ -126,7 +126,9 @@ export const createInitialState = ({
     retryDelay: 0,
     retryAttempts: migrationsConfig.retryAttempts,
     batchSize: migrationsConfig.batchSize,
+    maxBatchSize: migrationsConfig.batchSize,
     maxBatchSizeBytes: migrationsConfig.maxBatchSizeBytes.getValueInBytes(),
+    maxReadBatchSizeBytes: migrationsConfig.maxReadBatchSizeBytes.getValueInBytes(),
     discardUnknownObjects: migrationsConfig.discardUnknownObjects === kibanaVersion,
     discardCorruptObjects: migrationsConfig.discardCorruptObjects === kibanaVersion,
     logs: [],

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/kibana_migrator.test.ts

@@ -243,9 +243,9 @@ describe('KibanaMigrator', () => {
       const migrator = new KibanaMigrator(options);
       migrator.prepareMigrations();
       await expect(migrator.runMigrations()).rejects.toMatchInlineSnapshot(`
-              [Error: Unable to complete saved object migrations for the [.my-index] index. Error: Reindex failed with the following error:
-              {"_tag":"Some","value":{"type":"elasticsearch_exception","reason":"task failed with an error"}}]
-            `);
+        [Error: Unable to complete saved object migrations for the [.my-index] index. Error: Reindex failed with the following error:
+        {"_tag":"Some","value":{"type":"elasticsearch_exception","reason":"task failed with an error"}}]
+      `);
       expect(loggingSystemMock.collect(options.logger).error[0][0]).toMatchInlineSnapshot(`
         [Error: Reindex failed with the following error:
         {"_tag":"Some","value":{"type":"elasticsearch_exception","reason":"task failed with an error"}}]
@@ -533,6 +533,7 @@ const mockOptions = () => {
       algorithm: 'v2',
       batchSize: 20,
       maxBatchSizeBytes: ByteSizeValue.parse('20mb'),
+      maxReadBatchSizeBytes: new ByteSizeValue(536870888),
       pollInterval: 20000,
       scrollDuration: '10m',
       skip: false,

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/migrations_state_action_machine.test.ts

@@ -51,6 +51,7 @@ describe('migrationsStateActionMachine', () => {
       algorithm: 'v2',
       batchSize: 1000,
       maxBatchSizeBytes: new ByteSizeValue(1e8),
+      maxReadBatchSizeBytes: new ByteSizeValue(536870888),
       pollInterval: 0,
       scrollDuration: '0s',
       skip: false,

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/model/helpers.ts

@@ -17,6 +17,7 @@ import type { SavedObjectsRawDoc } from '@kbn/core-saved-objects-server';
 import type { IndexMapping } from '@kbn/core-saved-objects-base-server-internal';
 import type { AliasAction, FetchIndexResponse } from '../actions';
 import type { BulkIndexOperationTuple } from './create_batches';
+import { OutdatedDocumentsSearchRead, ReindexSourceToTempRead } from '../state';
 
 /** @internal */
 export type Aliases = Partial<Record<string, string>>;
@@ -285,3 +286,11 @@ export function getMigrationType({
  */
 export const getTempIndexName = (indexPrefix: string, kibanaVersion: string): string =>
   `${indexPrefix}_${kibanaVersion}_reindex_temp`;
+
+/** Increase batchSize by 20% until a maximum of maxBatchSize */
+export const increaseBatchSize = (
+  stateP: OutdatedDocumentsSearchRead | ReindexSourceToTempRead
+) => {
+  const increasedBatchSize = Math.floor(stateP.batchSize * 1.2);
+  return increasedBatchSize > stateP.maxBatchSize ? stateP.maxBatchSize : increasedBatchSize;
+};