diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/aws-lambda.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/aws-lambda.ts index 2c25347e54c..c6664b7bd63 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/aws-lambda.ts +++ b/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; + }); }; } 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 5767be9b318..8837b1e399f 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 @@ -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 () => { 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 e352c8d6081..b5a75d9c1c3 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 @@ -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; +}; diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/sync.js b/plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/sync.js index 0f65dce0ee9..168714abd3f 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/lambda-test/sync.js +++ b/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); +};