[LABS-311] Extract track_new_launcher_id from dl_track_new

chia/data_layer/data_layer_wallet.py

@@ -210,7 +210,7 @@ async def get_launcher_coin_state(self, launcher_id: bytes32, peer: WSChiaConnec
         return coin_states[0]
 
     # This is the entry point for non-owned singletons
-    async def track_new_launcher_id(
+    async def _track_new_launcher_id(
         self,
         launcher_id: bytes32,
         peer: WSChiaConnection,
@@ -863,6 +863,31 @@ async def stop_tracking_singleton(self, launcher_id: bytes32) -> None:
         await self.wallet_state_manager.dl_store.delete_singleton_records_by_launcher_id(launcher_id)
         await self.wallet_state_manager.dl_store.delete_launcher(launcher_id)
 
+    async def track_new_launcher_id(
+        self,
+        launcher_id: bytes32,
+        peer: WSChiaConnection | None = None,
+        spend: CoinSpend | None = None,
+        height: uint32 | None = None,
+    ) -> None:
+        if peer is None:
+            peer_list = self.wallet_state_manager.wallet_node.get_full_node_peers_in_order()
+        else:
+            peer_list = [peer]
+        peer_length = len(peer_list)
+        for i, peer_to_try in enumerate(peer_list):
+            try:
+                await self._track_new_launcher_id(
+                    launcher_id,
+                    peer_to_try,
+                    spend,
+                    height,
+                )
+            except LauncherCoinNotFoundError as e:
+                if i == peer_length - 1:
+                    raise e  # raise the error if we've tried all peers
+                continue  # try some other peers, maybe someone has it
+
     ###########
     # UTILITY #
     ###########

chia/wallet/wallet_rpc_api.py

@@ -14,9 +14,8 @@
 from clvm_tools.binutils import assemble
 
 from chia.consensus.block_rewards import calculate_base_farmer_reward
-from chia.data_layer.data_layer_errors import LauncherCoinNotFoundError
 from chia.data_layer.data_layer_util import DLProof, VerifyProofResponse, dl_verify_proof
-from chia.data_layer.data_layer_wallet import DataLayerWallet, Mirror
+from chia.data_layer.data_layer_wallet import Mirror
 from chia.pools.pool_wallet import PoolWallet
 from chia.pools.pool_wallet_info import FARMING_TO_POOL, PoolState, PoolWalletInfo, create_pool_state
 from chia.protocols.outbound_message import NodeType
@@ -3490,11 +3489,7 @@ async def create_new_dl(
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        try:
-            dl_wallet = self.service.wallet_state_manager.get_dl_wallet()
-        except ValueError:
-            async with self.service.wallet_state_manager.lock:
-                dl_wallet = await DataLayerWallet.create_new_dl_wallet(self.service.wallet_state_manager)
+        dl_wallet = await self.service.wallet_state_manager.get_dl_wallet(create_if_not_found=True)
 
         async with self.service.wallet_state_manager.lock:
             launcher_id = await dl_wallet.generate_new_reporter(
@@ -3512,25 +3507,11 @@ async def dl_track_new(self, request: DLTrackNew) -> Empty:
         """Initialize the DataLayer Wallet (only one can exist)"""
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
-        try:
-            dl_wallet = self.service.wallet_state_manager.get_dl_wallet()
-        except ValueError:
-            async with self.service.wallet_state_manager.lock:
-                dl_wallet = await DataLayerWallet.create_new_dl_wallet(
-                    self.service.wallet_state_manager,
-                )
-        peer_list = self.service.get_full_node_peers_in_order()
-        peer_length = len(peer_list)
-        for i, peer in enumerate(peer_list):
-            try:
-                await dl_wallet.track_new_launcher_id(
-                    request.launcher_id,
-                    peer,
-                )
-            except LauncherCoinNotFoundError as e:
-                if i == peer_length - 1:
-                    raise e  # raise the error if we've tried all peers
-                continue  # try some other peers, maybe someone has it
+
+        await (await self.service.wallet_state_manager.get_dl_wallet(create_if_not_found=True)).track_new_launcher_id(
+            request.launcher_id
+        )
+
         return Empty()
 
     @marshal
@@ -3539,7 +3520,7 @@ async def dl_stop_tracking(self, request: DLStopTracking) -> Empty:
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        dl_wallet = self.service.wallet_state_manager.get_dl_wallet()
+        dl_wallet = await self.service.wallet_state_manager.get_dl_wallet()
         await dl_wallet.stop_tracking_singleton(request.launcher_id)
         return Empty()
 
@@ -3549,7 +3530,7 @@ async def dl_latest_singleton(self, request: DLLatestSingleton) -> DLLatestSingl
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        wallet = self.service.wallet_state_manager.get_dl_wallet()
+        wallet = await self.service.wallet_state_manager.get_dl_wallet()
         record = await wallet.get_latest_singleton(request.launcher_id, request.only_confirmed)
         return DLLatestSingletonResponse(record)
 
@@ -3559,7 +3540,7 @@ async def dl_singletons_by_root(self, request: DLSingletonsByRoot) -> DLSingleto
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        wallet = self.service.wallet_state_manager.get_dl_wallet()
+        wallet = await self.service.wallet_state_manager.get_dl_wallet()
         records = await wallet.get_singletons_by_root(request.launcher_id, request.root)
         return DLSingletonsByRootResponse(records)
 
@@ -3575,7 +3556,7 @@ async def dl_update_root(
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        wallet = self.service.wallet_state_manager.get_dl_wallet()
+        wallet = await self.service.wallet_state_manager.get_dl_wallet()
         async with self.service.wallet_state_manager.lock:
             await wallet.create_update_state_spend(
                 request.launcher_id,
@@ -3604,7 +3585,7 @@ async def dl_update_multiple(
         if self.service.wallet_state_manager is None:
             raise RuntimeError("not initialized")
 
-        wallet = self.service.wallet_state_manager.get_dl_wallet()
+        wallet = await self.service.wallet_state_manager.get_dl_wallet()
         async with self.service.wallet_state_manager.lock:
             # TODO: This method should optionally link the singletons with announcements.
             #       Otherwise spends are vulnerable to signature subtraction.
@@ -3631,7 +3612,7 @@ async def dl_history(self, request: DLHistory) -> DLHistoryResponse:
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        wallet = self.service.wallet_state_manager.get_dl_wallet()
+        wallet = await self.service.wallet_state_manager.get_dl_wallet()
         additional_kwargs = {}
 
         if request.min_generation is not None:
@@ -3650,7 +3631,7 @@ async def dl_owned_singletons(self, request: Empty) -> DLOwnedSingletonsResponse
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        wallet = self.service.wallet_state_manager.get_dl_wallet()
+        wallet = await self.service.wallet_state_manager.get_dl_wallet()
         singletons = await wallet.get_owned_singletons()
 
         return DLOwnedSingletonsResponse(singletons, uint32(len(singletons)))
@@ -3661,7 +3642,7 @@ async def dl_get_mirrors(self, request: DLGetMirrors) -> DLGetMirrorsResponse:
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        wallet = self.service.wallet_state_manager.get_dl_wallet()
+        wallet = await self.service.wallet_state_manager.get_dl_wallet()
         return DLGetMirrorsResponse(await wallet.get_mirrors_for_launcher(request.launcher_id))
 
     @tx_endpoint(push=True)
@@ -3676,7 +3657,7 @@ async def dl_new_mirror(
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        dl_wallet = self.service.wallet_state_manager.get_dl_wallet()
+        dl_wallet = await self.service.wallet_state_manager.get_dl_wallet()
         async with self.service.wallet_state_manager.lock:
             await dl_wallet.create_new_mirror(
                 request.launcher_id,
@@ -3705,7 +3686,7 @@ async def dl_delete_mirror(
         if self.service.wallet_state_manager is None:
             raise ValueError("The wallet service is not currently initialized")
 
-        dl_wallet = self.service.wallet_state_manager.get_dl_wallet()
+        dl_wallet = await self.service.wallet_state_manager.get_dl_wallet()
         async with self.service.wallet_state_manager.lock:
             await dl_wallet.delete_mirror(
                 request.coin_id,

chia/wallet/wallet_state_manager.py

@@ -1700,7 +1700,7 @@ async def _add_coin_states(
                     elif coin_state.created_height is not None:
                         wallet_identifier, coin_data = await self.determine_coin_type(peer, coin_state, fork_height)
                         try:
-                            dl_wallet = self.get_dl_wallet()
+                            dl_wallet = await self.get_dl_wallet()
                         except ValueError:
                             pass
                         else:
@@ -2012,12 +2012,7 @@ async def _add_coin_states(
                                     and inner_puzhash is not None
                                     and (await self.puzzle_store.puzzle_hash_exists(inner_puzhash))
                                 ):
-                                    try:
-                                        dl_wallet = self.get_dl_wallet()
-                                    except ValueError:
-                                        dl_wallet = await DataLayerWallet.create_new_dl_wallet(
-                                            self,
-                                        )
+                                    dl_wallet = await self.get_dl_wallet(create_if_not_found=True)
                                     await dl_wallet.track_new_launcher_id(
                                         child.coin.name(),
                                         peer,
@@ -2591,14 +2586,18 @@ async def convert_puzzle_hash(self, wallet_id: uint32, puzzle_hash: bytes32) ->
 
         return puzzle_hash
 
-    def get_dl_wallet(self) -> DataLayerWallet:
+    async def get_dl_wallet(self, *, create_if_not_found: bool = False) -> DataLayerWallet:
         for wallet in self.wallets.values():
             if wallet.type() == WalletType.DATA_LAYER.value:
                 assert isinstance(wallet, DataLayerWallet), (
                     f"WalletType.DATA_LAYER should be a DataLayerWallet instance got: {type(wallet).__name__}"
                 )
                 return wallet
-        raise ValueError("DataLayerWallet not available")
+        if create_if_not_found:
+            async with self.lock:
+                return await DataLayerWallet.create_new_dl_wallet(self)
+        else:
+            raise ValueError("DataLayerWallet not available")
 
     async def get_or_create_vc_wallet(self) -> VCWallet:
         for _, wallet in self.wallets.items():