From a9375ffd42e01bcdf96e2cf19b0687c2706063d2 Mon Sep 17 00:00:00 2001 From: Jonathan Lee <107072447+jj22ee@users.noreply.github.com> Date: Mon, 16 Dec 2024 14:08:50 -0800 Subject: [PATCH] [Fix] Update AWS SDK Instrumentation to inject XRay trace context into HTTP Headers (#131) *Issue #, if available:* Fixes the absence of broken X-Ray context propagation when the underlying HTTP instrumentation is suppressed or disabled. Similarly to Java and Python, AWS SDK Js instrumentation itself should be able to inject the X-Ray Context into the HTTP Headers: - [Java example](https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/81c7713bb2f638c85006c3e152ad13d6e02f3259/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/TracingExecutionInterceptor.java#L308) - [Python example ](https://github.com/open-telemetry/opentelemetry-python-contrib/blob/main/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py#L163-L165) - See Specification that clarifies that AWS SDK instrumentations should use X-Ray propagator specifically. - https://github.com/open-telemetry/opentelemetry-specification/blob/v1.40.0/supplementary-guidelines/compatibility/aws.md?plain=1#L9-L12 Note - If the underlying HTTP instrumentation is enabled, then the underlying HTTP Child Span of the AWS SDK Span will overwrite the Trace Context to propagate through headers. *Description of changes:* - Move patched/extended instrumentations to an `patches/extended-instrumentations/` directory - Created `AwsSdkInstrumentationExtended` class that extends upstream AwsInstrumentation to override its patching mechanism of the `send` method. The overridden method will additionally update the AWS SDK middleware stack to inject the `X-Amzn-Trace-Id` HTTP header. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --- .../.eslintignore | 3 +- .../package.json | 11 +++- .../aws-lambda.ts | 0 .../aws-sdk-instrumentation-extended.ts | 42 ++++++++++++++ .../src/patches/instrumentation-patch.ts | 17 +++++- .../src/register.ts | 2 +- .../aws-lambda.test.ts | 2 +- .../aws-sdk-instrumentation-extended.test.ts | 48 +++++++++++++++ .../patches/instrumentation-patch.test.ts | 58 +++++++++++++++++++ 9 files changed, 176 insertions(+), 7 deletions(-) rename aws-distro-opentelemetry-node-autoinstrumentation/src/patches/{aws/services => extended-instrumentations}/aws-lambda.ts (100%) create mode 100644 aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-sdk-instrumentation-extended.ts rename aws-distro-opentelemetry-node-autoinstrumentation/test/patches/{aws/services => extended-instrumentations}/aws-lambda.test.ts (98%) create mode 100644 aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-sdk-instrumentation-extended.test.ts diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/.eslintignore b/aws-distro-opentelemetry-node-autoinstrumentation/.eslintignore index 34c98dc..8cbd431 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/.eslintignore +++ b/aws-distro-opentelemetry-node-autoinstrumentation/.eslintignore @@ -1,4 +1,5 @@ build node_modules .eslintrc.js -version.ts \ No newline at end of file +version.ts +src/third-party \ No newline at end of file diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/package.json b/aws-distro-opentelemetry-node-autoinstrumentation/package.json index 38b1d2a..4be5456 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/package.json +++ b/aws-distro-opentelemetry-node-autoinstrumentation/package.json @@ -31,9 +31,18 @@ "prepublishOnly": "npm run compile", "tdd": "yarn test -- --watch-extensions ts --watch", "test": "nyc ts-mocha --timeout 10000 -p tsconfig.json --require '@opentelemetry/contrib-test-utils' 'test/**/*.ts'", - "test:coverage": "nyc --all --check-coverage --functions 95 --lines 95 ts-mocha --timeout 10000 -p tsconfig.json --require '@opentelemetry/contrib-test-utils' 'test/**/*.ts'", + "test:coverage": "nyc --check-coverage --functions 95 --lines 95 ts-mocha --timeout 10000 -p tsconfig.json --require '@opentelemetry/contrib-test-utils' 'test/**/*.ts'", "watch": "tsc -w" }, + "nyc": { + "all": true, + "include": [ + "src/**/*.ts" + ], + "exclude": [ + "src/third-party/**/*.ts" + ] + }, "bugs": { "url": "https://github.com/aws-observability/aws-otel-js-instrumentation/issues" }, diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/aws-lambda.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-lambda.ts similarity index 100% rename from aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/aws-lambda.ts rename to aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-lambda.ts diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-sdk-instrumentation-extended.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-sdk-instrumentation-extended.ts new file mode 100644 index 0000000..b3044d6 --- /dev/null +++ b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-sdk-instrumentation-extended.ts @@ -0,0 +1,42 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { AwsInstrumentation } from '@opentelemetry/instrumentation-aws-sdk'; +import { context as otelContext, defaultTextMapSetter } from '@opentelemetry/api'; +import { AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray'; +import type { Command as AwsV3Command } from '@aws-sdk/types'; + +const awsXrayPropagator = new AWSXRayPropagator(); +const V3_CLIENT_CONFIG_KEY = Symbol('opentelemetry.instrumentation.aws-sdk.client.config'); +type V3PluginCommand = AwsV3Command & { + [V3_CLIENT_CONFIG_KEY]?: any; +}; + +// This class extends the upstream AwsInstrumentation to override its patching mechanism of the `send` method. +// The overriden method will additionally update the AWS SDK middleware stack to inject the `X-Amzn-Trace-Id` HTTP header. +// +// eslint-disable-next-line @typescript-eslint/ban-ts-comment +// @ts-ignore +export class AwsSdkInstrumentationExtended extends AwsInstrumentation { + // Override the upstream private _getV3SmithyClientSendPatch method to add middleware to inject X-Ray Trace Context into HTTP Headers + // https://github.com/open-telemetry/opentelemetry-js-contrib/blob/instrumentation-aws-sdk-v0.48.0/plugins/node/opentelemetry-instrumentation-aws-sdk/src/aws-sdk.ts#L373-L384 + override _getV3SmithyClientSendPatch(original: (...args: unknown[]) => Promise) { + return function send(this: any, command: V3PluginCommand, ...args: unknown[]): Promise { + this.middlewareStack?.add( + (next: any, context: any) => async (middlewareArgs: any) => { + awsXrayPropagator.inject(otelContext.active(), middlewareArgs.request.headers, defaultTextMapSetter); + const result = await next(middlewareArgs); + return result; + }, + { + step: 'build', + name: '_adotInjectXrayContextMiddleware', + override: true, + } + ); + + command[V3_CLIENT_CONFIG_KEY] = this.config; + return original.apply(this, [command, ...args]); + }; + } +} 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 7cc98bd..1302bb7 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts @@ -25,7 +25,9 @@ import { } from './aws/services/bedrock'; import { KinesisServiceExtension } from './aws/services/kinesis'; import { S3ServiceExtension } from './aws/services/s3'; -import { AwsLambdaInstrumentationPatch } from './aws/services/aws-lambda'; +import { AwsLambdaInstrumentationPatch } from './extended-instrumentations/aws-lambda'; +import { InstrumentationConfigMap } from '@opentelemetry/auto-instrumentations-node'; +import { AwsSdkInstrumentationExtended } from './extended-instrumentations/aws-sdk-instrumentation-extended'; export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID'; const awsPropagator = new AWSXRayPropagator(); @@ -38,7 +40,10 @@ export const headerGetter: TextMapGetter = { }, }; -export function applyInstrumentationPatches(instrumentations: Instrumentation[]): void { +export function applyInstrumentationPatches( + instrumentations: Instrumentation[], + instrumentationConfigs?: InstrumentationConfigMap +): void { /* Apply patches to upstream instrumentation libraries. @@ -50,10 +55,16 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[]) */ instrumentations.forEach((instrumentation, index) => { if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-sdk') { + diag.debug('Overriding aws sdk instrumentation'); + instrumentations[index] = new AwsSdkInstrumentationExtended( + instrumentationConfigs ? instrumentationConfigs['@opentelemetry/instrumentation-aws-sdk'] : undefined + ); + // Access private property servicesExtensions of AwsInstrumentation // eslint-disable-next-line @typescript-eslint/ban-ts-comment // @ts-ignore - const services: Map | undefined = (instrumentation as any).servicesExtensions?.services; + const services: Map | undefined = (instrumentations[index] as any).servicesExtensions + ?.services; if (services) { services.set('S3', new S3ServiceExtension()); services.set('Kinesis', new KinesisServiceExtension()); diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/register.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/register.ts index 823f164..751c0ff 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/src/register.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/src/register.ts @@ -62,7 +62,7 @@ const instrumentationConfigs: InstrumentationConfigMap = { const instrumentations: Instrumentation[] = getNodeAutoInstrumentations(instrumentationConfigs); // Apply instrumentation patches -applyInstrumentationPatches(instrumentations); +applyInstrumentationPatches(instrumentations, instrumentationConfigs); const configurator: AwsOpentelemetryConfigurator = new AwsOpentelemetryConfigurator(instrumentations, useXraySampler); const configuration: Partial = configurator.configure(); diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-lambda.test.ts similarity index 98% rename from aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts rename to aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-lambda.test.ts index d68e1e6..0580769 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-lambda.test.ts @@ -6,7 +6,7 @@ import * as path from 'path'; import * as fs from 'fs'; import { diag } from '@opentelemetry/api'; import { InstrumentationNodeModuleDefinition } from '@opentelemetry/instrumentation'; -import { AwsLambdaInstrumentationPatch } from '../../../../src/patches/aws/services/aws-lambda'; +import { AwsLambdaInstrumentationPatch } from '../../../src/patches/extended-instrumentations/aws-lambda'; describe('AwsLambdaInstrumentationPatch', () => { let instrumentation: AwsLambdaInstrumentationPatch; diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-sdk-instrumentation-extended.test.ts b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-sdk-instrumentation-extended.test.ts new file mode 100644 index 0000000..9751d21 --- /dev/null +++ b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-sdk-instrumentation-extended.test.ts @@ -0,0 +1,48 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +import * as sinon from 'sinon'; +import { AwsSdkInstrumentationExtended } from '../../../src/patches/extended-instrumentations/aws-sdk-instrumentation-extended'; +import expect from 'expect'; +import { AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray'; +import { Context, TextMapSetter } from '@opentelemetry/api'; + +describe('AwsSdkInstrumentationExtended', () => { + let instrumentation: AwsSdkInstrumentationExtended; + + beforeEach(() => { + instrumentation = new AwsSdkInstrumentationExtended({}); + }); + + afterEach(() => { + sinon.restore(); + }); + + it('overridden _getV3SmithyClientSendPatch updates MiddlewareStack', async () => { + const mockedMiddlewareStackInternal: any = []; + const mockedMiddlewareStack = { + add: (arg1: any, arg2: any) => mockedMiddlewareStackInternal.push([arg1, arg2]), + }; + const send = instrumentation + ._getV3SmithyClientSendPatch((...args: unknown[]) => Promise.resolve()) + .bind({ middlewareStack: mockedMiddlewareStack }); + sinon + .stub(AWSXRayPropagator.prototype, 'inject') + .callsFake((context: Context, carrier: unknown, setter: TextMapSetter) => { + (carrier as any)['isCarrierModified'] = 'carrierIsModified'; + }); + + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + await send({}, null); + + const middlewareArgs: any = { + request: { + headers: {}, + }, + }; + await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgs); + + expect(middlewareArgs.request.headers['isCarrierModified']).toEqual('carrierIsModified'); + expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware'); + }); +}); 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 fe946d6..9c11b06 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 @@ -23,6 +23,11 @@ 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 { S3 } from '@aws-sdk/client-s3'; +import nock = require('nock'); +import { ReadableSpan } from '@opentelemetry/sdk-trace-base'; +import { getTestSpans, registerInstrumentationTesting } from '@opentelemetry/contrib-test-utils'; +import { AwsSdkInstrumentationExtended } from '../../src/patches/extended-instrumentations/aws-sdk-instrumentation-extended'; const _STREAM_NAME: string = 'streamName'; const _BUCKET_NAME: string = 'bucketName'; @@ -45,6 +50,10 @@ const UNPATCHED_INSTRUMENTATIONS: Instrumentation[] = getNodeAutoInstrumentation const PATCHED_INSTRUMENTATIONS: Instrumentation[] = getNodeAutoInstrumentations(); applyInstrumentationPatches(PATCHED_INSTRUMENTATIONS); +const extendedAwsSdkInstrumentation: AwsInstrumentation = new AwsInstrumentation(); +applyInstrumentationPatches([extendedAwsSdkInstrumentation]); +registerInstrumentationTesting(extendedAwsSdkInstrumentation); + describe('InstrumentationPatchTest', () => { it('SanityTestUnpatchedAwsSdkInstrumentation', () => { const awsSdkInstrumentation: AwsInstrumentation = extractAwsSdkInstrumentation(UNPATCHED_INSTRUMENTATIONS); @@ -89,6 +98,9 @@ describe('InstrumentationPatchTest', () => { expect(services.get('BedrockRuntime')).toBeTruthy(); // Sanity check expect(services.has('InvalidService')).toBeFalsy(); + + // Check that the original AWS SDK Instrumentation is replaced with the extended version + expect(awsSdkInstrumentation).toBeInstanceOf(AwsSdkInstrumentationExtended); }); it('S3 without patching', () => { @@ -388,6 +400,52 @@ describe('InstrumentationPatchTest', () => { expect(filteredInstrumentations.length).toEqual(1); return filteredInstrumentations[0] as AwsLambdaInstrumentation; } + + describe('AwsSdkInstrumentationPatchTest', () => { + let s3: S3; + const region = 'us-east-1'; + + it('injects trace context header into request via propagator', async () => { + s3 = new S3({ + region: region, + credentials: { + accessKeyId: 'abcde', + secretAccessKey: 'abcde', + }, + }); + + const dummyBucketName: string = 'dummy-bucket-name'; + let reqHeaders: any = {}; + + nock(`https://${dummyBucketName}.s3.${region}.amazonaws.com`) + .get('/') + .reply(200, function (uri: any, requestBody: any) { + reqHeaders = this.req.headers; + return 'null'; + }); + + await s3 + .listObjects({ + Bucket: dummyBucketName, + }) + .catch((err: any) => {}); + + const testSpans: ReadableSpan[] = getTestSpans(); + const listObjectsSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => { + return s.name === 'S3.ListObjects'; + }); + + expect(listObjectsSpans.length).toBe(1); + + const traceId = listObjectsSpans[0].spanContext().traceId; + const spanId = listObjectsSpans[0].spanContext().spanId; + expect(reqHeaders['x-amzn-trace-id'] as string).toEqual( + `Root=1-${traceId.substring(0, 8)}-${listObjectsSpans[0] + .spanContext() + .traceId.substring(8, 32)};Parent=${spanId};Sampled=1` + ); + }); + }); }); describe('customExtractor', () => {