wallet: Some cleanup of WalletNode.fetch_puzzle_solution (#14896)

* wallet: Introduce `WalletNode.get_coin_spend_for_coin_state`

* Rename `fetch_puzzle_solution` and `get_coin_spend_for_coin_state`

* Free `fetch_coin_spend` and `fetch_coin_spend_for_coin_state`
This commit is contained in:
dustinface 2023-04-18 23:26:52 +07:00 committed by GitHub
parent 7c719c35c7
commit d9be2004c7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 62 additions and 64 deletions

View File

@ -50,6 +50,7 @@ from chia.wallet.transaction_record import TransactionRecord
from chia.wallet.util.compute_memos import compute_memos
from chia.wallet.util.merkle_utils import _simplify_merkle_proof
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.wallet_sync_utils import fetch_coin_spend, fetch_coin_spend_for_coin_state
from chia.wallet.util.wallet_types import AmountWithPuzzlehash, WalletType
from chia.wallet.wallet import Wallet
from chia.wallet.wallet_coin_record import WalletCoinRecord
@ -217,9 +218,7 @@ class DataLayerWallet:
await self.new_launcher_spend(spend, peer, height)
else:
launcher_state: CoinState = await self.get_launcher_coin_state(launcher_id, peer)
launcher_spend: CoinSpend = await self.wallet_state_manager.wallet_node.fetch_puzzle_solution(
launcher_state.spent_height, launcher_state.coin, peer
)
launcher_spend = await fetch_coin_spend_for_coin_state(launcher_state, peer)
await self.new_launcher_spend(launcher_spend, peer)
async def new_launcher_spend(
@ -279,9 +278,7 @@ class DataLayerWallet:
] = await self.wallet_state_manager.coin_store.get_coin_record(new_singleton.name())
while new_singleton_coin_record is not None and new_singleton_coin_record.spent_block_height > 0:
# We've already synced this before, so we need to sort of force a resync
parent_spend: CoinSpend = await self.wallet_state_manager.wallet_node.fetch_puzzle_solution(
new_singleton_coin_record.spent_block_height, new_singleton, peer
)
parent_spend = await fetch_coin_spend(new_singleton_coin_record.spent_block_height, new_singleton, peer)
await self.singleton_removed(parent_spend, new_singleton_coin_record.spent_block_height)
try:
additions = compute_additions(parent_spend)
@ -829,9 +826,7 @@ class DataLayerWallet:
parent_state: CoinState = (
await self.wallet_state_manager.wallet_node.get_coin_state([coin.parent_coin_info], peer=peer)
)[0]
parent_spend: Optional[CoinSpend] = await self.wallet_state_manager.wallet_node.fetch_puzzle_solution(
height, parent_state.coin, peer
)
parent_spend = await fetch_coin_spend(height, parent_state.coin, peer)
assert parent_spend is not None
launcher_id, urls = get_mirror_info(
parent_spend.puzzle_reveal.to_program(), parent_spend.solution.to_program()

View File

@ -70,6 +70,7 @@ from chia.wallet.util.address_type import AddressType, is_valid_address
from chia.wallet.util.compute_hints import compute_coin_hints
from chia.wallet.util.compute_memos import compute_memos
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.wallet_sync_utils import fetch_coin_spend_for_coin_state
from chia.wallet.util.wallet_types import AmountWithPuzzlehash, WalletType
from chia.wallet.wallet import CHIP_0002_SIGN_MESSAGE_PREFIX, Wallet
from chia.wallet.wallet_coin_record import WalletCoinRecord
@ -296,9 +297,7 @@ class WalletRpcApi:
if parent_coin_state_list is None or len(parent_coin_state_list) < 1:
raise ValueError(f"Parent coin record 0x{coin_state.coin.parent_coin_info.hex()} not found")
parent_coin_state: CoinState = parent_coin_state_list[0]
coin_spend: CoinSpend = await self.service.wallet_state_manager.wallet_node.fetch_puzzle_solution(
parent_coin_state.spent_height, parent_coin_state.coin, peer
)
coin_spend = await fetch_coin_spend_for_coin_state(parent_coin_state, peer)
return coin_spend, coin_state
##########################################################################################
@ -821,9 +820,7 @@ class WalletRpcApi:
[tr.additions[0].parent_coin_info], peer=peer
)
assert len(coin_state_list) == 1
coin_spend: CoinSpend = await self.service.wallet_state_manager.wallet_node.fetch_puzzle_solution(
coin_state_list[0].spent_height, coin_state_list[0].coin, peer
)
coin_spend = await fetch_coin_spend_for_coin_state(coin_state_list[0], peer)
tr = dataclasses.replace(tr, spend_bundle=SpendBundle([coin_spend], G2Element()))
else:
raise ValueError(f"Transaction 0x{transaction_id.hex()} doesn't have any coin spend.")

View File

@ -51,6 +51,7 @@ from chia.wallet.uncurried_puzzle import uncurry_puzzle
from chia.wallet.util.compute_memos import compute_memos
from chia.wallet.util.curry_and_treehash import calculate_hash_of_quoted_mod_hash, curry_and_treehash
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.wallet_sync_utils import fetch_coin_spend_for_coin_state
from chia.wallet.util.wallet_types import AmountWithPuzzlehash, WalletType
from chia.wallet.wallet import Wallet
from chia.wallet.wallet_coin_record import WalletCoinRecord
@ -347,9 +348,7 @@ class CATWallet:
[coin.parent_coin_info], peer=peer
)
assert coin_state[0].coin.name() == coin.parent_coin_info
coin_spend = await self.wallet_state_manager.wallet_node.fetch_puzzle_solution(
coin_state[0].spent_height, coin_state[0].coin, peer
)
coin_spend = await fetch_coin_spend_for_coin_state(coin_state[0], peer)
await self.puzzle_solution_received(coin_spend, parent_coin=coin_state[0].coin)
except Exception as e:
self.log.debug(f"Exception: {e}, traceback: {traceback.format_exc()}")

View File

@ -44,6 +44,7 @@ from chia.wallet.singleton import (
from chia.wallet.transaction_record import TransactionRecord
from chia.wallet.util.compute_memos import compute_memos
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.wallet_sync_utils import fetch_coin_spend
from chia.wallet.util.wallet_types import WalletType
from chia.wallet.wallet import CHIP_0002_SIGN_MESSAGE_PREFIX, Wallet
from chia.wallet.wallet_coin_record import WalletCoinRecord
@ -505,7 +506,7 @@ class DIDWallet:
await self.save_info(did_info)
assert children_state.created_height
parent_spend = await wallet_node.fetch_puzzle_solution(children_state.created_height, parent_coin, peer)
parent_spend = await fetch_coin_spend(uint32(children_state.created_height), parent_coin, peer)
assert parent_spend is not None
parent_innerpuz = get_inner_puzzle_from_singleton(parent_spend.puzzle_reveal.to_program())
assert parent_innerpuz is not None

View File

@ -50,6 +50,7 @@ from chia.wallet.transaction_record import TransactionRecord
from chia.wallet.uncurried_puzzle import uncurry_puzzle
from chia.wallet.util.compute_memos import compute_memos
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.wallet_sync_utils import fetch_coin_spend
from chia.wallet.util.wallet_types import AmountWithPuzzlehash, WalletType
from chia.wallet.wallet import CHIP_0002_SIGN_MESSAGE_PREFIX, Wallet
from chia.wallet.wallet_coin_record import WalletCoinRecord
@ -170,7 +171,7 @@ class NFTWallet:
return
assert coin_states
parent_coin = coin_states[0].coin
cs = await wallet_node.fetch_puzzle_solution(height, parent_coin, peer)
cs = await fetch_coin_spend(height, parent_coin, peer)
assert cs is not None
await self.puzzle_solution_received(cs, peer)

View File

@ -31,6 +31,7 @@ from chia.protocols.wallet_protocol import (
from chia.server.ws_connection import WSChiaConnection
from chia.types.blockchain_format.coin import Coin, hash_coin_ids
from chia.types.blockchain_format.sized_bytes import bytes32
from chia.types.coin_spend import CoinSpend
from chia.types.full_block import FullBlock
from chia.types.header_block import HeaderBlock
from chia.util.ints import uint32
@ -408,3 +409,25 @@ async def fetch_header_blocks_in_range(
assert res_h_blocks is not None
blocks.extend([bl for bl in res_h_blocks.header_blocks if bl.height >= start])
return blocks
async def fetch_coin_spend(height: uint32, coin: Coin, peer: WSChiaConnection) -> CoinSpend:
solution_response = await peer.call_api(
FullNodeAPI.request_puzzle_solution, wallet_protocol.RequestPuzzleSolution(coin.name(), height)
)
if solution_response is None or not isinstance(solution_response, wallet_protocol.RespondPuzzleSolution):
raise PeerRequestException(f"Was not able to obtain solution {solution_response}")
assert solution_response.response.puzzle.get_tree_hash() == coin.puzzle_hash
assert solution_response.response.coin_name == coin.name()
return CoinSpend(
coin,
solution_response.response.puzzle,
solution_response.response.solution,
)
async def fetch_coin_spend_for_coin_state(coin_state: CoinState, peer: WSChiaConnection) -> CoinSpend:
if coin_state.spent_height is None:
raise ValueError("coin_state.coin must be spent coin")
return await fetch_coin_spend(uint32(coin_state.spent_height), coin_state.coin, peer)

View File

@ -31,7 +31,6 @@ from chia.server.server import ChiaServer
from chia.server.ws_connection import WSChiaConnection
from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.sized_bytes import bytes32
from chia.types.coin_spend import CoinSpend
from chia.types.header_block import HeaderBlock
from chia.types.mempool_inclusion_status import MempoolInclusionStatus
from chia.types.spend_bundle import SpendBundle
@ -1556,21 +1555,6 @@ class WalletNode:
peer_request_cache.add_to_blocks_validated(reward_chain_hash, height)
return True
async def fetch_puzzle_solution(self, height: uint32, coin: Coin, peer: WSChiaConnection) -> CoinSpend:
solution_response = await peer.call_api(
FullNodeAPI.request_puzzle_solution, wallet_protocol.RequestPuzzleSolution(coin.name(), height)
)
if solution_response is None or not isinstance(solution_response, wallet_protocol.RespondPuzzleSolution):
raise PeerRequestException(f"Was not able to obtain solution {solution_response}")
assert solution_response.response.puzzle.get_tree_hash() == coin.puzzle_hash
assert solution_response.response.coin_name == coin.name()
return CoinSpend(
coin,
solution_response.response.puzzle,
solution_response.response.solution,
)
async def get_coin_state(
self, coin_names: List[bytes32], peer: WSChiaConnection, fork_height: Optional[uint32] = None
) -> List[CoinState]:

View File

@ -70,7 +70,11 @@ from chia.wallet.uncurried_puzzle import uncurry_puzzle
from chia.wallet.util.address_type import AddressType
from chia.wallet.util.compute_hints import compute_coin_hints
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.wallet_sync_utils import PeerRequestException, last_change_height_cs
from chia.wallet.util.wallet_sync_utils import (
PeerRequestException,
fetch_coin_spend_for_coin_state,
last_change_height_cs,
)
from chia.wallet.util.wallet_types import WalletIdentifier, WalletType
from chia.wallet.wallet import Wallet
from chia.wallet.wallet_blockchain import WalletBlockchain
@ -621,9 +625,7 @@ class WalletStateManager:
parent_coin_state = response[0]
assert parent_coin_state.spent_height == coin_state.created_height
coin_spend: Optional[CoinSpend] = await self.wallet_node.fetch_puzzle_solution(
parent_coin_state.spent_height, parent_coin_state.coin, peer
)
coin_spend = await fetch_coin_spend_for_coin_state(parent_coin_state, peer)
if coin_spend is None:
return None
@ -833,9 +835,7 @@ class WalletStateManager:
async def get_minter_did(self, launcher_coin: Coin, peer: WSChiaConnection) -> Optional[bytes32]:
# Get minter DID
eve_coin = (await self.wallet_node.fetch_children(launcher_coin.name(), peer=peer))[0]
eve_coin_spend: CoinSpend = await self.wallet_node.fetch_puzzle_solution(
eve_coin.spent_height, eve_coin.coin, peer
)
eve_coin_spend = await fetch_coin_spend_for_coin_state(eve_coin, peer)
eve_full_puzzle: Program = Program.from_bytes(bytes(eve_coin_spend.puzzle_reveal))
eve_uncurried_nft: Optional[UncurriedNFT] = UncurriedNFT.uncurry(*eve_full_puzzle.uncurry())
if eve_uncurried_nft is None:
@ -859,9 +859,7 @@ class WalletStateManager:
[launcher_parent[0].coin.parent_coin_info], peer=peer
)
assert did_coin is not None and len(did_coin) == 1 and did_coin[0].spent_height is not None
did_spend: CoinSpend = await self.wallet_node.fetch_puzzle_solution(
did_coin[0].spent_height, did_coin[0].coin, peer
)
did_spend = await fetch_coin_spend_for_coin_state(did_coin[0], peer)
puzzle = Program.from_bytes(bytes(did_spend.puzzle_reveal))
uncurried = uncurry_puzzle(puzzle)
did_curried_args = match_did_puzzle(uncurried.mod, uncurried.args)
@ -1210,9 +1208,7 @@ class WalletStateManager:
curr_coin_state: CoinState = coin_state
while curr_coin_state.spent_height is not None:
cs: CoinSpend = await self.wallet_node.fetch_puzzle_solution(
curr_coin_state.spent_height, curr_coin_state.coin, peer
)
cs = await fetch_coin_spend_for_coin_state(curr_coin_state, peer)
success = await pool_wallet.apply_state_transition(
cs, uint32(curr_coin_state.spent_height)
)
@ -1245,9 +1241,7 @@ class WalletStateManager:
assert len(new_coin_state) == 1
curr_coin_state = new_coin_state[0]
if record.wallet_type == WalletType.DATA_LAYER:
singleton_spend = await self.wallet_node.fetch_puzzle_solution(
coin_state.spent_height, coin_state.coin, peer
)
singleton_spend = await fetch_coin_spend_for_coin_state(coin_state, peer)
dl_wallet = self.get_wallet(id=uint32(record.wallet_id), required_type=DataLayerWallet)
await dl_wallet.singleton_removed(
singleton_spend,
@ -1268,9 +1262,7 @@ class WalletStateManager:
if child.spent_height is None:
# TODO handle spending launcher later block
continue
launcher_spend: Optional[CoinSpend] = await self.wallet_node.fetch_puzzle_solution(
child.spent_height, child.coin, peer
)
launcher_spend = await fetch_coin_spend_for_coin_state(child, peer)
if launcher_spend is None:
continue
try:

View File

@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import functools
import logging
from typing import List, Optional, Set
from unittest.mock import MagicMock
@ -1184,11 +1185,17 @@ class TestWalletSync:
return new_func
def flaky_fetch_puzzle_solution(node, func):
request_puzzle_solution_failure_tested = False
def flaky_request_puzzle_solution(func):
@functools.wraps(func)
async def new_func(*args, **kwargs):
if node.puzzle_solution_flaky:
node.puzzle_solution_flaky = False
raise PeerRequestException()
nonlocal request_puzzle_solution_failure_tested
if not request_puzzle_solution_failure_tested:
request_puzzle_solution_failure_tested = True
# This can just return None if we have `none_response` enabled.
reject = wallet_protocol.RejectPuzzleSolution(bytes32([0] * 32), uint32(0))
return make_msg(ProtocolMessageTypes.reject_puzzle_solution, reject)
else:
return await func(*args, **kwargs)
@ -1224,18 +1231,17 @@ class TestWalletSync:
return new_func
full_node_api.request_puzzle_solution = flaky_request_puzzle_solution(full_node_api.request_puzzle_solution)
for wallet_node, wallet_server in wallets:
wallet_node.coin_state_retry_seconds = 1
request_puzzle_solution_failure_tested = False
wallet_node.coin_state_flaky = True
wallet_node.puzzle_solution_flaky = True
wallet_node.fetch_children_flaky = True
wallet_node.get_timestamp_flaky = True
wallet_node.db_flaky = True
wallet_node.get_coin_state = flaky_get_coin_state(wallet_node, wallet_node.get_coin_state)
wallet_node.fetch_puzzle_solution = flaky_fetch_puzzle_solution(
wallet_node, wallet_node.fetch_puzzle_solution
)
wallet_node.fetch_children = flaky_fetch_children(wallet_node, wallet_node.fetch_children)
wallet_node.get_timestamp_for_height = flaky_get_timestamp(
wallet_node, wallet_node.get_timestamp_for_height
@ -1278,7 +1284,7 @@ class TestWalletSync:
await assert_coin_state_retry()
assert not wallet_node.coin_state_flaky
assert not wallet_node.puzzle_solution_flaky
assert request_puzzle_solution_failure_tested
assert not wallet_node.fetch_children_flaky
assert not wallet_node.get_timestamp_flaky
assert not wallet_node.db_flaky