fix(tracer): forward X-Amzn-Trace-Id header when instrumenting fetch (#3470)

Author: dreamorosi

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

@@ -19,7 +19,10 @@ import type {
 import type { Handler } from 'aws-lambda';
 import type { Segment, Subsegment } from 'aws-xray-sdk-core';
 import xraySdk from 'aws-xray-sdk-core';
-import { EnvironmentVariablesService } from './config/EnvironmentVariablesService.js';
+import {
+  type EnvironmentVariablesService,
+  environmentVariablesService,
+} from './config/EnvironmentVariablesService.js';
 import { ProviderService } from './provider/ProviderService.js';
 import type { ConfigServiceInterface } from './types/ConfigServiceInterface.js';
 import type { ProviderServiceInterface } from './types/ProviderService.js';
@@ -861,14 +864,6 @@ class Tracer extends Utility implements TracerInterface {
       : undefined;
   }
 
-  /**
-   * Set and initialize `envVarsService`.
-   * Used internally during initialization.
-   */
-  private setEnvVarsService(): void {
-    this.envVarsService = new EnvironmentVariablesService();
-  }
-
   /**
    * Method that reconciles the configuration passed with the environment variables.
    * Used internally during initialization.
@@ -879,7 +874,7 @@ class Tracer extends Utility implements TracerInterface {
     const { enabled, serviceName, captureHTTPsRequests, customConfigService } =
       options;
 
-    this.setEnvVarsService();
+    this.envVarsService = environmentVariablesService;
     this.setCustomConfigService(customConfigService);
     this.setTracingEnabled(enabled);
     this.setCaptureResponse();

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

@@ -69,4 +69,6 @@ class EnvironmentVariablesService
   }
 }
 
-export { EnvironmentVariablesService };
+const environmentVariablesService = new EnvironmentVariablesService();
+
+export { EnvironmentVariablesService, environmentVariablesService };

Diff for: packages/tracer/src/provider/ProviderService.ts

@@ -26,6 +26,7 @@ import http from 'node:http';
 import https from 'node:https';
 import { addUserAgentMiddleware } from '@aws-lambda-powertools/commons';
 import type { DiagnosticsChannel } from 'undici-types';
+import { environmentVariablesService } from '../config/EnvironmentVariablesService.js';
 import {
   findHeaderAndDecode,
   getRequestURL,
@@ -127,6 +128,13 @@ class ProviderService implements ProviderServiceInterface {
         );
         subsegment.addAttribute('namespace', 'remote');
 
+        // addHeader is not part of the type definition but it's available https://github.com/nodejs/undici/blob/main/docs/docs/api/DiagnosticsChannel.md#undicirequestcreate
+        // @ts-expect-error
+        request.addHeader(
+          'X-Amzn-Trace-Id',
+          `Root=${environmentVariablesService.getXrayTraceId()};Parent=${subsegment.id};Sampled=${subsegment.notTraced ? '0' : '1'}`
+        );
+
         (subsegment as HttpSubsegment).http = {
           request: {
             url: `${requestURL.protocol}//${requestURL.hostname}${requestURL.pathname}`,

Diff for: packages/tracer/tests/helpers/mockRequests.ts

@@ -1,5 +1,6 @@
 import { channel } from 'node:diagnostics_channel';
 import type { URL } from 'node:url';
+import { vi } from 'vitest';
 
 type MockFetchOptions = {
   origin?: string | URL;
@@ -31,7 +32,7 @@ const mockFetch = ({
   statusCode,
   headers,
   throwError,
-}: MockFetchOptions): void => {
+}: MockFetchOptions) => {
   const requestCreateChannel = channel('undici:request:create');
   const responseHeadersChannel = channel('undici:request:headers');
   const errorChannel = channel('undici:request:error');
@@ -40,6 +41,7 @@ const mockFetch = ({
     origin,
     method: method ?? 'GET',
     path,
+    addHeader: vi.fn(),
   };
 
   requestCreateChannel.publish({
@@ -70,6 +72,8 @@ const mockFetch = ({
       headers: encodedHeaders,
     },
   });
+
+  return request;
 };
 
 export { mockFetch };

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

@@ -364,7 +364,7 @@ describe('Class: ProviderService', () => {
 
       // Act
       provider.instrumentFetch();
-      mockFetch({
+      const mockRequest = mockFetch({
         origin: 'https://aws.amazon.com',
         path: '/blogs',
         headers: {
@@ -387,6 +387,12 @@ describe('Class: ProviderService', () => {
       });
       expect(subsegment.close).toHaveBeenCalledTimes(1);
       expect(provider.setSegment).toHaveBeenLastCalledWith(segment);
+      expect(mockRequest.addHeader).toHaveBeenLastCalledWith(
+        'X-Amzn-Trace-Id',
+        expect.stringMatching(
+          /Root=1-abcdef12-3456abcdef123456abcdef12;Parent=\S{16};Sampled=1/
+        )
+      );
     });
 
     it('excludes the content_length header when invalid or not found', async () => {
@@ -616,4 +622,39 @@ describe('Class: ProviderService', () => {
     expect(subsegment.close).toHaveBeenCalledTimes(1);
     expect(provider.setSegment).toHaveBeenLastCalledWith(segment);
   });
+
+  it('forwards the correct sampling decision in the request header', async () => {
+    // Prepare
+    const provider: ProviderService = new ProviderService();
+    const segment = new Subsegment('## dummySegment');
+    const subsegment = segment.addNewSubsegment('aws.amazon.com');
+    subsegment.notTraced = true;
+    vi.spyOn(segment, 'addNewSubsegment').mockImplementationOnce(
+      () => subsegment
+    );
+    vi.spyOn(provider, 'getSegment')
+      .mockImplementationOnce(() => segment)
+      .mockImplementationOnce(() => subsegment)
+      .mockImplementationOnce(() => subsegment);
+    vi.spyOn(subsegment, 'close');
+    vi.spyOn(provider, 'setSegment');
+
+    // Act
+    provider.instrumentFetch();
+    const mockRequest = mockFetch({
+      origin: 'https://aws.amazon.com',
+      path: '/blogs',
+      headers: {
+        'content-length': '100',
+      },
+    });
+
+    // Assess
+    expect(mockRequest.addHeader).toHaveBeenLastCalledWith(
+      'X-Amzn-Trace-Id',
+      expect.stringMatching(
+        /Root=1-abcdef12-3456abcdef123456abcdef12;Parent=\S{16};Sampled=0/
+      )
+    );
+  });
 });