From 2df6d5d977d6ffe4b2c3cda11758ab4aec259ee6 Mon Sep 17 00:00:00 2001 From: Min Xia Date: Thu, 10 Oct 2024 18:42:19 -0700 Subject: [PATCH] Patch aws-lambda instrumentation to support ESM (#101) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit *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. https://github.com/open-telemetry/opentelemetry-js/issues/4842 2. https://github.com/open-telemetry/opentelemetry-js-contrib/issues/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. --- .../src/patches/aws/services/aws-lambda.ts | 120 +++++++++++++ .../src/patches/instrumentation-patch.ts | 4 +- .../patches/aws/services/aws-lambda.test.ts | 160 ++++++++++++++++++ .../packages/layer/scripts/otel-instrument | 37 +++- .../layer/scripts/otel-instrument-esm | 46 +++++ 5 files changed, 364 insertions(+), 3 deletions(-) create mode 100644 aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/aws-lambda.ts create mode 100644 aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts create mode 100644 lambda-layer/packages/layer/scripts/otel-instrument-esm 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 new file mode 100644 index 0000000..6906bdd --- /dev/null +++ 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 index 776d4c3..7cc98bd 100644 --- 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 new file mode 100644 index 0000000..d68e1e6 --- /dev/null +++ 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 index fc34731..f5394e8 100644 --- 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 new file mode 100644 index 0000000..c6e9511 --- /dev/null +++ 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 "$@"