Keep 1.5s of data in cache

Author: emasab

lib/kafkajs/_consumer.js

@@ -155,6 +155,12 @@ class Consumer {
    */
   #messageCache = null;
 
+  /**
+   * The maximum size of the message cache.
+   * Will be adjusted dynamically.
+   */
+  #messageCacheMaxSize = 1;
+
   /**
    * Whether the user has enabled manual offset management (commits).
    */
@@ -175,6 +181,7 @@ class Consumer {
    * Maximum batch size passed in eachBatch calls.
    */
   #maxBatchSize = 32;
+  #maxBatchesSize = 32;
 
   /**
    * Whether worker termination has been scheduled.
@@ -191,6 +198,9 @@ class Consumer {
    */
   #concurrency = 1;
 
+
+  #runConfig = null;
+
   /**
    * Promise that resolves together with last in progress fetch.
    * It's set to null when no fetch is in progress.
@@ -228,7 +238,15 @@ class Consumer {
   /**
    * Last fetch real time clock in nanoseconds.
    */
-  #lastFetchClockNs = 0;
+  #lastFetchClockNs = 0n;
+  /**
+   * Last number of messages fetched.
+   */
+  #lastFetchedMessageCnt = 0n;
+  /**
+   * Last fetch concurrency used.
+   */
+  #lastFetchedConcurrency = 0n;
 
   /**
    * List of pending operations to be executed after
@@ -691,6 +709,10 @@ class Consumer {
           { code: error.ErrorCodes.ERR__INVALID_ARG });
       }
       this.#maxBatchSize = maxBatchSize;
+      this.#maxBatchesSize = maxBatchSize;
+      if (maxBatchSize === -1) {
+        this.#messageCacheMaxSize = Number.MAX_SAFE_INTEGER;
+      }
       delete rdKafkaConfig['js.max.batch.size'];
     }
 
@@ -933,6 +955,43 @@ class Consumer {
     return returnPayload;
   }
 
+  #updateMaxMessageCacheSize() {
+    if (this.#maxBatchSize === -1) {
+      return;
+    }
+
+    const nowNs = hrtime.bigint();
+    if (this.#lastFetchedMessageCnt > 0 && this.#lastFetchClockNs > 0n &&
+        nowNs > this.#lastFetchClockNs) {
+      const consumptionDurationSeconds = Number(nowNs - this.#lastFetchClockNs) / 1e9;
+      const messagesPerSecondSingleWorker = this.#lastFetchedMessageCnt / this.#lastFetchedConcurrency / consumptionDurationSeconds;
+      // Keep enough messages in the cache for 1.5 seconds of consumption.
+      this.#messageCacheMaxSize = Math.round(1.5 * messagesPerSecondSingleWorker) * this.#concurrency;
+      const minCacheSize = this.#runConfig.eachBatch ? this.#maxBatchesSize : this.#concurrency;
+      if (this.#messageCacheMaxSize < minCacheSize)
+        this.#messageCacheMaxSize = minCacheSize;
+      else if (this.#messageCacheMaxSize > minCacheSize * 10)
+        this.#messageCacheMaxSize = minCacheSize * 10;
+      if (this.#messageCacheMaxSize < 1024)
+        this.#messageCacheMaxSize = 1024;
+    }
+  }
+
+  #saveFetchStats(messages) {
+    this.#lastFetchClockNs = hrtime.bigint();
+    const partitionsNum = new Map();
+    for (const msg of messages) {
+      const key = partitionKey(msg);
+      partitionsNum.set(key, 1);
+      if (partitionsNum.size >= this.#concurrency) {
+        break;
+      }
+    }
+    this.#lastFetchedConcurrency = partitionsNum.size;
+    this.#lastFetchedMessageCnt = messages.length;
+  }
+
+
   async #fetchAndResolveWith(takeFromCache, size) {
     if (this.#fetchInProgress) {
       await this.#fetchInProgress;
@@ -959,6 +1018,8 @@ class Consumer {
       const fetchResult = new DeferredPromise();
       this.#logger.debug(`Attempting to fetch ${size} messages to the message cache`,
         this.#createConsumerBindingMessageMetadata());
+      
+      this.#updateMaxMessageCacheSize();
       this.#internalClient.consume(size, (err, messages) =>
         fetchResult.resolve([err, messages]));
 
@@ -975,7 +1036,7 @@ class Consumer {
 
       this.#messageCache.addMessages(messages);
       const res = takeFromCache();
-      this.#lastFetchClockNs = hrtime.bigint();
+      this.#saveFetchStats(messages);
       this.#maxPollIntervalRestart.resolve();
       return res;
     } finally {
@@ -1011,7 +1072,7 @@ class Consumer {
     }
 
     return this.#fetchAndResolveWith(() => this.#messageCache.next(),
-      Number.MAX_SAFE_INTEGER);
+      this.#messageCacheMaxSize);
   }
 
   /**
@@ -1042,7 +1103,7 @@ class Consumer {
 
     return this.#fetchAndResolveWith(() =>
         this.#messageCache.nextN(null, size),
-      Number.MAX_SAFE_INTEGER);
+      this.#messageCacheMaxSize);
   }
 
   /**
@@ -1528,7 +1589,9 @@ class Consumer {
    * @private
    */
   async #runInternal(config) {
+    this.#runConfig = config;
     this.#concurrency = config.partitionsConsumedConcurrently;
+    this.#maxBatchesSize = this.#maxBatchSize * this.#concurrency;
     const perMessageProcessor = config.eachMessage ? this.#messageProcessor : this.#batchProcessor;
     const fetcher = config.eachMessage
       ? (savedIdx) => this.#consumeSingleCached(savedIdx)
@@ -1543,6 +1606,8 @@ class Consumer {
 
         this.#workerTerminationScheduled = new DeferredPromise();
         this.#lastFetchClockNs = hrtime.bigint();
+        this.#lastFetchedMessageCnt = 0;
+        this.#lastFetchedConcurrency = 0;
         if (this.#pendingOperations.length === 0) {
           const workersToSpawn = Math.max(1, Math.min(this.#concurrency, this.#partitionCount));
           const cacheExpirationLoop = this.#cacheExpirationLoop();

test/promisified/consumer/consumeMessages.spec.js

@@ -642,20 +642,14 @@ describe.each(cases)('Consumer - partitionsConsumedConcurrently = %s -', (partit
 
         let errors = false;
         let receivedMessages = 0;
-        const batchLengths = [1, 1, 2,
-                              /* cache reset */
-                              1, 1];
         consumer.run({
             partitionsConsumedConcurrently,
             eachBatchAutoResolve: true,
             eachBatch: async (event) => {
                 receivedMessages++;
 
                 try {
-                    expect(event.batch.messages.length)
-                        .toEqual(batchLengths[receivedMessages - 1]);
-
-                    if (receivedMessages === 3) {
+                    if (event.batch.messages.length >= 32) {
                         expect(event.isStale()).toEqual(false);
                         await sleep(7500);
                         /* 7.5s 'processing'
@@ -732,31 +726,23 @@ describe.each(cases)('Consumer - partitionsConsumedConcurrently = %s -', (partit
 
         let errors = false;
         let receivedMessages = 0;
-        const batchLengths = [/* first we reach batches of 32 message and fetches of 64
-                               * max poll interval exceeded happens on second
-                               * 32 messages batch of the 64 msg fetch. */
-                              1, 1, 2, 2, 4, 4, 8, 8, 16, 16, 32, 32, 32, 32,
-                              /* max poll interval exceeded, 32 reprocessed +
-                               * 1 new message. */
-                              1, 1, 2, 2, 4, 4, 8, 8, 3];
+        let firstLongBatchProcessing;
         consumer.run({
             partitionsConsumedConcurrently,
             eachBatchAutoResolve: true,
             eachBatch: async (event) => {
                 receivedMessages++;
 
                 try {
-                    expect(event.batch.messages.length)
-                        .toEqual(batchLengths[receivedMessages - 1]);
-
-                    if (receivedMessages === 13) {
+                    if (!firstLongBatchProcessing && event.batch.messages.length >= 32) {
                         expect(event.isStale()).toEqual(false);
                         await sleep(6000);
                         /* 6s 'processing'
                          * cache clearance starts at 7000 */
                         expect(event.isStale()).toEqual(false);
+                        firstLongBatchProcessing = receivedMessages;
                     }
-                    if ( receivedMessages === 14) {
+                    if (firstLongBatchProcessing && receivedMessages === firstLongBatchProcessing + 1) {
                         expect(event.isStale()).toEqual(false);
                         await sleep(10000);
                         /* 10s 'processing'