feat(tracer): add isTraceSampled method (#1435)

Author: icholy

Diff for: packages/commons/src/config/EnvironmentVariablesService.ts

@@ -53,11 +53,23 @@ class EnvironmentVariablesService extends ConfigService {
    * @returns {string}
    */
   public getXrayTraceId(): string | undefined {
-    const xRayTraceId = this.get(this.xRayTraceIdVariable);
+    const xRayTraceData = this.getXrayTraceData();
 
-    if (xRayTraceId === '') return undefined;
+    return xRayTraceData?.Root;
+  }
 
-    return xRayTraceId.split(';')[0].replace('Root=', '');
+  /**
+   * It returns true if the Sampled flag is set in the _X_AMZN_TRACE_ID environment variable.
+   * 
+   * The AWS X-Ray Trace data available in the environment variable has this format:
+   * `Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1`,
+   *
+   * @returns {boolean}
+   */
+  public getXrayTraceSampled(): boolean {
+    const xRayTraceData = this.getXrayTraceData();
+
+    return xRayTraceData?.Sampled === '1';
   }
 
   /**
@@ -72,6 +84,28 @@ class EnvironmentVariablesService extends ConfigService {
     return truthyValues.includes(value.toLowerCase());
   }
 
+  /**
+   * It parses the key/value data present in the _X_AMZN_TRACE_ID environment variable
+   * and returns it as an object when available.
+   */
+  private getXrayTraceData(): Record<string, string> | undefined {
+    const xRayTraceEnv = this.get(this.xRayTraceIdVariable);
+
+    if (xRayTraceEnv === '') return undefined;
+
+    if (!xRayTraceEnv.includes('=')) return { Root: xRayTraceEnv };
+
+    const xRayTraceData: Record<string, string> = {};
+
+    xRayTraceEnv.split(';').forEach((field) => {
+
+      const [ key, value ] = field.split('=');
+
+      xRayTraceData[key] = value;
+    });
+
+    return xRayTraceData;
+  }
 }
 
 export {

Diff for: packages/commons/tests/unit/config/EnvironmentVariablesService.test.ts

@@ -110,6 +110,48 @@ describe('Class: EnvironmentVariablesService', () => {
 
   });
 
+  describe('Method: getXrayTraceSampled', () => {
+    
+    test('It returns true if the Sampled flag is set in the _X_AMZN_TRACE_ID environment variable', () => {
+
+      // Prepare
+      process.env._X_AMZN_TRACE_ID = 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047;Sampled=1';
+      const service = new EnvironmentVariablesService();
+
+      // Act
+      const value = service.getXrayTraceSampled();
+
+      // Assess
+      expect(value).toEqual(true);
+    });
+
+    test('It returns false if the Sampled flag is not set in the _X_AMZN_TRACE_ID environment variable', () => {
+
+      // Prepare
+      process.env._X_AMZN_TRACE_ID = 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=557abcec3ee5a047';
+      const service = new EnvironmentVariablesService();
+
+      // Act
+      const value = service.getXrayTraceSampled();
+
+      // Assess
+      expect(value).toEqual(false);
+    });
+
+    it('It returns false when no _X_AMZN_TRACE_ID environment variable is present', () => {
+
+      // Prepare
+      delete process.env._X_AMZN_TRACE_ID;
+      const service = new EnvironmentVariablesService();
+
+      // Act
+      const value = service.getXrayTraceSampled();
+
+      // Assess
+      expect(value).toEqual(false);
+    });
+  });
+
   describe('Method: isValueTrue', () => {
 
     const valuesToTest: Array<Array<string | boolean>> = [

Diff for: packages/tracer/src/Tracer.ts

@@ -508,6 +508,8 @@ class Tracer extends Utility implements TracerInterface {
    * @returns string - The root X-Ray trace id.
    */
   public getRootXrayTraceId(): string | undefined {
+    if (!this.isTracingEnabled()) return undefined;
+
     return this.envVarsService.getXrayTraceId();
   }
 
@@ -547,6 +549,21 @@ class Tracer extends Utility implements TracerInterface {
     return segment;
   }
 
+  /**
+   * Get the current value of the AWS X-Ray Sampled flag.
+   * 
+   * Utility method that returns the current AWS X-Ray Sampled flag.
+   *
+   * @see https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-traces
+   * 
+   * @returns boolean - `true` if the trace is sampled, `false` if tracing is disabled or the trace is not sampled.
+   */
+  public isTraceSampled(): boolean {
+    if (!this.isTracingEnabled()) return false;
+
+    return this.envVarsService.getXrayTraceSampled();
+  }
+
   /**
    * Get the current value of the `tracingEnabled` property.
    * 

Diff for: packages/tracer/src/TracerInterface.ts

@@ -13,6 +13,7 @@ interface TracerInterface {
   captureMethod(options?: CaptureMethodOptions): MethodDecorator
   getSegment(): Segment | Subsegment | undefined
   getRootXrayTraceId(): string | undefined
+  isTraceSampled(): boolean
   isTracingEnabled(): boolean
   putAnnotation: (key: string, value: string | number | boolean) => void
   putMetadata: (key: string, value: unknown, namespace?: string | undefined) => void

Diff for: packages/tracer/tests/unit/Tracer.test.ts

@@ -277,6 +277,48 @@ describe('Class: Tracer', () => {
 
     });
 
+    test('when called and Tracer is disabled, it returns undefined', () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer({ enabled: false });
+
+      // Act
+      const xRayTraceId = tracer.getRootXrayTraceId();
+
+      // Assess
+      expect(xRayTraceId).toBe(undefined);
+    });
+
+  });
+
+  describe('Method: isTraceSampled', () => {
+
+    test('when called, it returns true if the Sampled flag is set', () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer();
+
+      // Act
+      const xRayTraceSampled = tracer.isTraceSampled();
+
+      // Assess
+      expect(xRayTraceSampled).toBe(false);
+
+    });
+
+    test('when called and Trace is disabled, it returns false', () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer({ enabled: false });
+
+      // Act
+      const xRayTraceSampled = tracer.isTraceSampled();
+
+      // Assess
+      expect(xRayTraceSampled).toBe(false);
+
+    });
+
   });
 
   describe('Method: getSegment', () => {