Fixed trace headers and added unit tests

aws · Nov 4, 2022 · 5342ac9 · 5342ac9
1 parent 9b2660a
commit 5342ac9
Show file tree

Hide file tree

Showing 11 changed files with 567 additions and 46 deletions.
diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts
@@ -110,7 +110,14 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
     return next(args);
   }
 
-  const subsegment: Subsegment = segment.addNewSubsegment(service);
+  let subsegment: Subsegment;
+
+  if(segment.notTraced == false || segment.subsegments[segment.subsegments.length - 1].isSampled){
+    subsegment = segment.addNewSubsegment(service);
+  } else {
+    subsegment = segment.addNewSubsegmentWithoutSampling(service);
+  }
+
   subsegment.addAttribute('namespace', 'aws');
   const parent = (segment instanceof Subsegment ? segment.segment : segment);
 

diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js
@@ -76,7 +76,14 @@ function captureAWSRequest(req) {
   var throttledError = this.throttledError || throttledErrorDefault;
 
   var stack = (new Error()).stack;
-  var subsegment = parent.addNewSubsegment(this.serviceIdentifier);
+
+  let subsegment;
+  if(parent.notTraced == false || parent.subsegments[parent.subsegments.length - 1].isSampled){
+    subsegment = parent.addNewSubsegment(this.serviceIdentifier);
+  } else {
+    subsegment = parent.addNewSubsegmentWithoutSampling(this.serviceIdentifier);
+  }
+
   var traceId = parent.segment ? parent.segment.trace_id : parent.trace_id;
 
   var buildListener = function(req) {

diff --git a/packages/core/lib/patchers/http_p.js b/packages/core/lib/patchers/http_p.js
@@ -116,7 +116,13 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) {
       return baseFunc(...args);
     }
 
-    const subsegment = parent.addNewSubsegment(hostname);
+    let subsegment;
+    if(parent.notTraced == false || parent.subsegments[parent.subsegments.length - 1].isSampled){
+      subsegment = parent.addNewSubsegment(hostname);
+    } else {
+      subsegment = parent.addNewSubsegmentWithoutSampling(hostname);
+    }
+
     const root = parent.segment ? parent.segment : parent;
     subsegment.namespace = 'remote';
 

diff --git a/packages/core/lib/segments/attributes/subsegment.js b/packages/core/lib/segments/attributes/subsegment.js
@@ -44,15 +44,14 @@ Subsegment.prototype.addNewSubsegment = function addNewSubsegment(name) {
 };
 
 Subsegment.prototype.addSubsegmentWithoutSampling = function addSubsegmentWithoutSampling(subsegment){
-  subsegment.isSampled = false;
   this.addSubsegment(subsegment);
+  subsegment.isSampled = false;
 };
 
 Subsegment.prototype.addNewSubsegmentWithoutSampling = function addNewSubsegmentWithoutSampling(name){
   const subsegment = new Subsegment(name);
-  subsegment.isSampled = false;
   this.addSubsegment(subsegment);
-
+  subsegment.isSampled = false;
   return subsegment; 
 };
 
@@ -74,19 +73,13 @@ Subsegment.prototype.addSubsegment = function(subsegment) {
   subsegment.segment = this.segment;
   subsegment.parent = this;
 
-  if(subsegment.isSampled){ 
-    // if subsegment has not been set to false, use the sampling decision of the direct parent
-    subsegment.isSampled = subsegment.parent.isSampled;
-  }
+  subsegment.isSampled = subsegment.parent.isSampled;
 
-  if (subsegment.end_time === undefined && subsegment.isSampled) {
+  if (subsegment.end_time === undefined) {
     this.incrementCounter(subsegment.counter);
   }
-  // don't push to subsegment array if subsegment is not sampled
-  if(subsegment.isSampled){ 
-    this.subsegments.push(subsegment);
-  }
-
+  this.subsegments.push(subsegment);
+
 };
 
 /**
@@ -362,7 +355,7 @@ Subsegment.prototype.flush = function flush() {
   }
 
   if (this.segment.trace_id) {
-    if (this.segment.notTraced !== true) {
+    if (this.segment.notTraced !== true && this.isSampled) {
       SegmentEmitter.send(this);
     } else {
       logger.getLogger().debug('Ignoring flush on subsegment ' + this.id + '. Associated segment is marked as not sampled.');

diff --git a/packages/core/lib/segments/segment.d.ts b/packages/core/lib/segments/segment.d.ts
@@ -37,7 +37,7 @@ declare class Segment {
 
   addSubsegmentWithoutSampling(subsegment: Subsegment): void; 
 
-  addNewSubsegmentWithoutSampling(name: String): Subsegment
+  addNewSubsegmentWithoutSampling(name: string): Subsegment
 
   removeSubsegment(subsegment: Subsegment): void;
 

diff --git a/packages/core/lib/segments/segment.js b/packages/core/lib/segments/segment.js
@@ -218,16 +218,15 @@ Segment.prototype.addNewSubsegment = function addNewSubsegment(name) {
 };
 
 Segment.prototype.addSubsegmentWithoutSampling = function addSubsegmentWithoutSampling(subsegment){
-  subsegment.isSampled = false;
   this.addSubsegment(subsegment);
+  subsegment.isSampled = false;
 
 };
 
 Segment.prototype.addNewSubsegmentWithoutSampling = function addNewSubsegmentWithoutSampling(name){
   const subsegment = new Subsegment(name);
-  subsegment.isSampled = false;
   this.addSubsegment(subsegment);
-
+    subsegment.isSampled = false;
   return subsegment; 
 };
 
@@ -248,11 +247,10 @@ Segment.prototype.addSubsegment = function addSubsegment(subsegment) {
   subsegment.segment = this;
   subsegment.parent = this;
 
-  if(subsegment.isSampled){
-      this.subsegments.push(subsegment);
-  }
+  subsegment.isSampled = !subsegment.parent.notTraced;
+  this.subsegments.push(subsegment);
 
-  if (!subsegment.end_time && subsegment.isSampled) {
+  if (!subsegment.end_time) {
     this.incrementCounter(subsegment.counter);
   }
 };

diff --git a/packages/core/test/unit/patchers/aws3_p.test.js b/packages/core/test/unit/patchers/aws3_p.test.js
@@ -247,4 +247,87 @@ describe('AWS v3 patcher', function() {
       });
     });
   });
+
+
+  describe('#captureAWSRequest-Unsampled', function() {
+    var awsClient, awsRequest, sandbox, segment, stubResolve, addNewSubsegmentStub, sub;
+
+    before(function() {
+      awsClient = {
+        send: async (req) => {
+          const context = {
+            clientName: 'S3Client',
+            commandName: 'ListBucketsCommand',
+          };
+          const handler = awsClient.middlewareStack.resolve((args) => {
+            const error = req.response.error;
+            if (error) {
+              const err = new Error(error.message);
+              err.name = error.code;
+              err.$metadata = req.response.$metadata;
+              throw err;
+            }
+            return args;
+          }, context);
+          await handler(req);
+          return req.response;
+        },
+        config: {
+          region: async () => 'us-east-1',
+        },
+        middlewareStack: constructStack(),
+      };
+    });
+
+    beforeEach(function() {
+      sandbox = sinon.createSandbox();
+
+      awsRequest = new (class ListBucketsCommand {
+        constructor() {
+          this.request = {
+            method: 'GET',
+            url: '/',
+            connection: {
+              remoteAddress: 'localhost'
+            },
+            headers: {}
+          };
+          this.response = {};
+          this.output = {
+            $metadata: {
+              requestId: '123',
+              extendedRequestId: '456',
+            }
+          };
+        }
+      })();
+
+      segment = new Segment('testSegment', traceId);
+      sub = segment.addNewSubsegmentWithoutSampling('subseg');
+      stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(segment);
+      addNewSubsegmentStub = sandbox.stub(segment, 'addNewSubsegmentWithoutSampling').returns(sub);
+    });
+
+    afterEach(function() {
+      sandbox.restore();
+    });
+
+    describe('#automaticMode', () => {
+      beforeEach(() => {
+        sandbox.stub(contextUtils, 'isAutomaticMode').returns(true);
+      });
+
+      before(() => {
+        awsClient = awsPatcher.captureAWSClient(awsClient);
+      });
+
+
+      it('should inject the tracing headers', async function() {
+        await awsClient.send(awsRequest);
+
+        const expected = new RegExp('^Root=' + traceId + ';Parent=' + sub.id + ';Sampled=0$');
+        assert.match(awsRequest.request.headers['X-Amzn-Trace-Id'], expected);
+      });
+    });
+  });
 });
diff --git a/packages/core/test/unit/patchers/aws_p.test.js b/packages/core/test/unit/patchers/aws_p.test.js
@@ -259,4 +259,90 @@ describe('AWS patcher', function() {
       }, 50);
     });
   });
-});
+
+
+  describe('#captureAWSRequest-Unsampled', function() {
+    var awsClient, awsRequest, MyEmitter, sandbox, segment, stubResolve, stubResolveManual, sub;
+
+    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);
+      sub = segment.addNewSubsegmentWithoutSampling("subseg");
+
+      stubResolveManual = sandbox.stub(contextUtils, 'resolveManualSegmentParams');
+      stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(segment);
+      sandbox.stub(segment, 'addNewSubsegmentWithoutSampling').returns(sub);
+    });
+
+    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 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 + ';Parent=' + sub.id + ';Sampled=0$');
+        assert.match(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], expected);
+        done();
+      }, 50);
+    });
+
+  });
+});