fix: Handle case where Lambda handler is using callbacks instead of promises (#3404)

Author: astuyve

packages/dd-trace/src/lambda/handler.js

@@ -86,6 +86,13 @@ exports.datadog = function datadog (lambdaHandler) {
     const context = extractContext(args)
 
     checkTimeout(context)
-    return lambdaHandler.apply(this, args).then((res) => { clearTimeout(__lambdaTimeout); return res })
+    const result = lambdaHandler.apply(this, args)
+    if (result && typeof result.then === 'function') {
+      return result.then((res) => {
+        clearTimeout(__lambdaTimeout)
+        return res
+      })
+    }
+    return result
   }
 }

packages/dd-trace/test/lambda/fixtures/handler.js

@@ -20,6 +20,12 @@ const handler = async (_event, _context) => {
   return response
 }
 
+const callbackHandler = (_event, _context, callback) => {
+  const response = sampleResponse
+
+  callback('', response)
+}
+
 const timeoutHandler = async (...args) => {
   await _tracer.trace('self.sleepy', () => {
     return sleep(50)
@@ -63,5 +69,6 @@ module.exports = {
   handler,
   swappedArgsHandler,
   timeoutHandler,
-  errorHandler
+  errorHandler,
+  callbackHandler
 }

packages/dd-trace/test/lambda/index.spec.js

@@ -98,6 +98,43 @@ describe('lambda', () => {
       await checkTraces
     })
 
+    it('patches lambda function with callback correctly', async () => {
+      // Set the desired handler to patch
+      process.env.DD_LAMBDA_HANDLER = 'handler.callbackHandler'
+      // Load the agent and re-register hook for patching.
+      await loadAgent()
+
+      const _context = {
+        getRemainingTimeInMillis: () => 150
+      }
+      const _event = {}
+
+      // Mock `datadog-lambda` handler resolve and import.
+      const _handlerPath = path.resolve(__dirname, './fixtures/handler.js')
+      const app = require(_handlerPath)
+      datadog = require('./fixtures/datadog-lambda')
+      let result
+      const callback = (_error, res) => {
+        result = res
+      }
+      // Run the function.
+      datadog(app.callbackHandler)(_event, _context, callback)
+
+      expect(result).to.not.equal(undefined)
+      const body = JSON.parse(result.body)
+      expect(body.message).to.equal('hello!')
+
+      // Expect traces to be correct.
+      const checkTraces = agent.use((_traces) => {
+        const traces = _traces[0]
+        expect(traces).lengthOf(1)
+        traces.forEach((trace) => {
+          expect(trace.error).to.equal(0)
+        })
+      })
+      await checkTraces
+    })
+
     it('does wrap handler causing unhandled promise rejections', async () => {
       // Set the desired handler to patch
       process.env.DD_LAMBDA_HANDLER = 'handler.handler'