Add Patching Mechanism with AWS SDK telemetry improvements (#13)

*Issue #, if available:*

*Description of changes:*
1. `aws-attribute-keys.ts`:
- Update values for keys: `AWS_BUCKET_NAME, AWS_QUEUE_URL,
AWS_QUEUE_NAME, AWS_STREAM_NAME, AWS_TABLE_NAME` in order to match the
actual attribute collected from AWS SDK auto-instrumentation

2. `aws-metric-attribute-generator.ts`:
- [Similarly to
Python](https://github.com/aws-observability/aws-otel-python-instrumentation/blob/2c00bd07eaaf703880a24a2fcfe874cdb4196678/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py#L378-L380),
accommodates the fact that AWS_ATTRIBUTE_KEYS.AWS_TABLE_NAMES [has an
array of table
names](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/931318cac21ee3707f3735a64ac751566ee37182/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/dynamodb.ts#L99-L105)
from aws-sdk dynamodb client instrumentation

3. `aws-opentelemetry-configurator.ts`:
    - Takes and sets instrumentations in constructor.

4. `instrumentation-patch.ts`:
    - Applies patches for AWS SDK Instrumentation

5. `register.ts`:
- Applies patches from (4.) by default, unless `AWS_APPLY_PATCHES` env
var is not `'true'`

The following files are copied from upstream:
```
aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/ServiceExtension.ts
```
The following files are being contributed to upstream:
See:
open-telemetry/opentelemetry-js-contrib#2361
```
aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/kinesis.ts
aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/s3.ts
```

*Testing:*
- Add unit tests

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: jj22ee

aws-distro-opentelemetry-node-autoinstrumentation/package.json

@@ -30,13 +30,17 @@
     "prewatch": "npm run precompile",
     "prepublishOnly": "npm run compile",
     "tdd": "yarn test -- --watch-extensions ts --watch",
-    "test": "nyc ts-mocha --timeout 10000 -p tsconfig.json 'test/**/*.ts'",
+    "test": "nyc ts-mocha --timeout 10000 -p tsconfig.json --require '@opentelemetry/contrib-test-utils' 'test/**/*.ts'",
     "watch": "tsc -w"
   },
   "bugs": {
     "url": "https://github.com/aws-observability/aws-otel-js-instrumentation/issues"
   },
   "devDependencies": {
+    "@aws-sdk/client-kinesis": "3.85.0",
+    "@aws-sdk/client-s3": "3.85.0",
+    "@aws-sdk/client-sqs": "3.85.0",
+    "@opentelemetry/contrib-test-utils": "^0.40.0",
     "@types/mocha": "7.0.2",
     "@types/node": "18.6.5",
     "@types/sinon": "10.0.18",

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-attribute-keys.ts

@@ -1,6 +1,8 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import { SEMATTRS_AWS_DYNAMODB_TABLE_NAMES } from '@opentelemetry/semantic-conventions';
+
 // Utility class holding attribute keys with special meaning to AWS components
 export const AWS_ATTRIBUTE_KEYS: { [key: string]: string } = {
   AWS_SPAN_KIND: 'aws.span.kind',
@@ -19,14 +21,11 @@ export const AWS_ATTRIBUTE_KEYS: { [key: string]: string } = {
   // Used for JavaScript workaround - attribute for pre-calculated value of isLocalRoot
   AWS_IS_LOCAL_ROOT: 'aws.is.local.root',
 
-  // Divergence from Java/Python
-  // TODO: Audit this: These will most definitely be different in JavaScript.
-  //   For example:
-  //     - `messaging.url` for AWS_QUEUE_URL
-  //     - `aws.dynamodb.table_names` for AWS_TABLE_NAME
-  AWS_BUCKET_NAME: 'aws.bucket.name',
-  AWS_QUEUE_URL: 'aws.queue.url',
-  AWS_QUEUE_NAME: 'aws.queue.name',
-  AWS_STREAM_NAME: 'aws.stream.name',
-  AWS_TABLE_NAME: 'aws.table.name',
+  // AWS_#_NAME attributes are not supported in JavaScript as they are not part of the Semantic Conventions.
+  // TODO：Move to Semantic Conventions when these attributes are added.
+  AWS_S3_BUCKET: 'aws.s3.bucket',
+  AWS_SQS_QUEUE_URL: 'aws.sqs.queue.url',
+  AWS_SQS_QUEUE_NAME: 'aws.sqs.queue.name',
+  AWS_KINESIS_STREAM_NAME: 'aws.kinesis.stream.name',
+  AWS_DYNAMODB_TABLE_NAMES: SEMATTRS_AWS_DYNAMODB_TABLE_NAMES,
 };

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-metric-attribute-generator.ts

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { Attributes, AttributeValue, diag, SpanKind } from '@opentelemetry/api';
-import { Resource, defaultServiceName } from '@opentelemetry/resources';
+import { defaultServiceName, Resource } from '@opentelemetry/resources';
 import { ReadableSpan } from '@opentelemetry/sdk-trace-base';
 import {
   SEMATTRS_DB_CONNECTION_STRING,
@@ -340,30 +340,33 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
     let remoteResourceIdentifier: AttributeValue | undefined;
 
     if (AwsSpanProcessingUtil.isAwsSDKSpan(span)) {
-      if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_TABLE_NAME)) {
+      const awsTableNames: AttributeValue | undefined = span.attributes[AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_NAMES];
+      if (
+        AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_NAMES) &&
+        Array.isArray(awsTableNames) &&
+        awsTableNames.length === 1
+      ) {
         remoteResourceType = NORMALIZED_DYNAMO_DB_SERVICE_NAME + '::Table';
-        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_TABLE_NAME]
-        );
-      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_STREAM_NAME)) {
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(awsTableNames[0]);
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_NAME)) {
         remoteResourceType = NORMALIZED_KINESIS_SERVICE_NAME + '::Stream';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_STREAM_NAME]
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_NAME]
         );
-      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_BUCKET_NAME)) {
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_S3_BUCKET)) {
         remoteResourceType = NORMALIZED_S3_SERVICE_NAME + '::Bucket';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_BUCKET_NAME]
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_S3_BUCKET]
         );
-      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_QUEUE_NAME)) {
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_NAME)) {
         remoteResourceType = NORMALIZED_SQS_SERVICE_NAME + '::Queue';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_QUEUE_NAME]
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_NAME]
         );
-      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_QUEUE_URL)) {
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_URL)) {
         remoteResourceType = NORMALIZED_SQS_SERVICE_NAME + '::Queue';
         remoteResourceIdentifier = SqsUrlParser.getQueueName(
-          AwsMetricAttributeGenerator.escapeDelimiters(span.attributes[AWS_ATTRIBUTE_KEYS.AWS_QUEUE_URL])
+          AwsMetricAttributeGenerator.escapeDelimiters(span.attributes[AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_URL])
         );
       }
     } else if (AwsSpanProcessingUtil.isDBSpan(span)) {
@@ -479,8 +482,8 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
     return AwsMetricAttributeGenerator.escapeDelimiters(address) + (port !== '' ? '|' + port : '');
   }
 
-  private static escapeDelimiters(input: string | AttributeValue | undefined): string | undefined {
-    if (input === undefined) {
+  private static escapeDelimiters(input: string | AttributeValue | undefined | null): string | undefined {
+    if (typeof input !== 'string') {
       return undefined;
     }
 

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-opentelemetry-configurator.ts

@@ -4,10 +4,7 @@
 
 import { TextMapPropagator, diag } from '@opentelemetry/api';
 import { getPropagator } from '@opentelemetry/auto-configuration-propagators';
-import {
-  getNodeAutoInstrumentations,
-  getResourceDetectors as getResourceDetectorsFromEnv,
-} from '@opentelemetry/auto-instrumentations-node';
+import { getResourceDetectors as getResourceDetectorsFromEnv } from '@opentelemetry/auto-instrumentations-node';
 import { ENVIRONMENT, TracesSamplerValues, getEnv, getEnvWithoutDefaults } from '@opentelemetry/core';
 import { OTLPMetricExporter as OTLPGrpcOTLPMetricExporter } from '@opentelemetry/exporter-metrics-otlp-grpc';
 import {
@@ -85,7 +82,18 @@ export class AwsOpentelemetryConfigurator {
   private spanProcessors: SpanProcessor[];
   private propagator: TextMapPropagator;
 
-  constructor() {
+  /**
+   * The constructor will setup the AwsOpentelemetryConfigurator object to be able to provide a
+   * configuration for ADOT JavaScript Auto-Instrumentation.
+   *
+   * The `instrumentations` are the desired Node auto-instrumentations to be used when using ADOT JavaScript.
+   * The auto-Instrumentions are usually populated from OTel's `getNodeAutoInstrumentations()` method from the
+   * `@opentelemetry/auto-instrumentations-node` NPM package, and may have instrumentation patching applied.
+   *
+   * @constructor
+   * @param {Instrumentation[]} instrumentations - Auto-Instrumentations to be added to the ADOT Config
+   */
+  public constructor(instrumentations: Instrumentation[]) {
     /*
      * Set and Detect Resources via Resource Detectors
      *
@@ -135,7 +143,7 @@ export class AwsOpentelemetryConfigurator {
     autoResource = autoResource.merge(detectResourcesSync(internalConfig));
     this.resource = autoResource;
 
-    this.instrumentations = getNodeAutoInstrumentations();
+    this.instrumentations = instrumentations;
     this.propagator = getPropagator();
 
     // TODO: Consider removing AWSXRayIdGenerator as it is not needed

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/kinesis.ts

@@ -0,0 +1,41 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Attributes, SpanKind } from '@opentelemetry/api';
+import { AwsSdkInstrumentationConfig, NormalizedRequest } from '@opentelemetry/instrumentation-aws-sdk';
+import { AWS_ATTRIBUTE_KEYS } from '../../../aws-attribute-keys';
+import { RequestMetadata, ServiceExtension } from '../../../third-party/otel/aws/services/ServiceExtension';
+
+/*
+This file's contents are being contributed to upstream
+- https://github.com/open-telemetry/opentelemetry-js-contrib/pull/2361
+
+This class is a service extension to be used for the AWS JavaScript SDK instrumentation patch for Kinesis.
+The instrumentation patch adds this extension to the upstream's Map of known extension for Kinesis.
+Extensions allow for custom logic for adding service-specific information to spans, such as attributes.
+Specifically, we are adding logic to add the `aws.kinesis.stream.name` attribute, to be used to generate
+RemoteTarget and achieve parity with the Java/Python instrumentation.
+*/
+export class KinesisServiceExtension implements ServiceExtension {
+  requestPreSpanHook(request: NormalizedRequest, _config: AwsSdkInstrumentationConfig): RequestMetadata {
+    const streamName = request.commandInput?.StreamName;
+
+    const spanKind: SpanKind = SpanKind.CLIENT;
+    let spanName: string | undefined;
+
+    const spanAttributes: Attributes = {};
+
+    if (streamName) {
+      spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_NAME] = streamName;
+    }
+
+    const isIncoming = false;
+
+    return {
+      isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+}

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/s3.ts

@@ -0,0 +1,41 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Attributes, SpanKind } from '@opentelemetry/api';
+import { AwsSdkInstrumentationConfig, NormalizedRequest } from '@opentelemetry/instrumentation-aws-sdk';
+import { AWS_ATTRIBUTE_KEYS } from '../../../aws-attribute-keys';
+import { RequestMetadata, ServiceExtension } from '../../../third-party/otel/aws/services/ServiceExtension';
+
+/*
+This file's contents are being contributed to upstream
+- https://github.com/open-telemetry/opentelemetry-js-contrib/pull/2361
+
+This class is a service extension to be used for the AWS JavaScript SDK instrumentation patch for S3.
+The instrumentation patch adds this extension to the upstream's Map of known extension for S3.
+Extensions allow for custom logic for adding service-specific information to spans, such as attributes.
+Specifically, we are adding logic to add the `aws.s3.bucket` attribute, to be used to generate
+RemoteTarget and achieve parity with the Java/Python instrumentation.
+*/
+export class S3ServiceExtension implements ServiceExtension {
+  requestPreSpanHook(request: NormalizedRequest, _config: AwsSdkInstrumentationConfig): RequestMetadata {
+    const bucketName = request.commandInput?.Bucket;
+
+    const spanKind: SpanKind = SpanKind.CLIENT;
+    let spanName: string | undefined;
+
+    const spanAttributes: Attributes = {};
+
+    if (bucketName) {
+      spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_S3_BUCKET] = bucketName;
+    }
+
+    const isIncoming = false;
+
+    return {
+      isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+}

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts

@@ -0,0 +1,73 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Instrumentation } from '@opentelemetry/instrumentation';
+import { AwsSdkInstrumentationConfig, NormalizedRequest } from '@opentelemetry/instrumentation-aws-sdk';
+import { AWS_ATTRIBUTE_KEYS } from '../aws-attribute-keys';
+import { RequestMetadata } from '../third-party/otel/aws/services/ServiceExtension';
+import { KinesisServiceExtension } from './aws/services/kinesis';
+import { S3ServiceExtension } from './aws/services/s3';
+
+export function applyInstrumentationPatches(instrumentations: Instrumentation[]): void {
+  /*
+  Apply patches to upstream instrumentation libraries.
+
+  This method is invoked to apply changes to upstream instrumentation libraries, typically when changes to upstream
+  are required on a timeline that cannot wait for upstream release. Generally speaking, patches should be short-term
+  local solutions that are comparable to long-term upstream solutions.
+
+  Where possible, automated testing should be run to catch upstream changes resulting in broken patches
+  */
+  instrumentations.forEach(instrumentation => {
+    if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-sdk') {
+      // 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;
+      if (services) {
+        services.set('S3', new S3ServiceExtension());
+        services.set('Kinesis', new KinesisServiceExtension());
+        patchSqsServiceExtension(services.get('SQS'));
+      }
+    }
+  });
+}
+
+/*
+ * This patch extends the existing upstream extension for SQS. Extensions allow for custom logic for adding
+ * service-specific information to spans, such as attributes. Specifically, we are adding logic to add
+ * `aws.sqs.queue.url` and `aws.sqs.queue.name` attributes, to be used to generate RemoteTarget and achieve parity
+ * with the Java/Python instrumentation.
+ *
+ * Callout that today, the upstream logic adds `messaging.url` and `messaging.destination` but we feel that
+ * `aws.sqs` is more in line with existing AWS Semantic Convention attributes like `AWS_S3_BUCKET`, etc.
+ *
+ * @param sqsServiceExtension SQS Service Extension obtained the service extension list from the AWS SDK OTel Instrumentation
+ */
+function patchSqsServiceExtension(sqsServiceExtension: any): void {
+  // It is not expected that `sqsServiceExtension` is undefined
+  if (sqsServiceExtension) {
+    const requestPreSpanHook = sqsServiceExtension.requestPreSpanHook;
+    // Save original `requestPreSpanHook` under a similar name, to be invoked by the patched hook
+    sqsServiceExtension._requestPreSpanHook = requestPreSpanHook;
+    // The patched hook will populate the 'aws.sqs.queue.url' and 'aws.sqs.queue.name' attributes according to spec
+    // from the 'messaging.url' attribute
+    const patchedRequestPreSpanHook = (
+      request: NormalizedRequest,
+      _config: AwsSdkInstrumentationConfig
+    ): RequestMetadata => {
+      const requestMetadata: RequestMetadata = sqsServiceExtension._requestPreSpanHook(request, _config);
+      // It is not expected that `requestMetadata.spanAttributes` can possibly be undefined, but still be careful anyways
+      if (requestMetadata.spanAttributes) {
+        if (request.commandInput?.QueueUrl) {
+          requestMetadata.spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_URL] = request.commandInput.QueueUrl;
+        }
+        if (request.commandInput?.QueueName) {
+          requestMetadata.spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_NAME] = request.commandInput.QueueName;
+        }
+      }
+      return requestMetadata;
+    };
+    sqsServiceExtension.requestPreSpanHook = patchedRequestPreSpanHook;
+  }
+}

aws-distro-opentelemetry-node-autoinstrumentation/src/register.ts

@@ -3,8 +3,11 @@
 // Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
 
 import { DiagConsoleLogger, diag } from '@opentelemetry/api';
+import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
+import { Instrumentation } from '@opentelemetry/instrumentation';
 import * as opentelemetry from '@opentelemetry/sdk-node';
 import { AwsOpentelemetryConfigurator } from './aws-opentelemetry-configurator';
+import { applyInstrumentationPatches } from './patches/instrumentation-patch';
 
 diag.setLogger(new DiagConsoleLogger(), opentelemetry.core.getEnv().OTEL_LOG_LEVEL);
 
@@ -40,7 +43,12 @@ export function setAwsDefaultEnvironmentVariables(): void {
 }
 setAwsDefaultEnvironmentVariables();
 
-const configurator: AwsOpentelemetryConfigurator = new AwsOpentelemetryConfigurator();
+const instrumentations: Instrumentation[] = getNodeAutoInstrumentations();
+
+// Apply instrumentation patches
+applyInstrumentationPatches(instrumentations);
+
+const configurator: AwsOpentelemetryConfigurator = new AwsOpentelemetryConfigurator(instrumentations);
 const configuration: Partial<opentelemetry.NodeSDKConfiguration> = configurator.configure();
 
 const sdk: opentelemetry.NodeSDK = new opentelemetry.NodeSDK(configuration);

aws-distro-opentelemetry-node-autoinstrumentation/src/sqs-url-parser.ts

@@ -1,6 +1,8 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import { AttributeValue } from '@opentelemetry/api';
+
 const HTTP_SCHEMA: string = 'http://';
 const HTTPS_SCHEMA: string = 'https://';
 
@@ -16,8 +18,8 @@ export class SqsUrlParser {
    * /'s (excluding schema), the second part should be a 12-digit account id, and the third part
    * should be a valid queue name, per SQS naming conventions.
    */
-  public static getQueueName(url: string | undefined): string | undefined {
-    if (url === undefined) {
+  public static getQueueName(url: AttributeValue | undefined): string | undefined {
+    if (typeof url !== 'string') {
       return undefined;
     }
     url = url.replace(HTTP_SCHEMA, '').replace(HTTPS_SCHEMA, '');