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

Author: AbhiPrasad

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

@@ -1,8 +1,9 @@
-import type { Hub } from '@sentry/core';
-import type { EventProcessor, Integration } from '@sentry/types';
+import { getCurrentHub, getDynamicSamplingContextFromClient } from '@sentry/core';
+import type { EventProcessor, Integration, Span } from '@sentry/types';
 import {
   dynamicRequire,
   dynamicSamplingContextToSentryBaggageHeader,
+  generateSentryTraceHeader,
   getSanitizedUrlString,
   parseUrl,
   stringMatchesSomePattern,
@@ -12,7 +13,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
@@ -81,7 +88,7 @@ export class Undici implements Integration {
   /**
    * @inheritDoc
    */
-  public setupOnce(_addGlobalEventProcessor: (callback: EventProcessor) => void, getCurrentHub: () => Hub): void {
+  public setupOnce(_addGlobalEventProcessor: (callback: EventProcessor) => void): void {
     // Requires Node 16+ to use the diagnostics_channel API.
     if (NODE_VERSION.major && NODE_VERSION.major < 16) {
       return;
@@ -99,169 +106,205 @@ export class Undici implements Integration {
       return;
     }
 
-    const shouldCreateSpan = (url: string): boolean => {
-      if (this._options.shouldCreateSpanForRequest === undefined) {
+    // https://github.com/nodejs/undici/blob/e6fc80f809d1217814c044f52ed40ef13f21e43c/docs/api/DiagnosticsChannel.md
+    ds.subscribe(ChannelName.RequestCreate, this._onRequestCreate);
+    ds.subscribe(ChannelName.RequestEnd, this._onRequestEnd);
+    ds.subscribe(ChannelName.RequestError, this._onRequestError);
+  }
+
+  /** Helper that wraps shouldCreateSpanForRequest option */
+  private _shouldCreateSpan(url: string): boolean {
+    if (this._options.shouldCreateSpanForRequest === undefined) {
+      return true;
+    }
+
+    const cachedDecision = this._createSpanUrlMap.get(url);
+    if (cachedDecision !== undefined) {
+      return cachedDecision;
+    }
+
+    const decision = this._options.shouldCreateSpanForRequest(url);
+    this._createSpanUrlMap.set(url, decision);
+    return decision;
+  }
+
+  private _onRequestCreate = (message: unknown): void => {
+    const hub = getCurrentHub();
+    if (!hub.getIntegration(Undici)) {
+      return;
+    }
+
+    const { request } = message as RequestCreateMessage;
+
+    const stringUrl = request.origin ? request.origin.toString() + request.path : request.path;
+
+    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 parentSpan = scope.getSpan();
+
+    const span = this._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._createSpanUrlMap.get(url);
+      const cachedDecision = this._headersUrlMap.get(url);
       if (cachedDecision !== undefined) {
         return cachedDecision;
       }
 
-      const decision = this._options.shouldCreateSpanForRequest(url);
-      this._createSpanUrlMap.set(url, decision);
+      const decision = stringMatchesSomePattern(url, clientOptions.tracePropagationTargets);
+      this._headersUrlMap.set(url, decision);
       return decision;
     };
 
-    // https://github.com/nodejs/undici/blob/e6fc80f809d1217814c044f52ed40ef13f21e43c/docs/api/DiagnosticsChannel.md
-    ds.subscribe(ChannelName.RequestCreate, message => {
-      const hub = getCurrentHub();
-      if (!hub.getIntegration(Undici)) {
-        return;
+    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);
       }
+    }
+  };
 
-      const { request } = message as RequestCreateMessage;
+  private _onRequestEnd = (message: unknown): void => {
+    const hub = getCurrentHub();
+    if (!hub.getIntegration(Undici)) {
+      return;
+    }
 
-      const stringUrl = request.origin ? request.origin.toString() + request.path : request.path;
-      const url = parseUrl(stringUrl);
+    const { request, response } = message as RequestEndMessage;
 
-      if (isSentryRequest(stringUrl) || request.__sentry__ !== undefined) {
-        return;
-      }
+    const stringUrl = request.origin ? request.origin.toString() + request.path : request.path;
 
-      const client = hub.getClient<NodeClient>();
-      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);
-              }
-            }
-          }
-        }
-      }
-    });
+    if (isSentryRequest(stringUrl)) {
+      return;
+    }
 
-    ds.subscribe(ChannelName.RequestEnd, message => {
-      const hub = getCurrentHub();
-      if (!hub.getIntegration(Undici)) {
-        return;
-      }
+    const span = request.__sentry_span__;
+    if (span) {
+      span.setHttpStatus(response.statusCode);
+      span.finish();
+    }
 
-      const { request, response } = message as RequestEndMessage;
+    if (this._options.breadcrumbs) {
+      hub.addBreadcrumb(
+        {
+          category: 'http',
+          data: {
+            method: request.method,
+            status_code: response.statusCode,
+            url: stringUrl,
+          },
+          type: 'http',
+        },
+        {
+          event: 'response',
+          request,
+          response,
+        },
+      );
+    }
+  };
 
-      const stringUrl = request.origin ? request.origin.toString() + request.path : request.path;
+  private _onRequestError = (message: unknown): void => {
+    const hub = getCurrentHub();
+    if (!hub.getIntegration(Undici)) {
+      return;
+    }
 
-      if (isSentryRequest(stringUrl)) {
-        return;
-      }
+    const { request } = message as RequestErrorMessage;
 
-      const span = request.__sentry__;
-      if (span) {
-        span.setHttpStatus(response.statusCode);
-        span.finish();
-      }
+    const stringUrl = request.origin ? request.origin.toString() + request.path : request.path;
 
-      if (this._options.breadcrumbs) {
-        hub.addBreadcrumb(
-          {
-            category: 'http',
-            data: {
-              method: request.method,
-              status_code: response.statusCode,
-              url: stringUrl,
-            },
-            type: 'http',
-          },
-          {
-            event: 'response',
-            request,
-            response,
-          },
-        );
-      }
-    });
+    if (isSentryRequest(stringUrl)) {
+      return;
+    }
 
-    ds.subscribe(ChannelName.RequestError, message => {
-      const hub = getCurrentHub();
-      if (!hub.getIntegration(Undici)) {
-        return;
-      }
+    const span = request.__sentry_span__;
+    if (span) {
+      span.setStatus('internal_error');
+      span.finish();
+    }
 
-      const { request } = message as RequestErrorMessage;
+    if (this._options.breadcrumbs) {
+      hub.addBreadcrumb(
+        {
+          category: 'http',
+          data: {
+            method: request.method,
+            url: stringUrl,
+          },
+          level: 'error',
+          type: 'http',
+        },
+        {
+          event: 'error',
+          request,
+        },
+      );
+    }
+  };
+}
 
-      const stringUrl = request.origin ? request.origin.toString() + request.path : request.path;
+function setHeadersOnRequest(
+  request: RequestWithSentry,
+  sentryTrace: string,
+  sentryBaggageHeader: string | undefined,
+): void {
+  if (request.__sentry_has_headers__) {
+    return;
+  }
 
-      if (isSentryRequest(stringUrl)) {
-        return;
-      }
+  request.addHeader('sentry-trace', sentryTrace);
+  if (sentryBaggageHeader) {
+    request.addHeader('baggage', sentryBaggageHeader);
+  }
 
-      const span = request.__sentry__;
-      if (span) {
-        span.setStatus('internal_error');
-        span.finish();
-      }
+  request.__sentry_has_headers__ = true;
+}
 
-      if (this._options.breadcrumbs) {
-        hub.addBreadcrumb(
-          {
-            category: 'http',
-            data: {
-              method: request.method,
-              url: stringUrl,
-            },
-            level: 'error',
-            type: 'http',
-          },
-          {
-            event: 'error',
-            request,
-          },
-        );
-      }
-    });
+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 {