fix(withConcurrentAbruptClose): do not close unnecessarily (#4644)

do not call abrupt close repeatedly or if closed naturally

Author: yaacovCR

src/execution/__tests__/withConcurrentAbruptClose-test.ts

@@ -5,6 +5,9 @@ import { expectPromise } from '../../__testUtils__/expectPromise.js';
 
 import { withConcurrentAbruptClose } from '../withConcurrentAbruptClose.js';
 
+const asyncDispose: typeof Symbol.asyncDispose =
+  Symbol.asyncDispose ?? Symbol.for('Symbol.asyncDispose');
+
 /* eslint-disable @typescript-eslint/require-await */
 describe('withConcurrentAbruptClose', () => {
   it('calls function when returned', async () => {
@@ -164,6 +167,166 @@ describe('withConcurrentAbruptClose', () => {
     expect(returned).to.equal(true);
   });
 
+  it('calls the abrupt-close function at most once before completion is observed', async () => {
+    let resolveNext!: (result: IteratorResult<number, void>) => void;
+    const nextPromise = new Promise<IteratorResult<number, void>>((resolve) => {
+      resolveNext = resolve;
+    });
+
+    const source: AsyncGenerator<number, void, void> = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next(): Promise<IteratorResult<number, void>> {
+        return nextPromise;
+      },
+      return(): Promise<IteratorResult<number, void>> {
+        return Promise.resolve({ done: true, value: undefined });
+      },
+      throw(reason?: unknown): Promise<IteratorResult<number, void>> {
+        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
+        return Promise.reject(reason);
+      },
+      async [asyncDispose]() {
+        await this.return();
+      },
+    };
+
+    let cleanupCalls = 0;
+    const generator = withConcurrentAbruptClose(source, () => {
+      cleanupCalls += 1;
+    });
+
+    const pendingNext = generator.next();
+    await generator.return();
+    await generator[asyncDispose]();
+
+    resolveNext({ done: true, value: undefined });
+    await pendingNext;
+
+    expect(cleanupCalls).to.equal(1);
+  });
+
+  it('does not call cleanup function again when returned after completion', async () => {
+    let returned = false;
+
+    const source: AsyncGenerator<number, void, void> = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next(): Promise<IteratorResult<number, void>> {
+        return Promise.resolve({ done: true, value: undefined });
+      },
+      return(): Promise<IteratorResult<number, void>> {
+        returned = true;
+        return Promise.resolve({ done: true, value: undefined });
+      },
+      throw(reason?: unknown): Promise<IteratorResult<number, void>> {
+        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
+        return Promise.reject(reason);
+      },
+      async [asyncDispose]() {
+        await this.return();
+      },
+    };
+
+    let called = false;
+    const generator = withConcurrentAbruptClose(source, () => {
+      called = true;
+    });
+
+    expect(await generator.next()).to.deep.equal({
+      value: undefined,
+      done: true,
+    });
+    expect(await generator.return()).to.deep.equal({
+      value: undefined,
+      done: true,
+    });
+
+    expect(called).to.equal(false);
+    expect(returned).to.equal(true);
+  });
+
+  it('does not call cleanup function again when thrown after completion', async () => {
+    let thrownReason: unknown;
+
+    const source: AsyncGenerator<number, void, void> = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next(): Promise<IteratorResult<number, void>> {
+        return Promise.resolve({ done: true, value: undefined });
+      },
+      return(): Promise<IteratorResult<number, void>> {
+        return Promise.resolve({ done: true, value: undefined });
+      },
+      throw(reason?: unknown): Promise<IteratorResult<number, void>> {
+        thrownReason = reason;
+        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
+        return Promise.reject(reason);
+      },
+      async [Symbol.asyncDispose]() {
+        await this.return();
+      },
+    };
+
+    let called = false;
+    const generator = withConcurrentAbruptClose(source, () => {
+      called = true;
+    });
+
+    expect(await generator.next()).to.deep.equal({
+      value: undefined,
+      done: true,
+    });
+
+    const oops = new Error('Oops');
+    await expectPromise(generator.throw(oops)).toRejectWith('Oops');
+
+    expect(called).to.equal(false);
+    expect(thrownReason).to.equal(oops);
+  });
+
+  it('does not call cleanup function again when disposed after completion', async () => {
+    let returned = false;
+
+    const source: AsyncGenerator<number, void, void> = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next(): Promise<IteratorResult<number, void>> {
+        return Promise.resolve({ done: true, value: undefined });
+      },
+      return(): Promise<IteratorResult<number, void>> {
+        returned = true;
+        return Promise.resolve({ done: true, value: undefined });
+      },
+      throw(reason?: unknown): Promise<IteratorResult<number, void>> {
+        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
+        return Promise.reject(reason);
+      },
+      async [Symbol.asyncDispose]() {
+        await this.return();
+      },
+    };
+
+    let called = false;
+    {
+      await using generator = withConcurrentAbruptClose(source, () => {
+        called = true;
+      });
+
+      expect(await generator.next()).to.deep.equal({
+        value: undefined,
+        done: true,
+      });
+    }
+
+    expect(called).to.equal(false);
+    expect(returned).to.equal(true);
+  });
+
   it('returns the generator itself when the `Symbol.asyncIterator` method is called', async () => {
     async function* source() {
       yield 1;

src/execution/incremental/__tests__/stream-test.ts

@@ -2778,18 +2778,66 @@ describe('Execute: stream directive', () => {
 
     assert(returned);
   });
-  it('Returns underlying async iterables when uses resource is disposed', async () => {
+  it('Returns underlying async iterables when resource is disposed before source completion', async () => {
+    let returned = false;
+    const iterable = {
+      [Symbol.asyncIterator]: () => ({
+        next: () =>
+          new Promise(() => {
+            /* never resolves */
+          }),
+        return: () => {
+          returned = true;
+        },
+      }),
+    };
+
+    const document = parse(`
+      query {
+        friendList @stream(initialCount: 0) {
+          id
+        }
+      }
+    `);
+
+    const executeResult = await experimentalExecuteIncrementally({
+      schema,
+      document,
+      rootValue: {
+        friendList: iterable,
+      },
+    });
+    assert('initialResult' in executeResult);
+
+    {
+      await using iterator =
+        executeResult.subsequentResults[Symbol.asyncIterator]();
+      assert(iterator != null);
+
+      const result1 = executeResult.initialResult;
+      expectJSON(result1).toDeepEqual({
+        data: {
+          friendList: [],
+        },
+        pending: [{ id: '0', path: ['friendList'] }],
+        hasNext: true,
+      });
+    }
+
+    assert(returned);
+  });
+
+  it('Does not return underlying async iterables when resource is disposed after source completion', async () => {
     let index = 0;
     let returned = false;
+    const values = [friends[0]];
     const iterable = {
       [Symbol.asyncIterator]: () => ({
         next: () => {
-          if (index > 1) {
-            return new Promise(() => {
-              // never resolves
-            });
+          const friend = values[index++];
+          if (friend == null) {
+            return Promise.resolve({ done: true, value: undefined });
           }
-          const friend = friends[index++];
           return Promise.resolve({ done: false, value: friend });
         },
         return: () => {
@@ -2833,16 +2881,17 @@ describe('Execute: stream directive', () => {
         value: {
           incremental: [
             {
-              items: [{ id: '1' }, { id: '2' }],
+              items: [{ id: '1' }],
               id: '0',
             },
           ],
-          hasNext: true,
+          completed: [{ id: '0' }],
+          hasNext: false,
         },
       });
     }
 
-    assert(returned);
+    assert(!returned);
   });
   it('limits stream batches to the default capacity (100)', async () => {
     const document = parse(`

src/execution/legacyIncremental/__tests__/legacy-stream-test.ts

@@ -2697,18 +2697,65 @@ describe('Execute: stream directive (legacy)', () => {
 
     assert(returned);
   });
-  it('Returns underlying async iterables when uses resource is disposed', async () => {
+  it('Returns underlying async iterables when resource is disposed before source completion', async () => {
+    let returned = false;
+    const iterable = {
+      [Symbol.asyncIterator]: () => ({
+        next: () =>
+          new Promise(() => {
+            /* never resolves */
+          }),
+        return: () => {
+          returned = true;
+        },
+      }),
+    };
+
+    const document = parse(`
+      query {
+        friendList @stream(initialCount: 0) {
+          id
+        }
+      }
+    `);
+
+    const executeResult = await legacyExecuteIncrementally({
+      schema,
+      document,
+      rootValue: {
+        friendList: iterable,
+      },
+    });
+    assert('initialResult' in executeResult);
+
+    {
+      await using iterator =
+        executeResult.subsequentResults[Symbol.asyncIterator]();
+      assert(iterator != null);
+
+      const result1 = executeResult.initialResult;
+      expectJSON(result1).toDeepEqual({
+        data: {
+          friendList: [],
+        },
+        hasNext: true,
+      });
+    }
+
+    assert(returned);
+  });
+
+  it('Does not return underlying async iterables when resource is disposed after source completion', async () => {
     let index = 0;
     let returned = false;
+    const values = [friends[0]];
     const iterable = {
       [Symbol.asyncIterator]: () => ({
         next: () => {
-          if (index > 1) {
-            return new Promise(() => {
-              // never resolves
-            });
+          const friend = values[index++];
+          if (friend == null) {
+            return Promise.resolve({ done: true, value: undefined });
           }
-          const friend = friends[index++];
           return Promise.resolve({ done: false, value: friend });
         },
         return: () => {
@@ -2751,16 +2798,16 @@ describe('Execute: stream directive (legacy)', () => {
         value: {
           incremental: [
             {
-              items: [{ id: '1' }, { id: '2' }],
+              items: [{ id: '1' }],
               path: ['friendList', 0],
             },
           ],
-          hasNext: true,
+          hasNext: false,
         },
       });
     }
 
-    assert(returned);
+    assert(!returned);
   });
   it('limits stream batches to the default capacity (100)', async () => {
     const document = parse(`