feat(node): Add tracing without performance to Node Undici

fix type

force sampling decision

log this out

remove console log

Author: AbhiPrasad

packages/node/src/integrations/undici/index.ts

@@ -1,8 +1,10 @@
 import type { Hub } from '@sentry/core';
-import type { EventProcessor, Integration } from '@sentry/types';
+import { getDynamicSamplingContextFromClient } from '@sentry/core';
+import type { EventProcessor, Integration, Span } from '@sentry/types';
 import {
   dynamicRequire,
   dynamicSamplingContextToSentryBaggageHeader,
+  generateSentryTraceHeader,
   getSanitizedUrlString,
   parseUrl,
   stringMatchesSomePattern,
@@ -12,7 +14,13 @@ import { LRUMap } from 'lru_map';
 import type { NodeClient } from '../../client';
 import { NODE_VERSION } from '../../nodeVersion';
 import { isSentryRequest } from '../utils/http';
-import type { DiagnosticsChannel, RequestCreateMessage, RequestEndMessage, RequestErrorMessage } from './types';
+import type {
+  DiagnosticsChannel,
+  RequestCreateMessage,
+  RequestEndMessage,
+  RequestErrorMessage,
+  RequestWithSentry,
+} from './types';
 
 export enum ChannelName {
   // https://github.com/nodejs/undici/blob/e6fc80f809d1217814c044f52ed40ef13f21e43c/docs/api/DiagnosticsChannel.md#undicirequestcreate
@@ -124,63 +132,53 @@ export class Undici implements Integration {
       const { request } = message as RequestCreateMessage;
 
       const stringUrl = request.origin ? request.origin.toString() + request.path : request.path;
-      const url = parseUrl(stringUrl);
 
-      if (isSentryRequest(stringUrl) || request.__sentry__ !== undefined) {
+      if (isSentryRequest(stringUrl) || request.__sentry_span__ !== undefined) {
         return;
       }
 
       const client = hub.getClient<NodeClient>();
+      if (!client) {
+        return;
+      }
+
+      const clientOptions = client.getOptions();
       const scope = hub.getScope();
 
-      const activeSpan = scope.getSpan();
-
-      if (activeSpan && client) {
-        const clientOptions = client.getOptions();
-
-        if (shouldCreateSpan(stringUrl)) {
-          const method = request.method || 'GET';
-          const data: Record<string, unknown> = {
-            'http.method': method,
-          };
-          if (url.search) {
-            data['http.query'] = url.search;
-          }
-          if (url.hash) {
-            data['http.fragment'] = url.hash;
-          }
-          const span = activeSpan.startChild({
-            op: 'http.client',
-            description: `${method} ${getSanitizedUrlString(url)}`,
-            data,
-          });
-          request.__sentry__ = span;
-
-          const shouldAttachTraceData = (url: string): boolean => {
-            if (clientOptions.tracePropagationTargets === undefined) {
-              return true;
-            }
-
-            const cachedDecision = this._headersUrlMap.get(url);
-            if (cachedDecision !== undefined) {
-              return cachedDecision;
-            }
-
-            const decision = stringMatchesSomePattern(url, clientOptions.tracePropagationTargets);
-            this._headersUrlMap.set(url, decision);
-            return decision;
-          };
-
-          if (shouldAttachTraceData(stringUrl)) {
-            request.addHeader('sentry-trace', span.toTraceparent());
-            if (span.transaction) {
-              const dynamicSamplingContext = span.transaction.getDynamicSamplingContext();
-              const sentryBaggageHeader = dynamicSamplingContextToSentryBaggageHeader(dynamicSamplingContext);
-              if (sentryBaggageHeader) {
-                request.addHeader('baggage', sentryBaggageHeader);
-              }
-            }
-          }
+      const parentSpan = scope.getSpan();
+
+      const span = shouldCreateSpan(stringUrl) ? createRequestSpan(parentSpan, request, stringUrl) : undefined;
+      if (span) {
+        request.__sentry_span__ = span;
+      }
+
+      const shouldAttachTraceData = (url: string): boolean => {
+        if (clientOptions.tracePropagationTargets === undefined) {
+          return true;
+        }
+
+        const cachedDecision = this._headersUrlMap.get(url);
+        if (cachedDecision !== undefined) {
+          return cachedDecision;
+        }
+
+        const decision = stringMatchesSomePattern(url, clientOptions.tracePropagationTargets);
+        this._headersUrlMap.set(url, decision);
+        return decision;
+      };
+
+      if (shouldAttachTraceData(stringUrl)) {
+        if (span) {
+          const dynamicSamplingContext = span?.transaction?.getDynamicSamplingContext();
+          const sentryBaggageHeader = dynamicSamplingContextToSentryBaggageHeader(dynamicSamplingContext);
+
+          setHeadersOnRequest(request, span.toTraceparent(), sentryBaggageHeader);
+        } else {
+          const { traceId, sampled, dsc } = scope.getPropagationContext();
+          const sentryTrace = generateSentryTraceHeader(traceId, undefined, sampled);
+          const dynamicSamplingContext = dsc || getDynamicSamplingContextFromClient(traceId, client, scope);
+          const sentryBaggageHeader = dynamicSamplingContextToSentryBaggageHeader(dynamicSamplingContext);
+          setHeadersOnRequest(request, sentryTrace, sentryBaggageHeader);
         }
       }
     });
@@ -199,7 +197,7 @@ export class Undici implements Integration {
         return;
       }
 
-      const span = request.__sentry__;
+      const span = request.__sentry_span__;
       if (span) {
         span.setHttpStatus(response.statusCode);
         span.finish();
@@ -239,7 +237,7 @@ export class Undici implements Integration {
         return;
       }
 
-      const span = request.__sentry__;
+      const span = request.__sentry_span__;
       if (span) {
         span.setStatus('internal_error');
         span.finish();
@@ -265,3 +263,44 @@ export class Undici implements Integration {
     });
   }
 }
+
+function setHeadersOnRequest(
+  request: RequestWithSentry,
+  sentryTrace: string,
+  sentryBaggageHeader: string | undefined,
+): void {
+  if (request.__sentry_has_headers__) {
+    return;
+  }
+
+  request.addHeader('sentry-trace', sentryTrace);
+  if (sentryBaggageHeader) {
+    request.addHeader('baggage', sentryBaggageHeader);
+  }
+
+  request.__sentry_has_headers__ = true;
+}
+
+function createRequestSpan(
+  activeSpan: Span | undefined,
+  request: RequestWithSentry,
+  stringUrl: string,
+): Span | undefined {
+  const url = parseUrl(stringUrl);
+
+  const method = request.method || 'GET';
+  const data: Record<string, unknown> = {
+    'http.method': method,
+  };
+  if (url.search) {
+    data['http.query'] = url.search;
+  }
+  if (url.hash) {
+    data['http.fragment'] = url.hash;
+  }
+  return activeSpan?.startChild({
+    op: 'http.client',
+    description: `${method} ${getSanitizedUrlString(url)}`,
+    data,
+  });
+}

packages/node/src/integrations/undici/types.ts

@@ -234,7 +234,8 @@ export interface UndiciResponse {
 }
 
 export interface RequestWithSentry extends UndiciRequest {
-  __sentry__?: Span;
+  __sentry_span__?: Span;
+  __sentry_has_headers__?: boolean;
 }
 
 export interface RequestCreateMessage {

packages/node/test/integrations/undici.test.ts

@@ -28,7 +28,7 @@ beforeAll(async () => {
 
 const DEFAULT_OPTIONS = getDefaultNodeClientOptions({
   dsn: SENTRY_DSN,
-  tracesSampleRate: 1,
+  tracesSampler: () => true,
   integrations: [new Undici()],
 });
 
@@ -193,7 +193,7 @@ conditionalTest({ min: 16 })('Undici integration', () => {
     undoPatch();
   });
 
-  it('attaches the sentry trace and baggage headers', async () => {
+  it('attaches the sentry trace and baggage headers if there is an active span', async () => {
     const transaction = hub.startTransaction({ name: 'test-transaction' }) as Transaction;
     hub.getScope().setSpan(transaction);
 
@@ -208,22 +208,44 @@ conditionalTest({ min: 16 })('Undici integration', () => {
     );
   });
 
-  it('does not attach headers if `shouldCreateSpanForRequest` does not create a span', async () => {
+  it('attaches the sentry trace and baggage headers if there is no active span', async () => {
+    await fetch('http://localhost:18099', { method: 'POST' });
+
+    const propagationContext = hub.getScope().getPropagationContext();
+
+    expect(requestHeaders['sentry-trace'].includes(propagationContext.traceId)).toBe(true);
+    expect(requestHeaders['baggage']).toEqual(
+      `sentry-environment=production,sentry-public_key=0,sentry-trace_id=${propagationContext.traceId}`,
+    );
+  });
+
+  it('attaches headers if `shouldCreateSpanForRequest` does not create a span using propagation context', async () => {
     const transaction = hub.startTransaction({ name: 'test-transaction' }) as Transaction;
-    hub.getScope().setSpan(transaction);
+    const scope = hub.getScope();
+    const propagationContext = scope.getPropagationContext();
+
+    scope.setSpan(transaction);
 
     const undoPatch = patchUndici(hub, { shouldCreateSpanForRequest: url => url.includes('yes') });
 
     await fetch('http://localhost:18099/no', { method: 'POST' });
 
-    expect(requestHeaders['sentry-trace']).toBeUndefined();
-    expect(requestHeaders['baggage']).toBeUndefined();
+    expect(requestHeaders['sentry-trace']).toBeDefined();
+    expect(requestHeaders['baggage']).toBeDefined();
+
+    expect(requestHeaders['sentry-trace'].includes(propagationContext.traceId)).toBe(true);
+    const firstSpanId = requestHeaders['sentry-trace'].split('-')[1];
 
     await fetch('http://localhost:18099/yes', { method: 'POST' });
 
     expect(requestHeaders['sentry-trace']).toBeDefined();
     expect(requestHeaders['baggage']).toBeDefined();
 
+    expect(requestHeaders['sentry-trace'].includes(propagationContext.traceId)).toBe(false);
+
+    const secondSpanId = requestHeaders['sentry-trace'].split('-')[1];
+    expect(firstSpanId).not.toBe(secondSpanId);
+
     undoPatch();
   });
 
@@ -351,10 +373,10 @@ function setupTestServer() {
     res.end();
 
     // also terminate socket because keepalive hangs connection a bit
-    res.connection.end();
+    res.connection?.end();
   });
 
-  testServer.listen(18099, 'localhost');
+  testServer?.listen(18099, 'localhost');
 
   return new Promise(resolve => {
     testServer?.on('listening', resolve);