stoppable.js 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. 'use strict'
  2. const https = require('https')
  3. module.exports = (server, grace) => {
  4. grace = typeof grace === 'undefined' ? Infinity : grace
  5. const reqsPerSocket = new Map()
  6. let stopped = false
  7. let gracefully = true
  8. if (server instanceof https.Server) {
  9. server.on('secureConnection', onConnection)
  10. } else {
  11. server.on('connection', onConnection)
  12. }
  13. server.on('request', onRequest)
  14. server.stop = stop
  15. server._pendingSockets = reqsPerSocket
  16. return server
  17. function onConnection (socket) {
  18. reqsPerSocket.set(socket, 0)
  19. socket.once('close', () => reqsPerSocket.delete(socket))
  20. }
  21. function onRequest (req, res) {
  22. reqsPerSocket.set(req.socket, reqsPerSocket.get(req.socket) + 1)
  23. res.once('finish', () => {
  24. const pending = reqsPerSocket.get(req.socket) - 1
  25. reqsPerSocket.set(req.socket, pending)
  26. if (stopped && pending === 0) {
  27. req.socket.end()
  28. }
  29. })
  30. }
  31. function stop (callback) {
  32. // allow request handlers to update state before we act on that state
  33. setImmediate(() => {
  34. stopped = true
  35. if (grace < Infinity) {
  36. setTimeout(destroyAll, grace).unref()
  37. }
  38. server.close(e => {
  39. if (callback) {
  40. callback(e, gracefully)
  41. }
  42. })
  43. reqsPerSocket.forEach(endIfIdle)
  44. })
  45. }
  46. function endIfIdle (requests, socket) {
  47. if (requests === 0) socket.end()
  48. }
  49. function destroyAll () {
  50. gracefully = false
  51. reqsPerSocket.forEach((reqs, socket) => socket.end())
  52. setImmediate(() => {
  53. reqsPerSocket.forEach((reqs, socket) => socket.destroy())
  54. })
  55. }
  56. }