From 577c68bdd26519f8341fd1188ea4b8aabe357856 Mon Sep 17 00:00:00 2001
From: Jonathan Sun <bakafish@berkeley.edu>
Date: Thu, 27 Apr 2023 18:55:08 +0900
Subject: [PATCH] useSubscription: keep subscription Concast alive if new
 in-flight subscription subscribes at the same time last in-flight
 subscription unsubscribes (#10718)

* Nullify Concast.sub more carefully with rollback in event of last observer unsubscribe immediately followed by new observer subscribe

* Preserve concast.promise when deferring subscription unsubscribing

* add Concast test for added unsub/promise deferral

* Regenerate internal promise within Concast when rolling back deferred unsubscribe instead of delaying resolve/reject

* Improve reject message in new Concast test relating to concast.promise status

* Remove unneeded callback parameter from Concast.deferredUnsubscribe

* Avoid renewing Concast subscription and internalPromise if subscription is already closed at time of unsubscription

* add Concast test for added unsub/promise deferral

* Improve reject message in new Concast test relating to concast.promise status

* Always unsubscribe sub in deferredUnsubscribe if Concast handled an error

* Add changeset

* add Concast test for added unsub/promise deferral

* Improve reject message in new Concast test relating to concast.promise status

* Shift scope of Concast unsubscription fix from Concast to useSubscription

* Slight revision of changeset

* Remove leftover Concast test

* Use let variable instead of second subscription to carryout deferral of useSubscription unsub

---------

Co-authored-by: Jonathan <jonathan.sun@mujin.co.jp>
Co-authored-by: Lenz Weber-Tronic <lorenz.weber-tronic@apollographql.com>
---
 .changeset/sour-meals-compete.md              |  5 ++
 .../hooks/__tests__/useSubscription.test.tsx  | 69 +++++++++++++++++++
 src/react/hooks/useSubscription.ts            | 39 +++++++----
 3 files changed, 101 insertions(+), 12 deletions(-)
 create mode 100644 .changeset/sour-meals-compete.md

diff --git a/.changeset/sour-meals-compete.md b/.changeset/sour-meals-compete.md
new file mode 100644
index 00000000000..f6374e87ef8
--- /dev/null
+++ b/.changeset/sour-meals-compete.md
@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+Delay Concast subscription teardown slightly in `useSubscription` to prevent unexpected Concast teardown when one `useSubscription` hook tears down its in-flight Concast subscription immediately followed by another `useSubscription` hook reusing and subscribing to that same Concast
diff --git a/src/react/hooks/__tests__/useSubscription.test.tsx b/src/react/hooks/__tests__/useSubscription.test.tsx
index af802470c4d..ca84c5f35a9 100644
--- a/src/react/hooks/__tests__/useSubscription.test.tsx
+++ b/src/react/hooks/__tests__/useSubscription.test.tsx
@@ -996,4 +996,73 @@ describe('useSubscription Hook', () => {
       );
     });
   });
+
+  it('should handle simple subscription after old in-flight teardown immediately \
+followed by new in-flight setup', async () => {
+    const subscription = gql`
+      subscription {
+        car {
+          make
+        }
+      }
+    `;
+
+    const results = ['Audi', 'BMW'].map(make => ({
+      result: { data: { car: { make } } },
+    }));
+
+    const link = new MockSubscriptionLink();
+    const client = new ApolloClient({
+      link,
+      cache: new Cache({ addTypename: false })
+    });
+
+    const { result, unmount, rerender } = renderHook(
+      ({ coin }) => {
+        const heads = useSubscription(subscription, {
+          variables: {},
+          skip: coin === 'tails',
+          context: { coin: 'heads' }
+        });
+        const tails = useSubscription(subscription, {
+          variables: {},
+          skip: coin === 'heads',
+          context: { coin: 'tails' }
+        });
+        return { heads, tails };
+      },
+      {
+        initialProps: {
+          coin: 'heads'
+        },
+        wrapper: ({ children }) => (
+          <ApolloProvider client={client}>
+            {children}
+          </ApolloProvider>
+        )
+      },
+    );
+
+    rerender({ coin: 'tails' });
+
+    await new Promise(resolve => setTimeout(() => resolve('wait'), 20));
+
+    link.simulateResult(results[0]);
+
+    await waitFor(() => {
+      expect(result.current.tails.data).toEqual(results[0].result.data);
+    }, { interval: 1 });
+    expect(result.current.heads.data).toBeUndefined();
+
+    rerender({ coin: 'heads' });
+
+    link.simulateResult(results[1]);
+
+    await waitFor(() => {
+      expect(result.current.heads.data).toEqual(results[1].result.data);
+    }, { interval: 1 });
+    expect(result.current.tails.data).toBeUndefined();
+
+    unmount();
+  });
 });
diff --git a/src/react/hooks/useSubscription.ts b/src/react/hooks/useSubscription.ts
index bce2aa555b6..9cd10a3d14d 100644
--- a/src/react/hooks/useSubscription.ts
+++ b/src/react/hooks/useSubscription.ts
@@ -117,8 +117,13 @@ export function useSubscription<TData = any, TVariables extends OperationVariabl
       return;
     }
 
+    let subscriptionStopped = false;
     const subscription = observable.subscribe({
       next(fetchResult) {
+        if (subscriptionStopped) {
+          return;
+        }
+
         const result = {
           loading: false,
           // TODO: fetchResult.data can be null but SubscriptionResult.data
@@ -142,25 +147,35 @@ export function useSubscription<TData = any, TVariables extends OperationVariabl
         }
       },
       error(error) {
-        setResult({
-          loading: false,
-          data: void 0,
-          error,
-          variables: options?.variables,
-        });
-        ref.current.options?.onError?.(error);
+        if (!subscriptionStopped) {
+          setResult({
+            loading: false,
+            data: void 0,
+            error,
+            variables: options?.variables,
+          });
+          ref.current.options?.onError?.(error);
+        };
       },
       complete() {
-        if (ref.current.options?.onComplete) {
-          ref.current.options.onComplete();
-        } else if (ref.current.options?.onSubscriptionComplete) {
-          ref.current.options.onSubscriptionComplete();
+        if (!subscriptionStopped) {
+          if (ref.current.options?.onComplete) {
+            ref.current.options.onComplete();
+          } else if (ref.current.options?.onSubscriptionComplete) {
+            ref.current.options.onSubscriptionComplete();
+          }
         }
       },
     });
 
     return () => {
-      subscription.unsubscribe();
+      // immediately stop receiving subscription values, but do not unsubscribe
+      // until after a short delay in case another useSubscription hook is
+      // reusing the same underlying observable and is about to subscribe
+      subscriptionStopped = true;
+      setTimeout(() => {
+        subscription.unsubscribe();
+      });
     };
   }, [observable]);