Configurable max batch size

Author: emasab

CHANGELOG.md

@@ -1,3 +1,12 @@
+# confluent-kafka-javascript 1.6.1
+
+v1.6.1 is a maintenance release. It is supported for all usage.
+
+### Enhancements
+
+1. Configurable batch size through the `js.max.batch.size` property (#).
+
+
 # confluent-kafka-javascript 1.6.0
 
 v1.6.0 is a feature release. It is supported for all usage.

MIGRATION.md

@@ -303,9 +303,9 @@ producerRun().then(consumerRun).catch(console.error);
   - The `heartbeat()` no longer needs to be called by the user in the `eachMessage/eachBatch` callback.
     Heartbeats are automatically managed by librdkafka.
   - The `partitionsConsumedConcurrently` is supported by both `eachMessage` and `eachBatch`.
-  - An API compatible version of `eachBatch` is available, but the batch size calculation is not
-    as per configured parameters, rather, a constant maximum size is configured internally. This is subject
-    to change.
+  - An API compatible version of `eachBatch` is available, maximum batch size
+    can be configured through the `js.max.batch.size` configuration property
+    and defaults to 32.
     The property `eachBatchAutoResolve` is supported.
     Within the `eachBatch` callback, use of `uncommittedOffsets` is unsupported,
     and within the returned batch, `offsetLag` and `offsetLagLow` are supported.

examples/performance/performance-primitives.js

@@ -18,6 +18,9 @@ module.exports = {
     runLagMonitoring,
 };
 
+
+const MAX_BATCH_SIZE = process.env.MAX_BATCH_SIZE ? +process.env.MAX_BATCH_SIZE : null;
+
 function baseConfiguration(parameters) {
     let ret = {
         'client.id': 'kafka-test-performance',
@@ -141,6 +144,10 @@ function newCompatibleConsumer(parameters, eachBatch) {
     const autoCommitOpts = autoCommit > 0 ? 
         { 'enable.auto.commit': true, 'auto.commit.interval.ms': autoCommit } :
         { 'enable.auto.commit': false };
+    const jsOpts = {};
+    if (eachBatch && MAX_BATCH_SIZE !== null) {
+        jsOpts['js.max.batch.size'] = MAX_BATCH_SIZE;
+    }
 
     let groupId = eachBatch ? process.env.GROUPID_BATCH : process.env.GROUPID_MESSAGE;
     if (!groupId) {
@@ -152,6 +159,7 @@ function newCompatibleConsumer(parameters, eachBatch) {
         'auto.offset.reset': 'earliest',
         'fetch.queue.backoff.ms': '100',
         ...autoCommitOpts,
+        ...jsOpts,
     });
     return new CompatibleConsumer(consumer);
 }

lib/kafkajs/_consumer.js

@@ -155,17 +155,6 @@ class Consumer {
    */
   #messageCache = null;
 
-  /**
-   * The maximum size of the message cache.
-   * Will be adjusted dynamically.
-   */
-  #messageCacheMaxSize = 1;
-
-  /**
-   * Number of times we tried to increase the cache.
-   */
-  #increaseCount = 0;
-
   /**
    * Whether the user has enabled manual offset management (commits).
    */
@@ -182,6 +171,11 @@ class Consumer {
    */
   #partitionCount = 0;
 
+  /**
+   * Maximum batch size passed in eachBatch calls.
+   */
+  #maxBatchSize = 32;
+
   /**
    * Whether worker termination has been scheduled.
    */
@@ -311,8 +305,6 @@ class Consumer {
      * consumed messages upto N from the internalClient, but the user has stale'd the cache
      * after consuming just k (< N) messages. We seek back to last consumed offset + 1. */
     this.#messageCache.clear();
-    this.#messageCacheMaxSize = 1;
-    this.#increaseCount = 0;
     const clearPartitions = this.assignment();
     const seeks = [];
     for (const topicPartition of clearPartitions) {
@@ -691,6 +683,17 @@ class Consumer {
     this.#cacheExpirationTimeoutMs = this.#maxPollIntervalMs;
     rdKafkaConfig['max.poll.interval.ms'] = this.#maxPollIntervalMs * 2;
 
+    if (rdKafkaConfig['js.max.batch.size'] !== undefined) {
+      const maxBatchSize = +rdKafkaConfig['js.max.batch.size'];
+      if (Number.isInteger(maxBatchSize) && (maxBatchSize === -1 || maxBatchSize > 0)) {
+        throw new error.KafkaJSError(
+          "'js.max.batch.size' must be a positive integer or -1 for unlimited batch size.",
+          { code: error.ErrorCodes.ERR__INVALID_ARG });
+      }
+      this.#maxBatchSize = maxBatchSize;
+      delete rdKafkaConfig['js.max.batch.size'];
+    }
+
     return rdKafkaConfig;
   }
 
@@ -844,33 +847,6 @@ class Consumer {
     await this.commitOffsets();
   }
 
-  /**
-   * 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.
-   * @private
-   */
-   #increaseMaxSize() {
-     if (this.#messageCacheMaxSize === 1024)
-         return;
-     this.#increaseCount++;
-     if (this.#increaseCount <= 1)
-         return;
-     this.#messageCacheMaxSize = Math.min(this.#messageCacheMaxSize << 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.
-   * @private
-   */
-  #decreaseMaxSize(recvdSize) {
-      this.#messageCacheMaxSize = Math.max(Math.floor((recvdSize * 8) / 10), 1);
-      this.#increaseCount = 0;
-  }
-
   /**
    * Converts a list of messages returned by node-rdkafka into a message that can be used by the eachBatch callback.
    * @param {import("../..").Message[]} messages - must not be empty. Must contain messages from the same topic and partition.
@@ -1001,11 +977,6 @@ class Consumer {
       const res = takeFromCache();
       this.#lastFetchClockNs = hrtime.bigint();
       this.#maxPollIntervalRestart.resolve();
-      if (messages.length === this.#messageCacheMaxSize) {
-        this.#increaseMaxSize();
-      } else {
-        this.#decreaseMaxSize(messages.length);
-      }
       return res;
     } finally {
       this.#fetchInProgress.resolve();
@@ -1040,7 +1011,7 @@ class Consumer {
     }
 
     return this.#fetchAndResolveWith(() => this.#messageCache.next(),
-      this.#messageCacheMaxSize);
+      Number.MAX_SAFE_INTEGER);
   }
 
   /**
@@ -1071,7 +1042,7 @@ class Consumer {
 
     return this.#fetchAndResolveWith(() =>
         this.#messageCache.nextN(null, size),
-      this.#messageCacheMaxSize);
+      Number.MAX_SAFE_INTEGER);
   }
 
   /**
@@ -1559,11 +1530,9 @@ class Consumer {
   async #runInternal(config) {
     this.#concurrency = config.partitionsConsumedConcurrently;
     const perMessageProcessor = config.eachMessage ? this.#messageProcessor : this.#batchProcessor;
-    /* TODO: make this dynamic, based on max batch size / size of last message seen. */
-    const maxBatchSize = 32;
     const fetcher = config.eachMessage
       ? (savedIdx) => this.#consumeSingleCached(savedIdx)
-      : (savedIdx) => this.#consumeCachedN(savedIdx, maxBatchSize);
+      : (savedIdx) => this.#consumeCachedN(savedIdx, this.#maxBatchSize);
     this.#workers = [];
 
     await this.#lock.write(async () => {

types/kafkajs.d.ts

@@ -244,7 +244,17 @@ export interface ConsumerConfig {
   partitionAssignors?: PartitionAssignors[],
 }
 
-export type ConsumerGlobalAndTopicConfig = ConsumerGlobalConfig & ConsumerTopicConfig;
+export interface JSConsumerConfig {
+  /**
+   * Maximum batch size passed in eachBatch calls.
+   * A value of -1 means no limit.
+   *
+   * @default 32
+   */
+  'js.max.batch.size'?: string | number
+}
+
+export type ConsumerGlobalAndTopicConfig = ConsumerGlobalConfig & ConsumerTopicConfig & JSConsumerConfig;
 
 export interface ConsumerConstructorConfig extends ConsumerGlobalAndTopicConfig {
   kafkaJS?: ConsumerConfig;