apracjs/net/Pool.js

var mysql = require('../');
var Connection = require('./Connection');
var EventEmitter = require('events').EventEmitter;
var PoolConnection = require('./PoolConnection');
var PoolConfig = require('./PoolConfig')

class Pool extends EventEmitter {
  constructor() {
    super();
    this._acquiringConnections = [];
    this._allConnections = [];
    this._freeConnections = [];
    this._connectionQueue = [];
    this._closed = false;
    this.poolConfig = {
      config: new PoolConfig({})
    }
  }

  set poolConfig(options) {
    this.config = options.config;
    this.config.connectionConfig.pool = this;
  }


  getConnection(cb) {

    if (this._closed) {
      var err = new Error('Pool is closed.');
      err.code = 'POOL_CLOSED';
      process.nextTick(function () {
        cb(err);
      });
      return;
    }

    var connection;
    var pool = this;

    if (this._freeConnections.length > 0) {
      connection = this._freeConnections.shift();
      this.acquireConnection(connection, cb);
      return;
    }

    if (this.config.connectionLimit === 0 || this._allConnections.length < this.config.connectionLimit) {
      connection = new PoolConnection(this, { config: this.config.newConnectionConfig() });

      this._acquiringConnections.push(connection);
      this._allConnections.push(connection);

      connection.connect({ timeout: this.config.acquireTimeout }, function onConnect(err) {
        spliceConnection(pool._acquiringConnections, connection);

        if (pool._closed) {
          err = new Error('Pool is closed.');
          err.code = 'POOL_CLOSED';
        }

        if (err) {
          pool._purgeConnection(connection);
          cb(err);
          return;
        }

        pool.emit('connection', connection);
        pool.emit('acquire', connection);
        cb(null, connection);
      });
      return;
    }

    if (!this.config.waitForConnections) {
      process.nextTick(function () {
        var err = new Error('No connections available.');
        err.code = 'POOL_CONNLIMIT';
        cb(err);
      });
      return;
    }

    this._enqueueCallback(cb);
  };

  acquireConnection(connection, cb) {
    if (connection._pool !== this) {
      throw new Error('Connection acquired from wrong pool.');
    }

    var changeUser = this._needsChangeUser(connection);
    var pool = this;

    this._acquiringConnections.push(connection);

    function onOperationComplete(err) {
      spliceConnection(pool._acquiringConnections, connection);

      if (pool._closed) {
        err = new Error('Pool is closed.');
        err.code = 'POOL_CLOSED';
      }

      if (err) {
        pool._connectionQueue.unshift(cb);
        pool._purgeConnection(connection);
        return;
      }

      if (changeUser) {
        pool.emit('connection', connection);
      }

      pool.emit('acquire', connection);
      cb(null, connection);
    }

    if (changeUser) {
      // restore user back to pool configuration
      connection.config = this.config.newConnectionConfig();
      connection.changeUser({ timeout: this.config.acquireTimeout }, onOperationComplete);
    } else {
      // ping connection
      connection.ping({ timeout: this.config.acquireTimeout }, onOperationComplete);
    }
  };

  releaseConnection(connection) {

    if (this._acquiringConnections.indexOf(connection) !== -1) {
      // connection is being acquired
      return;
    }

    if (connection._pool) {
      if (connection._pool !== this) {
        throw new Error('Connection released to wrong pool');
      }

      if (this._freeConnections.indexOf(connection) !== -1) {
        // connection already in free connection pool
        // this won't catch all double-release cases
        throw new Error('Connection already released');
      } else {
        // add connection to end of free queue
        this._freeConnections.push(connection);
        this.emit('release', connection);
      }
    }

    if (this._closed) {
      // empty the connection queue
      this._connectionQueue.splice(0).forEach(function (cb) {
        var err = new Error('Pool is closed.');
        err.code = 'POOL_CLOSED';
        process.nextTick(function () {
          cb(err);
        });
      });
    } else if (this._connectionQueue.length) {
      // get connection with next waiting callback
      this.getConnection(this._connectionQueue.shift());
    }
  };

  async end() {
    this._closed = true;
    let pool = this;
    return await new Promise((resolve, reject) => {
      var calledBack = false;
      var waitingClose = 0;
      var connectionLength = this._allConnections.length;
      function onEnd(err) {
        if (!calledBack && (err || --waitingClose <= 0)) {
          pool._closed = false;//重启连接，方便下次重连
          calledBack = true;
          if(err){
            reject({
              code: -1,
              msg: err
            })
          } else {
            resolve({
              code: 0,
              msg: `成功关闭${connectionLength}条连接`
            });
          }
        }
      }

      while (this._allConnections.length !== 0) {
        waitingClose++;
        this._purgeConnection(this._allConnections[0], onEnd);
      }

      if (waitingClose === 0) {
        process.nextTick(onEnd);
      }
    });
  };

  async query(values) {
    return await new Promise((resolve, reject) => {
      var query = Connection.createQuery(values, (err, data) => {
        if (err) {
          reject(err);
        } else {
          resolve(data);
        }
      });

      this.getConnection(function (err, conn) {
        if (err) {
          query.on('error', function () { });
          query.end(err);
          reject(err);
          return;
        }

        // Release connection based off event
        query.once('end', function () {
          conn.release();
        });

        conn.query(query);
      });
    })
  };

  _enqueueCallback(callback) {

    if (this.config.queueLimit && this._connectionQueue.length >= this.config.queueLimit) {
      process.nextTick(function () {
        var err = new Error('Queue limit reached.');
        err.code = 'POOL_ENQUEUELIMIT';
        callback(err);
      });
      return;
    }

    // Bind to domain, as dequeue will likely occur in a different domain
    var cb = process.domain
      ? process.domain.bind(callback)
      : callback;

    this._connectionQueue.push(cb);
    this.emit('enqueue');
  };

  _needsChangeUser(connection) {
    var connConfig = connection.config;
    var poolConfig = this.config.connectionConfig;

    // check if changeUser values are different
    return connConfig.user !== poolConfig.user
      || connConfig.database !== poolConfig.database
      || connConfig.password !== poolConfig.password
      || connConfig.charsetNumber !== poolConfig.charsetNumber;
  };

  _purgeConnection(connection, callback) {
    var cb = callback || function () { };

    if (connection.state === 'disconnected') {
      connection.destroy();
    }

    this._removeConnection(connection);

    if (connection.state !== 'disconnected' && !connection._protocol._quitSequence) {
      connection._realEnd(cb);
      return;
    }

    process.nextTick(cb);
  };

  _removeConnection(connection) {
    connection._pool = null;

    // Remove connection from all connections
    spliceConnection(this._allConnections, connection);

    // Remove connection from free connections
    spliceConnection(this._freeConnections, connection);

    this.releaseConnection(connection);
  };

  escape(value) {
    return mysql.escape(value, this.config.connectionConfig.stringifyObjects, this.config.connectionConfig.timezone);
  };

  escapeId(value) {
    return mysql.escapeId(value, false);
  };
}

module.exports = new Pool();

function spliceConnection(array, connection) {
  var index;
  if ((index = array.indexOf(connection)) !== -1) {
    // Remove connection from all connections
    array.splice(index, 1);
  }
}