newrelic 4.6.0 → 4.9.0

package/lib/shim/promise-shim.js

@@ -0,0 +1,563 @@
+'use strict'
+
+const logger = require('../logger').child({component: 'PromiseShim'})
+const Shim = require('./shim')
+
+/**
+ * A helper class for wrapping promise modules.
+ *
+ * @extends Shim
+ */
+class PromiseShim extends Shim {
+  /**
+   * Constructs a shim associated with the given agent instance, specialized for
+   * instrumenting promise libraries.
+   *
+   * @param {Agent} agent
+   *  The agent this shim will use.
+   *
+   * @param {string} moduleName
+   *  The name of the module being instrumented.
+   *
+   * @param {string} resolvedName
+   *  The full path to the loaded module.
+   *
+   * @see Shim
+   */
+  constructor(agent, moduleName, resolvedName) {
+    super(agent, moduleName, resolvedName)
+    this._logger = logger.child({module: moduleName})
+    this._class = null
+  }
+
+  /**
+   * Grants access to the `Contextualizer` class used by the `PromiseShim` to
+   * propagate context down promise chains.
+   *
+   * @private
+   */
+  static get Contextualizer() {
+    return Contextualizer
+  }
+
+  /**
+   * Sets the class used to indentify promises from the wrapped promise library.
+   *
+   * @param {function} clss - The promise library's class.
+   */
+  setClass(clss) {
+    this._class = clss
+  }
+
+  /**
+   * Checks if the given object is an instance of a promise from the promise
+   * library being wrapped.
+   *
+   * @param {*} obj - The object to check the instance type of.
+   *
+   * @return {bool} True if the provided object is an instance of a promise from
+   *  this promise library.
+   *
+   * @see PromiseShim#setClass
+   */
+  isPromiseInstance(obj) {
+    return !!this._class && obj instanceof this._class
+  }
+
+  /**
+   * Wraps the given properties as constructors for the promise library.
+   *
+   * - `wrapConstructor(nodule, properties)`
+   * - `wrapConstructor(func)`
+   *
+   * It is only necessary to wrap the constructor for the class if there is no
+   * other way to access the executor function. Some libraries expose a separate
+   * method which is called to execute the executor. If that is available, it is
+   * better to wrap that using {@link PromiseShim#wrapExecutorCaller} than to
+   * use this method.
+   *
+   * @param {object|function} nodule
+   *  The source of the properties to wrap, or a single function to wrap.
+   *
+   * @param {string|array.<string>} [properties]
+   *  One or more properties to wrap. If omitted, the `nodule` parameter is
+   *  assumed to be the constructor to wrap.
+   *
+   * @return {object|function} The first parameter to this function, after
+   *  wrapping it or its properties.
+   *
+   * @see PromiseShim#wrapExecutorCaller
+   */
+  wrapConstructor(nodule, properties) {
+    return this.wrapClass(nodule, properties, {
+      pre: function prePromise(shim, Promise, name, args) {
+        // We are expecting one function argument for executor, anything else is
+        // non-standard, do not attempt to wrap. Also do not attempt to wrap if
+        // we are not in a transaction.
+        if (args.length !== 1 || !shim.isFunction(args[0]) || !shim.getActiveSegment()) {
+          return
+        }
+        _wrapExecutorContext(shim, args)
+      },
+      post: function postPromise(shim, Promise, name, args) {
+        // This extra property is added by `_wrapExecutorContext` in the pre step.
+        const executor = args[0]
+        const context = executor && executor.__NR_executorContext
+        if (!context || !shim.isFunction(context.executor)) {
+          return
+        }
+
+        context.promise = this
+        Contextualizer.link(null, this, shim.getSegment())
+        try {
+          // Must run after promise is defined so that `__NR_wrapper` can be set.
+          context.executor.apply(context.self, context.args)
+        } catch (e) {
+          const reject = context.args[1]
+          reject(e)
+        }
+      }
+    })
+  }
+
+  /**
+   * Wraps the given properties as the caller of promise executors.
+   *
+   * - `wrapExecutorCaller(nodule, properties)`
+   * - `wrapExecutorCaller(func)`
+   *
+   * Wrapping the executor caller method directly is preferable to wrapping
+   * the constructor of the promise class.
+   *
+   * @param {object|function} nodule
+   *  The source of the properties to wrap, or a single function to wrap.
+   *
+   * @param {string|array.<string>} [properties]
+   *  One or more properties to wrap. If omitted, the `nodule` parameter is
+   *  assumed to be the function to wrap.
+   *
+   * @return {object|function} The first parameter to this function, after
+   *  wrapping it or its properties.
+   *
+   * @see PromiseShim#wrapConstructor
+   */
+  wrapExecutorCaller(nodule, properties) {
+    return this.wrap(nodule, properties, function executorWrapper(shim, caller) {
+      if (!shim.isFunction(caller) || shim.isWrapped(caller)) {
+        return
+      }
+
+      return function wrappedExecutorCaller(executor) {
+        var parent = shim.getActiveSegment()
+        if (!this || !parent) {
+          return caller.apply(this, arguments)
+        }
+
+        if (!this.__NR_context) {
+          Contextualizer.link(null, this, parent)
+        }
+
+        const args = shim.argsToArray.apply(shim, arguments)
+        _wrapExecutorContext(shim, args)
+        const ret = caller.apply(this, args)
+        const context = args[0].__NR_executorContext
+        context.promise = this
+
+        // Bluebird catches executor errors and auto-rejects when it catches them,
+        // thus we need to do so as well.
+        //
+        // When adding new libraries, make sure to check that they behave the same
+        // way. We may need to enhance the promise spec to handle this variance.
+        try {
+          executor.apply(context.self, context.args)
+        } catch (e) {
+          const reject = context.args[1]
+          reject(e)
+        }
+        return ret
+      }
+    })
+  }
+
+  /**
+   * Wraps the given properties as methods which take is some value other than
+   * a function to call and return a promise.
+   *
+   * - `wrapCast(nodule, properties)`
+   * - `wrapCast(func)`
+   *
+   * Examples of promise cast methods include `Promise.resolve`, `Promise.all`,
+   * and Bluebird's `Promise.delay`. These are static methods which accept some
+   * arbitrary value and return a Promise instance.
+   *
+   * @param {object|function} nodule
+   *  The source of the properties to wrap, or a single function to wrap.
+   *
+   * @param {string|array.<string>} [properties]
+   *  One or more properties to wrap. If omitted, the `nodule` parameter is
+   *  assumed to be the function to wrap.
+   *
+   * @return {object|function} The first parameter to this function, after
+   *  wrapping it or its properties.
+   */
+  wrapCast(nodule, properties) {
+    return this.wrap(nodule, properties, function castWrapper(shim, cast) {
+      if (!shim.isFunction(cast) || shim.isWrapped(cast)) {
+        return
+      }
+
+      return function __NR_wrappedCast() {
+        const segment = shim.getSegment()
+        const prom = cast.apply(this, arguments)
+        if (segment) {
+          Contextualizer.link(null, prom, segment)
+        }
+        return prom
+      }
+    })
+  }
+
+  /**
+   * Wraps the given properties as promise chaining methods.
+   *
+   * - `wrapThen(nodule, properties)`
+   * - `wrapThen(func)`
+   *
+   * NOTE: You must set class used by the library before wrapping then-methods.
+   *
+   * Examples of promise then methods include `Promise#then`, `Promise#finally`,
+   * and Bluebird's `Promise#map`. These are methods which take a function to
+   * execute once the promise resolves and hands back a new promise.
+   *
+   * @param {object|function} nodule
+   *  The source of the properties to wrap, or a single function to wrap.
+   *
+   * @param {string|array.<string>} [properties]
+   *  One or more properties to wrap. If omitted, the `nodule` parameter is
+   *  assumed to be the function to wrap.
+   *
+   * @return {object|function} The first parameter to this function, after
+   *  wrapping it or its properties.
+   *
+   * @see PromiseShim#setClass
+   * @see PromiseShim#wrapCatch
+   */
+  wrapThen(nodule, properties) {
+    return this.wrap(nodule, properties, _wrapThen, [true])
+  }
+
+  /**
+   * Wraps the given properties as rejected promise chaining methods.
+   *
+   * - `wrapCatch(nodule, properties)`
+   * - `wrapCatch(func)`
+   *
+   * NOTE: You must set class used by the library before wrapping catch-methods.
+   *
+   * Promise catch methods differ from then methods in that only one function
+   * will be executed and only if the promise is rejected. Some libraries accept
+   * an additional argument to `Promise#catch` which is usually an error class
+   * to filter rejections by. This wrap method will handle that case.
+   *
+   * @param {object|function} nodule
+   *  The source of the properties to wrap, or a single function to wrap.
+   *
+   * @param {string|array.<string>} [properties]
+   *  One or more properties to wrap. If omitted, the `nodule` parameter is
+   *  assumed to be the function to wrap.
+   *
+   * @return {object|function} The first parameter to this function, after
+   *  wrapping it or its properties.
+   *
+   * @see PromiseShim#setClass
+   * @see PromiseShim#wrapThen
+   */
+  wrapCatch(nodule, properties) {
+    return this.wrap(nodule, properties, _wrapThen, [false])
+  }
+
+  /**
+   * Wraps the given properties as callback-to-promise conversion methods.
+   *
+   * - `wrapPromisify(nodule, properties)`
+   * - `wrapPromisify(func)`
+   *
+   * @param {object|function} nodule
+   *  The source of the properties to wrap, or a single function to wrap.
+   *
+   * @param {string|array.<string>} [properties]
+   *  One or more properties to wrap. If omitted, the `nodule` parameter is
+   *  assumed to be the function to wrap.
+   *
+   * @return {object|function} The first parameter to this function, after
+   *  wrapping it or its properties.
+   */
+  wrapPromisify(nodule, properties) {
+    return this.wrap(nodule, properties, function promisifyWrapper(shim, promisify) {
+      if (!shim.isFunction(promisify) || shim.isWrapped(promisify)) {
+        return
+      }
+
+      return function __NR_wrappedPromisify() {
+        const promisified = promisify.apply(this, arguments)
+        if (typeof promisified !== 'function') {
+          return promisified
+        }
+
+        Object.keys(promisified).forEach(function forEachProperty(prop) {
+          __NR_wrappedPromisified[prop] = promisified[prop]
+        })
+
+        return __NR_wrappedPromisified
+        function __NR_wrappedPromisified() {
+          const segment = shim.getActiveSegment()
+          if (!segment) {
+            return promisified.apply(this, arguments)
+          }
+
+          const prom = shim.applySegment(promisified, segment, true, this, arguments)
+          Contextualizer.link(null, prom, segment)
+          return prom
+        }
+      }
+    })
+  }
+}
+module.exports = PromiseShim
+
+// -------------------------------------------------------------------------- //
+
+/**
+ * @private
+ */
+function _wrapExecutorContext(shim, args) {
+  const context = {
+    executor: args[0],
+    promise: null,
+    self: null,
+    args: null
+  }
+  contextExporter.__NR_executorContext = context
+  args[0] = contextExporter
+
+  function contextExporter(resolve, reject) {
+    context.self = this
+    context.args = shim.argsToArray.apply(shim, arguments)
+    context.args[0] = _wrapResolver(context, resolve)
+    context.args[1] = _wrapResolver(context, reject)
+  }
+}
+
+/**
+ * @private
+ */
+function _wrapResolver(context, fn) {
+  return function wrappedResolveReject(val) {
+    const promise = context.promise
+    if (promise && promise.__NR_context) {
+      promise.__NR_context.getSegment().touch()
+    }
+    fn(val)
+  }
+}
+
+/**
+ * @private
+ */
+function _wrapThen(shim, fn, name, useAllParams) {
+  // Don't wrap non-functions.
+  if (shim.isWrapped(fn) || !shim.isFunction(fn)) {
+    return
+  }
+
+  return function __NR_wrappedThen() {
+    if (!(this instanceof shim._class)) {
+      return fn.apply(this, arguments)
+    }
+
+    const thenSegment = shim.getSegment()
+    const promise = this
+
+    // Wrap up the arguments and execute the real then.
+    let isWrapped = false
+    const args = new Array(arguments.length)
+    for (let i = 0; i < arguments.length; ++i) {
+      args[i] = wrapHandler(arguments[i], i, arguments.length)
+    }
+    const next = fn.apply(this, args)
+
+    // If we got a promise (which we should have), link the parent's context.
+    if (!isWrapped && next instanceof shim._class && next !== promise) {
+      Contextualizer.link(promise, next, thenSegment)
+    }
+    return next
+
+    function wrapHandler(handler, i, length) {
+      if (
+        !shim.isFunction(handler) ||          // Not a function
+        shim.isWrapped(handler)   ||          // Already wrapped
+        (!useAllParams && i !== (length - 1)) // Don't want all and not last
+      ) {
+        isWrapped = shim.isWrapped(handler)
+        return handler
+      }
+
+      return function __NR_wrappedThenHandler() {
+        if (!next || !next.__NR_context) {
+          return handler.apply(this, arguments)
+        }
+
+        let promSegment = next.__NR_context.getSegment()
+        const segment = promSegment || shim.getSegment()
+        if (segment && segment !== promSegment) {
+          next.__NR_context.setSegment(segment)
+          promSegment = segment
+        }
+
+        let ret = null
+        try {
+          ret = shim.applySegment(handler, promSegment, true, this, arguments)
+        } finally {
+          if (ret && typeof ret.then === 'function') {
+            ret = next.__NR_context.continueContext(ret)
+          }
+        }
+        return ret
+      }
+    }
+  }
+}
+
+/**
+ * @private
+ */
+class Context {
+  constructor(segment) {
+    this.segments = [segment]
+  }
+
+  branch() {
+    return this.segments.push(null) - 1
+  }
+}
+
+/**
+ * @private
+ */
+class Contextualizer {
+  constructor(idx, context) {
+    this.parentIdx = -1
+    this.idx = idx
+    this.context = context
+    this.child = null
+  }
+
+  static link(prev, next, segment) {
+    let ctxlzr = prev && prev.__NR_context
+    if (ctxlzr && !ctxlzr.isActive()) {
+      ctxlzr = prev.__NR_context = null
+    }
+
+    if (ctxlzr) {
+      // If prev has one child already, branch the context and update the child.
+      if (ctxlzr.child) {
+        // When the branch-point is the 2nd through nth link in the chain, it is
+        // necessary to track its segment separately so the branches can parent
+        // their segments on the branch-point.
+        if (ctxlzr.parentIdx !== -1) {
+          ctxlzr.idx = ctxlzr.context.branch()
+        }
+
+        // The first child needs to be updated to have its own branch as well. And
+        // each of that child's children must be updated with the new parent index.
+        // This is the only non-constant-time action for linking, but it only
+        // happens with branching promise chains specifically when the 2nd branch
+        // is added.
+        //
+        // Note: This does not account for branches of branches. That may result
+        // in improperly parented segments.
+        let parent = ctxlzr
+        let child = ctxlzr.child
+        const branchIdx = ctxlzr.context.branch()
+        do {
+          child.parentIdx = parent.idx
+          child.idx = branchIdx
+          parent = child
+          child = child.child
+        } while (child)
+
+        // We set the child to something falsey that isn't `null` so we can
+        // distinguish between having no child, having one child, and having
+        // multiple children.
+        ctxlzr.child = false
+      }
+
+      // If this is a branching link then create a new branch for the next promise.
+      // Otherwise, we can just piggy-back on the previous link's spot.
+      const idx = ctxlzr.child === false ? ctxlzr.context.branch() : ctxlzr.idx
+
+      // Create a new context for this next promise.
+      next.__NR_context = new Contextualizer(idx, ctxlzr.context)
+      next.__NR_context.parentIdx = ctxlzr.idx
+
+      // If this was our first child, remember it in case we have a 2nd.
+      if (ctxlzr.child === null) {
+        ctxlzr.child = next.__NR_context
+      }
+    } else if (segment) {
+      // This next promise is the root of a chain. Either there was no previous
+      // promise or the promise was created out of context.
+      next.__NR_context = new Contextualizer(0, new Context(segment))
+    }
+  }
+
+  isActive() {
+    const segments = this.context.segments
+    const segment = segments[this.idx] || segments[this.parentIdx] || segments[0]
+    return segment && segment.transaction.isActive()
+  }
+
+  getSegment() {
+    const segments = this.context.segments
+    let segment = segments[this.idx]
+    if (segment == null) {
+      segment = segments[this.idx] = segments[this.parentIdx] || segments[0]
+    }
+    return segment
+  }
+
+  setSegment(segment) {
+    return this.context.segments[this.idx] = segment
+  }
+
+  toJSON() {
+    // No-op.
+  }
+
+  continueContext(prom) {
+    const self = this
+    const nextContext = prom.__NR_context
+    if (!nextContext) {
+      return prom
+    }
+
+    // If we have `finally`, use that to sneak our context update.
+    if (typeof prom.finally === 'function') {
+      return prom.finally(__NR_continueContext)
+    }
+
+    // No `finally` means we need to hook into resolve and reject individually and
+    // pass through whatever happened.
+    return prom.then(function __NR_thenContext(val) {
+      __NR_continueContext()
+      return val
+    }, function __NR_catchContext(err) {
+      __NR_continueContext()
+      throw err // Re-throwing promise rejection, this is not New Relic's error.
+    })
+
+    function __NR_continueContext() {
+      self.setSegment(nextContext.getSegment())
+    }
+  }
+}