refactor: Ensure cancelable generator cancels pending next calls.

Author: timoxley

src/utils/PushQueue.ts

@@ -234,14 +234,18 @@ export default class PushQueue<T> {
     }
 
     async _cleanup() {
+        // capture error and pending next promises
         const { error, nextQueue } = this
         this.finished = true
         this.error = undefined
         this.pending = 0
+        // empty buffer then reassign
         this.buffer.length = 0
         this.buffer = []
+        // reassign nextQueue, emptying would mutate value we captured
         this.nextQueue = []
         const doneValue = { value: undefined, done: true }
+        // resolve all pending next promises
         while (nextQueue.length) {
             const p = nextQueue.shift()
             if (!p) { continue }

src/utils/iterators.js

@@ -198,6 +198,8 @@ export function CancelableGenerator(iterable, onFinally = () => {}, { timeout =
         return onDone
     }
 
+    let pendingNext = 0
+
     async function* CancelableGeneratorFn() {
         // manually iterate
         iterator = iterable[Symbol.asyncIterator]()
@@ -206,7 +208,8 @@ export function CancelableGenerator(iterable, onFinally = () => {}, { timeout =
             yield* {
                 // each next() races against cancel signal
                 next: async (...args) => {
-                    cancelSignal.removeAllListeners()
+                    pendingNext += 1
+                    cancelSignal.setMaxListeners(pendingNext)
                     // NOTE:
                     // Very easy to create a memleak here.
                     // Using a shared promise with Promise.race
@@ -226,7 +229,9 @@ export function CancelableGenerator(iterable, onFinally = () => {}, { timeout =
                         iterator.next(...args),
                         cancelPromise,
                     ]).finally(() => {
-                        cancelPromise.resolve({ value: undefined, done: true })
+                        pendingNext -= 1
+                        cancelSignal.setMaxListeners(pendingNext)
+                        cancelSignal.off('cancel', onCancel)
                     })
                 },
                 async throw(err) {

test/unit/iterators.test.js

@@ -8,14 +8,15 @@ import IteratorTest, { expected, MAX_ITEMS } from './IteratorTest'
 
 const WAIT = 20
 
-async function* generate(items = expected) {
-    await wait(WAIT * 0.1)
+async function* generate(items = expected, waitTime = WAIT) {
+    await wait(waitTime * 0.1)
     for await (const item of items) {
-        await wait(WAIT * 0.1)
+        console.log('next', item)
+        await wait(waitTime * 0.1)
         yield item
-        await wait(WAIT * 0.1)
+        await wait(waitTime * 0.1)
     }
-    await wait(WAIT * 0.1)
+    await wait(waitTime * 0.1)
 }
 
 async function* generateThrow(items = expected, { max = MAX_ITEMS, err = new Error('expected') }) {
@@ -669,6 +670,37 @@ describe('Iterator Utils', () => {
             }
         })
 
+        it('can handle errs when queued next calls', async () => {
+            const done = Defer()
+            const expectedError = new Error('expected')
+            try {
+                const itr = CancelableGenerator((async function* Gen() {
+                    yield* generate(expected, 1000)
+                }()), onFinally, {
+                    timeout: WAIT,
+                })
+
+                const tasks = expected.map(async () => itr.next())
+                await wait(100)
+                await itr.cancel(expectedError)
+                const result = Promise.allSettled(tasks)
+                // first is error
+                expect(result[0]).toEqual({ status: 'rejected', reason: expectedError })
+                // rest is undefined result
+                // not sure what good behaviour should be in this case
+                expect(result.slice(1)).toEqual(result.slice(1).map(() => ({
+                    status: 'fulfilled',
+                    value: {
+                        value: undefined,
+                        done: true
+                    }
+                })))
+                expect(itr.isCancelled()).toEqual(true)
+            } finally {
+                await done
+            }
+        }, 10000)
+
         it('can handle queued next calls resolving out of order', async () => {
             const done = Defer()
             try {