feat: support addStream and drain during iteration

Author: calebboyd

src/redis.ts

@@ -1,7 +1,7 @@
 import Redis, { ChainableCommander, RedisOptions } from 'ioredis'
 import { RedisStream } from './stream.js'
 import mkDebug from 'debug'
-import { XBatchResult, XStreamResult, env } from './types.js'
+import { XStreamResult, env } from './types.js'
 
 const debug = mkDebug('redis-x-stream')
 
@@ -16,8 +16,15 @@ function isNumber(num: number | string | undefined): num is number {
 export async function readAckDelete(
   stream: RedisStream
 ): Promise<IterableIterator<XStreamResult> | undefined> {
-  const pipeline = stream.client.pipeline(),
-    read = stream.group ? xreadgroup : xread
+  const pipeline = stream.client.pipeline()
+
+  if (stream.draining) {
+    ack(pipeline, stream)
+    await pipeline.exec()
+    return
+  }
+
+  const read = stream.group ? xreadgroup : xread
   xgroup(pipeline, stream)
   ack(pipeline, stream)
   read(pipeline, stream)
@@ -27,17 +34,17 @@ export async function readAckDelete(
     return
   }
 
-  //TODO: NOGROUP the consumer group this client was blocked on no longer exists
+  //TODO: FATAL - NOGROUP the consumer group this client was blocked on no longer exists
   for (const result of responses) {
     if (result[0] && !result[0]?.message.startsWith('BUSYGROUP')) {
       throw responses[0]
     }
   }
-  const result = responses[responses.length - 1][1] as XBatchResult | null
+  const result = responses[responses.length - 1][1] as XStreamResult[] | null
   if (!result) {
     return
   }
-  if (read === xreadgroup) {
+  if (stream.group) {
     for (const stream of result) {
       if (stream[1].length) {
         return result[Symbol.iterator]()
@@ -48,7 +55,10 @@ export async function readAckDelete(
   }
 }
 
-function ack(client: ChainableCommander, { deleteOnAck, pendingAcks, group }: RedisStream): void {
+export function ack(
+  client: ChainableCommander,
+  { deleteOnAck, pendingAcks, group }: RedisStream
+): void {
   if (!group || !pendingAcks.size) return
   for (const [stream, ids] of pendingAcks) {
     client.xack(stream, group, ...ids)
@@ -69,38 +79,29 @@ function xread(client: ChainableCommander, { block, count, streams, buffers }: R
   block = block === Infinity ? 0 : block
   const args: Parameters<typeof client['xread']> = ['COUNT', count] as IncrementalParameters
   if (isNumber(block)) args.unshift('BLOCK', block)
+  args.push('STREAMS', ...streams.keys(), ...streams.values())
   debug(`xread ${args.join(' ')}`)
-  client[buffers ? 'xreadBuffer' : 'xread'](
-    ...args,
-    'STREAMS',
-    ...streams.keys(),
-    ...streams.values()
-  )
+  client[buffers ? 'xreadBuffer' : 'xread'](...args)
 }
 
 function xreadgroup(
   client: ChainableCommander,
-  { block, count, group, consumer, noack, streams, buffers }: RedisStream
+  { block, count, first, group, consumer, noack, streams, buffers }: RedisStream
 ): void {
   block = block === Infinity ? 0 : block
   const args: Parameters<typeof client['xreadgroup']> = [
     'GROUP',
     group as string,
     consumer as string,
-    'COUNT',
-    count.toString(),
   ] as IncrementalParameters
+  if (!first) args.push('COUNT', count.toString())
   if (noack) args.push('NOACK')
   if (isNumber(block)) args.push('BLOCK', block.toString())
+  args.push('STREAMS', ...streams.keys(), ...streams.values())
   debug(`xreadgroup ${args.join(' ')}`)
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   // https://github.com/luin/ioredis/pull/1676#issue-1437398115
-  ;(client as any)[buffers ? 'xreadgroupBuffer' : 'xreadgroup'](
-    ...args,
-    'STREAMS',
-    ...streams.keys(),
-    ...streams.values()
-  )
+  ;(client as KindaAny)[buffers ? 'xreadgroupBuffer' : 'xreadgroup'](...args)
 }
 
 export function createClient(options?: Redis | string | RedisOptions) {

src/stream.spec.ts

@@ -1,3 +1,4 @@
+import { describe, it, expect } from 'vitest'
 import Redis from 'ioredis'
 import redisStream from './stream'
 import { rand } from './test.util.spec'
@@ -60,8 +61,8 @@ describe('RedisStream xread', () => {
   })
 
   it('should manage initial and finished state', async () => {
-    const stream = redisStream('my-stream' + rand())
-    expect(stream.count).toEqual(100)
+    const stream = redisStream({ streams: ['my-stream' + rand()], count: 2 })
+    expect(stream.count).toEqual(2)
     expect(stream.noack).toEqual(false)
     expect(stream.block).toBeUndefined()
     expect(stream.deleteOnAck).toBeFalsy()
@@ -88,7 +89,7 @@ describe('RedisStream xread', () => {
     return Promise.all([str.quit(), arr.quit(), rec.quit()])
   })
 
-  it('should start at zero by default for xread ', () => {
+  it('should start at zero by default for xread', () => {
     const myStream = 'key' + rand()
     const stream = redisStream(myStream)
     expect(stream.streams.get(myStream)).toEqual('0')

src/stream.ts

@@ -1,5 +1,5 @@
 import { hostname } from 'node:os'
-import { createClient, readAckDelete } from './redis.js'
+import { ack, createClient, readAckDelete } from './redis.js'
 import {
   RedisStreamOptions,
   RedisClient,
@@ -31,6 +31,7 @@ export class RedisStream {
   public readonly control?: RedisClient
   public readonly group?: string
   public readonly consumer?: string
+  public readonly blocked: boolean = false
 
   //xread options
   public streams: Map<string, string>
@@ -55,6 +56,12 @@ export class RedisStream {
    */
   public done = false
   public first = false
+  public draining = false
+  private unblocked = false
+  private reading = false
+
+  private readerId: number | null = null
+  private pendingId: Promise<number | null> | null = null
 
   private itr = {
     name: '',
@@ -88,6 +95,7 @@ export class RedisStream {
     }
 
     if (this.block === 0 || this.block === Infinity) {
+      this.blocked = true
       const { client, created } = createClient(options.redisControl)
       this.control = client
       this.createdControlConnection = created
@@ -101,6 +109,13 @@ export class RedisStream {
     this.createdConnection = created
     this.client = client
 
+    if (this.blocked) {
+      this.pendingId = this.client.client('ID').then(
+        (id) => (this.readerId = id),
+        () => (this.pendingId = this.readerId = null)
+      )
+    }
+
     if (options.consumer || options.group) {
       if (!options.group) {
         this.group = '_xs_g_' + options.consumer
@@ -152,9 +167,15 @@ export class RedisStream {
         this.itr.prev = this.ack(this.itr.name, this.itr.prev[0].toString())
       }
       if (!itr.stream) {
+        this.reading = true
         itr.stream = await readAckDelete(this)
+        this.reading = false
       }
       if (!itr.stream && !this.first) {
+        if (this.unblocked) {
+          this.unblocked = false
+          continue
+        }
         return this.quit()
       }
       if (this.first) {
@@ -180,6 +201,7 @@ export class RedisStream {
       if (result.done) {
         itr.entry = null
       } else {
+        //TODO add test case for this '>'
         this.streams.set(itr.name, this.group ? '>' : result.value[0].toString())
         if (this.ackOnIterate) itr.prev = result.value
         yield [itr.name, result.value]
@@ -211,6 +233,47 @@ export class RedisStream {
     return
   }
 
+  private async maybeUnblock() {
+    if (this.reading && !this.done) {
+      if (typeof this.readerId !== 'number') {
+        await this.pendingId
+        this.pendingId = null
+      }
+      if (typeof this.readerId !== 'number') {
+        throw new Error('Unable to read client id')
+      }
+      this.unblocked = true
+      await this.control?.client('UNBLOCK', this.readerId)
+    }
+  }
+
+  public async addStream(streamName: string) {
+    this.streams.set(streamName, '0')
+    await this.maybeUnblock()
+  }
+
+  /**
+   * Iterate through remaining items in the PEL and exit
+   */
+  public async drain() {
+    this.draining = true
+    await this.maybeUnblock()
+  }
+
+  /**
+   * Immediately stop processing entries
+   */
+  public async end() {
+    if (this.control && this.readerId) {
+      const pipeline = this.control.pipeline()
+      ack(pipeline, this)
+      pipeline.client('UNBLOCK', this.readerId)
+      await Promise.all([pipeline.exec(), this.quit()])
+    } else {
+      await this.quit()
+    }
+  }
+
   protected async return(): Promise<void> {
     await this.quit()
   }

src/test.util.spec.ts

@@ -1,11 +1,16 @@
 import { RedisClient, StreamEntry, XEntryResult } from './types.js'
+import mkDebug from 'debug'
+
+const debug = mkDebug('redis-x-stream')
 
 const delay = (ms: number): Promise<void> => new Promise((resolve) => setTimeout(resolve, ms)),
-  times = <T>(count: number, fn: () => T): Array<T> => Array.from(Array(count), fn) as T[],
+  times = <T>(count: number, fn: (_: undefined, i: number) => T): Array<T> =>
+    Array.from(Array(count), fn) as T[],
   quit = async (client: RedisClient): Promise<void> => {
     await client.quit()
     return new Promise((resolve) => client.once('end', resolve))
   },
+  randNum = (min: number, max: number) => Math.floor(Math.random() * (max - min) + min),
   rand = (): string => Math.random().toString(36).slice(6),
   drain = async (iterable: AsyncIterable<XEntryResult>): Promise<Map<string, StreamEntry[]>> => {
     const results = new Map<string, StreamEntry[]>()
@@ -17,15 +22,12 @@ const delay = (ms: number): Promise<void> => new Promise((resolve) => setTimeout
     return results
   },
   redisIdRegex = /\d+-\d/,
-  testEntries = [
-    ['1', 'hi'],
-    ['2', 'hello'],
-    ['3', 'hai'],
-  ]
+  testEntries = times(randNum(7, 23), (_, i) => [i.toString(), rand()])
 async function hydrateForTest(writer: RedisClient, stream: string, ...values: string[][]) {
   if (!values.length) values = testEntries
   const pipeline = writer.pipeline()
   for (const [key, value] of values) {
+    debug(`xadd ${stream} * ${key} ${value}`)
     pipeline.xadd(stream, '*', key, value)
   }
   await pipeline.exec()

src/types.ts

@@ -9,7 +9,6 @@ export type StreamEntry = [StreamEntryId, StreamEntryKeyValues]
 export type XEntryResult = [StreamKey, StreamEntry]
 //Result Types from iterator
 export type XStreamResult = [StreamKey, StreamEntry[]]
-export type XBatchResult = XStreamResult[]
 
 //Buffers
 export type StreamEntryKeyValueBuffers = Buffer[]

src/xread.spec.ts

@@ -1,4 +1,5 @@
 import Redis from 'ioredis'
+import { describe, it, expect, beforeAll, afterAll, beforeEach, afterEach } from 'vitest'
 import { RedisStream } from './stream.js'
 import { RedisClient } from './types.js'
 import { delay, hydrateForTest, quit, times, testEntries, redisIdRegex } from './test.util.spec.js'
@@ -15,7 +16,6 @@ describe('redis-x-stream xread', () => {
     writer = new Redis()
   })
   afterAll(() => quit(writer))
-
   beforeEach(() => {
     prefix = Math.random().toString(36).slice(6) + '_'
     reader = new Redis()
@@ -76,9 +76,10 @@ describe('redis-x-stream xread', () => {
     const stream = new RedisStream(streamName)
     const values = await hydrateForTest(writer, streamName)
     let ackAttempts = 0
-    for await (const [_, [id, __]] of stream) {
+    for await (const [_, [id]] of stream) {
+      void _
       try {
-        await stream.ack(id)
+        stream.ack(id)
       } catch (e: unknown) {
         if (e instanceof Error) {
           ackAttempts++
@@ -90,6 +91,36 @@ describe('redis-x-stream xread', () => {
     expect(ackAttempts).toBe(values.length)
   })
 
+  it('should allow adding a stream on a blocked iterable', async () => {
+    const myStream = key('my-stream')
+    const laterStream = key('later-stream')
+    await hydrateForTest(writer, myStream)
+    await hydrateForTest(writer, laterStream)
+    const stream = new RedisStream({ streams: [myStream], block: Infinity })
+    let i = 0
+    for await (const [streamName, _] of stream) {
+      i++
+      if (i === testEntries.length) {
+        expect(streamName).toEqual(myStream)
+        setTimeout(() => {
+          //TODO: expect stream is blocked?
+          stream.addStream(laterStream)
+        })
+      }
+      if (i > testEntries.length) {
+        expect(streamName).toEqual(laterStream)
+      }
+      if (i === testEntries.length * 2) {
+        setTimeout(() => {
+          i++
+          stream.end() //break;
+        })
+      }
+    }
+    //stream will block indefinitely (i++ in the future to assert after loop)
+    expect(i).toEqual(testEntries.length * 2 + 1)
+  })
+
   it('should not allow re-iteration (done is set)', async () => {
     const streamName = key('my-stream'),
       iterable = new RedisStream(streamName)