Oversampling Mitigation (#541)

* added implementation for oversampling mitigation * removed unnecessary comments and whitespaces * addressed feedback and fixed style issues * fixed sampling decision bug and updated unit tests * Make SqsMessageHelper static * Fixed trace headers and added unit tests * fixed style with eslint * renamed isSampled flag to notTraced and fixed styling with eslint * Revert "renamed isSampled flag to notTraced and fixed styling with eslint" This reverts commit f5ce869. * Revert "fixed style with eslint" This reverts commit 7cde0b2. * changed isSampled flag name to notTraced and fixed lint issues * fixed trace header bug
aws · Nov 10, 2022 · 93a4f31 · 93a4f31
1 parent 0db1b68
commit 93a4f31
Show file tree

Hide file tree

Showing 15 changed files with 823 additions and 12 deletions.
diff --git a/packages/core/lib/env/aws_lambda.js b/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;
 

diff --git a/packages/core/lib/env/sqs_message_helper.js b/packages/core/lib/env/sqs_message_helper.js
@@ -0,0 +1,12 @@
+class SqsMessageHelper {
+
+  static isSampled(message) {
+    const {attributes} = message; // extract attributes from message
+    if (!('AWSTraceHeader' in attributes)) {
+      return false;
+    }
+    return attributes['AWSTraceHeader'].includes('Sampled=1');
+  }
+}
+
+export default SqsMessageHelper;
diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts
@@ -110,15 +110,21 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
     return next(args);
   }
 
-  const subsegment: Subsegment = segment.addNewSubsegment(service);
+  let subsegment: Subsegment;
+
+  if (segment.notTraced) {
+    subsegment = segment.addNewSubsegmentWithoutSampling(service);
+  } else {
+    subsegment = segment.addNewSubsegment(service);
+  }
   subsegment.addAttribute('namespace', 'aws');
   const parent = (segment instanceof Subsegment ? segment.segment : segment);
 
   args.request.headers['X-Amzn-Trace-Id'] = stringify(
     {
       Root: parent.trace_id,
       Parent: subsegment.id,
-      Sampled: parent.notTraced ? '0' : '1',
+      Sampled: subsegment.notTraced ? '0' : '1',
     },
     ';',
   );

diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js
@@ -76,12 +76,19 @@ function captureAWSRequest(req) {
   var throttledError = this.throttledError || throttledErrorDefault;
 
   var stack = (new Error()).stack;
-  var subsegment = parent.addNewSubsegment(this.serviceIdentifier);
+
+  let subsegment;
+  if (parent.notTraced) {
+    subsegment = parent.addNewSubsegmentWithoutSampling(this.serviceIdentifier);
+  } else {
+    subsegment = parent.addNewSubsegment(this.serviceIdentifier);
+  }
+
   var traceId = parent.segment ? parent.segment.trace_id : parent.trace_id;
 
   var buildListener = function(req) {
     req.httpRequest.headers['X-Amzn-Trace-Id'] = 'Root=' + traceId + ';Parent=' + subsegment.id +
-      ';Sampled=' + (subsegment.segment.notTraced ? '0' : '1');
+      ';Sampled=' + (subsegment.notTraced ? '0' : '1');
   };
 
   var completeListener = function(res) {

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) {
+      subsegment = parent.addNewSubsegmentWithoutSampling(hostname);
+    } else {
+      subsegment = parent.addNewSubsegment(hostname);
+    }
+
     const root = parent.segment ? parent.segment : parent;
     subsegment.namespace = 'remote';
 
@@ -125,7 +131,7 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) {
     }
 
     options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
-      ';Sampled=' + (!root.notTraced ? '1' : '0');
+      ';Sampled=' + (subsegment.notTraced ? '0' : '1');
 
     const errorCapturer = function errorCapturer(e) {
       if (subsegmentCallback) {

diff --git a/packages/core/lib/segments/attributes/subsegment.d.ts b/packages/core/lib/segments/attributes/subsegment.d.ts
@@ -10,13 +10,18 @@ declare class Subsegment {
   parent: SegmentLike;
   segment: Segment;
   namespace?: string;
+  notTraced: 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;

diff --git a/packages/core/lib/segments/attributes/subsegment.js b/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.notTraced = false;
 };
 
 /**
@@ -37,11 +38,23 @@ Subsegment.prototype.init = function init(name) {
  */
 
 Subsegment.prototype.addNewSubsegment = function addNewSubsegment(name) {
-  var subsegment = new Subsegment(name);
+  const subsegment = new Subsegment(name);
   this.addSubsegment(subsegment);
   return subsegment;
 };
 
+Subsegment.prototype.addSubsegmentWithoutSampling = function addSubsegmentWithoutSampling(subsegment) {
+  this.addSubsegment(subsegment);
+  subsegment.notTraced = true;
+};
+
+Subsegment.prototype.addNewSubsegmentWithoutSampling = function addNewSubsegmentWithoutSampling(name) {
+  const subsegment = new Subsegment(name);
+  this.addSubsegment(subsegment);
+  subsegment.notTraced = true;
+  return subsegment;
+};
+
 /**
  * Adds a subsegment to the array of subsegments.
  * @param {Subsegment} subsegment - The subsegment to append.
@@ -60,11 +73,13 @@ Subsegment.prototype.addSubsegment = function(subsegment) {
   subsegment.segment = this.segment;
   subsegment.parent = this;
 
+  subsegment.notTraced = subsegment.parent.notTraced;
+
   if (subsegment.end_time === undefined) {
     this.incrementCounter(subsegment.counter);
   }
-
   this.subsegments.push(subsegment);
+
 };
 
 /**
@@ -340,7 +355,7 @@ Subsegment.prototype.flush = function flush() {
   }
 
   if (this.segment.trace_id) {
-    if (this.segment.notTraced !== true) {
+    if (this.segment.notTraced !== true && !this.notTraced) {
       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
@@ -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;

diff --git a/packages/core/lib/segments/segment.js b/packages/core/lib/segments/segment.js
@@ -217,6 +217,19 @@ Segment.prototype.addNewSubsegment = function addNewSubsegment(name) {
   return subsegment;
 };
 
+Segment.prototype.addSubsegmentWithoutSampling = function addSubsegmentWithoutSampling(subsegment) {
+  this.addSubsegment(subsegment);
+  subsegment.notTraced = true;
+
+};
+
+Segment.prototype.addNewSubsegmentWithoutSampling = function addNewSubsegmentWithoutSampling(name) {
+  const subsegment = new Subsegment(name);
+  this.addSubsegment(subsegment);
+  subsegment.notTraced = true;
+  return subsegment;
+};
+
 /**
  * Adds a subsegment to the array of subsegments.
  * @param {Subsegment} subsegment - The subsegment to append.
@@ -233,13 +246,17 @@ Segment.prototype.addSubsegment = function addSubsegment(subsegment) {
 
   subsegment.segment = this;
   subsegment.parent = this;
+
+  subsegment.notTraced = subsegment.parent.notTraced;
   this.subsegments.push(subsegment);
 
   if (!subsegment.end_time) {
     this.incrementCounter(subsegment.counter);
   }
 };
 
+
+
 /**
  * Removes the subsegment from the subsegments array, used in subsegment streaming.
  */

diff --git a/packages/core/test/unit/env/sqs_message_helper.test.js b/packages/core/test/unit/env/sqs_message_helper.test.js
@@ -0,0 +1,84 @@
+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() {
+
+    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);
+    });
+
+  });
+});