[otel] rely on context for parenting spans correctly

EmrysMyrddin · EmrysMyrddin · commit 6faa8672f3e7 · 2025-02-10T11:51:04.000+01:00
diff --git a/.changeset/@graphql-mesh_plugin-opentelemetry-532-dependencies.md b/.changeset/@graphql-mesh_plugin-opentelemetry-532-dependencies.md
@@ -0,0 +1,7 @@
+---
+'@graphql-mesh/plugin-opentelemetry': patch
+---
+
+dependencies updates: 
+
+- Added dependency [`@opentelemetry/context-async-hooks@^1.30.1` ↗︎](https://www.npmjs.com/package/@opentelemetry/context-async-hooks/v/1.30.1) (to `dependencies`)
diff --git a/e2e/opentelemetry/opentelemetry.e2e.ts b/e2e/opentelemetry/opentelemetry.e2e.ts
@@ -24,15 +24,18 @@ beforeAll(async () => {
 type JaegerTracesApiResponse = {
   data: Array<{
     traceID: string;
-    spans: Array<{
-      traceID: string;
-      spanID: string;
-      operationName: string;
-      tags: Array<{ key: string; value: string; type: string }>;
-    }>;
+    spans: JaegerTraceSpan[];
   }>;
 };
 
+type JaegerTraceSpan = {
+  traceID: string;
+  spanID: string;
+  operationName: string;
+  tags: Array<{ key: string; value: string; type: string }>;
+  references: Array<{ refType: string; spanID: string; traceID: string }>;
+};
+
 describe('OpenTelemetry', () => {
   (['grpc', 'http'] as const).forEach((OTLP_EXPORTER_TYPE) => {
     describe(`exporter > ${OTLP_EXPORTER_TYPE}`, () => {
@@ -77,6 +80,9 @@ describe('OpenTelemetry', () => {
             await checkFn(res);
             return;
           } catch (e) {
+            if (signal.aborted) {
+              throw err;
+            }
             err = e;
           }
         }
@@ -559,40 +565,52 @@ describe('OpenTelemetry', () => {
           expect(relevantTraces.length).toBe(1);
           const relevantTrace = relevantTraces[0];
           expect(relevantTrace).toBeDefined();
-          expect(relevantTrace?.spans.length).toBe(11);
+          expect(relevantTrace!.spans.length).toBe(18);
 
-          expect(relevantTrace?.spans).toContainEqual(
-            expect.objectContaining({ operationName: 'POST /graphql' }),
-          );
-          expect(relevantTrace?.spans).toContainEqual(
-            expect.objectContaining({ operationName: 'graphql.parse' }),
-          );
-          expect(relevantTrace?.spans).toContainEqual(
-            expect.objectContaining({ operationName: 'graphql.validate' }),
-          );
-          expect(relevantTrace?.spans).toContainEqual(
-            expect.objectContaining({ operationName: 'graphql.execute' }),
+          const spanTree = buildSpanTree(relevantTrace!.spans, 'POST /graphql');
+          expect(spanTree).toBeDefined();
+
+          const expectedHttpChildren = [
+            'graphql.parse',
+            'graphql.validate',
+            'graphql.execute',
+          ];
+          expect(spanTree!.children).toHaveLength(3);
+          for (const operationName of expectedHttpChildren) {
+            expect(spanTree?.children).toContainEqual(
+              expect.objectContaining({
+                span: expect.objectContaining({ operationName }),
+              }),
+            );
+          }
+
+          const executeSpan = spanTree?.children.find(
+            ({ span }) => span.operationName === 'graphql.execute',
           );
-          expect(
-            relevantTrace?.spans.filter(
-              (r) => r.operationName === 'subgraph.execute (accounts)',
-            ).length,
-          ).toBe(2);
-          expect(
-            relevantTrace?.spans.filter(
-              (r) => r.operationName === 'subgraph.execute (products)',
-            ).length,
-          ).toBe(2);
-          expect(
-            relevantTrace?.spans.filter(
-              (r) => r.operationName === 'subgraph.execute (inventory)',
-            ).length,
-          ).toBe(1);
-          expect(
-            relevantTrace?.spans.filter(
-              (r) => r.operationName === 'subgraph.execute (reviews)',
-            ).length,
-          ).toBe(2);
+
+          const expectedExecuteChildren = [
+            ['subgraph.execute (accounts)', 2],
+            ['subgraph.execute (products)', 2],
+            ['subgraph.execute (inventory)', 1],
+            ['subgraph.execute (reviews)', 2],
+          ] as const;
+
+          for (const [operationName, count] of expectedExecuteChildren) {
+            const matchingChildren = executeSpan!.children.filter(
+              ({ span }) => span.operationName === operationName,
+            );
+            expect(matchingChildren).toHaveLength(count);
+            for (const child of matchingChildren) {
+              expect(child.children).toHaveLength(1);
+              expect(child.children).toContainEqual(
+                expect.objectContaining({
+                  span: expect.objectContaining({
+                    operationName: 'http.fetch',
+                  }),
+                }),
+              );
+            }
+          }
         });
       });
 
@@ -1279,3 +1297,28 @@ describe('OpenTelemetry', () => {
     });
   });
 });
+
+type TraceTreeNode = {
+  span: JaegerTraceSpan;
+  children: TraceTreeNode[];
+};
+function buildSpanTree(
+  spans: JaegerTraceSpan[],
+  rootName: string,
+): TraceTreeNode | undefined {
+  function buildNode(root: JaegerTraceSpan): TraceTreeNode {
+    return {
+      span: root,
+      children: spans
+        .filter((span) =>
+          span.references.find(
+            (ref) => ref.refType === 'CHILD_OF' && ref.spanID === root.spanID,
+          ),
+        )
+        .map(buildNode),
+    };
+  }
+
+  const root = spans.find((span) => span.operationName === rootName);
+  return root && buildNode(root);
+}
diff --git a/package.json b/package.json
@@ -60,6 +60,7 @@
     "vitest": "^3.0.1"
   },
   "resolutions": {
+    "@envelop/core": "5.1.0-alpha-20250206123608-cd323a08c43c066eb34cba698d61b4f059ab5ee5",
     "@graphql-tools/delegate": "workspace:^",
     "@opentelemetry/exporter-trace-otlp-http": "patch:@opentelemetry/exporter-trace-otlp-http@npm%3A0.56.0#~/.yarn/patches/@opentelemetry-exporter-trace-otlp-http-npm-0.56.0-dddd282e41.patch",
     "@opentelemetry/otlp-exporter-base": "patch:@opentelemetry/otlp-exporter-base@npm%3A0.56.0#~/.yarn/patches/@opentelemetry-otlp-exporter-base-npm-0.56.0-ba3dc5f5c5.patch",
diff --git a/packages/plugins/opentelemetry/package.json b/packages/plugins/opentelemetry/package.json
@@ -50,6 +50,7 @@
     "@graphql-mesh/utils": "^0.103.6",
     "@graphql-tools/utils": "^10.7.0",
     "@opentelemetry/api": "^1.9.0",
+    "@opentelemetry/context-async-hooks": "^1.30.1",
     "@opentelemetry/exporter-trace-otlp-grpc": "^0.57.0",
     "@opentelemetry/exporter-trace-otlp-http": "^0.57.0",
     "@opentelemetry/exporter-zipkin": "^1.29.0",
@@ -62,6 +63,7 @@
     "tslib": "^2.8.1"
   },
   "devDependencies": {
+    "@whatwg-node/server": "^0.9.65",
     "graphql": "^16.9.0",
     "graphql-yoga": "^5.10.11",
     "pkgroll": "2.8.2"
diff --git a/packages/plugins/opentelemetry/src/contextManager.ts b/packages/plugins/opentelemetry/src/contextManager.ts
@@ -0,0 +1,42 @@
+import { trace, type Context } from '@opentelemetry/api';
+
+type Node = {
+  ctx: Context;
+  previous?: Node;
+};
+
+export class OtelContextStack {
+  #root: Node;
+  #current: Node;
+
+  constructor(root: Context) {
+    this.#root = { ctx: root };
+    this.#current = this.#root;
+  }
+
+  get current(): Context {
+    return this.#current.ctx;
+  }
+
+  get root(): Context {
+    return this.#root.ctx;
+  }
+
+  push = (ctx: Context) => {
+    this.#current = { ctx, previous: this.#current };
+  };
+
+  pop = () => {
+    this.#current = this.#current.previous ?? this.#root;
+  };
+
+  toString() {
+    let node: Node | undefined = this.#current;
+    const names = [];
+    while (node != undefined) {
+      names.push((trace.getSpan(node.ctx) as unknown as { name: string }).name);
+      node = node.previous;
+    }
+    return names.join(' -> ');
+  }
+}
diff --git a/packages/plugins/opentelemetry/src/plugin-utils.ts b/packages/plugins/opentelemetry/src/plugin-utils.ts
@@ -0,0 +1,104 @@
+import type { ExecutionRequest } from '@graphql-tools/utils';
+
+export function withState<
+  P,
+  GraphqlState = object,
+  HttpState = object,
+  SubExecState = object,
+>(plugin: WithState<P, HttpState, GraphqlState, SubExecState>): P {
+  const states: {
+    forRequest?: WeakMap<Request, Partial<HttpState>>;
+    forOperation?: WeakMap<any, Partial<GraphqlState>>;
+    forSubgraphExecution?: WeakMap<ExecutionRequest, Partial<SubExecState>>;
+  } = {};
+
+  function getProp(scope: keyof typeof states, key: any): PropertyDescriptor {
+    return {
+      get() {
+        if (!states[scope]) states[scope] = new WeakMap<any, any>();
+        let value = states[scope].get(key as any);
+        if (!value) states[scope].set(key, (value = {}));
+        return value;
+      },
+      enumerable: true,
+    };
+  }
+
+  const pluginWithState: Record<string, (payload: any) => unknown> = {};
+  for (const [hookName, hook] of Object.entries(plugin) as any) {
+    pluginWithState[hookName] = (payload) =>
+      hook({
+        ...payload,
+        get state() {
+          let { executionRequest, context, request } = payload;
+
+          const state = {};
+          const defineState = (scope: keyof typeof states, key: any) =>
+            Object.defineProperty(state, scope, getProp(scope, key));
+
+          if (executionRequest) {
+            defineState('forSubgraphExecution', executionRequest);
+            if (executionRequest.context) context = executionRequest.context;
+          }
+          if (context) {
+            defineState('forOperation', context);
+            if (context.request) request = context.request;
+          }
+          if (request) {
+            defineState('forRequest', request);
+          }
+          return state;
+        },
+      });
+  }
+
+  return pluginWithState as P;
+}
+
+export type HttpState<T> = {
+  forRequest: Partial<T>;
+};
+
+export type GraphQLState<T> = {
+  forOperation: Partial<T>;
+};
+
+export type GatewayState<T> = {
+  forSubgraphExecution: Partial<T>;
+};
+
+export function getMostSpecificState<T>(
+  state: Partial<HttpState<T> & GraphQLState<T> & GatewayState<T>> = {},
+): Partial<T> | undefined {
+  const { forOperation, forRequest, forSubgraphExecution } = state;
+  return forSubgraphExecution ?? forOperation ?? forRequest;
+}
+
+// Brace yourself! TS Wizardry is coming!
+
+type PayloadWithState<T, Http, GraphQL, Gateway> = T extends {
+  executionRequest: any;
+}
+  ? T & {
+      state: Partial<HttpState<Http> & GraphQLState<GraphQL>> &
+        GatewayState<Gateway>;
+    }
+  : T extends {
+        executionRequest?: any;
+      }
+    ? T & {
+        state: Partial<
+          HttpState<Http> & GraphQLState<GraphQL> & GatewayState<Gateway>
+        >;
+      }
+    : T extends { context: any }
+      ? T & { state: HttpState<Http> & GraphQLState<GraphQL> }
+      : T extends { request: any }
+        ? T & { state: HttpState<Http> }
+        : T;
+
+type WithState<P, Http = object, GraphQL = object, Gateway = object> = {
+  [K in keyof P]: P[K] extends ((payload: infer T) => infer R) | undefined
+    ? (payload: PayloadWithState<T, Http, GraphQL, Gateway>) => R | undefined
+    : P[K];
+};
diff --git a/packages/plugins/opentelemetry/src/plugin.ts b/packages/plugins/opentelemetry/src/plugin.ts
diff --git a/packages/plugins/opentelemetry/src/spans.ts b/packages/plugins/opentelemetry/src/spans.ts
diff --git a/packages/plugins/opentelemetry/src/utils.ts b/packages/plugins/opentelemetry/src/utils.ts
