From 9680b0661cf6022cf49865c12aaf8085e8a37847 Mon Sep 17 00:00:00 2001
From: Yaacov Rydzinski <yaacovCR@gmail.com>
Date: Fri, 19 Mar 2021 16:33:31 +0200
Subject: [PATCH] fix(race): concurrent next calls

---
 src/execution/__tests__/stream-test.js | 66 ++++++++++++++++++++
 src/execution/execute.js               | 85 +++++++++++++++-----------
 2 files changed, 117 insertions(+), 34 deletions(-)

diff --git a/src/execution/__tests__/stream-test.js b/src/execution/__tests__/stream-test.js
index 9e4796edaab..f8f57c254df 100644
--- a/src/execution/__tests__/stream-test.js
+++ b/src/execution/__tests__/stream-test.js
@@ -389,6 +389,72 @@ describe('Execute: stream directive', () => {
       },
     ]);
   });
+  it('Can stream a field that returns an async iterable', async () => {
+    const document = parse(`
+      query { 
+        asyncIterableList @stream(initialCount: 2) {
+          name
+          id
+        }
+      }
+    `);
+    const schema = new GraphQLSchema({ query });
+
+    const result = await execute({ schema, document, rootValue: {} });
+
+    const results = [];
+    if (isAsyncIterable(result)) {
+      const asyncResults = await Promise.all([
+        result.next(),
+        result.next(),
+        result.next(),
+        result.next(),
+      ]);
+      results.push(...asyncResults);
+    }
+
+    expect(results).to.deep.equal([
+      {
+        done: false,
+        value: {
+          data: {
+            asyncIterableList: [
+              {
+                name: 'Luke',
+                id: '1',
+              },
+              {
+                name: 'Han',
+                id: '2',
+              },
+            ],
+          },
+          hasNext: true,
+        },
+      },
+      {
+        done: false,
+        value: {
+          data: {
+            name: 'Leia',
+            id: '3',
+          },
+          path: ['asyncIterableList', 2],
+          hasNext: true,
+        },
+      },
+      {
+        done: false,
+        value: {
+          hasNext: false,
+        },
+      },
+      {
+        done: true,
+        value: undefined,
+      },
+    ]);
+  });
   it('Handles error thrown in async iterable before initialCount is reached', async () => {
     const document = parse(`
       query { 
diff --git a/src/execution/execute.js b/src/execution/execute.js
index a33353806b6..f8bff712fb9 100644
--- a/src/execution/execute.js
+++ b/src/execution/execute.js
@@ -1774,43 +1774,60 @@ export class Dispatcher {
       });
     }
     return new Promise((resolve) => {
+      let resolved = false;
       this._subsequentPayloads.forEach((promise) => {
-        promise.then(() => {
-          // resolve with actual promise, not resolved value of promise so we can remove it from this._subsequentPayloads
-          resolve({ promise });
-        });
-      });
-    })
-      .then(({ promise }) => {
-        this._subsequentPayloads.splice(
-          this._subsequentPayloads.indexOf(promise),
-          1,
-        );
-        return promise;
-      })
-      .then(({ value, done }) => {
-        if (done && this._subsequentPayloads.length === 0) {
-          // async iterable resolver just finished and no more pending payloads
-          return {
-            value: {
-              hasNext: false,
-            },
-            done: false,
+        promise.then((payload) => {
+          if (resolved) {
+            return;
+          }
+
+          resolved = true;
+
+          if (this._subsequentPayloads.length === 0) {
+            // a different call to next has exhausted all payloads
+            resolve({ value: undefined, done: true });
+            return;
+          }
+
+          const index = this._subsequentPayloads.indexOf(promise);
+
+          if (index === -1) {
+            // a different call to next has consumed this payload
+            resolve(this._race());
+            return;
+          }
+
+          this._subsequentPayloads.splice(index, 1);
+
+          const { value, done } = payload;
+
+          if (done && this._subsequentPayloads.length === 0) {
+            // async iterable resolver just finished and no more pending payloads
+            resolve({
+              value: {
+                hasNext: false,
+              },
+              done: false,
+            });
+            return;
+          } else if (done) {
+            // async iterable resolver just finished but there are pending payloads
+            // return the next one
+            resolve(this._race());
+            return;
+          }
+
+          const returnValue: ExecutionPatchResult = {
+            ...value,
+            hasNext: this._subsequentPayloads.length > 0,
           };
-        } else if (done) {
-          // async iterable resolver just finished but there are pending payloads
-          // return the next one
-          return this._race();
-        }
-        const returnValue: ExecutionPatchResult = {
-          ...value,
-          hasNext: this._subsequentPayloads.length > 0,
-        };
-        return {
-          value: returnValue,
-          done: false,
-        };
+          resolve({
+            value: returnValue,
+            done: false,
+          });
+        });
       });
+    });
   }
 
   _next(): Promise<IteratorResult<AsyncExecutionResult, void>> {