[Fix] Update AWS SDK Instrumentation to inject XRay trace context into HTTP Headers

aws-distro-opentelemetry-node-autoinstrumentation/.eslintignore

@@ -1,4 +1,5 @@
 build
 node_modules
 .eslintrc.js
-version.ts
+version.ts
+src/third-party

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"
   },

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<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]);
+    };
+  }
+}

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<APIGatewayProxyEventHeaders> = {
   },
 };
 
-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<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());

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<opentelemetry.NodeSDKConfiguration> = configurator.configure();

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts → 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;

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');
+  });
+});

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', () => {