From da792fe3c629354cf9e8faeca48c17e73dffc6be Mon Sep 17 00:00:00 2001 From: Ran Nozik Date: Wed, 2 Mar 2022 21:56:42 +0200 Subject: [PATCH] feat: support baggage propagation in aws lambda custom context extraction (#843) Co-authored-by: Nathaniel Ruiz Nowell Co-authored-by: Rauno Viskus --- .../src/instrumentation.ts | 2 +- .../test/integrations/lambda-handler.test.ts | 31 +++++++++++++++++++ .../test/lambda-test/async.js | 5 +++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts index 57ff18be582..f8d0c522893 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts +++ b/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 diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts index 68cc1231eab..cabb6d0f37a 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts +++ b/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 => { diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/async.js b/plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/async.js index fe8bc5f9ff5..c6c2e529cf6 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/async.js +++ b/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}`; +}