fix: Move publishRetryMessage to end of processing loop to avoid duplicate retries (#4966)

- [x] Understand the current flow and identify where
`publishRetryMessage` is called
- [x] Remove the `publishRetryMessage` call from line 360 (early in the
loop)
- [x] Add logic to track which messages should have retry messages
published
- [x] Call `publishRetryMessage` at the end of the loop for messages not
marked as invalid
- [x] Update tests to reflect the new behavior (publishRetryMessage
after runner creation)
- [x] Validate changes with linting and testing
- [x] Address code review feedback: use Set for performance, extract
helper function
- [x] Simplify naming and reduce complexity per reviewer feedback
- [x] Update lambda workflow to run on all PRs
- [x] Fix formatting issues
- [x] Fix merge conflicts from stu/fix_job_retry branch
- [x] Revert unnecessary changes to keep PR minimal and focused


Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: npalm <[email protected]>
Co-authored-by: Stuart Pearson <[email protected]>
Co-authored-by: Brend-Smits <[email protected]>

Author: 4 people

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

@@ -1711,6 +1711,7 @@ describe('Retry mechanism tests', () => {
 
   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);
 
@@ -1762,7 +1763,7 @@ describe('Retry mechanism tests', () => {
     );
   });
 
-  it('calls publishRetryMessage even when maximum runners is reached', async () => {
+  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);
@@ -1776,8 +1777,9 @@ describe('Retry mechanism tests', () => {
       runnerOwner: TEST_DATA_SINGLE.repositoryOwner,
     });
 
-    // publishRetryMessage should still be called even though no runners will be created
-    expect(mockPublishRetryMessage).toHaveBeenCalledTimes(2);
+    // 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();
   });
 
@@ -1801,6 +1803,7 @@ describe('Retry mechanism tests', () => {
 
   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);
 
@@ -1812,6 +1815,7 @@ describe('Retry mechanism tests', () => {
   });
 
   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);
@@ -1828,8 +1832,9 @@ describe('Retry mechanism tests', () => {
     });
   });
 
-  it('calls publishRetryMessage before runner creation', async () => {
+  it('calls publishRetryMessage after runner creation', async () => {
     const messages = createTestMessages(1);
+    mockCreateRunner.mockResolvedValue(['i-12345']); // Create the requested runner
 
     const callOrder: string[] = [];
     mockPublishRetryMessage.mockImplementation(() => {
@@ -1843,7 +1848,7 @@ describe('Retry mechanism tests', () => {
 
     await scaleUpModule.scaleUp(messages);
 
-    expect(callOrder).toEqual(['publishRetryMessage', 'createRunner']);
+    expect(callOrder).toEqual(['createRunner', 'publishRetryMessage']);
   });
 });
 

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

@@ -277,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') {
@@ -286,7 +286,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
         { eventType, messageId },
       );
 
-      invalidMessages.push(messageId);
+      rejectedMessageIds.add(messageId);
 
       continue;
     }
@@ -341,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({
@@ -359,7 +360,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
       }
 
       scaleUp++;
-      await publishRetryMessage(message);
+      queuedMessages.push(message);
     }
 
     if (scaleUp === 0) {
@@ -395,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 (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;
       }
     }
@@ -452,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() {