Account for lambda passthrough trace header and add unit tests

Author: majanjua-amzn

packages/core/lib/env/aws_lambda.js

@@ -175,13 +175,14 @@ var noOpSegment = function noOpSegment() {
     }
   };
 
-  // Test for valid trace data during SDK startup. It's likely we're still in the cold-start portion of the
-  // code at this point and a valid trace header has not been set
-  if (LambdaUtils.validTraceData(xAmznTraceId)) {
-    if (LambdaUtils.populateTraceData(segment, xAmznTraceId)) {
-      xAmznTraceIdPrev = xAmznTraceId;
-    }
+  // 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/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/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);