atatus-nodejs/lib/instrumentation/modules/amqplib/model.js

Atatus APM agent for Node.js

'use strict';

const semver = require('semver');

const constants = require('../../../constants');
const shimmer = require('../../shimmer');
const { redactKeysFromObject } = require('../../../filters/sanitize-field-names');
const { shouldIgnoreTopic } = require('./utils');

const NAME = 'RabbitMQ'
const TYPE = 'messaging';
const SUBTYPE = 'rabbitmq';

const CHANNEL_METHODS = [
    'close',
    'open',
    'assertQueue',
    'checkQueue',
    'deleteQueue',
    'bindQueue',
    'unbindQueue',
    'assertExchange',
    'checkExchange',
    'deleteExchange',
    'bindExchange',
    'unbindExchange',
    'purgeQueue',
    'cancel',
    'prefetch',
    'recover'
]

module.exports = function instrumentAMQPModel(mod, agent, { version, enabled }) {

    if (!enabled || !semver.satisfies(version, '>=0.5')) {
        return mod;
    }
    const config = agent._conf;
    const ins = agent._instrumentation;

    agent.logger.debug('shimming AMQP Channel Model');

    const proto = mod.Channel.prototype;

    shimmer.wrap(proto, 'consume', function (orgConsume) {
        return function wrappedConsume() {
            const [topic, cb] = arguments;
            agent.logger.debug(`Consuming from queue: ${topic}`);

            if (shouldIgnoreTopic(topic, config)) {
                return orgConsume.apply(this, arguments)
            }

            arguments[1] = function wrappedCallback() {
                const cbArgs = arguments

                const [msg] = cbArgs

                const traceHeaders = msg.properties.headers || {};
                const traceparent = traceHeaders['traceparent'];
                const tracestate = traceHeaders['tracestate'];

                const opts = {};

                if (typeof traceparent === 'string') {
                    opts.childOf = traceparent;
                } else if (traceparent instanceof Buffer) {
                    opts.childOf = traceparent.toString('utf-8');
                }

                if (typeof tracestate === 'string') {
                    opts.tracestate = tracestate;
                } else if (tracestate instanceof Buffer) {
                    opts.tracestate = tracestate.toString('utf-8');
                }


                const trans = ins.startTransaction(
                    `${NAME} RECEIVE from ${topic}`,
                    TYPE,
                    opts,
                );
                const messageCtx = { queue: { name: topic } };
                if (
                    config.captureBody === 'all' ||
                    config.captureBody === 'transactions'
                ) {
                    messageCtx.body = msg.content?.toString();
                }

                if (traceHeaders && config.captureHeaders) {
                    // Make sure there is no sensitive data
                    // and transform non-redacted buffers
                    messageCtx.headers = redactKeysFromObject(
                        traceHeaders,
                        config.sanitizeFieldNamesRegExp,
                    );
                    Object.keys(messageCtx.headers).forEach((key) => {
                        const value = messageCtx.headers[key];
                        if (value instanceof Buffer) {
                            messageCtx.headers[key] = value.toString('utf-8');
                        }
                    });
                }

                if (msg.properties.timestamp) {
                    messageCtx.age = {
                        ms: Date.now() - Number(msg.properties.timestamp),
                    };
                }
                trans.setMessageContext(messageCtx);

                agent.logger.debug(`Message consumed from queue: ${topic}`);

                let result, err;

                try {
                    result = cb.apply(this, cbArgs)
                } catch (ex) {
                    // Save the error for use in `finally` below, but re-throw it to
                    // not impact code flow.
                    err = ex;
                    agent.logger.error(`Error consuming message from ${topic}:`, err);
                    throw ex;
                } finally {
                    if (trans) {
                        trans.setOutcome(
                            err ? constants.OUTCOME_FAILURE : constants.OUTCOME_SUCCESS,
                        );
                        trans.end()
                    }
                }

                return result
            }

            let result, err;
            try {
                result = orgConsume.apply(this, arguments);
            } catch (ex) {
                // Save the error for use in `finally` below, but re-throw it to
                // not impact code flow.
                err = ex;
                agent.logger.error(`Error consuming message from ${topic}:`, err);
                throw ex;
            }

            return result;
        }
    })

    customMassWrap(proto, CHANNEL_METHODS, getWrapper);

    shimmer.wrap(proto, 'get', wrapperGet);

    return mod;

    function wrapperGet(orig) {
        return function wrappedGet() {
            const queue = arguments[0]
            const span = ins.createSpan(
                `Channel#get for ${queue}`,
                TYPE,
                SUBTYPE,
                'get',
                { exitSpan: true },
            );
            // console.log(span._parentSpan.type)

            if (!span) {
                return orig.apply(this, arguments);
            }
            agent.logger.debug(`Channel#get in ${NAME}`);

            let result, err;
            try {
                result = orig.apply(this, arguments);
            } catch (ex) {
                // Save the error for use in `finally` below, but re-throw it to
                // not impact code flow.
                err = ex;
                agent.logger.error(`Error in Channel#get:`, err);
                throw ex;
            } finally {
                span.setOutcome(
                    err ? constants.OUTCOME_FAILURE : constants.OUTCOME_SUCCESS,
                );
                span.end();
            }

            return result;
        }
    }

    function getWrapper(methodName) {
        return function wrapper(orig) {
            return function wrappedMethod() {

                let topic = ''

                if (methodName.endsWith("Queue") || methodName.endsWith("Exchange")) {
                    topic = ` of ${arguments[0]}`
                }

                const span = ins.createSpan(
                    `Channel#${methodName}${topic}`,
                    TYPE,
                    SUBTYPE,
                    methodName,
                    { exitSpan: true },
                );
                // console.log(span._parentSpan.type)

                if (!span) {
                    return orig.apply(this, arguments);
                }
                agent.logger.debug(`Channel#${methodName} in ${NAME}`);

                let result, err;
                try {
                    result = orig.apply(this, arguments);
                } catch (ex) {
                    // Save the error for use in `finally` below, but re-throw it to
                    // not impact code flow.
                    err = ex;
                    agent.logger.error(`Error in Channel#${methodName}:`, err);
                    throw ex;
                } finally {
                    span.setOutcome(
                        err ? constants.OUTCOME_FAILURE : constants.OUTCOME_SUCCESS,
                    );
                    span.end();
                }

                return result;
            }
        }
    }

    function customMassWrap(nodule, names, getWrapper) {
        if (!nodule) {
            logger().debug(
                { stack: new Error().stack },
                'must provide one or more modules to patch',
            );
            return;
        }
        if (!(names && Array.isArray(names))) {
            logger().debug('must provide one or more functions to wrap on modules');
            return;
        }

        for (const name of names) {
            const wrapper = getWrapper(name)
            shimmer.wrap(nodule, name, wrapper);
        }
    }
}