Skip to content

Commit a863763

Browse files
marshall-leemarshall-lee
committed
Enable tracing for AWSLambda.
1 parent fdfb63f commit a863763

File tree

3 files changed

+139
-34
lines changed

3 files changed

+139
-34
lines changed

‎packages/serverless/src/awslambda.ts‎

Lines changed: 49 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,15 @@
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';
12+
import { Transaction } from '@sentry/types';
213
import { addExceptionMechanism } from '@sentry/utils';
314
// 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
415
// eslint-disable-next-line import/no-unresolved
@@ -21,7 +32,8 @@ export type AsyncHandler<T extends Handler> = (
2132
context: Parameters<T>[1],
2233
) => Promise<NonNullable<Parameters<Parameters<T>[2]>[1]>>;
2334

24-
interface WrapperOptions {
35+
export interface WrapperOptions {
36+
tracing: boolean;
2537
flushTimeout: number;
2638
rethrowAfterCapture: boolean;
2739
callbackWaitsForEmptyEventLoop: boolean;
@@ -98,32 +110,35 @@ function enhanceScopeWithEnvironmentData(scope: Scope, context: Context): void {
98110
}
99111

100112
/**
101-
* Capture, flush the result down the network stream and await the response.
113+
* Capture exception with a a context.
102114
*
103115
* @param e exception to be captured
104-
* @param options WrapperOptions
116+
* @param context Context
105117
*/
106-
function captureExceptionAsync(e: unknown, context: Context, options: Partial<WrapperOptions>): Promise<boolean> {
118+
function captureExceptionWithContext(e: unknown, context: Context): void {
107119
withScope(scope => {
108120
addServerlessEventProcessor(scope);
109121
enhanceScopeWithEnvironmentData(scope, context);
110122
captureException(e);
111123
});
112-
return flush(options.flushTimeout);
113124
}
114125

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 => {
126+
/**
127+
* Wraps a lambda handler adding it error capture and tracing capabilities.
128+
*
129+
* @param handler Handler
130+
* @param options Options
131+
* @returns Handler
132+
*/
133+
export function wrapHandler<TEvent, TResult>(handler: Handler, wrapOptions: Partial<WrapperOptions> = {}): Handler {
120134
const options: WrapperOptions = {
135+
tracing: true,
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

@@ -165,11 +180,25 @@ export const wrapHandler = <TEvent = any, TResult = any>(
165180
if (args[0] === null || args[0] === undefined) {
166181
resolve(callback(...args));
167182
} else {
168-
captureExceptionAsync(args[0], context, options).finally(() => reject(callback(...args)));
183+
captureExceptionWithContext(args[0], context);
184+
reject(callback(...args));
169185
}
170186
};
171187
};
172188

189+
let transaction: Transaction | undefined;
190+
191+
if (options.tracing) {
192+
transaction = startTransaction({
193+
name: context.functionName,
194+
op: 'awslambda.handler',
195+
});
196+
// We put the transaction on the scope so users can attach children to it
197+
getCurrentHub().configureScope(scope => {
198+
scope.setSpan(transaction);
199+
});
200+
}
201+
173202
try {
174203
// AWSLambda is like Express. It makes a distinction about handlers based on it's last argument
175204
// async (event) => async handler
@@ -195,10 +224,15 @@ export const wrapHandler = <TEvent = any, TResult = any>(
195224
return handlerRv;
196225
} catch (e) {
197226
clearTimeout(timeoutWarningTimer);
198-
await captureExceptionAsync(e, context, options);
227+
captureExceptionWithContext(e, context);
199228
if (options.rethrowAfterCapture) {
200229
throw e;
201230
}
231+
} finally {
232+
if (transaction) {
233+
transaction.finish();
234+
}
235+
await flush(options.flushTimeout);
202236
}
203237
};
204-
};
238+
}

‎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: 73 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 () => {
@@ -142,22 +136,50 @@ describe('AWSLambda', () => {
142136
// @ts-ignore see "Why @ts-ignore" note
143137
expect(Sentry.fakeScope.setTag).toBeCalledWith('timeout', '1m40s');
144138
});
139+
140+
test('tracing enabled', async () => {
141+
expect.assertions(1);
142+
143+
const handler: Handler = (_event, _context, callback) => {
144+
callback(null, 42);
145+
};
146+
const wrappedHandler = wrapHandler(handler, { tracing: true });
147+
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
148+
expect(Sentry.startTransaction).toBeCalled();
149+
});
150+
151+
test('tracing disabled', async () => {
152+
expect.assertions(1);
153+
154+
const handler: Handler = (_event, _context, callback) => {
155+
callback(null, 42);
156+
};
157+
const wrappedHandler = wrapHandler(handler, { tracing: false });
158+
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
159+
expect(Sentry.startTransaction).not.toBeCalled();
160+
});
145161
});
146162

147163
describe('wrapHandler() on sync handler', () => {
148164
test('successful execution', async () => {
149-
expect.assertions(1);
165+
expect.assertions(5);
150166

151167
const handler: Handler = (_event, _context, callback) => {
152168
callback(null, 42);
153169
};
154170
const wrappedHandler = wrapHandler(handler);
155171
const rv = await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
156172
expect(rv).toStrictEqual(42);
173+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
174+
// @ts-ignore see "Why @ts-ignore" note
175+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
176+
// @ts-ignore see "Why @ts-ignore" note
177+
expect(Sentry.fakeTransaction.finish).toBeCalled();
178+
expect(Sentry.flush).toBeCalledWith(2000);
157179
});
158180

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

162184
const error = new Error('sorry');
163185
const handler: Handler = (_event, _context, callback) => {
@@ -168,7 +190,12 @@ describe('AWSLambda', () => {
168190
try {
169191
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
170192
} catch (e) {
193+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
194+
// @ts-ignore see "Why @ts-ignore" note
195+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
171196
expect(Sentry.captureException).toBeCalledWith(error);
197+
// @ts-ignore see "Why @ts-ignore" note
198+
expect(Sentry.fakeTransaction.finish).toBeCalled();
172199
expect(Sentry.flush).toBeCalledWith(2000);
173200
}
174201
});
@@ -186,7 +213,7 @@ describe('AWSLambda', () => {
186213
});
187214

188215
test('capture error', async () => {
189-
expect.assertions(2);
216+
expect.assertions(5);
190217

191218
const error = new Error('wat');
192219
const handler: Handler = (_event, _context, _callback) => {
@@ -197,22 +224,33 @@ describe('AWSLambda', () => {
197224
try {
198225
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
199226
} catch (e) {
200-
expect(Sentry.captureException).toBeCalled();
227+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
228+
// @ts-ignore see "Why @ts-ignore" note
229+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
230+
expect(Sentry.captureException).toBeCalledWith(e);
231+
// @ts-ignore see "Why @ts-ignore" note
232+
expect(Sentry.fakeTransaction.finish).toBeCalled();
201233
expect(Sentry.flush).toBeCalled();
202234
}
203235
});
204236
});
205237

206238
describe('wrapHandler() on async handler', () => {
207239
test('successful execution', async () => {
208-
expect.assertions(1);
240+
expect.assertions(5);
209241

210242
const handler: Handler = async (_event, _context) => {
211243
return 42;
212244
};
213245
const wrappedHandler = wrapHandler(handler);
214246
const rv = await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
215247
expect(rv).toStrictEqual(42);
248+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
249+
// @ts-ignore see "Why @ts-ignore" note
250+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
251+
// @ts-ignore see "Why @ts-ignore" note
252+
expect(Sentry.fakeTransaction.finish).toBeCalled();
253+
expect(Sentry.flush).toBeCalled();
216254
});
217255

218256
test('event and context are correctly passed to the original handler', async () => {
@@ -227,7 +265,7 @@ describe('AWSLambda', () => {
227265
});
228266

229267
test('capture error', async () => {
230-
expect.assertions(2);
268+
expect.assertions(5);
231269

232270
const error = new Error('wat');
233271
const handler: Handler = async (_event, _context) => {
@@ -238,22 +276,33 @@ describe('AWSLambda', () => {
238276
try {
239277
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
240278
} catch (e) {
241-
expect(Sentry.captureException).toBeCalled();
279+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
280+
// @ts-ignore see "Why @ts-ignore" note
281+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
282+
expect(Sentry.captureException).toBeCalledWith(error);
283+
// @ts-ignore see "Why @ts-ignore" note
284+
expect(Sentry.fakeTransaction.finish).toBeCalled();
242285
expect(Sentry.flush).toBeCalled();
243286
}
244287
});
245288
});
246289

247290
describe('wrapHandler() on async handler with a callback method (aka incorrect usage)', () => {
248291
test('successful execution', async () => {
249-
expect.assertions(1);
292+
expect.assertions(5);
250293

251294
const handler: Handler = async (_event, _context, _callback) => {
252295
return 42;
253296
};
254297
const wrappedHandler = wrapHandler(handler);
255298
const rv = await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
256299
expect(rv).toStrictEqual(42);
300+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
301+
// @ts-ignore see "Why @ts-ignore" note
302+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
303+
// @ts-ignore see "Why @ts-ignore" note
304+
expect(Sentry.fakeTransaction.finish).toBeCalled();
305+
expect(Sentry.flush).toBeCalled();
257306
});
258307

259308
test('event and context are correctly passed to the original handler', async () => {
@@ -268,7 +317,7 @@ describe('AWSLambda', () => {
268317
});
269318

270319
test('capture error', async () => {
271-
expect.assertions(2);
320+
expect.assertions(5);
272321

273322
const error = new Error('wat');
274323
const handler: Handler = async (_event, _context, _callback) => {
@@ -279,7 +328,12 @@ describe('AWSLambda', () => {
279328
try {
280329
await wrappedHandler(fakeEvent, fakeContext, fakeCallback);
281330
} catch (e) {
282-
expect(Sentry.captureException).toBeCalled();
331+
expect(Sentry.startTransaction).toBeCalledWith({ name: 'functionName', op: 'awslambda.handler' });
332+
// @ts-ignore see "Why @ts-ignore" note
333+
expect(Sentry.fakeParentScope.setSpan).toBeCalledWith(Sentry.fakeTransaction);
334+
expect(Sentry.captureException).toBeCalledWith(error);
335+
// @ts-ignore see "Why @ts-ignore" note
336+
expect(Sentry.fakeTransaction.finish).toBeCalled();
283337
expect(Sentry.flush).toBeCalled();
284338
}
285339
});

0 commit comments

Comments
 (0)