Allow cache to disburse multiple messages at once

milindl · milindl · commit 28d2253cdaac · 2024-06-13T11:29:57.000+05:30
diff --git a/LICENSE.heap-js b/LICENSE.heap-js
@@ -0,0 +1,36 @@
+Code from https://github.com/ignlg/heap-js/ | Commit Hash: a3bb403 | dist/heap-js.es5.js
+
+Code used in lib/kafkajs/_heap.js
+
+----
+
+
+BSD 3-Clause License
+
+Copyright (c) 2017, Ignacio Lago
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
diff --git a/lib/kafkajs/_consumer.js b/lib/kafkajs/_consumer.js
@@ -353,7 +353,7 @@ class Consumer {
 
     if (Object.hasOwn(kjsConfig, 'rebalanceTimeout')) {
       /* In librdkafka, we use the max poll interval as the rebalance timeout as well. */
-      rdKafkaConfig['max.poll.interval.ms'] = kjsConfig.rebalanceTimeout;
+      rdKafkaConfig['max.poll.interval.ms'] = +kjsConfig.rebalanceTimeout;
     } else {
       rdKafkaConfig['max.poll.interval.ms'] = 300000; /* librdkafka default */
     }
@@ -627,15 +627,24 @@ class Consumer {
 
   /**
    * Consumes a single message from the internal consumer.
+   * @param {number} savedIndex - the index of the message in the cache to return.
    * @returns {Promise<import("../..").Message>} a promise that resolves to a single message.
    * @note this method caches messages as well, but returns only a single message.
    */
-  async #consumeSingleCached() {
-    const msg = this.#messageCache.next();
+  async #consumeSingleCached(savedIndex) {
+    const msg = this.#messageCache.next(savedIndex);
     if (msg) {
       return msg;
     }
 
+    // TODO: Add this block for concurrency
+    // if (!msg) {
+    //   // it's possible that we get msg = null, but that's because partitionConcurrency
+    //   // exceeds the number of partitions containing messages. So in this case,
+    //   // we should not call for new fetches, rather, try to focus on what we have left.
+    //   return null;
+    // }
+
     return new Promise((resolve, reject) => {
       this.#internalClient.consume(this.#messageCache.maxSize, (err, messages) => {
         if (err) {
@@ -712,7 +721,8 @@ class Consumer {
     }
 
     const rdKafkaConfig = this.#config();
-    this.#messageCache = new MessageCache(Math.floor(rdKafkaConfig['max.poll.interval.ms'] * 0.8));
+    const maxPollInterval = rdKafkaConfig['max.poll.interval.ms'] ?? 300000;
+    this.#messageCache = new MessageCache(Math.floor(maxPollInterval * 0.8), 1);
 
     this.#state = ConsumerState.CONNECTING;
     this.#internalClient = new RdKafka.KafkaConsumer(rdKafkaConfig);
@@ -830,6 +840,7 @@ class Consumer {
   /* Internal polling loop.
    * It accepts the same config object that `run` accepts, but config.eachMessage must be set. */
   async #runInternalEachMessage(config) {
+    let savedIdx = -1;
     while (this.#state === ConsumerState.CONNECTED) {
 
       /* We need to acquire a lock here, because we need to ensure that we don't
@@ -840,26 +851,34 @@ class Consumer {
       /* Invalidate the message cache if needed */
       const locallyStale = this.#messageCache.popLocallyStale();
       if (this.#messageCache.isStale()) { /* global staleness */
+        // TODO: await all concurrent promises for eachMessage here.
         await this.#clearCacheAndResetPositions();
         await this.#lock.release();
         continue;
       } else if (locallyStale.length !== 0) { /* local staleness */
+        // TODO: is it correct to await some concurrent promises for eachMessage here?
+        // to be safe we can do it, but I don't think we really need to do that for
+        // correctness.
         await this.#clearCacheAndResetPositions(locallyStale);
         await this.#lock.release();
         continue;
       }
 
-      const m = await this.#consumeSingleCached().catch(e => {
+      const m = await this.#consumeSingleCached(savedIdx).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: ${e} : ${e.stack}`);
       });
 
       if (!m) {
+        // await all concurrency related promises right here if this is null, if any such promise exists.
+        // see note in consumeSingleCached
+        savedIdx = -1;
         await this.#lock.release();
         continue;
       }
+      savedIdx = m.index;
 
       /* TODO: add partitionsConsumedConcurrently-based concurrency here.
       * If we maintain a map of topic partitions to promises, and a counter,
@@ -944,6 +963,7 @@ class Consumer {
   /* Internal polling loop.
    * It accepts the same config object that `run` accepts, but config.eachBatch must be set. */
   async #runInternalEachBatch(config) {
+    let savedIdx = -1;
     while (this.#state === ConsumerState.CONNECTED) {
 
       /* We need to acquire a lock here, because we need to ensure that we don't
@@ -963,17 +983,19 @@ class Consumer {
         continue;
       }
 
-      const m = await this.#consumeSingleCached().catch(e => {
+      const m = await this.#consumeSingleCached(savedIdx).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)}`);
       });
 
       if (!m) {
+        savedIdx = -1;
         await this.#lock.release();
         continue;
       }
+      savedIdx = m.index;
 
       /* TODO: add partitionsConsumedConcurrently-based concurrency here.
       * If we maintain a map of topic partitions to promises, and a counter,
diff --git a/lib/kafkajs/_consumer_cache.js b/lib/kafkajs/_consumer_cache.js
@@ -2,6 +2,7 @@ const { hrtime } = require('process');
 const {
     partitionKey,
 } = require('./_common');
+const { Heap } = require('./_heap');
 
 /**
  * A PerPartitionMessageCache is a cache for messages for a single partition.
@@ -62,13 +63,13 @@ class PerPartitionMessageCache {
  */
 class MessageCache {
 
-    constructor(expiryDurationMs) {
+    constructor(expiryDurationMs, maxConcurrency) {
         /* Per partition cache list containing non-empty PPCs */
         this.ppcList = [];
         /* Map of topic+partition to PerPartitionMessageCache. */
         this.tpToPpc = new Map();
         /* Index of the current PPC in the ppcList. */
-        this.currentPpc = 0;
+        this.currentPpcTODO_remove_this = 0;
         /* Maximum size of the cache. (Capacity) */
         this.maxSize = 1;
         /* Number of times the size has been increased in a row, used for accounting for maxSize. */
@@ -81,6 +82,15 @@ class MessageCache {
         this.expiryDurationMs = expiryDurationMs;
         /* A list of caches which have been marked stale since the last call to popLocallyStale or addMessages. */
         this.locallyStaleCaches = [];
+        /* Max allowed concurrency */
+        this.maxConcurrency = maxConcurrency;
+        /* Contains a list of indices of ppcList from which we are allowed to consume. */
+        this.indices = new Heap();
+        /* Largest ppc index we are allowed to consume from (inclusive). */
+        this.maxIndicesIndex = 0;
+        /* Contains a list of indices of ppcList from which we have sent a message returned through next, but
+         * the user has not returned the index back to us via next(idx) */
+        this.pendingIndices = new Set();
     }
 
     /**
@@ -226,7 +236,6 @@ class MessageCache {
          * run out of messages. We need to clear them, else #add() will not add
          * them back to the ppcList since they're not empty. */
         this.ppcList.forEach(cache => cache.clear());
-        this.currentPpc = 0;
         this.ppcList = [];
 
         if (this.locallyStaleCaches.length !== 0 && this.locallyStaleCaches.some(tp => {
@@ -245,34 +254,72 @@ class MessageCache {
 
         // TODO: add ppcList sort step.
         // Rationale: ideally it's best to consume in the ascending order of timestamps.
+
+        /* Reset the indices and pendingIndices because ppcList is being created newly. */
+        this.indices.clear();
+        if (this.pendingIndices.size > 0) console.error('addMessages: pendingIndices = ', this.pendingIndices, console.trace());
+        this.pendingIndices.clear();
+        this.maxIndicesIndex = Math.min(this.maxConcurrency, this.ppcList.length - 1);
+        for (let i = 0; i <= this.maxIndicesIndex; i++) {
+            this.indices.push(i);
+        }
     }
 
     /**
      * Returns the next element in the cache, or null if none exists.
      *
      * If the current PPC is exhausted, it moves to the next PPC.
      * If all PPCs are exhausted, it returns null.
+     * @param {number} idx - after a consumer has consumed a message, it must return the index back to us via this parameter.
+     *                       otherwise, no messages from that topic partition will be consumed.
+     * @returns {Object} - the next message in the cache, or null if none exists. An `index` field is added to the message.
      * @warning Does not check for global staleness. That is left up to the user.
      *          Skips locally stale messages.
+     * The topicPartition, if provided, MUST be one such that the user has fetched
+     * the message from the same topicPartition earlier.
+     * @note Whenever making changes to this function, ensure that you benchmark perf.
      */
-    next() {
-        if (this.currentPpc >= this.ppcList.length) {
-            return null;
+    next(idx = -1) {
+        let index = idx;
+        if (!this.pendingIndices.has(index)) {
+            /* The user is behaving well by returning the index to us, but in the meanwhile, it's possible
+             * that we ran out of messages and fetched a new batch. So we just discard what the user is
+             * returning to us. */
+            index = -1;
+        } else {
+            this.pendingIndices.delete(index);
         }
 
-        let next = null;
-        while (next === null && this.currentPpc < this.ppcList.length) {
-            if (this.ppcList[this.currentPpc].isStale()) {
-                this.currentPpc++;
+        if (index === -1) {
+            if (this.indices.size() === 0)
+                return null;
+            index = this.indices.pop(); // index cannot be undefined here since indices.size > 0
+        }
+
+        while (true) {
+            const next = this.ppcList[index].next();
+            if (this.ppcList[index].isStale() || next === null) {
+                /* If the current PPC is stale or empty, then we move on to the next one.
+                 * It is equally valid to choose any PPC available within this.indices, or else
+                 * move on to the next PPC (maxIndicesIndex + 1) if available.
+                 * We prefer the second option a bit more since we don't have to do a heap operation. */
+                const toAdd = this.maxIndicesIndex + 1;
+                if (toAdd < this.ppcList.length) {
+                    this.maxIndicesIndex = toAdd;
+                    index = toAdd;
+                } else if (!this.indices.isEmpty()) {
+                    index = this.indices.pop()
+                } else {
+                    break; // nothing left.
+                }
                 continue;
             }
 
-            next = this.ppcList[this.currentPpc].next();
-            if (next !== null)
-                break;
-            this.currentPpc++;
+            this.pendingIndices.add(index);
+            next.index = index;
+            return next;
         }
-        return next; // Caller is responsible for triggering fetch logic here if next == null.
+        return null; // Caller is responsible for triggering fetch logic here if next == null.
     }
 
     /**
@@ -286,12 +333,15 @@ class MessageCache {
             cache.clear();
         }
         this.ppcList = [];
-        this.currentPpc = 0;
         this.maxSize = 1;
         this.increaseCount = 0;
         this.stale = false;
         this.cachedTime = hrtime();
         this.locallyStaleCaches = [];
+        this.indices.clear();
+        // if (this.pendingIndices.size > 0) console.log('clear: pendingIndices = ', this.pendingIndices, console.);
+        this.pendingIndices.clear();
+        this.currentIndex = 0;
     }
 }
 
diff --git a/lib/kafkajs/_heap.js b/lib/kafkajs/_heap.js
