improve

Author: unnoq

packages/shared/src/iterator.test.ts

@@ -297,9 +297,7 @@ describe('replicateAsyncIterator', async () => {
 
     const gen = async function* () {
       yield 1
-      await new Promise(resolve => setTimeout(resolve, 1))
-      yield 2
-      yield 3
+      await new Promise(resolve => setTimeout(resolve, 10))
       throw error
     }
 
@@ -310,30 +308,18 @@ describe('replicateAsyncIterator', async () => {
     expect(await iterators[0]!.next()).toEqual({ done: false, value: 1 })
     expect(await iterators[1]!.next()).toEqual({ done: false, value: 1 })
 
-    expect(await iterators[0]!.next()).toEqual({ done: false, value: 2 })
-    expect(await iterators[1]!.next()).toEqual({ done: false, value: 2 })
-
-    expect(await iterators[0]!.next()).toEqual({ done: false, value: 3 })
-    expect(await iterators[1]!.next()).toEqual({ done: false, value: 3 })
-    expect(await iterators[2]!.next()).toEqual({ done: false, value: 1 })
-
     await Promise.all([
       expect(iterators[0]!.next()).rejects.toThrow(error),
       expect(iterators[1]!.next()).rejects.toThrow(error),
     ])
-    expect(await iterators[2]!.next()).toEqual({ done: false, value: 2 })
 
     expect(await iterators[0]!.next()).toEqual({ done: true, value: undefined })
     expect(await iterators[1]!.next()).toEqual({ done: true, value: undefined })
-    expect(await iterators[2]!.next()).toEqual({ done: false, value: 3 })
+    expect(await iterators[2]!.next()).toEqual({ done: false, value: 1 })
 
     expect(await iterators[0]!.next()).toEqual({ done: true, value: undefined })
     expect(await iterators[1]!.next()).toEqual({ done: true, value: undefined })
     await expect(iterators[2]!.next()).rejects.toThrow(error)
-
-    expect(await iterators[0]!.next()).toEqual({ done: true, value: undefined })
-    expect(await iterators[1]!.next()).toEqual({ done: true, value: undefined })
-    expect(await iterators[2]!.next()).toEqual({ done: true, value: undefined })
   })
 
   it('on manual close', async () => {

packages/shared/src/iterator.ts

@@ -1,5 +1,5 @@
 import type { SetSpanErrorOptions } from './otel'
-import { defer, once, sequential } from './function'
+import { once, sequential } from './function'
 import { runInSpanContext, setSpanError, startSpan } from './otel'
 import { AsyncIdQueue } from './queue'
 
@@ -125,8 +125,12 @@ export function replicateAsyncIterator<T, TReturn, TNext>(
         }
       }
     }
-    catch (e) {
-      error = { value: e }
+    catch (reason) {
+      error = { value: reason }
+
+      queue.waiterIds.forEach((id) => {
+        queue.close({ id, reason })
+      })
     }
   })
 
@@ -139,15 +143,14 @@ export function replicateAsyncIterator<T, TReturn, TNext>(
         start()
 
         return new Promise((resolve, reject) => {
-          queue.pull(id)
-            .then(resolve)
-            .catch(reject)
-
-          defer(() => {
-            if (error) {
-              reject(error.value)
-            }
-          })
+          if (!error || queue.hasBufferedItems(id)) {
+            queue.pull(id)
+              .then(resolve)
+              .catch(reject)
+          }
+          else {
+            reject(error.value)
+          }
         })
       },
       async (reason) => {

packages/shared/src/queue.test.ts

@@ -132,4 +132,44 @@ describe('asyncIdQueue', () => {
     expect(queue.isOpen('2')).toBe(false)
     expect(queue.isOpen('3')).toBe(false)
   })
+
+  it('waiterIds', async () => {
+    queue.open('1')
+    queue.open('2')
+
+    const p1 = queue.pull('1')
+    const p2 = queue.pull('2')
+
+    expect(queue.waiterIds).toEqual(['1', '2'])
+
+    queue.push('1', 'item1')
+    queue.push('2', 'item2')
+
+    await expect(p1).resolves.toBe('item1')
+    await expect(p2).resolves.toBe('item2')
+
+    expect(queue.waiterIds).toEqual([])
+  })
+
+  it('hasBufferedItems', async () => {
+    queue.open('1')
+    queue.open('2')
+
+    expect(queue.hasBufferedItems('1')).toBe(false)
+    expect(queue.hasBufferedItems('2')).toBe(false)
+
+    queue.push('1', 'item1')
+    queue.push('2', 'item2')
+
+    expect(queue.hasBufferedItems('1')).toBe(true)
+    expect(queue.hasBufferedItems('2')).toBe(true)
+
+    await queue.pull('1')
+    expect(queue.hasBufferedItems('1')).toBe(false)
+    expect(queue.hasBufferedItems('2')).toBe(true)
+
+    await queue.pull('2')
+    expect(queue.hasBufferedItems('1')).toBe(false)
+    expect(queue.hasBufferedItems('2')).toBe(false)
+  })
 })

packages/shared/src/queue.ts

@@ -5,13 +5,21 @@ export interface AsyncIdQueueCloseOptions {
 
 export class AsyncIdQueue<T> {
   private readonly openIds = new Set<string>()
-  private readonly items = new Map<string, T[]>()
-  private readonly pendingPulls = new Map<string, (readonly [resolve: (item: T) => void, reject: (err: unknown) => void])[]>()
+  private readonly queues = new Map<string, T[]>()
+  private readonly waiters = new Map<string, (readonly [resolve: (item: T) => void, reject: (err: unknown) => void])[]>()
 
   get length(): number {
     return this.openIds.size
   }
 
+  get waiterIds(): string[] {
+    return Array.from(this.waiters.keys())
+  }
+
+  hasBufferedItems(id: string): boolean {
+    return Boolean(this.queues.get(id)?.length)
+  }
+
   open(id: string): void {
     this.openIds.add(id)
   }
@@ -23,77 +31,77 @@ export class AsyncIdQueue<T> {
   push(id: string, item: T): void {
     this.assertOpen(id)
 
-    const pending = this.pendingPulls.get(id)
+    const pending = this.waiters.get(id)
 
     if (pending?.length) {
       pending.shift()![0](item)
 
       if (pending.length === 0) {
-        this.pendingPulls.delete(id)
+        this.waiters.delete(id)
       }
     }
     else {
-      const items = this.items.get(id)
+      const items = this.queues.get(id)
 
       if (items) {
         items.push(item)
       }
       else {
-        this.items.set(id, [item])
+        this.queues.set(id, [item])
       }
     }
   }
 
   async pull(id: string): Promise<T> {
     this.assertOpen(id)
 
-    const items = this.items.get(id)
+    const items = this.queues.get(id)
 
     if (items?.length) {
       const item = items.shift()!
 
       if (items.length === 0) {
-        this.items.delete(id)
+        this.queues.delete(id)
       }
 
       return item
     }
 
     return new Promise<T>((resolve, reject) => {
-      const waitingPulls = this.pendingPulls.get(id)
+      const waitingPulls = this.waiters.get(id)
 
       const pending = [resolve, reject] as const
 
       if (waitingPulls) {
         waitingPulls.push(pending)
       }
       else {
-        this.pendingPulls.set(id, [pending])
+        this.waiters.set(id, [pending])
       }
     })
   }
 
   close({ id, reason }: AsyncIdQueueCloseOptions = {}): void {
     if (id === undefined) {
-      this.pendingPulls.forEach((pendingPulls, id) => {
+      this.waiters.forEach((pendingPulls, id) => {
         pendingPulls.forEach(([, reject]) => {
           reject(reason ?? new Error(`[AsyncIdQueue] Queue[${id}] was closed or aborted while waiting for pulling.`))
         })
       })
 
-      this.pendingPulls.clear()
+      this.waiters.clear()
       this.openIds.clear()
-      this.items.clear()
+      this.queues.clear()
       return
     }
 
-    this.pendingPulls.get(id)?.forEach(([, reject]) => {
+    this.waiters.get(id)?.forEach(([, reject]) => {
       reject(reason ?? new Error(`[AsyncIdQueue] Queue[${id}] was closed or aborted while waiting for pulling.`))
     })
 
-    this.pendingPulls.delete(id)
+    this.waiters.delete(id)
     this.openIds.delete(id)
-    this.items.delete(id)
+    this.queues.delete(id)
   }
 
   assertOpen(id: string): void {