Skip to content

Commit 99dace1

Browse files
jj22eejj22ee
committed
Support Trace Context extraction from Lambda Context object, and respect user-configured OTEL_PROPAGATORS
1 parent 74a1aa3 commit 99dace1

File tree

2 files changed

+205
-61
lines changed

2 files changed

+205
-61
lines changed

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

Lines changed: 37 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
import {
55
diag,
6-
isSpanContextValid,
76
Context as OtelContext,
87
context as otelContext,
98
propagation,
@@ -22,7 +21,7 @@ import {
2221
NormalizedRequest,
2322
NormalizedResponse,
2423
} from '@opentelemetry/instrumentation-aws-sdk';
25-
import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
24+
import { AWSXRAY_TRACE_ID_HEADER } from '@opentelemetry/propagator-aws-xray';
2625
import { APIGatewayProxyEventHeaders, Context } from 'aws-lambda';
2726
import { AWS_ATTRIBUTE_KEYS } from '../aws-attribute-keys';
2827
import { RequestMetadata } from '../third-party/otel/aws/services/ServiceExtension';
@@ -42,13 +41,26 @@ import { suppressTracing } from '@opentelemetry/core';
4241
export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
4342
export const AWSXRAY_TRACE_ID_HEADER_CAPITALIZED = 'X-Amzn-Trace-Id';
4443

45-
const awsPropagator = new AWSXRayPropagator();
46-
export const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
47-
keys(carrier: any): string[] {
48-
return Object.keys(carrier);
44+
interface LambdaCarrier {
45+
lambdaEventHeaders: APIGatewayProxyEventHeaders;
46+
lambdaContext?: Context;
47+
}
48+
49+
export const headerGetter: TextMapGetter<LambdaCarrier> = {
50+
keys(carrier: LambdaCarrier): string[] {
51+
const carrierKeys = Object.keys(carrier.lambdaEventHeaders);
52+
if (carrier.lambdaContext && typeof (carrier.lambdaContext as any)[lambdaContextXrayTraceIdKey] === 'string') {
53+
carrierKeys.push(AWSXRAY_TRACE_ID_HEADER);
54+
}
55+
return carrierKeys;
4956
},
50-
get(carrier: any, key: string) {
51-
return carrier[key];
57+
get(carrier: LambdaCarrier, key: string) {
58+
if (key === AWSXRAY_TRACE_ID_HEADER) {
59+
if (carrier.lambdaContext && typeof (carrier.lambdaContext as any)[lambdaContextXrayTraceIdKey] === 'string') {
60+
return (carrier.lambdaContext as any)[lambdaContextXrayTraceIdKey];
61+
}
62+
}
63+
return carrier.lambdaEventHeaders[key];
5264
},
5365
};
5466

@@ -94,29 +106,29 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])
94106

95107
/*
96108
* 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.
109+
* It first attempts to extract the trace context from the Lambda Handler Context object (_handlerContext.xRayTraceId)
110+
* If above approach fails, attempt to extract the trace context from the AWS X-Ray header, which is stored in the Lambda environment variables.
98111
* If a valid span context is extracted from the environment, it uses this as the parent context for the function's tracing.
99112
* If the X-Ray header is missing or invalid, it falls back to extracting trace context from the Lambda handler's event headers.
100113
* If neither approach succeeds, it defaults to using the root Otel context, ensuring the function is still instrumented for tracing.
101114
*/
115+
const lambdaContextXrayTraceIdKey = 'xRayTraceId';
102116
export const customExtractor = (event: any, _handlerContext: Context): OtelContext => {
103-
let parent: OtelContext | undefined = undefined;
104-
const lambdaTraceHeader = process.env[traceContextEnvironmentKey];
105-
if (lambdaTraceHeader) {
106-
parent = awsPropagator.extract(
107-
otelContext.active(),
108-
{ [AWSXRAY_TRACE_ID_HEADER]: lambdaTraceHeader },
109-
headerGetter
110-
);
111-
}
112-
if (parent) {
113-
const spanContext = trace.getSpan(parent)?.spanContext();
114-
if (spanContext && isSpanContextValid(spanContext)) {
115-
return parent;
116-
}
117-
}
117+
const xrayTraceIdFromLambdaEnv = process.env[traceContextEnvironmentKey];
118+
118119
const httpHeaders = event.headers || {};
119-
const extractedContext = propagation.extract(otelContext.active(), httpHeaders, headerGetter);
120+
if (xrayTraceIdFromLambdaEnv) {
121+
httpHeaders[AWSXRAY_TRACE_ID_HEADER] = xrayTraceIdFromLambdaEnv;
122+
}
123+
124+
const extractedContext = propagation.extract(
125+
otelContext.active(),
126+
{
127+
lambdaEventHeaders: httpHeaders,
128+
lambdaContext: _handlerContext,
129+
},
130+
headerGetter
131+
);
120132
if (trace.getSpan(extractedContext)?.spanContext()) {
121133
return extractedContext;
122134
}

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

Lines changed: 168 additions & 36 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,14 +36,14 @@ 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';
4242
import { getTestSpans } from '@opentelemetry/contrib-test-utils';
4343
import { instrumentationConfigs } from '../../src/register';
4444
import { LoggerProvider } from '@opentelemetry/api-logs';
4545
import { STS } from '@aws-sdk/client-sts';
46+
import { getPropagator } from '@opentelemetry/auto-configuration-propagators';
4647

4748
// It is assumed that bedrock.test.ts has already registered the
4849
// necessary instrumentations for testing by calling:
@@ -71,6 +72,7 @@ const mockHeaders = {
7172
'x-test-header': 'test-value',
7273
'content-type': 'application/json',
7374
};
75+
const mockCarrier = { lambdaEventHeaders: mockHeaders };
7476

7577
const UNPATCHED_INSTRUMENTATIONS: Instrumentation[] = getNodeAutoInstrumentations(instrumentationConfigs);
7678

@@ -965,22 +967,28 @@ describe('InstrumentationPatchTest', () => {
965967
});
966968

967969
describe('customExtractor', () => {
970+
let lambdaPropagator: api.TextMapPropagator;
968971
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';
972+
const MOCK_XRAY_TRACE_ID_0 = '8a3c0000d188f8fa79d48a391a770000';
973+
const MOCK_XRAY_TRACE_ID_1 = '8a3c0001d188f8fa79d48a391a770001';
974+
const MOCK_XRAY_TRACE_ID_STR_0 = '8a3c0000-d188f8fa79d48a391a770000';
975+
const MOCK_XRAY_TRACE_ID_STR_1 = '8a3c0001-d188f8fa79d48a391a770001';
976+
const MOCK_XRAY_PARENT_SPAN_ID_0 = '53995c3f42cd0000';
977+
const MOCK_XRAY_PARENT_SPAN_ID_1 = '53995c3f42cd0001';
972978
const MOCK_XRAY_LAMBDA_LINEAGE = 'Lineage=01cfa446:0';
973979

974980
const TRACE_ID_VERSION = '1'; // Assuming TRACE_ID_VERSION is defined somewhere in the code
975981

976982
// 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}`;
983+
const MOCK_XRAY_TRACE_CONTEXT_0_COMMON = `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR_0};Parent=${MOCK_XRAY_PARENT_SPAN_ID_0}`;
984+
const MOCK_XRAY_TRACE_CONTEXT_1_COMMON = `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR_1};Parent=${MOCK_XRAY_PARENT_SPAN_ID_1}`;
978985

979986
// 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}`;
987+
const MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_0_COMMON};Sampled=1;${MOCK_XRAY_LAMBDA_LINEAGE}`;
988+
const MOCK_XRAY_TRACE_CONTEXT_1_UNSAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_1_COMMON};Sampled=0;${MOCK_XRAY_LAMBDA_LINEAGE}`;
981989
// 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}`
990+
// `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR_0.slice(0, TRACE_ID_FIRST_PART_LENGTH)}` +
991+
// `-${MOCK_XRAY_TRACE_ID_STR_0.slice(TRACE_ID_FIRST_PART_LENGTH)};${MOCK_XRAY_LAMBDA_LINEAGE}`
984992
// );
985993

986994
// Create the W3C Trace Context (Sampled)
@@ -991,11 +999,22 @@ describe('customExtractor', () => {
991999
const MOCK_W3C_TRACE_STATE_VALUE = 'test_value';
9921000
const MOCK_TRACE_STATE = `${MOCK_W3C_TRACE_STATE_KEY}=${MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2`;
9931001

994-
let awsPropagatorStub: SinonStub;
995-
let traceGetSpanStub: SinonStub;
996-
// let propagationStub: SinonStub;
1002+
let awsPropagatorSpy: sinon.SinonSpy;
1003+
let traceGetSpanSpy: sinon.SinonSpy;
1004+
let propagationStub: sinon.SinonStub;
1005+
1006+
before(() => {
1007+
process.env.OTEL_PROPAGATORS = 'baggage,xray,tracecontext';
1008+
lambdaPropagator = getPropagator();
1009+
delete process.env.OTEL_PROPAGATORS;
1010+
});
9971011

9981012
beforeEach(() => {
1013+
propagationStub = sinon
1014+
.stub(propagation, 'extract')
1015+
.callsFake((context: OtelContext, carrier: unknown, getter?: api.TextMapGetter<unknown> | undefined) => {
1016+
return lambdaPropagator.extract(context, carrier, getter!);
1017+
});
9991018
// Clear environment variables before each test
10001019
delete process.env[traceContextEnvironmentKey];
10011020
});
@@ -1005,41 +1024,145 @@ describe('customExtractor', () => {
10051024
sinon.restore();
10061025
});
10071026

1008-
it('should extract context from lambda trace header when present', () => {
1009-
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_SAMPLED;
1010-
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
1027+
it('should extract context from handler context xRayTraceId property when present', () => {
1028+
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1029+
const mockHandlerContext = {
1030+
xRayTraceId: mockLambdaTraceHeader,
1031+
} as unknown as Context;
1032+
1033+
// Mock of the Span Context for validation
1034+
const mockParentSpanContext: api.SpanContext = {
1035+
isRemote: true,
1036+
traceFlags: 1,
1037+
traceId: MOCK_XRAY_TRACE_ID_0,
1038+
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
1039+
};
10111040

1012-
const mockParentContext = {} as OtelContext;
1041+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1042+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1043+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
10131044

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-
};
1045+
// Call the customExtractor function
1046+
const event = { headers: {} };
1047+
const result = customExtractor(event, mockHandlerContext);
10211048

1022-
// Stub awsPropagator.extract to return the mockParentContext
1023-
awsPropagatorStub = sinon.stub(AWSXRayPropagator.prototype, 'extract').returns(mockParentContext);
1049+
// Assertions
1050+
expect(awsPropagatorSpy.calledOnce).toBe(true);
1051+
expect(
1052+
awsPropagatorSpy.calledWith(
1053+
sinon.match.any,
1054+
{
1055+
lambdaEventHeaders: event.headers,
1056+
lambdaContext: mockHandlerContext,
1057+
},
1058+
sinon.match.any
1059+
)
1060+
).toBe(true);
1061+
expect(traceGetSpanSpy.calledOnce).toBe(true);
1062+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
1063+
});
1064+
1065+
it('should prioritize extract context from handler context xRayTraceId property instead of environment variable', () => {
1066+
const mockLambdaContextTraceHeader = MOCK_XRAY_TRACE_CONTEXT_1_UNSAMPLED;
1067+
const mockLambdEnvVarTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1068+
1069+
process.env[traceContextEnvironmentKey] = mockLambdEnvVarTraceHeader;
1070+
1071+
const mockHandlerContext = {
1072+
xRayTraceId: mockLambdaContextTraceHeader,
1073+
} as unknown as Context;
1074+
1075+
// Mock of the Span Context for validation
1076+
const mockParentSpanContext: api.SpanContext = {
1077+
isRemote: true,
1078+
traceFlags: 0,
1079+
traceId: MOCK_XRAY_TRACE_ID_1,
1080+
spanId: MOCK_XRAY_PARENT_SPAN_ID_1,
1081+
};
10241082

1025-
// Stub trace.getSpan to return the mock span
1026-
traceGetSpanStub = sinon.stub(trace, 'getSpan').returns(mockSpan as Span);
1083+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1084+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1085+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
10271086

10281087
// Call the customExtractor function
10291088
const event = { headers: {} };
1089+
const result = customExtractor(event, mockHandlerContext);
1090+
1091+
// Assertions
1092+
expect(awsPropagatorSpy.calledOnce).toBe(true);
1093+
expect(
1094+
awsPropagatorSpy.calledWith(
1095+
sinon.match.any,
1096+
{
1097+
lambdaEventHeaders: event.headers,
1098+
lambdaContext: mockHandlerContext,
1099+
},
1100+
sinon.match.any
1101+
)
1102+
).toBe(true);
1103+
expect(traceGetSpanSpy.calledOnce).toBe(true);
1104+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
1105+
});
1106+
1107+
it('should fallback to environment variable when handler context xRayTraceId is not available', () => {
1108+
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1109+
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
1110+
1111+
// Mock of the Span Context for validation
1112+
const mockParentSpanContext: api.SpanContext = {
1113+
isRemote: true,
1114+
traceFlags: 1,
1115+
traceId: MOCK_XRAY_TRACE_ID_0,
1116+
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
1117+
};
1118+
1119+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1120+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1121+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
1122+
1123+
// Call the customExtractor function with handler context without xRayTraceId
1124+
const event = { headers: {} };
10301125
const result = customExtractor(event, {} as Context);
10311126

10321127
// Assertions
1033-
expect(awsPropagatorStub.calledOnce).toBe(true);
1128+
expect(awsPropagatorSpy.calledOnce).toBe(true);
10341129
expect(
1035-
awsPropagatorStub.calledWith(
1130+
awsPropagatorSpy.calledWith(
10361131
sinon.match.any,
1037-
{ [AWSXRAY_TRACE_ID_HEADER]: mockLambdaTraceHeader },
1132+
{
1133+
lambdaEventHeaders: event.headers,
1134+
lambdaContext: {},
1135+
},
10381136
sinon.match.any
10391137
)
10401138
).toBe(true);
1041-
expect(traceGetSpanStub.calledOnce).toBe(true);
1042-
expect(result).toEqual(mockParentContext); // Should return the parent context when valid
1139+
expect(traceGetSpanSpy.calledOnce).toBe(true);
1140+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
1141+
});
1142+
1143+
it('should handle undefined handler context and fallback to environment variable', () => {
1144+
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
1145+
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
1146+
1147+
// Mock of the Span Context for validation
1148+
const mockParentSpanContext: api.SpanContext = {
1149+
isRemote: true,
1150+
traceFlags: 1,
1151+
traceId: MOCK_XRAY_TRACE_ID_0,
1152+
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
1153+
};
1154+
1155+
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
1156+
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
1157+
traceGetSpanSpy = sinon.spy(trace, 'getSpan');
1158+
1159+
// Call the customExtractor function with undefined handler context
1160+
const event = { headers: {} };
1161+
const result = customExtractor(event, undefined as any);
1162+
1163+
// Should only be called once (for environment variable)
1164+
expect(awsPropagatorSpy.calledOnce).toBe(true);
1165+
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the valid parent context
10431166
});
10441167

10451168
it('should extract context from HTTP headers when lambda trace header is not present', () => {
@@ -1056,27 +1179,36 @@ describe('customExtractor', () => {
10561179
}, // Empty function that returns undefined
10571180
} as unknown as OtelContext;
10581181

1059-
const propagationStub = sinon.stub(propagation, 'extract').returns(mockExtractedContext);
1182+
propagationStub.returns(mockExtractedContext);
10601183

10611184
// Call the customExtractor function
10621185
const mockHttpHeaders = event.headers;
10631186
customExtractor(event, {} as Context);
10641187

1065-
expect(propagationStub.calledWith(sinon.match.any, mockHttpHeaders, sinon.match.any)).toBe(true);
1188+
expect(
1189+
propagationStub.calledWith(
1190+
sinon.match.any,
1191+
{
1192+
lambdaEventHeaders: mockHttpHeaders,
1193+
lambdaContext: {},
1194+
},
1195+
sinon.match.any
1196+
)
1197+
).toBe(true);
10661198
});
10671199

10681200
it('should return all header keys from the carrier', () => {
1069-
const keys = headerGetter.keys(mockHeaders);
1201+
const keys = headerGetter.keys(mockCarrier);
10701202
expect(keys).toEqual(['x-test-header', 'content-type']);
10711203
});
10721204

10731205
it('should return the correct header value for a given key', () => {
1074-
const headerValue = headerGetter.get(mockHeaders, 'x-test-header');
1206+
const headerValue = headerGetter.get(mockCarrier, 'x-test-header');
10751207
expect(headerValue).toBe('test-value');
10761208
});
10771209

10781210
it('should return undefined for a key that does not exist', () => {
1079-
const headerValue = headerGetter.get(mockHeaders, 'non-existent-header');
1211+
const headerValue = headerGetter.get(mockCarrier, 'non-existent-header');
10801212
expect(headerValue).toBeUndefined();
10811213
});
10821214
});

0 commit comments

Comments
 (0)