Broker rediscovery with config.brokers parameter taking callback function

src/cluster/__tests__/brokerPool.spec.js

@@ -28,8 +28,8 @@ describe('Cluster > BrokerPool', () => {
   })
 
   describe('#connect', () => {
-    it('connects to the seed broker when the broker pool is created', async () => {
-      expect(brokerPool.seedBroker.isConnected()).toEqual(false)
+    it('when the broker pool is created seed broker is null', async () => {
+      expect(brokerPool.seedBroker).toEqual(undefined)
       await brokerPool.connect()
       expect(brokerPool.seedBroker.isConnected()).toEqual(true)
     })
@@ -41,6 +41,8 @@ describe('Cluster > BrokerPool', () => {
     })
 
     test('select a different seed broker on ILLEGAL_SASL_STATE error', async () => {
+      await brokerPool.createSeedBroker()
+
       const originalSeedPort = brokerPool.seedBroker.connection.port
       const illegalStateError = new KafkaJSProtocolError({
         message: 'ILLEGAL_SASL_STATE',
@@ -57,6 +59,8 @@ describe('Cluster > BrokerPool', () => {
     })
 
     test('select a different seed broker on connection errors', async () => {
+      await brokerPool.createSeedBroker()
+
       const originalSeedPort = brokerPool.seedBroker.connection.port
       brokerPool.seedBroker.connect = jest.fn(() => {
         throw new KafkaJSConnectionError('Test connection error')

src/cluster/__tests__/connectionBuilder.spec.js

@@ -1,7 +1,7 @@
 const { newLogger } = require('testHelpers')
 const connectionBuilder = require('../connectionBuilder')
 const Connection = require('../../network/connection')
-const { KafkaJSNonRetriableError } = require('../../errors')
+const { KafkaJSConnectionError } = require('../../errors')
 
 describe('Cluster > ConnectionBuilder', () => {
   let builder
@@ -28,8 +28,8 @@ describe('Cluster > ConnectionBuilder', () => {
     })
   })
 
-  test('creates a new connection using a random broker', () => {
-    const connection = builder.build()
+  test('creates a new connection using a random broker', async () => {
+    const connection = await builder.build()
     expect(connection).toBeInstanceOf(Connection)
     expect(connection.host).toBeOneOf(['host.test', 'host2.test', 'host3.test'])
     expect(connection.port).toBeOneOf([7777, 7778, 7779])
@@ -42,18 +42,18 @@ describe('Cluster > ConnectionBuilder', () => {
     expect(connection.socketFactory).toBe(socketFactory)
   })
 
-  test('when called without host and port iterates throught the seed brokers', () => {
-    const connections = Array(brokers.length)
-      .fill()
-      .map(() => {
-        const { host, port } = builder.build()
-        return `${host}:${port}`
-      })
+  test('when called without host and port picks random seed broker', async () => {
+    const connections = []
+    for (let i = 0; i < brokers.length; i++) {
+      const { host, port } = await builder.build()
+      connections.push(`${host}:${port}`)
+    }
+
     expect(connections).toIncludeSameMembers(brokers)
   })
 
-  test('accepts overrides for host, port and rack', () => {
-    const connection = builder.build({
+  test('accepts overrides for host, port and rack', async () => {
+    const connection = await builder.build({
       host: 'host.another',
       port: 8888,
       rack: 'rack',
@@ -63,9 +63,9 @@ describe('Cluster > ConnectionBuilder', () => {
     expect(connection.rack).toEqual('rack')
   })
 
-  it('throws an exception if brokers list is empty', () => {
-    expect(() => {
-      builder = connectionBuilder({
+  it('throws an exception if brokers list is empty', async () => {
+    await expect(
+      connectionBuilder({
         socketFactory,
         brokers: [],
         ssl,
@@ -74,16 +74,15 @@ describe('Cluster > ConnectionBuilder', () => {
         connectionTimeout,
         retry,
         logger,
-      })
-    }).toThrow(
-      KafkaJSNonRetriableError,
-      'Failed to connect: expected brokers array and got nothing'
+      }).build()
+    ).rejects.toEqual(
+      new KafkaJSConnectionError('Failed to connect: expected brokers array and got nothing')
     )
   })
 
-  it('throws an exception if brokers is null', () => {
-    expect(() => {
-      builder = connectionBuilder({
+  it('throws an exception if brokers is null', async () => {
+    await expect(
+      connectionBuilder({
         socketFactory,
         brokers: null,
         ssl,
@@ -92,10 +91,64 @@ describe('Cluster > ConnectionBuilder', () => {
         connectionTimeout,
         retry,
         logger,
-      })
-    }).toThrow(
-      KafkaJSNonRetriableError,
-      'Failed to connect: expected brokers array and got nothing'
+      }).build()
+    ).rejects.toEqual(
+      new KafkaJSConnectionError('Failed to connect: expected brokers array and got nothing')
     )
   })
+
+  it('brokers can be function that returns array of host:port strings', async () => {
+    const builder = connectionBuilder({
+      socketFactory,
+      brokers: () => ['host.test:7777'],
+      ssl,
+      sasl,
+      clientId,
+      connectionTimeout,
+      retry,
+      logger,
+    })
+
+    const connection = await builder.build()
+    expect(connection).toBeInstanceOf(Connection)
+    expect(connection.host).toBe('host.test')
+    expect(connection.port).toBe(7777)
+  })
+
+  it('brokers can be async function that returns array of host:port strings', async () => {
+    const builder = connectionBuilder({
+      socketFactory,
+      brokers: async () => ['host.test:7777'],
+      ssl,
+      sasl,
+      clientId,
+      connectionTimeout,
+      retry,
+      logger,
+    })
+
+    const connection = await builder.build()
+    expect(connection).toBeInstanceOf(Connection)
+    expect(connection.host).toBe('host.test')
+    expect(connection.port).toBe(7777)
+  })
+
+  it('sasl can be async function that returns sasl parameters', async () => {
+    const builder = connectionBuilder({
+      socketFactory,
+      brokers,
+      sasl: async () => ({
+        username: 'user',
+      }),
+      ssl,
+      clientId,
+      connectionTimeout,
+      retry,
+      logger,
+    })
+
+    const connection = await builder.build()
+    expect(connection).toBeInstanceOf(Connection)
+    expect(connection.sasl).toEqual({ username: 'user' })
+  })
 })

src/cluster/brokerPool.js

@@ -42,11 +42,6 @@ module.exports = class BrokerPool {
         ...options,
       })
 
-    this.seedBroker = this.createBroker({
-      connection: this.connectionBuilder.build(),
-      logger: this.rootLogger,
-    })
-
     this.brokers = {}
     this.metadata = null
     this.metadataExpireAt = null
@@ -60,7 +55,21 @@ module.exports = class BrokerPool {
    */
   hasConnectedBrokers() {
     const brokers = values(this.brokers)
-    return !!brokers.find(broker => broker.isConnected()) || this.seedBroker.isConnected()
+    return (
+      !!brokers.find(broker => broker.isConnected()) ||
+      (this.seedBroker ? this.seedBroker.isConnected() : false)
+    )
+  }
+
+  async createSeedBroker() {
+    if (this.seedBroker) {
+      await this.seedBroker.disconnect()
+    }
+
+    this.seedBroker = this.createBroker({
+      connection: await this.connectionBuilder.build(),
+      logger: this.rootLogger,
+    })
   }
 
   /**
@@ -72,17 +81,18 @@ module.exports = class BrokerPool {
       return
     }
 
+    if (!this.seedBroker) {
+      await this.createSeedBroker()
+    }
+
     return this.retrier(async (bail, retryCount, retryTime) => {
       try {
         await this.seedBroker.connect()
         this.versions = this.seedBroker.versions
       } catch (e) {
         if (e.name === 'KafkaJSConnectionError' || e.type === 'ILLEGAL_SASL_STATE') {
           // Connection builder will always rotate the seed broker
-          this.seedBroker = this.createBroker({
-            connection: this.connectionBuilder.build(),
-            logger: this.rootLogger,
-          })
+          await this.createSeedBroker()
           this.logger.error(
             `Failed to connect to seed broker, trying another broker from the list: ${e.message}`,
             { retryCount, retryTime }
@@ -102,7 +112,7 @@ module.exports = class BrokerPool {
    * @returns {Promise}
    */
   async disconnect() {
-    await this.seedBroker.disconnect()
+    this.seedBroker && (await this.seedBroker.disconnect())
     await Promise.all(values(this.brokers).map(broker => broker.disconnect()))
 
     this.brokers = {}
@@ -126,33 +136,39 @@ module.exports = class BrokerPool {
         this.metadataExpireAt = Date.now() + this.metadataMaxAge
 
         const replacedBrokers = []
-        this.brokers = this.metadata.brokers.reduce((result, { nodeId, host, port, rack }) => {
-          if (result[nodeId]) {
-            if (!hasBrokerBeenReplaced(result[nodeId], { host, port, rack })) {
-              return result
+
+        this.brokers = await this.metadata.brokers.reduce(
+          async (resultPromise, { nodeId, host, port, rack }) => {
+            const result = await resultPromise
+
+            if (result[nodeId]) {
+              if (!hasBrokerBeenReplaced(result[nodeId], { host, port, rack })) {
+                return result
+              }
+
+              replacedBrokers.push(result[nodeId])
             }
 
-            replacedBrokers.push(result[nodeId])
-          }
+            if (host === seedHost && port === seedPort) {
+              this.seedBroker.nodeId = nodeId
+              this.seedBroker.connection.rack = rack
+              return assign(result, {
+                [nodeId]: this.seedBroker,
+              })
+            }
 
-          if (host === seedHost && port === seedPort) {
-            this.seedBroker.nodeId = nodeId
-            this.seedBroker.connection.rack = rack
             return assign(result, {
-              [nodeId]: this.seedBroker,
+              [nodeId]: this.createBroker({
+                logger: this.rootLogger,
+                versions: this.versions,
+                supportAuthenticationProtocol: this.supportAuthenticationProtocol,
+                connection: await this.connectionBuilder.build({ host, port, rack }),
+                nodeId,
+              }),
             })
-          }
-
-          return assign(result, {
-            [nodeId]: this.createBroker({
-              logger: this.rootLogger,
-              versions: this.versions,
-              supportAuthenticationProtocol: this.supportAuthenticationProtocol,
-              connection: this.connectionBuilder.build({ host, port, rack }),
-              nodeId,
-            }),
-          })
-        }, this.brokers)
+          },
+          this.brokers
+        )
 
         const freshBrokerIds = this.metadata.brokers.map(({ nodeId }) => `${nodeId}`).sort()
         const currentBrokerIds = keys(this.brokers).sort()
@@ -283,7 +299,7 @@ module.exports = class BrokerPool {
           }
 
           // Rebuild the connection since it can't recover from illegal SASL state
-          broker.connection = this.connectionBuilder.build({
+          broker.connection = await this.connectionBuilder.build({
             host: broker.connection.host,
             port: broker.connection.port,
             rack: broker.connection.rack,

src/cluster/connectionBuilder.js

@@ -1,10 +1,9 @@
 const Connection = require('../network/connection')
-const { KafkaJSNonRetriableError } = require('../errors')
-const shuffle = require('../utils/shuffle')
+const { KafkaJSConnectionError } = require('../errors')
 
 const validateBrokers = brokers => {
   if (!brokers || brokers.length === 0) {
-    throw new KafkaJSNonRetriableError(`Failed to connect: expected brokers array and got nothing`)
+    throw new KafkaJSConnectionError(`Failed to connect: expected brokers array and got nothing`)
   }
 }
 
@@ -22,27 +21,26 @@ module.exports = ({
   logger,
   instrumentationEmitter = null,
 }) => {
-  validateBrokers(brokers)
-
-  const shuffledBrokers = shuffle(brokers)
-  const size = brokers.length
   let index = 0
 
   return {
-    build: ({ host, port, rack } = {}) => {
+    build: async ({ host, port, rack } = {}) => {
       if (!host) {
-        // Always rotate the seed broker
-        const [seedHost, seedPort] = shuffledBrokers[index++ % size].split(':')
-        host = seedHost
-        port = Number(seedPort)
+        const list = typeof brokers === 'function' ? await brokers() : brokers
+        validateBrokers(list)
+
+        const randomBroker = list[index++ % list.length]
+
+        host = randomBroker.split(':')[0]
+        port = Number(randomBroker.split(':')[1])
       }
 
       return new Connection({
         host,
         port,
         rack,
+        sasl: typeof sasl === 'function' ? await sasl() : sasl,
         ssl,
-        sasl,
         clientId,
         socketFactory,
         connectionTimeout,

types/index.d.ts

@@ -14,10 +14,13 @@ export class Kafka {
   logger(): Logger
 }
 
+export type BrokersFunction = () => Promise<string[]>;
+export type SaslFunction = () => Promise<SASLOptions>;
+
 export interface KafkaConfig {
-  brokers: string[]
+  brokers: string[] | BrokersFunction
   ssl?: tls.ConnectionOptions | boolean
-  sasl?: SASLOptions
+  sasl?: SASLOptions | SaslFunction
   clientId?: string
   connectionTimeout?: number
   authenticationTimeout?: number