Increase the retry timeout for scaleUp by using a stochastic overshoot to overcome SQSs limitations (#4051)

jeanschmidt · web-flow · commit 8b591dcf9307 · 2023-04-20T12:26:44.000+02:00
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/lambda.test.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/lambda.test.ts
@@ -96,6 +96,52 @@ describe('scaleUp', () => {
     expect(sqsDeleteMessageBatch).toBeCalledWith(metrics, evts);
   });
 
+  it('stochasticOvershoot when retryCount > 5', async () => {
+    const config = {
+      maxRetryScaleUpRecord: 1999,
+      retryScaleUpRecordDelayS: 20,
+      retryScaleUpRecordJitterPct: 0.2,
+      retryScaleUpRecordQueueUrl: 'asdf',
+    };
+    jest.spyOn(Config, 'Instance', 'get').mockImplementation(() => config as unknown as Config);
+    const records = [
+      { eventSource: 'aws:sqs', body: '{"id":1}', eventSourceARN: '1:2:3:4:5:6', receiptHandle: 'xxx', messageId: 1 },
+      {
+        eventSource: 'aws:sqs',
+        body: '{"id":2,"retryCount":3}',
+        eventSourceARN: '1:2:3:4:5:6:7',
+        receiptHandle: 'xxx',
+        messageId: 2,
+      },
+      {
+        eventSource: 'aws:sqs',
+        body: '{"id":3,"retryCount":12}',
+        eventSourceARN: '1:2:3:4:5:6:7',
+        receiptHandle: 'xxx',
+        messageId: 3,
+      },
+    ];
+
+    const mockedScaleUp = mocked(scaleUp).mockResolvedValue(undefined);
+    jest.spyOn(global.Math, 'random').mockReturnValue(0.5);
+
+    const callback = jest.fn();
+    await scaleUpL({ Records: records } as unknown as SQSEvent, {} as unknown as Context, callback);
+    expect(mockedScaleUp).toBeCalledTimes(2);
+
+    const expected = [
+      {
+        id: 3,
+        retryCount: 12,
+        delaySeconds: 900,
+      },
+    ];
+    expect(sqsSendMessages).toBeCalledTimes(1);
+    expect(sqsSendMessages).toBeCalledWith(metrics, expected, 'asdf');
+
+    expect(sqsDeleteMessageBatch).toBeCalledTimes(0);
+  });
+
   it('RetryableScalingError', async () => {
     const config = {
       maxRetryScaleUpRecord: 12,
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/lambda.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/lambda.ts
@@ -3,31 +3,35 @@ import { Context, SQSEvent, SQSRecord, ScheduledEvent } from 'aws-lambda';
 
 import { Config } from './scale-runners/config';
 import { ScaleUpMetrics, sendMetricsAtTimeout, sendMetricsTimeoutVars } from './scale-runners/metrics';
-import { getDelayWithJitterRetryCount } from './scale-runners/utils';
+import { getDelayWithJitterRetryCount, stochaticRunOvershoot } from './scale-runners/utils';
 import { scaleDown as scaleDownR } from './scale-runners/scale-down';
 import { sqsSendMessages, sqsDeleteMessageBatch } from './scale-runners/sqs';
 
-async function sendRetryEvents(evtFailed: Array<[SQSRecord, boolean]>, metrics: ScaleUpMetrics) {
+async function sendRetryEvents(evtFailed: Array<[SQSRecord, boolean, number]>, metrics: ScaleUpMetrics) {
   console.error(`Detected ${evtFailed.length} errors when processing messages, will retry relevant messages.`);
   metrics.exception();
   const messagesToSend: Array<ActionRequestMessage> = [];
   console.dir(evtFailed);
 
-  for (const [evt, retryable] of evtFailed) {
+  for (const [evt, retryable, incRetry] of evtFailed) {
     const body: ActionRequestMessage = JSON.parse(evt.body);
     const retryCount = body?.retryCount ?? 0;
 
     if (
       retryCount < Config.Instance.maxRetryScaleUpRecord &&
       (Config.Instance.retryScaleUpRecordQueueUrl?.length ?? 0) > 0
     ) {
-      if (retryable) {
-        metrics.scaleUpFailureRetryable(retryCount);
+      if (incRetry > 0) {
+        if (retryable) {
+          metrics.scaleUpFailureRetryable(retryCount);
+        } else {
+          metrics.scaleUpFailureNonRetryable(retryCount);
+        }
       } else {
-        metrics.scaleUpFailureNonRetryable(retryCount);
+        metrics.stochasticOvershoot();
       }
 
-      body.retryCount = retryCount + 1;
+      body.retryCount = retryCount + incRetry;
       body.delaySeconds = Math.min(
         900,
         getDelayWithJitterRetryCount(
@@ -64,24 +68,32 @@ export async function scaleUp(event: SQSEvent, context: Context, callback: any)
     (Config.Instance.lambdaTimeout - 10) * 1000,
   );
 
-  const evtFailed: Array<[SQSRecord, boolean]> = [];
+  const evtFailed: Array<[SQSRecord, boolean, number]> = [];
 
   try {
     recordsIterProcess: for (let i = 0; i < event.Records.length; i += 1) {
       const evt = event.Records[i];
 
       try {
-        await scaleUpR(evt.eventSource, JSON.parse(evt.body), metrics);
-        metrics.scaleUpSuccess();
+        const body = JSON.parse(evt.body) as ActionRequestMessage;
+        if (
+          !stochaticRunOvershoot(body?.retryCount ?? 0, 900, Math.max(Config.Instance.retryScaleUpRecordDelayS, 20))
+        ) {
+          console.warn(`Skipping message due to stochatic run overshoot: ${evt.body}`);
+          evtFailed.push([evt, true, 0]);
+        } else {
+          await scaleUpR(evt.eventSource, body, metrics);
+          metrics.scaleUpSuccess();
+        }
       } catch (e) {
         if (e instanceof RetryableScalingError) {
           console.error(`Retryable error thrown: "${e.message}"`);
-          evtFailed.push([evt, true]);
+          evtFailed.push([evt, true, 1]);
         } else {
           console.error(`Non-retryable error during request: "${e.message}"`);
           console.error(`All remaning '${event.Records.length - i}' messages will be scheduled to retry`);
           for (let ii = i; ii < event.Records.length; ii += 1) {
-            evtFailed.push([event.Records[ii], false]);
+            evtFailed.push([event.Records[ii], false, 1]);
           }
           break recordsIterProcess;
         }
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/metrics.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/metrics.ts
@@ -769,6 +769,11 @@ export class ScaleUpMetrics extends Metrics {
     this.countEntry('run.scaleup.success');
   }
 
+  /* istanbul ignore next */
+  stochasticOvershoot() {
+    this.countEntry('run.scaleup.stochasticOvershoot');
+  }
+
   /* istanbul ignore next */
   scaleUpFailureRetryable(retries: number) {
     this.countEntry('run.scaleup.failure.total.count');
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/utils.test.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/utils.test.ts
@@ -1,12 +1,14 @@
 import {
   expBackOff,
   getBoolean,
+  getDelay,
   getDelayWithJitter,
   getDelayWithJitterRetryCount,
   getRepo,
   getRepoKey,
   groupBy,
   shuffleArrayInPlace,
+  stochaticRunOvershoot,
 } from './utils';
 import nock from 'nock';
 
@@ -140,6 +142,20 @@ describe('./utils', () => {
     });
   });
 
+  describe('getDelay', () => {
+    it('test some numbers', () => {
+      expect(getDelay(0, 3)).toEqual(3);
+      expect(getDelay(1, 3)).toEqual(6);
+      expect(getDelay(2, 3)).toEqual(12);
+      expect(getDelay(3, 3)).toEqual(24);
+
+      expect(getDelay(0, 5)).toEqual(5);
+      expect(getDelay(1, 5)).toEqual(10);
+      expect(getDelay(2, 5)).toEqual(20);
+      expect(getDelay(3, 5)).toEqual(40);
+    });
+  });
+
   describe('getDelayWithJitter', () => {
     it('have jitter == 0', () => {
       expect(getDelayWithJitter(20, 0.0)).toEqual(20);
@@ -259,4 +275,41 @@ describe('shuffleArrayInPlace', () => {
       expect(arr).toContain(number);
     }
   });
+
+  describe('stochaticRunOvershoot', () => {
+    afterEach(() => {
+      jest.spyOn(global.Math, 'random').mockRestore();
+    });
+
+    it('test some values', () => {
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.0);
+      expect(stochaticRunOvershoot(0, 100, 10)).toBe(true);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.5);
+      expect(stochaticRunOvershoot(0, 100, 10)).toBe(true);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(1.0);
+      expect(stochaticRunOvershoot(0, 100, 10)).toBe(true);
+
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.0);
+      expect(stochaticRunOvershoot(4, 100, 10)).toBe(true);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.5);
+      expect(stochaticRunOvershoot(4, 100, 10)).toBe(true);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.6249);
+      expect(stochaticRunOvershoot(4, 100, 10)).toBe(true);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.6251);
+      expect(stochaticRunOvershoot(4, 100, 10)).toBe(false);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(1.0);
+      expect(stochaticRunOvershoot(4, 100, 10)).toBe(false);
+
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.0);
+      expect(stochaticRunOvershoot(12, 100, 10)).toBe(true);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.00244139);
+      expect(stochaticRunOvershoot(12, 100, 10)).toBe(true);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.00244141);
+      expect(stochaticRunOvershoot(12, 100, 10)).toBe(false);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(0.5);
+      expect(stochaticRunOvershoot(12, 100, 10)).toBe(false);
+      jest.spyOn(global.Math, 'random').mockReturnValueOnce(1.0);
+      expect(stochaticRunOvershoot(12, 100, 10)).toBe(false);
+    });
+  });
 });
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/utils.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/utils.ts
@@ -119,8 +119,17 @@ export function getDelayWithJitter(delayBase: number, jitter: number) {
   return Math.max(0, delayBase) * (1 + Math.random() * Math.max(0, jitter));
 }
 
+export function getDelay(retryCount: number, delayBase: number) {
+  return Math.max(0, delayBase) * Math.pow(2, Math.max(0, retryCount));
+}
+
 export function getDelayWithJitterRetryCount(retryCount: number, delayBase: number, jitter: number) {
-  return getDelayWithJitter(Math.max(0, delayBase) * Math.pow(2, Math.max(0, retryCount)), jitter);
+  return getDelayWithJitter(getDelay(retryCount, delayBase), jitter);
+}
+
+export function stochaticRunOvershoot(retryCount: number, maxTime: number, delayBase: number) {
+  const prob = maxTime / getDelay(retryCount, delayBase);
+  return Math.random() < prob;
 }
 
 /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
diff --git a/terraform-aws-github-runner/modules/runners/scale-up.tf b/terraform-aws-github-runner/modules/runners/scale-up.tf
@@ -46,11 +46,11 @@ resource "aws_lambda_function" "scale_up" {
       LAUNCH_TEMPLATE_VERSION_LINUX         = aws_launch_template.linux_runner.latest_version
       LAUNCH_TEMPLATE_VERSION_LINUX_NVIDIA  = aws_launch_template.linux_runner_nvidia.latest_version
       LAUNCH_TEMPLATE_VERSION_WINDOWS       = aws_launch_template.windows_runner.latest_version
-      MAX_RETRY_SCALEUP_RECORD              = "5"
+      MAX_RETRY_SCALEUP_RECORD              = "10"
       MUST_HAVE_ISSUES_LABELS               = join(",", var.must_have_issues_labels)
       REDIS_ENDPOINT                        = var.redis_endpoint
       REDIS_LOGIN                           = var.redis_login
-      RETRY_SCALE_UP_RECORD_DELAY_S         = "110"
+      RETRY_SCALE_UP_RECORD_DELAY_S         = "60"
       RETRY_SCALE_UP_RECORD_JITTER_PCT      = "0.5"
       RETRY_SCALE_UP_RECORD_QUEUE_URL       = var.sqs_build_queue_retry.url
       RUNNER_EXTRA_LABELS                   = var.runner_extra_labels
