chore: propagate context on AWS Lambda Instrumentation (#424)

Without this, spans created in the lambda function are not part of the
same trace as the handler span.

Co-authored-by: Valentin Marchaud <[email protected]>

plugins/node/opentelemetry-instrumentation-aws-lambda/src/aws-lambda.ts

@@ -24,7 +24,9 @@ import {
   safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
 import {
+  context as otelcontext,
   diag,
+  setSpan,
   Span,
   SpanKind,
   SpanStatusCode,
@@ -126,33 +128,35 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
         },
       });
 
-      // Lambda seems to pass a callback even if handler is of Promise form, so we wrap all the time before calling
-      // the handler and see if the result is a Promise or not. In such a case, the callback is usually ignored. If
-      // the handler happened to both call the callback and complete a returned Promise, whichever happens first will
-      // win and the latter will be ignored.
-      const wrappedCallback = plugin._wrapCallback(callback, span);
-      const maybePromise = safeExecuteInTheMiddle(
-        () => original.apply(this, [event, context, wrappedCallback]),
-        error => {
-          if (error != null) {
-            // Exception thrown synchronously before resolving callback / promise.
-            plugin._endSpan(span, error, () => {});
+      return otelcontext.with(setSpan(otelcontext.active(), span), () => {
+        // Lambda seems to pass a callback even if handler is of Promise form, so we wrap all the time before calling
+        // the handler and see if the result is a Promise or not. In such a case, the callback is usually ignored. If
+        // the handler happened to both call the callback and complete a returned Promise, whichever happens first will
+        // win and the latter will be ignored.
+        const wrappedCallback = plugin._wrapCallback(callback, span);
+        const maybePromise = safeExecuteInTheMiddle(
+          () => original.apply(this, [event, context, wrappedCallback]),
+          error => {
+            if (error != null) {
+              // Exception thrown synchronously before resolving callback / promise.
+              plugin._endSpan(span, error, () => {});
+            }
           }
+        ) as Promise<{}> | undefined;
+        if (typeof maybePromise?.then === 'function') {
+          return maybePromise.then(
+            value =>
+              new Promise(resolve =>
+                plugin._endSpan(span, undefined, () => resolve(value))
+              ),
+            (err: Error | string) =>
+              new Promise((resolve, reject) =>
+                plugin._endSpan(span, err, () => reject(err))
+              )
+          );
         }
-      ) as Promise<{}> | undefined;
-      if (typeof maybePromise?.then === 'function') {
-        return maybePromise.then(
-          value =>
-            new Promise(resolve =>
-              plugin._endSpan(span, undefined, () => resolve(value))
-            ),
-          (err: Error | string) =>
-            new Promise((resolve, reject) =>
-              plugin._endSpan(span, err, () => reject(err))
-            )
-        );
-      }
-      return maybePromise;
+        return maybePromise;
+      });
     };
   }
 

plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts

@@ -37,6 +37,7 @@ import {
 const memoryExporter = new InMemorySpanExporter();
 const provider = new NodeTracerProvider();
 provider.addSpanProcessor(new BatchSpanProcessor(memoryExporter));
+provider.register();
 
 const assertSpanSuccess = (span: ReadableSpan) => {
   assert.strictEqual(span.kind, SpanKind.SERVER);
@@ -145,6 +146,18 @@ describe('lambda handler', () => {
       const [span] = spans;
       assertSpanFailure(span);
     });
+
+    it('context should have parent trace', async () => {
+      initializeHandler('lambda-test/async.context');
+
+      const result = await lambdaRequire('lambda-test/async').context(
+        'arg',
+        ctx
+      );
+      const spans = memoryExporter.getFinishedSpans();
+      const [span] = spans;
+      assert.strictEqual(span.spanContext.traceId, result);
+    });
   });
 
   describe('sync success handler', () => {
@@ -238,6 +251,27 @@ describe('lambda handler', () => {
       assert.strictEqual(spans.length, 1);
       assertSpanFailure(span);
     });
+
+    it('context should have parent trace', async () => {
+      initializeHandler('lambda-test/sync.context');
+
+      const result = await new Promise((resolve, reject) => {
+        lambdaRequire('lambda-test/sync').context(
+          'arg',
+          ctx,
+          (err: Error, res: any) => {
+            if (err) {
+              reject(err);
+            } else {
+              resolve(res);
+            }
+          }
+        );
+      });
+      const spans = memoryExporter.getFinishedSpans();
+      const [span] = spans;
+      assert.strictEqual(span.spanContext.traceId, result);
+    });
   });
 
   it('should record string error in callback', async () => {

plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/async.js

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+const api = require('@opentelemetry/api');
 
 exports.handler = async function (event, context) {
   return 'ok';
@@ -25,3 +26,7 @@ exports.error = async function (event, context) {
 exports.stringerror = async function (event, context) {
   throw 'handler error';
 }
+
+exports.context = async function (event, context) {
+  return api.getSpanContext(api.context.active()).traceId;
+};

plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/sync.js

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+const api = require('@opentelemetry/api');
 
 exports.handler = function (event, context, callback) {
   callback(null, 'ok');
@@ -33,3 +34,7 @@ exports.stringerror = function (event, context, callback) {
 exports.callbackstringerror = function (event, context, callback) {
   callback('handler error');
 }
+
+exports.context = function (event, context, callback) {
+  callback(null, api.getSpanContext(api.context.active()).traceId);
+};