Fix awaiting for a flush (#20)

coralogix · Sep 13, 2024 · d2542ec · d2542ec
1 parent 64fe3ce
commit d2542ec
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
@@ -121,7 +121,7 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
       event: any,
       context: Context,
       callback: Callback
-    ) {
+    ): void {
 
       self._before_execution(event, context).then(
         (instrCtx) => {
@@ -140,25 +140,26 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
               } catch (err: any) {
                 // Catching synchronous failures
                 diag.debug('handler threw synchronously');
-                void self._after_execution(instrCtx, err, undefined);
-                context.callbackWaitsForEmptyEventLoop = false;
-                callback(err, undefined);
+                self._after_execution(instrCtx, err, undefined).then(() => {
+                  context.callbackWaitsForEmptyEventLoop = false;
+                  callback(err, undefined);
+                });
                 return;
               }
 
               if (typeof maybePromise?.then === 'function') {
                 diag.debug('handler returned a promise');
                 // Promise based async handler
                 maybePromise.then(
-                  (value: any) => {
+                  async (value: any) => {
                     diag.debug('handler promise completed');
-                    self._after_execution(instrCtx, undefined, value);
+                    await self._after_execution(instrCtx, undefined, value);
                     context.callbackWaitsForEmptyEventLoop = false;
                     callback(undefined, value);
                   },
-                  (err: Error | string) => {
+                  async (err: Error | string) => {
                     diag.debug('handler promise failed');
-                    self._after_execution(instrCtx, err, undefined);
+                    await self._after_execution(instrCtx, err, undefined);
                     context.callbackWaitsForEmptyEventLoop = false;
                     callback(err, undefined);
                   }
@@ -169,13 +170,14 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
             }
           );
         },
-        (err) => {
+        async (err) => {
           diag.error('_before_execution failed', err);
-          self._after_execution(undefined, err, undefined);
+          await self._after_execution(undefined, err, undefined);
           context.callbackWaitsForEmptyEventLoop = false;
           callback(err, undefined);
         }
       )
+
     }
   }
 
@@ -197,8 +199,6 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
 
     const invocationSpan = this._startInvocationSpan(event, context, invocationParentContext);
 
-    diag.info(`upstream: ${trace.getSpan(upstreamContext)?.spanContext().spanId} trigger: ${triggerSpan?.spanContext().spanId} invocationSpan: ${invocationSpan.spanContext().spanId}`)
-
     await this._sendEarlySpans(upstreamContext, triggerSpan, invocationParentContext, invocationSpan);
 
     return {triggerOrigin, triggerSpan, invocationSpan, invocationParentContext}