diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts index 93338a5b..a808fc74 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts @@ -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'; 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(), diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts index 2774e09d..4820e672 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts @@ -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'; @@ -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'; @@ -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) @@ -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(() => { @@ -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 = { - 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', () => {