Refactor polling hooks to use new hook

Author: MajorLift

ui/hooks/useMultiPolling.ts

@@ -1,6 +1,7 @@
 import { useEffect, useRef } from 'react';
 import { useSelector } from 'react-redux';
 import { getCompletedOnboarding } from '../ducks/metamask/metamask';
+import { useSyncEqualityCheck } from './useSyncEqualityCheck';
 
 type UseMultiPollingOptions<PollingInput> = {
   startPolling: (input: PollingInput) => Promise<string>;
@@ -16,28 +17,27 @@ const useMultiPolling = <PollingInput>(
 ) => {
   const completedOnboarding = useSelector(getCompletedOnboarding);
   const pollingTokens = useRef<Map<string, string>>(new Map());
+  const pollingInputs = useSyncEqualityCheck(usePollingOptions.input);
 
-  const prevPollingInputStringified = useRef<string | null>(null);
-  const hasPollingInputChanged =
-    JSON.stringify(usePollingOptions.input) !==
-    prevPollingInputStringified.current;
-
-  const isMounted = useRef(false);
+  const isMounted = useRef(true);
   useEffect(() => {
-    isMounted.current = true;
     return () => {
       isMounted.current = false;
+      // Stop all polling on unmount - access refs at cleanup time
+      for (const token of pollingTokens.current.values()) {
+        usePollingOptions.stopPollingByPollingToken(token);
+      }
     };
   }, []);
 
   useEffect(() => {
-    if (!completedOnboarding || !hasPollingInputChanged) {
-      // don't start polling if no selected account, or onboarding is incomplete, or polling inputs haven't changed
-      return () => undefined;
+    if (!completedOnboarding) {
+      // don't start polling if onboarding is incomplete
+      return;
     }
 
     // start new polls
-    for (const input of usePollingOptions.input) {
+    for (const input of pollingInputs) {
       const key = JSON.stringify(input);
       if (!pollingTokens.current.has(key)) {
         usePollingOptions.startPolling(input).then((token) => {
@@ -50,28 +50,21 @@ const useMultiPolling = <PollingInput>(
 
     // stop existing polls
     for (const [inputKey, token] of pollingTokens.current.entries()) {
-      const exists = usePollingOptions.input.some(
-        (i) => inputKey === JSON.stringify(i),
+      const exists = pollingInputs.some(
+        (i: PollingInput) => inputKey === JSON.stringify(i),
       );
 
       if (!exists) {
         usePollingOptions.stopPollingByPollingToken(token);
         pollingTokens.current.delete(inputKey);
       }
     }
-
-    prevPollingInputStringified.current = JSON.stringify(
-      usePollingOptions.input,
-    );
-
-    // stop all polling on dismount
-    return () => {
-      for (const token of pollingTokens.current.values()) {
-        usePollingOptions.stopPollingByPollingToken(token);
-      }
-      prevPollingInputStringified.current = null;
-    };
-  }, [usePollingOptions, hasPollingInputChanged, completedOnboarding]);
+  }, [
+    pollingInputs,
+    usePollingOptions.startPolling,
+    usePollingOptions.stopPollingByPollingToken,
+    completedOnboarding,
+  ]);
 };
 
 export default useMultiPolling;

ui/hooks/usePolling.ts

@@ -1,4 +1,5 @@
 import { useEffect, useRef } from 'react';
+import { useSyncEqualityCheck } from './useSyncEqualityCheck';
 
 type UsePollingOptions<PollingInput> = {
   callback?: (pollingToken: string) => (pollingToken: string) => void;
@@ -13,58 +14,50 @@ const usePolling = <PollingInput>(
 ) => {
   const pollTokenRef = useRef<null | string>(null);
   const cleanupRef = useRef<null | ((pollingToken: string) => void)>(null);
+  const pollingInput = useSyncEqualityCheck(usePollingOptions.input);
 
-  const prevPollingInputStringified = useRef<string | null>(null);
-  const hasPollingInputChanged =
-    JSON.stringify(usePollingOptions.input) !==
-    prevPollingInputStringified.current;
-
-  const isMounted = useRef(false);
+  const isMounted = useRef(true);
   useEffect(() => {
-    isMounted.current = true;
     return () => {
       isMounted.current = false;
     };
   }, []);
 
   useEffect(() => {
-    if (usePollingOptions.enabled === false || !hasPollingInputChanged) {
-      return () => {
-        // noop
-      };
-    }
-
     const cleanup = () => {
       if (pollTokenRef.current) {
         usePollingOptions.stopPollingByPollingToken(pollTokenRef.current);
         cleanupRef.current?.(pollTokenRef.current);
+        pollTokenRef.current = null;
+        cleanupRef.current = null;
       }
     };
 
-    // Start polling when the component mounts
-    usePollingOptions
-      .startPolling(usePollingOptions.input)
-      .then((pollToken) => {
-        pollTokenRef.current = pollToken;
-        cleanupRef.current = usePollingOptions.callback?.(pollToken) ?? null;
-        if (!isMounted.current) {
-          cleanup();
-        }
-      });
+    if (usePollingOptions.enabled === false) {
+      // Stop polling if it was previously enabled
+      if (pollTokenRef.current) {
+        cleanup();
+      }
+      return () => {
+        // noop
+      };
+    }
 
-    prevPollingInputStringified.current = JSON.stringify(
-      usePollingOptions.input,
-    );
+    // Start polling when the component mounts
+    usePollingOptions.startPolling(pollingInput).then((pollToken) => {
+      pollTokenRef.current = pollToken;
+      cleanupRef.current = usePollingOptions.callback?.(pollToken) ?? null;
+      if (!isMounted.current) {
+        cleanup();
+      }
+    });
 
-    // Return a cleanup function to stop polling when the component unmounts
-    return () => {
-      prevPollingInputStringified.current = null;
-      cleanup();
-    };
+    // Return a cleanup function to stop polling when the component unmounts or dependencies change
+    return cleanup;
   }, [
-    usePollingOptions.input,
-    hasPollingInputChanged,
+    pollingInput,
     usePollingOptions.enabled,
+    usePollingOptions.stopPollingByPollingToken,
   ]);
 };