feat(opentelemetry): add background spans (schema + init) (#1071)

Author: EmrysMyrddin

e2e/cloudflare-workers/cloudflare-workers.e2e.ts

@@ -135,7 +135,7 @@ describe.skipIf(gatewayRunner !== 'node')('Cloudflare Workers', () => {
     `);
 
     const traces = await getJaegerTraces(serviceName, 2);
-    expect(traces.data.length).toBe(3);
+    expect(traces.data.length).toBe(2);
     const relevantTraces = traces.data.filter((trace) =>
       trace.spans.some((span) => span.operationName === 'POST /graphql'),
     );
@@ -234,7 +234,7 @@ describe.skipIf(gatewayRunner !== 'node')('Cloudflare Workers', () => {
     );
 
     const traces = await getJaegerTraces(serviceName, 3);
-    expect(traces.data.length).toBe(4);
+    expect(traces.data.length).toBe(3);
 
     const relevantTraces = traces.data.filter((trace) =>
       trace.spans.some((span) => span.operationName === 'POST /graphql'),

packages/fusion-runtime/src/unifiedGraphManager.ts

@@ -117,6 +117,15 @@ export type Instrumentation = {
     payload: { executionRequest: ExecutionRequest; subgraphName: string },
     wrapped: () => MaybePromise<void>,
   ) => MaybePromise<void>;
+  /**
+   * Wrap each supergraph schema loading.
+   *
+   * Note: this span is only available when an Async compatible context manager is available
+   */
+  schema?: (
+    payload: null,
+    wrapped: () => MaybePromise<void>,
+  ) => MaybePromise<void>;
 };
 
 const UNIFIEDGRAPH_CACHE_KEY = 'hive-gateway:supergraph';

packages/plugins/opentelemetry/src/plugin.ts

@@ -54,6 +54,7 @@ import {
   createGraphQLSpan,
   createGraphQLValidateSpan,
   createHttpSpan,
+  createSchemaLoadingSpan,
   startSubgraphExecuteFetchSpan as createSubgraphExecuteFetchSpan,
   createUpstreamHttpFetchSpan,
   recordCacheError,
@@ -67,11 +68,15 @@ import {
   setGraphQLValidateAttributes,
   setParamsAttributes,
   setResponseAttributes,
+  setSchemaAttributes,
   setUpstreamFetchAttributes,
   setUpstreamFetchResponseAttributes,
 } from './spans';
 import { getEnvVar, tryContextManagerSetup } from './utils';
 
+const initializationTime =
+  'performance' in globalThis ? performance.now() : undefined;
+
 type BooleanOrPredicate<TInput = never> =
   | boolean
   | ((input: TInput) => boolean);
@@ -177,13 +182,6 @@ export type OpenTelemetryGatewayPluginOptions =
      * You may specify a boolean value to enable/disable all spans, or a function to dynamically enable/disable spans based on the input.
      */
     spans?: {
-      /**
-       * Enable/disable Spans of internal introspection queries in proxy mode (default: true).
-       */
-      introspection?: BooleanOrPredicate<{
-        executionRequest: ExecutionRequest;
-        subgraphName: string;
-      }>;
       /**
        * Enable/disable HTTP request spans (default: true).
        *
@@ -231,6 +229,16 @@ export type OpenTelemetryGatewayPluginOptions =
        * Enable/Disable cache related span events (default: true).
        */
       cache?: BooleanOrPredicate<{ key: string; action: 'read' | 'write' }>;
+      /**
+       * Enable/disable schema loading spans (default: true if context manager available).
+       *
+       * Note: This span requires an Async compatible context manager
+       */
+      schema?: boolean;
+      /**
+       * Enable/disable initialization span (default: true).
+       */
+      initialization?: boolean;
     };
   };
 
@@ -283,16 +291,25 @@ export function useOpenTelemetry(
   let provider: WebTracerProvider;
 
   const yogaVersion = createDeferred<string>();
+  let initSpan: Context | null;
 
   function isParentEnabled(state: State): boolean {
     const parentState = getMostSpecificState(state);
     return !parentState || !!parentState.otel;
   }
 
   function getContext(state?: State): Context {
-    return useContextManager
-      ? context.active()
-      : (getMostSpecificState(state)?.otel?.current ?? ROOT_CONTEXT);
+    const specificState = getMostSpecificState(state)?.otel;
+
+    if (initSpan && !specificState) {
+      return initSpan;
+    }
+
+    if (useContextManager) {
+      return context.active();
+    }
+
+    return specificState?.current ?? ROOT_CONTEXT;
   }
 
   const yogaLogger = createDeferred<YogaLogger>();
@@ -370,12 +387,27 @@ export function useOpenTelemetry(
   preparation$ = init().then((contextManager) => {
     useContextManager = contextManager;
     tracer = options.tracer || trace.getTracer('gateway');
+    initSpan = trace.setSpan(
+      context.active(),
+      tracer.startSpan('gateway.initialization', {
+        startTime: initializationTime,
+      }),
+    );
     preparation$ = fakePromise();
     return pluginLogger.then((logger) => {
       pluginLogger = fakePromise(logger);
       logger.debug(
         `context manager is ${useContextManager ? 'enabled' : 'disabled'}`,
       );
+      if (!useContextManager) {
+        if (options.spans?.schema) {
+          logger.warn(
+            'Schema loading spans are disabled because no context manager is available',
+          );
+        }
+        options.spans = options.spans ?? {};
+        options.spans.schema = false;
+      }
       diag.setLogger(
         {
           error: (message, ...args) =>
@@ -453,23 +485,25 @@ export function useOpenTelemetry(
           return wrapped();
         }
 
-        const ctx = getContext(parentState);
-        forOperation.otel = new OtelContextStack(
-          createGraphQLSpan({ tracer, ctx }),
-        );
-
-        if (useContextManager) {
-          wrapped = context.bind(forOperation.otel.current, wrapped);
-        }
-
         return unfakePromise(
-          fakePromise()
-            .then(wrapped)
-            .catch((err) => {
-              registerException(forOperation.otel?.current, err);
-              throw err;
-            })
-            .finally(() => trace.getSpan(forOperation.otel!.current)?.end()),
+          preparation$.then(() => {
+            const ctx = getContext(parentState);
+            forOperation.otel = new OtelContextStack(
+              createGraphQLSpan({ tracer, ctx }),
+            );
+
+            if (useContextManager) {
+              wrapped = context.bind(forOperation.otel.current, wrapped);
+            }
+
+            return fakePromise()
+              .then(wrapped)
+              .catch((err) => {
+                registerException(forOperation.otel?.current, err);
+                throw err;
+              })
+              .finally(() => trace.getSpan(forOperation.otel!.current)?.end());
+          }),
         );
       },
 
@@ -610,7 +644,7 @@ export function useOpenTelemetry(
           parentState.forOperation?.skipExecuteSpan ||
           !shouldTrace(
             isIntrospection
-              ? options.spans?.introspection
+              ? options.spans?.schema
               : options.spans?.subgraphExecute,
             {
               subgraphName,
@@ -625,7 +659,7 @@ export function useOpenTelemetry(
         // (such as Introspection requests in proxy mode), we don't want to use the active context,
         // we want the span to be in it's own trace.
         const parentContext = isIntrospection
-          ? ROOT_CONTEXT
+          ? context.active()
           : getContext(parentState);
 
         forSubgraphExecution.otel = new OtelContextStack(
@@ -671,29 +705,51 @@ export function useOpenTelemetry(
           return wrapped();
         }
 
-        const { forSubgraphExecution } = state;
-        const ctx = createUpstreamHttpFetchSpan({
-          ctx: getContext(state),
-          tracer,
-        });
-
-        forSubgraphExecution?.otel!.push(ctx);
+        return unfakePromise(
+          preparation$.then(() => {
+            const { forSubgraphExecution } = state;
+            const ctx = createUpstreamHttpFetchSpan({
+              ctx: getContext(state),
+              tracer,
+            });
+
+            forSubgraphExecution?.otel!.push(ctx);
+
+            if (useContextManager) {
+              wrapped = context.bind(ctx, wrapped);
+            }
+
+            return fakePromise()
+              .then(wrapped)
+              .catch((err) => {
+                registerException(ctx, err);
+                throw err;
+              })
+              .finally(() => {
+                trace.getSpan(ctx)?.end();
+                forSubgraphExecution?.otel!.pop();
+              });
+          }),
+        );
+      },
 
-        if (useContextManager) {
-          wrapped = context.bind(ctx, wrapped);
+      schema(_, wrapped) {
+        if (!shouldTrace(options.spans?.schema, null)) {
+          return wrapped();
         }
 
         return unfakePromise(
-          fakePromise()
-            .then(wrapped)
-            .catch((err) => {
-              registerException(ctx, err);
-              throw err;
-            })
-            .finally(() => {
-              trace.getSpan(ctx)?.end();
-              forSubgraphExecution?.otel!.pop();
-            }),
+          preparation$.then(() => {
+            const ctx = createSchemaLoadingSpan({ tracer });
+            return fakePromise()
+              .then(() => context.with(ctx, wrapped))
+              .catch((err) => {
+                trace.getSpan(ctx)?.recordException(err);
+              })
+              .finally(() => {
+                trace.getSpan(ctx)?.end();
+              });
+          }),
         );
       },
     },
@@ -863,6 +919,16 @@ export function useOpenTelemetry(
         setUpstreamFetchResponseAttributes({ ctx, response });
       };
     },
+
+    onSchemaChange(payload) {
+      setSchemaAttributes(payload);
+
+      if (initSpan) {
+        trace.getSpan(initSpan)?.end();
+        initSpan = null;
+      }
+    },
+
     async onDispose() {
       if (options.initializeNodeSDK) {
         await provider?.forceFlush?.();

packages/plugins/opentelemetry/src/spans.ts

@@ -7,6 +7,7 @@ import {
   type ExecutionResult,
 } from '@graphql-tools/utils';
 import {
+  ROOT_CONTEXT,
   SpanKind,
   SpanStatusCode,
   trace,
@@ -18,7 +19,7 @@ import {
   SEMATTRS_EXCEPTION_STACKTRACE,
   SEMATTRS_EXCEPTION_TYPE,
 } from '@opentelemetry/semantic-conventions';
-import type { ExecutionArgs } from 'graphql';
+import { printSchema, type ExecutionArgs, type GraphQLSchema } from 'graphql';
 import type { GraphQLParams } from 'graphql-yoga';
 import {
   getRetryInfo,
@@ -435,6 +436,24 @@ export function setExecutionResultAttributes(input: {
   }
 }
 
+export function createSchemaLoadingSpan(inputs: { tracer: Tracer }) {
+  const span = inputs.tracer.startSpan(
+    'gateway.schema',
+    { attributes: { 'gateway.schema.changed': false } },
+    ROOT_CONTEXT,
+  );
+  return trace.setSpan(ROOT_CONTEXT, span);
+}
+
+export function setSchemaAttributes(inputs: { schema: GraphQLSchema }) {
+  const span = trace.getActiveSpan();
+  if (!span) {
+    return;
+  }
+  span.setAttribute('gateway.schema.changed', true);
+  span.setAttribute('graphql.schema', printSchema(inputs.schema));
+}
+
 export function registerException(ctx: Context | undefined, error: any) {
   const span = ctx && trace.getSpan(ctx);
   if (!span) {