feat: pool adds gracefulExit to end gracefully

Readme.md

@@ -451,20 +451,37 @@ trying to gracefully shutdown a server. To end all the connections in the
 pool, use the `end` method on the pool:
 
 ```js
-pool.end(function (err) {
+pool.end(false, function (err) {
   // all connections in the pool have ended
 });
 ```
 
-The `end` method takes an _optional_ callback that you can use to know once
-all the connections have ended.
+The `end` method takes two _optional_ arguemnts:
+
+* `gracefulExit`: Determines whether to end gracefully. If `true`, every
+`pool.getConnection` or `pool.query` called before `pool.end` will complete.
+If `false`, only commands / queries already in progress will complete,
+others will fail. (Default: `false`)
+
+* `callback`: Will be called once all the connections have ended.
 
 **Once `pool.end()` has been called, `pool.getConnection` and other operations
 can no longer be performed**
 
-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.
+### Under the hood
+
+If `gracefulExit` is set to `true`, after calling `pool.end` the poll will
+enter into the `pendingClose` state, all former or queued queries will still
+complete. But the pool will no longer accept new queries.
+
+This works by not queueing the `QUIT` packet on all the connections until there
+is no connection in the aquiring state and no queued queries. All established
+connections will still queue queries which were added before calling `pool.end`.
+
+If `gracefulExit` is set to `false`, `pool.end` 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.
 
 Since this queues a `QUIT` packet on each connection, all commands / queries
 already in progress will complete, just like calling `connection.end()`. If

lib/Pool.js

@@ -17,9 +17,10 @@ function Pool(options) {
   this._freeConnections      = [];
   this._connectionQueue      = [];
   this._closed               = false;
+  this._pendingClose         = 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._pendingClose || _queued)) {
     connection = this._freeConnections.shift();
     this.acquireConnection(connection, cb);
     return;
   }
 
+  if (this._pendingClose) {
+    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._pendingClose) {
+      this.end(true, this._endCallback);
+    }
   }
 
   if (this._closed) {
@@ -154,19 +168,22 @@ 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;
-
+Pool.prototype.end = function (gracefulExit, cb) {
+  if (typeof gracefulExit === 'function') {
+    cb = gracefulExit;
+    gracefulExit = false;
+  }
   if (typeof cb !== 'function') {
     cb = function (err) {
       if (err) throw err;
     };
   }
 
+  var readyToEnd = false;
   var calledBack   = false;
   var waitingClose = 0;
 
@@ -177,14 +194,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 (!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._pendingClose = true;
+  this._endCallback = cb;
 };
 
 Pool.prototype.query = function (sql, values, cb) {

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

@@ -0,0 +1,45 @@
+var common = require('../../common');
+var assert = require('assert');
+var pool   = common.createPool({
+  connectionLimit    : 1,
+  port               : common.fakeServerPort,
+  queueLimit         : 5,
+  waitForConnections : 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(true, 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,36 @@
+var common = require('../../common');
+var assert = require('assert');
+var pool   = common.createPool({
+  connectionLimit    : 1,
+  port               : common.fakeServerPort,
+  queueLimit         : 5,
+  waitForConnections : true
+});
+
+var server = common.createFakeServer();
+
+server.listen(common.fakeServerPort, function (err) {
+  assert.ifError(err);
+
+  pool.getConnection(function (err, conn) {
+    assert.ifError(err);
+
+    pool.end(true, 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();
+  });
+});