Concurrent idempotent producer error handling fixes

src/cluster/index.js

@@ -1,5 +1,6 @@
 const BrokerPool = require('./brokerPool')
 const Lock = require('../utils/lock')
+const sharedPromiseTo = require('../utils/sharedPromiseTo')
 const createRetry = require('../retry')
 const connectionBuilder = require('./connectionBuilder')
 const flatten = require('../utils/flatten')
@@ -20,6 +21,13 @@ const mergeTopics = (obj, { topic, partitions }) => ({
   [topic]: [...(obj[topic] || []), ...partitions],
 })
 
+const PRIVATE = {
+  CONNECT: Symbol('private:Cluster:connect'),
+  REFRESH_METADATA: Symbol('private:Cluster:refreshMetadata'),
+  REFRESH_METADATA_IF_NECESSARY: Symbol('private:Cluster:refreshMetadataIfNecessary'),
+  FIND_CONTROLLER_BROKER: Symbol('private:Cluster:findControllerBroker'),
+}
+
 module.exports = class Cluster {
   /**
    * @param {Object} options
@@ -95,6 +103,36 @@ module.exports = class Cluster {
       metadataMaxAge,
     })
     this.committedOffsetsByGroup = offsets
+
+    this[PRIVATE.CONNECT] = sharedPromiseTo(async () => {
+      return await this.brokerPool.connect()
+    })
+
+    this[PRIVATE.REFRESH_METADATA] = sharedPromiseTo(async () => {
+      return await this.brokerPool.refreshMetadata(Array.from(this.targetTopics))
+    })
+
+    this[PRIVATE.REFRESH_METADATA_IF_NECESSARY] = sharedPromiseTo(async () => {
+      return await this.brokerPool.refreshMetadataIfNecessary(Array.from(this.targetTopics))
+    })
+
+    this[PRIVATE.FIND_CONTROLLER_BROKER] = sharedPromiseTo(async () => {
+      const { metadata } = this.brokerPool
+
+      if (!metadata || metadata.controllerId == null) {
+        throw new KafkaJSMetadataNotLoaded('Topic metadata not loaded')
+      }
+
+      const broker = await this.findBroker({ nodeId: metadata.controllerId })
+
+      if (!broker) {
+        throw new KafkaJSBrokerNotFound(
+          `Controller broker with id ${metadata.controllerId} not found in the cached metadata`
+        )
+      }
+
+      return broker
+    })
   }
 
   isConnected() {
@@ -106,7 +144,7 @@ module.exports = class Cluster {
    * @returns {Promise<void>}
    */
   async connect() {
-    await this.brokerPool.connect()
+    await this[PRIVATE.CONNECT]()
   }
 
   /**
@@ -132,15 +170,15 @@ module.exports = class Cluster {
    * @returns {Promise<void>}
    */
   async refreshMetadata() {
-    await this.brokerPool.refreshMetadata(Array.from(this.targetTopics))
+    await this[PRIVATE.REFRESH_METADATA]()
   }
 
   /**
    * @public
    * @returns {Promise<void>}
    */
   async refreshMetadataIfNecessary() {
-    await this.brokerPool.refreshMetadataIfNecessary(Array.from(this.targetTopics))
+    await this[PRIVATE.REFRESH_METADATA_IF_NECESSARY]()
   }
 
   /**
@@ -232,21 +270,7 @@ module.exports = class Cluster {
    * @returns {Promise<import("../../types").Broker>}
    */
   async findControllerBroker() {
-    const { metadata } = this.brokerPool
-
-    if (!metadata || metadata.controllerId == null) {
-      throw new KafkaJSMetadataNotLoaded('Topic metadata not loaded')
-    }
-
-    const broker = await this.findBroker({ nodeId: metadata.controllerId })
-
-    if (!broker) {
-      throw new KafkaJSBrokerNotFound(
-        `Controller broker with id ${metadata.controllerId} not found in the cached metadata`
-      )
-    }
-
-    return broker
+    return await this[PRIVATE.FIND_CONTROLLER_BROKER]()
   }
 
   /**

src/consumer/consumerGroup.js

@@ -38,8 +38,7 @@ const isRebalancing = e =>
 const PRIVATE = {
   JOIN: Symbol('private:ConsumerGroup:join'),
   SYNC: Symbol('private:ConsumerGroup:sync'),
-  HEARTBEAT: Symbol('private:ConsumerGroup:heartbeat'),
-  SHAREDHEARTBEAT: Symbol('private:ConsumerGroup:sharedHeartbeat'),
+  SHARED_HEARTBEAT: Symbol('private:ConsumerGroup:sharedHeartbeat'),
 }
 
 module.exports = class ConsumerGroup {
@@ -111,7 +110,7 @@ module.exports = class ConsumerGroup {
 
     this.lastRequest = Date.now()
 
-    this[PRIVATE.SHAREDHEARTBEAT] = sharedPromiseTo(async ({ interval }) => {
+    this[PRIVATE.SHARED_HEARTBEAT] = sharedPromiseTo(async ({ interval }) => {
       const { groupId, generationId, memberId } = this
       const now = Date.now()
 
@@ -404,7 +403,7 @@ module.exports = class ConsumerGroup {
   }
 
   async heartbeat({ interval }) {
-    return this[PRIVATE.SHAREDHEARTBEAT]({ interval })
+    return this[PRIVATE.SHARED_HEARTBEAT]({ interval })
   }
 
   async fetch() {

src/producer/__tests__/idempotentProduceMessages.spec.js

@@ -8,9 +8,8 @@ const {
   createCluster,
   createTopic,
   waitForMessages,
-  waitFor,
 } = require('testHelpers')
-const { KafkaJSError, KafkaJSProtocolError } = require('../../errors')
+const { KafkaJSError } = require('../../errors')
 
 const createProducer = require('../index')
 const createConsumer = require('../../consumer/index')
@@ -136,61 +135,56 @@ describe('Producer > Idempotent producer', () => {
     await consumer.run({ eachMessage: async message => messagesConsumed.push(message) })
 
     await waitForMessages(messagesConsumed, { number: messages.length })
-
-    expect(arrayUnique(messagesConsumed.map(({ message: { value } }) => value))).toHaveLength(
-      messages.length
-    )
+    expect(messagesConsumed).toHaveLength(messages.length)
   })
 
-  it('concurrent produce() calls > where produce() throws a retriable error on the first call, all subsequent calls throw UNKNOWN_PRODUCER_ID', async () => {
+  it('concurrent produce() calls > where produce() throws a retriable error on the first call, all messages are written to the partition once', async () => {
     for (const nodeId of [0, 1, 2]) {
       const broker = await cluster.findBroker({ nodeId })
 
       const brokerProduce = jest.spyOn(broker, 'produce')
       brokerProduce.mockImplementationOnce(async () => {
-        await waitFor(() => brokerProduce.mock.calls.length >= messages.length) // for all the other concurrent calls to have completed
         throw new KafkaJSError('retriable error')
       })
     }
 
-    const settlements = await PromiseAllSettled(
+    await PromiseAllSettled(
       messages.map(m => producer.send({ acks: -1, topic: topicName, messages: [m] }))
-    ).catch(e => e)
+    )
 
-    settlements
-      .filter(({ status }) => status === 'rejected')
-      .forEach(({ reason }) => {
-        expect(reason).toBeInstanceOf(KafkaJSProtocolError)
-        expect(reason.type).toBe('UNKNOWN_PRODUCER_ID')
-      })
+    const messagesConsumed = []
+    await consumer.run({ eachMessage: async message => messagesConsumed.push(message) })
 
-    expect(settlements.filter(({ status }) => status === 'fulfilled')).toHaveLength(1)
+    await waitForMessages(messagesConsumed, { number: messages.length })
+
+    expect(arrayUnique(messagesConsumed.map(({ message: { value } }) => value))).toHaveLength(
+      messages.length
+    )
   })
 
-  it('concurrent produce() calls > where produce() throws a retriable error on 2nd call, all subsequent calls throw OUT_OF_ORDER_SEQUENCE_NUMBER', async () => {
+  it('concurrent produce() calls > where produce() throws a retriable error on 2nd call, all messages are written to the partition once', async () => {
     for (const nodeId of [0, 1, 2]) {
       const broker = await cluster.findBroker({ nodeId })
 
       const brokerProduce = jest.spyOn(broker, 'produce')
       brokerProduce.mockImplementationOnce()
       brokerProduce.mockImplementationOnce(async () => {
-        await waitFor(() => brokerProduce.mock.calls.length >= messages.length) // for all the other concurrent calls to have completed
         throw new KafkaJSError('retriable error')
       })
     }
 
-    const settlements = await PromiseAllSettled(
+    await PromiseAllSettled(
       messages.map(m => producer.send({ acks: -1, topic: topicName, messages: [m] }))
-    ).catch(e => e)
+    )
 
-    settlements
-      .filter(({ status }) => status === 'rejected')
-      .forEach(({ reason }) => {
-        expect(reason).toBeInstanceOf(KafkaJSProtocolError)
-        expect(reason.type).toBe('OUT_OF_ORDER_SEQUENCE_NUMBER')
-      })
+    const messagesConsumed = []
+    await consumer.run({ eachMessage: async message => messagesConsumed.push(message) })
 
-    expect(settlements.filter(({ status }) => status === 'fulfilled')).toHaveLength(2)
+    await waitForMessages(messagesConsumed, { number: messages.length })
+
+    expect(arrayUnique(messagesConsumed.map(({ message: { value } }) => value))).toHaveLength(
+      messages.length
+    )
   })
 
   it('concurrent produce() calls > where produce() throws a retriable error after the message is written to the log, all messages are written to the partition once', async () => {

src/producer/eosManager/index.js

@@ -1,4 +1,5 @@
 const createRetry = require('../../retry')
+const Lock = require('../../utils/lock')
 const { KafkaJSNonRetriableError } = require('../../errors')
 const COORDINATOR_TYPES = require('../../protocol/coordinatorTypes')
 const createStateMachine = require('./transactionStateMachine')
@@ -64,6 +65,11 @@ module.exports = ({
    */
   let producerSequence = {}
 
+  /**
+   * Idempotent production requires a mutex lock per broker to serialize requests with sequence number handling
+   */
+  let brokerMutexLocks = {}
+
   /**
    * Topic partitions already participating in the transaction
    */
@@ -134,6 +140,7 @@ module.exports = ({
             producerId = result.producerId
             producerEpoch = result.producerEpoch
             producerSequence = {}
+            brokerMutexLocks = {}
 
             logger.debug('Initialized producer id & epoch', { producerId, producerEpoch })
           } catch (e) {
@@ -304,6 +311,18 @@ module.exports = ({
         return stateMachine.state() === STATES.TRANSACTING
       },
 
+      async acquireBrokerLock(broker) {
+        if (this.isInitialized()) {
+          brokerMutexLocks[broker.nodeId] =
+            brokerMutexLocks[broker.nodeId] || new Lock({ timeout: 0xffff })
+          await brokerMutexLocks[broker.nodeId].acquire()
+        }
+      },
+
+      releaseBrokerLock(broker) {
+        if (this.isInitialized()) brokerMutexLocks[broker.nodeId].release()
+      },
+
       /**
        * Mark the provided offsets as participating in the transaction for the given consumer group.
        *

src/producer/sendMessages.js

@@ -77,6 +77,7 @@ module.exports = ({ logger, cluster, partitioner, eosManager, retrier }) => {
 
         const topicData = createTopicData(topicDataForBroker)
 
+        await eosManager.acquireBrokerLock(broker)
         try {
           if (eosManager.isTransactional()) {
             await eosManager.addPartitionsToTransaction(topicData)
@@ -118,6 +119,8 @@ module.exports = ({ logger, cluster, partitioner, eosManager, retrier }) => {
         } catch (e) {
           responsePerBroker.delete(broker)
           throw e
+        } finally {
+          await eosManager.releaseBrokerLock(broker)
         }
       })
     }

src/producer/sendMessages.spec.js

@@ -85,6 +85,8 @@ describe('Producer > sendMessages', () => {
       updateSequence: jest.fn(),
       isTransactional: jest.fn().mockReturnValue(false),
       addPartitionsToTransaction: jest.fn(),
+      acquireBrokerLock: jest.fn(),
+      releaseBrokerLock: jest.fn(),
     }
 
     retrier = retry({ retries: 5 })