Prevent Express faulted segments from being closed twice

packages/express/lib/express_mw.js

@@ -1,3 +1,8 @@
+const AWSXRay = require('aws-xray-sdk-core');
+const { logger } = require('aws-xray-sdk-core/lib/middleware/sampling/service_connector');
+
+const mwUtils = AWSXRay.middleware;
+
 /**
  * Express middleware module.
  *
@@ -8,11 +13,6 @@
  * Otherwise, for manual mode, this appends the Segment object to the request object as req.segment.
  * @module express_mw
  */
-
-var AWSXRay = require('aws-xray-sdk-core');
-
-var mwUtils = AWSXRay.middleware;
-
 var expressMW = {
 
   /**
@@ -23,7 +23,6 @@ var expressMW = {
    * @alias module:express_mw.openSegment
    * @returns {function}
    */
-
   openSegment: function openSegment(defaultName) {
     if (!defaultName || typeof defaultName !== 'string')
       throw new Error('Default segment name was not supplied.  Please provide a string.');
@@ -52,23 +51,17 @@ var expressMW = {
 
   /**
    * After your routes, before any extra error handling middleware, use 'app.use(AWSXRayExpress.closeSegment())'.
+   * This is error-handling middleware, so it is called only when there is a server-side fault.
    * @alias module:express_mw.closeSegment
    * @returns {function}
    */
-
   closeSegment: function closeSegment() {
     return function close(err, req, res, next) {
       var segment = AWSXRay.resolveSegment(req.segment);
 
       if (segment && err) {
-        segment.close(err);
-
-        mwUtils.middlewareLog('Closed express segment with error', req.url, segment);
-
-      } else if (segment) {
-        segment.close();
-
-        mwUtils.middlewareLog('Closed express segment successfully', req.url, segment);
+        segment.addError(err);
+        logger.getLogger().debug('Added Express server fault to segment');
       }
 
       if (next)

packages/express/test/unit/express_mw.test.js

@@ -182,4 +182,32 @@ describe('Express middleware', function() {
       });
     });
   });
+
+  describe('#close', function() {
+    var segment, close;
+    beforeEach(function() {
+      xray.enableManualMode();
+      segment = new Segment('test');
+      sandbox = sinon.createSandbox();
+      close = expressMW.closeSegment();
+
+      req = {
+        segment: segment
+      };
+      res = {};
+    });
+
+    afterEach(function() {
+      xray.enableAutomaticMode();
+      sandbox.restore();
+    });
+
+    it('should add error using express middleware', function() {
+      const segment = req.segment;
+      close(new Error(), req, res);
+
+      assert.property(segment, 'cause');
+      assert.property(segment.cause, 'exceptions');
+    });
+  })
 });

packages/test_express/test/express.js

@@ -16,13 +16,14 @@ var ShakyStream = require('./shaky_stream').ShakyStream;
 
 var chai = require('chai');
 var assert = chai.assert;
+var sinon = require('sinon');
 chai.should();
 
 describe('Express', () => {
   /**
    * @type {dgram.Socket}
    */
-  var daemon;
+  var daemon, sandbox;
 
   before(() => {
     // force all requests to be sampled
@@ -39,6 +40,7 @@ describe('Express', () => {
 
   beforeEach((done) => {
     daemon = createDaemon(done);
+    sandbox = sinon.createSandbox();
   });
 
   afterEach((done) => {
@@ -47,6 +49,7 @@ describe('Express', () => {
     xray.middleware.hostPattern = null;
 
     daemon.close(done);
+    sandbox.restore();
   });
 
   it('should generate fix-named segments', (done) => {
@@ -137,6 +140,54 @@ describe('Express', () => {
     });
   });
 
+  it('should close segments with server faults properly', (done) => {
+    sandbox.spy(xray.Segment.prototype, 'close');
+    const daemonMessagesResolved = new Promise((resolve) => {
+      daemon.on('message', messageCounter(1, (resolve)));
+    });
+
+    const route = '/';
+    const name = 'test';
+
+    // Create express app that has a fault
+    const expressApp = createAppWithRoute({
+      name: name,
+      route: route,
+      handler: function(req, res) {
+        throw new Error('My test error');
+      }
+    });
+
+    const server = expressApp.listen(0, () => {
+      const url = 'http://127.0.0.1:' + server.address().port + route;
+
+      // wait for the express response, and the daemon to receive messages
+      Promise.all([triggerEndpoint(url), daemonMessagesResolved])
+      .then((data) => {
+        const result = data[0];
+        const messages = data[1];
+
+        assert.equal(result.status, 500);
+        assert.equal(messages.length, 1);
+
+        // verify the Segment has status code & exception populated, and that close is called once
+        const segment = parseMessage(messages[0]);
+        validateExpressSegment(segment, {
+          name: name,
+          responseStatus: result.status,
+          url: url
+        });
+        assert.property(segment, 'cause');
+        assert.property(segment.cause, 'exceptions');
+        assert(xray.Segment.prototype.close.calledOnce);
+        done();
+      }).catch(() => {
+        assert.fail();
+        done();
+      });
+    });
+  });
+
   describe('automatic mode', () => {
     before(() => {
       xray.enableAutomaticMode();