feat(instrumentation-aws-lambda): support esm handlers and all other patterns

Author: raphael-theriault-swi

plugins/node/opentelemetry-instrumentation-aws-lambda/README.md

@@ -35,6 +35,9 @@ const { registerInstrumentations } = require('@opentelemetry/instrumentation');
 const provider = new NodeTracerProvider();
 provider.register();
 
+// Note that this needs to appear after the tracer provider is registered,
+// otherwise the instrumentation will not properly flush data after each
+// lambda invocation.
 registerInstrumentations({
   instrumentations: [
     new AwsLambdaInstrumentation({
@@ -46,7 +49,7 @@ registerInstrumentations({
 
 In your Lambda function configuration, add or update the `NODE_OPTIONS` environment variable to require the wrapper, e.g.,
 
-`NODE_OPTIONS=--require lambda-wrapper`
+`NODE_OPTIONS=--require lambda-wrapper --experimental-loader @opentelemetry/instrumentation/hook.mjs`
 
 ## AWS Lambda Instrumentation Options
 

plugins/node/opentelemetry-instrumentation-aws-lambda/package.json

@@ -44,6 +44,7 @@
   },
   "devDependencies": {
     "@opentelemetry/api": "^1.3.0",
+    "@opentelemetry/contrib-test-utils": "^0.40.0",
     "@opentelemetry/core": "^1.8.0",
     "@opentelemetry/sdk-metrics": "^1.8.0",
     "@opentelemetry/sdk-trace-base": "^1.8.0",

plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts

@@ -14,8 +14,7 @@
  * limitations under the License.
  */
 
-import * as path from 'path';
-import * as fs from 'fs';
+import { env } from 'process';
 
 import {
   InstrumentationBase,
@@ -58,8 +57,13 @@ import {
 
 import { AwsLambdaInstrumentationConfig, EventContextExtractor } from './types';
 import { PACKAGE_NAME, PACKAGE_VERSION } from './version';
-import { env } from 'process';
-import { LambdaModule } from './internal-types';
+import {
+  isInvalidHandler,
+  moduleRootAndHandler,
+  resolveHandler,
+  splitHandlerString,
+  tryPath,
+} from './user-function';
 
 const awsPropagator = new AWSXRayPropagator();
 const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
@@ -92,6 +96,7 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
         this._config.disableAwsContextPropagation = true;
       }
     }
+    this._traceForceFlusher = this._traceForceFlush(trace.getTracerProvider());
   }
 
   override setConfig(config: AwsLambdaInstrumentationConfig = {}) {
@@ -110,61 +115,106 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
       );
       return [];
     }
+    if (isInvalidHandler(handlerDef)) {
+      this._diag.debug(
+        'Skipping lambda instrumentation: _HANDLER/lambdaHandler is invalid.',
+        { 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 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
-        filename += '.cjs';
-      }
+    const [moduleRoot, moduleAndHandler] = moduleRootAndHandler(handlerDef);
+    const [module, handlerPath] = splitHandlerString(moduleAndHandler);
+
+    if (!module || !handlerPath) {
+      this._diag.debug(
+        'Skipping lambda instrumentation: _HANDLER/lambdaHandler is invalid.',
+        { taskRoot, handlerDef, moduleRoot, module, handlerPath }
+      );
+      return [];
+    }
+
+    const filename = tryPath(taskRoot, moduleRoot, module);
+    if (!filename) {
+      this._diag.debug(
+        'Skipping lambda instrumentation: _HANDLER/lambdaHandler and LAMBDA_TASK_ROOT did not resolve to a file.',
+        {
+          taskRoot,
+          handlerDef,
+          moduleRoot,
+          module,
+          handlerPath,
+        }
+      );
+      return [];
     }
 
     diag.debug('Instrumenting lambda handler', {
       taskRoot,
       handlerDef,
-      handler,
+      filename,
       moduleRoot,
       module,
-      filename,
-      functionName,
+      handlerPath,
     });
 
+    const patch = (moduleExports: object) => {
+      const [container, functionName] = resolveHandler(
+        moduleExports,
+        handlerPath
+      );
+      if (
+        container == null ||
+        functionName == null ||
+        typeof container[functionName] !== 'function'
+      ) {
+        this._diag.debug(
+          'Skipping lambda instrumentation: _HANDLER/lambdaHandler did not resolve to a function.',
+          {
+            taskRoot,
+            handlerDef,
+            filename,
+            moduleRoot,
+            module,
+            handlerPath,
+          }
+        );
+        return moduleExports;
+      }
+
+      if (isWrapped(container[functionName])) {
+        this._unwrap(container, functionName);
+      }
+      this._wrap(container, functionName, this._getHandler());
+      return moduleExports;
+    };
+    const unpatch = (moduleExports?: object) => {
+      if (moduleExports == null) return;
+      const [container, functionName] = resolveHandler(
+        moduleExports,
+        handlerPath
+      );
+      if (
+        container == null ||
+        functionName == null ||
+        typeof container[functionName] !== 'function'
+      ) {
+        return;
+      }
+
+      this._unwrap(container, functionName);
+    };
+
     return [
       new InstrumentationNodeModuleDefinition(
-        // NB: The patching infrastructure seems to match names backwards, this must be the filename, while
-        // InstrumentationNodeModuleFile must be the module name.
+        // The patching infrastructure properly supports absolute paths when registering hooks but not when
+        // actually matching against filenames when patching, so we need to provide a file instrumentation
+        // that will actually match by using a relative path.
         filename,
         ['*'],
-        undefined,
-        undefined,
-        [
-          new InstrumentationNodeModuleFile(
-            module,
-            ['*'],
-            (moduleExports: LambdaModule) => {
-              if (isWrapped(moduleExports[functionName])) {
-                this._unwrap(moduleExports, functionName);
-              }
-              this._wrap(moduleExports, functionName, this._getHandler());
-              return moduleExports;
-            },
-            (moduleExports?: LambdaModule) => {
-              if (moduleExports == null) return;
-              this._unwrap(moduleExports, functionName);
-            }
-          ),
-        ]
+        patch,
+        unpatch,
+        [new InstrumentationNodeModuleFile(module, ['*'], patch, unpatch)]
       ),
     ];
   }

plugins/node/opentelemetry-instrumentation-aws-lambda/src/user-function.ts

@@ -0,0 +1,107 @@
+/**
+ * Adapted from https://github.com/aws/aws-lambda-nodejs-runtime-interface-client/blob/v3.1.0/src/UserFunction.js
+ */
+
+import * as path from 'path';
+import * as fs from 'fs';
+import { LambdaModule } from './internal-types';
+
+const FUNCTION_EXPR = /^([^.]*)\.(.*)$/;
+const RELATIVE_PATH_SUBSTRING = '..';
+
+/**
+ * Break the full handler string into two pieces, the module root and the actual
+ * handler string.
+ * Given './somepath/something/module.nestedobj.handler' this returns
+ * ['./somepath/something', 'module.nestedobj.handler']
+ */
+export function moduleRootAndHandler(
+  fullHandlerString: string
+): [moduleRoot: string, handler: string] {
+  const handlerString = path.basename(fullHandlerString);
+  const moduleRoot = fullHandlerString.substring(
+    0,
+    fullHandlerString.indexOf(handlerString)
+  );
+  return [moduleRoot, handlerString];
+}
+
+/**
+ * Split the handler string into two pieces: the module name and the path to
+ * the handler function.
+ */
+export function splitHandlerString(
+  handler: string
+): [module: string | undefined, functionPath: string | undefined] {
+  const match = handler.match(FUNCTION_EXPR);
+  if (match && match.length === 3) {
+    return [match[1], match[2]];
+  } else {
+    return [undefined, undefined];
+  }
+}
+
+/**
+ * Resolve the user's handler function key and its containing object from the module.
+ */
+export function resolveHandler(
+  object: object,
+  nestedProperty: string
+): [container: LambdaModule | undefined, handlerKey: string | undefined] {
+  const nestedPropertyKeys = nestedProperty.split('.');
+  const handlerKey = nestedPropertyKeys.pop();
+
+  const container = nestedPropertyKeys.reduce<object | undefined>(
+    (nested, key) => {
+      return nested && (nested as Partial<Record<string, object>>)[key];
+    },
+    object
+  );
+
+  if (container) {
+    return [container as LambdaModule, handlerKey];
+  } else {
+    return [undefined, undefined];
+  }
+}
+
+/**
+ * Attempt to determine the user's module path.
+ */
+export function tryPath(
+  appRoot: string,
+  moduleRoot: string,
+  module: string
+): string | undefined {
+  const lambdaStylePath = path.resolve(appRoot, moduleRoot, module);
+
+  const extensionless = fs.existsSync(lambdaStylePath);
+  if (extensionless) {
+    return lambdaStylePath;
+  }
+
+  const extensioned =
+    (fs.existsSync(lambdaStylePath + '.js') && lambdaStylePath + '.js') ||
+    (fs.existsSync(lambdaStylePath + '.mjs') && lambdaStylePath + '.mjs') ||
+    (fs.existsSync(lambdaStylePath + '.cjs') && lambdaStylePath + '.cjs');
+  if (extensioned) {
+    return extensioned;
+  }
+
+  try {
+    const nodeStylePath = require.resolve(module, {
+      paths: [appRoot, moduleRoot],
+    });
+    return nodeStylePath;
+  } catch {
+    return undefined;
+  }
+}
+
+export function isInvalidHandler(fullHandlerString: string): boolean {
+  if (fullHandlerString.includes(RELATIVE_PATH_SUBSTRING)) {
+    return true;
+  } else {
+    return false;
+  }
+}