fix race condition in setContext link that can leak subscription

CHANGELOG.md

@@ -1,3 +1,10 @@
+## Apollo Client 3.3.21 (not yet released)
+
+### Bug fixes
+
+- Fix race condition in `@apollo/client/link/context` that could leak subscriptions if the subscription is cancelled before `operation.setContext` is called. <br/>
+  [@sofianhn](https://github.com/sofianhn) in [#8399](https://github.com/apollographql/apollo-client/pull/8399)
+
 ## Apollo Client 3.3.20
 
 ### Bug fixes

src/link/context/__tests__/index.ts

@@ -205,3 +205,36 @@ it('unsubscribes without throwing before data', done => {
     done();
   }, 10);
 });
+
+it('does not start the next link subscription if the upstream subscription is already closed', done => {
+  let promiseResolved = false;
+  const withContext = setContext(() =>
+    sleep(5).then(() => {
+      promiseResolved = true;
+      return { dynamicallySet: true }
+    }),
+  );
+
+  let mockLinkCalled = false;
+  const mockLink = new ApolloLink(operation => {
+    mockLinkCalled = true;
+    fail('link should not be called');
+  });
+
+  const link = withContext.concat(mockLink);
+
+  let subscriptionReturnedData = false;
+  let handle = execute(link, { query }).subscribe(result => {
+    subscriptionReturnedData = true;
+    fail('subscription should not return data');
+  });
+
+  handle.unsubscribe();
+
+  setTimeout(() => {
+    expect(promiseResolved).toBe(true);
+    expect(mockLinkCalled).toBe(false);
+    expect(subscriptionReturnedData).toBe(false);
+    done();
+  }, 10);
+});

src/link/context/index.ts

@@ -12,10 +12,13 @@ export function setContext(setter: ContextSetter): ApolloLink {
 
     return new Observable(observer => {
       let handle: ZenObservable.Subscription;
+      let closed = false;
       Promise.resolve(request)
         .then(req => setter(req, operation.getContext()))
         .then(operation.setContext)
         .then(() => {
+          // if the observer is already closed, no need to subscribe.
+          if (closed) return;
           handle = forward(operation).subscribe({
             next: observer.next.bind(observer),
             error: observer.error.bind(observer),
@@ -25,6 +28,7 @@ export function setContext(setter: ContextSetter): ApolloLink {
         .catch(observer.error.bind(observer));
 
       return () => {
+        closed = true;
         if (handle) handle.unsubscribe();
       };
     });