Parse XRay header

thundra-io · Sep 12, 2019 · b9337a4 · b9337a4
1 parent e9b07ea
commit b9337a4
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 0 deletions.
diff --git a/src/Constants.ts b/src/Constants.ts
@@ -30,6 +30,7 @@ export const envVariableKeys = {
     THUNDRA_DISABLE_METRIC: 'thundra_agent_lambda_metric_disable',
     THUNDRA_DISABLE_LOG: 'thundra_agent_lambda_log_disable',
     THUNDRA_DISABLE_XRAY: 'thundra_agent_lambda_xray_disable',
+    _X_AMZN_TRACE_ID: '_X_AMZN_TRACE_ID',
 
     THUNDRA_LAMBDA_TRACE_REQUEST_SKIP: 'thundra_agent_lambda_trace_request_skip',
     THUNDRA_LAMBDA_TRACE_RESPONSE_SKIP: 'thundra_agent_lambda_trace_response_skip',

diff --git a/src/plugins/Invocation.ts b/src/plugins/Invocation.ts
@@ -67,6 +67,15 @@ class Invocation {
         this.invocationData.spanId = this.pluginContext.spanId;
         this.invocationData.traceId = this.pluginContext.traceId;
 
+        const xrayTraceInfo = Utils.getXRayTraceInfo();
+
+        if (xrayTraceInfo.traceID) {
+            this.invocationData.tags['aws.xray.trace.id'] = xrayTraceInfo.traceID;
+        }
+        if (xrayTraceInfo.segmentID) {
+            this.invocationData.tags['aws.xray.segment.id'] = xrayTraceInfo.segmentID;
+        }
+
         this.invocationData.tags['aws.lambda.memory_limit'] = this.pluginContext.maxMemory;
         this.invocationData.tags['aws.lambda.arn'] = originalContext.invokedFunctionArn;
         this.invocationData.tags['aws.account_no'] = Utils.getAWSAccountNo(originalContext.invokedFunctionArn);

diff --git a/src/plugins/utils/Utils.ts b/src/plugins/utils/Utils.ts
@@ -400,6 +400,32 @@ class Utils {
             return '';
         }
     }
+
+    static getXRayTraceInfo() {
+        let traceID: string = '';
+        let segmentID: string = '';
+        const xrayTraceHeader: string = Utils.getConfiguration(envVariableKeys._X_AMZN_TRACE_ID);
+        if (xrayTraceHeader) {
+            for (const traceHeaderPart of xrayTraceHeader.split(';')) {
+                const traceInfo = traceHeaderPart.split('=');
+                if (traceInfo.length !== 2) {
+                    continue;
+                }
+                const [traceInfoKey, traceInfoVal] = traceInfo;
+
+                if (traceInfoKey === 'Root') {
+                    traceID = traceInfoVal;
+                } else if (traceInfoKey === 'Parent') {
+                    segmentID = traceInfoVal;
+                }
+            }
+        }
+
+        return {
+            traceID,
+            segmentID,
+        };
+    }
 }
 
 export default Utils;