Add graceful shutdown for requesters and responders

examples/requester.js

@@ -32,3 +32,17 @@ function makeRequest() {
 makeRequest();
 
 setInterval(makeRequest, 5000);
+
+// Gracefully close responder after it completes any pending messages
+process.once('SIGINT', () => {
+    console.log('closing, press ctrl+c again to force exit');
+    randomRequest.close(() => {
+        console.log('exiting');
+        process.exit();
+    });
+
+    process.once('SIGINT', () => {
+        console.log('forced exit');
+        process.exit();
+    });
+});

examples/responder.js

@@ -27,3 +27,17 @@ randomResponder.on('promised request', function(req) {
         // reject(answer);
     });
 });
+
+// Gracefully close responder after it completes any pending messages
+process.once('SIGINT', () => {
+    console.log('closing, press ctrl+c again to force exit');
+    randomResponder.close(() => {
+        console.log('exiting');
+        process.exit();
+    });
+
+    process.once('SIGINT', () => {
+        console.log('forced exit');
+        process.exit();
+    });
+});

src/components/component.js

@@ -48,8 +48,28 @@ module.exports = class Component extends EventEmitter {
 
     onRemoved() { };
 
-    close() {
-        this.sock && this.sock.close();
+    close(cb) {
         this.discovery && this.discovery.stop();
+
+        if (typeof cb === 'function') {
+            const interval = setInterval(() => {
+                if (!this.messageIds || this.messageIds.length === 0) {
+                    if (this.sock) {
+                        if (this.sock.server) {
+                            this.sock.close(cb);
+                        } else {
+                            this.sock.close();
+                            cb();
+                        }
+                    } else {
+                        cb();
+                    }
+                    clearInterval(interval);
+                }
+            }, 100);
+            return;
+        } else {
+            this.sock && this.sock.close();
+        }
     }
 };

src/components/requester.js

@@ -3,6 +3,7 @@ const Monitorable = require('./monitorable');
 const Component = require('./component');
 const axon = require('@dashersw/axon');
 const debug = require('debug')('axon:req');
+const uuid = require('uuid');
 
 const SUBSET_IDENTIFIER = '__subset';
 
@@ -11,6 +12,8 @@ module.exports = class Requester extends Monitorable(Configurable(Component)) {
     constructor(advertisement, discoveryOptions) {
         super(advertisement, discoveryOptions);
 
+        this.messageIds = [];
+
         this.sock = new axon.types[this.type]();
         this.sock.set('retry timeout', 0);
         this.timeout = advertisement.timeout || process.env.COTE_REQUEST_TIMEOUT;
@@ -91,23 +94,50 @@ module.exports = class Requester extends Monitorable(Configurable(Component)) {
 
         if (alreadyConnected) return;
 
-        this.sock.connect(obj.advertisement.port, address);
+        this.sock.connect(obj.advertisement.port, address, sock => {
+            // Add the closing callback for the responder
+            const key = `closing__${address}:${sock.localPort}`;
+            this.sock.callbacks[key] = () => {
+                // Remove the socket so it no longer sends messages to it
+                const sockIndex = this.sock.socks.findIndex(sock => `closing__${sock._sockname.address}:${sock.localPort}` === key);
+                sockIndex >= 0 && this.sock.socks.splice(sockIndex, 1);
+            };
+        });
     }
 
     send(...args) {
         const hasCallback = typeof args[args.length - 1] == 'function';
         const timeout = args[0].__timeout || this.timeout;
 
-        if (hasCallback) return sendOverSocket(this.sock, timeout, ...args);
+        const cb = hasCallback && args.pop();
+        const messageId = uuid.v4();
+        this.messageIds.push(messageId);
+        const cbWithCounter = (...cbArgs) => {
+            cb && cb(...cbArgs);
+            const index = this.messageIds.indexOf(messageId);
+            index >= 0 && this.messageIds.splice(index, 1);
+        };
+
+        if (hasCallback) return sendOverSocket(this.sock, timeout, ...args, cbWithCounter);
 
         return new Promise((resolve, reject) => {
             sendOverSocket(this.sock, timeout, ...args, (err, res) => {
+                cbWithCounter();
                 if (err) return reject(err);
                 resolve(res);
             });
         });
     }
 
+    close(cb) {
+        if (cb) {
+            // Disable send so that it will stop queueing messages
+            this.send = () => {};
+        }
+
+        super.close(cb);
+    }
+
     get type() {
         return 'req';
     }

src/components/responder.js

@@ -2,6 +2,7 @@ const axon = require('@dashersw/axon');
 const portfinder = require('portfinder');
 const Configurable = require('./configurable');
 const Component = require('./component');
+const uuid = require('uuid');
 
 // eslint-disable-next-line
 const colors = require('colors');
@@ -10,6 +11,8 @@ module.exports = class Responder extends Configurable(Component) {
     constructor(advertisement, discoveryOptions) {
         super(advertisement, discoveryOptions);
 
+        this.messageIds = [];
+
         this.sock = new axon.types[this.type]();
         this.sock.on('bind', () => this.startDiscovery());
 
@@ -20,7 +23,15 @@ module.exports = class Responder extends Configurable(Component) {
                 this.discovery.log([this.advertisement.name, '>', `No listeners found for event: ${req.type}`.yellow]);
             }
 
-            this.emit(req.type, req, cb);
+            const messageId = uuid.v4();
+            this.messageIds.push(messageId);
+            const cbWithCounter = (...cbArgs) => {
+                cb(...cbArgs);
+                const index = this.messageIds.indexOf(messageId);
+                index >= 0 && this.messageIds.splice(index, 1);
+            };
+
+            this.emit(req.type, req, cbWithCounter);
         });
 
         const onPort = (err, port) => {
@@ -56,6 +67,18 @@ module.exports = class Responder extends Configurable(Component) {
         });
     }
 
+    close(cb) {
+        if (cb) {
+            // Send closing event to all requesters so they will stop sending messages to it
+            for (const sock of this.sock.socks) {
+                const key = `closing__${sock._peername.address}:${sock.remotePort}`;
+                sock.writable && sock.write(this.sock.pack([null, key]));
+            }
+        }
+
+        super.close(cb);
+    }
+
     get type() {
         return 'rep';
     }

test/component-lifecycle.js

@@ -7,6 +7,28 @@ const { Requester, Responder } = require('../')({ environment });
 
 LogSuppress.init(console);
 
+const setup = (t) => {
+    const key = r.generate();
+
+    const requester = new Requester({ name: `${t.title}: ${key} requester`, key });
+    const responder = new Responder({ name: `${t.title}: ${key} responder`, key });
+
+    let i = 1;
+    responder.on('ping', (req, cb) => {
+        i++;
+        setTimeout(() => {
+            cb(null, 'pong');
+        }, 500 * i);
+    });
+
+    requester.send({ type: 'ping' });
+    setInterval(() => {
+        requester.send({ type: 'ping' });
+    }, 500);
+
+    return { requester, responder };
+};
+
 test('Instantiate a requester', (t) => {
     const key = r.generate();
     const requester = new Requester({ name: `${t.title}: requester`, key });
@@ -29,3 +51,60 @@ test.cb('Discover and close a requester', (t) => {
         t.end();
     });
 });
+
+test.cb(`Component should call close callback if no sock`, (t) => {
+    const key = r.generate();
+    const requester = new Requester({ name: `${t.title}: ${key} requester`, key });
+
+    setTimeout(() => {
+        requester.sock = null;
+        requester.close(() => {
+            t.is(requester.sock, null);
+            t.end();
+        });
+    }, 750);
+});
+
+test.cb(`Requester should close socket immediately if no callback`, (t) => {
+    const { requester } = setup(t);
+
+    setTimeout(() => {
+        requester.close();
+        t.is(requester.messageIds.length, 2);
+        t.end();
+    }, 750);
+});
+
+test.cb(`Requester should wait for all messages to complete before close callback`, (t) => {
+    const { requester } = setup(t);
+
+    setTimeout(() => {
+        requester.close(() => {
+            t.is(requester.messageIds.length, 0);
+            t.end();
+        });
+        t.is(requester.messageIds.length, 2);
+    }, 750);
+});
+
+test.cb(`Responder should close socket immediately if no callback`, (t) => {
+    const { responder } = setup(t);
+
+    setTimeout(() => {
+        responder.close();
+        t.is(responder.messageIds.length, 2);
+        t.end();
+    }, 750);
+});
+
+test.cb(`Responder should wait for all messages to complete before close callback`, (t) => {
+    const { responder } = setup(t);
+
+    setTimeout(() => {
+        responder.close(() => {
+            t.is(responder.messageIds.length, 0);
+            t.end();
+        });
+        t.is(responder.messageIds.length, 2);
+    }, 750);
+});