deep-framework/vendor/aws-xray-sdk/lib/patchers/lambda_p.js

DEEP Framework

var CLSUtils = require('../utils').CLSUtils;
var Segment = require('../segments/segment');
var SegmentEmitter = require('../segment_emitter');

var logger = require('../logger');
var ns;

/**
* Used to initialize segments on AWS Lambda with extra data from the context.
* @param {function} fcn - The reference to the function to execute.
* @param {Bool} enableNestedMode - A boolean flag to enable nested mode.
*/

module.exports.captureLambda = function captureLambda(fcn) {
  if (!CLSUtils.isCLSMode()) {
    logger.warn('AWS Lambda does not support AWS X-Ray manual mode.  Switching to CLS mode.');
    CLSUtils.enableCLSMode();
  }

  ns = CLSUtils.getNamespace();

  return function captureLambdaFunction(event, context, callback) {
    SegmentEmitter.setDaemonAddress(process.env._AWS_XRAY_DAEMON_ADDRESS + ':' + process.env._AWS_XRAY_DAEMON_PORT);
    SegmentEmitter.disableReusableSocket();

    var xrayContext = context.xrayContext;
    var segment = new Segment(context.functionName, xrayContext.traceId, xrayContext.parentId);
    context.xrayContext.segment = segment;

    if (!segment.aws)
      segment.aws = {};

    segment.aws.lambda = context.invokedFunctionArn;
    segment.origin = 'AWS::Lambda::Function';

    if (!xrayContext.sampled)
      segment.notTraced = !xrayContext.sampled;

    ns.run(function () {
      CLSUtils.setSegment(segment);

      try {
        if (typeof callback === 'function') {
          fcn(event, context, captureCallbackFcn('lambdaCallback', callback, segment));
        } else {
          fcn(event, context);
        }

        segment.close();
      } catch (e) {
        segment.close(e);
        throw(e);
      }
    });
  };
};

function captureCallbackFcn (name, callback, segment) {
  return function captureCallback() {
    var sub = segment.addNewSubsegment(name);
    var args = [].slice.call(arguments);

    if (args && args[0]) {
      sub.addError(args[0]);
      segment.fault = true;
    }

    sub.close();
    callback.apply(null, args);
  };
}