appdynamics/lib/probes/rabbitmq-entry-probe.js

Performance Profiler and Monitor

'use strict';

exports.RabbitMQEntryProbe = RabbitMQEntryProbe;

const CHANNEL_MESSAGES = Symbol('appdynamics.channel_messages');
const ACTIVE_BT = Symbol('appdynamics.active_bt');

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

RabbitMQEntryProbe.prototype.init = function () {

};

RabbitMQEntryProbe.prototype.attach = function (obj, moduleName) {
  var self = this;
  if(moduleName == 'amqplib/lib/callback_model') {
    self.wrapMethods(obj);
  } else if(moduleName == 'amqplib/lib/channel_model') {
    self.wrapMethods(obj);
  }
};

RabbitMQEntryProbe.prototype.wrapMethods = function(obj) {
  if(!obj || !obj.Channel) return;
  var self = this;
  var proxy = self.agent.proxy;

  proxy.before(obj.Channel.prototype, 'consume', function(obj, args) {
    self.handleConsume(obj, args);
  }, false, false, self.agent.thread.current());
  proxy.after(obj.Channel.prototype, 'ack', function(obj, args, ret) {
    self.handleAck(false, "ACK", obj, args, ret);
  });
  proxy.after(obj.Channel.prototype, 'nack', function(obj, args, ret) {
    self.handleAck(true, "NACK", obj, args, ret);
  });
  proxy.after(obj.Channel.prototype, 'reject', function(obj, args, ret) {
    self.handleAck(true, "REJECT", obj, args, ret);
  });
  proxy.after(obj.Channel.prototype, 'ackAll', function(obj) {
    self.handleAckAll(obj);
  });
  proxy.after(obj.Channel.prototype, 'emit', function(obj, args, ret) {
    self.handleEmit(obj, args, ret);
  });
};

RabbitMQEntryProbe.prototype.endTransaction = function(message, isRejected, operation) {
  var self = this;
  const transaction = message[ACTIVE_BT];
  if (!transaction) return;
  if(isRejected != false) {
    var errorObject = (operation != "CHANNEL_ERROR" && operation != "CHANNEL_CLOSED")
                       ? `${operation} called on message` : `${operation} called on channel`;
    transaction.error = errorObject;
  } 
  self.agent.profiler.endTransaction(transaction.time, transaction);
  delete message[ACTIVE_BT];
}; 

RabbitMQEntryProbe.prototype.handleEmit = function(channel, args) {
  var self = this;
  var eventName = args[0];
  if(eventName === 'close') {
    self.endAllTransactionsOnChannel(channel, true, "CHANNEL_CLOSED");
  } else if (eventName === 'error') {
    self.endAllTransactionsOnChannel(channel, true, "CHANNEL_ERROR");
  }
};

RabbitMQEntryProbe.prototype.handleAckAll = function(channel) {
  var self = this;
  self.endAllTransactionsOnChannel(channel, false, "ACK");
};

RabbitMQEntryProbe.prototype.endAllTransactionsOnChannel = function(channel, isRejected, operation) {
  var self = this;
  const channelMessages = channel[CHANNEL_MESSAGES] ? channel[CHANNEL_MESSAGES] : [];
  channelMessages.forEach((message) => {
    self.endTransaction(message, false, operation);
  });
  Object.defineProperty(channel, CHANNEL_MESSAGES, {
    value: [],
    enumerable: false,
    configurable: true,
  });  
};

RabbitMQEntryProbe.prototype.handleAck = function(isRejected, operation, channel, args) {
  var self = this;
  if (!args || args.length < 1) return;

  const message = args[0];
  const allUpTo = args[1] || false;

  const channelMessages = channel[CHANNEL_MESSAGES] ? channel[CHANNEL_MESSAGES] : [];
  const msgIndex = channelMessages.findIndex((msgDetails) => msgDetails === message);
  if (msgIndex < 0) {
    // could be user double ack the same message? do nothing.
  } else if (operation != "REJECT" && allUpTo) {
    for (let i = 0; i <= msgIndex; i++) {
      self.endTransaction(channelMessages[i], isRejected, operation);
    }
    channelMessages.splice(0, msgIndex + 1);
  } else {
    self.endTransaction(message, isRejected, operation);
    channelMessages.splice(msgIndex, 1);
  }
};

RabbitMQEntryProbe.prototype.handleConsume = function(obj, args) {
  var self = this;
  var profiler = self.agent.profiler;
  if(!args || args.length < 2) {
    return;
  }

  // args = queue, callback, options
  args[1] = createWrapperFunc(obj, args[1]);

  var options = undefined;
  if(args.length >= 3) {
    options = args[2];
  } 

  function createWrapperFunc(channel, original) {
    return function() {
      var args = arguments;
      self.agent.context.runGeneric(handlerFunc, args);

      function handlerFunc() {
        var msg = undefined;
        if(arguments && arguments.length >= 1) {
          msg = arguments[0];
        }

        if (!msg) return original.apply(this, arguments);

        var exchange = "default";
        var routingKey = "default";
        var headers = {};
        if(msg && msg.properties && msg.properties.headers) {
          headers = msg.properties.headers;
        }

        if(msg && msg.fields && msg.fields.exchange) {
          exchange = msg.fields.exchange;
        }

        if(msg && msg.fields && msg.fields.routingKey) {
          routingKey = msg.fields.routingKey;
        }

        var req = {url: `${exchange}/${routingKey}`, method: null, headers: headers};
        var time = profiler.time(true);
        var transaction = profiler.startTransaction(time, req, 'NODEJS_WEB');
        self.agent.context.set('threadId', transaction.threadId);

        if (options.noAck) {
          profiler.endTransaction(time, transaction);
        } else {
          if (!channel.hasOwnProperty(CHANNEL_MESSAGES)) {
            Object.defineProperty(channel, CHANNEL_MESSAGES, {
              value: [],
              enumerable: false,
              configurable: true,
            });
          }

          Object.defineProperty(msg, ACTIVE_BT, {
            value: transaction,
            enumerable: false,
            configurable: true,
          });
          channel[CHANNEL_MESSAGES].push(msg);
        }

        return original.apply(this, arguments);
      }
    };
  }
};