fix: add timeout and close handling to connection.ready() (#253)

* fix: add timeout and close handling to connection.ready()

If neither 'connect' nor 'error' fires while waiting in ready(),
the callback never completes, causing pool.get() and kPerformWithRetry
to hang indefinitely. This adds a timeout using the existing
connectTimeout option and also listens for 'close' events.

Signed-off-by: Tristan Burch <tristan@day.ai>

* fix: resolve ready() timeout race conditions in connection

- Use host:port in timeout message when available for consistency with socket timeout
- Destroy socket when ready() timeout fires to prevent racing socket timeout from
  emitting an unhandled error event after listeners are cleaned up
- Skip timeout callback when connection is already in ERROR/CLOSED state to prevent
  double-invocation when connect() throws synchronously (e.g. invalid port)

Signed-off-by: Tristan Burch <tristan@tburch.com>
Signed-off-by: Tristan Burch <tristan@day.ai>

* test: add tests for ready() timeout and close handling

Signed-off-by: Tristan Burch <tristan@day.ai>

* 'TimeoutError' is not used
Signed-off-by: Tristan Burch <tristan@day.ai>

* fix: check membershipActive in updateAssignments metadata callback

When leaveGroup() is called during the metadata-insync phase of a
rebalance, the metadata callback in #updateAssignments() continued
unconditionally and overwrote assignments. Add a #membershipActive
guard consistent with #performJoinGroup() and #performSyncGroup().

Signed-off-by: Tristan Burch <tristan@day.ai>

---------

Signed-off-by: Tristan Burch <tristan@day.ai>
Signed-off-by: Tristan Burch <tristan@tburch.com>

Author: tburch

src/clients/consumer/consumer.ts

@@ -1361,6 +1361,11 @@ export class Consumer<Key = Buffer, Value = Buffer, HeaderKey = Buffer, HeaderVa
 
   #updateAssignments (newAssignments: TopicPartition[], callback: CallbackWithPromise<void>): void {
     this[kMetadata]({ topics: this.topics.current }, (error, metadata) => {
+      if (!this.#membershipActive) {
+        callback(null)
+        return
+      }
+
       if (error) {
         callback(error)
         return

src/network/connection.ts

@@ -308,20 +308,47 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
       callback = createPromisifiedCallback()
     }
 
-    const onConnect = () => {
+    const cleanup = () => {
+      clearTimeout(timeout)
+      this.removeListener('connect', onConnect)
       this.removeListener('error', onError)
+      this.removeListener('close', onClose)
+    }
 
+    const onConnect = () => {
+      cleanup()
       callback!(null)
     }
 
     const onError = (error: Error) => {
-      this.removeListener('connect', onConnect)
-
+      cleanup()
       callback!(error)
     }
 
+    const onClose = () => {
+      cleanup()
+      callback!(new NetworkError('Connection closed while waiting for ready.'))
+    }
+
+    const timeout = setTimeout(() => {
+      cleanup()
+
+      // If the connection already failed (e.g. invalid port), don't double-invoke the callback
+      if (this.#status === ConnectionStatuses.ERROR || this.#status === ConnectionStatuses.CLOSED) {
+        return
+      }
+
+      this.#socket?.destroy()
+      callback!(new TimeoutError(
+        this.#host
+          ? `Connection to ${this.#host}:${this.#port} timed out.`
+          : `Connection ready timed out after ${this.#options.connectTimeout}ms.`
+      ))
+    }, this.#options.connectTimeout)
+
     this.once('connect', onConnect)
     this.once('error', onError)
+    this.once('close', onClose)
     this.emit('ready')
 
     return callback[kCallbackPromise]

test/network/connection.test.ts

@@ -297,6 +297,30 @@ test('Connection.ready should reject when connection errors', async t => {
   await readyPromise
 })
 
+test('Connection.ready should reject with TimeoutError when connection does not become ready in time', async t => {
+  const connection = new Connection('test-client', { connectTimeout: 100 })
+  t.after(() => connection.close())
+
+  await rejects(() => connection.ready() as Promise<unknown>, {
+    message: 'Connection ready timed out after 100ms.'
+  })
+})
+
+test('Connection.ready should reject with NetworkError when connection closes while waiting', async t => {
+  const connection = new Connection('test-client', { connectTimeout: 5000 })
+  t.after(() => connection.close())
+
+  const readyPromise = connection.ready()
+
+  // Simulate a close event while waiting for ready
+  connection.emit('close')
+
+  await rejects(() => readyPromise as Promise<unknown>, {
+    code: 'PLT_KFK_NETWORK',
+    message: 'Connection closed while waiting for ready.'
+  })
+})
+
 test('Connection.close should close the connection', async t => {
   const { port } = await createServer(t)
   const connection = new Connection('test-client')