Merge pull request #72 from Workable/sigterm-handling

Handle SIGTERM signal

Author: nikosd23

test/rabbit.test.ts

@@ -7,11 +7,12 @@ import Exchange from '../ts/exchange';
 import * as ReplyQueue from '../ts/reply-queue';
 import * as DelayQueue from '../ts/delay-queue';
 import * as amqp from 'amqplib';
+import { randomUUID } from 'crypto';
 
 const sandbox = sinon.createSandbox();
 
 describe('Test rabbit class', function() {
-  let rabbit: Rabbit;
+  let rabbit: Rabbit | null;
 
   before(function() {
     this.url = process.env.RABBIT_URL || 'amqp://localhost';
@@ -364,4 +365,32 @@ describe('Test rabbit class', function() {
     await rabbit.destroyQueue(this.name);
     stub.calledWith(rabbit.consumeChannel, `test_${this.name}`).should.be.true();
   });
+
+  it('should register sigtermHandler in SIGTERM process listeners and remove on close', async function() {
+    rabbit = new Rabbit(this.url);
+    await rabbit.connected;
+
+    const handler = (rabbit as any).sigtermHandler;
+
+    process.listeners('SIGTERM').should.containEql(handler);
+
+    rabbit.close();
+    rabbit = null;
+
+    process.listeners('SIGTERM').should.not.containEql(handler);
+  });
+
+  it('should call unsubscribe for all queues on SIGTERM', async function() {
+    const stub = sandbox.stub(Queue.prototype, 'unsubscribe').resolves();
+
+    rabbit = new Rabbit(this.url);
+    await rabbit.connected;
+
+    const handler = () => ({});
+    await rabbit.createQueue(randomUUID().toString(), { exclusive: true }, handler);
+    await rabbit.createQueue(randomUUID().toString(), { exclusive: true }, handler);
+
+    (rabbit as any).sigtermHandler();
+    stub.calledTwice.should.be.true();
+  });
 });

ts/queue.ts

@@ -69,12 +69,13 @@ export default class Queue {
   async subscribe(handler: (msg: any, ack: (error?, reply?) => any) => any) {
     await this.created;
     this.handler = handler;
-    let tag = await this.channel.consume(this.name, this.onMessage.bind(this), { noAck: this.options.noAck });
-    this.tag = tag;
+    this.tag = await this.channel.consume(this.name, this.onMessage.bind(this), { noAck: this.options.noAck });
   }
 
   async unsubscribe() {
-    await this.channel.cancel(this.tag.consumerTag);
+    if (this.tag) {
+      await this.channel.cancel(this.tag.consumerTag);
+    }
     this.handler = null;
     this.tag = null;
   }

ts/rabbit.ts

@@ -27,6 +27,7 @@ export default class Rabbit extends EventEmitter {
   public socketOptions;
   private defaultQueueType: string | undefined;
   public logger: ReturnType<typeof getLogger>;
+  private readonly sigtermHandler: () => void;
 
   constructor(
     public url: string,
@@ -58,7 +59,10 @@ export default class Rabbit extends EventEmitter {
     this.scheduledPublish = scheduledPublish;
     this.socketOptions = socketOptions;
     this.defaultQueueType = defaultQueueType;
+    this.sigtermHandler = this.handleSigterm.bind(this);
     this.reconnect();
+
+    process.once('SIGTERM', this.sigtermHandler);
   }
 
   private async connect() {
@@ -239,7 +243,22 @@ export default class Rabbit extends EventEmitter {
   }
 
   async close() {
+    process.off('SIGTERM', this.sigtermHandler);
+
     await this.consumeConnection.close();
     await this.publishConnection.close();
   }
+
+  private handleSigterm() {
+    this.logger.info('Received SIGTERM. Cancelling consumers..');
+
+    for (const queue of Object.values(this.queues)) {
+      queue.unsubscribe()
+        .then(() => {
+          this.logger.debug(`Consumer for queue ${queue.name} cancelled.`);
+        }).catch(e => {
+          this.logger.error(`Consumer for queue ${queue.name} failed to cancel`, e);
+        });
+    }
+  }
 }