Allow lambda passthrough trace header

majanjua-amzn · Jun 7, 2024 · dd12415 · dd12415
1 parent 1502470
commit dd12415
Show file tree

Hide file tree

Showing 7 changed files with 104 additions and 3 deletions.
diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts
@@ -140,7 +140,9 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
     }
   }
 
-  args.request.headers['X-Amzn-Trace-Id'] = traceHeader;
+  if (!segment.noOp) {
+    args.request.headers['X-Amzn-Trace-Id'] = traceHeader;
+  }
 
   let res;
   try {

diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js
@@ -88,6 +88,9 @@ function captureAWSRequest(req) {
   const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
 
   var buildListener = function(req) {
+    if (parent.noOp) {
+      return;
+    }
     let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id +
       ';Sampled=' + (subsegment.notTraced ? '0' : '1');
     if (data != null) {

diff --git a/packages/core/lib/patchers/http_p.js b/packages/core/lib/patchers/http_p.js
@@ -130,8 +130,10 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) {
       options.headers = {};
     }
 
-    options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
-      ';Sampled=' + (subsegment.notTraced ? '0' : '1');
+    if (!parent.noOp) {
+      options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
+        ';Sampled=' + (subsegment.notTraced ? '0' : '1');
+    }
 
     const errorCapturer = function errorCapturer(e) {
       if (subsegmentCallback) {

diff --git a/packages/core/lib/segments/attributes/subsegment.js b/packages/core/lib/segments/attributes/subsegment.js
@@ -74,6 +74,7 @@ Subsegment.prototype.addSubsegment = function(subsegment) {
   subsegment.parent = this;
 
   subsegment.notTraced = subsegment.parent.notTraced;
+  subsegment.noOp = subsegment.parent.noOp;
 
   if (subsegment.end_time === undefined) {
     this.incrementCounter(subsegment.counter);

diff --git a/packages/core/lib/segments/segment.js b/packages/core/lib/segments/segment.js
@@ -248,6 +248,7 @@ Segment.prototype.addSubsegment = function addSubsegment(subsegment) {
   subsegment.parent = this;
 
   subsegment.notTraced = subsegment.parent.notTraced;
+  subsegment.noOp = subsegment.parent.noOp;
   this.subsegments.push(subsegment);
 
   if (!subsegment.end_time) {

diff --git a/packages/core/lib/utils.js b/packages/core/lib/utils.js
@@ -168,6 +168,10 @@ var utils = {
       if (!traceData) {
         traceData = {};
         logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format');
+      } else if (traceData.root && !traceData.parent && !traceData.sampled) {
+        // Lambda PassThrough only has root, treat as valid in this case and mark the segment
+        segment.noOp = true;
+        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/patchers/aws_p.test.js b/packages/core/test/unit/patchers/aws_p.test.js
@@ -348,4 +348,92 @@ describe('AWS patcher', function() {
     });
 
   });
+
+
+  describe('#captureAWSRequest-PassThrough-Header', function() {
+    var awsClient, awsRequest, MyEmitter, sandbox, segment, stubResolve, addNewSubsegmentStub, sub, addNewServiceSubsegmentStub, service;
+
+    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();
+
+      segment = new Segment('testSegment', traceId);
+      segment.noOp = true; // enforce passthrough behaviour
+      segment.additionalTraceData = {'Foo': 'bar'};
+      sub = segment.addNewSubsegmentWithoutSampling('subseg');
+      service = sub.addNewSubsegmentWithoutSampling('service');
+
+      stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(sub);
+      addNewSubsegmentStub = sandbox.stub(segment, 'addNewSubsegmentWithoutSampling').returns(sub);
+      addNewServiceSubsegmentStub = sandbox.stub(sub, 'addNewSubsegmentWithoutSampling').returns(service);
+    });
+
+    afterEach(function() {
+      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 not inject the tracing headers if passthrough mode', function(done) {
+      sandbox.stub(contextUtils, 'isAutomaticMode').returns(true);
+
+      awsClient.call(awsRequest);
+
+      awsRequest.emitter.emit('build');
+
+      setTimeout(function() {
+        assert.equal(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], undefined);
+        done();
+      }, 50);
+    });
+
+  });
 });