Skip to content

Commit

Permalink
chore: propagate context on AWS Lambda Instrumentation (open-telemetr…
Browse files Browse the repository at this point in the history
…y#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]>
  • Loading branch information
gregoryfranklin and vmarchaud authored Apr 22, 2021
1 parent 86708bc commit 21d1701
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 25 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,9 @@ import {
safeExecuteInTheMiddle,
} from '@opentelemetry/instrumentation';
import {
context as otelcontext,
diag,
setSpan,
Span,
SpanKind,
SpanStatusCode,
Expand Down Expand Up @@ -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;
});
};
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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', () => {
Expand Down Expand Up @@ -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 () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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;
};
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand All @@ -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);
};

0 comments on commit 21d1701

Please sign in to comment.