Merge remote-tracking branch 'curran/patch-12'

lib/agent.js

@@ -60,8 +60,15 @@ Agent.prototype.close = function(err) {
 };
 
 Agent.prototype._cleanup = function() {
+
+  // Only clean up once if the stream emits both 'end' and 'close'.
+  if (this.closed) return;
+
   this.closed = true;
 
+  this.backend.agentsCount--;
+  if (!this.stream.isServer) this.backend.remoteAgentsCount--;
+
   // Clean up doc subscription streams
   for (var collection in this.subscribedDocs) {
     var docs = this.subscribedDocs[collection];
@@ -269,12 +276,10 @@ Agent.prototype._open = function() {
       agent._handleMessage(request.data, callback);
     });
   });
-
-  this.stream.on('end', function() {
-    agent.backend.agentsCount--;
-    if (!agent.stream.isServer) agent.backend.remoteAgentsCount--;
-    agent._cleanup();
-  });
+
+  var cleanup = agent._cleanup.bind(agent);
+  this.stream.on('end', cleanup);
+  this.stream.on('close', cleanup);
 };
 
 // Check a request to see if its valid. Returns an error if there's a problem.

test/client/connection.js

@@ -101,6 +101,27 @@ describe('client connection', function() {
       });
     });
 
+    it('updates after connection socket stream emits "close"', function(done) {
+      var backend = this.backend;
+      var connection = backend.connect();
+      connection.on('connected', function() {
+        connection.socket.stream.emit('close')
+        expect(backend.agentsCount).equal(0);
+        done();
+      });
+    });
+
+    it('updates correctly after stream emits both "end" and "close"', function(done) {
+      var backend = this.backend;
+      var connection = backend.connect();
+      connection.on('connected', function() {
+        connection.socket.stream.emit('end')
+        connection.socket.stream.emit('close')
+        expect(backend.agentsCount).equal(0);
+        done();
+      });
+    });
+
     it('does not increment when agent connect is rejected', function() {
       var backend = this.backend;
       backend.use('connect', function(request, next) {