appdynamics/lib/core/context.js

Performance Profiler and Monitor

/*
Copyright (c) AppDynamics, Inc., and its affiliates
2018
All Rights Reserved
 */
'use strict';

function Context(agent) {
  this.agent = agent;
}

exports.Context = Context;

const EE_LISTENER_METHODS = [
  'addListener',
  'on',
  'once',
  'prependListener',
  'prependOnceListener',
];

Context.prototype.init = function () {
  this.enabled = false;
  this._kListeners = Symbol('kListeners');
  if (this.agent.opts.clsDisabled) {
    this.agent.logger.info('CLS disabled by configuration');
    return;
  }

  try {
    const { AsyncLocalStorage } = require("async_hooks");
    this.enabled = true;
    this.ns = new AsyncLocalStorage();
  } catch (e) {
    //this.agent.logger.info('CLS not supported for current node version');
  }
};

// run fn in CLS context if available; otherwise just run it without context
Context.prototype.run = function runInContext(fn, req, res) {
  var ns = this.ns;

  if (ns) {
    ns.run(this.active(), () => {
      if (req) this.bind(req);
      if (res) this.bind(res);
      fn.call(this, req, res);
    });
  } else {
    fn.call(this, req, res);
  }
};

// run fn in CLS context if available; otherwise just run it without context
Context.prototype.runGeneric = function runInContext(fn, args) {
  var ns = this.ns;

  if (ns) {
    ns.run(this.active(), () => {
      fn.apply(this, args);
    });
  } else {
    fn.apply(this, args);
  }
};

// Event emitter binding code is mostly based on
// Opentelemetry context manager
// https://github.com/open-telemetry/opentelemetry-js/blob/main/packages/opentelemetry-context-async-hooks/src/AbstractAsyncHooksContextManager.ts
Context.prototype.bindEmitter = function(ee) {
  if(!this.ns) return ee;

  EE_LISTENER_METHODS.forEach(methodName => {
    if(ee[methodName] === undefined) return;
    ee[methodName] = this.__wrapAddListener(ee, ee[methodName]);
  });

  if(typeof ee.removeListener === 'function') {
    ee.removeListener = this.__wrapRemoveListener(ee, ee.removeListener);
  }

  if(typeof ee.off === 'function') {
    ee.off = this.__wrapRemoveListener(ee, ee.off);
  }

  if(typeof ee.removeAllListeners === 'function') {
    ee.removeAllListeners = this.__wrapRemoveAllListeners(ee, ee.removeAllListeners);
  }
  return ee;
};


Context.prototype.bind = function(fn) {
  var self = this;
  if (self.ns) {
    var context = self.active();
    return function(...args) {
      return self.ns.run(context, () => {
        fn.apply(this, args);
      });
    };
  } else {
    return fn;
  }
};

// set a value in CLS context, if available
Context.prototype.set = function set(key, value) {
  if (this.ns) return this.active().set(key, value);
};

// get a value from CLS, if available
Context.prototype.get = function get(key) {
  if(this.ns) return this.active().get(key);
};

// get a value from CLS, if available
Context.prototype.active = function get() {
  if(this.ns && this.ns.getStore()) {
    return this.ns.getStore();
  } else {
    return new Map();
  }
};

Context.prototype.__wrapAddListener = function(ee, original) {
  var self = this;
  return function(event, listener) {

    if (self._wrapped) {
      return original.call(this, event, listener);
    }

    let state = self.__getState(ee);
    if(state === undefined) {
      state = self.__createState(ee);
    }

    let listeners = state[event];
    if(listeners === undefined) {
      listeners = new WeakMap();
      state[event] = listeners;
    }

    const wrappedListener = self.bind(listener);
    listeners.set(listener, wrappedListener);
    self._wrapped = true;
    try {
      return original.call(this, event, wrappedListener);
    } finally {
      self._wrapped = false;
    }
  };
};

Context.prototype.__wrapRemoveListener = function(ee, original) {
  var self = this;
  return function(event, listener) {
    const state  = self.__getState(ee);
    const events = state && state[event];
    if(events === undefined) {
      return original.call(this, event, listener);
    }
    const wrappedListener = events.get(listener);
    return original.call(this, event, wrappedListener || listener);
  };
};

Context.prototype.__wrapRemoveAllListeners = function(ee, original) {
  const self = this;
  return function(event) {
    const state = self.__getState(ee);
    if(state != undefined) {
      if(arguments.length === 0) {
        self.__createState(ee);
      } else if (state[event] !== undefined) {
        delete state[event];
      }
    }
    return original.apply(this, arguments);
  };
};

Context.prototype.__createState = function(ee) {
  const state = Object.create(null);
  ee[this._kListeners] = state;
  return state;
};

Context.prototype.__getState = function(ee) {
  return ee[this._kListeners];
};