fix(lambda_wrapper): fixing async callback handling (#18)

epsagon · Dec 6, 2018 · 2452110 · 2452110
1 parent f024be3
commit 2452110
Showing 1 changed file with 21 additions and 7 deletions.
diff --git 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;