Move common code into middleware utils to make creation of new middleware simpler (#306)

* Move common code into middleware utils to make creation of new middleware simpler

* Add middlewareLog function and add missing typescript types and tests

* Reinstate addError() to avoid closing segments multiple times

* Use middlewareLog in express middleware

Author: paulbrimicombe

packages/core/lib/middleware/mw_utils.d.ts

@@ -24,6 +24,10 @@ export function setDefaultName(name: string): void;
 
 export function disableCentralizedSampling(): void;
 
+export function middlewareLog(message: string, url: string, segment: Segment): void;
+
+export function traceRequestResponseCycle(req: http.IncomingMessage, res: http.ServerResponse): Segment;
+
 export interface BaseRuleConfig {
   http_method: string;
   url_path: string;

packages/core/lib/middleware/mw_utils.js

@@ -5,6 +5,11 @@
  * @module mw_utils
  */
 
+var Segment = require('../segments/segment');
+var IncomingRequestData = require('./incoming_request_data');
+var logger = require('../logger');
+var coreUtils = require('../utils');
+
 var wildcardMatch = require('../utils').wildcardMatch;
 var processTraceData = require('../utils').processTraceData;
 
@@ -134,6 +139,72 @@ var utils = {
       throw new Error('Please specify a path to the local sampling rules file, or supply an object containing the rules.');
 
     this.sampler.setLocalRules(source);
+  },
+
+  /**
+   * Logs a debug message including core request and segment information
+   * @param {string} message - The message to be logged
+   * @param {string} url - The request url being traced
+   * @param {Segment} - The current segment
+   */
+  middlewareLog: function middlewareLog(message, url, segment) {
+    logger.getLogger().debug(message + ': { url: ' + url + ', name: ' + segment.name + ', trace_id: ' +
+      segment.trace_id + ', id: ' + segment.id + ', sampled: ' + !segment.notTraced + ' }');
+  },
+
+  /**
+   * Traces the request/response cycle of an http.IncomingMessage / http.ServerResponse pair.
+   * Resolves sampling rules, creates a segment, adds the core request / response data adding
+   * throttling / error / fault flags based on the response status code.
+   * @param {http.IncomingMessage} req - The incoming request.
+   * @param {http.ServerResponse} res - The server response.
+   * @returns {Segment}
+   * @memberof AWSXRay
+   */
+  traceRequestResponseCycle: function traceRequestResponseCycle(req, res) {
+    var amznTraceHeader = this.processHeaders(req);
+    var name = this.resolveName(req.headers.host);
+    var segment = new Segment(name, amznTraceHeader.root, amznTraceHeader.parent);
+
+    var responseWithEmbeddedRequest = Object.assign({}, res, { req: req });
+    this.resolveSampling(amznTraceHeader, segment, responseWithEmbeddedRequest);
+
+    segment.addIncomingRequestData(new IncomingRequestData(req));
+
+    this.middlewareLog('Starting middleware segment', req.url, segment);
+
+    var middlewareLog = this.middlewareLog;
+    var didEnd = false;
+    var endSegment = function () {
+      // ensure `endSegment` is only called once
+      // in some versions of node.js 10.x and in all versions of node.js 11.x and higher,
+      // the 'finish' and 'close' event are BOTH triggered.
+      // Previously, only one or the other was triggered:
+      // https://github.com/nodejs/node/pull/20611
+      if (didEnd) return;
+      didEnd = true;
+
+      if (res.statusCode === 429) {
+        segment.addThrottleFlag();
+      }
+
+      const cause = coreUtils.getCauseTypeFromHttpStatus(
+        res.statusCode
+      );
+
+      if (cause) {
+        segment[cause] = true;
+      }
+
+      segment.http.close(res);
+      segment.close();
+
+      middlewareLog('Closed middleware segment successfully', req.url, segment);
+    };
+
+    res.on('finish', endSegment);
+    res.on('close', endSegment);
+    return segment;
   }
 };
 

packages/core/test-d/index.test-d.ts

@@ -5,6 +5,7 @@ import { Socket } from 'net';
 import { expectType, expectError } from 'tsd';
 import * as url from 'url';
 import * as AWSXRay from '../lib';
+import { Segment } from '../lib';
 
 expectType<void>(AWSXRay.plugins.EC2Plugin.getData((metadata?: AWSXRay.plugins.EC2Metadata) => { }));
 expectType<void>(AWSXRay.plugins.ECSPlugin.getData((metadata?: AWSXRay.plugins.ECSMetadata) => { }));
@@ -135,6 +136,7 @@ const rulesConfig: AWSXRay.middleware.RulesConfig = {
   }
 };
 expectType<void>(AWSXRay.middleware.setSamplingRules(rulesConfig));
+expectType<void>(AWSXRay.middleware.middlewareLog('log-message', 'https://example.org', segment));
 
 expectType<Namespace>(AWSXRay.getNamespace());
 expectType<AWSXRay.Segment | AWSXRay.Subsegment | undefined>(AWSXRay.resolveSegment(segment));
@@ -163,7 +165,9 @@ expectType<void>(segment.setMatchedSamplingRule('rule'));
 expectType<void>(segment.setServiceData({ version: '2.3.0', package: 'sample-app' }));
 expectType<void>(segment.addPluginData({ elastic_beanstalk: { environment: 'my_environment_name' } }));
 const incomingMessage = new http.IncomingMessage(new Socket());
+const serverResponse = new http.ServerResponse(incomingMessage);
 expectType<void>(segment.addIncomingRequestData(new AWSXRay.middleware.IncomingRequestData(incomingMessage)));
+expectType<Segment>(AWSXRay.middleware.traceRequestResponseCycle(incomingMessage, serverResponse))
 
 function testSegmentLike(segmentLike: AWSXRay.Segment | AWSXRay.Subsegment) {
   expectType<void>(segmentLike.addAnnotation('key', true));

packages/core/test/unit/middleware/mw_utils.test.js

@@ -4,6 +4,7 @@ var sinon = require('sinon');
 var MWUtils = require('../../../lib/middleware/mw_utils');
 var localSampler = require('../../../lib/middleware/sampling/local_sampler');
 var Segment = require('../../../lib/segments/segment');
+var logger = require('../../../lib/logger');
 
 //headers are case-insensitive
 var XRAY_HEADER = 'x-amzn-trace-id';
@@ -15,7 +16,11 @@ describe('Middleware utils', function() {
   var traceId = '1-f9194208-2c7ad569f5d6ff149137be86';
 
   function reloadMWUtils() {
-    var path = '../../../lib/middleware/mw_utils';
+    var path = '../../../lib/logger';
+    delete require.cache[require.resolve(path)];
+    logger = require(path);
+
+    path = '../../../lib/middleware/mw_utils';
     delete require.cache[require.resolve(path)];
     MWUtils = require(path);
   }
@@ -301,4 +306,35 @@ describe('Middleware utils', function() {
       assert.throws(function() { MWUtils.setSamplingRules(new String('')); });
     });
   });
+
+  describe('#middlewareLog', function () {
+    var sandbox, debugLoggerStub;
+
+    beforeEach(function() {
+      reloadMWUtils();
+      sandbox = sinon.sandbox.create();
+      debugLoggerStub = sandbox.stub(logger.getLogger(), 'debug');
+    });
+
+    afterEach(function() {
+      sandbox.restore();
+    });
+
+    it('should log the correct message', function () {
+
+      var message = 'Some message to be logged';
+      var url = 'https://example.org/some-path';
+      var segment = {
+        name: 'segment-name',
+        trace_id: 'some-trace-id',
+        id: 'segment-id',
+        notTraced: false
+      };
+      MWUtils.middlewareLog(message, url, segment);
+
+      var expectedMessage = message + ': { url: ' + url + ', name: ' + segment.name + ', trace_id: ' +
+        segment.trace_id + ', id: ' + segment.id + ', sampled: ' + !segment.notTraced + ' }';
+      debugLoggerStub.should.have.been.calledWithExactly(expectedMessage);
+    });
+  });
 });

packages/express/lib/express_mw.js

@@ -12,8 +12,6 @@
 var AWSXRay = require('aws-xray-sdk-core');
 
 var mwUtils = AWSXRay.middleware;
-var IncomingRequestData = mwUtils.IncomingRequestData;
-var Segment = AWSXRay.Segment;
 
 var expressMW = {
 
@@ -32,40 +30,8 @@ var expressMW = {
 
     mwUtils.setDefaultName(defaultName);
 
-    return function open(req, res, next) {
-      var amznTraceHeader = mwUtils.processHeaders(req);
-      var name = mwUtils.resolveName(req.headers.host);
-      var segment = new Segment(name, amznTraceHeader.root, amznTraceHeader.parent);
-
-      mwUtils.resolveSampling(amznTraceHeader, segment, res);
-      segment.addIncomingRequestData(new IncomingRequestData(req));
-
-      AWSXRay.getLogger().debug('Starting express segment: { url: ' + req.url + ', name: ' + segment.name + ', trace_id: ' +
-        segment.trace_id + ', id: ' + segment.id + ', sampled: ' + !segment.notTraced + ' }');
-
-      var didEnd = false;
-      var endSegment = function () {
-        // ensure `endSegment` is only called once
-        // in some versions of node.js 10.x and in all versions of node.js 11.x and higher,
-        // the 'finish' and 'close' event are BOTH triggered.
-        // Previously, only one or the other was triggered:
-        // https://github.com/nodejs/node/pull/20611
-        if (didEnd) return;
-        didEnd = true;
-        if (this.statusCode === 429)
-          segment.addThrottleFlag();
-        if (AWSXRay.utils.getCauseTypeFromHttpStatus(this.statusCode))
-          segment[AWSXRay.utils.getCauseTypeFromHttpStatus(this.statusCode)] = true;
-
-        segment.http.close(this);
-        segment.close();
-
-        AWSXRay.getLogger().debug('Closed express segment successfully: { url: ' + req.url + ', name: ' + segment.name + ', trace_id: ' +
-          segment.trace_id + ', id: ' + segment.id + ', sampled: ' + !segment.notTraced + ' }');
-      };
-
-      res.on('finish', endSegment);
-      res.on('close', endSegment);
+    return function (req, res, next) {
+      var segment = mwUtils.traceRequestResponseCycle(req, res);
 
       if (AWSXRay.isAutomaticMode()) {
         var ns = AWSXRay.getNamespace();
@@ -97,14 +63,12 @@ var expressMW = {
       if (segment && err) {
         segment.close(err);
 
-        AWSXRay.getLogger().debug('Closed express segment with error: { url: ' + req.url + ', name: ' + segment.name + ', trace_id: ' +
-          segment.trace_id + ', id: ' + segment.id + ', sampled: ' + !segment.notTraced + ' }');
+        mwUtils.middlewareLog('Closed express segment with error', req.url, segment);
 
       } else if (segment) {
         segment.close();
 
-        AWSXRay.getLogger().debug('Closed express segment successfully: { url: ' + req.url + ', name: ' + segment.name + ', trace_id: ' +
-          segment.trace_id + ', id: ' + segment.id + ', sampled: ' + !segment.notTraced + ' }');
+        mwUtils.middlewareLog('Closed express segment successfully', req.url, segment);
       }
 
       if (next)

packages/restify/lib/restify_mw.js

@@ -12,8 +12,6 @@
 var AWSXRay = require('aws-xray-sdk-core');
 
 var mwUtils = AWSXRay.middleware;
-var IncomingRequestData = mwUtils.IncomingRequestData;
-var Segment = AWSXRay.Segment;
 
 var restifyMW = {
 
@@ -40,22 +38,7 @@ var restifyMW = {
     var segment;
 
     server.use(function open(req, res, next) {
-      var amznTraceHeader = mwUtils.processHeaders(req);
-      var name = mwUtils.resolveName(req.headers.host);
-      segment = new Segment(name, amznTraceHeader.root, amznTraceHeader.parent);
-
-      mwUtils.resolveSampling(amznTraceHeader, segment, res);
-      segment.addIncomingRequestData(new IncomingRequestData(req));
-
-      res.on('finish', function() {
-        if (this.statusCode === 429)
-          segment.addThrottleFlag();
-        if (AWSXRay.utils.getCauseTypeFromHttpStatus(this.statusCode))
-          segment[AWSXRay.utils.getCauseTypeFromHttpStatus(this.statusCode)] = true;
-
-        segment.http.close(this);
-        segment.close();
-      });
+      segment = mwUtils.traceRequestResponseCycle(req, res);
 
       if (AWSXRay.isAutomaticMode()) {
         var ns = AWSXRay.getNamespace();

packages/restify/test/unit/restify_mw.test.js

@@ -19,7 +19,7 @@ chai.use(sinonChai);
 
 var utils = require('../test_utils');
 
-describe('Express middleware', function() {
+describe('Restify middleware', function() {
   var sandbox, server;
   var defaultName = 'defaultName';
   var hostName = 'expressMiddlewareTest';
@@ -152,10 +152,15 @@ describe('Express middleware', function() {
       it('should add a finish event to the response', function() {
         server.open(req, res);
 
-        onEventStub.should.have.been.calledOnce;
         onEventStub.should.have.been.calledWithExactly('finish', sinon.match.typeOf('function'));
       });
 
+      it('should add a close event to the response', function() {
+        server.open(req, res);
+
+        onEventStub.should.have.been.calledWithExactly('close', sinon.match.typeOf('function'));
+      });
+
       describe('in automatic mode', function() {
         it('should set the segment on the namespace if in automatic mode', function() {
           server.open(req, res);
@@ -185,9 +190,14 @@ describe('Express middleware', function() {
       it('should add a finish event to the response', function() {
         server.open(req, res);
 
-        onEventStub.should.have.been.calledOnce;
         onEventStub.should.have.been.calledWithExactly('finish', sinon.match.typeOf('function'));
       });
+
+      it('should add a close event to the response', function() {
+        server.open(req, res);
+
+        onEventStub.should.have.been.calledWithExactly('close', sinon.match.typeOf('function'));
+      });
     });
 
     describe('when the request completes', function() {