feat(server-utils): Rewrite @opentelemetry/instrumentation-dataloader to orchestrion

dev-packages/node-integration-tests/suites/tracing/dataloader/test.ts

@@ -1,7 +1,11 @@
 import { afterAll, describe, expect } from 'vitest';
+import { isOrchestrionEnabled } from '../../../utils';
 import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runner';
 
-const ORIGIN = 'auto.db.otel.dataloader';
+// The span origin depends on which instrumentation is active. When the generic orchestrion run is
+// enabled (via INJECT_ORCHESTRION) the OTel `Dataloader` integration is swapped for the
+// diagnostics-channel one, which stamps a different origin.
+const ORIGIN = isOrchestrionEnabled() ? 'auto.db.orchestrion.dataloader' : 'auto.db.otel.dataloader';
 const CACHE_GET_OP = 'cache.get';
 
 describe('dataloader auto-instrumentation', () => {
@@ -25,12 +29,15 @@ describe('dataloader auto-instrumentation', () => {
             expect(loadSpan?.status).toBe('ok');
             expect(loadSpan?.data?.['sentry.origin']).toBe(ORIGIN);
             expect(loadSpan?.data?.['sentry.op']).toBe(CACHE_GET_OP);
+            // A direct operation is a client call; the deferred `batch` below gets no kind
+            expect(loadSpan?.data?.['otel.kind']).toBe('CLIENT');
 
             const batchSpan = spans.find(span => span.description === 'dataloader.batch');
             expect(batchSpan).toBeDefined();
             expect(batchSpan?.op).toBe(CACHE_GET_OP);
             expect(batchSpan?.origin).toBe(ORIGIN);
             expect(batchSpan?.status).toBe('ok');
+            expect(batchSpan?.data?.['otel.kind']).toBeUndefined();
 
             // The batch span links back to the load span that triggered it
             expect(batchSpan?.links).toEqual([
@@ -39,6 +46,11 @@ describe('dataloader auto-instrumentation', () => {
                 span_id: loadSpan?.span_id,
               }),
             ]);
+
+            // Locks down the async behavior: `load` encloses the deferred `batch` span
+            expect(batchSpan?.parent_span_id).toBe(loadSpan?.span_id);
+            expect(loadSpan?.start_timestamp).toBeLessThanOrEqual(batchSpan?.start_timestamp ?? 0);
+            expect(loadSpan?.timestamp).toBeGreaterThanOrEqual(batchSpan?.timestamp ?? 0);
           },
         })
         .expect({

packages/deno/src/index.ts

@@ -114,6 +114,7 @@ export type { DenoRedisIntegrationOptions } from './integrations/redis';
 export { denoMysqlIntegration } from './integrations/mysql';
 export { denoPostgresIntegration } from './integrations/postgres';
 export { denoAmqplibIntegration } from './integrations/amqplib';
+export { denoDataloaderIntegration } from './integrations/dataloader';
 export { denoContextIntegration } from './integrations/context';
 export { globalHandlersIntegration } from './integrations/globalhandlers';
 export { normalizePathsIntegration } from './integrations/normalizepaths';

packages/deno/src/integrations/dataloader.ts

@@ -0,0 +1,34 @@
+import { dataloaderChannelIntegration } from '@sentry/server-utils/orchestrion';
+import type { Integration, IntegrationFn } from '@sentry/core';
+import { defineIntegration, extendIntegration } from '@sentry/core';
+import { setAsyncLocalStorageAsyncContextStrategy } from '../async';
+
+const INTEGRATION_NAME = 'DenoDataloader' as const;
+
+/**
+ * Create spans for `dataloader` load/batch operations under Deno.
+ *
+ * `dataloader` channels are injected by the orchestrion runtime hook at load time.
+ * The `@sentry/deno/import` loader must be active for this integration to
+ * record anything.
+ *
+ * The channel-subscription logic is shared with the other server runtimes in
+ * `@sentry/server-utils`. This just installs Deno's `AsyncLocalStorage` context
+ * strategy (so spans nest under the active span and survive dataloader's deferred
+ * batch dispatch) before delegating.
+ */
+const _denoDataloaderIntegration = (() => {
+  const inner = dataloaderChannelIntegration();
+
+  return extendIntegration(inner, {
+    name: INTEGRATION_NAME,
+    setupOnce() {
+      setAsyncLocalStorageAsyncContextStrategy();
+    },
+  });
+}) satisfies IntegrationFn;
+
+export const denoDataloaderIntegration = defineIntegration(_denoDataloaderIntegration) as () => Integration & {
+  name: 'DenoDataloader';
+  setupOnce: () => void;
+};

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

@@ -1,7 +1,8 @@
-import { DataloaderInstrumentation } from './vendored/instrumentation';
 import type { IntegrationFn } from '@sentry/core';
 import { defineIntegration } from '@sentry/core';
 import { generateInstrumentOnce } from '@sentry/node-core';
+import { dataloaderChannelIntegration, isOrchestrionInjected } from '@sentry/server-utils/orchestrion';
+import { DataloaderInstrumentation } from './vendored/instrumentation';
 
 const INTEGRATION_NAME = 'Dataloader' as const;
 
@@ -11,7 +12,15 @@ const _dataloaderIntegration = (() => {
   return {
     name: INTEGRATION_NAME,
     setupOnce() {
-      instrumentDataloader();
+      // Decide here, not in the factory: the runtime channel injection runs inside `Sentry.init()`,
+      // after the integrations array has already been built, so `isOrchestrionInjected()` is only
+      // reliable by `setupOnce`. When the diagnostics channels are injected (runtime hook or bundler
+      // plugin), subscribe to them; otherwise fall back to the vendored OTel instrumentation.
+      if (isOrchestrionInjected()) {
+        dataloaderChannelIntegration().setupOnce?.();
+      } else {
+        instrumentDataloader();
+      }
     },
   };
 }) satisfies IntegrationFn;

packages/server-utils/src/integrations/tracing-channel/dataloader.ts

@@ -0,0 +1,173 @@
+import * as diagnosticsChannel from 'node:diagnostics_channel';
+import type { IntegrationFn, Span, StartSpanOptions } from '@sentry/core';
+import {
+  debug,
+  defineIntegration,
+  SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
+  SPAN_KIND,
+  startInactiveSpan,
+  startSpan,
+  waitForTracingChannelBinding,
+} from '@sentry/core';
+import { DEBUG_BUILD } from '../../debug-build';
+import type { ChannelName } from '../../orchestrion/channels';
+import { CHANNELS } from '../../orchestrion/channels';
+import type { TracingChannelPayloadWithSpan } from '../../tracing-channel';
+import { bindTracingChannelToSpan } from '../../tracing-channel';
+
+// NOTE: this uses the same name as the OTel integration by design.
+// When enabled, the OTel 'Dataloader' integration is omitted from the default set.
+const INTEGRATION_NAME = 'Dataloader' as const;
+
+const MODULE_NAME = 'dataloader';
+const ORIGIN = 'auto.db.orchestrion.dataloader';
+
+// `load`, `loadMany` and `batch` are cache reads; the rest are cache mutations that get no `op`.
+const CACHE_GET_OP = 'cache.get';
+
+type Operation = 'load' | 'loadMany' | 'batch' | 'prime' | 'clear' | 'clearAll';
+
+// The link shape shared between a `load` span and the `batch` span it triggers.
+type DataLoaderSpanLink = { context: ReturnType<Span['spanContext']> };
+
+// The private batch object `dataloader` stores on the loader. We stash the pending `load` span links
+// here (matching the vendored OTel instrumentation) so the batch span can link back to them.
+interface DataLoaderBatch {
+  spanLinks?: DataLoaderSpanLink[];
+}
+
+interface DataLoaderInstance {
+  name?: string | null;
+  _batch?: DataLoaderBatch | null;
+}
+
+/**
+ * The shape orchestrion's transform attaches to the tracing-channel `context`. Documented here rather
+ * than imported because orchestrion's runtime doesn't export it.
+ */
+interface DataLoaderChannelContext {
+  arguments: unknown[];
+  self?: DataLoaderInstance;
+  result?: unknown;
+  error?: unknown;
+}
+
+// Marks a wrapped `batchLoadFn` so a re-used loader (or a double construct) isn't wrapped twice.
+const WRAPPED = Symbol('sentry.dataloader.wrapped');
+
+function getSpanName(loader: DataLoaderInstance | undefined, operation: Operation): string {
+  const name = loader?.name;
+
+  return name ? `${MODULE_NAME}.${operation} ${name}` : `${MODULE_NAME}.${operation}`;
+}
+
+function makeSpanOptions(loader: DataLoaderInstance | undefined, operation: Operation): StartSpanOptions {
+  const isCacheGet = operation === 'load' || operation === 'loadMany' || operation === 'batch';
+
+  return {
+    name: getSpanName(loader, operation),
+    // Every direct operation (`load`/`loadMany`/`prime`/`clear`/`clearAll`) is a client call, matching
+    // the vendored OTel instrumentation. The `batch` runs off a deferred tick with no obvious network
+    // peer, so it gets no kind.
+    kind: operation === 'batch' ? undefined : SPAN_KIND.CLIENT,
+    op: isCacheGet ? CACHE_GET_OP : undefined,
+    onlyIfParent: true,
+    attributes: {
+      [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: ORIGIN,
+    },
+  };
+}
+
+const _dataloaderChannelIntegration = (() => {
+  return {
+    name: INTEGRATION_NAME,
+    setupOnce() {
+      // `tracingChannel` is unavailable before Node 18.19 so do nothing in that case.
+      if (!diagnosticsChannel.tracingChannel) {
+        return;
+      }
+
+      DEBUG_BUILD && debug.log('[orchestrion:dataloader] subscribing to dataloader tracing channels');
+
+      waitForTracingChannelBinding(() => {
+        subscribeConstruct();
+        subscribeLoad();
+        subscribeSimpleOperation(CHANNELS.DATALOADER_LOAD_MANY, 'loadMany');
+        subscribeSimpleOperation(CHANNELS.DATALOADER_PRIME, 'prime');
+        subscribeSimpleOperation(CHANNELS.DATALOADER_CLEAR, 'clear');
+        subscribeSimpleOperation(CHANNELS.DATALOADER_CLEAR_ALL, 'clearAll');
+      });
+    },
+  };
+}) satisfies IntegrationFn;
+
+/**
+ * Wraps the user's `batchLoadFn` (constructor arg 0) so the batch span opens when it runs on the
+ * deferred dispatch tick. The span links back to the `load` calls that populated the batch.
+ */
+function subscribeConstruct(): void {
+  diagnosticsChannel
+    .tracingChannel<DataLoaderChannelContext>(CHANNELS.DATALOADER_CONSTRUCT)
+    .start.subscribe(message => {
+      const data = message as DataLoaderChannelContext;
+      const batchLoadFn = data.arguments[0];
+      if (typeof batchLoadFn !== 'function' || (batchLoadFn as { [WRAPPED]?: boolean })[WRAPPED]) {
+        return;
+      }
+
+      const original = batchLoadFn as (...args: unknown[]) => unknown;
+      const wrapped = function (this: DataLoaderInstance, ...args: unknown[]): unknown {
+        return startSpan({ ...makeSpanOptions(this, 'batch'), links: this._batch?.spanLinks }, () =>
+          original.apply(this, args),
+        );
+      };
+      (wrapped as { [WRAPPED]?: boolean })[WRAPPED] = true;
+      data.arguments[0] = wrapped;
+    });
+}
+
+/**
+ * `load` is a cache read that additionally records its span so the batch it feeds into can link back.
+ *
+ * The span itself is `Async` (ends on `asyncEnd`, so its duration covers the awaited load). But the
+ * link has to be recorded earlier, on the synchronous `end` (fired when `load` returns): the batch
+ * dispatches on a deferred tick and reads `spanLinks` at batch-span creation, which happens BEFORE
+ * `load`'s promise resolves (`asyncEnd`) — so recording at span end would be too late and drop the
+ * link. The sync `end` runs after `dataloader` has assigned `this._batch` in `load`'s body and before
+ * the deferred dispatch, so the link lands in time.
+ */
+function subscribeLoad(): void {
+  const channel = diagnosticsChannel.tracingChannel<DataLoaderChannelContext>(CHANNELS.DATALOADER_LOAD);
+
+  bindTracingChannelToSpan(channel, data => startInactiveSpanFor(data.self, 'load'), { requiresParentSpan: true });
+
+  channel.end.subscribe(message => {
+    const data = message as TracingChannelPayloadWithSpan<DataLoaderChannelContext>;
+    const span = data._sentrySpan;
+    const batch = data.self?._batch;
+    if (span && batch && span.isRecording()) {
+      (batch.spanLinks ??= []).push({ context: span.spanContext() });
+    }
+  });
+}
+
+function subscribeSimpleOperation(channelName: ChannelName, operation: Operation): void {
+  bindTracingChannelToSpan(
+    diagnosticsChannel.tracingChannel<DataLoaderChannelContext>(channelName),
+    data => startInactiveSpanFor(data.self, operation),
+    { requiresParentSpan: true },
+  );
+}
+
+function startInactiveSpanFor(loader: DataLoaderInstance | undefined, operation: Operation): Span {
+  return startInactiveSpan(makeSpanOptions(loader, operation));
+}
+
+/**
+ * EXPERIMENTAL: orchestrion-driven `dataloader` integration.
+ *
+ * Subscribes to the `orchestrion:dataloader:*` diagnostics_channels that the orchestrion code
+ * transform injects into `dataloader`'s constructor and prototype methods. Requires the orchestrion
+ * runtime hook or bundler plugin to be active.
+ */
+export const dataloaderChannelIntegration = defineIntegration(_dataloaderChannelIntegration);

packages/server-utils/src/orchestrion/config/dataloader.ts

@@ -1,6 +1,55 @@
 import type { InstrumentationConfig } from '@apm-js-collab/code-transformer';
 
-// TODO: Stub for the `dataloader` orchestrion integration (ports `@opentelemetry/instrumentation-dataloader`).
-export const dataloaderConfig: InstrumentationConfig[] = [];
+// `dataloader` ships a single transpiled CommonJS `index.js`. Its class methods are emitted as
+// `_proto.<name> = function <name>() {}` (named function *expressions*), so they match on
+// `expressionName` rather than `methodName`. The constructor is a named function declaration.
+// The version range mirrors `supportedVersions` in the vendored OTel instrumentation.
+const module = { name: 'dataloader', versionRange: '>=2.0.0 <3', filePath: 'index.js' } as const;
 
-export const dataloaderChannels = {} as const;
+export const dataloaderConfig = [
+  // Wrap the constructor so the subscriber can wrap the user's `batchLoadFn` (arg 0). The batch span
+  // is opened when that wrapped function actually runs (on the deferred dispatch tick), mirroring the
+  // vendored OTel instrumentation which also wraps `batchLoadFn` at construction time.
+  {
+    channelName: 'construct',
+    module,
+    functionQuery: { functionName: 'DataLoader', kind: 'Sync' },
+  },
+  // `load`/`loadMany` return Promises, so they're `Async`: the span ends on `asyncEnd` (when the
+  // load resolves), capturing the real latency and enclosing the deferred `batch` span — matching the
+  // vendored OTel `startSpan`. `prime`/`clear`/`clearAll` return `this` synchronously, so they stay `Sync`.
+  {
+    channelName: 'load',
+    module,
+    functionQuery: { expressionName: 'load', kind: 'Async' },
+  },
+  {
+    channelName: 'loadMany',
+    module,
+    functionQuery: { expressionName: 'loadMany', kind: 'Async' },
+  },
+  {
+    channelName: 'prime',
+    module,
+    functionQuery: { expressionName: 'prime', kind: 'Sync' },
+  },
+  {
+    channelName: 'clear',
+    module,
+    functionQuery: { expressionName: 'clear', kind: 'Sync' },
+  },
+  {
+    channelName: 'clearAll',
+    module,
+    functionQuery: { expressionName: 'clearAll', kind: 'Sync' },
+  },
+] satisfies InstrumentationConfig[];
+
+export const dataloaderChannels = {
+  DATALOADER_CONSTRUCT: 'orchestrion:dataloader:construct',
+  DATALOADER_LOAD: 'orchestrion:dataloader:load',
+  DATALOADER_LOAD_MANY: 'orchestrion:dataloader:loadMany',
+  DATALOADER_PRIME: 'orchestrion:dataloader:prime',
+  DATALOADER_CLEAR: 'orchestrion:dataloader:clear',
+  DATALOADER_CLEAR_ALL: 'orchestrion:dataloader:clearAll',
+} as const;

packages/server-utils/src/orchestrion/index.ts

@@ -1,5 +1,6 @@
 import { amqplibChannelIntegration } from '../integrations/tracing-channel/amqplib';
 import { anthropicChannelIntegration } from '../integrations/tracing-channel/anthropic';
+import { dataloaderChannelIntegration } from '../integrations/tracing-channel/dataloader';
 import { genericPoolChannelIntegration } from '../integrations/tracing-channel/generic-pool';
 import { googleGenAIChannelIntegration } from '../integrations/tracing-channel/google-genai';
 import {
@@ -24,6 +25,7 @@ export { nestjsChannels } from './config/nestjs';
 export {
   amqplibChannelIntegration,
   anthropicChannelIntegration,
+  dataloaderChannelIntegration,
   genericPoolChannelIntegration,
   googleGenAIChannelIntegration,
   graphqlChannelIntegration,
@@ -63,6 +65,11 @@ export type * from '../integrations/tracing-channel/graphql/graphql-types';
  * Framework SDKs that own their own channel listener (e.g. `@sentry/nestjs`'s `Nest`) are NOT here
  * either: their transform config is still in `SENTRY_INSTRUMENTATIONS`, but the listener lives in
  * their package and picks the channel-vs-OTel path itself at `setupOnce`, so it needs no central swap.
+ *
+ * NOTE: `dataloaderChannelIntegration` is also NOT here. Everything in this map is auto-appended to
+ * the default integrations, but the OTel `Dataloader` integration is opt-in (never a default). Like
+ * `@sentry/nestjs`'s `Nest`, its `@sentry/node` factory picks the channel-vs-OTel path itself at
+ * `setupOnce` (via `isOrchestrionInjected()`), so there's nothing for the central swap to do.
  */
 export const channelIntegrations = {
   postgresIntegration: postgresChannelIntegration,