Add tests for consumer.pause() and consumer.resume()

Author: milindl

MIGRATION.md

@@ -289,11 +289,13 @@
     * The `partitionsConsumedConcurrently` property is not supported (YET).
   * The `eachBatch` method is not supported.
   * `commitOffsets` does not (YET) support sending metadata for topic partitions being committed.
-  * `paused()` is not (YET) supported.
+  * `paused()` is supported without any changes.
   * Custom partition assignors are not supported.
   * Changes to `seek`:
-    * The restriction to call seek only after `run` is removed.
+    * The restriction to call seek only after `run` is removed. It can be called any time.
     * Rather than the `autoCommit` property of `run` deciding if the offset is committed, the librdkafka property `enable.auto.commit` of the consumer config is used.
+  * `pause` and `resume` MUST be called after the consumer group is joined. In practice, this means it can be called whenever `consumer.assignment()` has a non-zero size, or within the `eachMessage`
+    callback.
 
 ### Admin Client
 

lib/kafkajs/_consumer.js

@@ -67,6 +67,12 @@ class Consumer {
    */
   #pendingSeeks = new Map();
 
+  /**
+   * Stores the list of paused partitions, as a set of JSON.stringify'd TopicPartition objects.
+   * @type {Set<string>}
+   */
+  #pausedPartitions = new Set();
+
   /**
    * @constructor
    * @param {import("../../types/kafkajs").ConsumerConfig} kJSConfig
@@ -300,7 +306,7 @@ class Consumer {
         headers
       },
       heartbeat: async () => { /* no op */ },
-      pause: () => this.pause([{ topic: message.topic, partitions: [message.partition] }]),
+      pause: this.pause.bind(this, [{ topic: message.topic, partitions: [message.partition] }]),
     };
   }
 
@@ -443,33 +449,56 @@ class Consumer {
 
     while (this.#state === ConsumerState.CONNECTED) {
       const m = await this.#consumeSingle();
-      if (m) {
-        /* TODO: add partitionsConsumedConcurrently-based concurrency here.
-         * If we maintain a map of topic partitions to promises, and a counter,
-         * we can probably achieve it with the correct guarantees of ordering
-         * though to maximize performance, we need to consume only from partitions for which
-         * an eachMessage call is not already going.
-         * It's risky to consume, and then store the message in something like an
-         * array/list until it can be processed, because librdkafka marks it as
-         * 'stored'... but anyway - we can implement something like this.
-         */
-        if (this.#checkPendingSeeks) {
-          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. */
-            continue;
-          }
+      if (!m) {
+        continue;
+      }
+
+      /* TODO: add partitionsConsumedConcurrently-based concurrency here.
+        * If we maintain a map of topic partitions to promises, and a counter,
+        * we can probably achieve it with the correct guarantees of ordering
+        * though to maximize performance, we need to consume only from partitions for which
+        * an eachMessage call is not already going.
+        * It's risky to consume, and then store the message in something like an
+        * array/list until it can be processed, because librdkafka marks it as
+        * 'stored'... but anyway - we can implement something like this.
+        */
+
+      /* Make pending seeks 'concrete'. */
+      if (this.#checkPendingSeeks) {
+        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. */
+          continue;
         }
+      }
 
+      try {
         await config.eachMessage(
           this.#createPayload(m)
         )
-        /* TODO: another check we need to do here is to see how kafkaJS is handling
-         * commits. Are they commmitting after a message is _processed_?
-         * In that case we need to turn off librdkafka's auto-commit, and commit
-         * inside this function.
-         */
+      } 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
+         * is encouraged. To meet the API contract, we seek one offset backward at this point (which
+         * means seeking to the message offset). */
+        this.seek({
+          topic: m.topic,
+          partition: m.partition,
+          offset: m.offset,
+        });
       }
+
+      /* 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)
+        await this.#seekInternal();
+
+      /* TODO: another check we need to do here is to see how kafkaJS is handling
+       * commits. Are they commmitting after a message is _processed_?
+       * In that case we need to turn off librdkafka's auto-commit, and commit
+       * inside this function.
+       */
     }
   }
 
@@ -646,13 +675,18 @@ class Consumer {
    * all partitions for the given topic. If topic partition(s) are already paused
    * this method has no effect.
    * @param {{topic: string, partitions?: number[]}[]} topics
+   * @returns {Function} a function that can be called to resume the given topic partitions.
    */
   pause(topics) {
     if (this.#state !== ConsumerState.CONNECTED) {
       throw new error.KafkaJSError('Pause can only be called while connected.', { code: error.ErrorCodes.ERR__STATE });
     }
 
     for (let topic of topics) {
+      if (typeof topic.topic !== 'string') {
+        throw new error.KafkaJSError('Topic must be a string.', { code: error.ErrorCodes.ERR__INVALID_ARG });
+      }
+
       if (!topic.partitions) {
         topic.partitions = this.#getAllAssignedPartition(topic.topic);
       }
@@ -662,12 +696,31 @@ class Consumer {
     if (topics.length === 0) {
       return;
     }
-
     this.#internalClient.pause(topics);
+
+    topics.map(JSON.stringify).forEach(topicPartition => this.#pausedPartitions.add(topicPartition));
+
+    return () => this.resume(topics);
   }
 
+  /**
+   * Returns the list of paused topic partitions.
+   * @returns {{topic: string, partitions: number[]}[]} a list of paused topic partitions.
+   */
   paused() {
-    notImplemented();
+    const topicToPartitions = Array
+      .from(this.#pausedPartitions.values())
+      .map(JSON.parse)
+      .reduce(
+        (acc, { topic, partition }) => {
+          if (!acc[topic]) {
+            acc[topic] = [];
+          }
+          acc[topic].push(partition);
+          return acc;
+        },
+        {});
+    return Array.from(Object.entries(topicToPartitions), ([topic, partitions]) => ({ topic, partitions }));
   }
 
 
@@ -683,13 +736,22 @@ class Consumer {
     }
 
     for (let topic of topics) {
+      if (typeof topic.topic !== 'string') {
+        throw new error.KafkaJSError('Topic must be a string.', { code: error.ErrorCodes.ERR__INVALID_ARG });
+      }
+
       if (!topic.partitions) {
         topic.partitions = this.#getAllAssignedPartition(topic.topic);
       }
     }
 
     topics = this.#flattenTopicPartitions(topics);
+    if (topics.length === 0) {
+      return;
+    }
     this.#internalClient.resume(topics);
+
+    topics.map(JSON.stringify).forEach(topicPartition => this.#pausedPartitions.delete(topicPartition));
   }
 
   on(/* eventName, listener */) {