Add per-partition cache expiry logic

Author: milindl

lib/kafkajs/_consumer.js

@@ -19,7 +19,7 @@ const {
   partitionKey,
 } = require('./_common');
 const { Buffer } = require('buffer');
-const { hrtime } = require('process');
+const MessageCache = require('./_consumer_cache');
 
 const ConsumerState = Object.freeze({
   INIT: 0,
@@ -35,218 +35,6 @@ const PartitionAssigners = Object.freeze({
   cooperativeSticky: 'cooperative-sticky',
 });
 
-
-/**
- * A PerPartitionMessageCache is a cache for messages for a single partition.
- */
-class PerPartitionMessageCache {
-  /* The cache is a list of messages. */
-  cache = [];
-  /* Index of next element to be fetched in the cache. */
-  currentIndex = 0;
-  /* Whether the cache is stale. */
-  stale = false;
-
-  /**
-   * Returns the number of total elements in the cache.
-   */
-  size() {
-    return this.cache.length;
-  }
-
-  /**
-   * Clears the cache.
-   */
-  clear() {
-    this.cache = [];
-    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.
-   */
-  increaseMaxSize() {
-    if (this.maxSize === 1024)
-      return;
-
-    this.increaseCount++;
-    if (this.increaseCount <= 1)
-      return;
-
-    this.maxSize = Math.min(this.maxSize << 1, 1024);
-    this.increaseCount = 0;
-  }
-
-  /**
-   * Request a size decrease.
-   * It decreases the size to 80% of the last received size, with a minimum of 1.
-   * @param {number} recvdSize - the number of messages received in the last poll.
-   */
-  decreaseMaxSize(recvdSize) {
-    this.maxSize = Math.max(Math.floor((recvdSize * 8) / 10), 1);
-    this.increaseCount = 0;
-  }
-
-  /**
-   * Add a single message to the cache.
-   */
-  #add(message) {
-    const key = partitionKey(message)
-    const cache = this.tpToPpc.get(key);
-    cache.add(message);
-    if (cache.size() === 1) {
-      this.ppcList.push(cache);
-    }
-  }
-
-  /**
-   * Adds many messages into the cache, partitioning them as per their toppar.
-   */
-  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.
-  }
-
-  /**
-   * 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;
-    }
-
-    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 {
   /**
    * The config supplied by the user.
@@ -312,7 +100,6 @@ class Consumer {
   /**
    * A map of topic+partition to the offset that was last consumed.
    * The keys are of the type "<topic>|<partition>".
-   * This is only populated when we're in the kafkaJS compatibility mode.
    * @type {Map<string, number>}
    */
   #lastConsumedOffsets = new Map();
@@ -358,25 +145,25 @@ class Consumer {
   }
 
   /**
-   * Clear the message cache.
-   * For simplicity, this always clears the entire message cache rather than being selective.
+   * Clear the message cache, and reset to stored positions.
    *
-   * @param {boolean} seek - whether to seek to the stored offsets after clearing the cache.
-   *                         this should be set to true if partitions are retained after this operation.
+   * @param {Array<{topic: string, partition: number}>|null} topicPartitions to clear the cache for, if null, then clear all assigned.
    */
-  async #clearCacheAndResetPositions(seek = true) {
-    /* Seek to stored offset for each topic partition so that if
-     * we've gotten further along then they have, we can come back. */
-    if (seek) {
-      const assignment = this.assignment();
-      const seekPromises = [];
-      for (const topicPartitionOffset of assignment) {
-        const key = partitionKey(topicPartitionOffset);
-        if (!this.#lastConsumedOffsets.has(key))
-          continue;
+  async #clearCacheAndResetPositions(topicPartitions = null) {
+    /* Seek to stored offset for each topic partition. It's possible that we've
+     * consumed messages upto N from the internalClient, but the user has stale'd the cache
+     * after consuming just k (< N) messages. We seek to k+1. */
+
+    const clearPartitions = topicPartitions ? topicPartitions : this.assignment();
+    const seekPromises = [];
+    for (const topicPartitionOffset of clearPartitions) {
+      const key = partitionKey(topicPartitionOffset);
+      if (!this.#lastConsumedOffsets.has(key))
+        continue;
 
-        /* Fire off a seek */
-        const seekPromise = new Promise((resolve, reject) => this.#internalClient.seek({
+      /* Fire off a seek */
+      const seekPromise = new Promise((resolve, reject) => {
+        this.#internalClient.seek({
           topic: topicPartitionOffset.topic,
           partition: topicPartitionOffset.partition,
           offset: +this.#lastConsumedOffsets.get(key)
@@ -386,18 +173,24 @@ class Consumer {
           } else {
             resolve();
           }
-        }));
-        seekPromises.push(seekPromise);
-      }
+        });
 
-      /* TODO: we should cry more about this and render the consumer unusable. */
-      await Promise.all(seekPromises).catch(err => this.#logger.error("Seek error. This is effectively a fatal error:" + err));
+        this.#lastConsumedOffsets.delete(key);
+      });
+      seekPromises.push(seekPromise);
     }
 
-    /* Clear the cache. */
-    this.#messageCache.clear();
-    /* Clear the offsets - no need to keep them around. */
-    this.#lastConsumedOffsets.clear();
+    /* TODO: we should cry more about this and render the consumer unusable. */
+    await Promise.all(seekPromises).catch(err => this.#logger.error("Seek error. This is effectively a fatal error:" + err));
+
+
+    /* Clear the cache and stored offsets.
+     * We need to do this only if topicPartitions = null (global cache expiry).
+     * This is because in case of a local cache expiry, MessageCache handles
+     * skipping that (and clearing that later before getting new messages). */
+    if (!topicPartitions) {
+      this.#messageCache.clear();
+    }
   }
 
   /**
@@ -1044,9 +837,14 @@ class Consumer {
       if (!(await acquireOrLog(this.#lock, this.#logger)))
         continue;
 
-      /* Invalidate the message cache if needed. */
-      if (this.#messageCache.isStale()) {
-        await this.#clearCacheAndResetPositions(true);
+      /* Invalidate the message cache if needed */
+      const locallyStale = this.#messageCache.popLocallyStale();
+      if (this.#messageCache.isStale()) { /* global staleness */
+        await this.#clearCacheAndResetPositions();
+        await this.#lock.release();
+        continue;
+      } else if (locallyStale.length !== 0) { /* local staleness */
+        await this.#clearCacheAndResetPositions(locallyStale);
         await this.#lock.release();
         continue;
       }
@@ -1153,9 +951,14 @@ class Consumer {
       if (!(await acquireOrLog(this.#lock, this.#logger)))
         continue;
 
-      /* Invalidate the message cache if needed. */
-      if (this.#messageCache.isStale()) {
-        await this.#clearCacheAndResetPositions(true);
+      /* Invalidate the message cache if needed */
+      const locallyStale = this.#messageCache.popLocallyStale();
+      if (this.#messageCache.isStale()) { /* global staleness */
+        await this.#clearCacheAndResetPositions();
+        await this.#lock.release();
+        continue;
+      } else if (locallyStale.length !== 0) { /* local staleness */
+        await this.#clearCacheAndResetPositions(locallyStale);
         await this.#lock.release();
         continue;
       }
@@ -1441,14 +1244,21 @@ class Consumer {
         offset
       };
 
-      /* We need a complete reset of the cache if we're seeking to a different offset even for one partition.
-       * At a later point, this may be improved at the cost of added complexity of maintaining message generation,
-       * or else purging the cache of just those partitions which are seeked. */
-      await this.#clearCacheAndResetPositions(true);
+      /* The ideal sequence of events here is to:
+       * 1. Mark the cache as stale so we don't consume from it any further.
+       * 2. Call clearCacheAndResetPositions() for the topic partition, which is supposed
+       *    to be called after each cache invalidation.
+       *
+       * However, what (2) does is to pop lastConsumedOffsets[topic partition], and seeks to
+       * the said popped value. Seeking is redundant since we seek here anyway. So, we can skip
+       * the seek by just clearing the lastConsumedOffsets[topic partition].
+       */
+      this.#messageCache.markStale([topicPartition]);
+      this.#lastConsumedOffsets.delete(key);
 
       /* It's assumed that topicPartition is already assigned, and thus can be seeked to and committed to.
        * Errors are logged to detect bugs in the internal code. */
-      /* TODO: is it work awaiting seeks to finish? */
+      /* TODO: is it worth awaiting seeks to finish? */
       this.#internalClient.seek(topicPartitionOffset, 0, err => err ? this.#logger.error(err) : null);
       offsetsToCommit.push({
         topic: topicPartition.topic,
@@ -1567,8 +1377,9 @@ class Consumer {
     }
     this.#internalClient.pause(topics);
 
-    // TODO: make this staleness per-partition, not on a global cache level.
-    this.#messageCache.stale = true;
+    /* Mark the messages in the cache as stale, runInternal* will deal with
+     * making it unusable. */
+    this.#messageCache.markStale(topics);
 
     topics.map(JSON.stringify).forEach(topicPartition => this.#pausedPartitions.add(topicPartition));