Make messages stream breakable automatically

src/clients/consumer/messages-stream.ts

@@ -29,7 +29,9 @@ import {
   type ConsumeOptions,
   type CorruptedMessageHandler,
   type GroupAssignment,
-  type Offsets
+  type Offsets,
+  type MessagesStreamBreakModeValue,
+  MessagesStreamBreakModes
 } from './types.ts'
 
 // Don't move this function as being in the same file will enable V8 to remove.
@@ -46,6 +48,7 @@ export function defaultCorruptedMessageHandler (): boolean {
 export class MessagesStream<Key, Value, HeaderKey, HeaderValue> extends Readable {
   #consumer: Consumer<Key, Value, HeaderKey, HeaderValue>
   #mode: string
+  #breakMode: MessagesStreamBreakModeValue
   #fallbackMode: string
   #options: ConsumeOptions<Key, Value, HeaderKey, HeaderValue>
   #topics: string[]
@@ -82,6 +85,7 @@ export class MessagesStream<Key, Value, HeaderKey, HeaderValue> extends Readable
     super({ objectMode: true, highWaterMark: options.highWaterMark ?? defaultConsumerOptions.highWaterMark })
     this.#consumer = consumer
     this.#mode = mode ?? MessagesStreamModes.LATEST
+    this.#breakMode = options.breakMode ?? MessagesStreamBreakModes.MANUAL
     this.#fallbackMode = fallbackMode ?? MessagesStreamFallbackModes.LATEST
     this.#offsetsToCommit = new Map()
     this.#topics = structuredClone(options.topics)
@@ -283,6 +287,13 @@ export class MessagesStream<Key, Value, HeaderKey, HeaderValue> extends Readable
     return super[Symbol.asyncIterator]()
   }
 
+  collect (): Promise<Message<Key, Value, HeaderKey, HeaderValue>[]> {
+    if (this.#breakMode === MessagesStreamBreakModes.MANUAL) {
+      throw new UserError('Cannot collect messages when the stream is in MANUAL break mode.')
+    }
+    return Array.fromAsync(this)
+  }
+
   _construct (callback: (error?: Error) => void) {
     this.#refreshOffsets(callback as Callback<void>)
   }
@@ -402,6 +413,10 @@ export class MessagesStream<Key, Value, HeaderKey, HeaderValue> extends Readable
           }
 
           this.#pushRecords(metadata, topicIds, response, requestedOffsets)
+
+          if (this.#breakMode === MessagesStreamBreakModes.AFTER_FIRST_BATCH && this.#inflightNodes.size === 0) {
+            this.push(null)
+          }
         })
       }
     })
@@ -520,7 +535,7 @@ export class MessagesStream<Key, Value, HeaderKey, HeaderValue> extends Readable
       this.#autocommit()
     }
 
-    if (canPush && !(this.#shouldClose || this.closed || this.destroyed)) {
+    if (canPush && this.#breakMode !== MessagesStreamBreakModes.AFTER_FIRST_BATCH && !(this.#shouldClose || this.closed || this.destroyed)) {
       process.nextTick(() => {
         this.#fetch()
       })

src/clients/consumer/types.ts

@@ -62,6 +62,14 @@ export const MessagesStreamFallbackModes = {
 export type MessagesStreamFallbackMode = keyof typeof MessagesStreamFallbackModes
 export type MessagesStreamFallbackModeValue =
   (typeof MessagesStreamFallbackModes)[keyof typeof MessagesStreamFallbackModes]
+
+export const MessagesStreamBreakModes = {
+  MANUAL: 'manual',
+  AFTER_FIRST_BATCH: 'after-first-batch',
+} as const
+export type MessagesStreamBreakMode = keyof typeof MessagesStreamBreakModes
+export type MessagesStreamBreakModeValue = (typeof MessagesStreamBreakModes)[keyof typeof MessagesStreamBreakModes]
+
 export interface GroupOptions {
   sessionTimeout?: number
   rebalanceTimeout?: number
@@ -83,6 +91,7 @@ export interface ConsumeBaseOptions<Key, Value, HeaderKey, HeaderValue> {
 export interface StreamOptions {
   topics: string[]
   mode?: MessagesStreamModeValue
+  breakMode?: MessagesStreamBreakModeValue
   fallbackMode?: MessagesStreamFallbackModeValue
   offsets?: TopicWithPartitionAndOffset[]
   onCorruptedMessage?: CorruptedMessageHandler

test/clients/consumer/messages-stream.test.ts

@@ -1,4 +1,4 @@
-import { deepStrictEqual, ok, rejects, strictEqual, throws } from 'node:assert'
+import { deepStrictEqual, equal, ok, rejects, strictEqual, throws } from 'node:assert'
 import { randomUUID } from 'node:crypto'
 import { once } from 'node:events'
 import { Readable } from 'node:stream'
@@ -659,6 +659,83 @@ test('should support asyncIterator interface', async t => {
   }
 })
 
+test('should automatically break in non-manual break mode', async t => {
+  const groupId = createTestGroupId()
+  const topic = await createTopic(t, true)
+
+  // Produce test messages
+  await produceTestMessages(t, topic)
+
+  const consumer = createConsumer(t, groupId, { deserializers: stringDeserializers })
+
+  await consumer.topics.trackAll(topic)
+  await consumer.joinGroup({})
+
+  const stream = await consumer.consume({
+    topics: [topic],
+    mode: MessagesStreamModes.EARLIEST,
+    maxWaitTime: 1000,
+    breakMode: 'after-first-batch',
+  })
+
+  const messages = []
+  for await (const message of stream) {
+    messages.push(message)
+  }
+})
+
+test('should collect messages', async t => {
+  const groupId = createTestGroupId()
+  const topic = await createTopic(t, true)
+
+  // Produce test messages
+  await produceTestMessages(t, topic)
+
+  const consumer = createConsumer(t, groupId, { deserializers: stringDeserializers })
+
+  await consumer.topics.trackAll(topic)
+  await consumer.joinGroup({})
+
+  const stream = await consumer.consume({
+    topics: [topic],
+    mode: MessagesStreamModes.EARLIEST,
+    maxWaitTime: 1000,
+    breakMode: 'after-first-batch',
+  })
+
+  const messages = await stream.collect()
+  equal(messages.length > 0, true, 'Should collect messages from the stream')
+})
+
+test('collecting messages should throw an error if breakMode is manual', async t => {
+  const groupId = createTestGroupId()
+  const topic = await createTopic(t, true)
+
+  // Produce test messages
+  await produceTestMessages(t, topic)
+
+  const consumer = createConsumer(t, groupId, { deserializers: stringDeserializers })
+
+  await consumer.topics.trackAll(topic)
+  await consumer.joinGroup({})
+
+  const stream = await consumer.consume({
+    topics: [topic],
+    mode: MessagesStreamModes.EARLIEST,
+    maxWaitTime: 1000,
+    breakMode: 'manual',
+  })
+
+  await throws(
+    () => stream.collect(),
+    (error: any) => {
+      ok(error instanceof UserError)
+      strictEqual(error.message, 'Cannot collect messages when the stream is in MANUAL break mode.')
+      return true
+    }
+  )
+})
+
 test('should handle deserialization errors', async t => {
   const groupId = createTestGroupId()
   const topic = await createTopic(t, true)