From 365a4e3e09489013d295900a1beb2b29da03f5fb Mon Sep 17 00:00:00 2001 From: Min Xia Date: Wed, 9 Oct 2024 20:52:59 -0700 Subject: [PATCH] Patch aws-lambda instrumentation to support ESM commit d25c3c3bca13efe1a78eb5a8d3bd0813368084b2 Author: Min Xia Date: Tue Oct 8 15:45:36 2024 -0700 fix the lint error --- .../src/patches/aws/services/aws-lambda.ts | 125 ++++++++++++++ .../src/patches/instrumentation-patch.ts | 4 +- .../patches/aws/services/aws-lambda.test.ts | 153 ++++++++++++++++++ .../packages/layer/scripts/otel-instrument | 37 ++++- .../layer/scripts/otel-instrument-esm | 46 ++++++ 5 files changed, 362 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..36304e8 --- /dev/null +++ b/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); + } + ), + ] + ), + ]; + } + } +} 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..bf7c4a9 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..83596e0 --- /dev/null +++ b/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'); + + }); +}); 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 "$@"