From c2a88ea76e35e4005838bebce9f4e807cbe37b56 Mon Sep 17 00:00:00 2001
From: Trevor Scheer <trevor.scheer@gmail.com>
Date: Wed, 8 Dec 2021 17:42:58 -0800
Subject: [PATCH] Prevent re-entrant calls to update callback

---
 .../__tests__/gateway/supergraphSdl.test.ts   | 71 +++++++++++++++++++
 gateway-js/src/index.ts                       | 23 +++++-
 2 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/gateway-js/src/__tests__/gateway/supergraphSdl.test.ts b/gateway-js/src/__tests__/gateway/supergraphSdl.test.ts
index 023812a216..d1ae66fd8e 100644
--- a/gateway-js/src/__tests__/gateway/supergraphSdl.test.ts
+++ b/gateway-js/src/__tests__/gateway/supergraphSdl.test.ts
@@ -284,5 +284,76 @@ describe('Using supergraphSdl dynamic configuration', () => {
         /The gateway subgraphs health check failed\. Updating to the provided `supergraphSdl` will likely result in future request failures to subgraphs\. The following error occurred during the health check/,
       );
     });
+
+    it('throws an error when `update` is called after gateway fails to load', async () => {
+      let updateCallback: SupergraphSdlUpdateFunction;
+      const supergraphSdl = getTestingSupergraphSdl();
+      gateway = new ApolloGateway({
+        async supergraphSdl({ update }) {
+          updateCallback = update;
+          return {
+            supergraphSdl: 'invalid SDL',
+          };
+        },
+      });
+
+      try {
+        await gateway.load();
+      } catch {}
+
+      expect(() =>
+        updateCallback!(supergraphSdl),
+      ).toThrowErrorMatchingInlineSnapshot(
+        `"Can't call \`update\` callback after gateway failed to load."`,
+      );
+
+      // gateway failed to load, so we don't want the `afterEach` to call `gateway.stop()`
+      gateway = null;
+    });
+
+    it('throws an error when `update` is called while an update is in progress', async () => {
+      let updateCallback: SupergraphSdlUpdateFunction;
+      const supergraphSdl = getTestingSupergraphSdl();
+      gateway = new ApolloGateway({
+        async supergraphSdl({ update }) {
+          updateCallback = update;
+          return {
+            supergraphSdl,
+          };
+        },
+        experimental_didUpdateComposition() {
+          updateCallback(getTestingSupergraphSdl(fixturesWithUpdate));
+        },
+      });
+
+      await expect(gateway.load()).rejects.toThrowErrorMatchingInlineSnapshot(
+        `"Can't call \`update\` callback while supergraph update is in progress."`,
+      );
+
+      // gateway failed to load, so we don't want the `afterEach` to call `gateway.stop()`
+      gateway = null;
+    });
+
+    it('throws an error when `update` is called after gateway is stopped', async () => {
+      let updateCallback: SupergraphSdlUpdateFunction;
+      const supergraphSdl = getTestingSupergraphSdl();
+      gateway = new ApolloGateway({
+        async supergraphSdl({ update }) {
+          updateCallback = update;
+          return {
+            supergraphSdl,
+          };
+        },
+      });
+
+      await gateway.load();
+      await gateway.stop();
+
+      expect(() =>
+        updateCallback!(getTestingSupergraphSdl(fixturesWithUpdate)),
+      ).toThrowErrorMatchingInlineSnapshot(
+        `"Can't call \`update\` callback after gateway has been stopped."`,
+      );
+    });
   });
 });
diff --git a/gateway-js/src/index.ts b/gateway-js/src/index.ts
index 3621991bc7..dc92af7f33 100644
--- a/gateway-js/src/index.ts
+++ b/gateway-js/src/index.ts
@@ -133,7 +133,8 @@ type GatewayState =
       pollWaitTimer: NodeJS.Timer;
       doneWaiting: () => void;
     }
-  | { phase: 'polling'; pollingDonePromise: Promise<void> };
+  | { phase: 'polling'; pollingDonePromise: Promise<void> }
+  | { phase: 'updating schema' };
 
 // We want to be compatible with `load()` as called by both AS2 and AS3, so we
 // define its argument types ourselves instead of relying on imports.
@@ -614,10 +615,19 @@ export class ApolloGateway implements GraphQLService {
   }
 
   private externalSupergraphUpdateCallback(supergraphSdl: string) {
+    if (this.state.phase === "failed to load") {
+      throw new Error("Can't call `update` callback after gateway failed to load.");
+    } else if (this.state.phase === "updating schema") {
+      throw new Error("Can't call `update` callback while supergraph update is in progress.");
+    } else if (this.state.phase === "stopped") {
+      throw new Error("Can't call `update` callback after gateway has been stopped.");
+    }
+    this.state = { phase: "updating schema" };
     this.updateWithSupergraphSdl({
       supergraphSdl,
       id: this.getIdForSupergraphSdl(supergraphSdl),
     });
+    this.state = { phase: "loaded" };
   }
 
   private async externalSubgraphHealthCheckCallback(supergraphSdl: string) {
@@ -942,6 +952,11 @@ export class ApolloGateway implements GraphQLService {
       case 'loaded':
         // This is the normal case.
         break;
+      case 'updating schema':
+        // This should never happen
+        throw Error(
+          "ApolloGateway.pollServices called from an unexpected state 'updating schema'",
+        );
       default:
         throw new UnreachableCaseError(this.state);
     }
@@ -1306,7 +1321,7 @@ export class ApolloGateway implements GraphQLService {
   // schema polling). Can be called multiple times safely. Once it (async)
   // returns, all gateway background activity will be finished.
   public async stop() {
-    Promise.all(
+    await Promise.all(
       this.toDispose.map((p) =>
         p().catch((e) => {
           this.logger.error(
@@ -1368,6 +1383,10 @@ export class ApolloGateway implements GraphQLService {
         stoppingDone!();
         return;
       }
+      case "updating schema": {
+        // This should never happen
+        throw Error("`ApolloGateway.stop` called from an unexpected state `updating schema`");
+      }
       default:
         throw new UnreachableCaseError(this.state);
     }