Merge pull request #657 from wangzlei/master

Revert #651 [Lambda] Replace Facade with No-Op if trace header is missing data
aws · Jun 6, 2024 · 053b90f · 053b90f
2 parents 23ff8bb + 71b6f2c
commit 053b90f
Show file tree

Hide file tree

Showing 10 changed files with 24 additions and 224 deletions.
diff --git a/packages/core/lib/context_utils.js b/packages/core/lib/context_utils.js
@@ -99,10 +99,8 @@ var contextUtils = {
 
       if (!segment) {
         contextUtils.contextMissingStrategy.contextMissing('Failed to get the current sub/segment from the context.');
-      } else if (segment instanceof Segment && process.env.LAMBDA_TASK_ROOT) {
-        if (segment.facade == true || segment.noOp == true) {
-          segment.resolveLambdaTraceData();
-        }
+      } else if (segment instanceof Segment && process.env.LAMBDA_TASK_ROOT && segment.facade == true) {
+        segment.resolveLambdaTraceData();
       }
 
       return segment;

diff --git a/packages/core/lib/env/aws_lambda.js b/packages/core/lib/env/aws_lambda.js
@@ -36,12 +36,7 @@ module.exports.init = function init() {
 
   var namespace = contextUtils.getNamespace();
   namespace.enter(namespace.createContext());
-
-  if (LambdaUtils.validTraceData(process.env._X_AMZN_TRACE_ID)) {
-    contextUtils.setSegment(facadeSegment());
-  } else {
-    contextUtils.setSegment(noOpSegment());
-  }
+  contextUtils.setSegment(facadeSegment());
 };
 
 var facadeSegment = function facadeSegment() {
@@ -114,75 +109,3 @@ var facadeSegment = function facadeSegment() {
 
   return segment;
 };
-
-var noOpSegment = function noOpSegment() {
-  var segment = new Segment('no-op');
-  var whitelistFcn = [];
-  var silentFcn = ['addNewSubsegment', 'addSubsegment', 'removeSubsegment', 'toString', 'addSubsegmentWithoutSampling', 'addNewSubsegmentWithoutSampling', 'incrementCounter', 'decrementCounter', 'isClosed', 'close', 'format', 'flush'];
-  var xAmznTraceId = process.env._X_AMZN_TRACE_ID;
-
-  for (var key in segment) {
-    if (typeof segment[key] === 'function' && whitelistFcn.indexOf(key) === -1) {
-      if (silentFcn.indexOf(key) === -1) {
-        segment[key] = (function() {
-          var func = key;
-          return function noOp() {
-            logger.getLogger().warn('Function "' + func + '" cannot be called on an AWS Lambda segment. Please use a subsegment to record data.');
-            return;
-          };
-        })();
-      } else {
-        segment[key] = function noOp() {
-          return;
-        };
-      }
-    }
-  }
-
-  segment.trace_id = TraceID.Invalid().toString();
-  segment.isClosed = function() {
-    return true;
-  };
-  segment.in_progress = false;
-  segment.counter = 1;
-  segment.notTraced = true;
-  segment.noOp = true;
-
-  segment.reset = function reset() {
-    this.trace_id = TraceID.Invalid().toString();
-    this.id = '00000000';
-    delete this.subsegments;
-    this.notTraced = true;
-  };
-
-  segment.resolveLambdaTraceData = function resolveLambdaTraceData() {
-    var xAmznLambda = process.env._X_AMZN_TRACE_ID;
-
-    if (xAmznLambda) {
-
-      // This check resets the trace data whenever a new trace header is read to not leak data between invocations
-      if (xAmznLambda != xAmznTraceIdPrev) {
-        this.reset();
-
-        if (LambdaUtils.populateTraceData(segment, xAmznLambda)) {
-          xAmznTraceIdPrev = xAmznLambda;
-        }
-      }
-    } else {
-      this.reset();
-      contextUtils.contextMissingStrategy.contextMissing('Missing AWS Lambda trace data for X-Ray. ' +
-          'Ensure Active Tracing is enabled and no subsegments are created outside the function handler.');
-    }
-  };
-
-  // Since we're in a no-op segment, do not check if the trace data is valid; simply propagate the information
-  if (LambdaUtils.populateTraceData(segment, xAmznTraceId)) {
-    xAmznTraceIdPrev = xAmznTraceId;
-  }
-
-  return segment;
-};
-
-// For testing
-export const exportedFacadeSegment = { facadeSegment };
-export const exportedNoOpSegment = { noOpSegment };
diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts
@@ -125,10 +125,14 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
   const parent = (segment instanceof Subsegment ? segment.segment : segment);
   const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
 
-  let traceHeader = 'Root=' + parent.trace_id;
-  if (!(parent && parent.noOp)) {
-    traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1');
-  }
+  let traceHeader = stringify(
+    {
+      Root: parent.trace_id,
+      Parent: subsegment.id,
+      Sampled: subsegment.notTraced ? '0' : '1',
+    },
+    ';',
+  );
 
   if (data != null) {
     for (const [key, value] of Object.entries(data)) {

diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js
@@ -88,11 +88,8 @@ function captureAWSRequest(req) {
   const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
 
   var buildListener = function(req) {
-    let traceHeader = 'Root=' + traceId;
-    // Only append parent and sample if not in Lambda PassThrough mode
-    if (!(parent && parent.noOp)) {
-      traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1');
-    }
+    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;
@@ -102,9 +99,6 @@ function captureAWSRequest(req) {
   };
 
   var completeListener = function(res) {
-    if (subsegment == null) {
-      return;
-    }
     subsegment.addAttribute('namespace', 'aws');
     subsegment.addAttribute('aws', new Aws(res, subsegment.name));
 

diff --git a/packages/core/lib/patchers/http_p.js b/packages/core/lib/patchers/http_p.js
@@ -123,18 +123,15 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) {
       subsegment = parent.addNewSubsegment(hostname);
     }
 
+    const root = parent.segment ? parent.segment : parent;
     subsegment.namespace = 'remote';
 
     if (!options.headers) {
       options.headers = {};
     }
 
-    let traceHeader = 'Root=' + (parent.segment ? parent.segment : parent).trace_id;
-    if (!(parent && parent.noOp)) {
-      traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1');
-    }
-
-    options.headers['X-Amzn-Trace-Id'] = traceHeader;
+    options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
+      ';Sampled=' + (subsegment.notTraced ? '0' : '1');
 
     const errorCapturer = function errorCapturer(e) {
       if (subsegmentCallback) {

diff --git a/packages/core/lib/utils.js b/packages/core/lib/utils.js
@@ -168,8 +168,6 @@ var utils = {
       if (!traceData) {
         traceData = {};
         logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format');
-      } else if (segment.noOp == true && traceData.root) {
-        valid = true;
       } else if (!traceData.root || !traceData.parent || !traceData.sampled) {
         logger.getLogger().error('_X_AMZN_TRACE_ID is missing required information');
       } else {

diff --git a/packages/core/test/unit/env/aws_lambda.test.js b/packages/core/test/unit/env/aws_lambda.test.js
@@ -81,7 +81,7 @@ describe('AWSLambda', function() {
       assert.equal(facade.trace_id, TraceID.Invalid().toString());
     });
 
-    describe('the facade/no-op segment', function() {
+    describe('the facade segment', function() {
       afterEach(function() {
         populateStub.returns(true);
         delete process.env._X_AMZN_TRACE_ID;
@@ -95,25 +95,17 @@ describe('AWSLambda', function() {
         validateStub.should.have.been.calledWith(process.env._X_AMZN_TRACE_ID);
       });
 
-      it('should call populateTraceData on Facade if validTraceData returns true', function() {
+      it('should call populateTraceData if validTraceData returns true', function() {
         Lambda.init();
 
-        var segment = setSegmentStub.args[0][0];
-        assert.equal(segment.name, 'facade');
-        assert.isTrue(segment.facade);
-
         populateStub.should.have.been.calledOnce;
       });
 
-      it('should call populateTraceData on No-Op if validTraceData returns false', function() {
+      it('should not call populateTraceData if validTraceData returns false', function() {
         validateStub.returns(false);
         Lambda.init();
 
-        var segment = setSegmentStub.args[0][0];
-        assert.equal(segment.name, 'no-op');
-        assert.isTrue(segment.noOp);
-
-        populateStub.should.have.been.calledOnce;
+        populateStub.should.have.not.been.called;
       });
     });
   });

diff --git a/packages/core/test/unit/patchers/aws_p.test.js b/packages/core/test/unit/patchers/aws_p.test.js
@@ -13,10 +13,6 @@ var Utils = require('../../../lib/utils');
 
 var logger = require('../../../lib/logger').getLogger();
 
-import { exportedFacadeSegment, exportedNoOpSegment } from '../../../lib/env/aws_lambda';
-const { facadeSegment } = exportedFacadeSegment;
-const { noOpSegment } = exportedNoOpSegment;
-
 chai.should();
 chai.use(sinonChai);
 
@@ -352,91 +348,4 @@ describe('AWS patcher', function() {
     });
 
   });
-
-
-  describe('#captureAWSRequest-Lambda-PassThrough', function() {
-    var awsClient, awsRequest, MyEmitter, sandbox, segment, stubResolve, tempHeader;
-
-    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();
-
-      tempHeader = process.env._X_AMZN_TRACE_ID;
-      process.env._X_AMZN_TRACE_ID = 'Root=' + traceId + ';Foo=bar';
-
-      segment = noOpSegment();
-
-      stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(segment);
-    });
-
-    afterEach(function() {
-      process.env._X_AMZN_TRACE_ID = tempHeader;
-      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 + ';Foo=bar$');
-        assert.match(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], expected);
-        done();
-      }, 50);
-    });
-
-  });
 });
diff --git a/sdk_contrib/fetch/lib/fetch_p.js b/sdk_contrib/fetch/lib/fetch_p.js
@@ -104,12 +104,10 @@ const enableCapture = function enableCapture(baseFetchFunction, requestClass, do
 
     subsegment.namespace = 'remote';
 
-    let traceHeader = 'Root=' + (parent.segment ? parent.segment : parent).trace_id;
-    if (!(parent && parent.noOp)) {
-      traceHeader += ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1');
-    }
-
-    request.headers.set('X-Amzn-Trace-Id', traceHeader);
+    request.headers.set('X-Amzn-Trace-Id',
+      'Root=' + (parent.segment ? parent.segment : parent).trace_id +
+      ';Parent=' + subsegment.id +
+      ';Sampled=' + (subsegment.notTraced ? '0' : '1'));
 
     // Set up fetch call and capture any thrown errors
     const capturedFetch = async () => {

diff --git a/sdk_contrib/fetch/test/unit/fetch_p.test.js b/sdk_contrib/fetch/test/unit/fetch_p.test.js
@@ -273,19 +273,6 @@ describe('Unit tests', function () {
         'Root=12345;Parent=999;Sampled=1');
     });
 
-    it('adds X-Amzn-Trace-Id header with only root if noOp', async function () {
-      const activeFetch = captureFetch(true);
-      stubParentSegment.noOp = true;
-      stubParentSegment.trace_id = '12345';
-      stubSubsegment.notTraced = false;
-      stubSubsegment.id = '999';
-      const request = new FetchRequest('https://www.foo.com/test');
-      const requestHeadersSet = sandbox.stub(request.headers, 'set');
-      await activeFetch(request);
-      requestHeadersSet.should.have.been.calledOnceWith('X-Amzn-Trace-Id',
-        'Root=12345');
-    });
-
     it('calls subsegmentCallback on successful response', async function () {
       const spyCallback = sandbox.spy();
       const activeFetch = captureFetch(true, spyCallback);