fix(instrumentation/aws-lambda): Ensure callback is only called once (o…

…pen-telemetry#1384) * fix(instrumentation/aws-lambda): Ensure callback is only called once regardless how many providers are flushed Signed-off-by: Anthony J Mirabella <[email protected]> * Update plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts Co-authored-by: Daniel Dyla <[email protected]> * fix lint issue Signed-off-by: Anthony J Mirabella <[email protected]> --------- Signed-off-by: Anthony J Mirabella <[email protected]> Co-authored-by: Daniel Dyla <[email protected]>
dynatrace-oss-contrib · Feb 8, 2023 · d822f75 · d822f75
1 parent 096129c
commit d822f75
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 10 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
@@ -303,28 +303,23 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
 
     span.end();
 
+    const flushers = [];
     if (this._traceForceFlusher) {
-      this._traceForceFlusher().then(
-        () => callback(),
-        () => callback()
-      );
+      flushers.push(this._traceForceFlusher());
     } else {
       diag.error(
         'Spans may not be exported for the lambda function because we are not force flushing before callback.'
       );
-      callback();
     }
     if (this._metricForceFlusher) {
-      this._metricForceFlusher().then(
-        () => callback(),
-        () => callback()
-      );
+      flushers.push(this._metricForceFlusher());
     } else {
       diag.error(
         'Metrics may not be exported for the lambda function because we are not force flushing before callback.'
       );
-      callback();
     }
+
+    Promise.all(flushers).then(callback, callback);
   }
 
   private _applyResponseHook(

diff --git a/...telemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.force-flush.test.ts b/...telemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.force-flush.test.ts
@@ -181,4 +181,59 @@ describe('force flush', () => {
 
     assert.strictEqual(forceFlushed, true);
   });
+
+  it('should callback once after force flush providers', async () => {
+    const nodeTracerProvider = new NodeTracerProvider();
+    nodeTracerProvider.addSpanProcessor(
+      new BatchSpanProcessor(traceMemoryExporter)
+    );
+    nodeTracerProvider.register();
+    const tracerProvider = new ProxyTracerProvider();
+    tracerProvider.setDelegate(nodeTracerProvider);
+    let tracerForceFlushed = false;
+    const tracerForceFlush = () =>
+      new Promise<void>(resolve => {
+        tracerForceFlushed = true;
+        resolve();
+      });
+    nodeTracerProvider.forceFlush = tracerForceFlush;
+
+    const meterProvider = new MeterProvider();
+    meterProvider.addMetricReader(
+      new PeriodicExportingMetricReader({ exporter: metricMemoryExporter })
+    );
+    let meterForceFlushed = false;
+    const meterForceFlush = () =>
+      new Promise<void>(resolve => {
+        meterForceFlushed = true;
+        resolve();
+      });
+    meterProvider.forceFlush = meterForceFlush;
+
+    process.env._HANDLER = 'lambda-test/sync.handler';
+
+    instrumentation = new AwsLambdaInstrumentation();
+    instrumentation.setTracerProvider(tracerProvider);
+    instrumentation.setMeterProvider(meterProvider);
+
+    let callbackCount = 0;
+    await new Promise((resolve, reject) => {
+      lambdaRequire('lambda-test/sync').handler(
+        'arg',
+        ctx,
+        (err: Error, res: any) => {
+          callbackCount++;
+          if (err) {
+            reject(err);
+          } else {
+            resolve(res);
+          }
+        }
+      );
+    });
+
+    assert.strictEqual(tracerForceFlushed, true);
+    assert.strictEqual(meterForceFlushed, true);
+    assert.strictEqual(callbackCount, 1);
+  });
 });