Patch aws-lambda instrumentation to support ESM (#101)

mxiamxia · web-flow · commit 2df6d5d977d6 · 2024-10-10T18:42:19.000-07:00
*Description of changes:* 1. Add `IS_ESM` check in `otel_instrument` wrapper before calling lambda handler. If it is EMS format, add ESM node-options [supported](https://github.com/open-telemetry/opentelemetry-js/blob/966ac176af249d86de6cb10feac2306062846768/doc/esm-support.md#esm-options-for-different-versions-of-nodejs) by OTel community. Added a new wrapper script `otel_instrument_esm` if the esm auto detection logic is failed, customer can opt-in to tell tell layer go with ESM instrumentation path. 2. Set a new env var `HANDLER_IS_ESM` for lambda function when ESM is detected 3. Patch aws-lambda instrumentation, when `IS_ESM` env var is set, apply ESM supported `InstrumentationNodeModuleDefinition` to patch function handler, otherwise keep using the existing handler patcher. Note: this change add a new branch for supporting ESM, the existing CommonJS path should not be impacted. 1. open-telemetry/opentelemetry-js#4842 2. open-telemetry/opentelemetry-js-contrib#1942 *Test* ``` 2024-10-09T20:38:13.411-07:00 | Instrumenting lambda handler { -- | -- | 2024-10-09T20:38:13.411-07:00 | taskRoot: '/var/task', | 2024-10-09T20:38:13.411-07:00 | handlerDef: 'index.handler', | 2024-10-09T20:38:13.411-07:00 | handler: 'index.handler', | 2024-10-09T20:38:13.411-07:00 | moduleRoot: '', | 2024-10-09T20:38:13.411-07:00 | module: 'index', | 2024-10-09T20:38:13.411-07:00 | filename: '/var/task/index.mjs', | 2024-10-09T20:38:13.411-07:00 | functionName: 'handler' | 2024-10-09T20:38:13.411-07:00 | } | 2024-10-09T20:38:15.386-07:00 | 'cloud.account.id': '252610625673', -- | -- | -- | 2024-10-09T20:38:15.386-07:00 | 'aws.is.local.root': true, | 2024-10-09T20:38:15.386-07:00 | 'aws.local.service': 'TestESM', | 2024-10-09T20:38:15.386-07:00 | 'aws.local.operation': 'TestESM/Handler', | 2024-10-09T20:38:15.386-07:00 | 'aws.span.kind': 'LOCAL_ROOT' ``` By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/aws-lambda.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/aws-lambda.ts
@@ -0,0 +1,120 @@
+// 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);
+              }
+            ),
+          ]
+        ),
+      ];
+    }
+  }
+}
diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts b/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,
       });
diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts
@@ -0,0 +1,160 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+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');
+  });
+});
diff --git a/lambda-layer/packages/layer/scripts/otel-instrument b/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";
 
 
diff --git a/lambda-layer/packages/layer/scripts/otel-instrument-esm b/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 "$@"
