diff --git a/packages/tracer/src/Tracer.ts b/packages/tracer/src/Tracer.ts index b5310aa8e3..359cf128a1 100644 --- a/packages/tracer/src/Tracer.ts +++ b/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 --git a/packages/tracer/src/config/EnvironmentVariablesService.ts b/packages/tracer/src/config/EnvironmentVariablesService.ts index b1b1fcb65b..4f48eeef2c 100644 --- a/packages/tracer/src/config/EnvironmentVariablesService.ts +++ b/packages/tracer/src/config/EnvironmentVariablesService.ts @@ -69,4 +69,6 @@ class EnvironmentVariablesService } } -export { EnvironmentVariablesService }; +const environmentVariablesService = new EnvironmentVariablesService(); + +export { EnvironmentVariablesService, environmentVariablesService }; diff --git a/packages/tracer/src/provider/ProviderService.ts b/packages/tracer/src/provider/ProviderService.ts index fe83e8087d..a080a32278 100644 --- a/packages/tracer/src/provider/ProviderService.ts +++ b/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 --git a/packages/tracer/tests/helpers/mockRequests.ts b/packages/tracer/tests/helpers/mockRequests.ts index ff2d930f04..4385d06be4 100644 --- a/packages/tracer/tests/helpers/mockRequests.ts +++ b/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 --git a/packages/tracer/tests/unit/ProviderService.test.ts b/packages/tracer/tests/unit/ProviderService.test.ts index dc6fd36064..04134bde59 100644 --- a/packages/tracer/tests/unit/ProviderService.test.ts +++ b/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/ + ) + ); + }); });