Skip to content

Commit 87943da

Browse files
marshall-leemarshall-lee
committed
Enable tracing for AWSLambda.
1 parent ce16e33 commit 87943da

File tree

3 files changed

+137
-69
lines changed

3 files changed

+137
-69
lines changed

‎packages/serverless/src/awslambda.ts

Lines changed: 69 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,18 @@
1-
import { captureException, captureMessage, flush, Scope, SDK_VERSION, Severity, withScope } from '@sentry/node';
1+
import {
2+
captureException,
3+
captureMessage,
4+
flush,
5+
getCurrentHub,
6+
Scope,
7+
SDK_VERSION,
8+
Severity,
9+
startTransaction,
10+
withScope,
11+
} from '@sentry/node';
212
import { addExceptionMechanism } from '@sentry/utils';
313
// NOTE: I have no idea how to fix this right now, and don't want to waste more time, as it builds just fine — Kamil
414
// eslint-disable-next-line import/no-unresolved
5-
import { Callback, Context, Handler } from 'aws-lambda';
15+
import { Context, Handler } from 'aws-lambda';
616
import { hostname } from 'os';
717
import { performance } from 'perf_hooks';
818
import { types } from 'util';
@@ -21,7 +31,7 @@ export type AsyncHandler<T extends Handler> = (
2131
context: Parameters<T>[1],
2232
) => Promise<NonNullable<Parameters<Parameters<T>[2]>[1]>>;
2333

24-
interface WrapperOptions {
34+
export interface WrapperOptions {
2535
flushTimeout: number;
2636
rethrowAfterCapture: boolean;
2737
callbackWaitsForEmptyEventLoop: boolean;
@@ -98,36 +108,66 @@ function enhanceScopeWithEnvironmentData(scope: Scope, context: Context): void {
98108
}
99109

100110
/**
101-
* Capture, flush the result down the network stream and await the response.
111+
* Capture exception with a a context.
102112
*
103113
* @param e exception to be captured
104-
* @param options WrapperOptions
114+
* @param context Context
105115
*/
106-
function captureExceptionAsync(e: unknown, context: Context, options: Partial<WrapperOptions>): Promise<boolean> {
116+
function captureExceptionWithContext(e: unknown, context: Context): void {
107117
withScope(scope => {
108118
addServerlessEventProcessor(scope);
109119
enhanceScopeWithEnvironmentData(scope, context);
110120
captureException(e);
111121
});
112-
return flush(options.flushTimeout);
113122
}
114123

115-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
116-
export const wrapHandler = <TEvent = any, TResult = any>(
117-
handler: Handler,
118-
handlerOptions: Partial<WrapperOptions> = {},
119-
): Handler => {
124+
/**
125+
* Wraps a lambda handler adding it error capture and tracing capabilities.
126+
*
127+
* @param handler Handler
128+
* @param options Options
129+
* @returns Handler
130+
*/
131+
export function wrapHandler<TEvent, TResult>(
132+
handler: Handler<TEvent, TResult>,
133+
wrapOptions: Partial<WrapperOptions> = {},
134+
): Handler<TEvent, TResult> {
120135
const options: WrapperOptions = {
121136
flushTimeout: 2000,
122137
rethrowAfterCapture: true,
123138
callbackWaitsForEmptyEventLoop: false,
124139
captureTimeoutWarning: true,
125140
timeoutWarningLimit: 500,
126-
...handlerOptions,
141+
...wrapOptions,
127142
};
128143
let timeoutWarningTimer: NodeJS.Timeout;
129144

130-
return async (event: TEvent, context: Context, callback: Callback<TResult>) => {
145+
// AWSLambda is like Express. It makes a distinction about handlers based on it's last argument
146+
// async (event) => async handler
147+
// async (event, context) => async handler
148+
// (event, context, callback) => sync handler
149+
// Nevertheless whatever option is chosen by user, we convert it to async handler.
150+
const asyncHandler: AsyncHandler<typeof handler> =
151+
handler.length === 3
152+
? (event, context) =>
153+
new Promise((resolve, reject) => {
154+
const rv = (handler as SyncHandler<typeof handler>)(event, context, (error, result) => {
155+
if (error === null || error === undefined) {
156+
resolve(result!); // eslint-disable-line @typescript-eslint/no-non-null-assertion
157+
} else {
158+
reject(error);
159+
}
160+
});
161+
162+
// This should never happen, but still can if someone writes a handler as
163+
// `async (event, context, callback) => {}`
164+
if (isPromise(rv)) {
165+
((rv as unknown) as Promise<NonNullable<TResult>>).then(resolve, reject);
166+
}
167+
})
168+
: (handler as AsyncHandler<typeof handler>);
169+
170+
return async (event, context) => {
131171
context.callbackWaitsForEmptyEventLoop = options.callbackWaitsForEmptyEventLoop;
132172

133173
// In seconds. You cannot go any more granular than this in AWS Lambda.
@@ -155,50 +195,29 @@ export const wrapHandler = <TEvent = any, TResult = any>(
155195
}, timeoutWarningDelay);
156196
}
157197

158-
const callbackWrapper = <TResult>(
159-
callback: Callback<TResult>,
160-
resolve: (value?: unknown) => void,
161-
reject: (reason?: unknown) => void,
162-
): Callback<TResult> => {
163-
return (...args) => {
164-
clearTimeout(timeoutWarningTimer);
165-
if (args[0] === null || args[0] === undefined) {
166-
resolve(callback(...args));
167-
} else {
168-
captureExceptionAsync(args[0], context, options).finally(() => reject(callback(...args)));
169-
}
170-
};
171-
};
198+
const transaction = startTransaction({
199+
name: context.functionName,
200+
op: 'awslambda.handler',
201+
});
202+
// We put the transaction on the scope so users can attach children to it
203+
getCurrentHub().configureScope(scope => {
204+
scope.setSpan(transaction);
205+
});
172206

173207
try {
174-
// AWSLambda is like Express. It makes a distinction about handlers based on it's last argument
175-
// async (event) => async handler
176-
// async (event, context) => async handler
177-
// (event, context, callback) => sync handler
178-
const isSyncHandler = handler.length === 3;
179-
const handlerRv = isSyncHandler
180-
? await new Promise((resolve, reject) => {
181-
const rv = (handler as SyncHandler<Handler<TEvent, TResult>>)(
182-
event,
183-
context,
184-
callbackWrapper(callback, resolve, reject),
185-
);
186-
187-
// This should never happen, but still can if someone writes a handler as
188-
// `async (event, context, callback) => {}`
189-
if (isPromise(rv)) {
190-
((rv as unknown) as Promise<TResult>).then(resolve, reject);
191-
}
192-
})
193-
: await (handler as AsyncHandler<Handler<TEvent, TResult>>)(event, context);
208+
const handlerRv = await asyncHandler(event, context);
194209
clearTimeout(timeoutWarningTimer);
195210
return handlerRv;
196211
} catch (e) {
197212
clearTimeout(timeoutWarningTimer);
198-
await captureExceptionAsync(e, context, options);
213+
captureExceptionWithContext(e, context);
199214
if (options.rethrowAfterCapture) {
200215
throw e;
201216
}
217+
return (undefined as unknown) as TResult;
218+
} finally {
219+
transaction.finish();
220+
await flush(options.flushTimeout);
202221
}
203222
};
204-
};
223+
}

‎packages/serverless/test/__mocks__/@sentry/node.ts

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,23 +2,40 @@ export const SDK_VERSION = '6.6.6';
22
export const Severity = {
33
Warning: 'warning',
44
};
5+
export const fakeParentScope = {
6+
setSpan: jest.fn(),
7+
};
8+
export const fakeHub = {
9+
configureScope: jest.fn((fn: (arg: any) => any) => fn(fakeParentScope)),
10+
};
511
export const fakeScope = {
612
addEventProcessor: jest.fn(),
713
setTransactionName: jest.fn(),
814
setTag: jest.fn(),
915
setContext: jest.fn(),
1016
};
17+
export const fakeTransaction = {
18+
finish: jest.fn(),
19+
};
20+
export const getCurrentHub = jest.fn(() => fakeHub);
21+
export const startTransaction = jest.fn(_ => fakeTransaction);
1122
export const captureException = jest.fn();
1223
export const captureMessage = jest.fn();
1324
export const withScope = jest.fn(cb => cb(fakeScope));
1425
export const flush = jest.fn(() => Promise.resolve());
1526

1627
export const resetMocks = (): void => {
28+
fakeTransaction.finish.mockClear();
29+
fakeParentScope.setSpan.mockClear();
30+
fakeHub.configureScope.mockClear();
31+
1732
fakeScope.addEventProcessor.mockClear();
1833
fakeScope.setTransactionName.mockClear();
1934
fakeScope.setTag.mockClear();
2035
fakeScope.setContext.mockClear();
2136

37+
getCurrentHub.mockClear();
38+
startTransaction.mockClear();
2239
captureException.mockClear();
2340
captureMessage.mockClear();
2441
withScope.mockClear();

‎packages/serverless/test/awslambda.test.ts

Lines changed: 51 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -51,17 +51,11 @@ describe('AWSLambda', () => {
5151
test('flushTimeout', async () => {
5252
expect.assertions(1);
5353

54-
const error = new Error('wat');
55-
const handler = () => {
56-
throw error;
57-
};
54+
const handler = () => {};
5855
const wrappedHandler = wrapHandler(handler, { flushTimeout: 1337 });
5956

60-
try {
61-
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
62-
} catch (e) {
63-
expect(Sentry.flush).toBeCalledWith(1337);
64-
}
57+
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
58+
expect(Sentry.flush).toBeCalledWith(1337);
6559
});
6660

6761
test('rethrowAfterCapture', async () => {
@@ -146,18 +140,24 @@ describe('AWSLambda', () => {
146140

147141
describe('wrapHandler() on sync handler', () => {
148142
test('successful execution', async () => {
149-
expect.assertions(1);
143+
expect.assertions(5);
150144

151145
const handler: Handler = (_event, _context, callback) => {
152146
callback(null, 42);
153147
};
154148
const wrappedHandler = wrapHandler(handler);
155149
const rv = await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
156150
expect(rv).toStrictEqual(42);
151+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
152+
// @ts-ignore see "Why @ts-ignore" note
153+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
154+
// @ts-ignore see "Why @ts-ignore" note
155+
expect(Sentry.fakeTransaction.finish).toBeCalled();
156+
expect(Sentry.flush).toBeCalledWith(2000);
157157
});
158158

159159
test('unsuccessful execution', async () => {
160-
expect.assertions(2);
160+
expect.assertions(5);
161161

162162
const error = new Error('sorry');
163163
const handler: Handler = (_event, _context, callback) => {
@@ -168,7 +168,12 @@ describe('AWSLambda', () => {
168168
try {
169169
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
170170
} catch (e) {
171+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
172+
// @ts-ignore see "Why @ts-ignore" note
173+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
171174
expect(Sentry.captureException).toBeCalledWith(error);
175+
// @ts-ignore see "Why @ts-ignore" note
176+
expect(Sentry.fakeTransaction.finish).toBeCalled();
172177
expect(Sentry.flush).toBeCalledWith(2000);
173178
}
174179
});
@@ -186,7 +191,7 @@ describe('AWSLambda', () => {
186191
});
187192

188193
test('capture error', async () => {
189-
expect.assertions(2);
194+
expect.assertions(5);
190195

191196
const error = new Error('wat');
192197
const handler: Handler = (_event, _context, _callback) => {
@@ -197,22 +202,33 @@ describe('AWSLambda', () => {
197202
try {
198203
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
199204
} catch (e) {
200-
expect(Sentry.captureException).toBeCalled();
205+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
206+
// @ts-ignore see "Why @ts-ignore" note
207+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
208+
expect(Sentry.captureException).toBeCalledWith(e);
209+
// @ts-ignore see "Why @ts-ignore" note
210+
expect(Sentry.fakeTransaction.finish).toBeCalled();
201211
expect(Sentry.flush).toBeCalled();
202212
}
203213
});
204214
});
205215

206216
describe('wrapHandler() on async handler', () => {
207217
test('successful execution', async () => {
208-
expect.assertions(1);
218+
expect.assertions(5);
209219

210220
const handler: Handler = async (_event, _context) => {
211221
return 42;
212222
};
213223
const wrappedHandler = wrapHandler(handler);
214224
const rv = await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
215225
expect(rv).toStrictEqual(42);
226+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
227+
// @ts-ignore see "Why @ts-ignore" note
228+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
229+
// @ts-ignore see "Why @ts-ignore" note
230+
expect(Sentry.fakeTransaction.finish).toBeCalled();
231+
expect(Sentry.flush).toBeCalled();
216232
});
217233

218234
test('event and context are correctly passed to the original handler', async () => {
@@ -227,7 +243,7 @@ describe('AWSLambda', () => {
227243
});
228244

229245
test('capture error', async () => {
230-
expect.assertions(2);
246+
expect.assertions(5);
231247

232248
const error = new Error('wat');
233249
const handler: Handler = async (_event, _context) => {
@@ -238,22 +254,33 @@ describe('AWSLambda', () => {
238254
try {
239255
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
240256
} catch (e) {
241-
expect(Sentry.captureException).toBeCalled();
257+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
258+
// @ts-ignore see "Why @ts-ignore" note
259+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
260+
expect(Sentry.captureException).toBeCalledWith(error);
261+
// @ts-ignore see "Why @ts-ignore" note
262+
expect(Sentry.fakeTransaction.finish).toBeCalled();
242263
expect(Sentry.flush).toBeCalled();
243264
}
244265
});
245266
});
246267

247268
describe('wrapHandler() on async handler with a callback method (aka incorrect usage)', () => {
248269
test('successful execution', async () => {
249-
expect.assertions(1);
270+
expect.assertions(5);
250271

251272
const handler: Handler = async (_event, _context, _callback) => {
252273
return 42;
253274
};
254275
const wrappedHandler = wrapHandler(handler);
255276
const rv = await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
256277
expect(rv).toStrictEqual(42);
278+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
279+
// @ts-ignore see "Why @ts-ignore" note
280+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
281+
// @ts-ignore see "Why @ts-ignore" note
282+
expect(Sentry.fakeTransaction.finish).toBeCalled();
283+
expect(Sentry.flush).toBeCalled();
257284
});
258285

259286
test('event and context are correctly passed to the original handler', async () => {
@@ -268,7 +295,7 @@ describe('AWSLambda', () => {
268295
});
269296

270297
test('capture error', async () => {
271-
expect.assertions(2);
298+
expect.assertions(5);
272299

273300
const error = new Error('wat');
274301
const handler: Handler = async (_event, _context, _callback) => {
@@ -279,7 +306,12 @@ describe('AWSLambda', () => {
279306
try {
280307
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
281308
} catch (e) {
282-
expect(Sentry.captureException).toBeCalled();
309+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
310+
// @ts-ignore see "Why @ts-ignore" note
311+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
312+
expect(Sentry.captureException).toBeCalledWith(error);
313+
// @ts-ignore see "Why @ts-ignore" note
314+
expect(Sentry.fakeTransaction.finish).toBeCalled();
283315
expect(Sentry.flush).toBeCalled();
284316
}
285317
});

0 commit comments

Comments
 (0)