191 lines
5.1 KiB
JavaScript
191 lines
5.1 KiB
JavaScript
|
'use strict'
|
||
|
|
||
|
const {
|
||
|
BalancedPoolMissingUpstreamError,
|
||
|
InvalidArgumentError
|
||
|
} = require('../core/errors')
|
||
|
const {
|
||
|
PoolBase,
|
||
|
kClients,
|
||
|
kNeedDrain,
|
||
|
kAddClient,
|
||
|
kRemoveClient,
|
||
|
kGetDispatcher
|
||
|
} = require('./pool-base')
|
||
|
const Pool = require('./pool')
|
||
|
const { kUrl, kInterceptors } = require('../core/symbols')
|
||
|
const { parseOrigin } = require('../core/util')
|
||
|
const kFactory = Symbol('factory')
|
||
|
|
||
|
const kOptions = Symbol('options')
|
||
|
const kGreatestCommonDivisor = Symbol('kGreatestCommonDivisor')
|
||
|
const kCurrentWeight = Symbol('kCurrentWeight')
|
||
|
const kIndex = Symbol('kIndex')
|
||
|
const kWeight = Symbol('kWeight')
|
||
|
const kMaxWeightPerServer = Symbol('kMaxWeightPerServer')
|
||
|
const kErrorPenalty = Symbol('kErrorPenalty')
|
||
|
|
||
|
function getGreatestCommonDivisor (a, b) {
|
||
|
if (b === 0) return a
|
||
|
return getGreatestCommonDivisor(b, a % b)
|
||
|
}
|
||
|
|
||
|
function defaultFactory (origin, opts) {
|
||
|
return new Pool(origin, opts)
|
||
|
}
|
||
|
|
||
|
class BalancedPool extends PoolBase {
|
||
|
constructor (upstreams = [], { factory = defaultFactory, ...opts } = {}) {
|
||
|
super()
|
||
|
|
||
|
this[kOptions] = opts
|
||
|
this[kIndex] = -1
|
||
|
this[kCurrentWeight] = 0
|
||
|
|
||
|
this[kMaxWeightPerServer] = this[kOptions].maxWeightPerServer || 100
|
||
|
this[kErrorPenalty] = this[kOptions].errorPenalty || 15
|
||
|
|
||
|
if (!Array.isArray(upstreams)) {
|
||
|
upstreams = [upstreams]
|
||
|
}
|
||
|
|
||
|
if (typeof factory !== 'function') {
|
||
|
throw new InvalidArgumentError('factory must be a function.')
|
||
|
}
|
||
|
|
||
|
this[kInterceptors] = opts.interceptors?.BalancedPool && Array.isArray(opts.interceptors.BalancedPool)
|
||
|
? opts.interceptors.BalancedPool
|
||
|
: []
|
||
|
this[kFactory] = factory
|
||
|
|
||
|
for (const upstream of upstreams) {
|
||
|
this.addUpstream(upstream)
|
||
|
}
|
||
|
this._updateBalancedPoolStats()
|
||
|
}
|
||
|
|
||
|
addUpstream (upstream) {
|
||
|
const upstreamOrigin = parseOrigin(upstream).origin
|
||
|
|
||
|
if (this[kClients].find((pool) => (
|
||
|
pool[kUrl].origin === upstreamOrigin &&
|
||
|
pool.closed !== true &&
|
||
|
pool.destroyed !== true
|
||
|
))) {
|
||
|
return this
|
||
|
}
|
||
|
const pool = this[kFactory](upstreamOrigin, Object.assign({}, this[kOptions]))
|
||
|
|
||
|
this[kAddClient](pool)
|
||
|
pool.on('connect', () => {
|
||
|
pool[kWeight] = Math.min(this[kMaxWeightPerServer], pool[kWeight] + this[kErrorPenalty])
|
||
|
})
|
||
|
|
||
|
pool.on('connectionError', () => {
|
||
|
pool[kWeight] = Math.max(1, pool[kWeight] - this[kErrorPenalty])
|
||
|
this._updateBalancedPoolStats()
|
||
|
})
|
||
|
|
||
|
pool.on('disconnect', (...args) => {
|
||
|
const err = args[2]
|
||
|
if (err && err.code === 'UND_ERR_SOCKET') {
|
||
|
// decrease the weight of the pool.
|
||
|
pool[kWeight] = Math.max(1, pool[kWeight] - this[kErrorPenalty])
|
||
|
this._updateBalancedPoolStats()
|
||
|
}
|
||
|
})
|
||
|
|
||
|
for (const client of this[kClients]) {
|
||
|
client[kWeight] = this[kMaxWeightPerServer]
|
||
|
}
|
||
|
|
||
|
this._updateBalancedPoolStats()
|
||
|
|
||
|
return this
|
||
|
}
|
||
|
|
||
|
_updateBalancedPoolStats () {
|
||
|
this[kGreatestCommonDivisor] = this[kClients].map(p => p[kWeight]).reduce(getGreatestCommonDivisor, 0)
|
||
|
}
|
||
|
|
||
|
removeUpstream (upstream) {
|
||
|
const upstreamOrigin = parseOrigin(upstream).origin
|
||
|
|
||
|
const pool = this[kClients].find((pool) => (
|
||
|
pool[kUrl].origin === upstreamOrigin &&
|
||
|
pool.closed !== true &&
|
||
|
pool.destroyed !== true
|
||
|
))
|
||
|
|
||
|
if (pool) {
|
||
|
this[kRemoveClient](pool)
|
||
|
}
|
||
|
|
||
|
return this
|
||
|
}
|
||
|
|
||
|
get upstreams () {
|
||
|
return this[kClients]
|
||
|
.filter(dispatcher => dispatcher.closed !== true && dispatcher.destroyed !== true)
|
||
|
.map((p) => p[kUrl].origin)
|
||
|
}
|
||
|
|
||
|
[kGetDispatcher] () {
|
||
|
// We validate that pools is greater than 0,
|
||
|
// otherwise we would have to wait until an upstream
|
||
|
// is added, which might never happen.
|
||
|
if (this[kClients].length === 0) {
|
||
|
throw new BalancedPoolMissingUpstreamError()
|
||
|
}
|
||
|
|
||
|
const dispatcher = this[kClients].find(dispatcher => (
|
||
|
!dispatcher[kNeedDrain] &&
|
||
|
dispatcher.closed !== true &&
|
||
|
dispatcher.destroyed !== true
|
||
|
))
|
||
|
|
||
|
if (!dispatcher) {
|
||
|
return
|
||
|
}
|
||
|
|
||
|
const allClientsBusy = this[kClients].map(pool => pool[kNeedDrain]).reduce((a, b) => a && b, true)
|
||
|
|
||
|
if (allClientsBusy) {
|
||
|
return
|
||
|
}
|
||
|
|
||
|
let counter = 0
|
||
|
|
||
|
let maxWeightIndex = this[kClients].findIndex(pool => !pool[kNeedDrain])
|
||
|
|
||
|
while (counter++ < this[kClients].length) {
|
||
|
this[kIndex] = (this[kIndex] + 1) % this[kClients].length
|
||
|
const pool = this[kClients][this[kIndex]]
|
||
|
|
||
|
// find pool index with the largest weight
|
||
|
if (pool[kWeight] > this[kClients][maxWeightIndex][kWeight] && !pool[kNeedDrain]) {
|
||
|
maxWeightIndex = this[kIndex]
|
||
|
}
|
||
|
|
||
|
// decrease the current weight every `this[kClients].length`.
|
||
|
if (this[kIndex] === 0) {
|
||
|
// Set the current weight to the next lower weight.
|
||
|
this[kCurrentWeight] = this[kCurrentWeight] - this[kGreatestCommonDivisor]
|
||
|
|
||
|
if (this[kCurrentWeight] <= 0) {
|
||
|
this[kCurrentWeight] = this[kMaxWeightPerServer]
|
||
|
}
|
||
|
}
|
||
|
if (pool[kWeight] >= this[kCurrentWeight] && (!pool[kNeedDrain])) {
|
||
|
return pool
|
||
|
}
|
||
|
}
|
||
|
|
||
|
this[kCurrentWeight] = this[kClients][maxWeightIndex][kWeight]
|
||
|
this[kIndex] = maxWeightIndex
|
||
|
return this[kClients][maxWeightIndex]
|
||
|
}
|
||
|
}
|
||
|
|
||
|
module.exports = BalancedPool
|