fix: improve balance fetch and not reset to 0 (#6743)

salimtb · web-flow · commit c154788e0c78 · 2025-09-29T09:14:39.000Z
## Explanation ### Current State and Problem Previously, the `TokenBalancesController` would update token balances, native balances, and staked balances regardless of whether the values had actually changed. This led to several inefficiencies: 1. **Unnecessary state updates**: The controller would trigger state changes even when balance values remained the same 2. **Performance overhead**: Unchanged balances would still go through the entire update pipeline, causing unnecessary re-renders and processing 3. **Reset behavior**: Balances could be reset to zero and then immediately updated with the same value, creating unnecessary intermediate states ### Solution This PR implements value comparison logic to skip updates when balances haven't actually changed: 1. **Token Balance Updates**: Added comparison between `currentBalance` and `newBalance` before updating the state. Only updates are applied when values differ, preventing unnecessary state mutations. 2. **Native Balance Updates**: Enhanced the filter logic to compare existing balances from `AccountTrackerController` state with new balance values. Only balances that have actually changed are included in the update batch sent to `AccountTrackerController:updateNativeBalances`. 3. **Staked Balance Updates**: Similar optimization for staked balances, comparing current staked balance values with new ones before including them in the update batch sent to `AccountTrackerController:updateStakedBalances`. ### Key Improvements - **Reduced State Mutations**: Prevents unnecessary state updates when balance values are identical - **Better Performance**: Eliminates redundant processing and re-renders caused by unchanged balance updates - **Cleaner State Management**: Avoids intermediate states where balances are reset to zero before being set to the same value - **Batch Optimization**: Only sends actual changes to the AccountTracker, reducing the payload size and processing overhead ### Technical Details The implementation uses `isEqual` comparison for the main state update check and direct value comparison (`currentBalance !== newBalance`) for individual balance updates. This ensures that both the overall state update and individual balance updates are optimized to only occur when necessary. ## References  ## Checklist - [ ] I've updated the test suite for new or updated code as appropriate - [ ] I've updated documentation (JSDoc, Markdown, etc.) for new or updated code as appropriate - [ ] I've communicated my changes to consumers by [updating changelogs for packages I've changed](https://github.com/MetaMask/core/tree/main/docs/contributing.md#updating-changelogs), highlighting breaking changes as necessary - [ ] I've prepared draft pull requests for clients and consumer packages to resolve any breaking changes  --- > [!NOTE] > Avoids unnecessary token/native/staked balance updates by comparing existing values and only updating when changed, reducing state mutations and calls. > > - **Fixed: Balance update churn** > - `TokenBalancesController` > - Initialize missing entries without overwriting existing balances; only set balances when values change. > - Filter `updateNativeBalances`/`updateStakedBalances` calls by diffing against `AccountTrackerController:getState` and only send changed entries. > - Add `AccountTrackerController:getState` to allowed actions. > - `AccountTrackerController` > - Optimize `updateNativeBalances` and `updateStakedBalances` to clone state, compare current values, and call `update` only if changes exist. > - **Tests**: Permit and stub `AccountTrackerController:getState`; update messenger setups accordingly. > - **Changelog**: Note fix for unnecessary balance updates to improve performance. > > <sup>Written by [Cursor Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit 460183d. This will update automatically on new commits. Configure [here](https://cursor.com/dashboard?tab=bugbot).</sup>
diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md
@@ -11,6 +11,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Bump `@metamask/utils` from `^11.8.0` to `^11.8.1` ([#6708](https://github.com/MetaMask/core/pull/6708))
 
+### Fixed
+
+- Fix unnecessary balance updates in `TokenBalancesController` by skipping updates when values haven't changed ([#6743](https://github.com/MetaMask/core/pull/6743))
+  - Prevents unnecessary state mutations for token balances when values are identical
+  - Improves performance by reducing redundant processing and re-renders
+
 ## [77.0.0]
 
 ### Changed
diff --git a/packages/assets-controllers/src/AccountTrackerController.ts b/packages/assets-controllers/src/AccountTrackerController.ts
@@ -702,26 +702,46 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
   updateNativeBalances(
     balances: { address: string; chainId: Hex; balance: Hex }[],
   ) {
-    this.update((state) => {
-      balances.forEach(({ address, chainId, balance }) => {
-        const checksumAddress = toChecksumHexAddress(address);
+    const nextAccountsByChainId = cloneDeep(this.state.accountsByChainId);
+    let hasChanges = false;
 
-        // Ensure the chainId exists in the state
-        if (!state.accountsByChainId[chainId]) {
-          state.accountsByChainId[chainId] = {};
-        }
+    balances.forEach(({ address, chainId, balance }) => {
+      const checksumAddress = toChecksumHexAddress(address);
 
-        // Ensure the address exists for this chain
-        if (!state.accountsByChainId[chainId][checksumAddress]) {
-          state.accountsByChainId[chainId][checksumAddress] = {
-            balance: '0x0',
-          };
-        }
+      // Ensure the chainId exists in the state
+      if (!nextAccountsByChainId[chainId]) {
+        nextAccountsByChainId[chainId] = {};
+        hasChanges = true;
+      }
 
-        // Update the balance
-        state.accountsByChainId[chainId][checksumAddress].balance = balance;
-      });
+      // Check if the address exists for this chain
+      const accountExists = Boolean(
+        nextAccountsByChainId[chainId][checksumAddress],
+      );
+
+      // Ensure the address exists for this chain
+      if (!accountExists) {
+        nextAccountsByChainId[chainId][checksumAddress] = {
+          balance: '0x0',
+        };
+        hasChanges = true;
+      }
+
+      // Only update the balance if it has changed, or if this is a new account
+      const currentBalance =
+        nextAccountsByChainId[chainId][checksumAddress].balance;
+      if (!accountExists || currentBalance !== balance) {
+        nextAccountsByChainId[chainId][checksumAddress].balance = balance;
+        hasChanges = true;
+      }
     });
+
+    // Only call update if there are actual changes
+    if (hasChanges) {
+      this.update((state) => {
+        state.accountsByChainId = nextAccountsByChainId;
+      });
+    }
   }
 
   /**
@@ -738,27 +758,47 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
       stakedBalance: StakedBalance;
     }[],
   ) {
-    this.update((state) => {
-      stakedBalances.forEach(({ address, chainId, stakedBalance }) => {
-        const checksumAddress = toChecksumHexAddress(address);
+    const nextAccountsByChainId = cloneDeep(this.state.accountsByChainId);
+    let hasChanges = false;
 
-        // Ensure the chainId exists in the state
-        if (!state.accountsByChainId[chainId]) {
-          state.accountsByChainId[chainId] = {};
-        }
+    stakedBalances.forEach(({ address, chainId, stakedBalance }) => {
+      const checksumAddress = toChecksumHexAddress(address);
 
-        // Ensure the address exists for this chain
-        if (!state.accountsByChainId[chainId][checksumAddress]) {
-          state.accountsByChainId[chainId][checksumAddress] = {
-            balance: '0x0',
-          };
-        }
+      // Ensure the chainId exists in the state
+      if (!nextAccountsByChainId[chainId]) {
+        nextAccountsByChainId[chainId] = {};
+        hasChanges = true;
+      }
+
+      // Check if the address exists for this chain
+      const accountExists = Boolean(
+        nextAccountsByChainId[chainId][checksumAddress],
+      );
 
-        // Update the staked balance
-        state.accountsByChainId[chainId][checksumAddress].stakedBalance =
+      // Ensure the address exists for this chain
+      if (!accountExists) {
+        nextAccountsByChainId[chainId][checksumAddress] = {
+          balance: '0x0',
+        };
+        hasChanges = true;
+      }
+
+      // Only update the staked balance if it has changed, or if this is a new account
+      const currentStakedBalance =
+        nextAccountsByChainId[chainId][checksumAddress].stakedBalance;
+      if (!accountExists || !isEqual(currentStakedBalance, stakedBalance)) {
+        nextAccountsByChainId[chainId][checksumAddress].stakedBalance =
           stakedBalance;
-      });
+        hasChanges = true;
+      }
     });
+
+    // Only call update if there are actual changes
+    if (hasChanges) {
+      this.update((state) => {
+        state.accountsByChainId = nextAccountsByChainId;
+      });
+    }
   }
 
   #registerMessageHandlers() {
diff --git a/packages/assets-controllers/src/TokenBalancesController.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts
@@ -62,6 +62,7 @@ const setupController = ({
       'TokensController:getState',
       'AccountsController:getSelectedAccount',
       'AccountsController:listAccounts',
+      'AccountTrackerController:getState',
       'AccountTrackerController:updateNativeBalances',
       'AccountTrackerController:updateStakedBalances',
     ],
@@ -111,6 +112,13 @@ const setupController = ({
     jest.fn().mockImplementation(() => tokens),
   );
 
+  messenger.registerActionHandler(
+    'AccountTrackerController:getState',
+    jest.fn().mockImplementation(() => ({
+      accountsByChainId: {},
+    })),
+  );
+
   messenger.registerActionHandler(
     'AccountTrackerController:updateNativeBalances',
     jest.fn(),
diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts
@@ -32,6 +32,7 @@ import { produce } from 'immer';
 import { isEqual } from 'lodash';
 
 import type {
+  AccountTrackerControllerGetStateAction,
   AccountTrackerUpdateNativeBalancesAction,
   AccountTrackerUpdateStakedBalancesAction,
 } from './AccountTrackerController';
@@ -112,6 +113,7 @@ export type AllowedActions =
   | PreferencesControllerGetStateAction
   | AccountsControllerGetSelectedAccountAction
   | AccountsControllerListAccountsAction
+  | AccountTrackerControllerGetStateAction
   | AccountTrackerUpdateNativeBalancesAction
   | AccountTrackerUpdateStakedBalancesAction;
 
@@ -591,43 +593,55 @@ export class TokenBalancesController extends StaticIntervalPollingController<{
 
     const prev = this.state;
     const next = draft(prev, (d) => {
-      // First, initialize all tokens from allTokens state with balance 0
-      // for the accounts and chains we're processing
+      // Initialize account and chain structures if they don't exist, but preserve existing balances
       for (const chainId of targetChains) {
         for (const account of accountsToProcess) {
-          // Initialize tokens from allTokens
+          // Ensure the nested structure exists without overwriting existing balances
+          d.tokenBalances[account] ??= {};
+          d.tokenBalances[account][chainId] ??= {};
+          // Initialize tokens from allTokens only if they don't exist yet
           const chainTokens = this.#allTokens[chainId];
           if (chainTokens?.[account]) {
             Object.values(chainTokens[account]).forEach(
               (token: { address: string }) => {
                 const tokenAddress = checksum(token.address);
-                ((d.tokenBalances[account] ??= {})[chainId] ??= {})[
-                  tokenAddress
-                ] = '0x0';
+                // Only initialize if the token balance doesn't exist yet
+                if (!(tokenAddress in d.tokenBalances[account][chainId])) {
+                  d.tokenBalances[account][chainId][tokenAddress] = '0x0';
+                }
               },
             );
           }
 
-          // Initialize tokens from allDetectedTokens
+          // Initialize tokens from allDetectedTokens only if they don't exist yet
           const detectedChainTokens = this.#detectedTokens[chainId];
           if (detectedChainTokens?.[account]) {
             Object.values(detectedChainTokens[account]).forEach(
               (token: { address: string }) => {
                 const tokenAddress = checksum(token.address);
-                ((d.tokenBalances[account] ??= {})[chainId] ??= {})[
-                  tokenAddress
-                ] = '0x0';
+                // Only initialize if the token balance doesn't exist yet
+                if (!(tokenAddress in d.tokenBalances[account][chainId])) {
+                  d.tokenBalances[account][chainId][tokenAddress] = '0x0';
+                }
               },
             );
           }
         }
       }
 
-      // Then update with actual fetched balances where available
+      // Update with actual fetched balances only if the value has changed
       aggregated.forEach(({ success, value, account, token, chainId }) => {
         if (success && value !== undefined) {
-          ((d.tokenBalances[account] ??= {})[chainId] ??= {})[checksum(token)] =
-            toHex(value);
+          const newBalance = toHex(value);
+          const tokenAddress = checksum(token);
+          const currentBalance =
+            d.tokenBalances[account]?.[chainId]?.[tokenAddress];
+
+          // Only update if the balance has actually changed
+          if (currentBalance !== newBalance) {
+            ((d.tokenBalances[account] ??= {})[chainId] ??= {})[tokenAddress] =
+              newBalance;
+          }
         }
       });
     });
@@ -639,18 +653,34 @@ export class TokenBalancesController extends StaticIntervalPollingController<{
         (r) => r.success && r.token === ZERO_ADDRESS,
       );
 
-      // Update native token balances in a single batch operation for better performance
+      // Get current AccountTracker state to compare existing balances
+      const accountTrackerState = this.messagingSystem.call(
+        'AccountTrackerController:getState',
+      );
+
+      // Update native token balances only if they have changed
       if (nativeBalances.length > 0) {
-        const balanceUpdates = nativeBalances.map((balance) => ({
-          address: balance.account,
-          chainId: balance.chainId,
-          balance: balance.value ? BNToHex(balance.value) : '0x0',
-        }));
-
-        this.messagingSystem.call(
-          'AccountTrackerController:updateNativeBalances',
-          balanceUpdates,
-        );
+        const balanceUpdates = nativeBalances
+          .map((balance) => ({
+            address: balance.account,
+            chainId: balance.chainId,
+            balance: balance.value ? BNToHex(balance.value) : '0x0',
+          }))
+          .filter((update) => {
+            const currentBalance =
+              accountTrackerState.accountsByChainId[update.chainId]?.[
+                checksum(update.address)
+              ]?.balance;
+            // Only include if the balance has actually changed
+            return currentBalance !== update.balance;
+          });
+
+        if (balanceUpdates.length > 0) {
+          this.messagingSystem.call(
+            'AccountTrackerController:updateNativeBalances',
+            balanceUpdates,
+          );
+        }
       }
 
       // Get staking contract addresses for filtering
@@ -668,16 +698,27 @@ export class TokenBalancesController extends StaticIntervalPollingController<{
       });
 
       if (stakedBalances.length > 0) {
-        const stakedBalanceUpdates = stakedBalances.map((balance) => ({
-          address: balance.account,
-          chainId: balance.chainId,
-          stakedBalance: balance.value ? toHex(balance.value) : '0x0',
-        }));
-
-        this.messagingSystem.call(
-          'AccountTrackerController:updateStakedBalances',
-          stakedBalanceUpdates,
-        );
+        const stakedBalanceUpdates = stakedBalances
+          .map((balance) => ({
+            address: balance.account,
+            chainId: balance.chainId,
+            stakedBalance: balance.value ? toHex(balance.value) : '0x0',
+          }))
+          .filter((update) => {
+            const currentStakedBalance =
+              accountTrackerState.accountsByChainId[update.chainId]?.[
+                checksum(update.address)
+              ]?.stakedBalance;
+            // Only include if the staked balance has actually changed
+            return currentStakedBalance !== update.stakedBalance;
+          });
+
+        if (stakedBalanceUpdates.length > 0) {
+          this.messagingSystem.call(
+            'AccountTrackerController:updateStakedBalances',
+            stakedBalanceUpdates,
+          );
+        }
       }
     }
   }
