index.js 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467
  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.maxUses = this.options.maxUses || Infinity
  79. this.options.allowExitOnIdle = this.options.allowExitOnIdle || false
  80. this.options.maxLifetimeSeconds = this.options.maxLifetimeSeconds || 0
  81. this.log = this.options.log || function () {}
  82. this.Client = this.options.Client || Client || require('pg').Client
  83. this.Promise = this.options.Promise || global.Promise
  84. if (typeof this.options.idleTimeoutMillis === 'undefined') {
  85. this.options.idleTimeoutMillis = 10000
  86. }
  87. this._clients = []
  88. this._idle = []
  89. this._expired = new WeakSet()
  90. this._pendingQueue = []
  91. this._endCallback = undefined
  92. this.ending = false
  93. this.ended = false
  94. }
  95. _isFull() {
  96. return this._clients.length >= this.options.max
  97. }
  98. _pulseQueue() {
  99. this.log('pulse queue')
  100. if (this.ended) {
  101. this.log('pulse queue ended')
  102. return
  103. }
  104. if (this.ending) {
  105. this.log('pulse queue on ending')
  106. if (this._idle.length) {
  107. this._idle.slice().map((item) => {
  108. this._remove(item.client)
  109. })
  110. }
  111. if (!this._clients.length) {
  112. this.ended = true
  113. this._endCallback()
  114. }
  115. return
  116. }
  117. // if we don't have any waiting, do nothing
  118. if (!this._pendingQueue.length) {
  119. this.log('no queued requests')
  120. return
  121. }
  122. // if we don't have any idle clients and we have no more room do nothing
  123. if (!this._idle.length && this._isFull()) {
  124. return
  125. }
  126. const pendingItem = this._pendingQueue.shift()
  127. if (this._idle.length) {
  128. const idleItem = this._idle.pop()
  129. clearTimeout(idleItem.timeoutId)
  130. const client = idleItem.client
  131. client.ref && client.ref()
  132. const idleListener = idleItem.idleListener
  133. return this._acquireClient(client, pendingItem, idleListener, false)
  134. }
  135. if (!this._isFull()) {
  136. return this.newClient(pendingItem)
  137. }
  138. throw new Error('unexpected condition')
  139. }
  140. _remove(client) {
  141. const removed = removeWhere(this._idle, (item) => item.client === client)
  142. if (removed !== undefined) {
  143. clearTimeout(removed.timeoutId)
  144. }
  145. this._clients = this._clients.filter((c) => c !== client)
  146. client.end()
  147. this.emit('remove', client)
  148. }
  149. connect(cb) {
  150. if (this.ending) {
  151. const err = new Error('Cannot use a pool after calling end on the pool')
  152. return cb ? cb(err) : this.Promise.reject(err)
  153. }
  154. const response = promisify(this.Promise, cb)
  155. const result = response.result
  156. // if we don't have to connect a new client, don't do so
  157. if (this._isFull() || this._idle.length) {
  158. // if we have idle clients schedule a pulse immediately
  159. if (this._idle.length) {
  160. process.nextTick(() => this._pulseQueue())
  161. }
  162. if (!this.options.connectionTimeoutMillis) {
  163. this._pendingQueue.push(new PendingItem(response.callback))
  164. return result
  165. }
  166. const queueCallback = (err, res, done) => {
  167. clearTimeout(tid)
  168. response.callback(err, res, done)
  169. }
  170. const pendingItem = new PendingItem(queueCallback)
  171. // set connection timeout on checking out an existing client
  172. const tid = setTimeout(() => {
  173. // remove the callback from pending waiters because
  174. // we're going to call it with a timeout error
  175. removeWhere(this._pendingQueue, (i) => i.callback === queueCallback)
  176. pendingItem.timedOut = true
  177. response.callback(new Error('timeout exceeded when trying to connect'))
  178. }, this.options.connectionTimeoutMillis)
  179. this._pendingQueue.push(pendingItem)
  180. return result
  181. }
  182. this.newClient(new PendingItem(response.callback))
  183. return result
  184. }
  185. newClient(pendingItem) {
  186. const client = new this.Client(this.options)
  187. this._clients.push(client)
  188. const idleListener = makeIdleListener(this, client)
  189. this.log('checking client timeout')
  190. // connection timeout logic
  191. let tid
  192. let timeoutHit = false
  193. if (this.options.connectionTimeoutMillis) {
  194. tid = setTimeout(() => {
  195. this.log('ending client due to timeout')
  196. timeoutHit = true
  197. // force kill the node driver, and let libpq do its teardown
  198. client.connection ? client.connection.stream.destroy() : client.end()
  199. }, this.options.connectionTimeoutMillis)
  200. }
  201. this.log('connecting new client')
  202. client.connect((err) => {
  203. if (tid) {
  204. clearTimeout(tid)
  205. }
  206. client.on('error', idleListener)
  207. if (err) {
  208. this.log('client failed to connect', err)
  209. // remove the dead client from our list of clients
  210. this._clients = this._clients.filter((c) => c !== client)
  211. if (timeoutHit) {
  212. err.message = 'Connection terminated due to connection timeout'
  213. }
  214. // this client won’t be released, so move on immediately
  215. this._pulseQueue()
  216. if (!pendingItem.timedOut) {
  217. pendingItem.callback(err, undefined, NOOP)
  218. }
  219. } else {
  220. this.log('new client connected')
  221. if (this.options.maxLifetimeSeconds !== 0) {
  222. const maxLifetimeTimeout = setTimeout(() => {
  223. this.log('ending client due to expired lifetime')
  224. this._expired.add(client)
  225. const idleIndex = this._idle.findIndex((idleItem) => idleItem.client === client)
  226. if (idleIndex !== -1) {
  227. this._acquireClient(
  228. client,
  229. new PendingItem((err, client, clientRelease) => clientRelease()),
  230. idleListener,
  231. false
  232. )
  233. }
  234. }, this.options.maxLifetimeSeconds * 1000)
  235. maxLifetimeTimeout.unref()
  236. client.once('end', () => clearTimeout(maxLifetimeTimeout))
  237. }
  238. return this._acquireClient(client, pendingItem, idleListener, true)
  239. }
  240. })
  241. }
  242. // acquire a client for a pending work item
  243. _acquireClient(client, pendingItem, idleListener, isNew) {
  244. if (isNew) {
  245. this.emit('connect', client)
  246. }
  247. this.emit('acquire', client)
  248. client.release = this._releaseOnce(client, idleListener)
  249. client.removeListener('error', idleListener)
  250. if (!pendingItem.timedOut) {
  251. if (isNew && this.options.verify) {
  252. this.options.verify(client, (err) => {
  253. if (err) {
  254. client.release(err)
  255. return pendingItem.callback(err, undefined, NOOP)
  256. }
  257. pendingItem.callback(undefined, client, client.release)
  258. })
  259. } else {
  260. pendingItem.callback(undefined, client, client.release)
  261. }
  262. } else {
  263. if (isNew && this.options.verify) {
  264. this.options.verify(client, client.release)
  265. } else {
  266. client.release()
  267. }
  268. }
  269. }
  270. // returns a function that wraps _release and throws if called more than once
  271. _releaseOnce(client, idleListener) {
  272. let released = false
  273. return (err) => {
  274. if (released) {
  275. throwOnDoubleRelease()
  276. }
  277. released = true
  278. this._release(client, idleListener, err)
  279. }
  280. }
  281. // release a client back to the poll, include an error
  282. // to remove it from the pool
  283. _release(client, idleListener, err) {
  284. client.on('error', idleListener)
  285. client._poolUseCount = (client._poolUseCount || 0) + 1
  286. this.emit('release', err, client)
  287. // TODO(bmc): expose a proper, public interface _queryable and _ending
  288. if (err || this.ending || !client._queryable || client._ending || client._poolUseCount >= this.options.maxUses) {
  289. if (client._poolUseCount >= this.options.maxUses) {
  290. this.log('remove expended client')
  291. }
  292. this._remove(client)
  293. this._pulseQueue()
  294. return
  295. }
  296. const isExpired = this._expired.has(client)
  297. if (isExpired) {
  298. this.log('remove expired client')
  299. this._expired.delete(client)
  300. this._remove(client)
  301. this._pulseQueue()
  302. return
  303. }
  304. // idle timeout
  305. let tid
  306. if (this.options.idleTimeoutMillis) {
  307. tid = setTimeout(() => {
  308. this.log('remove idle client')
  309. this._remove(client)
  310. }, this.options.idleTimeoutMillis)
  311. if (this.options.allowExitOnIdle) {
  312. // allow Node to exit if this is all that's left
  313. tid.unref()
  314. }
  315. }
  316. if (this.options.allowExitOnIdle) {
  317. client.unref()
  318. }
  319. this._idle.push(new IdleItem(client, idleListener, tid))
  320. this._pulseQueue()
  321. }
  322. query(text, values, cb) {
  323. // guard clause against passing a function as the first parameter
  324. if (typeof text === 'function') {
  325. const response = promisify(this.Promise, text)
  326. setImmediate(function () {
  327. return response.callback(new Error('Passing a function as the first parameter to pool.query is not supported'))
  328. })
  329. return response.result
  330. }
  331. // allow plain text query without values
  332. if (typeof values === 'function') {
  333. cb = values
  334. values = undefined
  335. }
  336. const response = promisify(this.Promise, cb)
  337. cb = response.callback
  338. this.connect((err, client) => {
  339. if (err) {
  340. return cb(err)
  341. }
  342. let clientReleased = false
  343. const onError = (err) => {
  344. if (clientReleased) {
  345. return
  346. }
  347. clientReleased = true
  348. client.release(err)
  349. cb(err)
  350. }
  351. client.once('error', onError)
  352. this.log('dispatching query')
  353. try {
  354. client.query(text, values, (err, res) => {
  355. this.log('query dispatched')
  356. client.removeListener('error', onError)
  357. if (clientReleased) {
  358. return
  359. }
  360. clientReleased = true
  361. client.release(err)
  362. if (err) {
  363. return cb(err)
  364. }
  365. return cb(undefined, res)
  366. })
  367. } catch (err) {
  368. client.release(err)
  369. return cb(err)
  370. }
  371. })
  372. return response.result
  373. }
  374. end(cb) {
  375. this.log('ending')
  376. if (this.ending) {
  377. const err = new Error('Called end on pool more than once')
  378. return cb ? cb(err) : this.Promise.reject(err)
  379. }
  380. this.ending = true
  381. const promised = promisify(this.Promise, cb)
  382. this._endCallback = promised.callback
  383. this._pulseQueue()
  384. return promised.result
  385. }
  386. get waitingCount() {
  387. return this._pendingQueue.length
  388. }
  389. get idleCount() {
  390. return this._idle.length
  391. }
  392. get expiredCount() {
  393. return this._clients.reduce((acc, client) => acc + (this._expired.has(client) ? 1 : 0), 0)
  394. }
  395. get totalCount() {
  396. return this._clients.length
  397. }
  398. }
  399. module.exports = Pool