index.js 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479
  1. 'use strict'
  2. const EventEmitter = require('events').EventEmitter
  3. const NOOP = function () {}
  4. const removeWhere = (list, predicate) => {
  5. const i = list.findIndex(predicate)
  6. return i === -1 ? undefined : list.splice(i, 1)[0]
  7. }
  8. class IdleItem {
  9. constructor(client, idleListener, timeoutId) {
  10. this.client = client
  11. this.idleListener = idleListener
  12. this.timeoutId = timeoutId
  13. }
  14. }
  15. class PendingItem {
  16. constructor(callback) {
  17. this.callback = callback
  18. }
  19. }
  20. function throwOnDoubleRelease() {
  21. throw new Error('Release called on client which has already been released to the pool.')
  22. }
  23. function promisify(Promise, callback) {
  24. if (callback) {
  25. return { callback: callback, result: undefined }
  26. }
  27. let rej
  28. let res
  29. const cb = function (err, client) {
  30. err ? rej(err) : res(client)
  31. }
  32. const result = new Promise(function (resolve, reject) {
  33. res = resolve
  34. rej = reject
  35. }).catch((err) => {
  36. // replace the stack trace that leads to `TCP.onStreamRead` with one that leads back to the
  37. // application that created the query
  38. Error.captureStackTrace(err)
  39. throw err
  40. })
  41. return { callback: cb, result: result }
  42. }
  43. function makeIdleListener(pool, client) {
  44. return function idleListener(err) {
  45. err.client = client
  46. client.removeListener('error', idleListener)
  47. client.on('error', () => {
  48. pool.log('additional client error after disconnection due to error', err)
  49. })
  50. pool._remove(client)
  51. // TODO - document that once the pool emits an error
  52. // the client has already been closed & purged and is unusable
  53. pool.emit('error', err, client)
  54. }
  55. }
  56. class Pool extends EventEmitter {
  57. constructor(options, Client) {
  58. super()
  59. this.options = Object.assign({}, options)
  60. if (options != null && 'password' in options) {
  61. // "hiding" the password so it doesn't show up in stack traces
  62. // or if the client is console.logged
  63. Object.defineProperty(this.options, 'password', {
  64. configurable: true,
  65. enumerable: false,
  66. writable: true,
  67. value: options.password,
  68. })
  69. }
  70. if (options != null && options.ssl && options.ssl.key) {
  71. // "hiding" the ssl->key so it doesn't show up in stack traces
  72. // or if the client is console.logged
  73. Object.defineProperty(this.options.ssl, 'key', {
  74. enumerable: false,
  75. })
  76. }
  77. this.options.max = this.options.max || this.options.poolSize || 10
  78. this.options.min = this.options.min || 0
  79. this.options.maxUses = this.options.maxUses || Infinity
  80. this.options.allowExitOnIdle = this.options.allowExitOnIdle || false
  81. this.options.maxLifetimeSeconds = this.options.maxLifetimeSeconds || 0
  82. this.log = this.options.log || function () {}
  83. this.Client = this.options.Client || Client || require('pg').Client
  84. this.Promise = this.options.Promise || global.Promise
  85. if (typeof this.options.idleTimeoutMillis === 'undefined') {
  86. this.options.idleTimeoutMillis = 10000
  87. }
  88. this._clients = []
  89. this._idle = []
  90. this._expired = new WeakSet()
  91. this._pendingQueue = []
  92. this._endCallback = undefined
  93. this.ending = false
  94. this.ended = false
  95. }
  96. _isFull() {
  97. return this._clients.length >= this.options.max
  98. }
  99. _isAboveMin() {
  100. return this._clients.length > this.options.min
  101. }
  102. _pulseQueue() {
  103. this.log('pulse queue')
  104. if (this.ended) {
  105. this.log('pulse queue ended')
  106. return
  107. }
  108. if (this.ending) {
  109. this.log('pulse queue on ending')
  110. if (this._idle.length) {
  111. this._idle.slice().map((item) => {
  112. this._remove(item.client)
  113. })
  114. }
  115. if (!this._clients.length) {
  116. this.ended = true
  117. this._endCallback()
  118. }
  119. return
  120. }
  121. // if we don't have any waiting, do nothing
  122. if (!this._pendingQueue.length) {
  123. this.log('no queued requests')
  124. return
  125. }
  126. // if we don't have any idle clients and we have no more room do nothing
  127. if (!this._idle.length && this._isFull()) {
  128. return
  129. }
  130. const pendingItem = this._pendingQueue.shift()
  131. if (this._idle.length) {
  132. const idleItem = this._idle.pop()
  133. clearTimeout(idleItem.timeoutId)
  134. const client = idleItem.client
  135. client.ref && client.ref()
  136. const idleListener = idleItem.idleListener
  137. return this._acquireClient(client, pendingItem, idleListener, false)
  138. }
  139. if (!this._isFull()) {
  140. return this.newClient(pendingItem)
  141. }
  142. throw new Error('unexpected condition')
  143. }
  144. _remove(client, callback) {
  145. const removed = removeWhere(this._idle, (item) => item.client === client)
  146. if (removed !== undefined) {
  147. clearTimeout(removed.timeoutId)
  148. }
  149. this._clients = this._clients.filter((c) => c !== client)
  150. const context = this
  151. client.end(() => {
  152. context.emit('remove', client)
  153. if (typeof callback === 'function') {
  154. callback()
  155. }
  156. })
  157. }
  158. connect(cb) {
  159. if (this.ending) {
  160. const err = new Error('Cannot use a pool after calling end on the pool')
  161. return cb ? cb(err) : this.Promise.reject(err)
  162. }
  163. const response = promisify(this.Promise, cb)
  164. const result = response.result
  165. // if we don't have to connect a new client, don't do so
  166. if (this._isFull() || this._idle.length) {
  167. // if we have idle clients schedule a pulse immediately
  168. if (this._idle.length) {
  169. process.nextTick(() => this._pulseQueue())
  170. }
  171. if (!this.options.connectionTimeoutMillis) {
  172. this._pendingQueue.push(new PendingItem(response.callback))
  173. return result
  174. }
  175. const queueCallback = (err, res, done) => {
  176. clearTimeout(tid)
  177. response.callback(err, res, done)
  178. }
  179. const pendingItem = new PendingItem(queueCallback)
  180. // set connection timeout on checking out an existing client
  181. const tid = setTimeout(() => {
  182. // remove the callback from pending waiters because
  183. // we're going to call it with a timeout error
  184. removeWhere(this._pendingQueue, (i) => i.callback === queueCallback)
  185. pendingItem.timedOut = true
  186. response.callback(new Error('timeout exceeded when trying to connect'))
  187. }, this.options.connectionTimeoutMillis)
  188. if (tid.unref) {
  189. tid.unref()
  190. }
  191. this._pendingQueue.push(pendingItem)
  192. return result
  193. }
  194. this.newClient(new PendingItem(response.callback))
  195. return result
  196. }
  197. newClient(pendingItem) {
  198. const client = new this.Client(this.options)
  199. this._clients.push(client)
  200. const idleListener = makeIdleListener(this, client)
  201. this.log('checking client timeout')
  202. // connection timeout logic
  203. let tid
  204. let timeoutHit = false
  205. if (this.options.connectionTimeoutMillis) {
  206. tid = setTimeout(() => {
  207. this.log('ending client due to timeout')
  208. timeoutHit = true
  209. // force kill the node driver, and let libpq do its teardown
  210. client.connection ? client.connection.stream.destroy() : client.end()
  211. }, this.options.connectionTimeoutMillis)
  212. }
  213. this.log('connecting new client')
  214. client.connect((err) => {
  215. if (tid) {
  216. clearTimeout(tid)
  217. }
  218. client.on('error', idleListener)
  219. if (err) {
  220. this.log('client failed to connect', err)
  221. // remove the dead client from our list of clients
  222. this._clients = this._clients.filter((c) => c !== client)
  223. if (timeoutHit) {
  224. err = new Error('Connection terminated due to connection timeout', { cause: err })
  225. }
  226. // this client won’t be released, so move on immediately
  227. this._pulseQueue()
  228. if (!pendingItem.timedOut) {
  229. pendingItem.callback(err, undefined, NOOP)
  230. }
  231. } else {
  232. this.log('new client connected')
  233. if (this.options.maxLifetimeSeconds !== 0) {
  234. const maxLifetimeTimeout = setTimeout(() => {
  235. this.log('ending client due to expired lifetime')
  236. this._expired.add(client)
  237. const idleIndex = this._idle.findIndex((idleItem) => idleItem.client === client)
  238. if (idleIndex !== -1) {
  239. this._acquireClient(
  240. client,
  241. new PendingItem((err, client, clientRelease) => clientRelease()),
  242. idleListener,
  243. false
  244. )
  245. }
  246. }, this.options.maxLifetimeSeconds * 1000)
  247. maxLifetimeTimeout.unref()
  248. client.once('end', () => clearTimeout(maxLifetimeTimeout))
  249. }
  250. return this._acquireClient(client, pendingItem, idleListener, true)
  251. }
  252. })
  253. }
  254. // acquire a client for a pending work item
  255. _acquireClient(client, pendingItem, idleListener, isNew) {
  256. if (isNew) {
  257. this.emit('connect', client)
  258. }
  259. this.emit('acquire', client)
  260. client.release = this._releaseOnce(client, idleListener)
  261. client.removeListener('error', idleListener)
  262. if (!pendingItem.timedOut) {
  263. if (isNew && this.options.verify) {
  264. this.options.verify(client, (err) => {
  265. if (err) {
  266. client.release(err)
  267. return pendingItem.callback(err, undefined, NOOP)
  268. }
  269. pendingItem.callback(undefined, client, client.release)
  270. })
  271. } else {
  272. pendingItem.callback(undefined, client, client.release)
  273. }
  274. } else {
  275. if (isNew && this.options.verify) {
  276. this.options.verify(client, client.release)
  277. } else {
  278. client.release()
  279. }
  280. }
  281. }
  282. // returns a function that wraps _release and throws if called more than once
  283. _releaseOnce(client, idleListener) {
  284. let released = false
  285. return (err) => {
  286. if (released) {
  287. throwOnDoubleRelease()
  288. }
  289. released = true
  290. this._release(client, idleListener, err)
  291. }
  292. }
  293. // release a client back to the poll, include an error
  294. // to remove it from the pool
  295. _release(client, idleListener, err) {
  296. client.on('error', idleListener)
  297. client._poolUseCount = (client._poolUseCount || 0) + 1
  298. this.emit('release', err, client)
  299. // TODO(bmc): expose a proper, public interface _queryable and _ending
  300. if (err || this.ending || !client._queryable || client._ending || client._poolUseCount >= this.options.maxUses) {
  301. if (client._poolUseCount >= this.options.maxUses) {
  302. this.log('remove expended client')
  303. }
  304. return this._remove(client, this._pulseQueue.bind(this))
  305. }
  306. const isExpired = this._expired.has(client)
  307. if (isExpired) {
  308. this.log('remove expired client')
  309. this._expired.delete(client)
  310. return this._remove(client, this._pulseQueue.bind(this))
  311. }
  312. // idle timeout
  313. let tid
  314. if (this.options.idleTimeoutMillis && this._isAboveMin()) {
  315. tid = setTimeout(() => {
  316. this.log('remove idle client')
  317. this._remove(client, this._pulseQueue.bind(this))
  318. }, this.options.idleTimeoutMillis)
  319. if (this.options.allowExitOnIdle) {
  320. // allow Node to exit if this is all that's left
  321. tid.unref()
  322. }
  323. }
  324. if (this.options.allowExitOnIdle) {
  325. client.unref()
  326. }
  327. this._idle.push(new IdleItem(client, idleListener, tid))
  328. this._pulseQueue()
  329. }
  330. query(text, values, cb) {
  331. // guard clause against passing a function as the first parameter
  332. if (typeof text === 'function') {
  333. const response = promisify(this.Promise, text)
  334. setImmediate(function () {
  335. return response.callback(new Error('Passing a function as the first parameter to pool.query is not supported'))
  336. })
  337. return response.result
  338. }
  339. // allow plain text query without values
  340. if (typeof values === 'function') {
  341. cb = values
  342. values = undefined
  343. }
  344. const response = promisify(this.Promise, cb)
  345. cb = response.callback
  346. this.connect((err, client) => {
  347. if (err) {
  348. return cb(err)
  349. }
  350. let clientReleased = false
  351. const onError = (err) => {
  352. if (clientReleased) {
  353. return
  354. }
  355. clientReleased = true
  356. client.release(err)
  357. cb(err)
  358. }
  359. client.once('error', onError)
  360. this.log('dispatching query')
  361. try {
  362. client.query(text, values, (err, res) => {
  363. this.log('query dispatched')
  364. client.removeListener('error', onError)
  365. if (clientReleased) {
  366. return
  367. }
  368. clientReleased = true
  369. client.release(err)
  370. if (err) {
  371. return cb(err)
  372. }
  373. return cb(undefined, res)
  374. })
  375. } catch (err) {
  376. client.release(err)
  377. return cb(err)
  378. }
  379. })
  380. return response.result
  381. }
  382. end(cb) {
  383. this.log('ending')
  384. if (this.ending) {
  385. const err = new Error('Called end on pool more than once')
  386. return cb ? cb(err) : this.Promise.reject(err)
  387. }
  388. this.ending = true
  389. const promised = promisify(this.Promise, cb)
  390. this._endCallback = promised.callback
  391. this._pulseQueue()
  392. return promised.result
  393. }
  394. get waitingCount() {
  395. return this._pendingQueue.length
  396. }
  397. get idleCount() {
  398. return this._idle.length
  399. }
  400. get expiredCount() {
  401. return this._clients.reduce((acc, client) => acc + (this._expired.has(client) ? 1 : 0), 0)
  402. }
  403. get totalCount() {
  404. return this._clients.length
  405. }
  406. }
  407. module.exports = Pool