Replace concast.cleanup method with simpler concast.beforeNext API

package.json

@@ -56,7 +56,7 @@
     {
       "name": "apollo-client",
       "path": "./dist/apollo-client.min.cjs",
-      "maxSize": "29.45kB"
+      "maxSize": "29.5kB"
     }
   ],
   "engines": {

src/core/ObservableQuery.ts

@@ -831,12 +831,9 @@ Did you mean to call refetch(variables) instead of refetch({ variables })?`);
 
     if (!useDisposableConcast) {
       // We use the {add,remove}Observer methods directly to avoid wrapping
-      // observer with an unnecessary SubscriptionObserver object, in part so
-      // that we can remove it here without triggering any unsubscriptions,
-      // because we just want to ignore the old observable, not prematurely shut
-      // it down, since other consumers may be awaiting this.concast.promise.
+      // observer with an unnecessary SubscriptionObserver object.
       if (this.concast && this.observer) {
-        this.concast.removeObserver(this.observer, true);
+        this.concast.removeObserver(this.observer);
       }
 
       this.concast = concast;

src/core/QueryManager.ts

@@ -989,7 +989,7 @@ export class QueryManager<TStore> {
 
           byVariables.set(varJson, observable = concast);
 
-          concast.cleanup(() => {
+          concast.beforeNext(() => {
             if (byVariables.delete(varJson) &&
                 byVariables.size < 1) {
               inFlightLinkObservables.delete(serverQuery);
@@ -1161,7 +1161,7 @@ export class QueryManager<TStore> {
         : fromVariables(normalized.variables!)
     );
 
-    concast.cleanup(() => {
+    concast.beforeNext(() => {
       this.fetchCancelFns.delete(queryId);
 
       if (queryInfo.observableQuery) {

src/utilities/observables/Concast.ts

@@ -121,29 +121,24 @@ export class Concast<T> extends Observable<T> {
     }
   }
 
-  // Note: cleanup observers do not count towards this total.
-  private addCount = 0;
-
   public addObserver(observer: Observer<T>) {
     if (!this.observers.has(observer)) {
       // Immediately deliver the most recent message, so we can always
       // be sure all observers have the latest information.
       this.deliverLastMessage(observer);
       this.observers.add(observer);
-      ++this.addCount;
     }
   }
 
-  public removeObserver(
-    observer: Observer<T>,
-    quietly?: boolean,
-  ) {
-    if (this.observers.delete(observer) &&
-        --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
-      // have a chance to run and then remove themselves from this.observers.
+  public removeObserver(observer: Observer<T>) {
+    if (
+      this.observers.delete(observer) &&
+      this.observers.size < 1
+    ) {
+      // In case there are still any listeners in this.nextResultListeners, and
+      // no error or completion has been broadcast yet, make sure those
+      // observers have a chance to run and then remove themselves from
+      // this.observers.
       this.handlers.complete();
     }
   }
@@ -160,14 +155,15 @@ export class Concast<T> extends Observable<T> {
 
   // Name and argument of the most recently invoked observer method, used
   // to deliver latest results immediately to new observers.
-  private latest?: ["next" | "error", any];
+  private latest?: ["next", T] | ["error", any];
 
   // Bound handler functions that can be reused for every internal
   // subscription.
   private handlers = {
     next: (result: T) => {
       if (this.sub !== null) {
         this.latest = ["next", result];
+        this.notify("next", result);
         iterateObserversSafely(this.observers, "next", result);
       }
     },
@@ -182,6 +178,7 @@ export class Concast<T> extends Observable<T> {
         this.sub = null;
         this.latest = ["error", error];
         this.reject(error);
+        this.notify("error", error);
         iterateObserversSafely(this.observers, "error", error);
       }
     },
@@ -199,6 +196,7 @@ export class Concast<T> extends Observable<T> {
           } else {
             this.resolve();
           }
+          this.notify("complete");
           // We do not store this.latest = ["complete"], because doing so
           // discards useful information about the previous next (or
           // error) message. Instead, if new observers subscribe after
@@ -215,29 +213,35 @@ export class Concast<T> extends Observable<T> {
     },
   };
 
-  public cleanup(callback: () => any) {
+  private nextResultListeners = new Set<NextResultListener>();
+
+  private notify(
+    method: Parameters<NextResultListener>[0],
+    arg?: Parameters<NextResultListener>[1],
+  ) {
+    const { nextResultListeners } = this;
+    if (nextResultListeners.size) {
+      // Replacing this.nextResultListeners first ensures it does not grow while
+      // we are iterating over it, potentially leading to infinite loops.
+      this.nextResultListeners = new Set;
+      nextResultListeners.forEach(listener => listener(method, arg));
+    }
+  }
+
+  // We need a way to run callbacks just *before* the next result (or error or
+  // completion) is delivered by this Concast, so we can be sure any code that
+  // runs as a result of delivering that result/error observes the effects of
+  // running the callback(s). It was tempting to reuse the Observer type instead
+  // of introducing NextResultListener, but that messes with the sizing and
+  // maintenance of this.observers, and ends up being more code overall.
+  beforeNext(callback: NextResultListener) {
     let called = false;
-    const once = () => {
+    this.nextResultListeners.add((method, arg) => {
       if (!called) {
         called = true;
-        // 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();
+        callback(method, arg);
       }
-    }
-    const observer = {
-      next: once,
-      error: once,
-      complete: once,
-    };
-    const count = this.addCount;
-    this.addObserver(observer);
-    // Normally addObserver increments this.addCount, but we can "hide"
-    // cleanup observers by restoring this.addCount to its previous value
-    // after adding any cleanup observer.
-    this.addCount = count;
+    });
   }
 
   // A public way to abort observation and broadcast.
@@ -248,6 +252,11 @@ export class Concast<T> extends Observable<T> {
   }
 }
 
+type NextResultListener = (
+  method: "next" | "error" | "complete",
+  arg?: any,
+) => any;
+
 // Necessary because the Concast constructor has a different signature
 // than the Observable constructor.
 fixObservableSubclass(Concast);

src/utilities/observables/__tests__/Concast.ts

@@ -42,7 +42,7 @@ describe("Concast Observable (similar to Behavior Subject in RxJS)", () => {
       second: 0,
     };
 
-    concast.cleanup(() => {
+    concast.beforeNext(() => {
       ++cleanupCounts.first;
     });
 
@@ -58,7 +58,7 @@ describe("Concast Observable (similar to Behavior Subject in RxJS)", () => {
       },
     });
 
-    concast.cleanup(() => {
+    concast.beforeNext(() => {
       ++cleanupCounts.second;
     });
 
@@ -75,4 +75,71 @@ describe("Concast Observable (similar to Behavior Subject in RxJS)", () => {
       resolve();
     }).catch(reject);
   });
+
+  it("concast.beforeNext listeners run before next result/error", () => {
+    const log: Array<number | [string, any?]> = [];
+    let resolve7Promise: undefined | (() => void);
+
+    const concast = new Concast([
+      Observable.of(1, 2),
+
+      new Promise(resolve => setTimeout(resolve, 10)).then(() => {
+        enqueueListener();
+        return Observable.of(3, 4);
+      }),
+
+      Observable.of(5, 6),
+
+      new Promise<void>(resolve => {
+        resolve7Promise = resolve;
+      }).then(() => {
+        enqueueListener();
+        return Observable.of(7);
+      }),
+
+      Observable.of(8, 9),
+    ]);
+
+    function enqueueListener() {
+      concast.beforeNext((method, arg) => {
+        log.push([method, arg]);
+      });
+    }
+
+    const sub = concast.subscribe({
+      next(num) {
+        log.push(num);
+        if (num === 6) {
+          resolve7Promise!();
+        } else if (num === 8) {
+          enqueueListener();
+          // Prevent delivery of final 9 result.
+          sub.unsubscribe();
+        }
+      },
+    });
+
+    enqueueListener();
+
+    return concast.promise.then(lastResult => {
+      expect(lastResult).toBe(8);
+
+      expect(log).toEqual([
+        ["next", 1],
+        1,
+        2,
+        ["next", 3],
+        3,
+        4,
+        5,
+        6,
+        ["next", 7],
+        7,
+        8,
+        ["complete", void 0],
+      ]);
+
+      sub.unsubscribe();
+    });
+  });
 });