Migrations: dynamically adjust batchSize when reading

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

@@ -16,6 +16,7 @@ const migrationSchema = schema.object({
   }),
   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: 536870888, max: 536870888 }),
   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/index.ts

@@ -147,6 +147,10 @@ export interface RequestEntityTooLargeException {
   type: 'request_entity_too_large_exception';
 }
 
+export interface EsResponseTooLargeError {
+  type: 'es_response_too_large';
+}
+
 /** @internal */
 export interface AcknowledgeResponse {
   acknowledged: boolean;
@@ -169,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/pickup_updated_mappings.test.ts

@@ -29,7 +29,7 @@ describe('pickupUpdatedMappings', () => {
     elasticsearchClientMock.createErrorTransportRequestPromise(retryableError)
   );
   it('calls catchRetryableEsClientErrors when the promise rejects', async () => {
-    const task = pickupUpdatedMappings(client, 'my_index');
+    const task = pickupUpdatedMappings(client, 'my_index', 1000);
     try {
       await task();
     } catch (e) {

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

@@ -13,7 +13,6 @@ import {
   catchRetryableEsClientErrors,
   type RetryableEsClientError,
 } from './catch_retryable_es_client_errors';
-import { BATCH_SIZE } from './constants';
 
 export interface UpdateByQueryResponse {
   taskId: string;
@@ -35,7 +34,8 @@ export interface UpdateByQueryResponse {
 export const pickupUpdatedMappings =
   (
     client: ElasticsearchClient,
-    index: string
+    index: string,
+    batchSize: number
   ): TaskEither.TaskEither<RetryableEsClientError, UpdateByQueryResponse> =>
   () => {
     return client
@@ -46,7 +46,7 @@ export const pickupUpdatedMappings =
         allow_no_indices: false,
         index,
         // How many documents to update per batch
-        scroll_size: BATCH_SIZE,
+        scroll_size: batchSize,
         // force a refresh so that we can query the updated index immediately
         // after the operation completes
         refresh: true,

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: { 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,17 @@ 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,
+          });
+        } else {
+          throw e;
+        }
+      })
       .catch(catchRetryableEsClientErrors);
   };

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

@@ -36,6 +36,7 @@ describe('updateAndPickupMappings', () => {
       client,
       index: 'new_index',
       mappings: { properties: {} },
+      batchSize: 1000,
     });
     try {
       await task();
@@ -65,6 +66,7 @@ describe('updateAndPickupMappings', () => {
           },
         },
       },
+      batchSize: 1000,
     });
     try {
       await task();

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

@@ -28,6 +28,7 @@ export interface UpdateAndPickupMappingsParams {
   client: ElasticsearchClient;
   index: string;
   mappings: IndexMapping;
+  batchSize: number;
 }
 /**
  * Updates an index's mappings and runs an pickupUpdatedMappings task so that the mapping
@@ -37,6 +38,7 @@ export const updateAndPickupMappings = ({
   client,
   index,
   mappings,
+  batchSize,
 }: UpdateAndPickupMappingsParams): TaskEither.TaskEither<
   RetryableEsClientError,
   UpdateAndPickupMappingsResponse
@@ -74,7 +76,7 @@ export const updateAndPickupMappings = ({
   return pipe(
     putMappingTask,
     TaskEither.chain((res) => {
-      return pickupUpdatedMappings(client, index);
+      return pickupUpdatedMappings(client, index, batchSize);
     })
   );
 };