Allow lambda passthrough trace header

majanjua-amzn · Jun 7, 2024 · c49f12e · c49f12e
1 parent 1502470
commit c49f12e
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts
@@ -140,7 +140,9 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
     }
   }
 
-  args.request.headers['X-Amzn-Trace-Id'] = traceHeader;
+  if (!segment.noOp) {
+    args.request.headers['X-Amzn-Trace-Id'] = traceHeader;
+  }
 
   let res;
   try {

diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js
@@ -88,6 +88,9 @@ function captureAWSRequest(req) {
   const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
 
   var buildListener = function(req) {
+    if (parent.noOp) {
+      return;
+    }
     let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id +
       ';Sampled=' + (subsegment.notTraced ? '0' : '1');
     if (data != null) {

diff --git a/packages/core/lib/patchers/http_p.js b/packages/core/lib/patchers/http_p.js
@@ -130,8 +130,10 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) {
       options.headers = {};
     }
 
-    options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
-      ';Sampled=' + (subsegment.notTraced ? '0' : '1');
+    if (!parent.noOp) {
+      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,6 +168,10 @@ var utils = {
       if (!traceData) {
         traceData = {};
         logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format');
+      } else if (traceData.root && !traceData.parent && !traceData.sampled) {
+        // Lambda PassThrough only has root, treat as valid in this case and mark the segment
+        segment.noOp = true;
+        valid = true;
       } else if (!traceData.root || !traceData.parent || !traceData.sampled) {
         logger.getLogger().error('_X_AMZN_TRACE_ID is missing required information');
       } else {