fix: support new x ray format nodejs (#519)

* fix: support parsing new aws x-ray format * fix: better error msgs in case of invalid aws x-ray trace id * test: more coverage * test: edge cases for parsing x-ray trace id fields * ci: reduce test node 18 flakiness --------- Co-authored-by: Harel Moshe <[email protected]>
lumigo-io · Oct 16, 2024 · b754e4e · b754e4e
1 parent 13cf15f
commit b754e4e
Show file tree

Hide file tree

Showing 3 changed files with 159 additions and 11 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -125,7 +125,7 @@ jobs:
     environment:
       - TZ: Asia/Jerusalem
       - NODE_OPTIONS: --max_old_space_size=1500
-    resource_class: medium+
+    resource_class: large
     working_directory: ~/lumigo-node
     steps:
       - run:

diff --git a/src/utils.test.js b/src/utils.test.js
@@ -132,6 +132,80 @@ describe('utils', () => {
       Sampled: '1',
       transactionId: '613d623b633d643b653d6661',
     });
+
+    const fields = utils.getTraceId(
+      'Root=1-670d0060-1b85fdcd75ed1c2557382245;Lineage=1:aba0be3a:0'
+    );
+    expect(fields.Root).toEqual('1-670d0060-1b85fdcd75ed1c2557382245');
+    expect(fields.transactionId).toEqual('1b85fdcd75ed1c2557382245');
+    expect(fields.Lineage).toEqual('1:aba0be3a:0');
+    expect(fields.Parent).toBeTruthy();
+
+    // Invalid root value, the new parser should fail and then be handled by the legacy parser
+    const invalidRootTraceId = 'Root=6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1';
+    expect(utils.getTraceId(invalidRootTraceId)).toEqual({
+      Parent: '59fa1aeb03c2ec1f',
+      Root: '6ac46730d346cad0e53f89d0',
+      Sampled: '1',
+      // Note: the current implementation splits the root by `-` and uses the third part as the transactionId. If there isn't a third part it sets it as undefined.
+      transactionId: undefined,
+    });
+
+    // Invalid root value, the new parser should fail and then be handled by the legacy parser
+    const invalidAndShortTraceId = 'Root=6ac46730d346cad0e53f89d0';
+    expect(utils.getTraceId(invalidAndShortTraceId)).toEqual({
+      // Static values generated based on the x-ray trace id value
+      Parent: '526f6f743d36616334363733',
+      Root: '26f6f74',
+      Sampled: '1',
+      transactionId: '526f6f743d36616334363733',
+    });
+
+    // The x-ray value is too long, too many fields. We should parse it as usual
+    const longXrayId = [
+      'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1',
+      ...Array(1000)
+        .fill(0)
+        .map((_, i) => `${i}=${i}`),
+    ].join(';');
+    expect(utils.getTraceId(longXrayId)).toEqual({
+      Parent: '59fa1aeb03c2ec1f',
+      Root: '1-5b1d2450-6ac46730d346cad0e53f89d0',
+      Sampled: '1',
+      transactionId: '6ac46730d346cad0e53f89d0',
+    });
+  });
+
+  test.each`
+    xrayTraceId                                                                                          | expected
+    ${'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1'}                      | ${{ Root: '1-5b1d2450-6ac46730d346cad0e53f89d0', Parent: '59fa1aeb03c2ec1f', Sampled: '1' }}
+    ${'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f'}                                | ${{ Root: '1-5b1d2450-6ac46730d346cad0e53f89d0', Parent: '59fa1aeb03c2ec1f' }}
+    ${'Root=1-5b1d2450-6ac46730d346cad0e53f89d0'}                                                        | ${{ Root: '1-5b1d2450-6ac46730d346cad0e53f89d0' }}
+    ${'Root=1-670d0060-1b85fdcd75ed1c2557382245;Lineage=1:aba0be3a:0'}                                   | ${{ Root: '1-670d0060-1b85fdcd75ed1c2557382245', Lineage: '1:aba0be3a:0' }}
+    ${'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1;Lineage=1:aba0be3a:0'} | ${{ Root: '1-5b1d2450-6ac46730d346cad0e53f89d0', Parent: '59fa1aeb03c2ec1f', Sampled: '1', Lineage: '1:aba0be3a:0' }}
+    ${';;;;;;;;;'}                                                                                       | ${{}}
+    ${''}                                                                                                | ${{}}
+    ${'========='}                                                                                       | ${{}}
+    ${'=;=;=;=;'}                                                                                        | ${{}}
+    ${';;;;;===='}                                                                                       | ${{}}
+  `('splitXrayTraceIdToFields', ({ xrayTraceId, expected }) => {
+    expect(utils.splitXrayTraceIdToFields(xrayTraceId)).toEqual(expected);
+  });
+
+  test('splitXrayTraceIdToFields - too many fields', () => {
+    let longXrayId = Array(1000)
+      .fill(0)
+      .map((_, i) => `${i}=${i}`)
+      .join(';');
+
+    const expectedFields = {};
+    for (let i = 0; i < 100; i++) {
+      expectedFields[`${i}`] = `${i}`;
+    }
+
+    const fields = utils.splitXrayTraceIdToFields(longXrayId);
+
+    expect(fields).toEqual(expectedFields);
   });
 
   test('isObject', () => {
@@ -144,17 +218,17 @@ describe('utils', () => {
     expect(isObject({})).toEqual(true);
   });
 
-  test('getPatchedTraceId', () => {
-    const awsXAmznTraceId =
-      'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1';
-    const expectedRoot = 'Root=1';
-    const expectedSuffix = '6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1';
-
+  test.each`
+    awsXAmznTraceId                                                                                      | patchedTraceIdPrefix | patchedTraceIdSuffix
+    ${'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1'}                      | ${'Root=1'}          | ${'6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1'}
+    ${'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f'}                                | ${'Root=1'}          | ${'6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f'}
+    ${'Root=1-5b1d2450-6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1;Lineage=1:c01ac717:0'} | ${'Root=1'}          | ${'6ac46730d346cad0e53f89d0;Parent=59fa1aeb03c2ec1f;Sampled=1;Lineage=1:c01ac717:0'}
+  `('getPatchedTraceId', ({ awsXAmznTraceId, patchedTraceIdPrefix, patchedTraceIdSuffix }) => {
     const result = utils.getPatchedTraceId(awsXAmznTraceId);
 
     const [resultRoot, resultTime, resultSuffix] = result.split('-');
-    expect(resultRoot).toEqual(expectedRoot);
-    expect(resultSuffix).toEqual(expectedSuffix);
+    expect(resultRoot).toEqual(patchedTraceIdPrefix);
+    expect(resultSuffix).toEqual(patchedTraceIdSuffix);
 
     const timeDiff = Date.now() - 1000 * parseInt(resultTime, 16);
     expect(timeDiff).toBeGreaterThan(0);

diff --git a/src/utils.ts b/src/utils.ts
@@ -11,6 +11,14 @@ import { EdgeUrl } from './types/common/edgeTypes';
 import { CommonUtils } from '@lumigo/node-core';
 import { runOneTimeWrapper } from './utils/functionUtils';
 
+type xRayTraceIdFields = {
+  Root: String;
+  transactionId: String;
+  Parent: String;
+  Sampled?: String;
+  Lineage?: String;
+};
+
 export const getRandomId = CommonUtils.getRandomId;
 export const getRandomString = CommonUtils.getRandomString;
 export const md5Hash = CommonUtils.md5Hash;
@@ -97,7 +105,66 @@ export const getTracerInfo = (): { name: string; version: string } => {
   return { name, version };
 };
 
+const xRayFieldKeyValuePattern = /([^;=]+)=([^;=]+)/g;
+
+export const splitXrayTraceIdToFields = (
+  awsXAmznTraceId: string,
+  maxFields: number = 100
+): { [key: string]: string } => {
+  const matches = Array.from(awsXAmznTraceId.matchAll(xRayFieldKeyValuePattern)).slice(
+    0,
+    maxFields
+  );
+  return Object.fromEntries(matches.map((match) => [match[1], match[2]]));
+};
+
+export const getNewFormatTraceId = (awsXAmznTraceId: string): xRayTraceIdFields => {
+  const fields = splitXrayTraceIdToFields(awsXAmznTraceId);
+  const root = fields.Root;
+  if (!root) {
+    throw new Error(
+      `X-Ray trace ID is missing the Root field (_X_AMZN_TRACE_ID=${awsXAmznTraceId})`
+    );
+  }
+  const rootValueSplit = root.split('-');
+  if (rootValueSplit.length < 3) {
+    throw new Error(
+      `X-Ray trace ID Root field is not in the expected format (Root=${fields.Root}, _X_AMZN_TRACE_ID=${awsXAmznTraceId})`
+    );
+  }
+  const transactionId = fields.Root.split('-')[2];
+  // Note: we might not need to generate a Parent field if it's not present,
+  // but for now we'll keep it as is to minimize changes. The python tracer doesn't generate it FYI.
+  const { Parent: parent = getRandomString(16), Sampled: sampled, Lineage: lineage } = fields;
+
+  const parsedTraceId: xRayTraceIdFields = {
+    Root: root,
+    transactionId,
+    Parent: parent,
+  };
+  if (sampled) {
+    parsedTraceId.Sampled = sampled;
+  }
+  if (lineage) {
+    parsedTraceId.Lineage = lineage;
+  }
+
+  return parsedTraceId;
+};
+
 export const getTraceId = (awsXAmznTraceId) => {
+  try {
+    logger.debug(
+      `Parsing the _X_AMZN_TRACE_ID environment variable (_X_AMZN_TRACE_ID = ${awsXAmznTraceId})`
+    );
+    return getNewFormatTraceId(awsXAmznTraceId);
+  } catch (e) {
+    logger.warn(
+      `Failed parsing the _X_AMZN_TRACE_ID environment variable, falling back to legacy parsing implementation (_X_AMZN_TRACE_ID = ${awsXAmznTraceId})`,
+      e
+    );
+  }
+
   try {
     if (!awsXAmznTraceId) {
       throw new Error('Missing _X_AMZN_TRACE_ID environment variable.');
@@ -180,10 +247,17 @@ export const getTraceId = (awsXAmznTraceId) => {
 
 export const getPatchedTraceId = (awsXAmznTraceId): string => {
   // @ts-ignore
-  const { Root, Parent, Sampled, transactionId } = getTraceId(awsXAmznTraceId);
+  const { Root, Parent, Sampled, transactionId, Lineage } = getTraceId(awsXAmznTraceId);
   const rootArr = Root.split('-');
   const currentTime = Math.floor(Date.now() / 1000).toString(16);
-  return `Root=${rootArr[0]}-${currentTime}-${transactionId};Parent=${Parent};Sampled=${Sampled}`;
+  const patchedTraceId = [`Root=${rootArr[0]}-${currentTime}-${transactionId}`, `Parent=${Parent}`];
+  if (Sampled) {
+    patchedTraceId.push(`Sampled=${Sampled}`);
+  }
+  if (Lineage) {
+    patchedTraceId.push(`Lineage=${Lineage}`);
+  }
+  return patchedTraceId.join(';');
 };
 
 export const isPromise = (obj: any): boolean => typeof obj?.then === 'function';