Oversampling Mitigation

packages/core/lib/env/aws_lambda.js

@@ -41,7 +41,7 @@ module.exports.init = function init() {
 
 var facadeSegment = function facadeSegment() {
   var segment = new Segment('facade');
-  var whitelistFcn = ['addNewSubsegment', 'addSubsegment', 'removeSubsegment', 'toString'];
+  var whitelistFcn = ['addNewSubsegment', 'addSubsegment', 'removeSubsegment', 'toString', 'addSubsegmentWithoutSampling', 'addNewSubsegmentWithoutSampling'];
   var silentFcn = ['incrementCounter', 'decrementCounter', 'isClosed', 'close', 'format', 'flush'];
   var xAmznTraceId = process.env._X_AMZN_TRACE_ID;
 

packages/core/lib/env/sqs_message_helper.js

@@ -0,0 +1,14 @@
+class SqsMessageHelper {
+  constructor() {
+  }
+
+  isSampled(message) { 
+    const {attributes} = message; // extract attributes from message 
+    if (!('AWSTraceHeader' in attributes)) {
+      return false; 
+    }
+    return attributes['AWSTraceHeader'].includes('Sampled=1');
+  }
+}
+
+export default SqsMessageHelper;

packages/core/lib/patchers/aws_p.js

@@ -81,7 +81,8 @@ function captureAWSRequest(req) {
 
   var buildListener = function(req) {
     req.httpRequest.headers['X-Amzn-Trace-Id'] = 'Root=' + traceId + ';Parent=' + subsegment.id +
-      ';Sampled=' + (subsegment.segment.notTraced ? '0' : '1');
+      ';Sampled=' + (subsegment.isSampled ? '1' : '0');
+
   };
 
   var completeListener = function(res) {

packages/core/lib/patchers/http_p.js

@@ -123,9 +123,9 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) {
     if (!options.headers) {
       options.headers = {};
     }
-
+    
     options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
-      ';Sampled=' + (!root.notTraced ? '1' : '0');
+      ';Sampled=' + (subsegment.isSampled ? '1' : '0');
 
     const errorCapturer = function errorCapturer(e) {
       if (subsegmentCallback) {

packages/core/lib/segments/attributes/subsegment.d.ts

@@ -10,13 +10,18 @@ declare class Subsegment {
   parent: SegmentLike;
   segment: Segment;
   namespace?: string;
+  isSampled: boolean;
 
   constructor(name: string);
 
   addNewSubsegment(name: string): Subsegment;
 
   addSubsegment(subsegment: Subsegment): void;
 
+  addNewSubsegmentWithoutSampling(name: String): Subsegment;
+
+  addSubsegmentWithoutSampling(subsegment: Subsegment): void;
+
   removeSubsegment(subsegment: Subsegment): void;
 
   addAttribute(name: string, data: any): void;

packages/core/lib/segments/attributes/subsegment.js

@@ -28,6 +28,7 @@ Subsegment.prototype.init = function init(name) {
   this.start_time = SegmentUtils.getCurrentTime();
   this.in_progress = true;
   this.counter = 0;
+  this.isSampled = true;
 };
 
 /**
@@ -39,9 +40,23 @@ Subsegment.prototype.init = function init(name) {
 Subsegment.prototype.addNewSubsegment = function addNewSubsegment(name) {
   var subsegment = new Subsegment(name);
   this.addSubsegment(subsegment);
+  subsegment.isSampled = subsegment.segment? !subsegment.segment.notTraced : true;
   return subsegment;
 };
 
+Subsegment.prototype.addSubsegmentWithoutSampling = function addSubsegmentWithoutSampling(subsegment){
+  subsegment.isSampled = false;
+  this.addSubsegment(subsegment);
+};
+
+Subsegment.prototype.addNewSubsegmentWithoutSampling = function addNewSubsegmentWithoutSampling(name){
+  var subsegment = new Subsegment(name);
+  subsegment.isSampled = false;
+  this.addSubsegment(subsegment);
+
+  return subsegment; 
+};
+
 /**
  * Adds a subsegment to the array of subsegments.
  * @param {Subsegment} subsegment - The subsegment to append.
@@ -60,11 +75,14 @@ Subsegment.prototype.addSubsegment = function(subsegment) {
   subsegment.segment = this.segment;
   subsegment.parent = this;
 
-  if (subsegment.end_time === undefined) {
+  if (subsegment.end_time === undefined && subsegment.isSampled) {
     this.incrementCounter(subsegment.counter);
   }
-
-  this.subsegments.push(subsegment);
+  // don't push to subsegment array if subsegment is not sampled
+  if(subsegment.isSampled){ 
+    this.subsegments.push(subsegment);
+  }
+
 };
 
 /**
@@ -362,7 +380,7 @@ Subsegment.prototype.streamSubsegments = function streamSubsegments() {
     var open = [];
 
     this.subsegments.forEach(function(child) {
-      if (!child.streamSubsegments()) {
+      if (!child.streamSubsegments()) { 
         open.push(child);
       }
     });

packages/core/lib/segments/segment.d.ts

@@ -35,6 +35,10 @@ declare class Segment {
 
   addSubsegment(subsegment: Subsegment): void;
 
+  addSubsegmentWithoutSampling(subsegment: Subsegment): void; 
+
+  addNewSubsegmentWithoutSampling(name: String): Subsegment
+
   removeSubsegment(subsegment: Subsegment): void;
 
   addError(err: Error | string, remote?: boolean): void;

packages/core/lib/segments/segment.js

@@ -217,6 +217,20 @@ Segment.prototype.addNewSubsegment = function addNewSubsegment(name) {
   return subsegment;
 };
 
+Segment.prototype.addSubsegmentWithoutSampling = function addSubsegmentWithoutSampling(subsegment){
+  subsegment.isSampled = false;
+  this.addSubsegment(subsegment);
+
+};
+
+Segment.prototype.addNewSubsegmentWithoutSampling = function addNewSubsegmentWithoutSampling(name){
+  var subsegment = new Subsegment(name);
+  subsegment.isSampled = false;
+  this.addSubsegment(subsegment);
+
+  return subsegment; 
+};
+
 /**
  * Adds a subsegment to the array of subsegments.
  * @param {Subsegment} subsegment - The subsegment to append.
@@ -233,13 +247,18 @@ Segment.prototype.addSubsegment = function addSubsegment(subsegment) {
 
   subsegment.segment = this;
   subsegment.parent = this;
-  this.subsegments.push(subsegment);
 
-  if (!subsegment.end_time) {
+  if(subsegment.isSampled){
+      this.subsegments.push(subsegment);
+  }
+
+  if (!subsegment.end_time && subsegment.isSampled) {
     this.incrementCounter(subsegment.counter);
   }
 };
 
+
+
 /**
  * Removes the subsegment from the subsegments array, used in subsegment streaming.
  */

packages/core/test/unit/env/sqs_message_helper.test.js

@@ -0,0 +1,85 @@
+var assert = require('chai').assert;
+var chai = require('chai');
+var sinonChai = require('sinon-chai');
+
+import SqsMessageHelper from '../../../lib/env/sqs_message_helper';
+
+chai.should();
+chai.use(sinonChai);
+
+describe('#SqsMessageHelper', function (){
+
+    // sample records from https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html
+    const sampleSqsMessageEvent = {
+        "Records": [
+            {
+                "messageId": "059f36b4-87a3-44ab-83d2-661975830a7d",
+                "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a...",
+                "body": "Test message.",
+                "attributes": {
+                    "ApproximateReceiveCount": "1",
+                    "SentTimestamp": "1545082649183",
+                    "SenderId": "AIDAIENQZJOLO23YVJ4VO",
+                    "ApproximateFirstReceiveTimestamp": "1545082649185",
+                    "AWSTraceHeader":"Root=1-632BB806-bd862e3fe1be46a994272793;Sampled=1"
+                },
+                "messageAttributes": {},
+                "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+                "eventSource": "aws:sqs",
+                "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+                "awsRegion": "us-east-2"
+            },
+            {
+                "messageId": "2e1424d4-f796-459a-8184-9c92662be6da",
+                "receiptHandle": "AQEBzWwaftRI0KuVm4tP+/7q1rGgNqicHq...",
+                "body": "Test message.",
+                "attributes": {
+                    "ApproximateReceiveCount": "1",
+                    "SentTimestamp": "1545082650636",
+                    "SenderId": "AIDAIENQZJOLO23YVJ4VO",
+                    "ApproximateFirstReceiveTimestamp": "1545082650649",
+                    "AWSTraceHeader":"Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=0"
+                },
+                "messageAttributes": {},
+                "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+                "eventSource": "aws:sqs",
+                "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+                "awsRegion": "us-east-2"
+            },
+            {
+                "messageId": "2e1424d4-f796-459a-8184-9c92662be6da",
+                "receiptHandle": "AQEBzWwaftRI0KuVm4tP+/7q1rGgNqicHq...",
+                "body": "Test message.",
+                "attributes": {
+                    "ApproximateReceiveCount": "1",
+                    "SentTimestamp": "1545082650636",
+                    "SenderId": "AIDAIENQZJOLO23YVJ4VO",
+                    "ApproximateFirstReceiveTimestamp": "1545082650649",
+                    "AWSTraceHeader":"Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8"
+                },
+                "messageAttributes": {},
+                "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+                "eventSource": "aws:sqs",
+                "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+                "awsRegion": "us-east-2"
+            }
+        ]
+    }
+
+    describe('SqsMessageHelper isSampled', function(){
+        let sqsMessageHelper = new SqsMessageHelper();
+
+        it('should return true when AWSTraceHeader has Sampled=1', function(){
+            assert.equal(sqsMessageHelper.isSampled(sampleSqsMessageEvent.Records[0]), true)
+        });
+
+        it('should return false when AWSTraceHeader has Sampled=0', function(){
+            assert.equal(sqsMessageHelper.isSampled(sampleSqsMessageEvent.Records[1]), false)
+        });
+
+        it('should return false when AWSTraceHeader has no Sampled flag', function(){
+            assert.equal(sqsMessageHelper.isSampled(sampleSqsMessageEvent.Records[2]), false)
+        });
+
+    })
+  });

packages/core/test/unit/segments/attributes/subsegment.test.js

@@ -103,6 +103,33 @@ describe('Subsegment', function() {
     });
   });
 
+  describe('#addSubsegmentWithoutSampling', function (){
+    it('should have isSampled flag set to false', function(){
+      var subsegment = new Subsegment('test');
+      var child = new Subsegment('child')
+      subsegment.addSubsegmentWithoutSampling(child);
+
+      assert.equal(subsegment.isSampled, true);
+      assert.equal(child.isSampled, false);
+    })
+
+    it('should have isSampled flag set to false for new subsegment', function(){
+      var subsegment = new Subsegment('test');
+      var child = subsegment.addNewSubsegmentWithoutSampling('child');
+
+      assert.equal(subsegment.isSampled, true);
+      assert.equal(child.isSampled, false);
+    })
+
+    it('should not contain the child subsegment if not sampled', function(){
+      var subsegment = new Subsegment('test');
+      var child = new Subsegment('child')
+      subsegment.addSubsegmentWithoutSampling(child);
+
+      assert.notEqual(subsegment.subsegments[0], child);
+    })
+  });
+
   describe('#addError', function() {
     var err, exceptionStub, sandbox, subsegment;
 

packages/core/test/unit/segments/segment.test.js

@@ -9,6 +9,8 @@ var SegmentUtils = require('../../../lib/segments/segment_utils');
 var Segment = require('../../../lib/segments/segment');
 var Subsegment = require('../../../lib/segments/attributes/subsegment');
 var logger = require('../../../lib/logger');
+var Lambda = require('../../../lib/env/aws_lambda');
+var Context = require('../../../lib/context_utils');
 
 chai.should();
 chai.use(sinonChai);
@@ -262,6 +264,90 @@ describe('Segment', function() {
     });
   });
 
+  describe('#addSubsegmentWithoutSampling', function (){
+    let sandbox, setSegmentStub;
+
+    beforeEach(function() {
+      sandbox = sinon.createSandbox();
+      setSegmentStub = sandbox.stub(Context, 'setSegment');
+    });
+
+    afterEach(function() {
+      sandbox.restore();
+    });
+
+    it('should have isSampled flag set to false for subsegment of Lambda facade segment', function(){
+      process.env._X_AMZN_TRACE_ID;
+      Lambda.init();
+
+      setSegmentStub.should.have.been.calledOnce;
+
+      let facade = setSegmentStub.args[0][0];
+      let unsampledSegment = facade.addNewSubsegmentWithoutSampling('unsampled-subsegment');
+      assert.equal(unsampledSegment.isSampled, false);
+    })
+
+    it('should have isSampled flag set to true for subsegment of Lambda facade segment', function(){
+      process.env._X_AMZN_TRACE_ID;
+      Lambda.init();
+
+      setSegmentStub.should.have.been.calledOnce;
+
+      let facade = setSegmentStub.args[0][0];
+      let sampledSubsegment = facade.addNewSubsegment('sampled-subsegment');
+      assert.equal(sampledSubsegment.isSampled, true);
+    })
+
+    it('should have isSampled flag set to false', function(){
+      var segment = new Segment('parent');
+      var child = new Subsegment('child');
+      segment.addSubsegmentWithoutSampling(child);
+
+      assert.equal(child.isSampled, false);
+    })
+
+    it('should have isSampled flag set to false for new subsegment', function(){
+      var segment = new Segment('parent');
+      var child = segment.addNewSubsegmentWithoutSampling('child');
+
+      assert.equal(child.isSampled, false);
+    })
+
+    it('should not contain the child subsegment if not sampled', function(){
+      var segment = new Segment('parent');
+      var child = new Subsegment('child');
+      segment.addSubsegmentWithoutSampling(child);
+
+      assert.notEqual(segment.subsegments[0], child);
+    })
+
+
+    it('should not sample subsegment or subsegment of subsegment', function(){
+      var segment = new Segment('parent');
+      var child = new Subsegment('child');
+      var child2 = new Subsegment('child-2');
+      segment.addSubsegmentWithoutSampling(child);
+      child.addSubsegmentWithoutSampling(child2)
+
+      assert.equal(child.isSampled, false);
+      assert.equal(child2.isSampled, false);
+    })
+
+    it('should not sample subsegment or subsegment of subsegment - mix', function(){
+      var segment = new Segment('parent');
+      var child = new Subsegment('child');
+      var child2 = new Subsegment('child-2');
+      var child3 = new Subsegment('child-3');
+      segment.addSubsegmentWithoutSampling(child);
+      child.addSubsegment(child2)
+      child.addSubsegmentWithoutSampling(child3)
+
+      assert.equal(child.isSampled, false);
+      assert.equal(child2.isSampled, true);
+      assert.equal(child3.isSampled, false);
+    })
+  });
+
   describe('#addError', function() {
     var err, exceptionStub, sandbox, segment;