Propagate additional trace data into AWS requests on Lambda (#549)

aws · Dec 13, 2022 · b0aa443 · b0aa443
1 parent 8a4f584
commit b0aa443
Show file tree

Hide file tree

Showing 8 changed files with 85 additions and 28 deletions.
diff --git a/.github/workflows/pr-build.yml b/.github/workflows/pr-build.yml
@@ -11,6 +11,7 @@ jobs:
     name: Build Node ${{ matrix.node-version }} on ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
+      fail-fast: false
       matrix:
         os:
           - macos-latest

diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js
@@ -85,10 +85,17 @@ function captureAWSRequest(req) {
   }
 
   var traceId = parent.segment ? parent.segment.trace_id : parent.trace_id;
+  const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
 
   var buildListener = function(req) {
-    req.httpRequest.headers['X-Amzn-Trace-Id'] = 'Root=' + traceId + ';Parent=' + subsegment.id +
+    let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id +
       ';Sampled=' + (subsegment.notTraced ? '0' : '1');
+    if (data != null) {
+      for (const [key, value] of Object.entries(data)) {
+        traceHeader += ';' + key +'=' + value;
+      }
+    }
+    req.httpRequest.headers['X-Amzn-Trace-Id'] = traceHeader;
   };
 
   var completeListener = function(res) {

diff --git a/packages/core/lib/segments/segment.d.ts b/packages/core/lib/segments/segment.d.ts
@@ -13,6 +13,8 @@ declare class Segment {
   subsegments?: Array<Subsegment>;
   notTraced?: boolean;
 
+  additionalTraceData?: object
+
   constructor(name: string, rootId?: string | null, parentId?: string | null);
 
   addIncomingRequestData(data: IncomingRequestData): void;

diff --git a/packages/core/lib/utils.js b/packages/core/lib/utils.js
@@ -162,33 +162,37 @@ var utils = {
      */
     populateTraceData: function(segment, xAmznTraceId) {
       logger.getLogger().debug('Lambda trace data found: ' + xAmznTraceId);
-      var data = utils.processTraceData(xAmznTraceId);
+      let traceData = utils.processTraceData(xAmznTraceId);
       var valid = false;
 
-      if (!data) {
-        data = {};
+      if (!traceData) {
+        traceData = {};
         logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format');
-      } else if (!data.root || !data.parent || !data.sampled) {
+      } else if (!traceData.root || !traceData.parent || !traceData.sampled) {
         logger.getLogger().error('_X_AMZN_TRACE_ID is missing required information');
       } else {
         valid = true;
       }
 
-      segment.trace_id = TraceID.FromString(data.root).toString();  // Will always assign valid trace_id
-      segment.id = data.parent || crypto.randomBytes(8).toString('hex');
+      segment.trace_id = TraceID.FromString(traceData.root).toString();  // Will always assign valid trace_id
+      segment.id = traceData.parent || crypto.randomBytes(8).toString('hex');
 
-      if (data.root && segment.trace_id !== data.root)  {
+      if (traceData.root && segment.trace_id !== traceData.root)  {
         logger.getLogger().error('_X_AMZN_TRACE_ID contains invalid trace ID');
         valid = false;
       }
 
-      if (!parseInt(data.sampled)) {
+      if (!parseInt(traceData.sampled)) {
         segment.notTraced = true;
       } else {
         delete segment.notTraced;
       }
 
-      logger.getLogger().debug('Segment started: ' + JSON.stringify(data));
+      if (traceData.data) {
+        segment.userData = traceData.data;
+      }
+
+      logger.getLogger().debug('Segment started: ' + JSON.stringify(traceData));
       return valid;
     }
   },
@@ -202,6 +206,7 @@ var utils = {
 
   processTraceData: function processTraceData(traceData) {
     var amznTraceData = {};
+    var data = {};
     var reservedKeywords = ['root', 'parent', 'sampled', 'self'];
     var remainingBytes = 256;
 
@@ -217,19 +222,22 @@ var utils = {
       var pair = header.split('=');
 
       if (pair[0] && pair[1]) {
-        var key = pair[0].trim().toLowerCase();
-        var value = pair[1].trim().toLowerCase();
-        var reserved = reservedKeywords.indexOf(key) !== -1;
+        let key = pair[0].trim();
+        let value = pair[1].trim();
+        let lowerCaseKey = key.toLowerCase();
+        let reserved = reservedKeywords.indexOf(lowerCaseKey) !== -1;
 
         if (reserved) {
-          amznTraceData[key] = value;
-        } else if (!reserved && remainingBytes - (key.length + value.length) >= 0) {
-          amznTraceData[key] = value;
+          amznTraceData[lowerCaseKey] = value;
+        } else if (!reserved && remainingBytes - (lowerCaseKey.length + value.length) >= 0) {
+          data[key] = value;
           remainingBytes -= (key.length + value.length);
         }
       }
     });
 
+    amznTraceData['data'] = data;
+
     return amznTraceData;
   },
 

diff --git a/packages/core/test/unit/env/aws_lambda.test.js b/packages/core/test/unit/env/aws_lambda.test.js
@@ -199,4 +199,31 @@ describe('AWSLambda', function() {
       });
     });
   });
+
+  describe('PopulateAdditionalTraceData', function() {
+    var sandbox, setSegmentStub;
+
+    beforeEach(function() {
+      sandbox = sinon.createSandbox();
+      sandbox.stub(SegmentEmitter, 'disableReusableSocket');
+      sandbox.stub(LambdaUtils, 'validTraceData').returns(true);
+
+      setSegmentStub = sandbox.stub(contextUtils, 'setSegment');
+    });
+
+    afterEach(function() {
+      delete process.env._X_AMZN_TRACE_ID;
+      sandbox.restore();
+    });
+
+    it('should populate additional trace data', function() {
+      process.env._X_AMZN_TRACE_ID = 'Root=traceId;Lineage=1234abcd:4|3456abcd:6';
+      Lambda.init();
+
+      var facade = setSegmentStub.args[0][0];
+      facade.resolveLambdaTraceData();
+      var userData = facade.userData;
+      assert.equal(userData['Lineage'], '1234abcd:4|3456abcd:6');
+    });
+  });
 });
diff --git a/packages/core/test/unit/middleware/mw_utils.test.js b/packages/core/test/unit/middleware/mw_utils.test.js
@@ -72,23 +72,23 @@ describe('Middleware utils', function() {
       req.headers[XRAY_HEADER] = 'Root=' + traceId;
       var headers = MWUtils.processHeaders(req);
 
-      assert.deepEqual(headers, {root: traceId});
+      assert.deepEqual(headers, {root: traceId, data: {}});
     });
 
     it('should return a split array on an request with an "x-amzn-trace-id" header with a root ID and parent ID', function() {
       var req = { headers: {}};
       req.headers[XRAY_HEADER] = 'Root=' + traceId + '; Parent=' + parentId;
       var headers = MWUtils.processHeaders(req);
 
-      assert.deepEqual(headers, {root: traceId, parent: parentId});
+      assert.deepEqual(headers, {root: traceId, parent: parentId, data: {}});
     });
 
     it('should return a split array on an request with an "x-amzn-trace-id" header with a root ID, parent ID and sampling', function() {
       var req = { headers: {}};
       req.headers[XRAY_HEADER] = 'Root=' + traceId + '; Parent=' + parentId + '; Sampled=0';
       var headers = MWUtils.processHeaders(req);
 
-      assert.deepEqual(headers, {root: traceId, parent: parentId, sampled: '0'});
+      assert.deepEqual(headers, {root: traceId, parent: parentId, sampled: '0', data: {}});
     });
   });
 

diff --git a/packages/core/test/unit/patchers/aws_p.test.js b/packages/core/test/unit/patchers/aws_p.test.js
@@ -124,6 +124,7 @@ describe('AWS patcher', function() {
       awsRequest.emitter = new MyEmitter();
 
       segment = new Segment('testSegment', traceId);
+      segment.additionalTraceData = {'Foo': 'bar'};
       sub = segment.addNewSubsegment('subseg');
 
       stubResolveManual = sandbox.stub(contextUtils, 'resolveManualSegmentParams');
@@ -161,7 +162,7 @@ describe('AWS patcher', function() {
       awsRequest.emitter.emit('build');
 
       setTimeout(function() {
-        var expected = new RegExp('^Root=' + traceId + ';Parent=' + sub.id + ';Sampled=1$');
+        var expected = new RegExp('^Root=' + traceId + ';Parent=' + sub.id + ';Sampled=1' + ';Foo=bar$');
         assert.match(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], expected);
         done();
       }, 50);
@@ -307,6 +308,7 @@ describe('AWS patcher', function() {
       awsRequest.emitter = new MyEmitter();
 
       segment = new Segment('testSegment', traceId);
+      segment.additionalTraceData = {'Foo': 'bar'};
       sub = segment.addNewSubsegmentWithoutSampling('subseg');
       service = sub.addNewSubsegmentWithoutSampling('service');
 
@@ -339,7 +341,7 @@ describe('AWS patcher', function() {
       awsRequest.emitter.emit('build');
 
       setTimeout(function() {
-        var expected = new RegExp('^Root=' + traceId + ';Parent=' + service.id + ';Sampled=0$');
+        var expected = new RegExp('^Root=' + traceId + ';Parent=' + service.id + ';Sampled=0' + ';Foo=bar$');
         assert.match(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], expected);
         done();
       }, 50);

diff --git a/packages/core/test/unit/utils.test.js b/packages/core/test/unit/utils.test.js
@@ -75,48 +75,58 @@ describe('Utils', function() {
 
     it('should handle trace header values with excess semicolons correctly', function() {
       assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;'), {
-        root: '1-58ed6027-14afb2e09172c337713486c0'
+        root: '1-58ed6027-14afb2e09172c337713486c0',
+        data: {}
       });
     });
 
     it('should handle malformed key=value pairs correctly (missing value)', function() {
       assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;Parent'), {
-        root: '1-58ed6027-14afb2e09172c337713486c0'
+        root: '1-58ed6027-14afb2e09172c337713486c0',
+        data: {}
       });
     });
 
     it('should handle malformed key=value pairs correctly (empty key)', function() {
       assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;=48af77592b6dd73f'), {
-        root: '1-58ed6027-14afb2e09172c337713486c0'
+        root: '1-58ed6027-14afb2e09172c337713486c0',
+        data: {}
       });
     });
 
     it('should handle malformed key=value pairs correctly (empty value)', function() {
       assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;Parent='), {
-        root: '1-58ed6027-14afb2e09172c337713486c0'
+        root: '1-58ed6027-14afb2e09172c337713486c0',
+        data: {}
       });
     });
 
     it('should accept arbitrary key=value pairs', function() {
       assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;Foo=bar'), {
         root: '1-58ed6027-14afb2e09172c337713486c0',
-        foo: 'bar'
+        data: {
+          Foo: 'bar'
+        }
       });
     });
 
     it('should not accept arbitrary key=value pairs that exceed the 256 byte limit', function() {
       var longVal = 'a'.repeat(251);
       assert.deepEqual(Utils.processTraceData(`Root=1-58ed6027-14afb2e09172c337713486c0;Foo=bar;Baz=${longVal}`), {
         root: '1-58ed6027-14afb2e09172c337713486c0',
-        foo: 'bar'
+        data: {
+          Foo: 'bar'
+        }
       });
     });
 
     it('should always accept reserved keywords even if unreserved capacity exceeded', function() {
       var longVal = 'a'.repeat(251);
       assert.deepEqual(Utils.processTraceData(`Baz=${longVal};Root=1-58ed6027-14afb2e09172c337713486c0;Foo=bar`), {
         root: '1-58ed6027-14afb2e09172c337713486c0',
-        baz: longVal
+        data: {
+          Baz: longVal
+        }
       });
     });
   });