fix: job retry mechanism not triggering

lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts

@@ -10,6 +10,7 @@ import { createRunner, listEC2Runners } from './../aws/runners';
 import { RunnerInputParameters } from './../aws/runners.d';
 import * as scaleUpModule from './scale-up';
 import { getParameter } from '@aws-github-runner/aws-ssm-util';
+import { publishRetryMessage } from './job-retry';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 import type { Octokit } from '@octokit/rest';
 
@@ -33,6 +34,7 @@ const mockCreateRunner = vi.mocked(createRunner);
 const mockListRunners = vi.mocked(listEC2Runners);
 const mockSSMClient = mockClient(SSMClient);
 const mockSSMgetParameter = vi.mocked(getParameter);
+const mockPublishRetryMessage = vi.mocked(publishRetryMessage);
 
 vi.mock('@octokit/rest', () => ({
   Octokit: vi.fn().mockImplementation(function () {
@@ -63,6 +65,11 @@ vi.mock('@aws-github-runner/aws-ssm-util', async () => {
   };
 });
 
+vi.mock('./job-retry', () => ({
+  publishRetryMessage: vi.fn(),
+  checkAndRetryJob: vi.fn(),
+}));
+
 export type RunnerType = 'ephemeral' | 'non-ephemeral';
 
 // for ephemeral and non-ephemeral runners
@@ -1680,6 +1687,171 @@ describe('scaleUp with Github Data Residency', () => {
   });
 });
 
+describe('Retry mechanism tests', () => {
+  beforeEach(() => {
+    process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
+    process.env.ENABLE_EPHEMERAL_RUNNERS = 'true';
+    process.env.ENABLE_JOB_QUEUED_CHECK = 'true';
+    process.env.RUNNERS_MAXIMUM_COUNT = '10';
+    expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS };
+    mockSSMClient.reset();
+  });
+
+  const createTestMessages = (
+    count: number,
+    overrides: Partial<scaleUpModule.ActionRequestMessageSQS>[] = [],
+  ): scaleUpModule.ActionRequestMessageSQS[] => {
+    return Array.from({ length: count }, (_, i) => ({
+      ...TEST_DATA_SINGLE,
+      id: i + 1,
+      messageId: `message-${i + 1}`,
+      ...overrides[i],
+    }));
+  };
+
+  it('calls publishRetryMessage for each valid message when job is queued', async () => {
+    const messages = createTestMessages(3);
+    mockCreateRunner.mockResolvedValue(['i-12345', 'i-67890', 'i-abcdef']); // Create all requested runners
+
+    await scaleUpModule.scaleUp(messages);
+
+    expect(mockPublishRetryMessage).toHaveBeenCalledTimes(3);
+    expect(mockPublishRetryMessage).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({
+        id: 1,
+        messageId: 'message-1',
+      }),
+    );
+    expect(mockPublishRetryMessage).toHaveBeenNthCalledWith(
+      2,
+      expect.objectContaining({
+        id: 2,
+        messageId: 'message-2',
+      }),
+    );
+    expect(mockPublishRetryMessage).toHaveBeenNthCalledWith(
+      3,
+      expect.objectContaining({
+        id: 3,
+        messageId: 'message-3',
+      }),
+    );
+  });
+
+  it('does not call publishRetryMessage when job is not queued', async () => {
+    mockOctokit.actions.getJobForWorkflowRun.mockImplementation((params) => {
+      const isQueued = params.job_id === 1; // Only job 1 is queued
+      return {
+        data: {
+          status: isQueued ? 'queued' : 'completed',
+        },
+      };
+    });
+
+    const messages = createTestMessages(3);
+
+    await scaleUpModule.scaleUp(messages);
+
+    // Only message with id 1 should trigger retry
+    expect(mockPublishRetryMessage).toHaveBeenCalledTimes(1);
+    expect(mockPublishRetryMessage).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: 1,
+        messageId: 'message-1',
+      }),
+    );
+  });
+
+  it('does not call publishRetryMessage when maximum runners is reached and messages are marked invalid', async () => {
+    process.env.RUNNERS_MAXIMUM_COUNT = '0'; // No runners can be created
+
+    const messages = createTestMessages(2);
+
+    await scaleUpModule.scaleUp(messages);
+
+    // Verify listEC2Runners is called to check current runner count
+    expect(listEC2Runners).toHaveBeenCalledWith({
+      environment: 'unit-test-environment',
+      runnerType: 'Org',
+      runnerOwner: TEST_DATA_SINGLE.repositoryOwner,
+    });
+
+    // publishRetryMessage should NOT be called because messages are marked as invalid
+    // Invalid messages go back to the SQS queue and will be retried there
+    expect(mockPublishRetryMessage).not.toHaveBeenCalled();
+    expect(createRunner).not.toHaveBeenCalled();
+  });
+
+  it('calls publishRetryMessage with correct message structure including retry counter', async () => {
+    const message = {
+      ...TEST_DATA_SINGLE,
+      messageId: 'test-message-id',
+      retryCounter: 2,
+    };
+
+    await scaleUpModule.scaleUp([message]);
+
+    expect(mockPublishRetryMessage).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: message.id,
+        messageId: 'test-message-id',
+        retryCounter: 2,
+      }),
+    );
+  });
+
+  it('calls publishRetryMessage when ENABLE_JOB_QUEUED_CHECK is false', async () => {
+    process.env.ENABLE_JOB_QUEUED_CHECK = 'false';
+    mockCreateRunner.mockResolvedValue(['i-12345', 'i-67890']); // Create all requested runners
+
+    const messages = createTestMessages(2);
+
+    await scaleUpModule.scaleUp(messages);
+
+    // Should always call publishRetryMessage when queue check is disabled
+    expect(mockPublishRetryMessage).toHaveBeenCalledTimes(2);
+    expect(mockOctokit.actions.getJobForWorkflowRun).not.toHaveBeenCalled();
+  });
+
+  it('calls publishRetryMessage for each message in a multi-runner scenario', async () => {
+    mockCreateRunner.mockResolvedValue(['i-12345', 'i-67890', 'i-abcdef', 'i-11111', 'i-22222']); // Create all requested runners
+    const messages = createTestMessages(5);
+
+    await scaleUpModule.scaleUp(messages);
+
+    expect(mockPublishRetryMessage).toHaveBeenCalledTimes(5);
+    messages.forEach((msg, index) => {
+      expect(mockPublishRetryMessage).toHaveBeenNthCalledWith(
+        index + 1,
+        expect.objectContaining({
+          id: msg.id,
+          messageId: msg.messageId,
+        }),
+      );
+    });
+  });
+
+  it('calls publishRetryMessage after runner creation', async () => {
+    const messages = createTestMessages(1);
+    mockCreateRunner.mockResolvedValue(['i-12345']); // Create the requested runner
+
+    const callOrder: string[] = [];
+    mockPublishRetryMessage.mockImplementation(() => {
+      callOrder.push('publishRetryMessage');
+      return Promise.resolve();
+    });
+    mockCreateRunner.mockImplementation(async () => {
+      callOrder.push('createRunner');
+      return ['i-12345'];
+    });
+
+    await scaleUpModule.scaleUp(messages);
+
+    expect(callOrder).toEqual(['createRunner', 'publishRetryMessage']);
+  });
+});
+
 function defaultOctokitMockImpl() {
   mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({
     data: {

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

@@ -7,6 +7,7 @@ import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient
 import { createRunner, listEC2Runners, tag } from './../aws/runners';
 import { RunnerInputParameters } from './../aws/runners.d';
 import { metricGitHubAppRateLimit } from '../github/rate-limit';
+import { publishRetryMessage } from './job-retry';
 
 const logger = createChildLogger('scale-up');
 
@@ -276,7 +277,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
   };
 
   const validMessages = new Map<string, MessagesWithClient>();
-  const invalidMessages: string[] = [];
+  const rejectedMessageIds = new Set<string>();
   for (const payload of payloads) {
     const { eventType, messageId, repositoryName, repositoryOwner } = payload;
     if (ephemeralEnabled && eventType !== 'workflow_job') {
@@ -285,7 +286,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
         { eventType, messageId },
       );
 
-      invalidMessages.push(messageId);
+      rejectedMessageIds.add(messageId);
 
       continue;
     }
@@ -340,6 +341,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
   for (const [group, { githubInstallationClient, messages }] of validMessages.entries()) {
     // Work out how much we want to scale up by.
     let scaleUp = 0;
+    const queuedMessages: ActionRequestMessageSQS[] = [];
 
     for (const message of messages) {
       const messageLogger = logger.createChild({
@@ -358,6 +360,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
       }
 
       scaleUp++;
+      queuedMessages.push(message);
     }
 
     if (scaleUp === 0) {
@@ -393,11 +396,18 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
       if (ephemeralEnabled) {
         // This removes `missingInstanceCount` items from the start of the array
         // so that, if we retry more messages later, we pick fresh ones.
-        invalidMessages.push(...messages.splice(0, missingInstanceCount).map(({ messageId }) => messageId));
+        const removedMessages = messages.splice(0, missingInstanceCount);
+        removedMessages.forEach(({ messageId }) => rejectedMessageIds.add(messageId));
       }
 
       // No runners will be created, so skip calling the EC2 API.
-      if (missingInstanceCount === scaleUp) {
+      if (newRunners <= 0) {
+        // Publish retry messages for messages that are not rejected
+        for (const message of queuedMessages) {
+          if (!rejectedMessageIds.has(message.messageId)) {
+            await publishRetryMessage(message as ActionRequestMessageRetry);
+          }
+        }
         continue;
       }
     }
@@ -450,11 +460,19 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
         failedInstanceCount,
       });
 
-      invalidMessages.push(...messages.slice(0, failedInstanceCount).map(({ messageId }) => messageId));
+      const failedMessages = messages.slice(0, failedInstanceCount);
+      failedMessages.forEach(({ messageId }) => rejectedMessageIds.add(messageId));
+    }
+
+    // Publish retry messages for messages that are not rejected
+    for (const message of queuedMessages) {
+      if (!rejectedMessageIds.has(message.messageId)) {
+        await publishRetryMessage(message as ActionRequestMessageRetry);
+      }
     }
   }
 
-  return invalidMessages;
+  return Array.from(rejectedMessageIds);
 }
 
 export function getGitHubEnterpriseApiUrl() {