Fix cleanup observer blocking unsubscribe (2) (#6985)

src/core/__tests__/QueryManager/index.ts

@@ -16,7 +16,7 @@ import {
 // mocks
 import mockQueryManager from '../../../utilities/testing/mocking/mockQueryManager';
 import mockWatchQuery from '../../../utilities/testing/mocking/mockWatchQuery';
-import { MockApolloLink, mockSingleLink } from '../../../utilities/testing/mocking/mockLink';
+import { MockApolloLink, mockSingleLink, MockLink } from '../../../utilities/testing/mocking/mockLink';
 
 // core
 import { ApolloQueryResult } from '../../types';
@@ -460,6 +460,77 @@ describe('QueryManager', () => {
     expect(subscription.unsubscribe).not.toThrow();
   });
 
+  // Query should be aborted on last .unsubscribe()
+  itAsync('causes link unsubscription if unsubscribed', (resolve, reject) => {
+    const expResult = {
+      data: {
+        allPeople: {
+          people: [
+            {
+              name: 'Luke Skywalker',
+            },
+          ],
+        },
+      },
+    };
+
+    const request = {
+      query: gql`
+        query people {
+          allPeople(first: 1) {
+            people {
+              name
+            }
+          }
+        }
+      `,
+      variables: undefined
+    };
+
+    const mockedResponse = {
+      request,
+      result: expResult
+    };
+
+    const onRequestSubscribe = jest.fn();
+    const onRequestUnsubscribe = jest.fn();
+
+    const mockedSingleLink = new MockLink([mockedResponse], {
+      addTypename: true,
+      onSubscribe: onRequestSubscribe,
+      onUnsubscribe: onRequestUnsubscribe
+    });
+
+    const mockedQueryManger = new QueryManager({
+      link: mockedSingleLink,
+      cache: new InMemoryCache({ addTypename: false }),
+    });
+
+    const observableQuery = mockedQueryManger.watchQuery({
+      query: request.query,
+      variables: request.variables,
+      notifyOnNetworkStatusChange: false
+    });
+
+    const subscription = observableQuery.subscribe({
+      next: wrap(reject, () => {
+        reject(new Error('Link subscriptions should have been cancelled'));
+      }),
+    });
+
+    subscription.unsubscribe();
+
+    return new Promise(
+      // Unsubscribing from the link happens after a microtask
+      // (Promise.resolve().then) delay, so we need to wait at least that
+      // long before verifying onRequestUnsubscribe was called.
+      resolve => setTimeout(resolve, 0)
+    ).then(() => {
+      expect(onRequestSubscribe).toHaveBeenCalledTimes(1);
+      expect(onRequestUnsubscribe).toHaveBeenCalledTimes(1);
+    }).then(resolve, reject);
+  });
+
   itAsync('supports interoperability with other Observable implementations like RxJS', (resolve, reject) => {
     const expResult = {
       data: {

src/react/components/__tests__/client/Mutation.test.tsx

@@ -108,7 +108,7 @@ describe('General Mutation testing', () => {
 
     function mockClient(m: any) {
       return new ApolloClient({
-        link: new MockLink(m, false),
+        link: new MockLink(m, { addTypename: false }),
         cache: new Cache({ addTypename: false })
       });
     }

src/utilities/observables/Concast.ts

@@ -92,7 +92,7 @@ export class Concast<T> extends Observable<T> {
     // source observable. It's tempting to do this step lazily in
     // addObserver, but this.promise can be accessed without calling
     // addObserver, so consumption needs to begin eagerly.
-    this.handlers.complete!();
+    this.handlers.complete();
   }
 
   private deliverLastMessage(observer: Observer<T>) {
@@ -131,15 +131,12 @@ export class Concast<T> extends Observable<T> {
     quietly?: boolean,
   ) {
     if (this.observers.delete(observer) &&
-        this.observers.size < 1) {
-      if (quietly) return;
-      if (this.sub) {
-        this.sub.unsubscribe();
-        // In case anyone happens to be listening to this.promise, after
-        // this.observers has become empty.
-        this.reject(new Error("Observable cancelled prematurely"));
-      }
-      this.sub = null;
+        --this.addCount < 1 &&
+        !quietly) {
+      // In case there are still any cleanup observers in this.observers,
+      // and no error or completion has been broadcast yet, make sure
+      // those observers receive an error that terminates them.
+      this.handlers.error(new Error("Observable cancelled prematurely"));
     }
   }
 
@@ -159,17 +156,21 @@ export class Concast<T> extends Observable<T> {
 
   // Bound handler functions that can be reused for every internal
   // subscription.
-  private handlers: Observer<T> = {
-    next: result => {
+  private handlers = {
+    next: (result: T) => {
       if (this.sub !== null) {
         this.latest = ["next", result];
         iterateObserversSafely(this.observers, "next", result);
       }
     },
 
-    error: error => {
-      if (this.sub !== null) {
-        if (this.sub) this.sub.unsubscribe();
+    error: (error: any) => {
+      const { sub } = this;
+      if (sub !== null) {
+        // Delay unsubscribing from the underlying subscription slightly,
+        // so that immediately subscribing another observer can keep the
+        // subscription active.
+        if (sub) Promise.resolve().then(() => sub.unsubscribe());
         this.sub = null;
         this.latest = ["error", error];
         this.reject(error);
@@ -209,13 +210,10 @@ export class Concast<T> extends Observable<T> {
     const once = () => {
       if (!called) {
         called = true;
-        // If there have been no other (non-cleanup) observers added, pass
-        // true for the quietly argument, so the removal of the cleanup
-        // observer does not call this.sub.unsubscribe. If a cleanup
-        // observer is added and removed before any other observers
-        // subscribe, we do not want to prevent other observers from
-        // subscribing later.
-        this.removeObserver(observer, !this.addCount);
+        // Removing a cleanup observer should not unsubscribe from the
+        // underlying Observable, so the only removeObserver behavior we
+        // need here is to delete observer from this.observers.
+        this.observers.delete(observer);
         callback();
       }
     }
@@ -236,7 +234,7 @@ export class Concast<T> extends Observable<T> {
   public cancel = (reason: any) => {
     this.reject(reason);
     this.sources = [];
-    this.handlers.complete!();
+    this.handlers.complete();
   }
 }
 

src/utilities/testing/mocking/MockedProvider.tsx

@@ -47,7 +47,7 @@ export class MockedProvider extends React.Component<
       defaultOptions,
       link: link || new MockLink(
         mocks || [],
-        addTypename,
+        { addTypename },
       ),
       resolvers,
     });

src/utilities/testing/mocking/mockLink.ts

@@ -35,17 +35,25 @@ function requestToKey(request: GraphQLRequest, addTypename: Boolean): string {
   return JSON.stringify(requestKey);
 }
 
+interface MockLinkOptions {
+  addTypename?: boolean;
+  onSubscribe?: () => void;
+  onUnsubscribe?: () => void;
+}
+
 export class MockLink extends ApolloLink {
   public operation: Operation;
   public addTypename: Boolean = true;
   private mockedResponsesByKey: { [key: string]: MockedResponse[] } = {};
+  private options: MockLinkOptions;
 
   constructor(
     mockedResponses: ReadonlyArray<MockedResponse>,
-    addTypename: Boolean = true
+    options: MockLinkOptions = {}
   ) {
     super();
-    this.addTypename = addTypename;
+    this.options = options;
+    this.addTypename = options.addTypename ?? true;
     if (mockedResponses) {
       mockedResponses.forEach(mockedResponse => {
         this.addMockedResponse(mockedResponse);
@@ -109,7 +117,9 @@ export class MockLink extends ApolloLink {
       }
     }
 
-    return new Observable(observer => {
+    const requestObservable = new Observable<FetchResult>(observer => {
+      this.options.onSubscribe?.();
+
       const timer = setTimeout(() => {
         if (configError) {
           try {
@@ -141,9 +151,12 @@ export class MockLink extends ApolloLink {
       }, response && response.delay || 0);
 
       return () => {
+        this.options.onUnsubscribe?.();
         clearTimeout(timer);
       };
     });
+
+    return requestObservable;
   }
 
   private normalizeMockedResponse(
@@ -183,5 +196,5 @@ export function mockSingleLink(
     maybeTypename = true;
   }
 
-  return new MockLink(mocks, maybeTypename);
+  return new MockLink(mocks, { addTypename: maybeTypename });
 }