Skip to content

Commit

Permalink
WIP: list runners, create token if needed
Browse files Browse the repository at this point in the history
  • Loading branch information
@gertjanmaas
gertjanmaas committed May 7, 2020
1 parent 503543c commit cf7124c
Show file tree
Hide file tree
Showing 5 changed files with 230 additions and 411 deletions.
117 changes: 77 additions & 40 deletions modules/runners/lambdas/scale-runners/src/scale-runners/handler.test.ts
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import { mocked } from 'ts-jest/utils';
import { ActionRequestMessage, handle } from './handler';

import { createAppAuth } from '@octokit/auth-app';
import { Octokit } from '@octokit/rest';
import { listRunners } from './runners';

jest.mock('@octokit/auth-app', () => ({
createAppAuth: jest.fn().mockImplementation(() => jest.fn().mockImplementation(() => ({ token: 'Blaat' }))),
Expand All @@ -10,14 +12,16 @@ const mockOctokit = {
checks: { get: jest.fn() },
actions: {
listRepoWorkflowRuns: jest.fn(),
listSelfHostedRunnersForOrg: jest.fn(),
listSelfHostedRunnersForRepo: jest.fn(),
createRegistrationTokenForOrg: jest.fn(),
createRegistrationTokenForRepo: jest.fn(),
},
};
jest.mock('@octokit/rest', () => ({
Octokit: jest.fn().mockImplementation(() => mockOctokit),
}));

jest.mock('./runners');

const TEST_DATA: ActionRequestMessage = {
id: 1,
eventType: 'check_run',
Expand All @@ -32,27 +36,29 @@ describe('handler', () => {
process.env.GITHUB_APP_ID = '1337';
process.env.GITHUB_APP_CLIENT_ID = 'TEST_CLIENT_ID';
process.env.GITHUB_APP_CLIENT_SECRET = 'TEST_CLIENT_SECRET';
process.env.RUNNERS_MAXIMUM_COUNT = '3';
jest.clearAllMocks();
mockOctokit.actions.listRepoWorkflowRuns.mockImplementation(() => ({
data: {
total_count: 1,
},
}));
const mockRunnersReturnValue = {
const mockTokenReturnValue = {
data: {
total_count: 1,
runners: [
{
id: 23,
name: 'Test Runner',
status: 'online',
os: 'linux',
},
],
token: '1234abcd',
},
};
mockOctokit.actions.listSelfHostedRunnersForOrg.mockImplementation(() => mockRunnersReturnValue);
mockOctokit.actions.listSelfHostedRunnersForRepo.mockImplementation(() => mockRunnersReturnValue);
mockOctokit.actions.createRegistrationTokenForOrg.mockImplementation(() => mockTokenReturnValue);
mockOctokit.actions.createRegistrationTokenForRepo.mockImplementation(() => mockTokenReturnValue);
const mockListRunners = mocked(listRunners);
mockListRunners.mockImplementation(async () => [
{
instanceId: 'i-1234',
launchTime: new Date(),
repo: `${TEST_DATA.repositoryOwner}/${TEST_DATA.repositoryName}`,
org: TEST_DATA.repositoryOwner,
},
]);
});

it('ignores non-sqs events', async () => {
Expand All @@ -69,30 +75,61 @@ describe('handler', () => {
});
});

// describe('on org level', () => {
// beforeAll(() => {
// process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
// });

// it('gets the current org level runners', async () => {
// await handle('aws:sqs', TEST_DATA);
// expect(mockOctokit.actions.listSelfHostedRunnersForOrg).toBeCalledWith({
// org: TEST_DATA.repositoryOwner,
// });
// });
// });

// describe('on repo level', () => {
// beforeAll(() => {
// delete process.env.ENABLE_ORGANIZATION_RUNNERS;
// });

// it('gets the current repo level runners', async () => {
// await handle('aws:sqs', TEST_DATA);
// expect(mockOctokit.actions.listSelfHostedRunnersForRepo).toBeCalledWith({
// owner: TEST_DATA.repositoryOwner,
// repo: TEST_DATA.repositoryName,
// });
// });
// });
it('does not list runners when no workflows are queued', async () => {
mockOctokit.actions.listRepoWorkflowRuns.mockImplementation(() => ({
data: { total_count: 0, runners: [] },
}));
await handle('aws:sqs', TEST_DATA);
expect(listRunners).not.toBeCalled();
});

describe('on org level', () => {
beforeAll(() => {
process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
});

it('gets the current org level runners', async () => {
await handle('aws:sqs', TEST_DATA);
expect(listRunners).toBeCalledWith({ repoName: undefined });
});

it('does not create a token when maximum runners has been reached', async () => {
process.env.RUNNERS_MAXIMUM_COUNT = '1';
await handle('aws:sqs', TEST_DATA);
expect(mockOctokit.actions.createRegistrationTokenForOrg).not.toBeCalled();
});

it('creates a token when maximum runners has not been reached', async () => {
await handle('aws:sqs', TEST_DATA);
expect(mockOctokit.actions.createRegistrationTokenForOrg).toBeCalled();
expect(mockOctokit.actions.createRegistrationTokenForOrg).toBeCalledWith({
org: TEST_DATA.repositoryOwner,
});
});
});

describe('on repo level', () => {
beforeAll(() => {
process.env.ENABLE_ORGANIZATION_RUNNERS = 'false';
});

it('gets the current repo level runners', async () => {
await handle('aws:sqs', TEST_DATA);
expect(listRunners).toBeCalledWith({ repoName: `${TEST_DATA.repositoryOwner}/${TEST_DATA.repositoryName}` });
});

it('does not create a token when maximum runners has been reached', async () => {
process.env.RUNNERS_MAXIMUM_COUNT = '1';
await handle('aws:sqs', TEST_DATA);
expect(mockOctokit.actions.createRegistrationTokenForRepo).not.toBeCalled();
});

it('creates a token when maximum runners has not been reached', async () => {
await handle('aws:sqs', TEST_DATA);
expect(mockOctokit.actions.createRegistrationTokenForRepo).toBeCalledWith({
owner: TEST_DATA.repositoryOwner,
repo: TEST_DATA.repositoryName,
});
});
});
});
38 changes: 26 additions & 12 deletions modules/runners/lambdas/scale-runners/src/scale-runners/handler.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { createAppAuth } from '@octokit/auth-app';
import { Octokit } from '@octokit/rest';
import { AppAuth } from '@octokit/auth-app/dist-types/types';
import { listRunners } from './runners';
import yn from 'yn';

export interface ActionRequestMessage {
Expand Down Expand Up @@ -34,29 +35,42 @@ async function createInstallationClient(githubAppAuth: AppAuth): Promise<Octokit
export const handle = async (eventSource: string, payload: ActionRequestMessage): Promise<void> => {
if (eventSource !== 'aws:sqs') throw Error('Cannot handle non-SQS events!');
const enableOrgLevel = yn(process.env.ENABLE_ORGANIZATION_RUNNERS);
const maximumRunners = parseInt(process.env.RUNNERS_MAXIMUM_COUNT || '3');
const githubAppAuth = createGithubAppAuth(payload.installationId);
const githubInstallationClient = await createInstallationClient(githubAppAuth);
const queuedWorkflows = await githubInstallationClient.actions.listRepoWorkflowRuns({
owner: payload.repositoryOwner,
repo: payload.repositoryName,
// @ts-ignore (typing is incorrect)
// @ts-ignore (typing of the 'status' field is incorrect)
status: 'queued',
});
console.info(
`Repo ${payload.repositoryOwner}/${payload.repositoryName} has ${queuedWorkflows.data.total_count} queued workflow runs`,
);

if (queuedWorkflows.data.total_count > 0) {
// console.log(enableOrgLevel);
// const currentRunners = enableOrgLevel
// ? await githubInstallationClient.actions.listSelfHostedRunnersForOrg({
// org: payload.repositoryOwner,
// })
// : await githubInstallationClient.actions.listSelfHostedRunnersForRepo({
// owner: payload.repositoryOwner,
// repo: payload.repositoryName,
// });
// // const currentOnlineRunners = currentRunners.data.runners.filter((r) => r.status === 'online');
// // if (currentOnlineRunners.length > 0)
const currentRunners = await listRunners({
repoName: enableOrgLevel ? undefined : `${payload.repositoryOwner}/${payload.repositoryName}`,
});
console.info(
`${
enableOrgLevel
? `Organization ${payload.repositoryOwner}`
: `Repo ${payload.repositoryOwner}/${payload.repositoryName}`
} has ${currentRunners.length}/${maximumRunners} runners`,
);
console.log(currentRunners.length);
console.log(maximumRunners);
if (currentRunners.length < maximumRunners) {
// create token
const registrationToken = enableOrgLevel
? await githubInstallationClient.actions.createRegistrationTokenForOrg({ org: payload.repositoryOwner })
: await githubInstallationClient.actions.createRegistrationTokenForRepo({
owner: payload.repositoryOwner,
repo: payload.repositoryName,
});
const token = registrationToken.data.token;
// create runner
}
}
};
92 changes: 86 additions & 6 deletions modules/runners/lambdas/scale-runners/src/scale-runners/runners.test.ts
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,98 @@
import { listRunners } from './runners';
import { handle } from './handler';
import { listRunners, RunnerInfo } from './runners';
import { EC2 } from 'aws-sdk';

jest.mock('./handler');
const mockEC2 = { describeInstances: jest.fn() };
jest.mock('aws-sdk', () => ({
EC2: jest.fn().mockImplementation(() => mockEC2),
}));

describe('list instances', () => {
beforeAll(() => {
const mockDescribeInstances = { promise: jest.fn() };
beforeEach(() => {
jest.clearAllMocks();
mockEC2.describeInstances.mockImplementation(() => mockDescribeInstances);
const mockRunningInstances: AWS.EC2.DescribeInstancesResult = {
Reservations: [
{
Instances: [
{
LaunchTime: new Date('2020-10-10T14:48:00.000+09:00'),
InstanceId: 'i-1234',
Tags: [
{ Key: 'Repo', Value: 'CoderToCat/hello-world' },
{ Key: 'Org', Value: 'CoderToCat' },
{ Key: 'Application', Value: 'github-action-runner' },
],
},
{
LaunchTime: new Date('2020-10-11T14:48:00.000+09:00'),
InstanceId: 'i-5678',
Tags: [
{ Key: 'Repo', Value: 'SomeAwesomeCoder/some-amazing-library' },
{ Key: 'Org', Value: 'SomeAwesomeCoder' },
{ Key: 'Application', Value: 'github-action-runner' },
],
},
],
},
],
};
mockDescribeInstances.promise.mockReturnValue(mockRunningInstances);
});
it('returns a list of instances', () => {
listRunners();

it('returns a list of instances', async () => {
const resp = await listRunners();
expect(resp.length).toBe(2);
expect(resp).toContainEqual({
instanceId: 'i-1234',
launchTime: new Date('2020-10-10T14:48:00.000+09:00'),
repo: 'CoderToCat/hello-world',
org: 'CoderToCat',
});
expect(resp).toContainEqual({
instanceId: 'i-5678',
launchTime: new Date('2020-10-11T14:48:00.000+09:00'),
repo: 'SomeAwesomeCoder/some-amazing-library',
org: 'SomeAwesomeCoder',
});
});

it('calls EC2 describe instances', async () => {
await listRunners();
expect(mockEC2.describeInstances).toBeCalled();
});

it('filters instances on repo name', async () => {
await listRunners({ repoName: 'SomeAwesomeCoder/some-amazing-library' });
expect(mockEC2.describeInstances).toBeCalledWith({
Filters: [
{ Name: 'tag:Application', Values: ['github-action-runner'] },
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
{ Name: 'tag:Repo', Values: ['SomeAwesomeCoder/some-amazing-library'] },
],
});
});

it('filters instances on org name', async () => {
await listRunners({ orgName: 'SomeAwesomeCoder' });
expect(mockEC2.describeInstances).toBeCalledWith({
Filters: [
{ Name: 'tag:Application', Values: ['github-action-runner'] },
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
{ Name: 'tag:Org', Values: ['SomeAwesomeCoder'] },
],
});
});

it('filters instances on both org name and repo name', async () => {
await listRunners({ orgName: 'SomeAwesomeCoder', repoName: 'SomeAwesomeCoder/some-amazing-library' });
expect(mockEC2.describeInstances).toBeCalledWith({
Filters: [
{ Name: 'tag:Application', Values: ['github-action-runner'] },
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
{ Name: 'tag:Repo', Values: ['SomeAwesomeCoder/some-amazing-library'] },
{ Name: 'tag:Org', Values: ['SomeAwesomeCoder'] },
],
});
});
});
49 changes: 34 additions & 15 deletions modules/runners/lambdas/scale-runners/src/scale-runners/runners.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,26 +2,45 @@ import { EC2 } from 'aws-sdk';

export interface RunnerInfo {
instanceId: string;
launchTime: Date;
repo: string;
org: string;
launchTime: Date | undefined;
repo: string | undefined;
org: string | undefined;
}

const ec2 = new EC2();
export async function listRunners(
repoName: string | undefined = undefined,
orgName: string | undefined = undefined,
): Promise<RunnerInfo[]> {
let filters = [
export interface ListRunnerFilters {
repoName?: string;
orgName?: string;
}

export async function listRunners(filters: ListRunnerFilters | undefined = undefined): Promise<RunnerInfo[]> {
const ec2 = new EC2();
let ec2Filters = [
{ Name: 'tag:Application', Values: ['github-action-runner'] },
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
];
if (repoName !== undefined) {
filters.push({ Name: 'tag:Repo', Values: [repoName] });
if (filters) {
if (filters.repoName !== undefined) {
ec2Filters.push({ Name: 'tag:Repo', Values: [filters.repoName] });
}
if (filters.orgName !== undefined) {
ec2Filters.push({ Name: 'tag:Org', Values: [filters.orgName] });
}
}
if (orgName !== undefined) {
filters.push({ Name: 'tag:Org', Values: [orgName] });
const runningInstances = await ec2.describeInstances({ Filters: ec2Filters }).promise();
const runners: RunnerInfo[] = [];
if (runningInstances.Reservations) {
for (const r of runningInstances.Reservations) {
if (r.Instances) {
for (const i of r.Instances) {
runners.push({
instanceId: i.InstanceId as string,
launchTime: i.LaunchTime,
repo: i.Tags?.find((e) => e.Key === 'Repo')?.Value,
org: i.Tags?.find((e) => e.Key === 'Org')?.Value,
});
}
}
}
}
const runningInstances = await ec2.describeInstances({ Filters: filters }).promise();
return [{ instanceId: 'i-123', launchTime: new Date(), repo: 'bla', org: 'bla' }];
return runners;
}
Loading

0 comments on commit cf7124c

Please sign in to comment.