feat: multi chain polling for token prices

app/scripts/metamask-controller.js

@@ -1007,6 +1007,7 @@ export default class MetamaskController extends EventEmitter {
       state: initState.TokenRatesController,
       messenger: tokenRatesMessenger,
       tokenPricesService: new CodefiTokenPricesServiceV2(),
+      disabled: !this.preferencesController.state.useCurrencyRateCheck,
     });
 
     this.controllerMessenger.subscribe(
@@ -1015,9 +1016,9 @@ export default class MetamaskController extends EventEmitter {
         const { useCurrencyRateCheck: prevUseCurrencyRateCheck } = prevState;
         const { useCurrencyRateCheck: currUseCurrencyRateCheck } = currState;
         if (currUseCurrencyRateCheck && !prevUseCurrencyRateCheck) {
-          this.tokenRatesController.start();
+          this.tokenRatesController.enable();
         } else if (!currUseCurrencyRateCheck && prevUseCurrencyRateCheck) {
-          this.tokenRatesController.stop();
+          this.tokenRatesController.disable();
         }
       }, this.preferencesController.state),
     );
@@ -2590,12 +2591,6 @@ export default class MetamaskController extends EventEmitter {
 
     const preferencesControllerState = this.preferencesController.state;
 
-    const { useCurrencyRateCheck } = preferencesControllerState;
-
-    if (useCurrencyRateCheck) {
-      this.tokenRatesController.start();
-    }
-
     if (this.#isTokenListPollingRequired(preferencesControllerState)) {
       this.tokenListController.start();
     }
@@ -2608,12 +2603,6 @@ export default class MetamaskController extends EventEmitter {
 
     const preferencesControllerState = this.preferencesController.state;
 
-    const { useCurrencyRateCheck } = preferencesControllerState;
-
-    if (useCurrencyRateCheck) {
-      this.tokenRatesController.stop();
-    }
-
     if (this.#isTokenListPollingRequired(preferencesControllerState)) {
       this.tokenListController.stop();
     }
@@ -3250,6 +3239,7 @@ export default class MetamaskController extends EventEmitter {
       backup,
       approvalController,
       phishingController,
+      tokenRatesController,
       // Notification Controllers
       authenticationController,
       userStorageController,
@@ -4016,6 +4006,13 @@ export default class MetamaskController extends EventEmitter {
           currencyRateController,
         ),
 
+      tokenRatesStartPolling:
+        tokenRatesController.startPolling.bind(tokenRatesController),
+      tokenRatesStopPollingByPollingToken:
+        tokenRatesController.stopPollingByPollingToken.bind(
+          tokenRatesController,
+        ),
+
       // GasFeeController
       gasFeeStartPollingByNetworkClientId:
         gasFeeController.startPollingByNetworkClientId.bind(gasFeeController),
@@ -6641,12 +6638,13 @@ export default class MetamaskController extends EventEmitter {
 
   /**
    * A method that is called by the background when all instances of metamask are closed.
-   * Currently used to stop polling in the gasFeeController.
+   * Currently used to stop controller polling.
    */
   onClientClosed() {
     try {
       this.gasFeeController.stopAllPolling();
       this.currencyRateController.stopAllPolling();
+      this.tokenRatesController.stopAllPolling();
       this.appStateController.clearPollingTokens();
     } catch (error) {
       console.error(error);

package.json

@@ -286,7 +286,7 @@
     "@metamask/address-book-controller": "^6.0.0",
     "@metamask/announcement-controller": "^7.0.0",
     "@metamask/approval-controller": "^7.0.0",
-    "@metamask/assets-controllers": "patch:@metamask/assets-controllers@npm%3A41.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-41.0.0-57b3d695bb.patch",
+    "@metamask/assets-controllers": "patch:@metamask/assets-controllers@npm%3A42.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-42.0.0-57b3d695bb.patch",
     "@metamask/base-controller": "^7.0.0",
     "@metamask/bitcoin-wallet-snap": "^0.8.2",
     "@metamask/browser-passworder": "^4.3.0",

ui/contexts/assetPolling.tsx

@@ -0,0 +1,13 @@
+import React, { ReactNode } from 'react';
+import useCurrencyRatePolling from '../hooks/useCurrencyRatePolling';
+import useTokenRatesPolling from '../hooks/useTokenRatesPolling';
+
+// This provider is a step towards making controller polling fully UI based.
+// Eventually, individual UI components will call the use*Polling hooks to
+// poll and return particular data. This polls globally in the meantime.
+export const AssetPollingProvider = ({ children }: { children: ReactNode }) => {
+  useCurrencyRatePolling();
+  useTokenRatesPolling();
+
+  return <>{children}</>;
+};

ui/hooks/useMultiPolling.ts

@@ -0,0 +1,57 @@
+import { useEffect, useState } from 'react';
+
+type UseMultiPollingOptions<PollingInput> = {
+  startPolling: (input: PollingInput) => Promise<string>;
+  stopPollingByPollingToken: (pollingToken: string) => void;
+  input: PollingInput[];
+};
+
+// A hook that manages multiple polling loops of a polling controller.
+// Callers provide an array of inputs, and the hook manages starting
+// and stopping polling loops for each input.
+const useMultiPolling = <PollingInput>(
+  usePollingOptions: UseMultiPollingOptions<PollingInput>,
+) => {
+  const [polls, setPolls] = useState(new Map());
+
+  useEffect(() => {
+    // start new polls
+    for (const input of usePollingOptions.input) {
+      const key = JSON.stringify(input);
+      if (!polls.has(key)) {
+        usePollingOptions
+          .startPolling(input)
+          .then((token) =>
+            setPolls((prevPolls) => new Map(prevPolls).set(key, token)),
+          );
+      }
+    }
+
+    // stop existing polls
+    for (const [inputKey, token] of polls.entries()) {
+      const exists = usePollingOptions.input.some(
+        (i) => inputKey === JSON.stringify(i),
+      );
+
+      if (!exists) {
+        usePollingOptions.stopPollingByPollingToken(token);
+        setPolls((prevPolls) => {
+          const newPolls = new Map(prevPolls);
+          newPolls.delete(inputKey);
+          return newPolls;
+        });
+      }
+    }
+  }, [usePollingOptions.input && JSON.stringify(usePollingOptions.input)]);
+
+  // stop all polling on dismount
+  useEffect(() => {
+    return () => {
+      for (const token of polls.values()) {
+        usePollingOptions.stopPollingByPollingToken(token);
+      }
+    };
+  }, []);
+};
+
+export default useMultiPolling;

ui/hooks/useTokenRatesPolling.ts

@@ -0,0 +1,40 @@
+import { useSelector } from 'react-redux';
+import {
+  getMarketData,
+  getNetworkConfigurationsByChainId,
+  getTokenExchangeRates,
+  getTokensMarketData,
+  getUseCurrencyRateCheck,
+} from '../selectors';
+import {
+  tokenRatesStartPolling,
+  tokenRatesStopPollingByPollingToken,
+} from '../store/actions';
+import useMultiPolling from './useMultiPolling';
+
+const useTokenRatesPolling = ({ chainIds }: { chainIds?: string[] } = {}) => {
+  // Selectors to determine polling input
+  const useCurrencyRateCheck = useSelector(getUseCurrencyRateCheck);
+  const networkConfigurations = useSelector(getNetworkConfigurationsByChainId);
+
+  // Selectors returning state updated by the polling
+  const tokenExchangeRates = useSelector(getTokenExchangeRates);
+  const tokensMarketData = useSelector(getTokensMarketData);
+  const marketData = useSelector(getMarketData);
+
+  useMultiPolling({
+    startPolling: tokenRatesStartPolling,
+    stopPollingByPollingToken: tokenRatesStopPollingByPollingToken,
+    input: useCurrencyRateCheck
+      ? chainIds ?? Object.keys(networkConfigurations)
+      : [],
+  });
+
+  return {
+    tokenExchangeRates,
+    tokensMarketData,
+    marketData,
+  };
+};
+
+export default useTokenRatesPolling;

ui/pages/index.js

@@ -10,7 +10,7 @@ import {
   LegacyMetaMetricsProvider,
 } from '../contexts/metametrics';
 import { MetamaskNotificationsProvider } from '../contexts/metamask-notifications';
-import { CurrencyRateProvider } from '../contexts/currencyRate';
+import { AssetPollingProvider } from '../contexts/assetPolling';
 import ErrorPage from './error';
 import Routes from './routes';
 
@@ -49,11 +49,11 @@ class Index extends PureComponent {
               <LegacyMetaMetricsProvider>
                 <I18nProvider>
                   <LegacyI18nProvider>
-                    <CurrencyRateProvider>
+                    <AssetPollingProvider>
                       <MetamaskNotificationsProvider>
                         <Routes />
                       </MetamaskNotificationsProvider>
-                    </CurrencyRateProvider>
+                    </AssetPollingProvider>
                   </LegacyI18nProvider>
                 </I18nProvider>
               </LegacyMetaMetricsProvider>

ui/selectors/selectors.js

@@ -582,11 +582,27 @@ export const getTokenExchangeRates = (state) => {
   );
 };
 
+/**
+ * Get market data for tokens on the current chain
+ *
+ * @param state
+ * @returns {Record<Hex, import('@metamask/assets-controllers').MarketDataDetails>}
+ */
 export const getTokensMarketData = (state) => {
   const chainId = getCurrentChainId(state);
   return state.metamask.marketData?.[chainId];
 };
 
+/**
+ * Get market data for tokens across all chains
+ *
+ * @param state
+ * @returns {Record<Hex, Record<Hex, import('@metamask/assets-controllers').MarketDataDetails>>}
+ */
+export const getMarketData = (state) => {
+  return state.metamask.marketData;
+};
+
 export function getAddressBook(state) {
   const chainId = getCurrentChainId(state);
   if (!state.metamask.addressBook[chainId]) {

ui/store/actions.ts

@@ -4555,6 +4555,37 @@ export async function currencyRateStopPollingByPollingToken(
   await removePollingTokenFromAppState(pollingToken);
 }
 
+/**
+ * Informs the TokenRatesController that the UI requires
+ * token rate polling for the given chain id.
+ *
+ * @param chainId - The chain id to poll token rates on.
+ * @returns polling token that can be used to stop polling
+ */
+export async function tokenRatesStartPolling(chainId: string): Promise<string> {
+  const pollingToken = await submitRequestToBackground(
+    'tokenRatesStartPolling',
+    [{ chainId }],
+  );
+  await addPollingTokenToAppState(pollingToken);
+  return pollingToken;
+}
+
+/**
+ * Informs the TokenRatesController that the UI no longer
+ * requires token rate polling for the given chain id.
+ *
+ * @param pollingToken -
+ */
+export async function tokenRatesStopPollingByPollingToken(
+  pollingToken: string,
+) {
+  await submitRequestToBackground('tokenRatesStopPollingByPollingToken', [
+    pollingToken,
+  ]);
+  await removePollingTokenFromAppState(pollingToken);
+}
+
 /**
  * Informs the GasFeeController that the UI requires gas fee polling
  *

yarn.lock

@@ -4772,9 +4772,9 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@metamask/assets-controllers@npm:41.0.0":
-  version: 41.0.0
-  resolution: "@metamask/assets-controllers@npm:41.0.0"
+"@metamask/assets-controllers@npm:42.0.0":
+  version: 42.0.0
+  resolution: "@metamask/assets-controllers@npm:42.0.0"
   dependencies:
     "@ethereumjs/util": "npm:^8.1.0"
     "@ethersproject/address": "npm:^5.7.0"
@@ -4806,13 +4806,13 @@ __metadata:
     "@metamask/keyring-controller": ^17.0.0
     "@metamask/network-controller": ^22.0.0
     "@metamask/preferences-controller": ^13.0.0
-  checksum: 10/63f1a9605d692217889511ca161ee614d8e12d7f7233773afb34c4fb6323fad1c29b3a4ee920ef6f84e4b165ffb8764dfd105bdc9bad75084f52a7c876faa4f5
+  checksum: 10/64d2bd43139ee5c19bd665b07212cd5d5dd41b457dedde3b5db31442292c4d064dc015011f5f001bb423683675fb20898ff652e91d2339ad1d21cc45fa93487a
   languageName: node
   linkType: hard
 
-"@metamask/assets-controllers@patch:@metamask/assets-controllers@npm%3A41.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-41.0.0-57b3d695bb.patch":
-  version: 41.0.0
-  resolution: "@metamask/assets-controllers@patch:@metamask/assets-controllers@npm%3A41.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-41.0.0-57b3d695bb.patch::version=41.0.0&hash=e14ff8"
+"@metamask/assets-controllers@patch:@metamask/assets-controllers@npm%3A42.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-42.0.0-57b3d695bb.patch":
+  version: 42.0.0
+  resolution: "@metamask/assets-controllers@patch:@metamask/assets-controllers@npm%3A42.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-42.0.0-57b3d695bb.patch::version=42.0.0&hash=e14ff8"
   dependencies:
     "@ethereumjs/util": "npm:^8.1.0"
     "@ethersproject/address": "npm:^5.7.0"
@@ -4844,7 +4844,7 @@ __metadata:
     "@metamask/keyring-controller": ^17.0.0
     "@metamask/network-controller": ^22.0.0
     "@metamask/preferences-controller": ^13.0.0
-  checksum: 10/f7d609be61f4e952abd78d996a44131941f1fcd476066d007bed5047d1c887d38e9e9cf117eeb963148674fd9ad6ae87c8384bc8a21d4281628aaab1b60ce7a8
+  checksum: 10/9a6727b28f88fd2df3f4b1628dd5d8c2f3e73fd4b9cd090f22d175c2522faa6c6b7e9a93d0ec2b2d123a263c8f4116fbfe97f196b99401b28ac8597f522651eb
   languageName: node
   linkType: hard
 
@@ -26397,7 +26397,7 @@ __metadata:
     "@metamask/announcement-controller": "npm:^7.0.0"
     "@metamask/api-specs": "npm:^0.9.3"
     "@metamask/approval-controller": "npm:^7.0.0"
-    "@metamask/assets-controllers": "patch:@metamask/assets-controllers@npm%3A41.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-41.0.0-57b3d695bb.patch"
+    "@metamask/assets-controllers": "patch:@metamask/assets-controllers@npm%3A42.0.0#~/.yarn/patches/@metamask-assets-controllers-npm-42.0.0-57b3d695bb.patch"
     "@metamask/auto-changelog": "npm:^2.1.0"
     "@metamask/base-controller": "npm:^7.0.0"
     "@metamask/bitcoin-wallet-snap": "npm:^0.8.2"