Use only 1.5 seconds cache size estimation aligned at batch size * concurrency

Author: emasab

lib/kafkajs/_consumer.js

@@ -197,10 +197,6 @@ class Consumer {
    * The number of partitions to consume concurrently as set by the user, or 1.
    */
   #concurrency = 1;
-
-
-  #runConfig = null;
-
   /**
    * Promise that resolves together with last in progress fetch.
    * It's set to null when no fetch is in progress.
@@ -958,7 +954,7 @@ class Consumer {
   #updateMaxMessageCacheSize() {
     if (this.#maxBatchSize === -1) {
       // In case of unbounded max batch size it returns all available messages
-      // for a partition in each batch. Cache is unbounded as well as
+      // for a partition in each batch. Cache is unbounded given that
       // it takes only one call to process each partition.
       return;
     }
@@ -969,15 +965,11 @@ class Consumer {
       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 concurrent consumption.
-      this.#messageCacheMaxSize = Math.round(1.5 * messagesPerSecondSingleWorker) * this.#concurrency;
-      const minCacheSize = this.#runConfig.eachBatch ? this.#maxBatchesSize : this.#concurrency;
-      if (this.#messageCacheMaxSize < minCacheSize)
-        // Keep at least one batch or one message per worker.
-        // It's possible less workers than requested were active in previous run.
-        this.#messageCacheMaxSize = minCacheSize;
-      else if (this.#messageCacheMaxSize > minCacheSize * 10)
-        // Keep at most 10 messages or batches per requested worker.
-        this.#messageCacheMaxSize = minCacheSize * 10;
+      // Round up to the nearest multiple of `#maxBatchesSize`.
+      this.#messageCacheMaxSize = Math.ceil(
+        Math.round(1.5 * messagesPerSecondSingleWorker) * this.#concurrency
+        / this.#maxBatchesSize
+      ) * this.#maxBatchesSize;
     }
   }
 
@@ -1110,24 +1102,6 @@ class Consumer {
       this.#messageCacheMaxSize);
   }
 
-  /**
-   * Consumes n messages from the internal consumer.
-   * @returns {Promise<import("../..").Message[]>} A promise that resolves to a list of messages. The size of this list is guaranteed to be less than or equal to n.
-   * @note this method cannot be used in conjunction with #consumeSingleCached.
-   * @private
-   */
-  async #consumeN(n) {
-    return new Promise((resolve, reject) => {
-      this.#internalClient.consume(n, (err, messages) => {
-        if (err) {
-          reject(createKafkaJsErrorFromLibRdKafkaError(err));
-          return;
-        }
-        resolve(messages);
-      });
-    });
-  }
-
   /**
    * Flattens a list of topics with partitions into a list of topic, partition.
    * @param {Array<({topic: string, partitions: Array<number>}|{topic: string, partition: number})>} topics
@@ -1593,16 +1567,18 @@ 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)
       : (savedIdx) => this.#consumeCachedN(savedIdx, this.#maxBatchSize);
-    this.#workers = [];
 
     await this.#lock.write(async () => {
+      this.#workers = [];
+      this.#concurrency = config.partitionsConsumedConcurrently;
+      this.#maxBatchesSize = (
+        config.eachBatch && this.#maxBatchSize > 0 ?
+        this.#maxBatchSize :
+        1) * this.#concurrency;
 
       while (!this.#disconnectStarted) {
         if (this.#maxPollIntervalRestart.resolved)
@@ -1639,38 +1615,6 @@ class Consumer {
     this.#maxPollIntervalRestart.resolve();
   }
 
-  /**
-   * Consumes a single message from the consumer within the given timeout.
-   * THIS METHOD IS NOT IMPLEMENTED.
-   * @note This method cannot be used with run(). Either that, or this must be used.
-   *
-   * @param {any} args
-   * @param {number} args.timeout - the timeout in milliseconds, defaults to 1000.
-   * @returns {import("../..").Message|null} a message, or null if the timeout was reached.
-   * @private
-   */
-  async consume({ timeout } = { timeout: 1000 }) {
-    if (this.#state !== ConsumerState.CONNECTED) {
-      throw new error.KafkaJSError('consume can only be called while connected.', { code: error.ErrorCodes.ERR__STATE });
-    }
-
-    if (this.#running) {
-      throw new error.KafkaJSError('consume() and run() cannot be used together.', { code: error.ErrorCodes.ERR__CONFLICT });
-    }
-
-    this.#internalClient.setDefaultConsumeTimeout(timeout);
-    let m = null;
-
-    try {
-      const ms = await this.#consumeN(1);
-      m = ms[0];
-    } finally {
-      this.#internalClient.setDefaultConsumeTimeout(undefined);
-    }
-
-    throw new error.KafkaJSError('consume() is not implemented.' + m, { code: error.ErrorCodes.ERR__NOT_IMPLEMENTED });
-  }
-
   async #commitOffsetsUntilNoStateErr(offsetsToCommit) {
     let err = { code: error.ErrorCodes.ERR_NO_ERROR };
     do {

test/promisified/consumer/consumeMessages.spec.js

@@ -434,9 +434,8 @@ describe.each(cases)('Consumer - partitionsConsumedConcurrently = %s -', (partit
                 inProgressMaxValue = Math.max(inProgress, inProgressMaxValue);
                 if (inProgressMaxValue >= expectedMaxConcurrentWorkers) {
                     maxConcurrentWorkersReached.resolve();
-                } else if (messagesConsumed.length > 2048) {
-                    await sleep(1000);
                 }
+                await sleep(100);
                 inProgress--;
             },
         });

test/promisified/consumer/consumerCacheTests.spec.js

@@ -142,6 +142,8 @@ describe.each(cases)('Consumer message cache - isAutoCommit = %s - partitionsCon
          * the consumers are created with the same groupId, we create them here.
          * TODO: verify correctness of theory. It's conjecture... which solves flakiness. */
         let groupId = `consumer-group-id-${secureRandom()}`;
+        const multiplier = 9;
+        const numMessages = 16 * multiplier;
         consumer = createConsumer({
             groupId,
             maxWaitTimeInMs: 100,
@@ -164,7 +166,6 @@ describe.each(cases)('Consumer message cache - isAutoCommit = %s - partitionsCon
         const messagesConsumed = [];
         const messagesConsumedConsumer1 = [];
         const messagesConsumedConsumer2 = [];
-        let consumer2ConsumeRunning = false;
 
         consumer.run({
             partitionsConsumedConcurrently,
@@ -176,18 +177,13 @@ describe.each(cases)('Consumer message cache - isAutoCommit = %s - partitionsCon
                         { topic: event.topic, partition: event.partition, offset: Number(event.message.offset) + 1 },
                     ]);
 
-                /* Until the second consumer joins, consume messages slowly so as to not consume them all
-                 * before the rebalance triggers. */
-                if (messagesConsumed.length > 1024 && !consumer2ConsumeRunning) {
-                    await sleep(10);
-                }
+                await sleep(100);
             }
         });
 
         /* Evenly distribute 1024*9 messages across 3 partitions */
         let i = 0;
-        const multiplier = 9;
-        const messages = Array(1024 * multiplier)
+        const messages = Array(numMessages)
             .fill()
             .map(() => {
                 const value = secureRandom();
@@ -198,7 +194,7 @@ describe.each(cases)('Consumer message cache - isAutoCommit = %s - partitionsCon
 
         // Wait for the messages - some of them, before starting the
         // second consumer.
-        await waitForMessages(messagesConsumed, { number: 1024 });
+        await waitForMessages(messagesConsumed, { number: 16 });
 
         await consumer2.connect();
         await consumer2.subscribe({ topic: topicName });
@@ -210,18 +206,17 @@ describe.each(cases)('Consumer message cache - isAutoCommit = %s - partitionsCon
         });
 
         await waitFor(() => consumer2.assignment().length > 0, () => null);
-        consumer2ConsumeRunning = true;
 
         /* Now that both consumers have joined, wait for all msgs to be consumed */
-        await waitForMessages(messagesConsumed, { number: 1024 * multiplier });
+        await waitForMessages(messagesConsumed, { number: numMessages });
 
         /* No extra messages should be consumed. */
         await sleep(1000);
-        expect(messagesConsumed.length).toEqual(1024 * multiplier);
+        expect(messagesConsumed.length).toEqual(numMessages);
 
         /* Check if all messages were consumed. */
         expect(messagesConsumed.map(event => (+event.message.offset)).sort((a, b) => a - b))
-            .toEqual(Array(1024 * multiplier).fill().map((_, i) => Math.floor(i / 3)));
+            .toEqual(Array(numMessages).fill().map((_, i) => Math.floor(i / 3)));
 
         /* Consumer2 should have consumed at least one message. */
         expect(messagesConsumedConsumer2.length).toBeGreaterThan(0);