Wallet consistancy (#10532)

* use db transaction, -1 in synced up to height, delete unused funcitons

* use transaction info in key-val-store/pool-store

* cat stores

* db lock

* remove unused lock, set synced not always in transaction

* fix store tests

Co-authored-by: wjblanke <wjb98672@gmail.com>

chia/pools/pool_wallet.py

@@ -280,7 +280,7 @@ class PoolWallet:
                 )
             return False
 
-        await self.wallet_state_manager.pool_store.add_spend(self.wallet_id, new_state, block_height)
+        await self.wallet_state_manager.pool_store.add_spend(self.wallet_id, new_state, block_height, True)
         tip_spend = (await self.get_tip())[1]
         self.log.info(f"New PoolWallet singleton tip_coin: {tip_spend} farmed at height {block_height}")
 
@@ -350,12 +350,16 @@ class PoolWallet:
             if spend.coin.name() == launcher_coin_id:
                 launcher_spend = spend
         assert launcher_spend is not None
-        await self.wallet_state_manager.pool_store.add_spend(self.wallet_id, launcher_spend, block_height)
+        await self.wallet_state_manager.pool_store.add_spend(
+            self.wallet_id, launcher_spend, block_height, in_transaction
+        )
         await self.update_pool_config()
 
         p2_puzzle_hash: bytes32 = (await self.get_current_state()).p2_singleton_puzzle_hash
-        await self.wallet_state_manager.add_new_wallet(self, self.wallet_info.id, create_puzzle_hashes=False)
-        await self.wallet_state_manager.add_interested_puzzle_hashes([p2_puzzle_hash], [self.wallet_id], False)
+        await self.wallet_state_manager.add_new_wallet(
+            self, self.wallet_info.id, create_puzzle_hashes=False, in_transaction=in_transaction
+        )
+        await self.wallet_state_manager.add_interested_puzzle_hashes([p2_puzzle_hash], [self.wallet_id], in_transaction)
         return self
 
     @staticmethod

chia/wallet/cat_wallet/cat_wallet.py

@@ -164,7 +164,11 @@ class CATWallet:
 
     @staticmethod
     async def create_wallet_for_cat(
-        wallet_state_manager: Any, wallet: Wallet, limitations_program_hash_hex: str, name=None
+        wallet_state_manager: Any,
+        wallet: Wallet,
+        limitations_program_hash_hex: str,
+        name=None,
+        in_transaction=False,
     ) -> CATWallet:
         self = CATWallet()
         self.cost_of_single_tx = None
@@ -189,12 +193,16 @@ class CATWallet:
         limitations_program_hash = bytes32(hexstr_to_bytes(limitations_program_hash_hex))
         self.cat_info = CATInfo(limitations_program_hash, None)
         info_as_string = bytes(self.cat_info).hex()
-        self.wallet_info = await wallet_state_manager.user_store.create_wallet(name, WalletType.CAT, info_as_string)
+        self.wallet_info = await wallet_state_manager.user_store.create_wallet(
+            name, WalletType.CAT, info_as_string, in_transaction=in_transaction
+        )
         if self.wallet_info is None:
             raise Exception("wallet_info is None")
 
-        self.lineage_store = await CATLineageStore.create(self.wallet_state_manager.db_wrapper, self.get_asset_id())
-        await self.wallet_state_manager.add_new_wallet(self, self.id())
+        self.lineage_store = await CATLineageStore.create(
+            self.wallet_state_manager.db_wrapper, self.get_asset_id(), in_transaction=in_transaction
+        )
+        await self.wallet_state_manager.add_new_wallet(self, self.id(), in_transaction=in_transaction)
         return self
 
     @staticmethod
@@ -220,7 +228,7 @@ class CATWallet:
             self.cat_info = CATInfo(cat_info.limitations_program_hash, cat_info.my_tail)
             self.lineage_store = await CATLineageStore.create(self.wallet_state_manager.db_wrapper, self.get_asset_id())
             for coin_id, lineage in cat_info.lineage_proofs:
-                await self.add_lineage(coin_id, lineage)
+                await self.add_lineage(coin_id, lineage, False)
             await self.save_info(self.cat_info, False)
 
         return self
@@ -310,7 +318,7 @@ class CATWallet:
 
         inner_puzzle = await self.inner_puzzle_for_cat_puzhash(coin.puzzle_hash)
         lineage_proof = LineageProof(coin.parent_coin_info, inner_puzzle.get_tree_hash(), coin.amount)
-        await self.add_lineage(coin.name(), lineage_proof)
+        await self.add_lineage(coin.name(), lineage_proof, True)
 
         lineage = await self.get_lineage_proof_for_coin(coin)
 
@@ -349,7 +357,9 @@ class CATWallet:
             if parent_coin is None:
                 raise ValueError("Error in finding parent")
             await self.add_lineage(
-                coin_name, LineageProof(parent_coin.parent_coin_info, inner_puzzle.get_tree_hash(), parent_coin.amount)
+                coin_name,
+                LineageProof(parent_coin.parent_coin_info, inner_puzzle.get_tree_hash(), parent_coin.amount),
+                True,
             )
         else:
             # The parent is not a CAT which means we need to scrub all of its children from our DB
@@ -778,14 +788,14 @@ class CATWallet:
 
         return tx_list
 
-    async def add_lineage(self, name: bytes32, lineage: Optional[LineageProof]):
+    async def add_lineage(self, name: bytes32, lineage: Optional[LineageProof], in_transaction):
         """
         Lineage proofs are stored as a list of parent coins and the lineage proof you will need if they are the
         parent of the coin you are trying to spend. 'If I'm your parent, here's the info you need to spend yourself'
         """
         self.log.info(f"Adding parent {name}: {lineage}")
         if lineage is not None:
-            await self.lineage_store.add_lineage_proof(name, lineage)
+            await self.lineage_store.add_lineage_proof(name, lineage, in_transaction)
 
     async def remove_lineage(self, name: bytes32):
         self.log.info(f"Removing parent {name} (probably had a non-CAT parent)")

chia/wallet/cat_wallet/lineage_store.py

@@ -22,16 +22,21 @@ class CATLineageStore:
     table_name: str
 
     @classmethod
-    async def create(cls, db_wrapper: DBWrapper, asset_id: str):
+    async def create(cls, db_wrapper: DBWrapper, asset_id: str, in_transaction=False):
         self = cls()
         self.table_name = f"lineage_proofs_{asset_id}"
         self.db_wrapper = db_wrapper
         self.db_connection = self.db_wrapper.db
-        await self.db_connection.execute(
-            (f"CREATE TABLE IF NOT EXISTS {self.table_name}(" " coin_id text PRIMARY KEY," " lineage blob)")
-        )
-
-        await self.db_connection.commit()
+        if not in_transaction:
+            await self.db_wrapper.lock.acquire()
+        try:
+            await self.db_connection.execute(
+                (f"CREATE TABLE IF NOT EXISTS {self.table_name}(" " coin_id text PRIMARY KEY," " lineage blob)")
+            )
+        finally:
+            if not in_transaction:
+                await self.db_connection.commit()
+                self.db_wrapper.lock.release()
         return self
 
     async def close(self):
@@ -42,23 +47,35 @@ class CATLineageStore:
         await cursor.close()
         await self.db_connection.commit()
 
-    async def add_lineage_proof(self, coin_id: bytes32, lineage: LineageProof) -> None:
-        cursor = await self.db_connection.execute(
-            f"INSERT OR REPLACE INTO {self.table_name} VALUES(?, ?)",
-            (coin_id.hex(), bytes(lineage)),
-        )
+    async def add_lineage_proof(self, coin_id: bytes32, lineage: LineageProof, in_transaction) -> None:
+        if not in_transaction:
+            await self.db_wrapper.lock.acquire()
+        try:
+            cursor = await self.db_connection.execute(
+                f"INSERT OR REPLACE INTO {self.table_name} VALUES(?, ?)",
+                (coin_id.hex(), bytes(lineage)),
+            )
 
-        await cursor.close()
-        await self.db_connection.commit()
+            await cursor.close()
+        finally:
+            if not in_transaction:
+                await self.db_connection.commit()
+                self.db_wrapper.lock.release()
 
-    async def remove_lineage_proof(self, coin_id: bytes32) -> None:
-        cursor = await self.db_connection.execute(
-            f"DELETE FROM {self.table_name} WHERE coin_id=?;",
-            (coin_id.hex(),),
-        )
+    async def remove_lineage_proof(self, coin_id: bytes32, in_transaction=True) -> None:
+        if not in_transaction:
+            await self.db_wrapper.lock.acquire()
+        try:
+            cursor = await self.db_connection.execute(
+                f"DELETE FROM {self.table_name} WHERE coin_id=?;",
+                (coin_id.hex(),),
+            )
 
-        await cursor.close()
-        await self.db_connection.commit()
+            await cursor.close()
+        finally:
+            if not in_transaction:
+                await self.db_connection.commit()
+                self.db_wrapper.lock.release()
 
     async def get_lineage_proof(self, coin_id: bytes32) -> Optional[LineageProof]:
 

chia/wallet/key_val_store.py

@@ -46,19 +46,31 @@ class KeyValStore:
 
         return object_type.from_bytes(row[1])
 
-    async def set_object(self, key: str, obj: Any):
+    async def set_object(self, key: str, obj: Any, in_transaction=False):
         """
         Adds object to key val store. Obj MUST support __bytes__ and bytes() methods.
         """
-        async with self.db_wrapper.lock:
+        if not in_transaction:
+            await self.db_wrapper.lock.acquire()
+
+        try:
             cursor = await self.db_connection.execute(
                 "INSERT OR REPLACE INTO key_val_store VALUES(?, ?)",
                 (key, bytes(obj)),
             )
             await cursor.close()
-            await self.db_connection.commit()
+        finally:
+            if not in_transaction:
+                await self.db_connection.commit()
+                self.db_wrapper.lock.release()
 
-    async def remove_object(self, key: str):
-        cursor = await self.db_connection.execute("DELETE FROM key_val_store where key=?", (key,))
-        await cursor.close()
-        await self.db_connection.commit()
+    async def remove_object(self, key: str, in_transaction=False):
+        if not in_transaction:
+            await self.db_wrapper.lock.acquire()
+        try:
+            cursor = await self.db_connection.execute("DELETE FROM key_val_store where key=?", (key,))
+            await cursor.close()
+        finally:
+            if not in_transaction:
+                await self.db_connection.commit()
+                self.db_wrapper.lock.release()

chia/wallet/puzzles/genesis_checkers.py

@@ -72,7 +72,7 @@ class GenesisById(LimitationsProgram):
         origin_id = origin.name()
 
         cat_inner: Program = await wallet.get_new_inner_puzzle()
-        await wallet.add_lineage(origin_id, LineageProof())
+        await wallet.add_lineage(origin_id, LineageProof(), False)
         genesis_coin_checker: Program = cls.construct([Program.to(origin_id)])
 
         minted_cat_puzzle_hash: bytes32 = construct_cat_puzzle(

chia/wallet/puzzles/tails.py

@@ -72,7 +72,7 @@ class GenesisById(LimitationsProgram):
         origin_id = origin.name()
 
         cat_inner: Program = await wallet.get_new_inner_puzzle()
-        await wallet.add_lineage(origin_id, LineageProof())
+        await wallet.add_lineage(origin_id, LineageProof(), False)
         tail: Program = cls.construct([Program.to(origin_id)])
 
         minted_cat_puzzle_hash: bytes32 = construct_cat_puzzle(CAT_MOD, tail.get_tree_hash(), cat_inner).get_tree_hash()

chia/wallet/trade_manager.py

@@ -122,7 +122,7 @@ class TradeManager:
             for tx in tx_records:
                 if TradeStatus(trade.status) == TradeStatus.PENDING_ACCEPT:
                     await self.wallet_state_manager.add_transaction(
-                        dataclasses.replace(tx, confirmed_at_height=height, confirmed=True)
+                        dataclasses.replace(tx, confirmed_at_height=height, confirmed=True), in_transaction=True
                     )
 
             self.log.info(f"Trade with id: {trade.trade_id} confirmed at height: {height}")

chia/wallet/trading/trade_store.py

@@ -205,31 +205,6 @@ class TradeStore:
         await self.add_trade_record(tx, False)
         return True
 
-    async def set_not_sent(self, id: bytes32):
-        """
-        Updates trade sent count to 0.
-        """
-
-        current: Optional[TradeRecord] = await self.get_trade_record(id)
-        if current is None:
-            return None
-
-        tx: TradeRecord = TradeRecord(
-            confirmed_at_index=current.confirmed_at_index,
-            accepted_at_time=current.accepted_at_time,
-            created_at_time=current.created_at_time,
-            is_my_offer=current.is_my_offer,
-            sent=uint32(0),
-            offer=current.offer,
-            taken_offer=current.taken_offer,
-            coins_of_interest=current.coins_of_interest,
-            trade_id=current.trade_id,
-            status=uint32(TradeStatus.PENDING_CONFIRM.value),
-            sent_to=[],
-        )
-
-        await self.add_trade_record(tx, False)
-
     async def get_trades_count(self) -> Tuple[int, int, int]:
         """
         Returns the number of trades in the database broken down by is_my_offer status
@@ -447,10 +422,10 @@ class TradeStore:
         return records
 
     async def rollback_to_block(self, block_index):
-
-        # Delete from storage
-        cursor = await self.db_connection.execute(
-            "DELETE FROM trade_records WHERE confirmed_at_index>?", (block_index,)
-        )
-        await cursor.close()
-        await self.db_connection.commit()
+        async with self.db_wrapper.lock:
+            # Delete from storage
+            cursor = await self.db_connection.execute(
+                "DELETE FROM trade_records WHERE confirmed_at_index>?", (block_index,)
+            )
+            await cursor.close()
+            await self.db_connection.commit()

chia/wallet/wallet_blockchain.py

@@ -184,9 +184,9 @@ class WalletBlockchain(BlockchainInterface):
             return self._peak
         return await self._basic_store.get_object("PEAK_BLOCK", HeaderBlock)
 
-    async def set_finished_sync_up_to(self, height: int):
+    async def set_finished_sync_up_to(self, height: int, in_transaction=False):
         if height > await self.get_finished_sync_up_to():
-            await self._basic_store.set_object("FINISHED_SYNC_UP_TO", uint32(height))
+            await self._basic_store.set_object("FINISHED_SYNC_UP_TO", uint32(height), in_transaction)
             await self.clean_block_records()
 
     async def get_finished_sync_up_to(self):

chia/wallet/wallet_node.py

@@ -97,7 +97,6 @@ class WalletNode:
     node_peaks: Dict[bytes32, Tuple[uint32, bytes32]]
     validation_semaphore: Optional[asyncio.Semaphore]
     local_node_synced: bool
-    new_state_lock: Optional[asyncio.Lock]
 
     def __init__(
         self,
@@ -120,7 +119,6 @@ class WalletNode:
         self.proof_hashes: List = []
         self.state_changed_callback = None
         self.wallet_state_manager = None
-        self.new_state_lock = None
         self.server = None
         self.wsm_close_task = None
         self.sync_task: Optional[asyncio.Task] = None
@@ -591,12 +589,11 @@ class WalletNode:
         # TODO: optimize fetching
         if self.validation_semaphore is None:
             self.validation_semaphore = asyncio.Semaphore(6)
-        if self.new_state_lock is None:
-            self.new_state_lock = asyncio.Lock()
 
         # If there is a fork, we need to ensure that we roll back in trusted mode to properly handle reorgs
         if trusted and fork_height is not None and height is not None and fork_height != height - 1:
             await self.wallet_state_manager.reorg_rollback(fork_height)
+            await self.wallet_state_manager.blockchain.set_finished_sync_up_to(fork_height)
         cache: PeerRequestCache = self.get_cache_for_peer(peer)
         if fork_height is not None:
             cache.clear_after_height(fork_height)
@@ -610,6 +607,7 @@ class WalletNode:
         items = sorted(items_input, key=last_change_height_cs)
 
         async def receive_and_validate(inner_states: List[CoinState], inner_idx_start: int, cs_heights: List[uint32]):
+            assert self.wallet_state_manager is not None
             try:
                 assert self.validation_semaphore is not None
                 async with self.validation_semaphore:
@@ -624,25 +622,37 @@ class WalletNode:
                         if await self.validate_received_state_from_peer(inner_state, peer, cache, fork_height)
                     ]
                     if len(valid_states) > 0:
-                        assert self.new_state_lock is not None
-                        async with self.new_state_lock:
+                        async with self.wallet_state_manager.db_wrapper.lock:
                             self.log.info(
                                 f"new coin state received ({inner_idx_start}-"
                                 f"{inner_idx_start + len(inner_states) - 1}/ {len(items)})"
                             )
                             if self.wallet_state_manager is None:
                                 return
-                            await self.wallet_state_manager.new_coin_state(valid_states, peer, fork_height)
+                            try:
+                                await self.wallet_state_manager.db_wrapper.commit_transaction()
+                                await self.wallet_state_manager.db_wrapper.begin_transaction()
+                                await self.wallet_state_manager.new_coin_state(valid_states, peer, fork_height)
 
-                            if update_finished_height:
-                                if len(cs_heights) == 1:
-                                    # We have processed all past tasks, so we can increase the height safely
-                                    synced_up_to = last_change_height_cs(valid_states[-1]) - 1
-                                else:
-                                    # We know we have processed everything before this min height
-                                    synced_up_to = min(cs_heights)
-                                await self.wallet_state_manager.blockchain.set_finished_sync_up_to(synced_up_to)
+                                if update_finished_height:
+                                    if len(cs_heights) == 1:
+                                        # We have processed all past tasks, so we can increase the height safely
+                                        synced_up_to = last_change_height_cs(valid_states[-1]) - 1
+                                    else:
+                                        # We know we have processed everything before this min height
+                                        synced_up_to = min(cs_heights) - 1
+                                    await self.wallet_state_manager.blockchain.set_finished_sync_up_to(
+                                        synced_up_to, in_transaction=True
+                                    )
+                                await self.wallet_state_manager.db_wrapper.commit_transaction()
 
+                            except Exception as e:
+                                tb = traceback.format_exc()
+                                self.log.error(f"Exception while adding state: {e} {tb}")
+                                await self.wallet_state_manager.db_wrapper.rollback_transaction()
+                                await self.wallet_state_manager.coin_store.rebuild_wallet_cache()
+                                await self.wallet_state_manager.tx_store.rebuild_tx_cache()
+                                await self.wallet_state_manager.pool_store.rebuild_cache()
             except Exception as e:
                 tb = traceback.format_exc()
                 self.log.error(f"Exception while adding state: {e} {tb}")
@@ -661,16 +671,24 @@ class WalletNode:
                 self.log.error(f"Disconnected from peer {peer.peer_node_id} host {peer.peer_host}")
                 return False
             if trusted:
-                try:
-                    self.log.info(f"new coin state received ({idx}-" f"{idx + len(states) - 1}/ {len(items)})")
-                    await self.wallet_state_manager.new_coin_state(states, peer, fork_height)
-                    await self.wallet_state_manager.blockchain.set_finished_sync_up_to(
-                        last_change_height_cs(states[-1]) - 1
-                    )
-                except Exception as e:
-                    tb = traceback.format_exc()
-                    self.log.error(f"Error adding states.. {e} {tb}")
-                    return False
+                async with self.wallet_state_manager.db_wrapper.lock:
+                    try:
+                        self.log.info(f"new coin state received ({idx}-" f"{idx + len(states) - 1}/ {len(items)})")
+                        await self.wallet_state_manager.db_wrapper.commit_transaction()
+                        await self.wallet_state_manager.db_wrapper.begin_transaction()
+                        await self.wallet_state_manager.new_coin_state(states, peer, fork_height)
+                        await self.wallet_state_manager.db_wrapper.commit_transaction()
+                        await self.wallet_state_manager.blockchain.set_finished_sync_up_to(
+                            last_change_height_cs(states[-1]) - 1, in_transaction=True
+                        )
+                    except Exception as e:
+                        await self.wallet_state_manager.db_wrapper.rollback_transaction()
+                        await self.wallet_state_manager.coin_store.rebuild_wallet_cache()
+                        await self.wallet_state_manager.tx_store.rebuild_tx_cache()
+                        await self.wallet_state_manager.pool_store.rebuild_cache()
+                        tb = traceback.format_exc()
+                        self.log.error(f"Error adding states.. {e} {tb}")
+                        return False
             else:
                 while len(concurrent_tasks_cs_heights) >= target_concurrent_tasks:
                     await asyncio.sleep(0.1)

chia/wallet/wallet_pool_store.py

@@ -40,6 +40,7 @@ class WalletPoolStore:
         wallet_id: int,
         spend: CoinSpend,
         height: uint32,
+        in_transaction=False,
     ) -> None:
         """
         Appends (or replaces) entries in the DB. The new list must be at least as long as the existing list, and the
@@ -47,31 +48,38 @@ class WalletPoolStore:
         until db_wrapper.commit() is called. However it is written to the cache, so it can be fetched with
         get_all_state_transitions.
         """
-        if wallet_id not in self._state_transitions_cache:
-            self._state_transitions_cache[wallet_id] = []
-        all_state_transitions: List[Tuple[uint32, CoinSpend]] = self.get_spends_for_wallet(wallet_id)
+        if not in_transaction:
+            await self.db_wrapper.lock.acquire()
+        try:
+            if wallet_id not in self._state_transitions_cache:
+                self._state_transitions_cache[wallet_id] = []
+            all_state_transitions: List[Tuple[uint32, CoinSpend]] = self.get_spends_for_wallet(wallet_id)
 
-        if (height, spend) in all_state_transitions:
-            return
+            if (height, spend) in all_state_transitions:
+                return
 
-        if len(all_state_transitions) > 0:
-            if height < all_state_transitions[-1][0]:
-                raise ValueError("Height cannot go down")
-            if spend.coin.parent_coin_info != all_state_transitions[-1][1].coin.name():
-                raise ValueError("New spend does not extend")
+            if len(all_state_transitions) > 0:
+                if height < all_state_transitions[-1][0]:
+                    raise ValueError("Height cannot go down")
+                if spend.coin.parent_coin_info != all_state_transitions[-1][1].coin.name():
+                    raise ValueError("New spend does not extend")
 
-        all_state_transitions.append((height, spend))
+            all_state_transitions.append((height, spend))
 
-        cursor = await self.db_connection.execute(
-            "INSERT OR REPLACE INTO pool_state_transitions VALUES (?, ?, ?, ?)",
-            (
-                len(all_state_transitions) - 1,
-                wallet_id,
-                height,
-                bytes(spend),
-            ),
-        )
-        await cursor.close()
+            cursor = await self.db_connection.execute(
+                "INSERT OR REPLACE INTO pool_state_transitions VALUES (?, ?, ?, ?)",
+                (
+                    len(all_state_transitions) - 1,
+                    wallet_id,
+                    height,
+                    bytes(spend),
+                ),
+            )
+            await cursor.close()
+        finally:
+            if not in_transaction:
+                await self.db_connection.commit()
+                self.db_wrapper.lock.release()
 
     def get_spends_for_wallet(self, wallet_id: int) -> List[Tuple[uint32, CoinSpend]]:
         """

chia/wallet/wallet_state_manager.py

@@ -598,7 +598,7 @@ class WalletStateManager:
                     "automatically_add_unknown_cats", False
                 ):
                     cat_wallet = await CATWallet.create_wallet_for_cat(
-                        self, self.main_wallet, bytes(tail_hash).hex()[2:]
+                        self, self.main_wallet, bytes(tail_hash).hex()[2:], in_transaction=True
                     )
                     wallet_id = cat_wallet.id()
                     wallet_type = WalletType(cat_wallet.type())
@@ -731,7 +731,7 @@ class WalletStateManager:
                             name=bytes32(token_bytes()),
                             memos=[],
                         )
-                        await self.tx_store.add_transaction_record(tx_record, False)
+                        await self.tx_store.add_transaction_record(tx_record, True)
 
                     children = await self.wallet_node.fetch_children(peer, coin_state.coin.name(), fork_height)
                     assert children is not None
@@ -786,7 +786,7 @@ class WalletStateManager:
                                 memos=[],
                             )
 
-                            await self.tx_store.add_transaction_record(tx_record, False)
+                            await self.tx_store.add_transaction_record(tx_record, True)
                 else:
                     await self.coin_store.set_spent(coin_state.coin.name(), coin_state.spent_height)
                     rem_tx_records: List[TransactionRecord] = []
@@ -870,7 +870,7 @@ class WalletStateManager:
                         child.coin.name(),
                         [launcher_spend],
                         child.spent_height,
-                        False,
+                        True,
                         "pool_wallet",
                     )
                     coin_added = launcher_spend.additions()[0]
@@ -1030,7 +1030,7 @@ class WalletStateManager:
             wallet = self.wallets[wallet_id]
             await wallet.coin_added(coin, height)
 
-        await self.create_more_puzzle_hashes()
+        await self.create_more_puzzle_hashes(in_transaction=True)
         return coin_record_1
 
     async def add_pending_transaction(self, tx_record: TransactionRecord):
@@ -1047,11 +1047,11 @@ class WalletStateManager:
         self.tx_pending_changed()
         self.state_changed("pending_transaction", tx_record.wallet_id)
 
-    async def add_transaction(self, tx_record: TransactionRecord):
+    async def add_transaction(self, tx_record: TransactionRecord, in_transaction=False):
         """
         Called from wallet to add transaction that is not being set to full_node
         """
-        await self.tx_store.add_transaction_record(tx_record, False)
+        await self.tx_store.add_transaction_record(tx_record, in_transaction)
         self.state_changed("pending_transaction", tx_record.wallet_id)
 
     async def remove_from_queue(
@@ -1123,7 +1123,7 @@ class WalletStateManager:
                 if remove:
                     remove_ids.append(wallet_id)
         for wallet_id in remove_ids:
-            await self.user_store.delete_wallet(wallet_id, in_transaction=True)
+            await self.user_store.delete_wallet(wallet_id, in_transaction=False)
             self.wallets.pop(wallet_id)
 
     async def _await_closed(self) -> None:
@@ -1153,10 +1153,10 @@ class WalletStateManager:
                     return wallet
         return None
 
-    async def add_new_wallet(self, wallet: Any, wallet_id: int, create_puzzle_hashes=True):
+    async def add_new_wallet(self, wallet: Any, wallet_id: int, create_puzzle_hashes=True, in_transaction=False):
         self.wallets[uint32(wallet_id)] = wallet
         if create_puzzle_hashes:
-            await self.create_more_puzzle_hashes()
+            await self.create_more_puzzle_hashes(in_transaction=in_transaction)
         self.state_changed("wallet_created")
 
     async def get_spendable_coins_for_wallet(self, wallet_id: int, records=None) -> Set[WalletCoinRecord]:
@@ -1191,9 +1191,6 @@ class WalletStateManager:
         await self.action_store.create_action(name, wallet_id, wallet_type, callback, done, data, in_transaction)
         self.tx_pending_changed()
 
-    async def set_action_done(self, action_id: int):
-        await self.action_store.action_done(action_id)
-
     async def generator_received(self, height: uint32, header_hash: uint32, program: Program):
 
         actions: List[WalletAction] = await self.action_store.get_all_pending_actions()

tests/pools/test_wallet_pool_store.py

@@ -64,15 +64,15 @@ class TestWalletPoolStore:
             assert store.get_spends_for_wallet(0) == []
             assert store.get_spends_for_wallet(1) == []
 
-            await store.add_spend(1, solution_1, 100)
+            await store.add_spend(1, solution_1, 100, True)
             assert store.get_spends_for_wallet(1) == [(100, solution_1)]
 
             # Idempotent
-            await store.add_spend(1, solution_1, 100)
+            await store.add_spend(1, solution_1, 100, True)
             assert store.get_spends_for_wallet(1) == [(100, solution_1)]
 
             with pytest.raises(ValueError):
-                await store.add_spend(1, solution_1, 101)
+                await store.add_spend(1, solution_1, 101, True)
 
             # Rebuild cache, no longer present
             await db_wrapper.rollback_transaction()
@@ -80,18 +80,18 @@ class TestWalletPoolStore:
             assert store.get_spends_for_wallet(1) == []
 
             await store.rebuild_cache()
-            await store.add_spend(1, solution_1, 100)
+            await store.add_spend(1, solution_1, 100, False)
             assert store.get_spends_for_wallet(1) == [(100, solution_1)]
 
             solution_1_alt: CoinSpend = make_child_solution(solution_0_alt)
 
             with pytest.raises(ValueError):
-                await store.add_spend(1, solution_1_alt, 100)
+                await store.add_spend(1, solution_1_alt, 100, False)
 
             assert store.get_spends_for_wallet(1) == [(100, solution_1)]
 
             solution_2: CoinSpend = make_child_solution(solution_1)
-            await store.add_spend(1, solution_2, 100)
+            await store.add_spend(1, solution_2, 100, False)
             await store.rebuild_cache()
             solution_3: CoinSpend = make_child_solution(solution_2)
             await store.add_spend(1, solution_3, 100)