feat(runners): add origin property to runner parameters and update related logic

Author: Brend-Smits

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

@@ -44,4 +44,5 @@ export interface RunnerInputParameters {
   amiIdSsmParameterName?: string;
   tracingEnabled?: boolean;
   onDemandFailoverOnError?: string[];
+  origin?: string;
 }

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

@@ -256,6 +256,10 @@ async function createInstances(
     { Key: 'ghr:Owner', Value: runnerParameters.runnerOwner },
   ];
 
+  if (runnerParameters.origin) {
+    tags.push({ Key: 'ghr:origin', Value: runnerParameters.origin });
+  }
+
   if (runnerParameters.tracingEnabled) {
     const traceId = tracer.getRootXrayTraceId();
     tags.push({ Key: 'ghr:trace_id', Value: traceId! });

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

@@ -168,6 +168,7 @@ describe(`Test job retry check`, () => {
     expect(publishMessage).toHaveBeenCalledWith(
       JSON.stringify({
         ...message,
+        origin: 'job-retry',
       }),
       'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue',
     );
@@ -206,6 +207,7 @@ describe(`Test job retry check`, () => {
     expect(publishMessage).toHaveBeenCalledWith(
       JSON.stringify({
         ...message,
+        origin: 'job-retry',
       }),
       'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue',
     );
@@ -270,6 +272,41 @@ describe(`Test job retry check`, () => {
     // assert
     expect(publishMessage).not.toHaveBeenCalled();
   });
+
+  it('should set origin to job-retry when publishing message for retry', async () => {
+    // setup
+    mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({
+      data: {
+        status: 'queued',
+      },
+    }));
+
+    const message: ActionRequestMessageRetry = {
+      eventType: 'workflow_job',
+      id: 0,
+      installationId: 0,
+      repositoryName: 'test',
+      repositoryOwner: 'github-aws-runners',
+      repoOwnerType: 'Organization',
+      retryCounter: 0,
+    };
+    process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
+    process.env.RUNNER_NAME_PREFIX = 'test';
+    process.env.JOB_QUEUE_SCALE_UP_URL =
+      'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue';
+
+    // act
+    await checkAndRetryJob(message);
+
+    // assert
+    expect(publishMessage).toHaveBeenCalledWith(
+      JSON.stringify({
+        ...message,
+        origin: 'job-retry',
+      }),
+      'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue',
+    );
+  });
 });
 
 describe('Test job retry handler (batch processing)', () => {

lambdas/functions/control-plane/src/scale-runners/job-retry.ts

@@ -64,7 +64,9 @@ export async function checkAndRetryJob(payload: ActionRequestMessageRetry): Prom
 
   // check job is still queued
   if (await isJobQueued(ghClient, payload)) {
-    await publishMessage(JSON.stringify(payload), jobQueueUrl);
+    // Set the origin to 'job-retry' when publishing back to the scale-up queue
+    const retryPayload = { ...payload, origin: 'job-retry' };
+    await publishMessage(JSON.stringify(retryPayload), jobQueueUrl);
     createMetric(enableMetrics, environment, payload);
     logger.info(`Job is still queued, message published to build queue and will be handled by scale-up.`, { payload });
   } else {

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

@@ -105,6 +105,7 @@ const EXPECTED_RUNNER_PARAMS: RunnerInputParameters = {
   subnets: ['subnet-123'],
   tracingEnabled: false,
   onDemandFailoverOnError: [],
+  origin: 'build-queue',
 };
 let expectedRunnerParams: RunnerInputParameters;
 
@@ -187,6 +188,27 @@ describe('scaleUp with GHES', () => {
     expect(listEC2Runners).not.toBeCalled();
   });
 
+  it('creates a runner with origin set to build-queue when no origin is provided', async () => {
+    process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
+    await scaleUpModule.scaleUp('aws:sqs', TEST_DATA);
+    expect(createRunner).toBeCalledWith(
+      expect.objectContaining({
+        origin: 'build-queue',
+      }),
+    );
+  });
+
+  it('creates a runner with origin from payload when provided', async () => {
+    process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
+    const testDataWithOrigin = { ...TEST_DATA, origin: 'job-retry' };
+    await scaleUpModule.scaleUp('aws:sqs', testDataWithOrigin);
+    expect(createRunner).toBeCalledWith(
+      expect.objectContaining({
+        origin: 'job-retry',
+      }),
+    );
+  });
+
   describe('on org level', () => {
     beforeEach(() => {
       process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';

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

@@ -29,6 +29,7 @@ export interface ActionRequestMessage {
   installationId: number;
   repoOwnerType: string;
   retryCounter?: number;
+  origin?: string;
 }
 
 export interface ActionRequestMessageSQS extends ActionRequestMessage {
@@ -62,6 +63,7 @@ interface CreateEC2RunnerConfig {
   amiIdSsmParameterName?: string;
   tracingEnabled?: boolean;
   onDemandFailoverOnError?: string[];
+  origin?: string;
 }
 
 function generateRunnerServiceConfig(githubRunnerConfig: CreateGitHubRunnerConfig, token: string) {
@@ -301,6 +303,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
     }
 
     const key = enableOrgLevel ? payload.repositoryOwner : `${payload.repositoryOwner}/${payload.repositoryName}`;
+    const origin = payload.origin || 'build-queue';
 
     let entry = validMessages.get(key);
 
@@ -431,6 +434,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
         amiIdSsmParameterName,
         tracingEnabled,
         onDemandFailoverOnError,
+        origin,
       },
       newRunners,
       githubInstallationClient,