fix race condition in nft tab

brave · Nov 25, 2023 · 142e0d1 · 142e0d1
1 parent 9c20c33
commit 142e0d1
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 66 deletions.
diff --git a/Sources/BraveWallet/Crypto/Stores/NFTStore.swift b/Sources/BraveWallet/Crypto/Stores/NFTStore.swift
@@ -153,6 +153,8 @@ public class NFTStore: ObservableObject, WalletObserverStore {
   private var metadataCache: [String: NFTMetadata] = [:]
   /// Spam from SimpleHash in form of `NetworkAssets`
   private var simpleHashSpamNFTs: [NetworkAssets] = []
+  /// All User NFTs that includes visible, hidden, spam
+  private var allUserNFTs: [BraveWallet.BlockchainToken] = []
 
   var isObserving: Bool {
     rpcServiceObserver != nil && keyringServiceObserver != nil && walletServiveObserber != nil
@@ -261,24 +263,6 @@ public class NFTStore: ObservableObject, WalletObserverStore {
       let selectedAccounts = filters.accounts.filter(\.isSelected).map(\.model)
       let selectedNetworks = filters.networks.filter(\.isSelected).map(\.model)
 
-      // user visible NFTs
-      let userVisibleNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: true)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
-      // user hidden NFTs
-      let userHiddenNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: false)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
       // all spam NFTs marked by SimpleHash
       simpleHashSpamNFTs = await walletService.simpleHashSpamNFTs(for: selectedAccounts, on: selectedNetworks)
       let unionedSpamNFTs = computeSpamNFTs(
@@ -287,29 +271,20 @@ public class NFTStore: ObservableObject, WalletObserverStore {
         simpleHashSpamNFTs: simpleHashSpamNFTs
       )
 
-      let allNetworkNFTs = generateAllNFTsInNetworks(
-        userVisibleNFTs: userVisibleNFTs,
-        userHiddenNFTs: userHiddenNFTs,
-        computedSpamNFTs: unionedSpamNFTs
-      )
-      userNFTGroups = buildNFTGroupModels(
+      (userNFTGroups, allUserNFTs) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )
 
-      var allNFTs: [BraveWallet.BlockchainToken] = []
-      for networkAssets in [userVisibleNFTs, userHiddenNFTs, unionedSpamNFTs] {
-        allNFTs.append(contentsOf: networkAssets.flatMap(\.tokens))
-      }
       // if we're not hiding unowned or grouping by account, balance isn't needed
       if filters.isHidingUnownedNFTs || filters.groupBy == .accounts {
         let allAccounts = filters.accounts.map(\.model)
         nftBalancesCache = await withTaskGroup(
           of: [String: [String: Int]].self,
           body: { @MainActor [nftBalancesCache, rpcService] group in
-            for nft in allNFTs { // for each NFT
+            for nft in allUserNFTs { // for each NFT
               guard let networkForNFT = allNetworks.first(where: { $0.chainId == nft.chainId }) else {
                 continue
               }
@@ -343,22 +318,22 @@ public class NFTStore: ObservableObject, WalletObserverStore {
       }
 
       guard !Task.isCancelled else { return }
-      userNFTGroups = buildNFTGroupModels(
+      (userNFTGroups, allUserNFTs) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )
 
       // fetch nft metadata for all NFTs
-      let allMetadata = await rpcService.fetchNFTMetadata(tokens: allNFTs, ipfsApi: ipfsApi)
+      let allMetadata = await rpcService.fetchNFTMetadata(tokens: allUserNFTs, ipfsApi: ipfsApi)
       for (key, value) in allMetadata { // update cached values
         metadataCache[key] = value
       }
       guard !Task.isCancelled else { return }
-      userNFTGroups = buildNFTGroupModels(
+      (userNFTGroups, allUserNFTs) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )
@@ -520,29 +495,60 @@ public class NFTStore: ObservableObject, WalletObserverStore {
 
   private func buildNFTGroupModels(
     groupBy: GroupBy,
-    allUserNFTs: [NetworkAssets],
+    spams: [NetworkAssets],
     selectedAccounts: [BraveWallet.AccountInfo],
     selectedNetworks: [BraveWallet.NetworkInfo]
-  ) -> [NFTGroupViewModel] {
+  ) -> ([NFTGroupViewModel], [BraveWallet.BlockchainToken]) {
+
+    // user visible NFTs
+    let userVisibleNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: true)
+      .map { networkAssets in
+        NetworkAssets(
+          network: networkAssets.network,
+          tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
+          sortOrder: networkAssets.sortOrder
+        )
+      }
+    // user hidden NFTs
+    let userHiddenNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: false)
+      .map { networkAssets in
+        NetworkAssets(
+          network: networkAssets.network,
+          tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
+          sortOrder: networkAssets.sortOrder
+        )
+      }
+
+    let allUserNFTsInNetworks = generateAllNFTsInNetworks(
+      userVisibleNFTs: userVisibleNFTs,
+      userHiddenNFTs: userHiddenNFTs,
+      computedSpamNFTs: spams
+    )
+
+    var allUserNFTs: [BraveWallet.BlockchainToken] = []
+    for networkAssets in [userVisibleNFTs, userHiddenNFTs, spams] {
+      allUserNFTs.append(contentsOf: networkAssets.flatMap(\.tokens))
+    }
+
     let groups: [NFTGroupViewModel]
     switch filters.groupBy {
     case .none:
       let assets = buildNFTAssetViewModels(
         for: .none,
-        allUserNFTs: allUserNFTs
+        allUserNFTs: allUserNFTsInNetworks
       )
-      return [
+      return ([
         .init(
           groupType: .none,
           assets: assets
         )
-      ]
+      ], allUserNFTs)
     case .accounts:
       groups = selectedAccounts.map { account in
         let groupType: AssetGroupType = .account(account)
         let assets = buildNFTAssetViewModels(
           for: .account(account),
-          allUserNFTs: allUserNFTs
+          allUserNFTs: allUserNFTsInNetworks
         )
         return NFTGroupViewModel(
           groupType: groupType,
@@ -554,15 +560,15 @@ public class NFTStore: ObservableObject, WalletObserverStore {
         let groupType: AssetGroupType = .network(network)
         let assets = buildNFTAssetViewModels(
           for: .network(network),
-          allUserNFTs: allUserNFTs
+          allUserNFTs: allUserNFTsInNetworks
         )
         return NFTGroupViewModel(
           groupType: groupType,
           assets: assets
         )
       }
     }
-    return groups
+    return (groups, allUserNFTs)
   }
 
   @MainActor func isNFTDiscoveryEnabled() async -> Bool {
@@ -589,36 +595,16 @@ public class NFTStore: ObservableObject, WalletObserverStore {
       guard let self else { return }
       let selectedAccounts = self.filters.accounts.filter(\.isSelected).map(\.model)
       let selectedNetworks = self.filters.networks.filter(\.isSelected).map(\.model)
-      let userVisibleAssets = self.assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: true)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
-      let userHiddenAssets = self.assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: false)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
+
       let unionedSpamNFTs = computeSpamNFTs(
         selectedNetworks: selectedNetworks,
         selectedAccounts: selectedAccounts,
         simpleHashSpamNFTs: simpleHashSpamNFTs
       )
 
-      let allNetworkNFTs = generateAllNFTsInNetworks(
-        userVisibleNFTs: userVisibleAssets,
-        userHiddenNFTs: userHiddenAssets,
-        computedSpamNFTs: unionedSpamNFTs
-      )
-      userNFTGroups = buildNFTGroupModels(
+      (userNFTGroups, _) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )

diff --git a/Sources/BraveWallet/Extensions/RpcServiceExtensions.swift b/Sources/BraveWallet/Extensions/RpcServiceExtensions.swift
@@ -103,6 +103,8 @@ extension BraveWalletJsonRpcService {
       }
     case .btc:
       completion(nil)
+    case .zec:
+      completion(nil)
     @unknown default:
       completion(nil)
     }

diff --git a/Tests/BraveWalletTests/NFTStoreTests.swift b/Tests/BraveWalletTests/NFTStoreTests.swift
@@ -73,6 +73,8 @@ class NFTStoreTests: XCTestCase {
         completion([.mockFilecoinTestnet])
       case .btc:
         XCTFail("Should not fetch btc network")
+      case .zec:
+        XCTFail("Should not fetch zec network")
       @unknown default:
         XCTFail("Should not fetch unknown network")
       }