remove async storage from graphql

Author: wconti27

packages/datadog-instrumentations/src/graphql.js

@@ -2,8 +2,7 @@
 
 const {
   addHook,
-  channel,
-  AsyncResource
+  channel
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -89,10 +88,8 @@ function wrapParse (parse) {
       return parse.apply(this, arguments)
     }
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-    return asyncResource.runInAsyncScope(() => {
-      parseStartCh.publish()
+    const ctx = {}
+    return parseStartCh.runStores(ctx, () => {
       let document
       try {
         document = parse.apply(this, arguments)
@@ -107,11 +104,12 @@ function wrapParse (parse) {
         return document
       } catch (err) {
         err.stack
-        parseErrorCh.publish(err)
+        ctx.error = err
+        parseErrorCh.publish(ctx)
 
         throw err
       } finally {
-        parseFinishCh.publish({ source, document, docSource: documentSources.get(document) })
+        parseFinishCh.publish({ source, document, docSource: documentSources.get(document), ...ctx })
       }
     })
   }
@@ -123,25 +121,24 @@ function wrapValidate (validate) {
       return validate.apply(this, arguments)
     }
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-    return asyncResource.runInAsyncScope(() => {
-      validateStartCh.publish({ docSource: documentSources.get(document), document })
-
+    const ctx = { docSource: documentSources.get(document), document }
+    return validateStartCh.runStores(ctx, () => {
       let errors
       try {
         errors = validate.apply(this, arguments)
         if (errors && errors[0]) {
-          validateErrorCh.publish(errors && errors[0])
+          ctx.error = errors && errors[0]
+          validateErrorCh.publish(ctx)
         }
         return errors
       } catch (err) {
         err.stack
-        validateErrorCh.publish(err)
+        ctx.error = err
+        validateErrorCh.publish(ctx)
 
         throw err
       } finally {
-        validateFinishCh.publish({ document, errors })
+        validateFinishCh.publish({ errors, ...ctx })
       }
     })
   }
@@ -155,15 +152,15 @@ function wrapExecute (execute) {
         return exe.apply(this, arguments)
       }
 
-      const asyncResource = new AsyncResource('bound-anonymous-fn')
-      return asyncResource.runInAsyncScope(() => {
-        const args = normalizeArgs(arguments, defaultFieldResolver)
-        const schema = args.schema
-        const document = args.document
-        const source = documentSources.get(document)
-        const contextValue = args.contextValue
-        const operation = getOperation(document, args.operationName)
+      const args = normalizeArgs(arguments, defaultFieldResolver)
+      const schema = args.schema
+      const document = args.document
+      const source = documentSources.get(document)
+      const contextValue = args.contextValue
+      const operation = getOperation(document, args.operationName)
 
+      const ctx = { operation, args, docSource: documentSources.get(document) }
+      return startExecuteCh.runStores(ctx, () => {
         if (contexts.has(contextValue)) {
           return exe.apply(this, arguments)
         }
@@ -173,26 +170,22 @@ function wrapExecute (execute) {
           wrapFields(schema._mutationType)
         }
 
-        startExecuteCh.publish({
-          operation,
-          args,
-          docSource: documentSources.get(document)
-        })
-
-        const context = { source, asyncResource, fields: {}, abortController: new AbortController() }
+        const context = { source, fields: {}, abortController: new AbortController(), ...ctx }
 
         contexts.set(contextValue, context)
 
-        return callInAsyncScope(exe, asyncResource, this, arguments, context.abortController, (err, res) => {
+        return callInAsyncScope(exe, this, arguments, context.abortController, (err, res) => {
           if (finishResolveCh.hasSubscribers) finishResolvers(context)
 
           const error = err || (res && res.errors && res.errors[0])
 
           if (error) {
-            executeErrorCh.publish(error)
+            ctx.error = error
+            executeErrorCh.publish(ctx)
           }
 
-          finishExecuteCh.publish({ res, args, context })
+          ctx.res = res
+          finishExecuteCh.publish(ctx)
         })
       })
     }
@@ -211,8 +204,8 @@ function wrapResolve (resolve) {
 
     const field = assertField(context, info, args)
 
-    return callInAsyncScope(resolve, field.asyncResource, this, arguments, context.abortController, (err) => {
-      updateFieldCh.publish({ field, info, err })
+    return callInAsyncScope(resolve, this, arguments, context.abortController, (err) => {
+      updateFieldCh.publish({ field, info, err, ...field.ctx })
     })
   }
 
@@ -221,32 +214,30 @@ function wrapResolve (resolve) {
   return resolveAsync
 }
 
-function callInAsyncScope (fn, aR, thisArg, args, abortController, cb) {
+function callInAsyncScope (fn, thisArg, args, abortController, cb) {
   cb = cb || (() => {})
 
-  return aR.runInAsyncScope(() => {
-    if (abortController?.signal.aborted) {
-      cb(null, null)
-      throw new AbortError('Aborted')
-    }
+  if (abortController?.signal.aborted) {
+    cb(null, null)
+    throw new AbortError('Aborted')
+  }
 
-    try {
-      const result = fn.apply(thisArg, args)
-      if (result && typeof result.then === 'function') {
-        // bind callback to this scope
-        result.then(
-          aR.bind(res => cb(null, res)),
-          aR.bind(err => cb(err))
-        )
-      } else {
-        cb(null, result)
-      }
-      return result
-    } catch (err) {
-      cb(err)
-      throw err
+  try {
+    const result = fn.apply(thisArg, args)
+    if (result && typeof result.then === 'function') {
+      // bind callback to this scope
+      result.then(
+        res => cb(null, res),
+        err => cb(err)
+      )
+    } else {
+      cb(null, result)
     }
-  })
+    return result
+  } catch (err) {
+    cb(err)
+    throw err
+  }
 }
 
 function pathToArray (path) {
@@ -271,27 +262,15 @@ function assertField (context, info, args) {
 
   if (!field) {
     const parent = getParentField(context, path)
-
-    // we want to spawn the new span off of the parent, not a new async resource
-    parent.asyncResource.runInAsyncScope(() => {
-      /* this child resource will run a branched scope off of the parent resource, which
-      accesses the parent span from the storage unit in its own scope */
-      const childResource = new AsyncResource('bound-anonymous-fn')
-
-      childResource.runInAsyncScope(() => {
-        startResolveCh.publish({
-          info,
-          context,
-          args
-        })
-      })
-
-      field = fields[pathString] = {
-        parent,
-        asyncResource: childResource,
-        error: null
-      }
-    })
+    // we need to pass the parent span to the field if it exists for correct span parenting
+    // of nested fields
+    const ctx = { info, context, args, childOf: parent?.ctx?.currentStore?.span }
+    startResolveCh.publish(ctx)
+    field = fields[pathString] = {
+      parent,
+      error: null,
+      ctx
+    }
   }
 
   return field
@@ -349,20 +328,16 @@ function wrapFieldType (field) {
 function finishResolvers ({ fields }) {
   Object.keys(fields).reverse().forEach(key => {
     const field = fields[key]
-    const asyncResource = field.asyncResource
-    asyncResource.runInAsyncScope(() => {
-      if (field.error) {
-        resolveErrorCh.publish(field.error)
-      }
-      finishResolveCh.publish(field.finishTime)
-    })
+    const ctx = { field, finishTime: field.finishTime, ...field.ctx }
+    if (field.error) {
+      ctx.error = field.error
+      resolveErrorCh.publish(ctx)
+    }
+    finishResolveCh.publish(ctx)
   })
 }
 
-addHook({ name: '@graphql-tools/executor', file: 'cjs/execution/execute.js', versions: ['>=0.0.14'] }, execute => {
-  shimmer.wrap(execute, 'execute', wrapExecute(execute))
-  return execute
-})
+addHook({ name: '@graphql-tools/executor', file: 'cjs/execution/execute.js', versions: ['>=0.0.14'] }, execute => {})
 
 addHook({ name: 'graphql', file: 'execution/execute.js', versions: ['>=0.10'] }, execute => {
   shimmer.wrap(execute, 'execute', wrapExecute(execute))

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

@@ -11,7 +11,9 @@ class GraphQLExecutePlugin extends TracingPlugin {
   static get type () { return 'graphql' }
   static get kind () { return 'server' }
 
-  start ({ operation, args, docSource }) {
+  bindStart (ctx) {
+    const { operation, args, docSource } = ctx
+
     const type = operation && operation.operation
     const name = operation && operation.name && operation.name.value
     const document = args.document
@@ -27,20 +29,25 @@ class GraphQLExecutePlugin extends TracingPlugin {
         'graphql.operation.name': name,
         'graphql.source': source
       }
-    })
+    }, ctx)
 
     addVariableTags(this.config, span, args.variableValues)
+
+    return ctx.currentStore
   }
 
-  finish ({ res, args }) {
-    const span = this.activeSpan
+  finish (ctx) {
+    const { res, args } = ctx
+    const span = ctx?.currentStore?.span || this.activeSpan
     this.config.hooks.execute(span, args, res)
     if (res?.errors) {
       for (const err of res.errors) {
         extractErrorIntoSpanEvent(this._tracerConfig, span, err)
       }
     }
-    super.finish()
+    super.finish(ctx)
+
+    return ctx.parentStore
   }
 }
 

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

@@ -6,24 +6,29 @@ class GraphQLParsePlugin extends TracingPlugin {
   static get id () { return 'graphql' }
   static get operation () { return 'parser' }
 
-  start () {
+  bindStart (ctx) {
     this.startSpan('graphql.parse', {
       service: this.config.service,
       type: 'graphql',
       meta: {}
-    })
+    }, ctx)
+
+    return ctx.currentStore
   }
 
-  finish ({ source, document, docSource }) {
-    const span = this.activeSpan
+  finish (ctx) {
+    const { source, document, docSource } = ctx
+    const span = ctx?.currentStore?.span || this.activeSpan
 
     if (this.config.source && document) {
       span.setTag('graphql.source', docSource)
     }
 
     this.config.hooks.parse(span, source, document)
 
-    super.finish()
+    super.finish(ctx)
+
+    return ctx.parentStore
   }
 }
 

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

@@ -9,7 +9,8 @@ class GraphQLResolvePlugin extends TracingPlugin {
   static get id () { return 'graphql' }
   static get operation () { return 'resolve' }
 
-  start ({ info, context, args }) {
+  start (ctx) {
+    const { info, context, args, childOf } = ctx
     const path = getPath(info, this.config)
 
     if (!shouldInstrument(this.config, path)) return
@@ -35,14 +36,15 @@ class GraphQLResolvePlugin extends TracingPlugin {
     const span = this.startSpan('graphql.resolve', {
       service: this.config.service,
       resource: `${info.fieldName}:${info.returnType}`,
+      childOf,
       type: 'graphql',
       meta: {
         'graphql.field.name': info.fieldName,
         'graphql.field.path': computedPathString,
         'graphql.field.type': info.returnType.name,
         'graphql.source': source
       }
-    })
+    }, ctx)
 
     if (fieldNode && this.config.variables && fieldNode.arguments) {
       const variables = this.config.variables(info.variableValues)
@@ -58,17 +60,21 @@ class GraphQLResolvePlugin extends TracingPlugin {
     if (this.resolverStartCh.hasSubscribers) {
       this.resolverStartCh.publish({ context, resolverInfo: getResolverInfo(info, args) })
     }
+
+    return ctx.currentStore
   }
 
   constructor (...args) {
     super(...args)
 
-    this.addTraceSub('updateField', ({ field, info, err }) => {
+    this.addTraceSub('updateField', (ctx) => {
+      const { field, info, err } = ctx
+
       const path = getPath(info, this.config)
 
       if (!shouldInstrument(this.config, path)) return
 
-      const span = this.activeSpan
+      const span = ctx?.currentStore?.span || this.activeSpan
       field.finishTime = span._getTime ? span._getTime() : 0
       field.error = field.error || err
     })
@@ -81,8 +87,13 @@ class GraphQLResolvePlugin extends TracingPlugin {
     super.configure(config.depth === 0 ? false : config)
   }
 
-  finish (finishTime) {
-    this.activeSpan.finish(finishTime)
+  bindFinish (ctx) {
+    const { finishTime } = ctx
+
+    const span = ctx?.currentStore?.span || this.activeSpan
+    span.finish(finishTime)
+
+    return ctx.parentStore
   }
 }