diff --git a/packages/core/lib/context_utils.js b/packages/core/lib/context_utils.js index fe20f440..d6c47594 100644 --- a/packages/core/lib/context_utils.js +++ b/packages/core/lib/context_utils.js @@ -99,8 +99,10 @@ var contextUtils = { if (!segment) { contextUtils.contextMissingStrategy.contextMissing('Failed to get the current sub/segment from the context.'); - } else if (segment instanceof Segment && process.env.LAMBDA_TASK_ROOT && segment.facade == true) { - segment.resolveLambdaTraceData(); + } else if (segment instanceof Segment && process.env.LAMBDA_TASK_ROOT) { + if (segment.facade == true || segment.noOp == true) { + segment.resolveLambdaTraceData(); + } } return segment; diff --git a/packages/core/lib/env/aws_lambda.js b/packages/core/lib/env/aws_lambda.js index 18ed0f12..8473d7bc 100644 --- a/packages/core/lib/env/aws_lambda.js +++ b/packages/core/lib/env/aws_lambda.js @@ -36,7 +36,12 @@ module.exports.init = function init() { var namespace = contextUtils.getNamespace(); namespace.enter(namespace.createContext()); - contextUtils.setSegment(facadeSegment()); + + if (LambdaUtils.validTraceData(process.env._X_AMZN_TRACE_ID)) { + contextUtils.setSegment(facadeSegment()); + } else { + contextUtils.setSegment(noOpSegment()); + } }; var facadeSegment = function facadeSegment() { @@ -109,3 +114,75 @@ var facadeSegment = function facadeSegment() { return segment; }; + +var noOpSegment = function noOpSegment() { + var segment = new Segment('no-op'); + var whitelistFcn = []; + var silentFcn = ['addNewSubsegment', 'addSubsegment', 'removeSubsegment', 'toString', 'addSubsegmentWithoutSampling', 'addNewSubsegmentWithoutSampling', 'incrementCounter', 'decrementCounter', 'isClosed', 'close', 'format', 'flush']; + var xAmznTraceId = process.env._X_AMZN_TRACE_ID; + + for (var key in segment) { + if (typeof segment[key] === 'function' && whitelistFcn.indexOf(key) === -1) { + if (silentFcn.indexOf(key) === -1) { + segment[key] = (function() { + var func = key; + return function noOp() { + logger.getLogger().warn('Function "' + func + '" cannot be called on an AWS Lambda segment. Please use a subsegment to record data.'); + return; + }; + })(); + } else { + segment[key] = function noOp() { + return; + }; + } + } + } + + segment.trace_id = TraceID.Invalid().toString(); + segment.isClosed = function() { + return true; + }; + segment.in_progress = false; + segment.counter = 1; + segment.notTraced = true; + segment.noOp = true; + + segment.reset = function reset() { + this.trace_id = TraceID.Invalid().toString(); + this.id = '00000000'; + delete this.subsegments; + this.notTraced = true; + }; + + segment.resolveLambdaTraceData = function resolveLambdaTraceData() { + var xAmznLambda = process.env._X_AMZN_TRACE_ID; + + if (xAmznLambda) { + + // This check resets the trace data whenever a new trace header is read to not leak data between invocations + if (xAmznLambda != xAmznTraceIdPrev) { + this.reset(); + + if (LambdaUtils.populateTraceData(segment, xAmznLambda)) { + xAmznTraceIdPrev = xAmznLambda; + } + } + } else { + this.reset(); + contextUtils.contextMissingStrategy.contextMissing('Missing AWS Lambda trace data for X-Ray. ' + + 'Ensure Active Tracing is enabled and no subsegments are created outside the function handler.'); + } + }; + + // Since we're in a no-op segment, do not check if the trace data is valid; simply propagate the information + if (LambdaUtils.populateTraceData(segment, xAmznTraceId)) { + xAmznTraceIdPrev = xAmznTraceId; + } + + return segment; +}; + +// For testing +export const exportedFacadeSegment = { facadeSegment }; +export const exportedNoOpSegment = { noOpSegment }; diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts index 782317e4..bf38ff45 100644 --- a/packages/core/lib/patchers/aws3_p.ts +++ b/packages/core/lib/patchers/aws3_p.ts @@ -125,14 +125,10 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment const parent = (segment instanceof Subsegment ? segment.segment : segment); const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData; - let traceHeader = stringify( - { - Root: parent.trace_id, - Parent: subsegment.id, - Sampled: subsegment.notTraced ? '0' : '1', - }, - ';', - ); + let traceHeader = 'Root=' + parent.trace_id; + if (!(parent && parent.noOp)) { + traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + } if (data != null) { for (const [key, value] of Object.entries(data)) { diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js index 776d33ef..746b9547 100644 --- a/packages/core/lib/patchers/aws_p.js +++ b/packages/core/lib/patchers/aws_p.js @@ -88,8 +88,11 @@ function captureAWSRequest(req) { const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData; var buildListener = function(req) { - let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id + - ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + let traceHeader = 'Root=' + traceId; + // Only append parent and sample if not in Lambda PassThrough mode + if (!(parent && parent.noOp)) { + traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + } if (data != null) { for (const [key, value] of Object.entries(data)) { traceHeader += ';' + key +'=' + value; @@ -99,6 +102,9 @@ function captureAWSRequest(req) { }; var completeListener = function(res) { + if (subsegment == null) { + return; + } subsegment.addAttribute('namespace', 'aws'); subsegment.addAttribute('aws', new Aws(res, subsegment.name)); diff --git a/packages/core/lib/patchers/http_p.js b/packages/core/lib/patchers/http_p.js index fcf2f16f..a4519e5f 100644 --- a/packages/core/lib/patchers/http_p.js +++ b/packages/core/lib/patchers/http_p.js @@ -123,15 +123,18 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) { subsegment = parent.addNewSubsegment(hostname); } - const root = parent.segment ? parent.segment : parent; subsegment.namespace = 'remote'; if (!options.headers) { options.headers = {}; } - options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id + - ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + let traceHeader = 'Root=' + (parent.segment ? parent.segment : parent).trace_id; + if (!(parent && parent.noOp)) { + traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + } + + options.headers['X-Amzn-Trace-Id'] = traceHeader; const errorCapturer = function errorCapturer(e) { if (subsegmentCallback) { diff --git a/packages/core/lib/utils.js b/packages/core/lib/utils.js index 93efcd3b..d907f5fb 100644 --- a/packages/core/lib/utils.js +++ b/packages/core/lib/utils.js @@ -168,6 +168,8 @@ var utils = { if (!traceData) { traceData = {}; logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format'); + } else if (segment.noOp == true && traceData.root) { + valid = true; } else if (!traceData.root || !traceData.parent || !traceData.sampled) { logger.getLogger().error('_X_AMZN_TRACE_ID is missing required information'); } else { diff --git a/packages/core/test/unit/env/aws_lambda.test.js b/packages/core/test/unit/env/aws_lambda.test.js index eb7e6fb1..dc66764d 100644 --- a/packages/core/test/unit/env/aws_lambda.test.js +++ b/packages/core/test/unit/env/aws_lambda.test.js @@ -81,7 +81,7 @@ describe('AWSLambda', function() { assert.equal(facade.trace_id, TraceID.Invalid().toString()); }); - describe('the facade segment', function() { + describe('the facade/no-op segment', function() { afterEach(function() { populateStub.returns(true); delete process.env._X_AMZN_TRACE_ID; @@ -95,17 +95,25 @@ describe('AWSLambda', function() { validateStub.should.have.been.calledWith(process.env._X_AMZN_TRACE_ID); }); - it('should call populateTraceData if validTraceData returns true', function() { + it('should call populateTraceData on Facade if validTraceData returns true', function() { Lambda.init(); + var segment = setSegmentStub.args[0][0]; + assert.equal(segment.name, 'facade'); + assert.isTrue(segment.facade); + populateStub.should.have.been.calledOnce; }); - it('should not call populateTraceData if validTraceData returns false', function() { + it('should call populateTraceData on No-Op if validTraceData returns false', function() { validateStub.returns(false); Lambda.init(); - populateStub.should.have.not.been.called; + var segment = setSegmentStub.args[0][0]; + assert.equal(segment.name, 'no-op'); + assert.isTrue(segment.noOp); + + populateStub.should.have.been.calledOnce; }); }); }); diff --git a/packages/core/test/unit/patchers/aws_p.test.js b/packages/core/test/unit/patchers/aws_p.test.js index fdd48d93..3603fb30 100644 --- a/packages/core/test/unit/patchers/aws_p.test.js +++ b/packages/core/test/unit/patchers/aws_p.test.js @@ -13,6 +13,10 @@ var Utils = require('../../../lib/utils'); var logger = require('../../../lib/logger').getLogger(); +import { exportedFacadeSegment, exportedNoOpSegment } from '../../../lib/env/aws_lambda'; +const { facadeSegment } = exportedFacadeSegment; +const { noOpSegment } = exportedNoOpSegment; + chai.should(); chai.use(sinonChai); @@ -348,4 +352,91 @@ describe('AWS patcher', function() { }); }); + + + describe('#captureAWSRequest-Lambda-PassThrough', function() { + var awsClient, awsRequest, MyEmitter, sandbox, segment, stubResolve, tempHeader; + + before(function() { + MyEmitter = function() { + EventEmitter.call(this); + }; + + awsClient = { + customizeRequests: function customizeRequests(captureAWSRequest) { + this.call = captureAWSRequest; + }, + throttledError: function throttledError() {} + }; + awsClient = awsPatcher.captureAWSClient(awsClient); + + util.inherits(MyEmitter, EventEmitter); + }); + + beforeEach(function() { + sandbox = sinon.createSandbox(); + + awsRequest = { + httpRequest: { + method: 'GET', + url: '/', + connection: { + remoteAddress: 'localhost' + }, + headers: {} + }, + response: {} + }; + + awsRequest.on = function(event, fcn) { + if (event === 'complete') { + this.emitter.on(event, fcn.bind(this, this.response)); + } else { + this.emitter.on(event, fcn.bind(this, this)); + } + return this; + }; + + awsRequest.emitter = new MyEmitter(); + + tempHeader = process.env._X_AMZN_TRACE_ID; + process.env._X_AMZN_TRACE_ID = 'Root=' + traceId + ';Foo=bar'; + + segment = noOpSegment(); + + stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(segment); + }); + + afterEach(function() { + process.env._X_AMZN_TRACE_ID = tempHeader; + sandbox.restore(); + }); + + it('should log an info statement and exit if parent is not found on the context or on the call params', function(done) { + stubResolve.returns(); + var logStub = sandbox.stub(logger, 'info'); + + awsClient.call(awsRequest); + + setTimeout(function() { + logStub.should.have.been.calledOnce; + done(); + }, 50); + }); + + it('should inject the tracing headers', function(done) { + sandbox.stub(contextUtils, 'isAutomaticMode').returns(true); + + awsClient.call(awsRequest); + + awsRequest.emitter.emit('build'); + + setTimeout(function() { + var expected = new RegExp('^Root=' + traceId + ';Foo=bar$'); + assert.match(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], expected); + done(); + }, 50); + }); + + }); }); diff --git a/sdk_contrib/fetch/lib/fetch_p.js b/sdk_contrib/fetch/lib/fetch_p.js index 4666a771..d4ba1769 100644 --- a/sdk_contrib/fetch/lib/fetch_p.js +++ b/sdk_contrib/fetch/lib/fetch_p.js @@ -104,10 +104,12 @@ const enableCapture = function enableCapture(baseFetchFunction, requestClass, do subsegment.namespace = 'remote'; - request.headers.set('X-Amzn-Trace-Id', - 'Root=' + (parent.segment ? parent.segment : parent).trace_id + - ';Parent=' + subsegment.id + - ';Sampled=' + (subsegment.notTraced ? '0' : '1')); + let traceHeader = 'Root=' + (parent.segment ? parent.segment : parent).trace_id; + if (!(parent && parent.noOp)) { + traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + } + + request.headers.set('X-Amzn-Trace-Id', traceHeader); // Set up fetch call and capture any thrown errors const capturedFetch = async () => { diff --git a/sdk_contrib/fetch/test/unit/fetch_p.test.js b/sdk_contrib/fetch/test/unit/fetch_p.test.js index 4dde8d52..9dbc5128 100644 --- a/sdk_contrib/fetch/test/unit/fetch_p.test.js +++ b/sdk_contrib/fetch/test/unit/fetch_p.test.js @@ -273,6 +273,19 @@ describe('Unit tests', function () { 'Root=12345;Parent=999;Sampled=1'); }); + it('adds X-Amzn-Trace-Id header with only root if noOp', async function () { + const activeFetch = captureFetch(true); + stubParentSegment.noOp = true; + stubParentSegment.trace_id = '12345'; + stubSubsegment.notTraced = false; + stubSubsegment.id = '999'; + const request = new FetchRequest('https://www.foo.com/test'); + const requestHeadersSet = sandbox.stub(request.headers, 'set'); + await activeFetch(request); + requestHeadersSet.should.have.been.calledOnceWith('X-Amzn-Trace-Id', + 'Root=12345'); + }); + it('calls subsegmentCallback on successful response', async function () { const spyCallback = sandbox.spy(); const activeFetch = captureFetch(true, spyCallback);