Skip to content

Support Trace Context extraction from Lambda Context object #254

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

Support Trace Context extraction from Lambda Context object #254

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f76b467
Support Trace Context extraction from Lambda Context object
jj22ee Aug 21, 2025
4f2b29c
Merge branch 'main' into xray-header-from-context-support
jj22ee Sep 4, 2025
255b7b3
Merge branch 'main' into xray-header-from-context-support
jj22ee Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,14 +94,33 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])

/*
* This function `customExtractor` is used to extract SpanContext for AWS Lambda functions.
* It first attempts to extract the trace context from the AWS X-Ray header, which is stored in the Lambda environment variables.
* It first attempts to extract the trace context from the Lambda Handler Context object (_handlerContext.xRayTraceId)
* If above approach fails, attempt to extract the trace context from the AWS X-Ray header, which is stored in the Lambda environment variables.
* If a valid span context is extracted from the environment, it uses this as the parent context for the function's tracing.
* If the X-Ray header is missing or invalid, it falls back to extracting trace context from the Lambda handler's event headers.
* If neither approach succeeds, it defaults to using the root Otel context, ensuring the function is still instrumented for tracing.
*/
const lambdaContextXrayTraceIdKey = 'xRayTraceId';
Copy link
Contributor Author

@jj22ee jj22ee Aug 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Later we can remove this key once the Context type from aws-lambda package has this.

export const customExtractor = (event: any, _handlerContext: Context): OtelContext => {
let parent: OtelContext | undefined = undefined;
const lambdaTraceHeader = process.env[traceContextEnvironmentKey];

let lambdaTraceHeader;
if (_handlerContext && typeof (_handlerContext as any)[lambdaContextXrayTraceIdKey] === 'string') {
lambdaTraceHeader = (_handlerContext as any)[lambdaContextXrayTraceIdKey];
parent = awsPropagator.extract(
otelContext.active(),
{ [AWSXRAY_TRACE_ID_HEADER]: lambdaTraceHeader },
headerGetter
);
}
if (parent) {
const spanContext = trace.getSpan(parent)?.spanContext();
if (spanContext && isSpanContextValid(spanContext)) {
return parent;
}
}

lambdaTraceHeader = process.env[traceContextEnvironmentKey];
if (lambdaTraceHeader) {
parent = awsPropagator.extract(
otelContext.active(),
Expand Down
187 changes: 158 additions & 29 deletions aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import {
SpanStatusCode,
ROOT_CONTEXT,
} from '@opentelemetry/api';
import * as api from '@opentelemetry/api';
import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
import { Instrumentation } from '@opentelemetry/instrumentation';
import { AwsInstrumentation, NormalizedRequest, NormalizedResponse } from '@opentelemetry/instrumentation-aws-sdk';
Expand All @@ -35,7 +36,6 @@ import {
import * as sinon from 'sinon';
import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
import { Context } from 'aws-lambda';
import { SinonStub } from 'sinon';
import { Lambda } from '@aws-sdk/client-lambda';
import * as nock from 'nock';
import { ReadableSpan, Span as SDKSpan } from '@opentelemetry/sdk-trace-base';
Expand Down Expand Up @@ -966,21 +966,26 @@ describe('InstrumentationPatchTest', () => {

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

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

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

// Different versions of the XRAY trace context
const MOCK_XRAY_TRACE_CONTEXT_SAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=1;${MOCK_XRAY_LAMBDA_LINEAGE}`;
const MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_0_COMMON};Sampled=1;${MOCK_XRAY_LAMBDA_LINEAGE}`;
const MOCK_XRAY_TRACE_CONTEXT_1_UNSAMPLED = `${MOCK_XRAY_TRACE_CONTEXT_1_COMMON};Sampled=0;${MOCK_XRAY_LAMBDA_LINEAGE}`;
// const MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH = (
// `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR.slice(0, TRACE_ID_FIRST_PART_LENGTH)}` +
// `-${MOCK_XRAY_TRACE_ID_STR.slice(TRACE_ID_FIRST_PART_LENGTH)};${MOCK_XRAY_LAMBDA_LINEAGE}`
// `Root=${TRACE_ID_VERSION}-${MOCK_XRAY_TRACE_ID_STR_0.slice(0, TRACE_ID_FIRST_PART_LENGTH)}` +
// `-${MOCK_XRAY_TRACE_ID_STR_0.slice(TRACE_ID_FIRST_PART_LENGTH)};${MOCK_XRAY_LAMBDA_LINEAGE}`
// );

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

let awsPropagatorStub: SinonStub;
let traceGetSpanStub: SinonStub;
let awsPropagatorSpy: sinon.SinonSpy;
let traceGetSpanSpy: sinon.SinonSpy;
// let propagationStub: SinonStub;

beforeEach(() => {
Expand All @@ -1005,41 +1010,165 @@ describe('customExtractor', () => {
sinon.restore();
});

it('should extract context from lambda trace header when present', () => {
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_SAMPLED;
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;
it('should extract context from handler context xRayTraceId property when present', () => {
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
const mockHandlerContext = {
xRayTraceId: mockLambdaTraceHeader,
} as unknown as Context;

// Mock of the Span Context for validation
const mockParentSpanContext: api.SpanContext = {
isRemote: true,
traceFlags: 1,
traceId: MOCK_XRAY_TRACE_ID_0,
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
};

const mockParentContext = {} as OtelContext;
sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
traceGetSpanSpy = sinon.spy(trace, 'getSpan');

// Partial mock of the Span object
const mockSpan: Partial<Span> = {
spanContext: sinon.stub().returns({
traceId: MOCK_XRAY_TRACE_ID,
spanId: MOCK_XRAY_PARENT_SPAN_ID,
}),
};
// Call the customExtractor function
const event = { headers: {} };
const result = customExtractor(event, mockHandlerContext);

// Assertions
expect(awsPropagatorSpy.calledOnce).toBe(true);
expect(
awsPropagatorSpy.calledWith(
sinon.match.any,
{ [AWSXRAY_TRACE_ID_HEADER]: mockLambdaTraceHeader },
sinon.match.any
)
).toBe(true);
expect(traceGetSpanSpy.calledOnce).toBe(true);
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
});

it('should prioritize extract context from handler context xRayTraceId property instead of environment variable', () => {
const mockLambdaContextTraceHeader = MOCK_XRAY_TRACE_CONTEXT_1_UNSAMPLED;
const mockLambdEnvVarTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;

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

// Stub trace.getSpan to return the mock span
traceGetSpanStub = sinon.stub(trace, 'getSpan').returns(mockSpan as Span);
const mockHandlerContext = {
xRayTraceId: mockLambdaContextTraceHeader,
} as unknown as Context;

// Mock of the Span Context for validation
const mockParentSpanContext: api.SpanContext = {
isRemote: true,
traceFlags: 0,
traceId: MOCK_XRAY_TRACE_ID_1,
spanId: MOCK_XRAY_PARENT_SPAN_ID_1,
};

sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
traceGetSpanSpy = sinon.spy(trace, 'getSpan');

// Call the customExtractor function
const event = { headers: {} };
const result = customExtractor(event, mockHandlerContext);

// Assertions
expect(awsPropagatorSpy.calledOnce).toBe(true);
expect(
awsPropagatorSpy.calledWith(
sinon.match.any,
{ [AWSXRAY_TRACE_ID_HEADER]: mockLambdaContextTraceHeader },
sinon.match.any
)
).toBe(true);
expect(traceGetSpanSpy.calledOnce).toBe(true);
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
});

it('should fallback to environment variable when handler context xRayTraceId is not available', () => {
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;

// Mock of the Span Context for validation
const mockParentSpanContext: api.SpanContext = {
isRemote: true,
traceFlags: 1,
traceId: MOCK_XRAY_TRACE_ID_0,
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
};

sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
traceGetSpanSpy = sinon.spy(trace, 'getSpan');

// Call the customExtractor function with handler context without xRayTraceId
const event = { headers: {} };
const result = customExtractor(event, {} as Context);

// Assertions
expect(awsPropagatorStub.calledOnce).toBe(true);
expect(awsPropagatorSpy.calledOnce).toBe(true);
expect(
awsPropagatorStub.calledWith(
awsPropagatorSpy.calledWith(
sinon.match.any,
{ [AWSXRAY_TRACE_ID_HEADER]: mockLambdaTraceHeader },
sinon.match.any
)
).toBe(true);
expect(traceGetSpanStub.calledOnce).toBe(true);
expect(result).toEqual(mockParentContext); // Should return the parent context when valid
expect(traceGetSpanSpy.calledOnce).toBe(true);
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the parent context when valid
});

it('should handle invalid span context from handler context xRayTraceId and fallback to environment variable', () => {
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
const mockHandlerContext = {
xRayTraceId: 'invalid-trace-header',
} as unknown as Context;
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;

// Mock of the Span Context for validation
const mockParentSpanContext: api.SpanContext = {
isRemote: true,
traceFlags: 1,
traceId: MOCK_XRAY_TRACE_ID_0,
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
};

sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
traceGetSpanSpy = sinon.spy(trace, 'getSpan');

// Call the customExtractor function
const event = { headers: {} };
const result = customExtractor(event, mockHandlerContext);

// Assertions - should be called twice (handler context + environment)
expect(awsPropagatorSpy.calledTwice).toBe(true);
expect(traceGetSpanSpy.calledTwice).toBe(true);
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the valid parent context
});

it('should handle undefined handler context and fallback to environment variable', () => {
const mockLambdaTraceHeader = MOCK_XRAY_TRACE_CONTEXT_0_SAMPLED;
process.env[traceContextEnvironmentKey] = mockLambdaTraceHeader;

// Mock of the Span Context for validation
const mockParentSpanContext: api.SpanContext = {
isRemote: true,
traceFlags: 1,
traceId: MOCK_XRAY_TRACE_ID_0,
spanId: MOCK_XRAY_PARENT_SPAN_ID_0,
};

sinon.stub(otelContext, 'active').returns(ROOT_CONTEXT);
awsPropagatorSpy = sinon.spy(AWSXRayPropagator.prototype, 'extract');
traceGetSpanSpy = sinon.spy(trace, 'getSpan');

// Call the customExtractor function with undefined handler context
const event = { headers: {} };
const result = customExtractor(event, undefined as any);

// Should only be called once (for environment variable)
expect(awsPropagatorSpy.calledOnce).toBe(true);
expect(trace.getSpan(result)?.spanContext()).toEqual(mockParentSpanContext); // Should return the valid parent context
});

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