Skip to content

Commit

Permalink
[Fix] Update AWS SDK Instrumentation to inject XRay trace context int…
Browse files Browse the repository at this point in the history
…o 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.
  • Loading branch information
jj22ee authored Dec 16, 2024
1 parent 694e4f4 commit a9375ff
Show file tree
Hide file tree
Showing 9 changed files with 176 additions and 7 deletions.
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
build
node_modules
.eslintrc.js
version.ts
version.ts
src/third-party
11 changes: 10 additions & 1 deletion aws-distro-opentelemetry-node-autoinstrumentation/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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"
},
Expand Down
Original file line number Diff line number Diff line change
@@ -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<any, any, any, any, any> & {
[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<any>) {
return function send(this: any, command: V3PluginCommand, ...args: unknown[]): Promise<any> {
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]);
};
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -38,7 +40,10 @@ export const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
},
};

export function applyInstrumentationPatches(instrumentations: Instrumentation[]): void {
export function applyInstrumentationPatches(
instrumentations: Instrumentation[],
instrumentationConfigs?: InstrumentationConfigMap
): void {
/*
Apply patches to upstream instrumentation libraries.
Expand All @@ -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<string, ServiceExtension> | undefined = (instrumentation as any).servicesExtensions?.services;
const services: Map<string, ServiceExtension> | undefined = (instrumentations[index] as any).servicesExtensions
?.services;
if (services) {
services.set('S3', new S3ServiceExtension());
services.set('Kinesis', new KinesisServiceExtension());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<opentelemetry.NodeSDKConfiguration> = configurator.configure();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
@@ -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');
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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);
Expand Down Expand Up @@ -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', () => {
Expand Down Expand Up @@ -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', () => {
Expand Down

0 comments on commit a9375ff

Please sign in to comment.