Refactor graphql plugin

CONTRIBUTING.md

@@ -407,6 +407,8 @@ You can also run the tests for multiple plugins at once by separating them with
 PLUGINS="amqplib|bluebird" yarn test:plugins
 ```
 
+The necessary shell commands for the setup can also be executed at once by the `yarn env` script.
+
 ### Other Unit Tests
 
 There are several types of unit tests, for various types of components. The

index.d.ts

@@ -2048,16 +2048,31 @@ declare namespace tracer {
     */
    interface google_genai extends Integration {}
 
-   /** @hidden */
-   interface ExecutionArgs {
-     schema: any,
-     document: any,
-     rootValue?: any,
-     contextValue?: any,
-     variableValues?: any,
-     operationName?: string,
-     fieldResolver?: any,
-     typeResolver?: any,
+    /** @hidden - the `graphql.ExecutionArgs` passed to the `execute` call */
+    interface ExecutionArgs {
+      schema: any,
+      document: any,
+      rootValue?: any,
+      contextValue?: any,
+      variableValues?: any,
+      operationName?: string,
+      fieldResolver?: any,
+      typeResolver?: any,
+    }
+
+    interface FieldContext {
+      /** The `graphql.GraphQLResolveInfo` for the resolver call */
+      info: any;
+      /** The arguments passed to the resolver */
+      args: any;
+      /** The error thrown by the resolver, if any */
+      error: null | Error;
+      /** The result returned by the resolver, if any */
+      res: unknown;
+      /** The field context from the resolver of the parent field (a level up on the path) */
+      parentField: FieldContext | null;
+      /** The nesting depth of the field in the query */
+      depth: number;
     }
 
     /**
@@ -2138,6 +2153,7 @@ declare namespace tracer {
         execute?: (span?: Span, args?: ExecutionArgs, res?: any) => void;
         validate?: (span?: Span, document?: any, errors?: any) => void;
         parse?: (span?: Span, source?: any, document?: any) => void;
+        resolve?: (span?: Span, field: FieldContext) => void;
       }
     }
 

packages/datadog-instrumentations/src/graphql.js

@@ -25,7 +25,6 @@ const executeErrorCh = channel('apm:graphql:execute:error')
 // resolve channels
 const startResolveCh = channel('apm:graphql:resolve:start')
 const finishResolveCh = channel('apm:graphql:resolve:finish')
-const updateFieldCh = channel('apm:graphql:resolve:updateField')
 const resolveErrorCh = channel('apm:graphql:resolve:error')
 
 // parse channels
@@ -38,13 +37,6 @@ const validateStartCh = channel('apm:graphql:validate:start')
 const validateFinishCh = channel('apm:graphql:validate:finish')
 const validateErrorCh = channel('apm:graphql:validate:error')
 
-class AbortError extends Error {
-  constructor (message) {
-    super(message)
-    this.name = 'AbortError'
-  }
-}
-
 const types = new Set(['query', 'mutation', 'subscription'])
 
 function getOperation (document, operationName) {
@@ -171,8 +163,9 @@ function wrapExecute (execute) {
         args,
         docSource: documentSources.get(document),
         source,
-        fields: {},
-        abortController: new AbortController(),
+        fields: new WeakMap(), // fields keyed by their Path object
+        finalizations: [], // fields whose `.finalize()` method needs to be invoked before finishing execution
+        abortController: new AbortController(), // allow startExecuteCh/startResolveCh subscribers to block execution
       }
 
       return startExecuteCh.runStores(ctx, () => {
@@ -183,8 +176,10 @@ function wrapExecute (execute) {
 
         contexts.set(contextValue, ctx)
 
-        return callInAsyncScope(exe, this, arguments, ctx.abortController, (err, res) => {
-          if (finishResolveCh.hasSubscribers) finishResolvers(ctx)
+        return callInAsyncScope(exe, this, arguments, ctx.abortController.signal, (err, res) => {
+          if (ctx.finalizations.length) {
+            finalizeResolvers(ctx.finalizations)
+          }
 
           const error = err || (res && res.errors && res.errors[0])
 
@@ -211,13 +206,23 @@ function wrapResolve (resolve) {
 
     if (!ctx) return resolve.apply(this, arguments)
 
-    const field = assertField(ctx, info, args)
+    const field = createField(ctx, info, args)
+    ctx.fields.set(info.path, field)
+
+    startResolveCh.publish(field)
 
-    return callInAsyncScope(resolve, this, arguments, ctx.abortController, (err) => {
-      field.ctx.error = err
-      field.ctx.info = info
-      field.ctx.field = field
-      updateFieldCh.publish(field.ctx)
+    if (field.finalize) {
+      // register for `.finalize()` invocation before execution finishes
+      ctx.finalizations.push(field)
+    }
+
+    return callInAsyncScope(resolve, this, arguments, ctx.abortController.signal, (err, res) => {
+      if (err) {
+        field.error = err
+        resolveErrorCh.publish(field)
+      }
+      field.res = res
+      finishResolveCh.publish(field)
     })
   }
 
@@ -226,16 +231,15 @@ function wrapResolve (resolve) {
   return resolveAsync
 }
 
-function callInAsyncScope (fn, thisArg, args, abortController, cb) {
-  cb = cb || (() => {})
-
-  if (abortController?.signal.aborted) {
+function callInAsyncScope (fn, thisArg, args, abortSignal, cb) {
+  if (abortSignal.aborted) {
     cb(null, null)
-    throw new AbortError('Aborted')
+    throw abortSignal.reason
   }
 
+  let result
   try {
-    const result = fn.apply(thisArg, args)
+    result = fn.apply(thisArg, args)
     if (result && typeof result.then === 'function') {
       return result.then(
         res => {
@@ -248,44 +252,38 @@ function callInAsyncScope (fn, thisArg, args, abortController, cb) {
         }
       )
     }
-    cb(null, result)
-    return result
   } catch (err) {
     cb(err)
     throw err
-  }
+  } // else
+  cb(null, result)
+  return result
 }
 
-function pathToArray (path) {
-  const flattened = []
-  let curr = path
-  while (curr) {
-    flattened.push(curr.key)
-    curr = curr.prev
+function createField (rootCtx, info, args) {
+  const parentField = getParentField(rootCtx, info.path)
+  return {
+    rootCtx,
+    info,
+    args,
+    error: null,
+    res: null,
+    parentField,
+    depth: parentField ? parentField.depth + 1 : 1,
+    finalize: null, // sometimes populated by GraphQLResolvePlugin in `resolve:start` handler
+    finishTime: 0, // populated by GraphQLResolvePlugin in `resolve:updateField` handler
+    // currentStore, parentStore - sometimes populated by GraphQLResolvePlugin.startSpan in `resolve:start` handler
   }
-  return flattened.reverse()
 }
 
-function assertField (rootCtx, info, args) {
-  const pathInfo = info && info.path
-
-  const path = pathToArray(pathInfo)
-
-  const pathString = path.join('.')
-  const fields = rootCtx.fields
-
-  let field = fields[pathString]
-
-  if (!field) {
-    const fieldCtx = { info, rootCtx, args }
-    startResolveCh.publish(fieldCtx)
-    field = fields[pathString] = {
-      error: null,
-      ctx: fieldCtx,
-    }
+function getParentField (rootCtx, path) {
+  let curr = path.prev
+  // Skip segments in (nested) lists
+  // Could also be done by `while (!rootCtx.fields.has(curr))`
+  while (curr && typeof curr.key === 'number') {
+    curr = curr.prev
   }
-
-  return field
+  return rootCtx.fields.get(curr) ?? null
 }
 
 function wrapFields (type) {
@@ -320,16 +318,9 @@ function wrapFieldType (field) {
   wrapFields(unwrappedType)
 }
 
-function finishResolvers ({ fields }) {
-  for (const key of Object.keys(fields).reverse()) {
-    const field = fields[key]
-    field.ctx.finishTime = field.finishTime
-    field.ctx.field = field
-    if (field.error) {
-      field.ctx.error = field.error
-      resolveErrorCh.publish(field.ctx)
-    }
-    finishResolveCh.publish(field.ctx)
+function finalizeResolvers (contexts) {
+  for (const fieldCtx of contexts) {
+    fieldCtx.finalize()
   }
 }
 

packages/datadog-plugin-graphql/src/execute.js

@@ -47,7 +47,7 @@ class GraphQLExecutePlugin extends TracingPlugin {
     }
     super.finish(ctx)
 
-    return ctx.parentStore
+    // return ctx.parentStore
   }
 }
 

packages/datadog-plugin-graphql/src/index.js

@@ -61,12 +61,14 @@ function getVariablesFilter (config) {
 
 const noop = () => {}
 
+// FIXME: should not be necessary given `TracingPlugin.configure` already does this defaulting
 function getHooks ({ hooks }) {
   const execute = hooks?.execute ?? noop
   const parse = hooks?.parse ?? noop
   const validate = hooks?.validate ?? noop
+  const resolve = hooks?.resolve ?? noop
 
-  return { execute, parse, validate }
+  return { execute, parse, validate, resolve }
 }
 
 module.exports = GraphQLPlugin