remove extended-instrumentations, use patching instead

Author: jj22ee

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-lambda.ts

@@ -12,9 +12,12 @@ import {
   trace,
   Span,
   Tracer,
+  SpanStatusCode,
+  defaultTextMapSetter,
 } from '@opentelemetry/api';
 import { Instrumentation } from '@opentelemetry/instrumentation';
 import {
+  AwsInstrumentation,
   AwsSdkInstrumentationConfig,
   NormalizedRequest,
   NormalizedResponse,
@@ -31,9 +34,8 @@ import {
 } from './aws/services/bedrock';
 import { SecretsManagerServiceExtension } from './aws/services/secretsmanager';
 import { StepFunctionsServiceExtension } from './aws/services/step-functions';
-import { InstrumentationConfigMap } from '@opentelemetry/auto-instrumentations-node';
-import { AwsSdkInstrumentationExtended } from './extended-instrumentations/aws-sdk-instrumentation-extended';
-import { AwsLambdaInstrumentationPatch } from './extended-instrumentations/aws-lambda';
+import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lambda';
+import type { Command as AwsV3Command } from '@aws-sdk/types';
 
 export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
 const awsPropagator = new AWSXRayPropagator();
@@ -46,11 +48,7 @@ export const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
   },
 };
 
-export function applyInstrumentationPatches(
-  instrumentations: Instrumentation[],
-  instrumentationConfigs?: InstrumentationConfigMap,
-  usePatchedAwsLambdaInstrumentation: boolean = false
-): void {
+export function applyInstrumentationPatches(instrumentations: Instrumentation[]): void {
   /*
   Apply patches to upstream instrumentation libraries.
 
@@ -62,10 +60,8 @@ export function applyInstrumentationPatches(
   */
   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
-      );
+      diag.debug('Patching aws sdk instrumentation');
+      patchAwsSdkInstrumentation(instrumentation);
 
       // Access private property servicesExtensions of AwsInstrumentation
       // eslint-disable-next-line @typescript-eslint/ban-ts-comment
@@ -83,18 +79,11 @@ export function applyInstrumentationPatches(
         patchSnsServiceExtension(services.get('SNS'));
         patchLambdaServiceExtension(services.get('Lambda'));
       }
+    } else if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-lambda') {
+      diag.debug('Patching aws lambda instrumentation');
+      patchAwsLambdaInstrumentation(instrumentation);
     }
   });
-
-  if (usePatchedAwsLambdaInstrumentation) {
-    diag.debug('Overriding aws lambda instrumentation');
-    const lambdaInstrumentation = new AwsLambdaInstrumentationPatch({
-      // When the `AwsLambdaInstrumentationPatch` is removed in favor of the patched one,
-      // setting `eventContextExtractor` should be in the `instrumentationConfigs` map.
-      eventContextExtractor: customExtractor,
-    });
-    instrumentations.push(lambdaInstrumentation);
-  }
 }
 
 /*
@@ -259,3 +248,78 @@ function patchLambdaServiceExtension(lambdaServiceExtension: any): void {
     }
   }
 }
+
+// Override the upstream private _endSpan method to remove the unnecessary metric force-flush error message
+// https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts#L358-L398
+function patchAwsLambdaInstrumentation(instrumentation: Instrumentation): void {
+  if (instrumentation) {
+    (instrumentation as AwsLambdaInstrumentation)['_endSpan'] = function (
+      span: Span,
+      err: string | Error | null | undefined,
+      callback: () => void
+    ) {
+      if (err) {
+        span.recordException(err);
+      }
+
+      let errMessage;
+      if (typeof err === 'string') {
+        errMessage = err;
+      } else if (err) {
+        errMessage = err.message;
+      }
+      if (errMessage) {
+        span.setStatus({
+          code: SpanStatusCode.ERROR,
+          message: errMessage,
+        });
+      }
+
+      span.end();
+
+      const flushers = [];
+      if ((this as any)._traceForceFlusher) {
+        flushers.push((this as any)._traceForceFlusher());
+      } else {
+        diag.error(
+          'Spans may not be exported for the lambda function because we are not force flushing before callback.'
+        );
+      }
+
+      Promise.all(flushers).then(callback, callback);
+    };
+  }
+}
+
+// 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
+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;
+};
+function patchAwsSdkInstrumentation(instrumentation: Instrumentation): void {
+  if (instrumentation) {
+    (instrumentation as AwsInstrumentation)['_getV3SmithyClientSendPatch'] = function (
+      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/extended-instrumentations/aws-sdk-instrumentation-extended.ts

@@ -17,8 +17,8 @@ import { diag, DiagConsoleLogger, trace } from '@opentelemetry/api';
 import { getNodeAutoInstrumentations, InstrumentationConfigMap } from '@opentelemetry/auto-instrumentations-node';
 import { Instrumentation } from '@opentelemetry/instrumentation';
 import * as opentelemetry from '@opentelemetry/sdk-node';
-import { AwsOpentelemetryConfigurator, isLambdaEnvironment } from './aws-opentelemetry-configurator';
-import { applyInstrumentationPatches } from './patches/instrumentation-patch';
+import { AwsOpentelemetryConfigurator } from './aws-opentelemetry-configurator';
+import { applyInstrumentationPatches, customExtractor } from './patches/instrumentation-patch';
 
 diag.setLogger(new DiagConsoleLogger(), opentelemetry.core.getEnv().OTEL_LOG_LEVEL);
 
@@ -36,8 +36,6 @@ This file may also be used to apply patches to upstream instrumentation - usuall
 long-term changes to upstream.
 */
 
-let usePatchedAwsLambdaInstrumentation = false;
-
 export function setAwsDefaultEnvironmentVariables(): void {
   if (!process.env.OTEL_EXPORTER_OTLP_PROTOCOL) {
     process.env.OTEL_EXPORTER_OTLP_PROTOCOL = 'http/protobuf';
@@ -51,18 +49,13 @@ export function setAwsDefaultEnvironmentVariables(): void {
   if (!process.env.OTEL_NODE_DISABLED_INSTRUMENTATIONS) {
     process.env.OTEL_NODE_DISABLED_INSTRUMENTATIONS = 'fs,dns';
   }
-
-  // Temporary solution to never initialize the unpatched AWS Lambda Instrumentation,
-  // and only initialize the patched one in `src/patches/extended-instrumentations/aws-lambda.ts`
-  // This workaround is for until the patched Lambda Instrumentation is removed
-  if (!process.env.OTEL_NODE_DISABLED_INSTRUMENTATIONS.includes('aws-lambda') && isLambdaEnvironment()) {
-    process.env.OTEL_NODE_DISABLED_INSTRUMENTATIONS += ',aws-lambda';
-    usePatchedAwsLambdaInstrumentation = true;
-  }
 }
 setAwsDefaultEnvironmentVariables();
 
-const instrumentationConfigs: InstrumentationConfigMap = {
+export const instrumentationConfigs: InstrumentationConfigMap = {
+  '@opentelemetry/instrumentation-aws-lambda': {
+    eventContextExtractor: customExtractor,
+  },
   '@opentelemetry/instrumentation-aws-sdk': {
     suppressInternalInstrumentation: true,
   },
@@ -73,7 +66,7 @@ const instrumentationConfigs: InstrumentationConfigMap = {
 const instrumentations: Instrumentation[] = getNodeAutoInstrumentations(instrumentationConfigs);
 
 // Apply instrumentation patches
-applyInstrumentationPatches(instrumentations, instrumentationConfigs, usePatchedAwsLambdaInstrumentation);
+applyInstrumentationPatches(instrumentations);
 
 const configurator: AwsOpentelemetryConfigurator = new AwsOpentelemetryConfigurator(instrumentations, useXraySampler);
 const configuration: Partial<opentelemetry.NodeSDKConfiguration> = configurator.configure();