fix: fast refresh issue with hooks

Closes apollographql#5870

config/jest.config.js

@@ -6,6 +6,7 @@ module.exports = {
     '.(ts|tsx)': 'ts-jest',
   },
   globals: {
+    __DEV__: true,
     'ts-jest': {
       diagnostics: true,
       tsconfig: {

scripts/memory/tests.js

@@ -6,6 +6,8 @@ const {
   makeVar,
 } = require("@apollo/client/core");
 
+global.__DEV__ = true;
+
 function itAsync(message, testFn) {
   const start = Date.now();
   let timeout;

src/react/hooks/useSubscription.ts

@@ -1,16 +1,18 @@
-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 }
@@ -37,8 +39,19 @@ export function useSubscription<TData = any, TVariables = OperationVariables>(
   subscriptionData.setOptions(updatedOptions, true);
   subscriptionData.context = context;
 
+  // @ts-expect-error: __DEV__ is a global exposed by react
+  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 = undefined;
+    };
+  }, []);
 
   return subscriptionData.execute(result);
 }

src/react/hooks/utils/useAfterFastRefresh.ts

@@ -0,0 +1,30 @@
+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) {
+  // @ts-expect-error: __DEV__ is a global exposed by react
+  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>,
@@ -70,8 +71,19 @@ export function useBaseQuery<TData = any, TVariables = OperationVariables>(
     ? (result as QueryTuple<TData, TVariables>)[1]
     : (result as QueryResult<TData, TVariables>);
 
+  // @ts-expect-error: __DEV__ is a global exposed by react
+  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 = undefined;
+    }
   }, []);
 
   useEffect(() => queryData.afterExecute({ lazy }), [