withConcurrentAbruptClose should await handlers (#4532)

this makes our helper function more robust/predictable

Author: yaacovCR

src/execution/__tests__/mapAsyncIterable-test.ts

@@ -2,6 +2,7 @@ import { expect } from 'chai';
 import { describe, it } from 'mocha';
 
 import { expectPromise } from '../../__testUtils__/expectPromise.js';
+import { resolveOnNextTick } from '../../__testUtils__/resolveOnNextTick.js';
 
 import { mapAsyncIterable } from '../mapAsyncIterable.js';
 
@@ -248,6 +249,66 @@ describe('mapAsyncIterable', () => {
     });
   });
 
+  it('waits for source handlers before actually throwing', async () => {
+    const abortReason = new Error('aborted');
+    let storedReason: unknown;
+
+    const iterable: AsyncIterableIterator<number> = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next() {
+        return Promise.resolve({ value: 1, done: false });
+      },
+      async throw(reason?: unknown) {
+        if (storedReason === undefined) {
+          await resolveOnNextTick();
+          // eslint-disable-next-line require-atomic-updates
+          storedReason = reason;
+          return { value: undefined, done: true };
+        }
+        // eslint-disable-next-line @typescript-eslint/only-throw-error
+        throw storedReason;
+      },
+      return() {
+        return Promise.resolve({ value: undefined, done: true });
+      },
+    };
+
+    const mapped = mapAsyncIterable(iterable, (x) => x);
+
+    expect(await mapped.next()).to.deep.equal({ value: 1, done: false });
+
+    const thrown = mapped.throw(abortReason);
+    await expectPromise(thrown).toRejectWith('aborted');
+    expect(storedReason).to.equal(abortReason);
+  });
+
+  it('throws given reason, ignoring source throw result', async () => {
+    const iterable: AsyncIterableIterator<number> = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next() {
+        return Promise.resolve({ value: 1, done: false });
+      },
+      throw(_reason?: unknown) {
+        return Promise.resolve({ value: 1, done: false });
+      },
+      return() {
+        return Promise.resolve({ value: undefined, done: true });
+      },
+    };
+
+    const mapped = mapAsyncIterable(iterable, (x) => x);
+
+    expect(await mapped.next()).to.deep.equal({ value: 1, done: false });
+
+    const abortReason = new Error('aborted');
+    const thrown = mapped.throw(abortReason);
+    await expectPromise(thrown).toRejectWith('aborted');
+  });
+
   it('passes through caught errors through async generators', async () => {
     async function* source() {
       yield 1;

src/execution/__tests__/withConcurrentAbruptClose-test.ts

@@ -175,4 +175,42 @@ describe('withConcurrentAbruptClose', () => {
 
     expect(generator[Symbol.asyncIterator]()).to.equal(generator);
   });
+
+  it('awaits beforeThrow so an abrupt close can set the rejection reason', async () => {
+    const abortReason = new Error('aborted');
+    let storedReason: unknown;
+
+    const generator = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next() {
+        return Promise.resolve({ value: undefined, done: true });
+      },
+      throw() {
+        return storedReason === undefined
+          ? Promise.resolve({ value: undefined, done: true })
+          : // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
+            Promise.reject(storedReason);
+      },
+      return() {
+        return Promise.resolve({ value: undefined, done: true });
+      },
+      async [Symbol.asyncDispose]() {
+        await this.return();
+      },
+    };
+
+    const wrapped = withConcurrentAbruptClose(
+      generator,
+      () => undefined,
+      async () => {
+        await Promise.resolve();
+        storedReason = abortReason;
+      },
+    );
+
+    await expectPromise(wrapped.throw(abortReason)).toRejectWith('aborted');
+    expect(storedReason).to.equal(abortReason);
+  });
 });

src/execution/withConcurrentAbruptClose.ts

@@ -32,27 +32,29 @@ export function withConcurrentAbruptClose<T>(
       return generator.next();
     },
     async return() {
-      ignoreErrors(beforeReturn);
+      await ignoreErrors(beforeReturn);
       return generator.return();
     },
     async throw(error?: unknown) {
-      ignoreErrors(() => beforeThrow(error));
+      await ignoreErrors(() => beforeThrow(error));
       return generator.throw(error);
     },
     async [asyncDispose]() {
-      ignoreErrors(beforeReturn);
+      await ignoreErrors(beforeReturn);
       if (typeof generator[asyncDispose] === 'function') {
         await generator[asyncDispose]();
       }
     },
   };
 }
 
-function ignoreErrors(fn: () => PromiseOrValue<unknown>): void {
+function ignoreErrors(
+  fn: () => PromiseOrValue<unknown>,
+): PromiseOrValue<unknown> {
   try {
     const result = fn();
     if (isPromise(result)) {
-      result.catch(() => {
+      return result.catch(() => {
         // ignore error
       });
     }