feat: pool adds gracefulExit to end gracefully

.gitignore

@@ -2,3 +2,4 @@
 coverage/
 node_modules/
 npm-debug.log
+.DS_Store

Readme.md

@@ -392,6 +392,7 @@ constructor. In addition to those options pools accept a few extras:
 * `queueLimit`: The maximum number of connection requests the pool will queue
   before returning an error from `getConnection`. If set to `0`, there is no
   limit to the number of queued connection requests. (Default: `0`)
+* `gracefulExit`: Determines whether to end gracefully. If `true`, every `pool.getConnection` or `pool.query` called before `pool.end` will success. If `false`, only commands / queries already in progress will complete, others will throw an error.
 
 ## Pool events
 
@@ -462,6 +463,11 @@ all the connections have ended.
 **Once `pool.end()` has been called, `pool.getConnection` and other operations
 can no longer be performed**
 
+If `gracefulExit` is set to `true`, the connections end _gracefully_, so all
+ -pending queries will still complete and the time to end the pool will vary.
+
+The default `gracefulExit` is `false`, the following behavior will take effect.
+
 This works by calling `connection.end()` on every active connection in the
 pool, which queues a `QUIT` packet on the connection. And sets a flag to
 prevent `pool.getConnection` from continuing to create any new connections.

lib/Pool.js

@@ -17,9 +17,10 @@ function Pool(options) {
   this._freeConnections      = [];
   this._connectionQueue      = [];
   this._closed               = false;
+  this._pendingClosing       = false;
 }
 
-Pool.prototype.getConnection = function (cb) {
+Pool.prototype.getConnection = function (cb, queued) {
 
   if (this._closed) {
     var err = new Error('Pool is closed.');
@@ -33,12 +34,21 @@ Pool.prototype.getConnection = function (cb) {
   var connection;
   var pool = this;
 
-  if (this._freeConnections.length > 0) {
+  if (this._freeConnections.length > 0 && (!this._pendingClosing || queued)) {
     connection = this._freeConnections.shift();
     this.acquireConnection(connection, cb);
     return;
   }
 
+  if (this._pendingClosing) {
+    var err = new Error('Pool is closed.');
+    err.code = 'POOL_CLOSED';
+    process.nextTick(function () {
+      cb(err);
+    });
+    return;
+  }
+
   if (this.config.connectionLimit === 0 || this._allConnections.length < this.config.connectionLimit) {
     connection = new PoolConnection(this, { config: this.config.newConnectionConfig() });
 
@@ -141,6 +151,10 @@ Pool.prototype.releaseConnection = function releaseConnection(connection) {
       this._freeConnections.push(connection);
       this.emit('release', connection);
     }
+
+    if (this._pendingClosing) {
+      this.end(this._endCallback);
+    }
   }
 
   if (this._closed) {
@@ -154,19 +168,18 @@ Pool.prototype.releaseConnection = function releaseConnection(connection) {
     });
   } else if (this._connectionQueue.length) {
     // get connection with next waiting callback
-    this.getConnection(this._connectionQueue.shift());
+    this.getConnection(this._connectionQueue.shift(), true);
   }
 };
 
 Pool.prototype.end = function (cb) {
-  this._closed = true;
-
   if (typeof cb !== 'function') {
     cb = function (err) {
       if (err) throw err;
     };
   }
 
+  var readyToEnd = false;
   var calledBack   = false;
   var waitingClose = 0;
 
@@ -177,14 +190,26 @@ Pool.prototype.end = function (cb) {
     }
   }
 
-  while (this._allConnections.length !== 0) {
-    waitingClose++;
-    this._purgeConnection(this._allConnections[0], onEnd);
+  if (this._acquiringConnections.length === 0 && this._connectionQueue.length === 0) {
+    readyToEnd = true;
   }
 
-  if (waitingClose === 0) {
-    process.nextTick(onEnd);
+  if (!this.config.gracefulExit || readyToEnd) {
+    this._closed = true;
+
+    while (this._allConnections.length !== 0) {
+      waitingClose++;
+      this._purgeConnection(this._allConnections[0], onEnd);
+    }
+
+    if (waitingClose === 0) {
+      process.nextTick(onEnd);
+    }
+    return;
   }
+
+  this._pendingClosing = true;
+  this._endCallback = cb;
 };
 
 Pool.prototype.query = function (sql, values, cb) {

lib/PoolConfig.js

@@ -20,6 +20,9 @@ function PoolConfig(options) {
   this.queueLimit         = (options.queueLimit === undefined)
     ? 0
     : Number(options.queueLimit);
+  this.gracefulExit       = (options.gracefulExit === undefined)
+    ? false
+    : Boolean(options.gracefulExit);
 }
 
 PoolConfig.prototype.newConnectionConfig = function newConnectionConfig() {

test/unit/pool/test-graceful-exit-ping.js

@@ -0,0 +1,46 @@
+var common = require('../../common');
+var assert = require('assert');
+var pool   = common.createPool({
+  connectionLimit    : 1,
+  port               : common.fakeServerPort,
+  queueLimit         : 5,
+  waitForConnections : true,
+  gracefulExit       : true
+});
+
+var server = common.createFakeServer();
+
+server.listen(common.fakeServerPort, function (err) {
+  assert.ifError(err);
+
+  pool.getConnection(function (err, conn) {
+    assert.ifError(err);
+    conn.release();
+
+    pool.getConnection(function (err, conn) {
+      assert.ifError(err);
+      conn.release();
+    });
+
+    pool.end(function (err) {
+      assert.ifError(err);
+      server.destroy();
+    });
+
+    pool.getConnection(function (err) {
+      assert.ok(err);
+      assert.equal(err.message, 'Pool is closed.');
+      assert.equal(err.code, 'POOL_CLOSED');
+    });
+  });
+});
+
+server.on('connection', function (conn) {
+  conn.handshake();
+  conn.on('ping', function () {
+    setTimeout(function () {
+      conn._sendPacket(new common.Packets.OkPacket());
+      conn._parser.resetPacketNumber();
+    }, 100);
+  });
+});

test/unit/pool/test-graceful-exit-queued.js

@@ -0,0 +1,37 @@
+var common = require('../../common');
+var assert = require('assert');
+var pool   = common.createPool({
+  connectionLimit    : 1,
+  port               : common.fakeServerPort,
+  queueLimit         : 5,
+  waitForConnections : true,
+  gracefulExit       : true
+});
+
+var server = common.createFakeServer();
+
+server.listen(common.fakeServerPort, function (err) {
+  assert.ifError(err);
+
+  pool.getConnection(function (err, conn) {
+    assert.ifError(err);
+
+    pool.end(function (err) {
+      assert.ifError(err);
+      server.destroy();
+    });
+
+    pool.getConnection(function (err) {
+      assert.ok(err);
+      assert.equal(err.message, 'Pool is closed.');
+      assert.equal(err.code, 'POOL_CLOSED');
+    });
+
+    conn.release();
+  });
+
+  pool.getConnection(function (err, conn) {
+    assert.ifError(err);
+    conn.release();
+  });
+});