lib/instrumentation/modules/mongodb-core.js

/*
 * Copyright Elasticsearch B.V. and other contributors where applicable.
 * Licensed under the BSD 2-Clause License; you may not use this file except in
 * compliance with the BSD 2-Clause License.
 */

'use strict';

var semver = require('semver');

const { getDBDestination } = require('../context');
var shimmer = require('../shimmer');

var SERVER_FNS = ['insert', 'update', 'remove', 'auth'];
var CURSOR_FNS_FIRST = ['next', '_getmore'];

const firstSpan = Symbol('first-span');

module.exports = function (mongodb, agent, { version, enabled }) {
  if (!enabled) return mongodb;
  if (!semver.satisfies(version, '>=1.2.19 <4')) {
    agent.logger.debug(
      'mongodb-core version %s not supported - aborting...',
      version,
    );
    return mongodb;
  }

  const ins = agent._instrumentation;

  if (mongodb.Server) {
    agent.logger.debug('shimming mongodb-core.Server.prototype.command');
    shimmer.wrap(mongodb.Server.prototype, 'command', wrapCommand);
    agent.logger.debug(
      'shimming mongodb-core.Server.prototype functions: %j',
      SERVER_FNS,
    );
    shimmer.massWrap(mongodb.Server.prototype, SERVER_FNS, wrapQuery);
  }

  if (mongodb.Cursor) {
    agent.logger.debug(
      'shimming mongodb-core.Cursor.prototype functions: %j',
      CURSOR_FNS_FIRST,
    );
    shimmer.massWrap(mongodb.Cursor.prototype, CURSOR_FNS_FIRST, wrapCursor);
  }

  return mongodb;

  function wrapCommand(orig) {
    return function wrappedFunction(ns, cmd) {
      var trans = agent._instrumentation.currTransaction();
      var id = trans && trans.id;
      var span;

      agent.logger.debug(
        'intercepted call to mongodb-core.Server.prototype.command %o',
        { id, ns },
      );

      if (trans && arguments.length > 0) {
        var index = arguments.length - 1;
        var cb = arguments[index];
        if (typeof cb === 'function') {
          var type;
          if (cmd.findAndModify) type = 'findAndModify';
          else if (cmd.createIndexes) type = 'createIndexes';
          else if (cmd.ismaster) type = 'ismaster';
          else if (cmd.count) type = 'count';
          else type = 'command';

          span = ins.createSpan(ns + '.' + type, 'db', 'mongodb', 'query', {
            exitSpan: true,
          });
          if (span) {
            span.setDbContext({ type: 'mongodb', instance: ns });
            arguments[index] = ins.bindFunctionToRunContext(
              ins.currRunContext(),
              wrappedCallback,
            );
          }
        }
      }

      return orig.apply(this, arguments);

      function wrappedCallback(_err, commandResult) {
        agent.logger.debug(
          'intercepted mongodb-core.Server.prototype.command callback %o',
          { id },
        );
        if (commandResult && commandResult.connection) {
          span._setDestinationContext(
            getDBDestination(
              commandResult.connection.host,
              commandResult.connection.port,
            ),
          );
        }
        span.end();
        return cb.apply(this, arguments);
      }
    };
  }

  function wrapQuery(orig, name) {
    return function wrappedFunction(ns) {
      var trans = agent._instrumentation.currTransaction();
      var id = trans && trans.id;
      var span;

      agent.logger.debug(
        'intercepted call to mongodb-core.Server.prototype.%s %o',
        name,
        { id, ns },
      );

      if (trans && arguments.length > 0) {
        var index = arguments.length - 1;
        var cb = arguments[index];
        if (typeof cb === 'function') {
          span = ins.createSpan(ns + '.' + name, 'db', 'mongodb', 'query', {
            exitSpan: true,
          });
          if (span) {
            span.setDbContext({ type: 'mongodb', instance: ns });
            arguments[index] = ins.bindFunctionToRunContext(
              ins.currRunContext(),
              wrappedCallback,
            );
          }
        }
      }

      return orig.apply(this, arguments);

      function wrappedCallback(_err, commandResult) {
        agent.logger.debug(
          'intercepted mongodb-core.Server.prototype.%s callback %o',
          name,
          { id },
        );
        if (commandResult && commandResult.connection) {
          span._setDestinationContext(
            getDBDestination(
              commandResult.connection.host,
              commandResult.connection.port,
            ),
          );
        }
        span.end();
        return cb.apply(this, arguments);
      }
    };
  }
  function wrapCursor(orig, name) {
    return function wrappedFunction() {
      var trans = agent._instrumentation.currTransaction();
      var id = trans && trans.id;
      var span;

      agent.logger.debug(
        'intercepted call to mongodb-core.Cursor.prototype.%s %o',
        name,
        { id },
      );

      if (trans && arguments.length > 0) {
        var cb = arguments[0];
        if (typeof cb === 'function') {
          if (name !== 'next' || !this[firstSpan]) {
            var spanName = `${this.ns}.${this.cmd.find ? 'find' : name}`;
            span = ins.createSpan(spanName, 'db', 'mongodb', 'query', {
              exitSpan: true,
            });
          }
          if (span) {
            span.setDbContext({ type: 'mongodb', instance: this.ns });
            // Limitation: Currently not getting destination address/port for
            // cursor calls.
            arguments[0] = ins.bindFunctionToRunContext(
              ins.currRunContext(),
              wrappedCallback,
            );
            if (name === 'next') {
              this[firstSpan] = true;
            }
          }
        }
      }

      return orig.apply(this, arguments);

      function wrappedCallback() {
        agent.logger.debug(
          'intercepted mongodb-core.Cursor.prototype.%s callback %o',
          name,
          { id },
        );
        span.end();
        return cb.apply(this, arguments);
      }
    };
  }
};