feat(opentelemetry): add cache events and attributes (#993)

EmrysMyrddin · enisdenjo · commit 838e9854645f · 2025-04-21T14:30:28.000+02:00
diff --git a/packages/plugins/opentelemetry/src/plugin-utils.ts b/packages/plugins/opentelemetry/src/plugin-utils.ts
@@ -56,19 +56,23 @@ export function withState<
   function addStateGetters(src: any) {
     const result: any = {};
     for (const [hookName, hook] of Object.entries(src) as any) {
-      result[hookName] =
-        typeof hook !== 'function'
-          ? hook
-          : (payload: any, ...args: any[]) =>
-              hook(
-                {
-                  ...payload,
-                  get state() {
-                    return getState(payload);
-                  },
+      if (typeof hook !== 'function') {
+        result[hookName] = hook;
+      } else {
+        result[hookName] = {
+          [hook.name](payload: any, ...args: any[]) {
+            return hook(
+              {
+                ...payload,
+                get state() {
+                  return getState(payload);
                 },
-                ...args,
-              );
+              },
+              ...args,
+            );
+          },
+        }[hook.name];
+      }
     }
     return result;
   }
diff --git a/packages/plugins/opentelemetry/src/plugin.ts b/packages/plugins/opentelemetry/src/plugin.ts
@@ -51,8 +51,11 @@ import {
   createHttpSpan,
   startSubgraphExecuteFetchSpan as createSubgraphExecuteFetchSpan,
   createUpstreamHttpFetchSpan,
+  recordCacheError,
+  recordCacheEvent,
   registerException,
   setExecutionAttributesOnOperationSpan,
+  setExecutionResultAttributes,
   setGraphQLExecutionAttributes,
   setGraphQLExecutionResultAttributes,
   setGraphQLParseAttributes,
@@ -205,6 +208,10 @@ export type OpenTelemetryGatewayPluginOptions =
        * Enable/disable upstream HTTP fetch calls spans (default: true).
        */
       upstreamFetch?: BooleanOrPredicate<ExecutionRequest | undefined>;
+      /**
+       * Enable/Disable cache related span events (default: true).
+       */
+      cache?: BooleanOrPredicate<{ key: string; action: 'read' | 'write' }>;
     };
   };
 
@@ -677,6 +684,25 @@ export function useOpenTelemetry(
       });
     },
 
+    onCacheGet: (payload) =>
+      shouldTrace(options.spans?.cache, { key: payload.key, action: 'read' })
+        ? {
+            onCacheMiss: () => recordCacheEvent('miss', payload),
+            onCacheHit: () => recordCacheEvent('hit', payload),
+            onCacheGetError: ({ error }) =>
+              recordCacheError('read', error, payload),
+          }
+        : undefined,
+
+    onCacheSet: (payload) =>
+      shouldTrace(options.spans?.cache, { key: payload.key, action: 'write' })
+        ? {
+            onCacheSetDone: () => recordCacheEvent('write', payload),
+            onCacheSetError: ({ error }) =>
+              recordCacheError('write', error, payload),
+          }
+        : undefined,
+
     onResponse({ response, state }) {
       try {
         state.forRequest.otel &&
@@ -686,7 +712,7 @@ export function useOpenTelemetry(
       }
     },
 
-    onParams({ state, context: gqlCtx, params }) {
+    onParams: function onParamsOTEL({ state, context: gqlCtx, params }) {
       if (
         !isParentEnabled(state) ||
         !shouldTrace(options.spans?.graphql, gqlCtx)
@@ -698,6 +724,21 @@ export function useOpenTelemetry(
       setParamsAttributes({ ctx, params });
     },
 
+    onExecutionResult: function onExeResOTEL({
+      result,
+      context: gqlCtx,
+      state,
+    }) {
+      if (
+        !isParentEnabled(state) ||
+        !shouldTrace(options.spans?.graphql, gqlCtx)
+      ) {
+        return;
+      }
+
+      setExecutionResultAttributes({ ctx: getContext(state), result });
+    },
+
     onParse({ state, context: gqlCtx }) {
       if (
         !isParentEnabled(state) ||
diff --git a/packages/plugins/opentelemetry/src/spans.ts b/packages/plugins/opentelemetry/src/spans.ts
@@ -1,3 +1,4 @@
+import { OnCacheGetHookEventPayload } from '@graphql-hive/gateway-runtime';
 import { defaultPrintFn } from '@graphql-mesh/transport-common';
 import {
   getOperationASTFromDocument,
@@ -12,6 +13,11 @@ import {
   type Context,
   type Tracer,
 } from '@opentelemetry/api';
+import {
+  SEMATTRS_EXCEPTION_MESSAGE,
+  SEMATTRS_EXCEPTION_STACKTRACE,
+  SEMATTRS_EXCEPTION_TYPE,
+} from '@opentelemetry/semantic-conventions';
 import type { ExecutionArgs } from 'graphql';
 import type { GraphQLParams } from 'graphql-yoga';
 import {
@@ -78,6 +84,10 @@ export function setResponseAttributes(ctx: Context, response: Response) {
   const span = trace.getSpan(ctx);
   if (span) {
     span.setAttribute(SEMATTRS_HTTP_STATUS_CODE, response.status);
+    span.setAttribute(
+      'gateway.cache.response_cache',
+      response.status === 304 && response.headers.get('ETag') ? 'hit' : 'miss',
+    );
     span.setStatus({
       code: response.ok ? SpanStatusCode.OK : SpanStatusCode.ERROR,
       message: response.ok ? undefined : response.statusText,
@@ -385,6 +395,46 @@ export function setUpstreamFetchResponseAttributes(input: {
   });
 }
 
+export function recordCacheEvent(
+  event: string,
+  payload: OnCacheGetHookEventPayload,
+) {
+  trace.getActiveSpan()?.addEvent('gateway.cache.' + event, {
+    'gateway.cache.key': payload.key,
+    'gateway.cache.ttl': payload.ttl,
+  });
+}
+
+export function recordCacheError(
+  action: 'read' | 'write',
+  error: Error,
+  payload: OnCacheGetHookEventPayload,
+) {
+  trace.getActiveSpan()?.addEvent('gateway.cache.error', {
+    'gateway.cache.key': payload.key,
+    'gateway.cache.ttl': payload.ttl,
+    'gateway.cache.action': action,
+    [SEMATTRS_EXCEPTION_TYPE]:
+      'code' in error ? (error.code as string) : error.message,
+    [SEMATTRS_EXCEPTION_MESSAGE]: error.message,
+    [SEMATTRS_EXCEPTION_STACKTRACE]: error.stack,
+  });
+}
+
+const responseCacheSymbol = Symbol.for('servedFromResponseCache');
+export function setExecutionResultAttributes(input: {
+  ctx: Context;
+  result?: any; // We don't need a proper type here because we rely on Symbol mark from response cache plugin
+}) {
+  const span = trace.getSpan(input.ctx);
+  if (input.result && span) {
+    span.setAttribute(
+      'gateway.cache.response_cache',
+      input.result[responseCacheSymbol] ? 'hit' : 'miss',
+    );
+  }
+}
+
 export function registerException(ctx: Context | undefined, error: any) {
   const span = ctx && trace.getSpan(ctx);
   if (!span) {
diff --git a/packages/plugins/opentelemetry/tests/useOpenTelemetry.spec.ts b/packages/plugins/opentelemetry/tests/useOpenTelemetry.spec.ts
@@ -532,5 +532,59 @@ describe('useOpenTelemetry', () => {
         children.forEach(spanTree.expectChild);
       }
     });
+
+    it('should have a response cache attribute', async () => {
+      function checkCacheAttributes(attrs: {
+        http: 'hit' | 'miss';
+        operation?: 'hit' | 'miss';
+      }) {
+        const { span: httpSpan } = spanExporter.assertRoot('POST /graphql');
+        const operationSpan = spanExporter.spans.find(({ name }) =>
+          name.startsWith('graphql.operation'),
+        );
+
+        expect(httpSpan.attributes['gateway.cache.response_cache']).toBe(
+          attrs.http,
+        );
+        if (attrs.operation) {
+          expect(operationSpan).toBeDefined();
+          expect(
+            operationSpan!.attributes['gateway.cache.response_cache'],
+          ).toBe(attrs.operation);
+        }
+      }
+      await using gateway = await buildTestGateway({
+        gatewayOptions: {
+          cache: await import('@graphql-mesh/cache-localforage').then(
+            ({ default: Cache }) => new Cache(),
+          ),
+          responseCaching: {
+            session: () => '1',
+          },
+        },
+      });
+      await gateway.query();
+
+      checkCacheAttributes({ http: 'miss', operation: 'miss' });
+
+      spanExporter.reset();
+      await gateway.query();
+
+      checkCacheAttributes({ http: 'miss', operation: 'hit' });
+
+      spanExporter.reset();
+      const response = await gateway.fetch('http://gateway/graphql', {
+        method: 'POST',
+        headers: {
+          'content-type': 'application/json',
+          'If-None-Match':
+            'c2f6fb105ef60ccc99dd6725b55939742e69437d4f85d52bf4664af3799c49fa',
+          'If-Modified-Since': new Date(),
+        },
+      });
+      expect(response.status).toBe(304);
+
+      checkCacheAttributes({ http: 'hit' }); // There is no graphql operation span when cached by HTTP
+    });
   });
 });
diff --git a/packages/plugins/opentelemetry/tests/utils.ts b/packages/plugins/opentelemetry/tests/utils.ts
@@ -81,24 +81,39 @@ export async function buildTestGateway(
 
   return {
     otelPlugin: otelPlugin!,
-    query: async (
-      body: GraphQLParams = {
+    query: async ({
+      shouldReturnErrors,
+      body = {
         query: /* GraphQL */ `
           query {
             hello
           }
         `,
       },
-    ) => {
+    }: {
+      body?: GraphQLParams;
+      shouldReturnErrors?: boolean;
+    } = {}) => {
       const response = await gateway.fetch('http://localhost:4000/graphql', {
         method: 'POST',
         headers: {
           'content-type': 'application/json',
         },
         body: JSON.stringify(body),
       });
-      return response.json();
+
+      const result = await response.json();
+      if (shouldReturnErrors) {
+        expect(result.errors).toBeDefined();
+      } else {
+        if (result.errors) {
+          console.error(result.errors);
+        }
+        expect(result.errors).not.toBeDefined();
+      }
+      return result;
     },
+    fetch: gateway.fetch,
     [Symbol.asyncDispose]: () => {
       diag.disable();
       return stack.disposeAsync();
