Add async lock for consume/disconnect

Author: milindl

LICENSE.kafkajs

@@ -1,5 +1,6 @@
 The promisified API (lib/kafkajs) is inspired by kafkajs (github.com/tulios/kafkajs).
 The promisified tests (test/promisified) are also adapted from there.
+An async lock implementation and many error types are also adapted from there.
 The license notice is reproduced below.
 
 ----

lib/kafkajs/_common.js

@@ -540,6 +540,83 @@ function notImplemented(msg = 'Not implemented') {
   throw new error.KafkaJSError(msg, { code: error.ErrorCodes.ERR__NOT_IMPLEMENTED });
 }
 
+/* Code from the async lock is from github.com/tulios/kafkajs.
+ * For more details, see LICENSE.kafkajs */
+const LockStates = Object.freeze({
+  LOCKED: 'locked',
+  TIMEOUT: 'timeout',
+  WAITING: 'waiting',
+  TIMEOUT_ERROR_MESSAGE: 'timeoutErrorMessage',
+});
+
+class Lock {
+  constructor({ timeout, description = null } = {}) {
+    if (typeof timeout !== 'number') {
+      throw new TypeError(`'timeout' is not a number, received '${typeof timeout}'`);
+    }
+
+    this[LockStates.LOCKED] = false;
+    this[LockStates.TIMEOUT] = timeout;
+    this[LockStates.WAITING] = new Set();
+    this[LockStates.TIMEOUT_ERROR_MESSAGE] = () => {
+      const timeoutMessage = `Timeout while acquiring lock (${this[LockStates.WAITING].size} waiting locks)`;
+      return description ? `${timeoutMessage}: "${description}"` : timeoutMessage;
+    }
+  }
+
+  async acquire() {
+    return new Promise((resolve, reject) => {
+      if (!this[LockStates.LOCKED]) {
+        this[LockStates.LOCKED] = true;
+        return resolve();
+      }
+
+      let timeoutId = null;
+      const tryToAcquire = async () => {
+        if (!this[LockStates.LOCKED]) {
+          this[LockStates.LOCKED] = true;
+          clearTimeout(timeoutId);
+          this[LockStates.WAITING].delete(tryToAcquire);
+          return resolve();
+        }
+      }
+
+      this[LockStates.WAITING].add(tryToAcquire);
+      timeoutId = setTimeout(() => {
+        // The message should contain the number of waiters _including_ this one
+        const e = new error.KafkaJSLockTimeout(this[LockStates.TIMEOUT_ERROR_MESSAGE]());
+        this[LockStates.WAITING].delete(tryToAcquire);
+        reject(e);
+      }, this[LockStates.TIMEOUT]);
+    })
+  }
+
+  async release() {
+    this[LockStates.LOCKED] = false;
+    const waitingLock = this[LockStates.WAITING].values().next().value;
+
+    if (waitingLock) {
+      return waitingLock();
+    }
+  }
+}
+
+/**
+  * Acquires a lock, or logs an error if it fails.
+  * @param {Lock} lock
+  * @param {import("../../types/kafkajs").Logger} logger
+  * @returns {boolean} true if the lock was acquired, false otherwise.
+  */
+async function acquireOrLog(lock, logger) {
+  try {
+    await lock.acquire();
+    return true;
+  } catch (e) {
+    logger.error(`Failed to acquire lock: ${e.message}`);
+  }
+  return false;
+}
+
 module.exports = {
   kafkaJSToRdKafkaConfig,
   topicPartitionOffsetToRdKafka,
@@ -553,5 +630,7 @@ module.exports = {
   CompatibilityErrorMessages,
   severityToLogLevel,
   checkAllowedKeys,
-  checkIfKafkaJsKeysPresent
+  checkIfKafkaJsKeysPresent,
+  Lock,
+  acquireOrLog,
 };

lib/kafkajs/_consumer.js

@@ -11,7 +11,9 @@ const {
   CompatibilityErrorMessages,
   severityToLogLevel,
   checkAllowedKeys,
-  logLevel
+  logLevel,
+  Lock,
+  acquireOrLog
 } = require('./_common');
 const { Buffer } = require('buffer');
 
@@ -91,6 +93,15 @@ class Consumer {
    */
   #logger = new DefaultLogger();
 
+  /**
+   * A lock for consuming and disconnecting.
+   * This lock should be held whenever we want to change the state from CONNECTED to any state other than CONNECTED.
+   * In practical terms, this lock is held whenever we're consuming a message, or disconnecting.
+   * We set the timeout to 5 seconds, after which we log an error, but keep trying to acquire the lock.
+   * @type {Lock}
+   */
+  #lock = new Lock({ timeout: 5000 });
+
   /**
    * @constructor
    * @param {import("../../types/kafkajs").ConsumerConfig} kJSConfig
@@ -575,21 +586,21 @@ class Consumer {
   /* Internal polling loop. It accepts the same config object that `run` accepts. */
   async #runInternal(config) {
     while (this.#state === ConsumerState.CONNECTED) {
+
+      /* We need to acquire a lock here, because we need to ensure that we don't
+      * disconnect while in the middle of processing a message. */
+      if (!(await acquireOrLog(this.#lock, this.#logger)))
+        continue;
+
       const m = await this.#consumeSingle().catch(e => {
         /* Since this error cannot be exposed to the user in the current situation, just log and retry.
          * This is due to restartOnFailure being set to always true. */
         if (this.#logger)
           this.#logger.error(`Consumer encountered error while consuming. Retrying. Error details: ${JSON.stringify(e)}`);
       });
 
-      /* It's possible for the state to change as we await something.
-       * Unfortunately, we have no alternative but to break and drop the message, if it exists.
-       * TODO: fix this, possibly with a flag in disconnect that waits until the run loop winds down.
-       */
-      if (this.#state !== ConsumerState.CONNECTED)
-        break;
-
       if (!m) {
+        this.#lock.release();
         continue;
       }
 
@@ -608,27 +619,15 @@ class Consumer {
         const invalidateMessage = await this.#seekInternal({ topic: m.topic, partition: m.partition });
         if (invalidateMessage) {
           /* Don't pass this message on to the user if this topic partition was seeked to. */
+          this.#lock.release();
           continue;
         }
-
-        /* It's possible for the state to change as we await something.
-         * Unfortunately, we have no alternative but to break and drop the message.
-         * TODO: fix this, possibly with a flag in disconnect that waits until the run loop winds down.
-         */
-        if (this.#state !== ConsumerState.CONNECTED)
-          break;
       }
 
       try {
         await config.eachMessage(
           this.#createPayload(m)
         )
-
-        /* It's possible for the state to change as we await something.
-         * Unfortunately, we have no alternative but to break without taking any action that the user might need.
-         */
-        if (this.#state !== ConsumerState.CONNECTED)
-          break;
       } catch (e) {
         /* It's not only possible, but expected that an error will be thrown by eachMessage.
          * This is especially true since the pattern of pause() followed by throwing an error
@@ -645,7 +644,7 @@ class Consumer {
       /* Force a immediate seek here. It's possible that there are no more messages to be passed to the user,
        * but the user seeked in the call to eachMessage, or else we encountered the error catch block.
        * In that case, the results of that seek will never be reflected unless we do this. */
-      if (this.#checkPendingSeeks && this.#state === ConsumerState.CONNECTED)
+      if (this.#checkPendingSeeks)
         await this.#seekInternal();
 
       /* TODO: another check we need to do here is to see how kafkaJS is handling
@@ -654,13 +653,8 @@ class Consumer {
        * inside this function.
        */
 
-      /* Yield for a bit to allow other scheduled tasks on the event loop to run.
-       * For instance, if disconnect() is called during/after we await eachMessage, and
-       * we don't await anything else after that, this loop will run despite needing to
-       * disconnect.
-       * It's better than any pending tasks be processed here, while we've processed one
-       * message completely, rather than between message processing. */
-      await new Promise((resolve) => setTimeout(resolve, 0));
+      /* Release the lock so that any pending disconnect can go through. */
+      await this.#lock.release();
     }
   }
 
@@ -952,7 +946,13 @@ class Consumer {
       return;
     }
 
+    while (!(await acquireOrLog(this.#lock, this.#logger))); /* Just retry... */
+
     this.#state = ConsumerState.DISCONNECTING;
+
+    /* Since there are state-checks before everything, we are safe to proceed without the lock. */
+    await this.#lock.release();
+
     await new Promise((resolve, reject) => {
       const cb = (err) => {
         if (err) {

lib/kafkajs/_error.js

@@ -130,6 +130,13 @@ class KafkaJSTimeout extends KafkaJSError {
     }
 }
 
+class KafkaJSLockTimeout extends KafkaJSTimeout {
+    constructor() {
+        super(...arguments)
+        this.name = 'KafkaJSLockTimeout'
+    }
+}
+
 /**
  * @typedef {Object} KafkaJSAggregateError represents an error raised when multiple errors occur at once.
  */
@@ -179,6 +186,7 @@ module.exports = {
     KafkaJSGroupCoordinatorNotFound,
     KafkaJSNotImplemented,
     KafkaJSTimeout,
+    KafkaJSLockTimeout,
     KafkaJSAggregateError,
     KafkaJSNoBrokerAvailableError,
     isRebalancing,

test/promisified/consumer/consumeMessages.spec.js

@@ -2,6 +2,7 @@ jest.setTimeout(30000)
 
 const { ErrorCodes, CompressionTypes } = require('../../../lib').KafkaJS;
 
+const { doesNotMatch } = require('assert');
 const {
     secureRandom,
     createTopic,
@@ -443,6 +444,43 @@ describe('Consumer', () => {
         expect(offsetsConsumed.length).toEqual(messages.length)
     });
 
+    it('does not disconnect in the middle of message processing', async () => {
+        await producer.connect();
+        await consumer.connect();
+        await consumer.subscribe({ topic: topicName });
+
+        let calls = 0;
+        let failedSeek = false;
+        consumer.run({
+            eachMessage: async ({ message }) => {
+                /* Take a long time to process the message. */
+                await sleep(7000);
+                try {
+                    consumer.seek({ topic: topicName, partition: 0, offset: message.offset });
+                } catch (e) {
+                    failedSeek = true;
+                }
+                calls++;
+            }
+        });
+
+        await producer.send({
+            topic: topicName,
+            messages: [{ key: '1', value: '1' }],
+        });
+
+        /* Waiting for assignment and then a bit more means that the first eachMessage starts running. */
+        await waitFor(() => consumer.assignment().length > 0, () => { }, { delay: 50 });
+        await sleep(200);
+        await consumer.disconnect();
+
+        /* Even without explicitly waiting for it, a pending call to eachMessage must complete before disconnect does. */
+        expect(calls).toEqual(1);
+        expect(failedSeek).toEqual(false);
+
+        await producer.disconnect();
+    });
+
     describe('transactions', () => {
         it('accepts messages from an idempotent producer', async () => {
             producer = createProducer({

test/promisified/testhelpers.js

@@ -1,6 +1,5 @@
 const crypto = require('crypto');
 const process = require('process');
-const { logLevel } = require('../../lib/kafkajs');
 const { Kafka } = require('../../lib').KafkaJS;
 
 // TODO: pick this up from a file