handle ScaleError

Author: npalm

lambdas/functions/control-plane/src/aws/runners.ts

@@ -205,7 +205,6 @@ async function processFleetResult(
       instances.push(...instancesOnDemand);
       return instances;
     } else if ((failedCount = countScaleErrors(errors, scaleErrors)) > 0) {
-
       logger.warn('Create fleet failed, ScaleError will be thrown to trigger retry for ephemeral runners.');
       logger.debug('Create fleet failed.', { data: fleet.Errors });
       throw new ScaleError('Failed to create instance, create fleet failed.', failedCount);
@@ -218,7 +217,7 @@ async function processFleetResult(
 }
 
 function countScaleErrors(errors: string[], scaleErrors: string[]): number {
-  return errors.reduce((acc, e) => scaleErrors.includes(e) ? acc + 1 : acc, 0);
+  return errors.reduce((acc, e) => (scaleErrors.includes(e) ? acc + 1 : acc), 0);
 }
 
 async function getAmiIdOverride(runnerParameters: Runners.RunnerInputParameters): Promise<string | undefined> {

lambdas/functions/control-plane/src/lambda.test.ts

@@ -28,11 +28,11 @@ const sqsRecord: SQSRecord = {
   },
   awsRegion: '',
   body: JSON.stringify(body),
-  eventSource: 'aws:SQS',
+  eventSource: 'aws:sqs',
   eventSourceARN: '',
   md5OfBody: '',
   messageAttributes: {},
-  messageId: '',
+  messageId: 'abcd1234',
   receiptHandle: '',
 };
 
@@ -109,14 +109,16 @@ describe('Test scale up lambda wrapper.', () => {
     await expect(scaleUpHandler(sqsEvent, context)).resolves.not.toThrow();
   });
 
-  it('Scale should be rejected', async () => {
+  it('Scale should create a batch failure message', async () => {
     const error = new ScaleError('Scale should be rejected');
     const mock = vi.fn() as MockedFunction<typeof scaleUp>;
     mock.mockImplementation(() => {
       return Promise.reject(error);
     });
     vi.mocked(scaleUp).mockImplementation(mock);
-    await expect(scaleUpHandler(sqsEvent, context)).rejects.toThrow(error);
+    await expect(scaleUpHandler(sqsEvent, context)).resolves.toEqual({
+      batchItemFailures: [{ itemIdentifier: sqsRecord.messageId }],
+    });
   });
 
   describe('Batch processing', () => {
@@ -240,12 +242,14 @@ describe('Test scale up lambda wrapper.', () => {
       const records = createMultipleRecords(2);
       const multiRecordEvent: SQSEvent = { Records: records };
 
-      const error = new ScaleError('Critical scaling error');
+      const error = new ScaleError('Critical scaling error', 2);
       const mock = vi.fn(scaleUp);
       mock.mockImplementation(() => Promise.reject(error));
       vi.mocked(scaleUp).mockImplementation(mock);
 
-      await expect(scaleUpHandler(multiRecordEvent, context)).rejects.toThrow(error);
+      await expect(scaleUpHandler(multiRecordEvent, context)).resolves.toEqual({
+        batchItemFailures: [{ itemIdentifier: 'message-0' }, { itemIdentifier: 'message-1' }],
+      });
     });
   });
 });

lambdas/functions/control-plane/src/lambda.ts

@@ -62,10 +62,21 @@ export async function scaleUpHandler(event: SQSEvent, context: Context): Promise
     return { batchItemFailures };
   } catch (e) {
     if (e instanceof ScaleError) {
-      throw e;
+      for (let i = 0; i < e.failedInstanceCount; i++) {
+        batchItemFailures.push({
+          itemIdentifier: sqsMessages[i].messageId,
+        });
+      }
+      logger.warn(
+        `ScaleError detected, ${e.failedInstanceCount} could not be created. A retry will be attempted via SQS.`,
+        { error: e },
+      );
+    } else {
+      logger.error(`Error processing batch (size: ${sqsMessages.length}): ${(e as Error).message}, ignoring batch`, {
+        error: e,
+      });
     }
 
-    logger.warn(`Will retry error: ${e}`);
     return { batchItemFailures };
   }
 }

lambdas/functions/control-plane/src/scale-runners/ScaleError.ts

@@ -1,7 +1,7 @@
 class ScaleError extends Error {
   constructor(
     public message: string,
-    public failedInstanceCount: number = 1
+    public failedInstanceCount: number = 1,
   ) {
     super(message);
     this.name = 'ScaleError';