Expose stoppable.increment and stoppable.decrement

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) {

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."