Skip to content

Commit f76b467

Browse files
jj22eejj22ee
committed
Support Trace Context extraction from Lambda Context object
1 parent 8661ac1 commit f76b467

File tree

2 files changed

+179
-31
lines changed

2 files changed

+179
-31
lines changed

‎aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,14 +94,33 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])
9494

9595
/*
9696
* This function `customExtractor` is used to extract SpanContext for AWS Lambda functions.
97-
* It first attempts to extract the trace context from the AWS X-Ray header, which is stored in the Lambda environment variables.
97+
* It first attempts to extract the trace context from the Lambda Handler Context object (_handlerContext.xRayTraceId)
98+
* If above approach fails, attempt to extract the trace context from the AWS X-Ray header, which is stored in the Lambda environment variables.
9899
* If a valid span context is extracted from the environment, it uses this as the parent context for the function's tracing.
99100
* If the X-Ray header is missing or invalid, it falls back to extracting trace context from the Lambda handler's event headers.
100101
* If neither approach succeeds, it defaults to using the root Otel context, ensuring the function is still instrumented for tracing.
101102
*/
103+
const lambdaContextXrayTraceIdKey = 'xRayTraceId';
102104
export const customExtractor = (event: any, _handlerContext: Context): OtelContext => {
103105
let parent: OtelContext | undefined = undefined;
104-
const lambdaTraceHeader = process.env[traceContextEnvironmentKey];
106+
107+
let lambdaTraceHeader;
108+
if (_handlerContext && typeof (_handlerContext as any)[lambdaContextXrayTraceIdKey] === 'string') {
109+
lambdaTraceHeader = (_handlerContext as any)[lambdaContextXrayTraceIdKey];
110+
parent = awsPropagator.extract(
111+
otelContext.active(),
112+
{ [AWSXRAY_TRACE_ID_HEADER]: lambdaTraceHeader },
113+
headerGetter
114+
);
115+
}
116+
if (parent) {
117+
const spanContext = trace.getSpan(parent)?.spanContext();
118+
if (spanContext && isSpanContextValid(spanContext)) {
119+
return parent;
120+
}
121+
}
122+
123+
lambdaTraceHeader = process.env[traceContextEnvironmentKey];
105124
if (lambdaTraceHeader) {
106125
parent = awsPropagator.extract(
107126
otelContext.active(),

‎aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts‎

Lines changed: 158 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
SpanStatusCode,
1818
ROOT_CONTEXT,
1919
} from '@opentelemetry/api';
20+
import * as api from '@opentelemetry/api';
2021
import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
2122
import { Instrumentation } from '@opentelemetry/instrumentation';
2223
import { AwsInstrumentation, NormalizedRequest, NormalizedResponse } from '@opentelemetry/instrumentation-aws-sdk';
@@ -35,7 +36,6 @@ import {
3536
import * as sinon from 'sinon';
3637
import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
3738
import { Context } from 'aws-lambda';
38-
import { SinonStub } from 'sinon';
3939
import { Lambda } from '@aws-sdk/client-lambda';
4040
import * as nock from 'nock';
4141
import { ReadableSpan, Span as SDKSpan } from '@opentelemetry/sdk-trace-base';
@@ -966,21 +966,26 @@ describe('InstrumentationPatchTest', () => {
966966

967967
describe('customExtractor', () => {
968968
const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
969-
const MOCK_XRAY_TRACE_ID = '8a3c60f7d188f8fa79d48a391a778fa6';
970-
const MOCK_XRAY_TRACE_ID_STR = '1-8a3c60f7-d188f8fa79d48a391a778fa6';
971-
const MOCK_XRAY_PARENT_SPAN_ID = '53995c3f42cd8ad8';
969+
const MOCK_XRAY_TRACE_ID_0 = '8a3c0000d188f8fa79d48a391a770000';
970+
const MOCK_XRAY_TRACE_ID_1 = '8a3c0001d188f8fa79d48a391a770001';
971+
const MOCK_XRAY_TRACE_ID_STR_0 = '8a3c0000-d188f8fa79d48a391a770000';
972+
const MOCK_XRAY_TRACE_ID_STR_1 = '8a3c0001-d188f8fa79d48a391a770001';
973+
const MOCK_XRAY_PARENT_SPAN_ID_0 = '53995c3f42cd0000';
974+
const MOCK_XRAY_PARENT_SPAN_ID_1 = '53995c3f42cd0001';
972975
const MOCK_XRAY_LAMBDA_LINEAGE = 'Lineage=01cfa446:0';
973976

974977
const TRACE_ID_VERSION = '1'; // Assuming TRACE_ID_VERSION is defined somewhere in the code
975978

976979
// Common part of the XRAY trace context
977-
const MOCK_XRAY_TRACE_CONTEXT_COMMON = `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR};Parent=${MOCK_XRAY_PARENT_SPAN_ID}`;
980+
const MOCK_XRAY_TRACE_CONTEXT_0_COMMON = `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR_0};Parent=${MOCK_XRAY_PARENT_SPAN_ID_0}`;
981+
const MOCK_XRAY_TRACE_CONTEXT_1_COMMON = `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR_1};Parent=${MOCK_XRAY_PARENT_SPAN_ID_1}`;
978982

979983
// Different versions of the XRAY trace context
980-
const MOCK_XRAY_TRACE_CONTEXT_SAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=1;${MOCK_XRAY_LAMBDA_LINEAGE}`;
984+
const MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_0_COMMON};Sampled=1;${MOCK_XRAY_LAMBDA_LINEAGE}`;
985+
const MOCK_XRAY_TRACE_CONTEXT_1_UNSAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_1_COMMON};Sampled=0;${MOCK_XRAY_LAMBDA_LINEAGE}`;
981986
// const MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH = (
982-
// `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR.slice(0, TRACE_ID_FIRST_PART_LENGTH)}` +
983-
// `-${MOCK_XRAY_TRACE_ID_STR.slice(TRACE_ID_FIRST_PART_LENGTH)};${MOCK_XRAY_LAMBDA_LINEAGE}`
987+
// `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR_0.slice(0, TRACE_ID_FIRST_PART_LENGTH)}` +
988+
// `-${MOCK_XRAY_TRACE_ID_STR_0.slice(TRACE_ID_FIRST_PART_LENGTH)};${MOCK_XRAY_LAMBDA_LINEAGE}`
984989
// );
985990

986991
// Create the W3C Trace Context (Sampled)
@@ -991,8 +996,8 @@ describe('customExtractor', () => {
991996
const MOCK_W3C_TRACE_STATE_VALUE = 'test_value';
992997
const MOCK_TRACE_STATE = `${MOCK_W3C_TRACE_STATE_KEY}=${MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2`;
993998

994-
let awsPropagatorStub: SinonStub;
995-
let traceGetSpanStub: SinonStub;
999+
let awsPropagatorSpy: sinon.SinonSpy;
1000+
let traceGetSpanSpy: sinon.SinonSpy;
9961001
// let propagationStub: SinonStub;
9971002

9981003
beforeEach(() => {
@@ -1005,41 +1010,165 @@ describe('customExtractor', () => {
10051010
sinon.restore();
10061011
});
10071012

1008-
it('should extract context from lambda trace header when present', () => {
1009-
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_SAMPLED;
1010-
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
1013+
it('should extract context from handler context xRayTraceId property when present', () => {
1014+
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1015+
const mockHandlerContext = {
1016+
xRayTraceId: mockLambdaTraceHeader,
1017+
} as unknown as Context;
1018+
1019+
// Mock of the Span Context for validation
1020+
const mockParentSpanContext: api.SpanContext = {
1021+
isRemote: true,
1022+
traceFlags: 1,
1023+
traceId: MOCK_XRAY_TRACE_ID_0,
1024+
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
1025+
};
10111026

1012-
const mockParentContext = {} as OtelContext;
1027+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1028+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1029+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
10131030

1014-
// Partial mock of the Span object
1015-
const mockSpan: Partial<Span> = {
1016-
spanContext: sinon.stub().returns({
1017-
traceId: MOCK_XRAY_TRACE_ID,
1018-
spanId: MOCK_XRAY_PARENT_SPAN_ID,
1019-
}),
1020-
};
1031+
// Call the customExtractor function
1032+
const event = { headers: {} };
1033+
const result = customExtractor(event, mockHandlerContext);
1034+
1035+
// Assertions
1036+
expect(awsPropagatorSpy.calledOnce).toBe(true);
1037+
expect(
1038+
awsPropagatorSpy.calledWith(
1039+
sinon.match.any,
1040+
{ [AWSXRAY_TRACE_ID_HEADER]: mockLambdaTraceHeader },
1041+
sinon.match.any
1042+
)
1043+
).toBe(true);
1044+
expect(traceGetSpanSpy.calledOnce).toBe(true);
1045+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
1046+
});
1047+
1048+
it('should prioritize extract context from handler context xRayTraceId property instead of environment variable', () => {
1049+
const mockLambdaContextTraceHeader = MOCK_XRAY_TRACE_CONTEXT_1_UNSAMPLED;
1050+
const mockLambdEnvVarTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
10211051

1022-
// Stub awsPropagator.extract to return the mockParentContext
1023-
awsPropagatorStub = sinon.stub(AWSXRayPropagator.prototype, 'extract').returns(mockParentContext);
1052+
process.env[traceContextEnvironmentKey] = mockLambdEnvVarTraceHeader;
10241053

1025-
// Stub trace.getSpan to return the mock span
1026-
traceGetSpanStub = sinon.stub(trace, 'getSpan').returns(mockSpan as Span);
1054+
const mockHandlerContext = {
1055+
xRayTraceId: mockLambdaContextTraceHeader,
1056+
} as unknown as Context;
1057+
1058+
// Mock of the Span Context for validation
1059+
const mockParentSpanContext: api.SpanContext = {
1060+
isRemote: true,
1061+
traceFlags: 0,
1062+
traceId: MOCK_XRAY_TRACE_ID_1,
1063+
spanId: MOCK_XRAY_PARENT_SPAN_ID_1,
1064+
};
1065+
1066+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1067+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1068+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
10271069

10281070
// Call the customExtractor function
10291071
const event = { headers: {} };
1072+
const result = customExtractor(event, mockHandlerContext);
1073+
1074+
// Assertions
1075+
expect(awsPropagatorSpy.calledOnce).toBe(true);
1076+
expect(
1077+
awsPropagatorSpy.calledWith(
1078+
sinon.match.any,
1079+
{ [AWSXRAY_TRACE_ID_HEADER]: mockLambdaContextTraceHeader },
1080+
sinon.match.any
1081+
)
1082+
).toBe(true);
1083+
expect(traceGetSpanSpy.calledOnce).toBe(true);
1084+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
1085+
});
1086+
1087+
it('should fallback to environment variable when handler context xRayTraceId is not available', () => {
1088+
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1089+
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
1090+
1091+
// Mock of the Span Context for validation
1092+
const mockParentSpanContext: api.SpanContext = {
1093+
isRemote: true,
1094+
traceFlags: 1,
1095+
traceId: MOCK_XRAY_TRACE_ID_0,
1096+
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
1097+
};
1098+
1099+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1100+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1101+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
1102+
1103+
// Call the customExtractor function with handler context without xRayTraceId
1104+
const event = { headers: {} };
10301105
const result = customExtractor(event, {} as Context);
10311106

10321107
// Assertions
1033-
expect(awsPropagatorStub.calledOnce).toBe(true);
1108+
expect(awsPropagatorSpy.calledOnce).toBe(true);
10341109
expect(
1035-
awsPropagatorStub.calledWith(
1110+
awsPropagatorSpy.calledWith(
10361111
sinon.match.any,
10371112
{ [AWSXRAY_TRACE_ID_HEADER]: mockLambdaTraceHeader },
10381113
sinon.match.any
10391114
)
10401115
).toBe(true);
1041-
expect(traceGetSpanStub.calledOnce).toBe(true);
1042-
expect(result).toEqual(mockParentContext); // Should return the parent context when valid
1116+
expect(traceGetSpanSpy.calledOnce).toBe(true);
1117+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
1118+
});
1119+
1120+
it('should handle invalid span context from handler context xRayTraceId and fallback to environment variable', () => {
1121+
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1122+
const mockHandlerContext = {
1123+
xRayTraceId: 'invalid-trace-header',
1124+
} as unknown as Context;
1125+
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
1126+
1127+
// Mock of the Span Context for validation
1128+
const mockParentSpanContext: api.SpanContext = {
1129+
isRemote: true,
1130+
traceFlags: 1,
1131+
traceId: MOCK_XRAY_TRACE_ID_0,
1132+
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
1133+
};
1134+
1135+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1136+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1137+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
1138+
1139+
// Call the customExtractor function
1140+
const event = { headers: {} };
1141+
const result = customExtractor(event, mockHandlerContext);
1142+
1143+
// Assertions - should be called twice (handler context + environment)
1144+
expect(awsPropagatorSpy.calledTwice).toBe(true);
1145+
expect(traceGetSpanSpy.calledTwice).toBe(true);
1146+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the valid parent context
1147+
});
1148+
1149+
it('should handle undefined handler context and fallback to environment variable', () => {
1150+
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1151+
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
1152+
1153+
// Mock of the Span Context for validation
1154+
const mockParentSpanContext: api.SpanContext = {
1155+
isRemote: true,
1156+
traceFlags: 1,
1157+
traceId: MOCK_XRAY_TRACE_ID_0,
1158+
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
1159+
};
1160+
1161+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1162+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1163+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
1164+
1165+
// Call the customExtractor function with undefined handler context
1166+
const event = { headers: {} };
1167+
const result = customExtractor(event, undefined as any);
1168+
1169+
// Should only be called once (for environment variable)
1170+
expect(awsPropagatorSpy.calledOnce).toBe(true);
1171+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the valid parent context
10431172
});
10441173

10451174
it('should extract context from HTTP headers when lambda trace header is not present', () => {

0 commit comments

Comments
 (0)