Fix connection / disconnection issues

brianc · brianc · commit bd87cddc72bb · 2017-06-18T16:48:14.000-05:00
diff --git a/lib/client.js b/lib/client.js
@@ -109,12 +109,35 @@ Client.prototype.connect = function(callback) {
     self.secretKey = msg.secretKey;
   });
 
+  const connectingErrorHandler = (err) => {
+    if (this._connectionError) {
+      return;
+    }
+    this._connectionError = true
+    if (callback) {
+      return callback(err)
+    }
+    this.emit('error', err)
+  }
+
+  const connectedErrorHandler = (err) => {
+    if(this.activeQuery) {
+      var activeQuery = self.activeQuery;
+      this.activeQuery = null;
+      return activeQuery.handleError(err, con);
+    }
+    this.emit('error', err)
+  }
+
+  con.on('error', connectingErrorHandler)
+
   //hook up query handling events to connection
   //after the connection initially becomes ready for queries
   con.once('readyForQuery', function() {
     self._connecting = false;
     self._attachListeners(con);
-
+    con.removeListener('error', connectingErrorHandler);
+    con.on('error', connectedErrorHandler)
 
     //process possible callback argument to Client#connect
     if (callback) {
@@ -136,37 +159,7 @@ Client.prototype.connect = function(callback) {
     self._pulseQueryQueue();
   });
 
-  con.on('error', function(error) {
-    if(this.activeQuery) {
-      var activeQuery = self.activeQuery;
-      this.activeQuery = null;
-      return activeQuery.handleError(error, con);
-    }
-
-    if (this._connecting) {
-      // set a flag indicating we've seen an error during connection
-      // the backend will terminate the connection and we don't want
-      // to throw a second error when the connection is terminated
-      this._connectionError = true;
-    }
-
-    if(!callback) {
-      return this.emit('error', error);
-    }
-
-    con.end(); // make sure ECONNRESET errors don't cause error events
-    callback(error);
-    callback = null;
-  }.bind(this));
-
-  con.once('end', function() {
-    if (callback) {
-      // haven't received a connection message yet!
-      var err = new Error('Connection terminated');
-      callback(err);
-      callback = null;
-      return;
-    }
+  con.once('end', () => {
     if(this.activeQuery) {
       var disconnectError = new Error('Connection terminated');
       this.activeQuery.handleError(disconnectError, con);
@@ -177,12 +170,19 @@ Client.prototype.connect = function(callback) {
       // on this client then we have an unexpected disconnection
       // treat this as an error unless we've already emitted an error
       // during connection.
-      if (!this._connectionError) {
-        this.emit('error', new Error('Connection terminated unexpectedly'));
+      const error = new Error('Connection terminated unexpectedly')
+      if (this._connecting && !this._connectionError) {
+        if (callback) {
+          callback(error)
+        } else {
+          this.emit('error', error)
+        }
+      } else if (!this._connectionError) {
+        this.emit('error', error);
       }
     }
     this.emit('end');
-  }.bind(this));
+  });
 
 
   con.on('notice', function(msg) {
diff --git a/test/integration/client/error-handling-tests.js b/test/integration/client/error-handling-tests.js
@@ -1,10 +1,10 @@
 "use strict";
-"use strict";
 
 var helper = require('./test-helper');
 var util = require('util');
 
 var pg = helper.pg
+const Client = pg.Client
 
 var createErorrClient = function() {
   var client = helper.client();
@@ -84,6 +84,7 @@ suite.test('non-query error with callback', function(done) {
     user:'asldkfjsadlfkj'
   });
   client.connect(assert.calls(function(error, client) {
+    console.log('error!', error.stack)
     assert(error instanceof Error)
     done()
   }));
@@ -142,3 +143,19 @@ suite.test('within a simple query', (done) => {
     done();
   });
 });
+
+suite.test('connected, idle client error', (done) => {
+  const client = new Client()
+  client.connect((err) => {
+    if (err) {
+      throw new Error('Should not receive error callback after connection')
+    }
+    setImmediate(() => {
+      (client.connection || client.native).emit('error', new Error('expected'))
+    })
+  })
+  client.on('error', (err) => {
+    assert.equal(err.message, 'expected')
+    client.end(done)
+  })
+})
diff --git a/test/integration/connection-pool/error-tests.js b/test/integration/connection-pool/error-tests.js
@@ -8,6 +8,11 @@ pg.defaults.poolSize = 2;
 const pool = new pg.Pool()
 
 const suite = new helper.Suite()
+suite.test('connecting to invalid port', (cb) => {
+  const pool = new pg.Pool({ port: 13801 })
+  pool.connect().catch(e => cb())
+})
+
 suite.test('errors emitted on pool', (cb) => {
   //get first client
   pool.connect(assert.success(function (client, done) {
diff --git a/test/integration/test-helper.js b/test/integration/test-helper.js
@@ -1,8 +1,8 @@
 "use strict";
-var helper = require(__dirname + '/../test-helper');
+var helper = require('./../test-helper');
 
 if(helper.args.native) {
-  Client = require(__dirname + '/../../lib/native');
+  Client = require('./../../lib/native');
   helper.Client = Client;
   helper.pg = helper.pg.native;
 }
diff --git a/test/unit/client/early-disconnect-tests.js b/test/unit/client/early-disconnect-tests.js
@@ -1,7 +1,7 @@
 "use strict";
-var helper = require(__dirname + '/test-helper');
+var helper = require('./test-helper');
 var net = require('net');
-var pg = require('../../..//lib/index.js');
+var pg = require('../../../lib/index.js');
 
 /* console.log() messages show up in `make test` output. TODO: fix it. */
 var server = net.createServer(function(c) {
