feat(node): WIP Only add span listeners conditionally

Author: mydea

packages/node/src/integrations/tracing/connect.ts

@@ -104,7 +104,7 @@ function addConnectSpanAttributes(span: Span): void {
     [SEMANTIC_ATTRIBUTE_SENTRY_OP]: `${type}.connect`,
   });
 
-  // Also update the name, we don't need to "middleware - " prefix
+  // Also update the name, we don't need the "middleware - " prefix
   const name = attributes['connect.name'];
   if (typeof name === 'string') {
     span.updateName(name);

packages/node/src/integrations/tracing/dataloader.ts

@@ -3,10 +3,11 @@ import {
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   defineIntegration,
+  getClient,
   spanToJSON,
 } from '@sentry/core';
 import type { IntegrationFn } from '@sentry/core';
-import { generateInstrumentOnce } from '../../otel/instrument';
+import { callWhenWrapped, generateInstrumentOnce } from '../../otel/instrument';
 
 const INTEGRATION_NAME = 'Dataloader';
 
@@ -19,31 +20,38 @@ export const instrumentDataloader = generateInstrumentOnce(
 );
 
 const _dataloaderIntegration = (() => {
+  let hookCallback: undefined | (() => void);
+
   return {
     name: INTEGRATION_NAME,
     setupOnce() {
-      instrumentDataloader();
-    },
+      const instrumentation = instrumentDataloader();
 
-    setup(client) {
-      client.on('spanStart', span => {
-        const spanJSON = spanToJSON(span);
-        if (spanJSON.description?.startsWith('dataloader')) {
-          span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.dataloader');
+      callWhenWrapped(instrumentation, () => {
+        const client = getClient();
+        if (hookCallback || !client) {
+          return;
         }
 
-        // These are all possible dataloader span descriptions
-        // Still checking for the future versions
-        // in case they add support for `clear` and `prime`
-        if (
-          spanJSON.description === 'dataloader.load' ||
-          spanJSON.description === 'dataloader.loadMany' ||
-          spanJSON.description === 'dataloader.batch'
-        ) {
-          span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'cache.get');
-          // TODO: We can try adding `key` to the `data` attribute upstream.
-          // Or alternatively, we can add `requestHook` to the dataloader instrumentation.
-        }
+        hookCallback = client.on('spanStart', span => {
+          const spanJSON = spanToJSON(span);
+          if (spanJSON.description?.startsWith('dataloader')) {
+            span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.db.otel.dataloader');
+          }
+
+          // These are all possible dataloader span descriptions
+          // Still checking for the future versions
+          // in case they add support for `clear` and `prime`
+          if (
+            spanJSON.description === 'dataloader.load' ||
+            spanJSON.description === 'dataloader.loadMany' ||
+            spanJSON.description === 'dataloader.batch'
+          ) {
+            span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'cache.get');
+            // TODO: We can try adding `key` to the `data` attribute upstream.
+            // Or alternatively, we can add `requestHook` to the dataloader instrumentation.
+          }
+        });
       });
     },
   };

packages/node/src/otel/instrument.ts

@@ -7,19 +7,22 @@ export const INSTRUMENTED: Record<string, Instrumentation> = {};
  * Instrument an OpenTelemetry instrumentation once.
  * This will skip running instrumentation again if it was already instrumented.
  */
-export function generateInstrumentOnce<Options = unknown>(
+export function generateInstrumentOnce<
+  Options = unknown,
+  InstrumentationInstance extends Instrumentation = Instrumentation,
+>(
   name: string,
-  creator: (options?: Options) => Instrumentation,
-): ((options?: Options) => void) & { id: string } {
+  creator: (options?: Options) => InstrumentationInstance,
+): ((options?: Options) => InstrumentationInstance) & { id: string } {
   return Object.assign(
     (options?: Options) => {
-      const instrumented = INSTRUMENTED[name];
+      const instrumented = INSTRUMENTED[name] as InstrumentationInstance | undefined;
       if (instrumented) {
         // If options are provided, ensure we update them
         if (options) {
           instrumented.setConfig(options);
         }
-        return;
+        return instrumented;
       }
 
       const instrumentation = creator(options);
@@ -28,7 +31,35 @@ export function generateInstrumentOnce<Options = unknown>(
       registerInstrumentations({
         instrumentations: [instrumentation],
       });
+
+      return instrumentation;
     },
     { id: name },
   );
 }
+
+/**
+ * Ensure a given callback is called when the instrumentation is actually wrapping something.
+ * This can be used to ensure some logic is only called when the instrumentation is actually active.
+ * This depends on wrapping `_wrap` (inception!). If this is not possible (e.g. the property name is mangled, ...)
+ * the callback will be called immediately.
+ */
+export function callWhenWrapped<T extends Instrumentation>(instrumentation: T, callback: () => void): void {
+  if (!hasWrap(instrumentation)) {
+    callback();
+    return;
+  }
+
+  const originalWrap = instrumentation['_wrap'];
+
+  instrumentation['_wrap'] = (...args: Parameters<typeof originalWrap>) => {
+    callback();
+    return originalWrap(...args);
+  };
+}
+
+function hasWrap<T extends Instrumentation>(
+  instrumentation: T,
+): instrumentation is T & { _wrap: (...args: unknown[]) => unknown } {
+  return typeof (instrumentation as T & { _wrap?: (...args: unknown[]) => unknown })['_wrap'] === 'function';
+}