Skip to content

Commit 365a4e3

Browse files
mxiamxiamxiamxia
committed
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
1 parent 032677f commit 365a4e3

File tree

5 files changed

+362
-3
lines changed

5 files changed

+362
-3
lines changed

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

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
// Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
4+
5+
import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lambda';
6+
import * as path from 'path';
7+
import * as fs from 'fs';
8+
import { InstrumentationNodeModuleDefinition, InstrumentationNodeModuleFile, isWrapped } from '@opentelemetry/instrumentation';
9+
import { diag } from '@opentelemetry/api';
10+
11+
export class AwsLambdaInstrumentationPatch extends AwsLambdaInstrumentation {
12+
13+
override init() {
14+
// Custom logic before calling the original implementation
15+
diag.debug('Initializing AwsLambdaInstrumentationPatch');
16+
const taskRoot = process.env.LAMBDA_TASK_ROOT;
17+
const handlerDef = this._config.lambdaHandler ?? process.env._HANDLER;
18+
19+
// _HANDLER and LAMBDA_TASK_ROOT are always defined in Lambda but guard bail out if in the future this changes.
20+
if (!taskRoot || !handlerDef) {
21+
this._diag.debug(
22+
'Skipping lambda instrumentation: no _HANDLER/lambdaHandler or LAMBDA_TASK_ROOT.',
23+
{ taskRoot, handlerDef }
24+
);
25+
return [];
26+
}
27+
28+
const handler = path.basename(handlerDef);
29+
const moduleRoot = handlerDef.substr(0, handlerDef.length - handler.length);
30+
31+
const [module, functionName] = handler.split('.', 2);
32+
33+
// Lambda loads user function using an absolute path.
34+
let filename = path.resolve(taskRoot, moduleRoot, module);
35+
if (!filename.endsWith('.js')) {
36+
// its impossible to know in advance if the user has a cjs or js or mjs file.
37+
// check that the .js file exists otherwise fallback to next known possibility
38+
try {
39+
fs.statSync(`${filename}.js`);
40+
filename += '.js';
41+
} catch (e) {
42+
// fallback to .cjs
43+
try {
44+
fs.statSync(`${filename}.cjs`);
45+
filename += '.cjs';
46+
} catch (e) {
47+
// fall back to .mjs
48+
filename += '.mjs';
49+
}
50+
}
51+
}
52+
53+
diag.debug('Instrumenting lambda handler', {
54+
taskRoot,
55+
handlerDef,
56+
handler,
57+
moduleRoot,
58+
module,
59+
filename,
60+
functionName,
61+
});
62+
63+
if (filename.endsWith('.mjs') || process.env.HANDLER_IS_ESM) {
64+
return [
65+
new InstrumentationNodeModuleDefinition(
66+
// NB: The patching infrastructure seems to match names backwards, this must be the filename, while
67+
// InstrumentationNodeModuleFile must be the module name.
68+
filename,
69+
['*'],
70+
(moduleExports: any) => {
71+
diag.debug('Applying patch for lambda esm handler');
72+
if (isWrapped(moduleExports[functionName])) {
73+
this._unwrap(moduleExports, functionName);
74+
}
75+
this._wrap(
76+
moduleExports,
77+
functionName,
78+
(this as any)._getHandler()
79+
);
80+
return moduleExports;
81+
},
82+
(moduleExports?: any) => {
83+
if (moduleExports == null) return;
84+
diag.debug('Removing patch for lambda esm handler');
85+
this._unwrap(moduleExports, functionName);
86+
}
87+
)
88+
];
89+
} else {
90+
return [
91+
new InstrumentationNodeModuleDefinition(
92+
// NB: The patching infrastructure seems to match names backwards, this must be the filename, while
93+
// InstrumentationNodeModuleFile must be the module name.
94+
filename,
95+
['*'],
96+
undefined,
97+
undefined,
98+
[
99+
new InstrumentationNodeModuleFile(
100+
module,
101+
['*'],
102+
(moduleExports: any) => {
103+
diag.debug('Applying patch for lambda handler');
104+
if (isWrapped(moduleExports[functionName])) {
105+
this._unwrap(moduleExports, functionName);
106+
}
107+
this._wrap(
108+
moduleExports,
109+
functionName,
110+
(this as any)._getHandler()
111+
);
112+
return moduleExports;
113+
},
114+
(moduleExports?: any) => {
115+
if (moduleExports == null) return;
116+
diag.debug('Removing patch for lambda handler');
117+
this._unwrap(moduleExports, functionName);
118+
}
119+
),
120+
]
121+
),
122+
];
123+
}
124+
}
125+
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import {
1212
trace,
1313
} from '@opentelemetry/api';
1414
import { Instrumentation } from '@opentelemetry/instrumentation';
15-
import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lambda';
1615
import { AwsSdkInstrumentationConfig, NormalizedRequest } from '@opentelemetry/instrumentation-aws-sdk';
1716
import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
1817
import { APIGatewayProxyEventHeaders, Context } from 'aws-lambda';
@@ -26,6 +25,7 @@ import {
2625
} from './aws/services/bedrock';
2726
import { KinesisServiceExtension } from './aws/services/kinesis';
2827
import { S3ServiceExtension } from './aws/services/s3';
28+
import { AwsLambdaInstrumentationPatch } from "./aws/services/aws-lambda";
2929

3030
export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
3131
const awsPropagator = new AWSXRayPropagator();
@@ -65,7 +65,7 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])
6565
}
6666
} else if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-lambda') {
6767
diag.debug('Overriding aws lambda instrumentation');
68-
const lambdaInstrumentation = new AwsLambdaInstrumentation({
68+
const lambdaInstrumentation = new AwsLambdaInstrumentationPatch({
6969
eventContextExtractor: customExtractor,
7070
disableAwsContextPropagation: true,
7171
});

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

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
import * as assert from 'assert';
2+
import * as sinon from 'sinon';
3+
import * as path from 'path';
4+
import * as fs from 'fs';
5+
import { diag } from '@opentelemetry/api';
6+
import { InstrumentationNodeModuleDefinition } from '@opentelemetry/instrumentation';
7+
import { AwsLambdaInstrumentationPatch } from "../../../../src/patches/aws/services/aws-lambda";
8+
9+
describe('AwsLambdaInstrumentationPatch', () => {
10+
let instrumentation: AwsLambdaInstrumentationPatch;
11+
12+
beforeEach(() => {
13+
instrumentation = new AwsLambdaInstrumentationPatch({});
14+
});
15+
16+
afterEach(() => {
17+
sinon.restore();
18+
});
19+
20+
describe('init', () => {
21+
it('should skip instrumentation when LAMBDA_TASK_ROOT and _HANDLER are not set', () => {
22+
process.env.LAMBDA_TASK_ROOT = '';
23+
process.env._HANDLER = '';
24+
25+
const result = instrumentation.init();
26+
27+
assert.strictEqual(result.length, 0);
28+
});
29+
30+
it('should fallback to .cjs if .js does not exist', () => {
31+
process.env.LAMBDA_TASK_ROOT = '/var/task';
32+
process.env._HANDLER = 'src/index.handler';
33+
34+
sinon.stub(path, 'basename').returns('index.handler');
35+
sinon.stub(fs, 'statSync')
36+
.onFirstCall().throws(new Error('File not found')) // .js file does not exist
37+
.onSecondCall().returns({} as any); // .cjs file exists
38+
39+
const result = instrumentation.init();
40+
41+
assert.strictEqual(result[0].name, '/var/task/src/index.cjs');
42+
});
43+
44+
it('should fallback to .mjs when .js and .cjs do not exist', () => {
45+
process.env.LAMBDA_TASK_ROOT = '/var/task';
46+
process.env._HANDLER = 'src/index.handler';
47+
48+
sinon.stub(path, 'basename').returns('index.handler');
49+
sinon.stub(fs, 'statSync')
50+
.onFirstCall().throws(new Error('File not found')) // .js not found
51+
.onSecondCall().throws(new Error('File not found')) // .cjs not found
52+
53+
const result = instrumentation.init();
54+
55+
assert.strictEqual(result[0].name, '/var/task/src/index.mjs');
56+
});
57+
58+
it('should instrument CommonJS handler correctly', () => {
59+
process.env.LAMBDA_TASK_ROOT = '/var/task';
60+
process.env._HANDLER = 'src/index.handler';
61+
62+
sinon.stub(path, 'basename').returns('index.handler');
63+
sinon.stub(fs, 'statSync').returns({} as any); // Mock that the .js file exists
64+
const debugStub = sinon.stub(diag, 'debug');
65+
66+
const result = instrumentation.init();
67+
68+
assert.strictEqual(result.length, 1);
69+
assert.strictEqual(result[0].name, '/var/task/src/index.js');
70+
assert(result[0] instanceof InstrumentationNodeModuleDefinition);
71+
assert.strictEqual(result[0].files.length, 1);
72+
assert(debugStub.calledWithMatch('Instrumenting lambda handler', sinon.match.object));
73+
});
74+
75+
it('should return ESM instrumentation for .mjs files or when HANDLER_IS_ESM is set', () => {
76+
process.env.LAMBDA_TASK_ROOT = '/var/task';
77+
process.env._HANDLER = 'src/index.handler';
78+
process.env.HANDLER_IS_ESM = 'true'; // ESM environment variable set
79+
80+
sinon.stub(path, 'basename').returns('index.handler');
81+
sinon.stub(fs, 'statSync').throws(new Error('File not found')); // No .js or .cjs file exists
82+
83+
const result = instrumentation.init();
84+
85+
assert.strictEqual(result.length, 1);
86+
assert.strictEqual(result[0].name, '/var/task/src/index.mjs');
87+
assert(result[0] instanceof InstrumentationNodeModuleDefinition);
88+
assert.strictEqual(result[0].files.length, 0); //
89+
delete process.env.HANDLER_IS_ESM;
90+
});
91+
});
92+
93+
it('should apply and remove patches correctly for a MJS handler', () => {
94+
process.env.LAMBDA_TASK_ROOT = '/var/task';
95+
process.env._HANDLER = 'src/index.handler';
96+
process.env.HANDLER_IS_ESM = 'true'; // ESM environment variable set
97+
98+
// Mock the module exports object with a sample function
99+
const fakeModuleExports = { handler: sinon.stub() };
100+
101+
const wrapSpy = sinon.spy(instrumentation, '_wrap' as any);
102+
const unwrapSpy = sinon.spy(instrumentation, '_unwrap' as any);
103+
104+
const result = instrumentation.init()[0];
105+
// Ensure result contains patch and unpatch functions
106+
assert(result.patch, 'patch function should be defined');
107+
assert(result.unpatch, 'unpatch function should be defined');
108+
109+
// Call the patch function with the mocked module exports
110+
result.patch(fakeModuleExports);
111+
112+
// Assert that wrap is called after patching
113+
assert(wrapSpy.calledOnce, '_wrap should be called once when patch is applied');
114+
115+
// Call the unpatch function with the mocked module exports
116+
result.unpatch(fakeModuleExports);
117+
118+
// Assert that unwrap is called after unpatching
119+
assert(unwrapSpy.calledOnce, '_unwrap should be called once when unpatch is called');
120+
121+
delete process.env.HANDLER_IS_ESM;
122+
});
123+
124+
it('should apply and remove patches correctly for a CJS handler', () => {
125+
process.env.LAMBDA_TASK_ROOT = '/var/task';
126+
process.env._HANDLER = 'src/index.handler';
127+
128+
// Mock the module exports object with a sample function
129+
const fakeModuleExports = { handler: sinon.stub() };
130+
sinon.stub(fs, 'statSync').returns({} as any); // Mock that the .js file exists
131+
132+
const wrapSpy = sinon.spy(instrumentation, '_wrap' as any);
133+
const unwrapSpy = sinon.spy(instrumentation, '_unwrap' as any);
134+
135+
const result = instrumentation.init()[0];
136+
// Ensure result contains patch and unpatch functions
137+
assert(result.files[0].patch, 'patch function should be defined');
138+
assert(result.files[0].unpatch, 'unpatch function should be defined');
139+
140+
// Call the patch function with the mocked module exports
141+
result.files[0].patch(fakeModuleExports);
142+
143+
// Assert that wrap is called after patching
144+
assert(wrapSpy.calledOnce, '_wrap should be called once when patch is applied');
145+
146+
// Call the unpatch function with the mocked module exports
147+
result.files[0].unpatch(fakeModuleExports);
148+
149+
// Assert that unwrap is called after unpatching
150+
assert(unwrapSpy.calledOnce, '_unwrap should be called once when unpatch is called');
151+
152+
});
153+
});

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

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,40 @@
11
#!/bin/bash
2-
export NODE_OPTIONS="${NODE_OPTIONS} --require /opt/wrapper.js"
2+
isESMScript() {
3+
# Lambda function root directory
4+
TASK_DIR="/var/task"
5+
6+
# Flag variables to track conditions
7+
local found_mjs=false
8+
local is_module=false
9+
10+
# Check for any files ending with `.mjs`
11+
if ls "$TASK_DIR"/*.mjs &>/dev/null; then
12+
found_mjs=true
13+
fi
14+
15+
# Check if `package.json` exists and if it contains `"type": "module"`
16+
if [ -f "$TASK_DIR/package.json" ]; then
17+
# Check for the `"type": "module"` attribute in `package.json`
18+
if grep -q '"type": *"module"' "$TASK_DIR/package.json"; then
19+
is_module=true
20+
fi
21+
fi
22+
23+
# Return true if both conditions are met
24+
if $found_mjs || $is_module; then
25+
return 0 # 0 in bash means true
26+
else
27+
return 1 # 1 in bash means false
28+
fi
29+
}
30+
31+
if isESMScript; then
32+
export NODE_OPTIONS="${NODE_OPTIONS} --import @aws/aws-distro-opentelemetry-node-autoinstrumentation/register --experimental-loader=@opentelemetry/instrumentation/hook.mjs"
33+
export HANDLER_IS_ESM=true
34+
else
35+
export NODE_OPTIONS="${NODE_OPTIONS} --require /opt/wrapper.js"
36+
fi
37+
338
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";
439

540

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
#!/bin/bash
2+
export NODE_OPTIONS="${NODE_OPTIONS} --import @aws/aws-distro-opentelemetry-node-autoinstrumentation/register --experimental-loader=@opentelemetry/instrumentation/hook.mjs"
3+
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";
4+
export HANDLER_IS_ESM=true
5+
6+
# - If OTEL_EXPORTER_OTLP_PROTOCOL is not set by user, the default exporting protocol is http/protobuf.
7+
if [ -z "${OTEL_EXPORTER_OTLP_PROTOCOL}" ]; then
8+
export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf
9+
fi
10+
11+
# - If OTEL_NODE_ENABLED_INSTRUMENTATIONS is not set by user, use default instrumentation
12+
if [ -z "${OTEL_NODE_ENABLED_INSTRUMENTATIONS}" ]; then
13+
export OTEL_NODE_ENABLED_INSTRUMENTATIONS="aws-lambda,aws-sdk"
14+
fi
15+
16+
# - Set the service name
17+
if [ -z "${OTEL_SERVICE_NAME}" ]; then
18+
export OTEL_SERVICE_NAME=$AWS_LAMBDA_FUNCTION_NAME;
19+
fi
20+
21+
# - Set the propagators
22+
if [[ -z "$OTEL_PROPAGATORS" ]]; then
23+
export OTEL_PROPAGATORS="tracecontext,baggage,xray"
24+
fi
25+
26+
# - Set Application Signals configuration
27+
if [ -z "${OTEL_AWS_APPLICATION_SIGNALS_ENABLED}" ]; then
28+
export OTEL_AWS_APPLICATION_SIGNALS_ENABLED="true";
29+
fi
30+
31+
if [ -z "${OTEL_METRICS_EXPORTER}" ]; then
32+
export OTEL_METRICS_EXPORTER="none";
33+
fi
34+
35+
# - Append Lambda Resource Attributes to OTel Resource Attribute List
36+
if [ -z "${OTEL_RESOURCE_ATTRIBUTES}" ]; then
37+
export OTEL_RESOURCE_ATTRIBUTES=$LAMBDA_RESOURCE_ATTRIBUTES;
38+
else
39+
export OTEL_RESOURCE_ATTRIBUTES="$LAMBDA_RESOURCE_ATTRIBUTES,$OTEL_RESOURCE_ATTRIBUTES";
40+
fi
41+
42+
if [[ $OTEL_RESOURCE_ATTRIBUTES != *"service.name="* ]]; then
43+
export OTEL_RESOURCE_ATTRIBUTES="service.name=${AWS_LAMBDA_FUNCTION_NAME},${OTEL_RESOURCE_ATTRIBUTES}"
44+
fi
45+
46+
exec "$@"

0 commit comments

Comments
 (0)