Fix Fast Refresh/Live Reload interaction with React hooks (#7952)

Closes #5870

CHANGELOG.md

@@ -112,6 +112,9 @@
 - Suppress noisy `Missing cache result fields...` warnings by default unless `setLogVerbosity("debug")` called. <br/>
   [@benjamn](https://github.com/benjamn) in [#8489](https://github.com/apollographql/apollo-client/pull/8489)
 
+- Improve interaction between React hooks and React Fast Refresh in development. <br/>
+  [@andreialecu](https://github.com/andreialecu) in [#7952](https://github.com/apollographql/apollo-client/pull/7952)
+
 ### Documentation
 TBD
 

src/react/hooks/useSubscription.ts

@@ -1,24 +1,26 @@
-import { useContext, useState, useRef, useEffect } from 'react';
+import { useContext, useState, useRef, useEffect, useReducer } from 'react';
 import { DocumentNode } from 'graphql';
 import { TypedDocumentNode } from '@graphql-typed-document-node/core';
 
 import { SubscriptionHookOptions } from '../types/types';
 import { SubscriptionData } from '../data';
 import { OperationVariables } from '../../core';
 import { getApolloContext } from '../context';
+import { useAfterFastRefresh } from './utils/useAfterFastRefresh';
 
 export function useSubscription<TData = any, TVariables = OperationVariables>(
   subscription: DocumentNode | TypedDocumentNode<TData, TVariables>,
   options?: SubscriptionHookOptions<TData, TVariables>
 ) {
+  const [, forceUpdate] = useReducer(x => x + 1, 0);
   const context = useContext(getApolloContext());
   const updatedOptions = options
     ? { ...options, subscription }
     : { subscription };
   const [result, setResult] = useState({
     loading: !updatedOptions.skip,
-    error: undefined,
-    data: undefined
+    error: void 0,
+    data: void 0,
   });
 
   const subscriptionDataRef = useRef<SubscriptionData<TData, TVariables>>();
@@ -37,8 +39,18 @@ export function useSubscription<TData = any, TVariables = OperationVariables>(
   subscriptionData.setOptions(updatedOptions, true);
   subscriptionData.context = context;
 
+  if (__DEV__) {
+    // ensure we run an update after refreshing so that we can resubscribe
+    useAfterFastRefresh(forceUpdate);
+  }
+
   useEffect(() => subscriptionData.afterExecute());
-  useEffect(() => subscriptionData.cleanup.bind(subscriptionData), []);
+  useEffect(() => {
+    return () => {
+      subscriptionData.cleanup();
+      subscriptionDataRef.current = void 0;
+    };
+  }, []);
 
   return subscriptionData.execute(result);
 }

src/react/hooks/utils/useAfterFastRefresh.ts

@@ -0,0 +1,29 @@
+import { useEffect, useRef } from "react";
+
+/**
+ * This hook allows running a function only immediately after a react
+ * fast refresh or live reload.
+ *
+ * Useful in order to ensure that we can reinitialize things that have been
+ * disposed by cleanup functions in `useEffect`.
+ * @param effectFn a function to run immediately after a fast refresh
+ */
+export function useAfterFastRefresh(effectFn: () => unknown) {
+  if (__DEV__) {
+    const didRefresh = useRef(false);
+    useEffect(() => {
+      return () => {
+        // Detect fast refresh, only runs multiple times in fast refresh
+        didRefresh.current = true;
+      };
+    }, []);
+
+    useEffect(() => {
+      if (didRefresh.current === true) {
+        // This block only runs after a fast refresh
+        didRefresh.current = false;
+        effectFn();
+      }
+    }, [])
+  }
+}

src/react/hooks/utils/useBaseQuery.ts

@@ -12,6 +12,7 @@ import { QueryData } from '../../data';
 import { useDeepMemo } from './useDeepMemo';
 import { OperationVariables } from '../../../core';
 import { getApolloContext } from '../../context';
+import { useAfterFastRefresh } from './useAfterFastRefresh';
 
 export function useBaseQuery<TData = any, TVariables = OperationVariables>(
   query: DocumentNode | TypedDocumentNode<TData, TVariables>,
@@ -54,8 +55,8 @@ export function useBaseQuery<TData = any, TVariables = OperationVariables>(
   const memo = {
     options: {
       ...updatedOptions,
-      onError: undefined,
-      onCompleted: undefined
+      onError: void 0,
+      onCompleted: void 0
     } as QueryHookOptions<TData, TVariables>,
     context,
     tick
@@ -70,8 +71,18 @@ export function useBaseQuery<TData = any, TVariables = OperationVariables>(
     ? (result as QueryTuple<TData, TVariables>)[1]
     : (result as QueryResult<TData, TVariables>);
 
+  if (__DEV__) {
+    // ensure we run an update after refreshing so that we reinitialize
+    useAfterFastRefresh(forceUpdate);
+  }
+
   useEffect(() => {
-    return () => queryData.cleanup();
+    return () => {
+      queryData.cleanup();
+      // this effect can run multiple times during a fast-refresh
+      // so make sure we clean up the ref
+      queryDataRef.current = void 0;
+    }
   }, []);
 
   useEffect(() => queryData.afterExecute({ lazy }), [