Merge pull request #46 from DataDog/darcy.rayner/use-xray-env-var

Darcy.rayner/use xray env var
DataDog · Jan 21, 2020 · 61b7890 · 61b7890
2 parents 884334a + 0fd8f63
commit 61b7890
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 64 deletions.
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "datadog-lambda-js",
-  "version": "0.12.0",
+  "version": "0.13.0",
   "description": "Lambda client library that supports hybrid tracing in node js",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

diff --git a/src/trace/constants.ts b/src/trace/constants.ts
@@ -16,3 +16,4 @@ export const xraySubsegmentName = "datadog-metadata";
 export const xraySubsegmentKey = "trace";
 export const xrayBaggageSubsegmentKey = "root_span_metadata";
 export const xraySubsegmentNamespace = "datadog";
+export const xrayTraceEnvVar = "_X_AMZN_TRACE_ID";
diff --git a/src/trace/context.spec.ts b/src/trace/context.spec.ts
@@ -22,12 +22,6 @@ jest.mock("aws-xray-sdk-core", () => {
         throw Error("Unimplemented");
       }
     },
-    getSegment: () => {
-      if (currentSegment === undefined) {
-        throw Error("Empty");
-      }
-      return currentSegment;
-    },
   };
 });
 
@@ -127,38 +121,34 @@ describe("convertTraceContext", () => {
 });
 
 describe("readTraceContextFromXray", () => {
-  it("will parse a trace context from the xray", () => {
-    currentSegment = {
-      id: "0b11cc4230d3e09e",
-      trace_id: "1-5ce31dc2-2c779014b90ce44db5e03875",
-    };
-
-    const traceContext = readTraceContextFromXray();
-    expect(traceContext).toEqual({
-      parentID: "797643193680388254",
-      sampleMode: SampleMode.USER_KEEP,
-      traceID: "4110911582297405557",
-      source: Source.Xray,
-    });
-  });
-  it("will ignore a trace context from the xray, when sampling is turned off", () => {
-    currentSegment = {
-      id: "0b11cc4230d3e09e",
-      notTraced: true,
-      trace_id: "1-5ce31dc2-2c779014b90ce44db5e03875",
-    };
-
-    const traceContext = readTraceContextFromXray();
-    expect(traceContext).toEqual({
-      parentID: "797643193680388254",
-      sampleMode: SampleMode.USER_REJECT,
-      traceID: "4110911582297405557",
-      source: Source.Xray,
+  afterEach(() => {
+    process.env["_X_AMZN_TRACE_ID"] = undefined;
+  });
+  it("returns a trace context from a valid env var", () => {
+    process.env["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
+    const context = readTraceContextFromXray();
+    expect(context).toEqual({
+      parentID: "10713633173203262661",
+      sampleMode: 2,
+      source: "xray",
+      traceID: "3995693151288333088",
     });
   });
-  it("returns undefined when trace header isn't in environment", () => {
-    const traceContext = readTraceContextFromXray();
-    expect(traceContext).toBeUndefined();
+  it("returns undefined when given an invalid env var", () => {
+    const badCases = [
+      "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5",
+      "Root=1-5e272390-8c398be037738dc042009320",
+      "1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1",
+      "Root=1-5e272390-8c398be037738dc042009320;94ae789b969f1cc5;Sampled=1",
+      "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;1",
+      "Root=a;Parent=94ae789b969f1cc5;Sampled=1",
+      "Root=1-5e272390-8c398be037738dc042009320;Parent=b;Sampled=1",
+      undefined,
+    ];
+    for (const badCase of badCases) {
+      process.env["_X_AMZN_TRACE_ID"] = badCase;
+      expect(readTraceContextFromXray()).toBeUndefined();
+    }
   });
 });
 
@@ -361,11 +351,11 @@ describe("readStepFunctionContextFromEvent", () => {
 });
 
 describe("extractTraceContext", () => {
+  afterEach(() => {
+    process.env["_X_AMZN_TRACE_ID"] = undefined;
+  });
   it("returns trace read from header as highest priority", () => {
-    currentSegment = {
-      parent_id: "0b11cc4230d3e09e",
-      trace_id: "1-5ce31dc2-2c779014b90ce44db5e03875",
-    };
+    process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
     const result = extractTraceContext({
       headers: {
@@ -382,10 +372,7 @@ describe("extractTraceContext", () => {
     });
   });
   it("returns trace read from env if no headers present", () => {
-    currentSegment = {
-      id: "0b11cc4230d3e09e",
-      trace_id: "1-5ce31dc2-2c779014b90ce44db5e03875",
-    };
+    process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
     const result = extractTraceContext({});
     expect(result).toEqual({
@@ -396,10 +383,7 @@ describe("extractTraceContext", () => {
     });
   });
   it("returns trace read from env if no headers present", () => {
-    currentSegment = {
-      id: "0b11cc4230d3e09e",
-      trace_id: "1-5ce31dc2-2c779014b90ce44db5e03875",
-    };
+    process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
     const result = extractTraceContext({});
     expect(result).toEqual({

diff --git a/src/trace/context.ts b/src/trace/context.ts
@@ -1,4 +1,4 @@
-import { captureFunc, getSegment } from "aws-xray-sdk-core";
+import { captureFunc } from "aws-xray-sdk-core";
 import { BigNumber } from "bignumber.js";
 
 import { logDebug, logError } from "../utils";
@@ -12,6 +12,7 @@ import {
   xraySubsegmentKey,
   xraySubsegmentName,
   xraySubsegmentNamespace,
+  xrayTraceEnvVar,
 } from "./constants";
 
 export interface XRayTraceHeader {
@@ -117,20 +118,55 @@ export function readTraceFromEvent(event: any): TraceContext | undefined {
   };
 }
 
-export function readTraceContextFromXray() {
-  try {
-    const segment = getSegment();
-    logDebug(`Setting X-Ray parent trace to segment with ${segment.id} and trace ${segment.trace_id}`);
-    const traceHeader = {
-      parentID: segment.id,
-      sampled: segment.notTraced ? 0 : 1,
-      traceID: segment.trace_id,
-    };
-    return convertTraceContext(traceHeader);
-  } catch (error) {
-    logError("couldn't read xray trace header", { innerError: error });
-  }
-  return undefined;
+export function readTraceContextFromXray(): TraceContext | undefined {
+  const header = process.env[xrayTraceEnvVar];
+  if (header === undefined) {
+    logError("couldn't read xray trace header from env");
+    return;
+  }
+  const context = parseTraceContextHeader(header);
+  if (context === undefined) {
+    logError("couldn't read xray trace context from env, variable had invalid format");
+  } else {
+    logDebug("read trace context from environment", context);
+  }
+  return context;
+}
+
+export function parseTraceContextHeader(header: string): TraceContext | undefined {
+  // Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
+  logDebug(`Reading trace context from env var ${header}`);
+  const [root, parent, sampled] = header.split(";");
+  if (parent === undefined || sampled === undefined) {
+    return;
+  }
+  const [, rawTraceID] = root.split("=");
+  if (rawTraceID === undefined) {
+    return;
+  }
+  const traceID = convertToAPMTraceID(rawTraceID);
+  if (traceID === undefined) {
+    return;
+  }
+  const [, rawParentID] = parent.split("=");
+  if (rawParentID === undefined) {
+    return;
+  }
+  const parentID = convertToAPMParentID(rawParentID);
+  if (parentID === undefined) {
+    return;
+  }
+  const [, rawSampled] = sampled.split("=");
+  if (rawSampled === undefined) {
+    return;
+  }
+  const sampleMode = convertToSampleMode(parseInt(rawSampled, 10));
+  return {
+    parentID,
+    sampleMode,
+    source: Source.Xray,
+    traceID,
+  };
 }
 
 export function readStepFunctionContextFromEvent(event: any): StepFunctionContext | undefined {