Add per-partition cache with global expiry

Author: milindl

lib/kafkajs/_common.js

@@ -689,6 +689,15 @@ async function acquireOrLog(lock, logger) {
   return false;
 }
 
+/**
+ * Creates a key for maps from a topicPartition object.
+ * @param {{topic: string, partition: number}} topicPartition Any object which can be treated as a topic partition.
+ * @returns {string} The created key.
+ */
+function partitionKey(topicPartition) {
+  return topicPartition.topic + '|'+ (topicPartition.partition);
+}
+
 module.exports = {
   kafkaJSToRdKafkaConfig,
   topicPartitionOffsetToRdKafka,
@@ -707,4 +716,5 @@ module.exports = {
   checkIfKafkaJsKeysPresent,
   Lock,
   acquireOrLog,
+  partitionKey,
 };

lib/kafkajs/_consumer.js

@@ -15,7 +15,8 @@ const {
   checkAllowedKeys,
   logLevel,
   Lock,
-  acquireOrLog
+  acquireOrLog,
+  partitionKey,
 } = require('./_common');
 const { Buffer } = require('buffer');
 const { hrtime } = require('process');
@@ -36,47 +37,129 @@ const PartitionAssigners = Object.freeze({
 
 
 /**
- * MessageCache represents a cache of messages that have been consumed,
- * but not yet passed to the user.
- * It has a dynamic capacity, increased or decreased based on requirement.
+ * A PerPartitionMessageCache is a cache for messages for a single partition.
  */
-class MessageCache {
+class PerPartitionMessageCache {
   /* The cache is a list of messages. */
   cache = [];
-  /* The maximum size of the cache. Set to 1 initially. */
-  maxSize = 1;
   /* Index of next element to be fetched in the cache. */
-  currentIndex = this.maxSize;
+  currentIndex = 0;
   /* Whether the cache is stale. */
   stale = false;
-  /* Number of times the cache has been requested to be increased in size. */
-  increaseCount = 0;
-  /* Last cached time */
-  cachedTime = hrtime();
-  /* Expiry duration for this cache */
-  expiryDurationMs = 500;
 
-  constructor(expiryDurationMs) {
-    this.expiryDurationMs = expiryDurationMs;
+  /**
+   * Returns the number of total elements in the cache.
+   */
+  size() {
+    return this.cache.length;
   }
 
   /**
    * Clears the cache.
    */
   clear() {
     this.cache = [];
-    this.maxSize = 1;
-    this.currentIndex = this.maxSize;
+    this.currentIndex = 0;
     this.stale = false;
+  }
+
+  /**
+   * Adds a message to the cache.
+   */
+  add(message) {
+    this.cache.push(message);
+  }
+
+  /**
+   * Returns whether the cache is stale.
+   */
+  isStale() {
+    return this.stale;
+  }
+
+  /**
+   * @returns The next element in the cache or null if none exists.
+   * @warning Does not check for staleness.
+   */
+  next() {
+    return this.currentIndex < this.cache.length ? this.cache[this.currentIndex++] : null;
+  }
+}
+
+
+/**
+ * MessageCache defines a dynamically sized cache for messages.
+ * Internally, it uses PerPartitionMessageCache to store messages for each partition.
+ * The capacity is increased or decreased according to whether the last fetch of messages
+ * was less than the current capacity or saturated the current capacity.
+ */
+class MessageCache {
+
+  constructor(expiryDurationMs) {
+    /* 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;
+    /* 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. */
     this.increaseCount = 0;
+    /* Last cached time */
     this.cachedTime = hrtime();
+    /* Whether the cache is stale. */
+    this.stale = false;
+    /* Expiry duration for this cache */
+    this.expiryDurationMs = expiryDurationMs;
+  }
+
+  addTopicPartitions(topicPartitions) {
+    if (this.ppcList.length !== 0) {
+      throw new Error('Cannot add topic partitions to a non-empty cache.');
+    }
+    for (const topicPartition of topicPartitions) {
+      const key = partitionKey(topicPartition);
+      this.tpToPpc.set(key, new PerPartitionMessageCache());
+    }
+  }
+
+  removeTopicPartitions(topicPartitions = null) {
+    if (this.ppcList.length !== 0) {
+      throw new Error('Cannot remove topic partitions from a non-empty cache.');
+    }
+
+    if (topicPartitions === null) {
+      this.tpToPpc.clear();
+      return;
+    }
+    for (const topicPartition of assignment) {
+      const key = partitionKey(topicPartition);
+      this.tpToPpc.delete(key);
+    }
+  }
+
+  /**
+   * Returns whether the cache is stale.
+   */
+  isStale() {
+    if (this.stale)
+      return true;
+
+    const cacheTime = hrtime(this.cachedTime);
+    const cacheTimeMs = Math.floor(cacheTime[0] * 1000 + cacheTime[1] / 1000000);
+    this.stale = cacheTimeMs > this.expiryDurationMs;
+
+    // TODO: ideally, local staleness should not lead to global staleness.
+    // But for now, make it so because seeking to stored offset on local staleness is tricky.
+    this.stale = this.stale || this.ppcList.some(cache => cache.isStale());
+    return this.stale;
   }
 
   /**
    * Request a size increase.
    * It increases the size by 2x, but only if the size is less than 1024,
    * only if the size has been requested to be increased twice in a row.
-   * @returns
    */
   increaseMaxSize() {
     if (this.maxSize === 1024)
@@ -101,33 +184,67 @@ class MessageCache {
   }
 
   /**
-   * Sets cache and resets all the indices and timer.
-   * @param {*} messages
+   * Add a single message to the cache.
    */
-  setCache(messages) {
-    this.cache = messages;
-    this.currentIndex = 1;
-    this.cachedTime = hrtime();
+  #add(message) {
+    const key = partitionKey(message)
+    const cache = this.tpToPpc.get(key);
+    cache.add(message);
+    if (cache.size() === 1) {
+      this.ppcList.push(cache);
+    }
   }
 
   /**
-   * @returns The next element in the cache or null if none exists.
-   * @warning Does not check for staleness.
+   * Adds many messages into the cache, partitioning them as per their toppar.
    */
-  next() {
-    return this.currentIndex < this.cache.length ? this.cache[this.currentIndex++] : null;
+  addMessages(messages) {
+    this.stale = false;
+    this.cachedTime = hrtime();
+    this.currentPpc = 0;
+    for (const message of messages)
+      this.#add(message);
+
+    // TODO: add ppcList sort step.
+    // Rationale: ideally it's best to consume in the ascending order of timestamps.
   }
 
-  /* Whether the cache is stale. */
-  isStale() {
-    if (this.stale)
-      return true;
+  /**
+   * 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.
+   * @warning Does not check for staleness. That is left up to the user.
+   */
+  next() {
+    if (this.currentPpc >= this.ppcList.length) {
+      return null;
+    }
 
-    const cacheTime = hrtime(this.cachedTime);
-    const cacheTimeMs = Math.floor(cacheTime[0] * 1000 + cacheTime[1] / 1000000);
-    return cacheTimeMs > this.expiryDurationMs;
+    let next = null;
+    while (next === null && this.currentPpc < this.ppcList.length) {
+      next = this.ppcList[this.currentPpc].next();
+      if (next !== null)
+        break;
+      this.currentPpc++;
+    }
+    return next; // Caller is responsible for triggering fetch logic here if next == null.
   }
 
+  /**
+   * Clears cache completely.
+   */
+  clear() {
+    for (const cache of this.ppcList) {
+      cache.clear();
+    }
+    this.ppcList = [];
+    this.currentPpc = 0;
+    this.maxSize = 1;
+    this.increaseCount = 0;
+    this.stale = false;
+    this.cachedTime = hrtime();
+  }
 }
 
 class Consumer {
@@ -254,7 +371,7 @@ class Consumer {
       const assignment = this.assignment();
       const seekPromises = [];
       for (const topicPartitionOffset of assignment) {
-        const key = `${topicPartitionOffset.topic}|${topicPartitionOffset.partition}`;
+        const key = partitionKey(topicPartitionOffset);
         if (!this.#lastConsumedOffsets.has(key))
           continue;
 
@@ -376,11 +493,19 @@ class Consumer {
               }
             }
 
+            // Populate per-partion caches.
+            // For cooperative sticky, just add the newly recieved partitions.
+            // If it's eager, it's already empty, so we can add all the partitions.
+            this.#messageCache.addTopicPartitions(assignment);
+
           } else {
-            if (this.#internalClient.rebalanceProtocol() === "EAGER")
+            if (this.#internalClient.rebalanceProtocol() === "EAGER") {
               this.#internalClient.unassign();
-            else
+              this.#messageCache.removeTopicPartitions();
+            } else {
               this.#internalClient.incrementalUnassign(assignment);
+              this.#messageCache.removeTopicPartitions(assignment);
+            }
           }
         } catch (e) {
           // Ignore exceptions if we are not connected
@@ -724,8 +849,8 @@ class Consumer {
           reject(createKafkaJsErrorFromLibRdKafkaError(err));
           return;
         }
-        this.#messageCache.setCache(messages);
-        const message = messages[0];
+        this.#messageCache.addMessages(messages);
+        const message = this.#messageCache.next();
         if (messages.length === this.#messageCache.maxSize) {
           this.#messageCache.increaseMaxSize();
         } else {
@@ -930,7 +1055,7 @@ class Consumer {
         /* 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)}`);
+          this.#logger.error(`Consumer encountered error while consuming. Retrying. Error details: ${e} : ${e.stack}`);
       });
 
       if (!m) {
@@ -973,7 +1098,7 @@ class Consumer {
          *
          * So - do nothing but a debug log, but at this point eachMessageProcessed is false.
          */
-        this.#logger.debug(`Consumer encountered error while processing message. Error details: ${JSON.stringify(e)}. The same message may be reprocessed.`);
+        this.#logger.debug(`Consumer encountered error while processing message. Error details: ${e}: ${e.stack}. The same message may be reprocessed.`);
       }
 
       /* If the message is unprocessed, due to an error, or because the user has not resolved it, we seek back. */
@@ -993,7 +1118,7 @@ class Consumer {
               topic: m.topic, partition: m.partition, offset: Number(m.offset) + 1, leaderEpoch: m.leaderEpoch
             }]);
           }
-          this.#lastConsumedOffsets.set(`${m.topic}|${m.partition}`, Number(m.offset) + 1);
+          this.#lastConsumedOffsets.set(partitionKey(m), Number(m.offset) + 1);
         } catch (e) {
           /* Not much we can do, except log the error. */
           if (this.#logger)
@@ -1089,7 +1214,7 @@ class Consumer {
          * So - do nothing but a debug log, but at this point eachMessageProcessed needs to be false unless
          * the user has explicitly marked it as true.
          */
-        this.#logger.debug(`Consumer encountered error while processing message. Error details: ${JSON.stringify(e)}. The same message may be reprocessed.`);
+        this.#logger.debug(`Consumer encountered error while processing message. Error details: ${e}: ${e.stack}. The same message may be reprocessed.`);
 
         /* The value of eachBatchAutoResolve is not important. The only place where a message is marked processed
          * despite an error is if the user says so, and the user can use resolveOffsets for both the possible
@@ -1115,7 +1240,7 @@ class Consumer {
               topic: m.topic, partition: m.partition, offset: Number(m.offset) + 1, leaderEpoch: m.leaderEpoch
             }]);
           }
-          this.#lastConsumedOffsets.set(`${m.topic}|${m.partition}`, Number(m.offset) + 1);
+          this.#lastConsumedOffsets.set(partitionKey(m), Number(m.offset) + 1);
         } catch (e) {
           /* Not much we can do, except log the error. */
           if (this.#logger)
@@ -1273,7 +1398,7 @@ class Consumer {
 
     for (let i = 0; i < assignment.length; i++) {
       const topicPartition = assignment[i];
-      const key = `${topicPartition.topic}|${topicPartition.partition}`;
+      const key = partitionKey(topicPartition);
       if (!this.#pendingSeeks.has(key))
         continue;
 
@@ -1303,7 +1428,7 @@ class Consumer {
     let invalidateMessage = false;
 
     for (const topicPartition of assignment) {
-      const key = `${topicPartition.topic}|${topicPartition.partition}`;
+      const key = partitionKey(topicPartition);
       if (!this.#pendingSeeks.has(key))
         continue;
 
@@ -1372,7 +1497,7 @@ class Consumer {
     }
 
     this.#checkPendingSeeks = true;
-    this.#pendingSeeks.set(`${rdKafkaTopicPartitionOffset.topic}|${rdKafkaTopicPartitionOffset.partition}`, rdKafkaTopicPartitionOffset.offset);
+    this.#pendingSeeks.set(partitionKey(rdKafkaTopicPartitionOffset), rdKafkaTopicPartitionOffset.offset);
   }
 
   async describeGroup() {
@@ -1441,6 +1566,8 @@ class Consumer {
       return;
     }
     this.#internalClient.pause(topics);
+
+    // TODO: make this staleness per-partition, not on a global cache level.
     this.#messageCache.stale = true;
 
     topics.map(JSON.stringify).forEach(topicPartition => this.#pausedPartitions.add(topicPartition));