diff --git a/lib/stoppable.js b/lib/stoppable.js index 33736fb..f3733a2 100644 --- a/lib/stoppable.js +++ b/lib/stoppable.js @@ -6,6 +6,10 @@ module.exports = (server, grace = Infinity) => { server.on('secureConnection', onConnection) server.on('request', onRequest) server.stop = stop + server.stoppable = { + increment, + decrement + } server._pendingSockets = reqsPerSocket return server @@ -15,14 +19,23 @@ module.exports = (server, grace = Infinity) => { } function onRequest (req, res) { - reqsPerSocket.set(req.socket, reqsPerSocket.get(req.socket) + 1) - res.once('finish', () => { - const pending = reqsPerSocket.get(req.socket) - 1 - reqsPerSocket.set(req.socket, pending) - if (stopped && pending === 0) { - req.socket.end() - } - }) + increment(req.socket) + res.once('finish', () => decrement(req.socket)) + } + + function increment (socket) { + const counter = reqsPerSocket.get(socket) + 1 + reqsPerSocket.set(socket, counter) + return counter + } + + function decrement (socket, callback) { + const counter = reqsPerSocket.get(socket) - 1 + reqsPerSocket.set(socket, counter) + if (stopped && counter === 0) { + (callback || (() => socket.end()))() + } + return counter } function stop(callback) { diff --git a/readme.md b/readme.md index 37084c8..7ffc1c5 100644 --- a/readme.md +++ b/readme.md @@ -48,6 +48,37 @@ Closes the server. - callback: passed along to the existing `server.close` function to auto-register a 'close' event +## Manual operations on pending counts + +It's also possible to manually increment and decrement a pending count of a specific socket. +This can be helpful when you want to lock the server while doing some jobs. +For example, you can increment a pending count when you receive a request via Websocket +and decrement it once you have processed it. + +**stoppable.increment** + +```js +server.stoppable.increment(socket) +``` + +Increment a pending count of a specified socket. This method returns the new pending count. + +- socket: A socket of which you want to increment a pending count + +**stoppable.decrement** + +```js +server.stoppable.decrement(socket, callback) +``` + +Decrement a pending count of a specified socket. If the pending count becomes 0 and the server +has already been stopped, it closes the socket. When you specify a callback, it calls +the callback instead of closing the socket. You can do some cleanups in this callback, +but it's your responsibility to close the socket. This method returns the new pending count. + +- socket: A socket of which you want to decrement a pending count +- callback: A cleanup function. The first argument is a socket to be closed. + ## Design decisions - Monkey patching generally sucks, but in this case it's the nicest API. Let's call it "decorating."