diff --git a/packages/dd-trace/src/lambda/handler.js b/packages/dd-trace/src/lambda/handler.js index 6279c3d2d20..88457ce9d07 100644 --- a/packages/dd-trace/src/lambda/handler.js +++ b/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 } } diff --git a/packages/dd-trace/test/lambda/fixtures/handler.js b/packages/dd-trace/test/lambda/fixtures/handler.js index 1f5640f8f56..2541b0cd1cc 100644 --- a/packages/dd-trace/test/lambda/fixtures/handler.js +++ b/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 } diff --git a/packages/dd-trace/test/lambda/index.spec.js b/packages/dd-trace/test/lambda/index.spec.js index f265baa0a46..8d22949dd76 100644 --- a/packages/dd-trace/test/lambda/index.spec.js +++ b/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'