Skip to content

Commit d8516ef

Browse files
committed
add cloudwatch events scheduler
1 parent 726cd87 commit d8516ef

File tree

2 files changed

+246
-0
lines changed

2 files changed

+246
-0
lines changed

src/scheduler.ts

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
import { v4 as uuidv4 } from 'uuid';
2+
import CloudWatchEvents from 'aws-sdk/clients/cloudwatchevents';
3+
4+
import { SessionProxy } from './proxy';
5+
import { HandlerRequest, minToCron } from './utils';
6+
7+
8+
const LOGGER = console;
9+
10+
/**
11+
* Schedule a re-invocation of the executing handler no less than 1 minute from
12+
* now
13+
*
14+
* @param session AWS session where to retrieve CloudWatchEvents client
15+
* @param functionArn the ARN of the Lambda function to be invoked
16+
* @param minutesFromNow the minimum minutes from now that the re-invocation
17+
* will occur. CWE provides only minute-granularity
18+
* @param handlerRequest additional context which the handler can provide itself
19+
* for re-invocation
20+
*/
21+
export function rescheduleAfterMinutes(
22+
session: SessionProxy,
23+
functionArn: string,
24+
minutesFromNow: number,
25+
handlerRequest: HandlerRequest,
26+
): void {
27+
const client: CloudWatchEvents = session.client('CloudWatchEvents') as CloudWatchEvents;
28+
const cron = minToCron(Math.max(minutesFromNow, 1));
29+
const identifier = uuidv4();
30+
const ruleName = `reinvoke-handler-${identifier}`;
31+
const targetId = `reinvoke-target-${identifier}`;
32+
handlerRequest.requestContext.cloudWatchEventsRuleName = ruleName;
33+
handlerRequest.requestContext.cloudWatchEventsTargetId = targetId;
34+
const jsonRequest = JSON.stringify(handlerRequest);
35+
LOGGER.debug(`Scheduling re-invoke at ${cron} (${identifier})`);
36+
client.putRule({
37+
Name: ruleName,
38+
ScheduleExpression: cron,
39+
State: 'ENABLED',
40+
});
41+
client.putTargets({
42+
Rule: ruleName,
43+
Targets: [{
44+
Id: targetId,
45+
Arn: functionArn,
46+
Input: jsonRequest,
47+
}],
48+
});
49+
}
50+
51+
/**
52+
* After a re-invocation, the CWE rule which generated the reinvocation should
53+
* be scrubbed
54+
*
55+
* @param session AWS session where to retrieve CloudWatchEvents client
56+
* @param ruleName the name of the CWE rule which triggered a re-invocation
57+
* @param targetId the target of the CWE rule which triggered a re-invocation
58+
*/
59+
export function cleanupCloudwatchEvents(
60+
session: SessionProxy, ruleName: string, targetId: string
61+
): void {
62+
const client: CloudWatchEvents = session.client('CloudWatchEvents') as CloudWatchEvents;
63+
try {
64+
if (targetId && ruleName) {
65+
client.removeTargets({
66+
Rule: ruleName,
67+
Ids: [targetId],
68+
});
69+
}
70+
} catch(err) {
71+
LOGGER.error(
72+
`Error cleaning CloudWatchEvents Target (targetId=${targetId}): ${err.message}`
73+
);
74+
}
75+
76+
try {
77+
if (ruleName) {
78+
client.deleteRule({
79+
Name: ruleName,
80+
Force: true,
81+
});
82+
}
83+
} catch(err) {
84+
LOGGER.error(
85+
`Error cleaning CloudWatchEvents Rule (ruleName=${ruleName}): ${err.message}`
86+
);
87+
}
88+
}

tests/lib/scheduler.test.ts

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
import CloudWatchEvents from 'aws-sdk/clients/cloudwatchevents';
2+
import awsUtil = require('aws-sdk/lib/util');
3+
4+
import { cleanupCloudwatchEvents, rescheduleAfterMinutes } from '../../src/scheduler';
5+
import { SessionProxy } from '../../src/proxy';
6+
import { RequestContext } from '../../src/interface';
7+
import * as utils from '../../src/utils';
8+
9+
10+
const mockResult = (output: any): jest.Mock => {
11+
return jest.fn().mockReturnValue({
12+
promise: jest.fn().mockResolvedValue(output)
13+
});
14+
};
15+
16+
const IDENTIFIER: string = 'f3390613-b2b5-4c31-a4c6-66813dff96a6';
17+
18+
jest.mock('aws-sdk/clients/cloudwatchevents');
19+
jest.mock('uuid', () => {
20+
return {
21+
v4: () => IDENTIFIER
22+
};
23+
});
24+
25+
describe('when getting scheduler', () => {
26+
27+
let session: SessionProxy;
28+
let handlerRequest: utils.HandlerRequest;
29+
let cwEvents: jest.Mock;
30+
let spyConsoleError: jest.SpyInstance;
31+
let spyMinToCron: jest.SpyInstance;
32+
let mockPutRule: jest.Mock;
33+
let mockPutTargets: jest.Mock;
34+
let mockRemoveTargets: jest.Mock;
35+
let mockDeleteRule: jest.Mock;
36+
37+
beforeEach(() => {
38+
spyConsoleError = jest.spyOn(global.console, 'error').mockImplementation(() => {});
39+
spyMinToCron = jest.spyOn(utils, 'minToCron')
40+
.mockReturnValue('cron(30 16 21 11 ? 2019)');
41+
mockPutRule = mockResult({ ResponseMetadata: { RequestId: 'mock-request' }});
42+
mockPutTargets = mockResult({ ResponseMetadata: { RequestId: 'mock-request' }});
43+
mockRemoveTargets = mockResult({ ResponseMetadata: { RequestId: 'mock-request' }});
44+
mockDeleteRule = mockResult({ ResponseMetadata: { RequestId: 'mock-request' }});
45+
46+
cwEvents = (CloudWatchEvents as unknown) as jest.Mock;
47+
cwEvents.mockImplementation(() => {
48+
const returnValue = {
49+
deleteRule: mockDeleteRule,
50+
putRule: mockPutRule,
51+
putTargets: mockPutTargets,
52+
removeTargets: mockRemoveTargets,
53+
};
54+
return {
55+
...returnValue,
56+
makeRequest: (operation: string, params?: {[key: string]: any}) => {
57+
return returnValue[operation](params);
58+
}
59+
};
60+
});
61+
session = new SessionProxy({});
62+
session['client'] = cwEvents;
63+
64+
handlerRequest = new utils.HandlerRequest()
65+
handlerRequest.requestContext = {} as RequestContext<Map<string, any>>;
66+
handlerRequest.toJSON = jest.fn(() => new Object());
67+
});
68+
69+
afterEach(() => {
70+
jest.clearAllMocks();
71+
jest.restoreAllMocks();
72+
});
73+
74+
test('reschedule after minutes zero', () => {
75+
// if called with zero, should call cron with a 1
76+
rescheduleAfterMinutes(session, 'arn:goes:here', 0, handlerRequest);
77+
78+
expect(cwEvents).toHaveBeenCalledTimes(1);
79+
expect(cwEvents).toHaveBeenCalledWith('CloudWatchEvents');
80+
expect(spyMinToCron).toHaveBeenCalledTimes(1);
81+
expect(spyMinToCron).toHaveBeenCalledWith(1);
82+
});
83+
84+
test('reschedule after minutes not zero', () => {
85+
// if called with another number, should use that
86+
rescheduleAfterMinutes(session, 'arn:goes:here', 2, handlerRequest);
87+
88+
expect(cwEvents).toHaveBeenCalledTimes(1);
89+
expect(cwEvents).toHaveBeenCalledWith('CloudWatchEvents');
90+
expect(spyMinToCron).toHaveBeenCalledTimes(1);
91+
expect(spyMinToCron).toHaveBeenCalledWith(2);
92+
});
93+
94+
test('reschedule after minutes success', () => {
95+
rescheduleAfterMinutes(session, 'arn:goes:here', 2, handlerRequest);
96+
97+
expect(cwEvents).toHaveBeenCalledTimes(1);
98+
expect(cwEvents).toHaveBeenCalledWith('CloudWatchEvents');
99+
expect(mockPutRule).toHaveBeenCalledTimes(1);
100+
expect(mockPutRule).toHaveBeenCalledWith({
101+
Name: `reinvoke-handler-${IDENTIFIER}`,
102+
ScheduleExpression: 'cron(30 16 21 11 ? 2019)',
103+
State: 'ENABLED',
104+
});
105+
expect(mockPutTargets).toHaveBeenCalledTimes(1);
106+
expect(mockPutTargets).toHaveBeenCalledWith({
107+
Rule: `reinvoke-handler-${IDENTIFIER}`,
108+
Targets: [
109+
{
110+
Id: `reinvoke-target-${IDENTIFIER}`,
111+
Arn: 'arn:goes:here',
112+
Input: '{}',
113+
}
114+
],
115+
});
116+
});
117+
118+
test('cleanup cloudwatch events empty', () => {
119+
// cleanup should silently pass if rule/target are empty
120+
cleanupCloudwatchEvents(session, '', '');
121+
122+
expect(cwEvents).toHaveBeenCalledTimes(1);
123+
expect(cwEvents).toHaveBeenCalledWith('CloudWatchEvents');
124+
expect(mockRemoveTargets).toHaveBeenCalledTimes(0);
125+
expect(mockDeleteRule).toHaveBeenCalledTimes(0);
126+
expect(spyConsoleError).toHaveBeenCalledTimes(0);
127+
});
128+
129+
test('cleanup cloudwatch events success', () => {
130+
// when rule_name and target_id are provided we should call events client and not
131+
// log errors if the deletion succeeds
132+
cleanupCloudwatchEvents(session, 'rulename', 'targetid');
133+
134+
expect(spyConsoleError).toHaveBeenCalledTimes(0);
135+
expect(cwEvents).toHaveBeenCalledTimes(1);
136+
expect(cwEvents).toHaveBeenCalledWith('CloudWatchEvents');
137+
expect(mockRemoveTargets).toHaveBeenCalledTimes(1);
138+
expect(mockDeleteRule).toHaveBeenCalledTimes(1);
139+
expect(mockPutRule).toHaveBeenCalledTimes(0);
140+
expect(spyConsoleError).toHaveBeenCalledTimes(0);
141+
});
142+
143+
test('cleanup cloudwatch events client error', () => {
144+
// cleanup should catch and log client failures
145+
const error = awsUtil.error(new Error(), { code: '1' });
146+
mockRemoveTargets.mockImplementation(() => {throw error});
147+
mockDeleteRule.mockImplementation(() => {throw error});
148+
149+
cleanupCloudwatchEvents(session, 'rulename', 'targetid');
150+
151+
expect(cwEvents).toHaveBeenCalledTimes(1);
152+
expect(cwEvents).toHaveBeenCalledWith('CloudWatchEvents');
153+
expect(spyConsoleError).toHaveBeenCalledTimes(2);
154+
expect(mockRemoveTargets).toHaveBeenCalledTimes(1);
155+
expect(mockDeleteRule).toHaveBeenCalledTimes(1);
156+
expect(mockPutTargets).toHaveBeenCalledTimes(0);
157+
});
158+
});

0 commit comments

Comments
 (0)