feat: support baggage propagation in aws lambda custom context extrac…

…tion (open-telemetry#843)

Co-authored-by: Nathaniel Ruiz Nowell <enruizno@uwaterloo.ca>
Co-authored-by: Rauno Viskus <Rauno56@users.noreply.github.com>

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

@@ -190,7 +190,7 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
         );
       }
 
-      return otelContext.with(trace.setSpan(otelContext.active(), span), () => {
+      return otelContext.with(trace.setSpan(parent, 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

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

@@ -626,6 +626,37 @@ describe('lambda handler', () => {
       assert.strictEqual(span.parentSpanId, sampledGenericSpanContext.spanId);
     });
 
+    it('prefers to extract baggage over sampled lambda context if "eventContextExtractor" is defined', async () => {
+      process.env[traceContextEnvironmentKey] = sampledAwsHeader;
+      const customExtractor = (event: any): OtelContext => {
+        return propagation.extract(
+          context.active(),
+          event.customContextCarrier
+        );
+      };
+
+      initializeHandler('lambda-test/async.handler_return_baggage', {
+        disableAwsContextPropagation: true,
+        eventContextExtractor: customExtractor,
+      });
+
+      const baggage = 'abcd=1234';
+      const customRemoteEvent = {
+        customContextCarrier: {
+          traceparent: sampledGenericSpan,
+          baggage,
+        },
+      };
+
+      const lambdaTestAsync = lambdaRequire('lambda-test/async');
+      const actual = await lambdaTestAsync.handler_return_baggage(
+        customRemoteEvent,
+        ctx
+      );
+
+      assert.strictEqual(actual, baggage);
+    });
+
     it('creates trace from ROOT_CONTEXT when "disableAwsContextPropagation" is true, eventContextExtractor is provided, and no custom context is found', async () => {
       process.env[traceContextEnvironmentKey] = sampledAwsHeader;
       const customExtractor = (event: any): OtelContext => {

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

@@ -30,3 +30,8 @@ exports.stringerror = async function (event, context) {
 exports.context = async function (event, context) {
   return api.trace.getSpan(api.context.active()).spanContext().traceId;
 };
+
+exports.handler_return_baggage = async function (event, context) {
+  const [baggageEntryKey, baggageEntryValue] =  api.propagation.getBaggage(api.context.active()).getAllEntries()[0];
+  return `${baggageEntryKey}=${baggageEntryValue.value}`;
+}