remove fs dependency from lambda hook (#5404)

Author: rochdev

packages/dd-trace/src/lambda/runtime/patch.js

@@ -2,7 +2,7 @@
 
 const path = require('path')
 
-const { _extractModuleNameAndHandlerPath, _extractModuleRootAndHandler, _getLambdaFilePath } = require('./ritm')
+const { _extractModuleNameAndHandlerPath, _extractModuleRootAndHandler, _getLambdaFilePaths } = require('./ritm')
 const { datadog } = require('../handler')
 const { addHook } = require('../../../../datadog-instrumentations/src/helpers/instrument')
 const shimmer = require('../../../../datadog-shimmer')
@@ -65,9 +65,11 @@ if (originalLambdaHandler !== undefined) {
   const [_module, handlerPath] = _extractModuleNameAndHandlerPath(moduleAndHandler)
 
   const lambdaStylePath = path.resolve(lambdaTaskRoot, moduleRoot, _module)
-  const lambdaFilePath = _getLambdaFilePath(lambdaStylePath)
+  const lambdaFilePaths = _getLambdaFilePaths(lambdaStylePath)
 
-  addHook({ name: lambdaFilePath }, patchLambdaModule(handlerPath))
+  for (const lambdaFilePath of lambdaFilePaths) {
+    addHook({ name: lambdaFilePath }, patchLambdaModule(handlerPath))
+  }
 } else {
   // Instrumentation is done manually.
   addHook({ name: 'datadog-lambda-js' }, patchDatadogLambdaModule)

packages/dd-trace/src/lambda/runtime/ritm.js

@@ -7,7 +7,6 @@
  */
 'use strict'
 
-const fs = require('fs')
 const path = require('path')
 
 const log = require('../../log')
@@ -60,23 +59,18 @@ function _extractModuleNameAndHandlerPath (handler) {
 }
 
 /**
- * Returns the correct path of the file to be patched
- * when required.
+ * Returns all possible paths of the files to be patched when required.
  *
  * @param {*} lambdaStylePath the path comprised of the `LAMBDA_TASK_ROOT`,
  * the root of the module of the Lambda handler, and the module name.
- * @returns the lambdaStylePath with the appropiate extension for the hook.
+ * @returns the lambdaStylePath with appropiate extensions for the hook.
  */
-function _getLambdaFilePath (lambdaStylePath) {
-  let lambdaFilePath = lambdaStylePath
-  if (fs.existsSync(lambdaStylePath + '.js')) {
-    lambdaFilePath += '.js'
-  } else if (fs.existsSync(lambdaStylePath + '.mjs')) {
-    lambdaFilePath += '.mjs'
-  } else if (fs.existsSync(lambdaStylePath + '.cjs')) {
-    lambdaFilePath += '.cjs'
-  }
-  return lambdaFilePath
+function _getLambdaFilePaths (lambdaStylePath) {
+  return [
+    `${lambdaStylePath}.js`,
+    `${lambdaStylePath}.mjs`,
+    `${lambdaStylePath}.cjs`
+  ]
 }
 
 /**
@@ -92,12 +86,13 @@ const registerLambdaHook = () => {
     const [_module] = _extractModuleNameAndHandlerPath(moduleAndHandler)
 
     const lambdaStylePath = path.resolve(lambdaTaskRoot, moduleRoot, _module)
-    const lambdaFilePath = _getLambdaFilePath(lambdaStylePath)
+    const lambdaFilePaths = _getLambdaFilePaths(lambdaStylePath)
 
-    Hook([lambdaFilePath], (moduleExports) => {
+    // TODO: Redo this like any other instrumentation.
+    Hook(lambdaFilePaths, (moduleExports, name) => {
       require('./patch')
 
-      for (const { hook } of instrumentations[lambdaFilePath]) {
+      for (const { hook } of instrumentations[name]) {
         try {
           moduleExports = hook(moduleExports)
         } catch (e) {
@@ -133,6 +128,6 @@ const registerLambdaHook = () => {
 module.exports = {
   _extractModuleRootAndHandler,
   _extractModuleNameAndHandlerPath,
-  _getLambdaFilePath,
+  _getLambdaFilePaths,
   registerLambdaHook
 }