fix race condition in setContext link that can leak subscription

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();
       };
     });