Merge pull request #722 from jj22ee/xray-header-from-context-support

Support extracting X-Ray Trace Header Lambda Invoke Store

Author: wangzlei

package-lock.json

@@ -8,6 +8,8 @@ var SegmentUtils = require('../segments/segment_utils');
 var logger = require('../logger');
 const TraceID = require('../segments/attributes/trace_id');
 
+const { InvokeStore } = require('@aws/lambda-invoke-store');
+
 /**
  * @namespace
  * @ignore
@@ -80,7 +82,9 @@ var facadeSegment = function facadeSegment() {
   };
 
   segment.resolveLambdaTraceData = function resolveLambdaTraceData() {
-    var xAmznLambda = process.env._X_AMZN_TRACE_ID;
+    const traceIdFromInvokeStore = InvokeStore.getXRayTraceId();
+    const traceIdFromEnv = process.env._X_AMZN_TRACE_ID;
+    var xAmznLambda = traceIdFromInvokeStore ?? traceIdFromEnv;
 
     if (xAmznLambda) {
 

packages/core/lib/env/aws_lambda.js

@@ -24,6 +24,7 @@
   "//": "@types/cls-hooked is exposed in API so must be in dependencies, not devDependencies",
   "dependencies": {
     "@aws-sdk/types": "^3.4.1",
+    "@aws/lambda-invoke-store": "^0.0.1",
     "@smithy/service-error-classification": "^2.0.4",
     "@types/cls-hooked": "^4.3.3",
     "atomic-batcher": "^1.0.2",
@@ -55,4 +56,4 @@
   ],
   "license": "Apache-2.0",
   "repository": "https://github.com/aws/aws-xray-sdk-node/tree/master/packages/core"
-}
+}

packages/core/package.json

@@ -3,6 +3,7 @@ var chai = require('chai');
 var sinon = require('sinon');
 var sinonChai = require('sinon-chai');
 
+chai.should();
 chai.use(sinonChai);
 
 var contextUtils = require('../../../lib/context_utils');
@@ -13,6 +14,7 @@ var Segment = require('../../../lib/segments/segment');
 var SegmentUtils = require('../../../lib/segments/segment_utils');
 var SegmentEmitter = require('../../../lib/segment_emitter');
 const TraceID = require('../../../lib/segments/attributes/trace_id');
+const { InvokeStore } = require('@aws/lambda-invoke-store');
 
 describe('AWSLambda', function() {
   var sandbox;
@@ -147,6 +149,8 @@ describe('AWSLambda', function() {
 
     describe('#resolveLambdaTraceData', function() {
       var sandbox, traceId;
+      var INVOKE_STORE_TRACE_ID = 'Root=1-12345678-12345678901234567890abcd;Parent=abcdef0123456789;Sampled=1';
+      var ENV_TRACE_ID = 'Root=1-87654321-09876543210987654321dcba;Parent=fedcba9876543210;Sampled=0';
 
       beforeEach(function() {
         sandbox = sinon.createSandbox();
@@ -197,6 +201,47 @@ describe('AWSLambda', function() {
         facade.resolveLambdaTraceData();
         populateStub.should.have.not.been.called;
       });
+
+      it('should prioritize InvokeStore trace ID over environment variable if both are defined', function() {
+        process.env._X_AMZN_TRACE_ID = ENV_TRACE_ID;
+        Lambda.init();
+
+        sandbox.stub(InvokeStore, 'getXRayTraceId').returns(INVOKE_STORE_TRACE_ID);
+        populateStub.reset();
+
+        var facade = setSegmentStub.args[0][0];
+        facade.resolveLambdaTraceData();
+
+        populateStub.should.have.been.calledWith(facade, INVOKE_STORE_TRACE_ID);
+      });
+
+      it('should use InvokeStore trace ID when environment variable is undefined', function() {
+        delete process.env._X_AMZN_TRACE_ID;
+
+        Lambda.init();
+
+        sandbox.stub(InvokeStore, 'getXRayTraceId').returns(INVOKE_STORE_TRACE_ID);
+        populateStub.reset();
+
+        var facade = setSegmentStub.args[0][0];
+        facade.resolveLambdaTraceData();
+
+        populateStub.should.have.been.calledWith(facade, INVOKE_STORE_TRACE_ID);
+      });
+
+      it('should use environment variable when InvokeStore returns undefined', function() {
+        process.env._X_AMZN_TRACE_ID = INVOKE_STORE_TRACE_ID;
+        Lambda.init();
+
+        process.env._X_AMZN_TRACE_ID = ENV_TRACE_ID;
+        sandbox.stub(InvokeStore, 'getXRayTraceId').returns(undefined);
+        populateStub.reset();
+
+        var facade = setSegmentStub.args[0][0];
+        facade.resolveLambdaTraceData();
+
+        populateStub.should.have.been.calledWith(facade, ENV_TRACE_ID);
+      });
     });
   });