Patch aws-lambda instrumentation to support ESM

commit d25c3c3
Author: Min Xia <[email protected]>
Date:   Tue Oct 8 15:45:36 2024 -0700

    fix the lint error

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

@@ -0,0 +1,125 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+// Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
+
+import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lambda';
+import * as path from 'path';
+import * as fs from 'fs';
+import { InstrumentationNodeModuleDefinition, InstrumentationNodeModuleFile, isWrapped } from '@opentelemetry/instrumentation';
+import { diag } from '@opentelemetry/api';
+
+export class AwsLambdaInstrumentationPatch extends AwsLambdaInstrumentation {
+
+  override init() {
+    // Custom logic before calling the original implementation
+    diag.debug('Initializing AwsLambdaInstrumentationPatch');
+    const taskRoot = process.env.LAMBDA_TASK_ROOT;
+    const handlerDef = this._config.lambdaHandler ?? process.env._HANDLER;
+
+    // _HANDLER and LAMBDA_TASK_ROOT are always defined in Lambda but guard bail out if in the future this changes.
+    if (!taskRoot || !handlerDef) {
+      this._diag.debug(
+        'Skipping lambda instrumentation: no _HANDLER/lambdaHandler or LAMBDA_TASK_ROOT.',
+        { taskRoot, handlerDef }
+      );
+      return [];
+    }
+
+    const handler = path.basename(handlerDef);
+    const moduleRoot = handlerDef.substr(0, handlerDef.length - handler.length);
+
+    const [module, functionName] = handler.split('.', 2);
+
+    // Lambda loads user function using an absolute path.
+    let filename = path.resolve(taskRoot, moduleRoot, module);
+    if (!filename.endsWith('.js')) {
+      // its impossible to know in advance if the user has a cjs or js or mjs file.
+      // check that the .js file exists otherwise fallback to next known possibility
+      try {
+        fs.statSync(`${filename}.js`);
+        filename += '.js';
+      } catch (e) {
+        // fallback to .cjs
+        try {
+          fs.statSync(`${filename}.cjs`);
+          filename += '.cjs';
+        } catch (e) {
+          // fall back to .mjs
+          filename += '.mjs';
+        }
+      }
+    }
+
+    diag.debug('Instrumenting lambda handler', {
+      taskRoot,
+      handlerDef,
+      handler,
+      moduleRoot,
+      module,
+      filename,
+      functionName,
+    });
+
+    if (filename.endsWith('.mjs') || process.env.HANDLER_IS_ESM) {
+      return [
+        new InstrumentationNodeModuleDefinition(
+          // NB: The patching infrastructure seems to match names backwards, this must be the filename, while
+          // InstrumentationNodeModuleFile must be the module name.
+          filename,
+          ['*'],
+          (moduleExports: any) => {
+            diag.debug('Applying patch for lambda esm handler');
+            if (isWrapped(moduleExports[functionName])) {
+              this._unwrap(moduleExports, functionName);
+            }
+            this._wrap(
+              moduleExports,
+              functionName,
+              (this as any)._getHandler()
+            );
+            return moduleExports;
+          },
+          (moduleExports?: any) => {
+            if (moduleExports == null) return;
+            diag.debug('Removing patch for lambda esm handler');
+            this._unwrap(moduleExports, functionName);
+          }
+        )
+      ];
+    } else {
+      return [
+        new InstrumentationNodeModuleDefinition(
+          // NB: The patching infrastructure seems to match names backwards, this must be the filename, while
+          // InstrumentationNodeModuleFile must be the module name.
+          filename,
+          ['*'],
+          undefined,
+          undefined,
+          [
+            new InstrumentationNodeModuleFile(
+              module,
+              ['*'],
+              (moduleExports: any) => {
+                diag.debug('Applying patch for lambda handler');
+                if (isWrapped(moduleExports[functionName])) {
+                  this._unwrap(moduleExports, functionName);
+                }
+                this._wrap(
+                  moduleExports,
+                  functionName,
+                  (this as any)._getHandler()
+                );
+                return moduleExports;
+              },
+              (moduleExports?: any) => {
+                if (moduleExports == null) return;
+                diag.debug('Removing patch for lambda handler');
+                this._unwrap(moduleExports, functionName);
+              }
+            ),
+          ]
+        ),
+      ];
+    }
+  }
+}

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

@@ -12,7 +12,6 @@ import {
   trace,
 } from '@opentelemetry/api';
 import { Instrumentation } from '@opentelemetry/instrumentation';
-import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lambda';
 import { AwsSdkInstrumentationConfig, NormalizedRequest } from '@opentelemetry/instrumentation-aws-sdk';
 import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
 import { APIGatewayProxyEventHeaders, Context } from 'aws-lambda';
@@ -26,6 +25,7 @@ import {
 } from './aws/services/bedrock';
 import { KinesisServiceExtension } from './aws/services/kinesis';
 import { S3ServiceExtension } from './aws/services/s3';
+import { AwsLambdaInstrumentationPatch } from "./aws/services/aws-lambda";
 
 export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
 const awsPropagator = new AWSXRayPropagator();
@@ -65,7 +65,7 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])
       }
     } else if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-lambda') {
       diag.debug('Overriding aws lambda instrumentation');
-      const lambdaInstrumentation = new AwsLambdaInstrumentation({
+      const lambdaInstrumentation = new AwsLambdaInstrumentationPatch({
         eventContextExtractor: customExtractor,
         disableAwsContextPropagation: true,
       });

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts

@@ -0,0 +1,153 @@
+import * as assert from 'assert';
+import * as sinon from 'sinon';
+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";
+
+describe('AwsLambdaInstrumentationPatch', () => {
+  let instrumentation: AwsLambdaInstrumentationPatch;
+
+  beforeEach(() => {
+    instrumentation = new AwsLambdaInstrumentationPatch({});
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  describe('init', () => {
+    it('should skip instrumentation when LAMBDA_TASK_ROOT and _HANDLER are not set', () => {
+      process.env.LAMBDA_TASK_ROOT = '';
+      process.env._HANDLER = '';
+
+      const result = instrumentation.init();
+
+      assert.strictEqual(result.length, 0);
+    });
+
+    it('should fallback to .cjs if .js does not exist', () => {
+      process.env.LAMBDA_TASK_ROOT = '/var/task';
+      process.env._HANDLER = 'src/index.handler';
+
+      sinon.stub(path, 'basename').returns('index.handler');
+      sinon.stub(fs, 'statSync')
+        .onFirstCall().throws(new Error('File not found')) // .js file does not exist
+        .onSecondCall().returns({} as any); // .cjs file exists
+
+      const result = instrumentation.init();
+
+      assert.strictEqual(result[0].name, '/var/task/src/index.cjs');
+    });
+
+    it('should fallback to .mjs when .js and .cjs do not exist', () => {
+      process.env.LAMBDA_TASK_ROOT = '/var/task';
+      process.env._HANDLER = 'src/index.handler';
+
+      sinon.stub(path, 'basename').returns('index.handler');
+      sinon.stub(fs, 'statSync')
+        .onFirstCall().throws(new Error('File not found')) // .js not found
+        .onSecondCall().throws(new Error('File not found')) // .cjs not found
+
+      const result = instrumentation.init();
+
+      assert.strictEqual(result[0].name, '/var/task/src/index.mjs');
+    });
+
+    it('should instrument CommonJS handler correctly', () => {
+      process.env.LAMBDA_TASK_ROOT = '/var/task';
+      process.env._HANDLER = 'src/index.handler';
+
+      sinon.stub(path, 'basename').returns('index.handler');
+      sinon.stub(fs, 'statSync').returns({} as any); // Mock that the .js file exists
+      const debugStub = sinon.stub(diag, 'debug');
+
+      const result = instrumentation.init();
+
+      assert.strictEqual(result.length, 1);
+      assert.strictEqual(result[0].name, '/var/task/src/index.js');
+      assert(result[0] instanceof InstrumentationNodeModuleDefinition);
+      assert.strictEqual(result[0].files.length, 1);
+      assert(debugStub.calledWithMatch('Instrumenting lambda handler', sinon.match.object));
+    });
+
+    it('should return ESM instrumentation for .mjs files or when HANDLER_IS_ESM is set', () => {
+      process.env.LAMBDA_TASK_ROOT = '/var/task';
+      process.env._HANDLER = 'src/index.handler';
+      process.env.HANDLER_IS_ESM = 'true'; // ESM environment variable set
+
+      sinon.stub(path, 'basename').returns('index.handler');
+      sinon.stub(fs, 'statSync').throws(new Error('File not found')); // No .js or .cjs file exists
+
+      const result = instrumentation.init();
+
+      assert.strictEqual(result.length, 1);
+      assert.strictEqual(result[0].name, '/var/task/src/index.mjs');
+      assert(result[0] instanceof InstrumentationNodeModuleDefinition);
+      assert.strictEqual(result[0].files.length, 0); //
+      delete process.env.HANDLER_IS_ESM;
+    });
+  });
+
+  it('should apply and remove patches correctly for a MJS handler', () => {
+    process.env.LAMBDA_TASK_ROOT = '/var/task';
+    process.env._HANDLER = 'src/index.handler';
+    process.env.HANDLER_IS_ESM = 'true'; // ESM environment variable set
+
+    // Mock the module exports object with a sample function
+    const fakeModuleExports = { handler: sinon.stub() };
+
+    const wrapSpy = sinon.spy(instrumentation, '_wrap' as any);
+    const unwrapSpy = sinon.spy(instrumentation, '_unwrap' as any);
+
+    const result = instrumentation.init()[0];
+    // Ensure result contains patch and unpatch functions
+    assert(result.patch, 'patch function should be defined');
+    assert(result.unpatch, 'unpatch function should be defined');
+
+    // Call the patch function with the mocked module exports
+    result.patch(fakeModuleExports);
+
+    // Assert that wrap is called after patching
+    assert(wrapSpy.calledOnce, '_wrap should be called once when patch is applied');
+
+    // Call the unpatch function with the mocked module exports
+    result.unpatch(fakeModuleExports);
+
+    // Assert that unwrap is called after unpatching
+    assert(unwrapSpy.calledOnce, '_unwrap should be called once when unpatch is called');
+
+    delete process.env.HANDLER_IS_ESM;
+  });
+
+  it('should apply and remove patches correctly for a CJS handler', () => {
+    process.env.LAMBDA_TASK_ROOT = '/var/task';
+    process.env._HANDLER = 'src/index.handler';
+
+    // Mock the module exports object with a sample function
+    const fakeModuleExports = { handler: sinon.stub() };
+    sinon.stub(fs, 'statSync').returns({} as any); // Mock that the .js file exists
+
+    const wrapSpy = sinon.spy(instrumentation, '_wrap' as any);
+    const unwrapSpy = sinon.spy(instrumentation, '_unwrap' as any);
+
+    const result = instrumentation.init()[0];
+    // Ensure result contains patch and unpatch functions
+    assert(result.files[0].patch, 'patch function should be defined');
+    assert(result.files[0].unpatch, 'unpatch function should be defined');
+
+    // Call the patch function with the mocked module exports
+    result.files[0].patch(fakeModuleExports);
+
+    // Assert that wrap is called after patching
+    assert(wrapSpy.calledOnce, '_wrap should be called once when patch is applied');
+
+    // Call the unpatch function with the mocked module exports
+    result.files[0].unpatch(fakeModuleExports);
+
+    // Assert that unwrap is called after unpatching
+    assert(unwrapSpy.calledOnce, '_unwrap should be called once when unpatch is called');
+
+  });
+});

lambda-layer/packages/layer/scripts/otel-instrument

@@ -1,5 +1,40 @@
 #!/bin/bash
-export NODE_OPTIONS="${NODE_OPTIONS} --require /opt/wrapper.js"
+isESMScript() {
+  # Lambda function root directory
+  TASK_DIR="/var/task"
+
+  # Flag variables to track conditions
+  local found_mjs=false
+  local is_module=false
+
+  # Check for any files ending with `.mjs`
+  if ls "$TASK_DIR"/*.mjs &>/dev/null; then
+      found_mjs=true
+    fi
+
+  # Check if `package.json` exists and if it contains `"type": "module"`
+  if [ -f "$TASK_DIR/package.json" ]; then
+    # Check for the `"type": "module"` attribute in `package.json`
+    if grep -q '"type": *"module"' "$TASK_DIR/package.json"; then
+      is_module=true
+    fi
+  fi
+
+  # Return true if both conditions are met
+  if $found_mjs || $is_module; then
+    return 0  # 0 in bash means true
+  else
+    return 1  # 1 in bash means false
+  fi
+}
+
+if isESMScript; then
+  export NODE_OPTIONS="${NODE_OPTIONS} --import @aws/aws-distro-opentelemetry-node-autoinstrumentation/register --experimental-loader=@opentelemetry/instrumentation/hook.mjs"
+  export HANDLER_IS_ESM=true
+else
+  export NODE_OPTIONS="${NODE_OPTIONS} --require /opt/wrapper.js"
+fi
+
 export LAMBDA_RESOURCE_ATTRIBUTES="cloud.region=$AWS_REGION,cloud.provider=aws,faas.name=$AWS_LAMBDA_FUNCTION_NAME,faas.version=$AWS_LAMBDA_FUNCTION_VERSION,faas.instance=$AWS_LAMBDA_LOG_STREAM_NAME,aws.log.group.names=$AWS_LAMBDA_LOG_GROUP_NAME";
 
 

lambda-layer/packages/layer/scripts/otel-instrument-esm

@@ -0,0 +1,46 @@
+#!/bin/bash
+export NODE_OPTIONS="${NODE_OPTIONS} --import @aws/aws-distro-opentelemetry-node-autoinstrumentation/register --experimental-loader=@opentelemetry/instrumentation/hook.mjs"
+export LAMBDA_RESOURCE_ATTRIBUTES="cloud.region=$AWS_REGION,cloud.provider=aws,faas.name=$AWS_LAMBDA_FUNCTION_NAME,faas.version=$AWS_LAMBDA_FUNCTION_VERSION,faas.instance=$AWS_LAMBDA_LOG_STREAM_NAME,aws.log.group.names=$AWS_LAMBDA_LOG_GROUP_NAME";
+export HANDLER_IS_ESM=true
+
+# - If OTEL_EXPORTER_OTLP_PROTOCOL is not set by user, the default exporting protocol is http/protobuf.
+if [ -z "${OTEL_EXPORTER_OTLP_PROTOCOL}" ]; then
+    export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf
+fi
+
+# - If OTEL_NODE_ENABLED_INSTRUMENTATIONS is not set by user, use default instrumentation
+if [ -z "${OTEL_NODE_ENABLED_INSTRUMENTATIONS}" ]; then
+    export OTEL_NODE_ENABLED_INSTRUMENTATIONS="aws-lambda,aws-sdk"
+fi
+
+# - Set the service name
+if [ -z "${OTEL_SERVICE_NAME}" ]; then
+    export OTEL_SERVICE_NAME=$AWS_LAMBDA_FUNCTION_NAME;
+fi
+
+# - Set the propagators
+if [[ -z "$OTEL_PROPAGATORS" ]]; then
+  export OTEL_PROPAGATORS="tracecontext,baggage,xray"
+fi
+
+# - Set Application Signals configuration
+if [ -z "${OTEL_AWS_APPLICATION_SIGNALS_ENABLED}" ]; then
+    export OTEL_AWS_APPLICATION_SIGNALS_ENABLED="true";
+fi
+
+if [ -z "${OTEL_METRICS_EXPORTER}" ]; then
+    export OTEL_METRICS_EXPORTER="none";
+fi
+
+# - Append Lambda Resource Attributes to OTel Resource Attribute List
+if [ -z "${OTEL_RESOURCE_ATTRIBUTES}" ]; then
+    export OTEL_RESOURCE_ATTRIBUTES=$LAMBDA_RESOURCE_ATTRIBUTES;
+else
+    export OTEL_RESOURCE_ATTRIBUTES="$LAMBDA_RESOURCE_ATTRIBUTES,$OTEL_RESOURCE_ATTRIBUTES";
+fi
+
+if [[ $OTEL_RESOURCE_ATTRIBUTES != *"service.name="* ]]; then
+  export OTEL_RESOURCE_ATTRIBUTES="service.name=${AWS_LAMBDA_FUNCTION_NAME},${OTEL_RESOURCE_ATTRIBUTES}"
+fi
+
+exec "$@"