[Lambda] Replace Facade segment with No-Op if trace header is missing data

Author: majanjua-amzn

packages/core/lib/context_utils.js

@@ -99,8 +99,10 @@ 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 && segment.facade == true) {
-        segment.resolveLambdaTraceData();
+      } else if (segment instanceof Segment && process.env.LAMBDA_TASK_ROOT) {
+        if (segment.facade == true || segment.noOp == true) {
+          segment.resolveLambdaTraceData();
+        }
       }
 
       return segment;

packages/core/lib/env/aws_lambda.js

@@ -36,7 +36,13 @@ module.exports.init = function init() {
 
   var namespace = contextUtils.getNamespace();
   namespace.enter(namespace.createContext());
-  contextUtils.setSegment(facadeSegment());
+
+  if (LambdaUtils.validTraceData(process.env._X_AMZN_TRACE_ID)) {
+    contextUtils.setSegment(facadeSegment());
+  }
+  else {
+    contextUtils.setSegment(noOpSegment());
+  }
 };
 
 var facadeSegment = function facadeSegment() {
@@ -109,3 +115,74 @@ 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.');
+    }
+  };
+
+  // Test for valid trace data during SDK startup. It's likely we're still in the cold-start portion of the
+  // code at this point and a valid trace header has not been set
+  if (LambdaUtils.validTraceData(xAmznTraceId)) {
+    if (LambdaUtils.populateTraceData(segment, xAmznTraceId)) {
+      xAmznTraceIdPrev = xAmznTraceId;
+    }
+  }
+
+  return segment;
+};

packages/core/lib/patchers/aws_p.js

@@ -88,8 +88,11 @@ function captureAWSRequest(req) {
   const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
 
   var buildListener = function(req) {
-    let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id +
-      ';Sampled=' + (subsegment.notTraced ? '0' : '1');
+    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');
+    }
     if (data != null) {
       for (const [key, value] of Object.entries(data)) {
         traceHeader += ';' + key +'=' + value;
@@ -99,6 +102,9 @@ function captureAWSRequest(req) {
   };
 
   var completeListener = function(res) {
+    if (subsegment == null) {
+      return
+    }
     subsegment.addAttribute('namespace', 'aws');
     subsegment.addAttribute('aws', new Aws(res, subsegment.name));
 

packages/core/lib/utils.js

@@ -168,6 +168,8 @@ 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 {