Merge pull request #651 from majanjua-amzn/master

[Lambda] Replace Facade with No-Op if trace header is missing data

Author: wangzlei

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;

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 };

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)) {

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));
 

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) {

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 {

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;
       });
     });
   });

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);
+    });
+
+  });
 });

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 () => {

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);