feat: Wait for msgs to be processed before emitting stopped (#454)

* feat: Wait for msgs to be processed before emitting stopped

* chore: Update README with waitForInFlightMessages info

* fix: Await final poll for graceful shutdown

* refactor: Rename waitForInFlightMessagesMs to pollingCompleteWaitTimeMs and move to constructor

* chore: Rename waitForInFlightMessagesToComplete to waitForPollingToComplete

* refactor: Use setTimeout instead of Promise/loop in waitForPollingToComplete

* chore: Remove outdated info from README

* chore: Rename beganWaitingForStop to stopRequestedAtTimestamp

Author: barwin

src/consumer.ts

@@ -54,7 +54,10 @@ export class Consumer extends TypedEventEmitter {
   private waitTimeSeconds: number;
   private authenticationErrorTimeout: number;
   private pollingWaitTimeMs: number;
+  private pollingCompleteWaitTimeMs: number;
   private heartbeatInterval: number;
+  private isPolling: boolean;
+  private stopRequestedAtTimestamp: number;
   public abortController: AbortController;
 
   constructor(options: ConsumerOptions) {
@@ -77,6 +80,7 @@ export class Consumer extends TypedEventEmitter {
     this.authenticationErrorTimeout =
       options.authenticationErrorTimeout ?? 10000;
     this.pollingWaitTimeMs = options.pollingWaitTimeMs ?? 0;
+    this.pollingCompleteWaitTimeMs = options.pollingCompleteWaitTimeMs ?? 0;
     this.shouldDeleteMessages = options.shouldDeleteMessages ?? true;
     this.alwaysAcknowledge = options.alwaysAcknowledge ?? false;
     this.sqs =
@@ -142,7 +146,31 @@ export class Consumer extends TypedEventEmitter {
       this.emit('aborted');
     }
 
-    this.emit('stopped');
+    this.stopRequestedAtTimestamp = Date.now();
+    this.waitForPollingToComplete();
+  }
+
+  /**
+   * Wait for final poll and in flight messages to complete.
+   * @private
+   */
+  private waitForPollingToComplete(): void {
+    if (!this.isPolling || !(this.pollingCompleteWaitTimeMs > 0)) {
+      this.emit('stopped');
+      return;
+    }
+
+    const exceededTimeout =
+      Date.now() - this.stopRequestedAtTimestamp >
+      this.pollingCompleteWaitTimeMs;
+    if (exceededTimeout) {
+      logger.debug('waiting_for_polling_to_complete_timeout_exceeded');
+      this.emit('stopped');
+      return;
+    }
+
+    logger.debug('waiting_for_polling_to_complete');
+    setTimeout(this.waitForPollingToComplete, 1000);
   }
 
   /**
@@ -198,6 +226,8 @@ export class Consumer extends TypedEventEmitter {
 
     logger.debug('polling');
 
+    this.isPolling = true;
+
     let currentPollingTimeout = this.pollingWaitTimeMs;
     this.receiveMessage({
       QueueUrl: this.queueUrl,
@@ -227,6 +257,9 @@ export class Consumer extends TypedEventEmitter {
       })
       .catch((err) => {
         this.emitError(err);
+      })
+      .finally(() => {
+        this.isPolling = false;
       });
   }
 

src/types.ts

@@ -46,6 +46,12 @@ export interface ConsumerOptions {
    * @defaultvalue `0`
    */
   pollingWaitTimeMs?: number;
+  /**
+   * If you want the stop action to wait for the final poll to complete and in-flight messages
+   * to be processed before emitting 'stopped' set this to the max amount of time to wait.
+   * @defaultvalue `0`
+   */
+  pollingCompleteWaitTimeMs?: number;
   /**
    * If true, sets the message visibility timeout to 0 after a `processing_error`.
    * @defaultvalue `false`

test/features/gracefulShutdown.feature

@@ -0,0 +1,6 @@
+Feature: Graceful shutdown
+
+  Scenario: Several messages in flight
+    Given Several messages are sent to the SQS queue
+    Then the application is stopped while messages are in flight
+    Then the in-flight messages should be processed before stopped is emitted

test/features/step_definitions/gracefulShutdown.js

@@ -0,0 +1,56 @@
+const { Given, Then, After } = require('@cucumber/cucumber');
+const assert = require('assert');
+const { PurgeQueueCommand } = require('@aws-sdk/client-sqs');
+const pEvent = require('p-event');
+
+const { consumer } = require('../utils/consumer/gracefulShutdown');
+const { producer } = require('../utils/producer');
+const { sqs, QUEUE_URL } = require('../utils/sqs');
+
+Given('Several messages are sent to the SQS queue', async () => {
+  const params = {
+    QueueUrl: QUEUE_URL
+  };
+  const command = new PurgeQueueCommand(params);
+  const response = await sqs.send(command);
+
+  assert.strictEqual(response['$metadata'].httpStatusCode, 200);
+
+  const size = await producer.queueSize();
+  assert.strictEqual(size, 0);
+
+  await producer.send(['msg1', 'msg2', 'msg3']);
+
+  const size2 = await producer.queueSize();
+
+  assert.strictEqual(size2, 3);
+});
+
+Then('the application is stopped while messages are in flight', async () => {
+  consumer.start();
+
+  consumer.stop();
+
+  assert.strictEqual(consumer.isRunning, false);
+});
+
+Then(
+  'the in-flight messages should be processed before stopped is emitted',
+  async () => {
+    let numProcessed = 0;
+    consumer.on('message_processed', () => {
+      numProcessed++;
+    });
+
+    await pEvent(consumer, 'stopped');
+
+    assert.strictEqual(numProcessed, 3);
+
+    const size = await producer.queueSize();
+    assert.strictEqual(size, 0);
+  }
+);
+
+After(() => {
+  return consumer.stop();
+});

test/features/utils/consumer/gracefulShutdown.js

@@ -0,0 +1,17 @@
+const { Consumer } = require('../../../../dist/consumer');
+
+const { QUEUE_URL, sqs } = require('../sqs');
+
+const consumer = Consumer.create({
+  queueUrl: QUEUE_URL,
+  sqs,
+  pollingWaitTimeMs: 1000,
+  pollingCompleteWaitTimeMs: 5000,
+  batchSize: 10,
+  handleMessage: async (message) => {
+    await new Promise((resolve) => setTimeout(resolve, 1500));
+    return message;
+  }
+});
+
+exports.consumer = consumer;

test/tests/consumer.test.ts

@@ -1481,6 +1481,79 @@ describe('Consumer', () => {
       sandbox.assert.calledOnce(handleAbort);
       sandbox.assert.calledOnce(handleStop);
     });
+
+    it('waits for in-flight messages before emitting stopped (within timeout)', async () => {
+      const handleStop = sandbox.stub().returns(null);
+      const handleResponseProcessed = sandbox.stub().returns(null);
+
+      // A slow message handler
+      handleMessage = sandbox.stub().callsFake(async () => {
+        await new Promise((resolve) => setTimeout(resolve, 5000));
+      });
+
+      consumer = new Consumer({
+        queueUrl: QUEUE_URL,
+        region: REGION,
+        handleMessage,
+        sqs,
+        pollingCompleteWaitTimeMs: 5000,
+        authenticationErrorTimeout: AUTHENTICATION_ERROR_TIMEOUT
+      });
+
+      consumer.on('stopped', handleStop);
+      consumer.on('response_processed', handleResponseProcessed);
+
+      consumer.start();
+      await clock.nextAsync();
+      consumer.stop();
+
+      await clock.runAllAsync();
+
+      sandbox.assert.calledOnce(handleStop);
+      sandbox.assert.calledOnce(handleResponseProcessed);
+      sandbox.assert.calledOnce(handleMessage);
+      assert.ok(handleMessage.calledBefore(handleStop));
+
+      // handleResponseProcessed is called after handleMessage, indicating
+      // messages were allowed to complete before 'stopped' was emitted
+      assert.ok(handleResponseProcessed.calledBefore(handleStop));
+    });
+
+    it('waits for in-flight messages before emitting stopped (timeout reached)', async () => {
+      const handleStop = sandbox.stub().returns(null);
+      const handleResponseProcessed = sandbox.stub().returns(null);
+
+      // A slow message handler
+      handleMessage = sandbox.stub().callsFake(async () => {
+        await new Promise((resolve) => setTimeout(resolve, 5000));
+      });
+
+      consumer = new Consumer({
+        queueUrl: QUEUE_URL,
+        region: REGION,
+        handleMessage,
+        sqs,
+        pollingCompleteWaitTimeMs: 500,
+        authenticationErrorTimeout: AUTHENTICATION_ERROR_TIMEOUT
+      });
+
+      consumer.on('stopped', handleStop);
+      consumer.on('response_processed', handleResponseProcessed);
+
+      consumer.start();
+      await clock.nextAsync();
+      consumer.stop();
+
+      await clock.runAllAsync();
+
+      sandbox.assert.calledOnce(handleStop);
+      sandbox.assert.calledOnce(handleResponseProcessed);
+      sandbox.assert.calledOnce(handleMessage);
+      assert(handleMessage.calledBefore(handleStop));
+
+      // Stop was called before the message could be processed, because we reached timeout.
+      assert(handleStop.calledBefore(handleResponseProcessed));
+    });
   });
 
   describe('isRunning', async () => {