feat: add pause/resume consumer api

Signed-off-by: Mirza Brunjadze <[email protected]>

Author: xxzefgh

docs/consumer.md

@@ -183,6 +183,43 @@ If `force` is not `true`, then the method will throw an error if any `MessagesSt
 
 The return value is `void`.
 
+### `pause(partitions)`
+
+Pauses message consumption for specific topic-partitions.
+
+The parameter is an array of `TopicPartition` objects, where each object has:
+
+| Property  | Type     | Description                   |
+| --------- | -------- | ----------------------------- |
+| topic     | `string` | The topic name.               |
+| partition | `number` | The partition number to pause |
+
+**Important:** Pausing only prevents new fetch requests for the specified partitions. Messages that are already buffered in the stream's internal buffer may still be emitted after calling `pause()`.
+
+The return value is `void`.
+
+### `resume(partitions)`
+
+Resumes message consumption for specific topic-partitions that were previously paused.
+
+The parameter is an array of `TopicPartition` objects with the same structure as `pause()`.
+
+The return value is `void`.
+
+### `paused()`
+
+Returns an array of all currently paused topic-partitions.
+
+The return value is an array of `TopicPartition` objects.
+
+### `isPaused(partition)`
+
+Checks if a specific topic-partition is currently paused.
+
+The parameter is a `TopicPartition` object.
+
+The return value is `boolean`.
+
 ## FAQs
 
 ### My consumer is not receiving any message when the application restarts

src/clients/consumer/consumer.ts

@@ -98,7 +98,8 @@ import {
   type ListCommitsOptions,
   type ListOffsetsOptions,
   type Offsets,
-  type OffsetsWithTimestamps
+  type OffsetsWithTimestamps,
+  type TopicPartition
 } from './types.ts'
 
 export class Consumer<Key = Buffer, Value = Buffer, HeaderKey = Buffer, HeaderValue = Buffer> extends Base<
@@ -118,7 +119,8 @@ export class Consumer<Key = Buffer, Value = Buffer, HeaderKey = Buffer, HeaderVa
   #heartbeatInterval: NodeJS.Timeout | null
   #lastHeartbeat: Date | null
   #streams: Set<MessagesStream<Key, Value, HeaderKey, HeaderValue>>
-  #partitionsAssigner: GroupPartitionsAssigner;
+  #partitionsAssigner: GroupPartitionsAssigner
+  #pausedPartitions: Map<string, Set<number>>;
   /*
     The following requests are blocking in Kafka:
 
@@ -158,6 +160,7 @@ export class Consumer<Key = Buffer, Value = Buffer, HeaderKey = Buffer, HeaderVa
     this.#heartbeatInterval = null
     this.#lastHeartbeat = null
     this.#streams = new Set()
+    this.#pausedPartitions = new Map()
     this.#partitionsAssigner = this[kOptions].partitionAssigner ?? roundRobinAssigner
 
     this.#validateGroupOptions(this[kOptions], groupIdAndOptionsValidator)
@@ -297,6 +300,58 @@ export class Consumer<Key = Buffer, Value = Buffer, HeaderKey = Buffer, HeaderVa
     return callback![kCallbackPromise]
   }
 
+  pause (partitions: TopicPartition[]): void {
+    if (!this.assignments) {
+      throw new UserError('Cannot pause partitions before joining a consumer group.')
+    }
+
+    for (const { topic, partition } of partitions) {
+      const assignment = this.assignments.find(a => a.topic === topic)
+      if (!assignment) {
+        throw new UserError(`Topic '${topic}' is not assigned to this consumer.`, { topic })
+      }
+
+      const existing = this.#pausedPartitions.get(topic)
+      if (existing) {
+        existing.add(partition)
+      } else {
+        this.#pausedPartitions.set(topic, new Set([partition]))
+      }
+    }
+  }
+
+  resume (partitions: TopicPartition[]): void {
+    if (!this.assignments) {
+      throw new UserError('Cannot resume partitions before joining a consumer group.')
+    }
+
+    for (const { topic, partition } of partitions) {
+      const assignment = this.assignments.find(a => a.topic === topic)
+      if (!assignment) {
+        throw new UserError(`Topic '${topic}' is not assigned to this consumer.`, { topic })
+      }
+
+      const existing = this.#pausedPartitions.get(topic)
+      if (existing) {
+        existing.delete(partition)
+        if (existing.size === 0) {
+          this.#pausedPartitions.delete(topic)
+        }
+      }
+    }
+  }
+
+  paused (): TopicPartition[] {
+    return this.#pausedPartitions
+      .entries()
+      .flatMap(([topic, partitions]) => Array.from(partitions).map(partition => ({ topic, partition })))
+      .toArray()
+  }
+
+  isPaused (partition: TopicPartition): boolean {
+    return !!this.#pausedPartitions.get(partition.topic)?.has(partition.partition)
+  }
+
   fetch (options: FetchOptions<Key, Value, HeaderKey, HeaderValue>, callback: CallbackWithPromise<FetchResponse>): void
   fetch (options: FetchOptions<Key, Value, HeaderKey, HeaderValue>): Promise<FetchResponse>
   fetch (

src/clients/consumer/messages-stream.ts

@@ -372,6 +372,10 @@ export class MessagesStream<Key, Value, HeaderKey, HeaderValue> extends Readable
         const partitions = assignment.partitions
 
         for (const partition of partitions) {
+          if (this.#consumer.isPaused({ topic, partition })) {
+            continue
+          }
+
           const leader = metadata.topics.get(topic)!.partitions[partition].leader
 
           if (this.#inflightNodes.has(leader)) {

src/clients/consumer/types.ts

@@ -10,6 +10,11 @@ export interface GroupProtocolSubscription {
   metadata?: Buffer | string
 }
 
+export interface TopicPartition {
+  topic: string
+  partition: number
+}
+
 export interface GroupAssignment {
   topic: string
   partitions: number[]

test/clients/consumer/consumer.test.ts

@@ -1,4 +1,4 @@
-import { deepStrictEqual, ok, strictEqual } from 'node:assert'
+import { deepStrictEqual, ok, strictEqual, doesNotThrow } from 'node:assert'
 import { randomUUID } from 'node:crypto'
 import { once } from 'node:events'
 import { test, type TestContext } from 'node:test'
@@ -3264,3 +3264,222 @@ test('metrics should track the number of active topics', async t => {
     deepStrictEqual(activeTopics.values[0].value, 0)
   }
 })
+
+test('pause should prevent fetches from paused partitions during consumption', async t => {
+  const topic = await createTopic(t, true, 2)
+  const messages = Array.from({ length: 10 }, (_, i) => ({
+    topic,
+    key: `key-${i}`,
+    value: `value-${i}`,
+    partition: i % 2
+  }))
+
+  await produceTestMessages({ t, messages })
+
+  const consumer = createConsumer(t)
+  const stream = await consumer.consume({
+    topics: [topic],
+    mode: MessagesStreamModes.EARLIEST,
+    autocommit: true
+  })
+
+  consumer.pause([{ topic, partition: 0 }])
+
+  const received: number[] = []
+  let count = 0
+
+  for await (const message of stream) {
+    received.push(message.partition)
+    count++
+
+    if (count === 5) {
+      break
+    }
+  }
+
+  strictEqual(
+    received.every(p => p === 1),
+    true
+  )
+})
+
+test('resume should allow fetches from previously paused partitions', async t => {
+  const topic = await createTopic(t, true, 2)
+  const messages = Array.from({ length: 10 }, (_, i) => ({
+    topic,
+    key: `key-${i}`,
+    value: `value-${i}`
+  }))
+
+  await produceTestMessages({ t, messages: messages.slice(0, 5), overrideOptions: { partitioner: () => 0 } })
+  await produceTestMessages({ t, messages: messages.slice(5, 10), overrideOptions: { partitioner: () => 1 } })
+
+  const consumer = createConsumer(t)
+  const stream = await consumer.consume({
+    topics: [topic],
+    mode: MessagesStreamModes.EARLIEST,
+    autocommit: true,
+    maxWaitTime: 100
+  })
+
+  consumer.pause([{ topic, partition: 0 }])
+
+  const received: { key: string; partition: number }[] = []
+  let count = 0
+
+  for await (const message of stream) {
+    received.push({ key: message.key.toString(), partition: message.partition })
+    count++
+
+    if (count === 5) {
+      consumer.resume([{ topic, partition: 0 }])
+    }
+
+    if (count === 10) {
+      break
+    }
+  }
+
+  strictEqual(
+    received.slice(0, 5).every(m => m.partition === 1),
+    true
+  )
+  strictEqual(
+    received.slice(5, 10).every(m => m.partition === 0),
+    true
+  )
+})
+
+test('resume should handle resuming non-paused partitions gracefully', async t => {
+  const topic = await createTopic(t, true)
+  const consumer = createConsumer(t)
+
+  consumer.topics.trackAll(topic)
+  await consumer.joinGroup({})
+
+  doesNotThrow(() => consumer.resume([{ topic, partition: 0 }]))
+})
+
+test('pause/resume should throw error if consumer has not joined a group', async t => {
+  const topic = await createTopic(t, true)
+  const consumer = createConsumer(t)
+
+  try {
+    consumer.pause([{ topic, partition: 0 }])
+    throw new Error('Expected error not thrown')
+  } catch (error) {
+    strictEqual(error instanceof UserError, true)
+    strictEqual(error.message, 'Cannot pause partitions before joining a consumer group.')
+  }
+
+  try {
+    consumer.resume([{ topic, partition: 0 }])
+    throw new Error('Expected error not thrown')
+  } catch (error) {
+    strictEqual(error instanceof UserError, true)
+    strictEqual(error.message, 'Cannot resume partitions before joining a consumer group.')
+  }
+})
+
+test('pause/resume should throw error if topic is not assigned to consumer', async t => {
+  const topic1 = await createTopic(t, true)
+  const topic2 = await createTopic(t, true)
+  const consumer = createConsumer(t)
+
+  consumer.topics.trackAll(topic1)
+  await consumer.joinGroup({})
+
+  try {
+    consumer.pause([{ topic: topic2, partition: 0 }])
+    throw new Error('Expected error not thrown')
+  } catch (error) {
+    strictEqual(error instanceof UserError, true)
+    strictEqual(error.message, `Topic '${topic2}' is not assigned to this consumer.`)
+  }
+
+  try {
+    consumer.resume([{ topic: topic2, partition: 0 }])
+    throw new Error('Expected error not thrown')
+  } catch (error) {
+    strictEqual(error instanceof UserError, true)
+    strictEqual(error.message, `Topic '${topic2}' is not assigned to this consumer.`)
+  }
+})
+
+test('pause/resume should handle multiple topic-partitions', async t => {
+  const topic1 = await createTopic(t, true)
+  const topic2 = await createTopic(t, true)
+  const consumer = createConsumer(t)
+
+  consumer.topics.trackAll(topic1, topic2)
+  await consumer.joinGroup({})
+
+  consumer.pause([
+    { topic: topic1, partition: 0 },
+    { topic: topic2, partition: 0 }
+  ])
+
+  strictEqual(consumer.isPaused({ topic: topic1, partition: 0 }), true)
+  strictEqual(consumer.isPaused({ topic: topic2, partition: 0 }), true)
+
+  consumer.resume([
+    { topic: topic1, partition: 0 },
+    { topic: topic2, partition: 0 }
+  ])
+
+  strictEqual(consumer.isPaused({ topic: topic1, partition: 0 }), false)
+  strictEqual(consumer.isPaused({ topic: topic2, partition: 0 }), false)
+})
+
+test('paused should return all paused topic-partitions', async t => {
+  const topic1 = await createTopic(t, true, 2)
+  const topic2 = await createTopic(t, true, 2)
+  const consumer = createConsumer(t)
+
+  consumer.topics.trackAll(topic1, topic2)
+  await consumer.joinGroup({})
+
+  consumer.pause([
+    { topic: topic1, partition: 0 },
+    { topic: topic1, partition: 1 },
+    { topic: topic2, partition: 0 }
+  ])
+
+  deepStrictEqual(consumer.paused(), [
+    { topic: topic1, partition: 0 },
+    { topic: topic1, partition: 1 },
+    { topic: topic2, partition: 0 }
+  ])
+})
+
+test('isPaused should return true for paused topic-partitions', async t => {
+  const topic = await createTopic(t, true)
+  const consumer = createConsumer(t)
+
+  consumer.topics.trackAll(topic)
+  await consumer.joinGroup({})
+
+  consumer.pause([{ topic, partition: 0 }])
+  strictEqual(consumer.isPaused({ topic, partition: 0 }), true)
+})
+
+test('isPaused should return false for non-paused topic-partitions', async t => {
+  const topic = await createTopic(t, true)
+  const consumer = createConsumer(t)
+
+  consumer.topics.trackAll(topic)
+  await consumer.joinGroup({})
+
+  strictEqual(consumer.isPaused({ topic, partition: 0 }), false)
+})
+
+test('isPaused should return false for topic-partitions not assigned to consumer', async t => {
+  const topic1 = await createTopic(t, true)
+  const topic2 = await createTopic(t, true)
+  const consumer = createConsumer(t)
+
+  consumer.topics.trackAll(topic1)
+  await consumer.joinGroup({})
+
+  strictEqual(consumer.isPaused({ topic: topic2, partition: 0 }), false)
+})