diff --git a/src/wrappers/lambda.js b/src/wrappers/lambda.js index 419d2483..4d37e36b 100644 --- a/src/wrappers/lambda.js +++ b/src/wrappers/lambda.js @@ -87,6 +87,7 @@ function baseLambdaWrapper( if (callbackCalled) { return Promise.resolve(); } + callbackCalled = true; if (error) { // not catching false here, but that seems OK eventInterface.setException(runner, error); } @@ -117,15 +118,17 @@ function baseLambdaWrapper( return tracer.sendTraceSync(); }; + let waitForOriginalCallbackPromise = Promise.resolve(); const wrappedCallback = (error, result) => { - handleUserExecutionDone(error, result).then(() => { - utils.debugLog('calling User\'s callback'); - return originalCallback(error, result); + waitForOriginalCallbackPromise = new Promise((resolve) => { + handleUserExecutionDone(error, result).then(() => { + utils.debugLog('calling User\'s callback'); + originalCallback(error, result); + resolve(); + }); }); - callbackCalled = true; }; - try { runner.setStartTime(utils.createTimestampFromTime(startTime)); let result = ( @@ -135,12 +138,23 @@ function baseLambdaWrapper( ); if (result instanceof Promise) { + let raisedError; + let returnValue; result = result .then((res) => { - handleUserExecutionDone(null, res).then(() => res); + returnValue = res; + return handleUserExecutionDone(null, res); }) .catch((err) => { - handleUserExecutionDone(err).then(() => err); + raisedError = err; + return handleUserExecutionDone(err); + }) + .then(() => waitForOriginalCallbackPromise) + .then(() => { + if (raisedError) { + throw raisedError; + } + return returnValue; }); } return result;