Skip to content

Commit c82d0fe

Browse files
dreamorosidreamorosi
committed
Revisited middleware implementation
1 parent 019655c commit c82d0fe

File tree

5 files changed

+86
-93
lines changed

5 files changed

+86
-93
lines changed

Diff for: ‎packages/tracing/src/Tracer.ts

+2-2
Original file line numberDiff line numberDiff line change
@@ -321,13 +321,13 @@ class Tracer implements TracerInterface {
321321
return originalMethod?.apply(target, [ event, context, callback ]);
322322
}
323323

324-
return this.provider.captureAsyncFunc(`## ${context.functionName}`, async subsegment => {
324+
return this.provider.captureAsyncFunc(`## ${process.env._HANDLER}`, async subsegment => {
325325
this.annotateColdStart();
326326
this.addServiceNameAnnotation();
327327
let result: unknown;
328328
try {
329329
result = await originalMethod?.apply(target, [ event, context, callback ]);
330-
this.addResponseAsMetadata(result, context.functionName);
330+
this.addResponseAsMetadata(result, process.env._HANDLER);
331331
} catch (error) {
332332
this.addErrorAsMetadata(error as Error);
333333
throw error;

Diff for: ‎packages/tracing/src/middleware/middy.ts

+22-13
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import middy from '@middy/core';
2-
import { Subsegment } from 'aws-xray-sdk-core';
32
import { Tracer } from '../Tracer';
3+
import { Segment, Subsegment } from 'aws-xray-sdk-core';
44

55
/**
66
* A middy middleware automating capture of metadata and annotations on segments or subsegments ofr a Lambda Handler.
@@ -27,32 +27,41 @@ import { Tracer } from '../Tracer';
2727
* @returns middleware object - The middy middleware object
2828
*/
2929
const captureLambdaHandler = (target: Tracer): middy.MiddlewareObj => {
30-
const captureLambdaHandlerBefore = async (request: middy.Request): Promise<void> => {
30+
let lambdaSegment: Subsegment | Segment;
31+
32+
const open = (): void => {
33+
lambdaSegment = target.getSegment();
34+
const handlerSegment = lambdaSegment.addNewSubsegment(`## ${process.env._HANDLER}`);
35+
target.setSegment(handlerSegment);
36+
};
37+
38+
const close = (): void => {
39+
const subsegment = target.getSegment();
40+
subsegment?.close();
41+
target.setSegment(lambdaSegment as Segment);
42+
};
43+
44+
const captureLambdaHandlerBefore = async (_request: middy.Request): Promise<void> => {
3145
if (target.isTracingEnabled()) {
32-
const subsegment = new Subsegment(`## ${request.context.functionName}`);
33-
target.setSegment(subsegment);
46+
open();
3447
target.annotateColdStart();
3548
target.addServiceNameAnnotation();
3649
}
3750
};
38-
51+
3952
const captureLambdaHandlerAfter = async (request: middy.Request): Promise<void> => {
4053
if (target.isTracingEnabled()) {
41-
const subsegment = target.getSegment();
42-
target.addResponseAsMetadata(request.response, request.context.functionName);
43-
subsegment.close();
54+
target.addResponseAsMetadata(request.response, process.env._HANDLER);
55+
close();
4456
}
4557
};
46-
58+
4759
const captureLambdaHandlerError = async (request: middy.Request): Promise<void> => {
4860
if (target.isTracingEnabled()) {
49-
const subsegment = target.getSegment();
5061
target.addErrorAsMetadata(request.error as Error);
51-
5262
// TODO: should this error be thrown?? I.e. should we stop the event flow & return?
5363
// throw request.error;
54-
55-
subsegment?.close();
64+
close();
5665
}
5766
};
5867

Diff for: ‎packages/tracing/tests/helpers/populateEnvironmentVariables.ts

+1
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ process.env.AWS_LAMBDA_FUNCTION_NAME = 'my-lambda-function';
44
process.env.AWS_EXECUTION_ENV = 'nodejs14.x';
55
process.env.AWS_LAMBDA_FUNCTION_MEMORY_SIZE = '128';
66
process.env.AWS_REGION = 'eu-central-1';
7+
process.env._HANDLER = 'index.handler';
78

89
// Powertools variables
910
process.env.POWERTOOLS_SERVICE_NAME = 'hello-world';

Diff for: ‎packages/tracing/tests/unit/Tracer.test.ts

+17-17
Original file line numberDiff line numberDiff line change
@@ -568,7 +568,7 @@ describe('Class: Tracer', () => {
568568
// Prepare
569569
process.env.POWERTOOLS_TRACER_CAPTURE_RESPONSE = 'false';
570570
const tracer: Tracer = new Tracer();
571-
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
571+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## index.handler');
572572
jest.spyOn(tracer.provider, 'getSegment').mockImplementation(() => newSubsegment);
573573
setContextMissingStrategy(() => null);
574574
const captureAsyncFuncSpy = jest.spyOn(tracer.provider, 'captureAsyncFunc');
@@ -599,7 +599,7 @@ describe('Class: Tracer', () => {
599599

600600
// Prepare
601601
const tracer: Tracer = new Tracer();
602-
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
602+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## index.handler');
603603
jest.spyOn(tracer.provider, 'getSegment')
604604
.mockImplementation(() => newSubsegment);
605605
setContextMissingStrategy(() => null);
@@ -622,12 +622,12 @@ describe('Class: Tracer', () => {
622622

623623
// Assess
624624
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
625-
expect(captureAsyncFuncSpy).toHaveBeenCalledWith('## foo-bar-function', expect.anything());
625+
expect(captureAsyncFuncSpy).toHaveBeenCalledWith('## index.handler', expect.anything());
626626
expect(newSubsegment).toEqual(expect.objectContaining({
627-
name: '## foo-bar-function',
627+
name: '## index.handler',
628628
metadata: {
629629
'hello-world': {
630-
'foo-bar-function response': {
630+
'index.handler response': {
631631
foo: 'bar',
632632
},
633633
},
@@ -641,7 +641,7 @@ describe('Class: Tracer', () => {
641641
// Prepare
642642
process.env.POWERTOOLS_TRACER_CAPTURE_ERROR = 'false';
643643
const tracer: Tracer = new Tracer();
644-
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
644+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## index.handler');
645645
jest.spyOn(tracer.provider, 'getSegment')
646646
.mockImplementation(() => newSubsegment);
647647
setContextMissingStrategy(() => null);
@@ -663,7 +663,7 @@ describe('Class: Tracer', () => {
663663
await expect(new Lambda().handler({}, context, () => console.log('Lambda invoked!'))).rejects.toThrowError(Error);
664664
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
665665
expect(newSubsegment).toEqual(expect.objectContaining({
666-
name: '## foo-bar-function',
666+
name: '## index.handler',
667667
}));
668668
expect('cause' in newSubsegment).toBe(false);
669669
expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
@@ -678,7 +678,7 @@ describe('Class: Tracer', () => {
678678

679679
// Prepare
680680
const tracer: Tracer = new Tracer();
681-
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
681+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## index.handler');
682682
jest.spyOn(tracer.provider, 'getSegment')
683683
.mockImplementation(() => newSubsegment);
684684
setContextMissingStrategy(() => null);
@@ -699,7 +699,7 @@ describe('Class: Tracer', () => {
699699
await expect(new Lambda().handler({}, context, () => console.log('Lambda invoked!'))).rejects.toThrowError(Error);
700700
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
701701
expect(newSubsegment).toEqual(expect.objectContaining({
702-
name: '## foo-bar-function',
702+
name: '## index.handler',
703703
}));
704704
expect('cause' in newSubsegment).toBe(true);
705705
expect(addErrorSpy).toHaveBeenCalledTimes(1);
@@ -712,8 +712,8 @@ describe('Class: Tracer', () => {
712712

713713
// Prepare
714714
const tracer: Tracer = new Tracer();
715-
const newSubsegmentFirstInvocation: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
716-
const newSubsegmentSecondInvocation: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
715+
const newSubsegmentFirstInvocation: Segment | Subsegment | undefined = new Subsegment('## index.handler');
716+
const newSubsegmentSecondInvocation: Segment | Subsegment | undefined = new Subsegment('## index.handler');
717717
jest.spyOn(tracer.provider, 'getSegment')
718718
.mockImplementationOnce(() => newSubsegmentFirstInvocation)
719719
.mockImplementation(() => newSubsegmentSecondInvocation);
@@ -739,21 +739,21 @@ describe('Class: Tracer', () => {
739739

740740
// Assess
741741
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(2);
742-
expect(captureAsyncFuncSpy).toHaveBeenCalledWith('## foo-bar-function', expect.anything());
742+
expect(captureAsyncFuncSpy).toHaveBeenCalledWith('## index.handler', expect.anything());
743743
expect(putAnnotationSpy.mock.calls.filter(call =>
744744
call[0] === 'ColdStart'
745745
)).toEqual([
746746
[ 'ColdStart', true ],
747747
[ 'ColdStart', false ],
748748
]);
749749
expect(newSubsegmentFirstInvocation).toEqual(expect.objectContaining({
750-
name: '## foo-bar-function',
750+
name: '## index.handler',
751751
annotations: expect.objectContaining({
752752
'ColdStart': true,
753753
})
754754
}));
755755
expect(newSubsegmentSecondInvocation).toEqual(expect.objectContaining({
756-
name: '## foo-bar-function',
756+
name: '## index.handler',
757757
annotations: expect.objectContaining({
758758
'ColdStart': false,
759759
})
@@ -765,7 +765,7 @@ describe('Class: Tracer', () => {
765765

766766
// Prepare
767767
const tracer: Tracer = new Tracer();
768-
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
768+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## index.handler');
769769
jest.spyOn(tracer.provider, 'getSegment')
770770
.mockImplementation(() => newSubsegment);
771771
setContextMissingStrategy(() => null);
@@ -788,9 +788,9 @@ describe('Class: Tracer', () => {
788788

789789
// Assess
790790
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
791-
expect(captureAsyncFuncSpy).toHaveBeenCalledWith('## foo-bar-function', expect.anything());
791+
expect(captureAsyncFuncSpy).toHaveBeenCalledWith('## index.handler', expect.anything());
792792
expect(newSubsegment).toEqual(expect.objectContaining({
793-
name: '## foo-bar-function',
793+
name: '## index.handler',
794794
annotations: expect.objectContaining({
795795
'Service': 'hello-world',
796796
})

0 commit comments

Comments
 (0)