From 69ea31f509f82c2a2352a3c50561607348800ce2 Mon Sep 17 00:00:00 2001 From: Mariano Sorgente <3069354+mariano54@users.noreply.github.com> Date: Wed, 8 Jun 2022 11:19:23 -0400 Subject: [PATCH] Ms.optimize request additions (#11669) * Don't return all coins in respond_additions * Optimize request_additions significantly by using a cache * Fix bug with respond_additions and add some tests * Test none header hash, and test returning empty response --- chia/full_node/coin_store.py | 10 ++ chia/full_node/full_node_api.py | 31 ++---- tests/wallet/sync/test_wallet_sync.py | 138 +++++++++++++++++++++++++- 3 files changed, 155 insertions(+), 24 deletions(-) diff --git a/chia/full_node/coin_store.py b/chia/full_node/coin_store.py index 56d822215b..5008d3da40 100644 --- a/chia/full_node/coin_store.py +++ b/chia/full_node/coin_store.py @@ -12,6 +12,8 @@ from chia.util.chunks import chunks import time import logging +from chia.util.lru_cache import LRUCache + log = logging.getLogger(__name__) @@ -21,12 +23,14 @@ class CoinStore: """ db_wrapper: DBWrapper2 + coins_added_at_height_cache: LRUCache @classmethod async def create(cls, db_wrapper: DBWrapper2): self = cls() self.db_wrapper = db_wrapper + self.coins_added_at_height_cache = LRUCache(capacity=100) async with self.db_wrapper.write_db() as conn: @@ -197,6 +201,10 @@ class CoinStore: return coins async def get_coins_added_at_height(self, height: uint32) -> List[CoinRecord]: + coins_added: Optional[List[CoinRecord]] = self.coins_added_at_height_cache.get(height) + if coins_added is not None: + return coins_added + async with self.db_wrapper.read_db() as conn: async with conn.execute( "SELECT confirmed_index, spent_index, coinbase, puzzle_hash, " @@ -208,6 +216,7 @@ class CoinStore: for row in rows: coin = self.row_to_coin(row) coins.append(CoinRecord(coin, row[0], row[1], row[2], row[6])) + self.coins_added_at_height_cache.put(height, coins) return coins async def get_coins_removed_at_height(self, height: uint32) -> List[CoinRecord]: @@ -461,6 +470,7 @@ class CoinStore: await conn.execute( "UPDATE coin_record SET spent_index = 0, spent = 0 WHERE spent_index>?", (block_index,) ) + self.coins_added_at_height_cache = LRUCache(self.coins_added_at_height_cache.capacity) return list(coin_changes.values()) # Store CoinRecord in DB diff --git a/chia/full_node/full_node_api.py b/chia/full_node/full_node_api.py index 8ec145c317..86f4a27984 100644 --- a/chia/full_node/full_node_api.py +++ b/chia/full_node/full_node_api.py @@ -1113,26 +1113,11 @@ class FullNodeAPI: if header_hash is None: raise ValueError(f"Block at height {request.height} not found") - block: Optional[FullBlock] = await self.full_node.block_store.get_full_block(header_hash) - - # We lock so that the coin store does not get modified - if ( - block is None - or block.is_transaction_block() is False - or self.full_node.blockchain.height_to_hash(block.height) != request.header_hash - ): - reject = wallet_protocol.RejectAdditionsRequest(request.height, header_hash) - - msg = make_msg(ProtocolMessageTypes.reject_additions_request, reject) - return msg - - assert block is not None and block.foliage_transaction_block is not None - # Note: this might return bad data if there is a reorg in this time - additions = await self.full_node.coin_store.get_coins_added_at_height(block.height) + additions = await self.full_node.coin_store.get_coins_added_at_height(request.height) - if self.full_node.blockchain.height_to_hash(block.height) != request.header_hash: - raise ValueError(f"Block {block.header_hash} no longer in chain") + if self.full_node.blockchain.height_to_hash(request.height) != header_hash: + raise ValueError(f"Block {header_hash} no longer in chain, or invalid header_hash") puzzlehash_coins_map: Dict[bytes32, List[Coin]] = {} for coin_record in additions: @@ -1147,7 +1132,7 @@ class FullNodeAPI: if request.puzzle_hashes is None: for puzzle_hash, coins in puzzlehash_coins_map.items(): coins_map.append((puzzle_hash, coins)) - response = wallet_protocol.RespondAdditions(block.height, block.header_hash, coins_map, None) + response = wallet_protocol.RespondAdditions(request.height, header_hash, coins_map, None) else: # Create addition Merkle set addition_merkle_set = MerkleSet() @@ -1156,12 +1141,13 @@ class FullNodeAPI: addition_merkle_set.add_already_hashed(puzzle) addition_merkle_set.add_already_hashed(hash_coin_ids([c.name() for c in coins])) - assert addition_merkle_set.get_root() == block.foliage_transaction_block.additions_root for puzzle_hash in request.puzzle_hashes: + # This is a proof of inclusion if it's in (result==True), or exclusion of it's not in result, proof = addition_merkle_set.is_included_already_hashed(puzzle_hash) if puzzle_hash in puzzlehash_coins_map: coins_map.append((puzzle_hash, puzzlehash_coins_map[puzzle_hash])) hash_coin_str = hash_coin_ids([c.name() for c in puzzlehash_coins_map[puzzle_hash]]) + # This is a proof of inclusion of all coin ids that have this ph result_2, proof_2 = addition_merkle_set.is_included_already_hashed(hash_coin_str) assert result assert result_2 @@ -1170,9 +1156,8 @@ class FullNodeAPI: coins_map.append((puzzle_hash, [])) assert not result proofs_map.append((puzzle_hash, proof, None)) - response = wallet_protocol.RespondAdditions(block.height, block.header_hash, coins_map, proofs_map) - msg = make_msg(ProtocolMessageTypes.respond_additions, response) - return msg + response = wallet_protocol.RespondAdditions(request.height, header_hash, coins_map, proofs_map) + return make_msg(ProtocolMessageTypes.respond_additions, response) @api_request async def request_removals(self, request: wallet_protocol.RequestRemovals) -> Optional[Message]: diff --git a/tests/wallet/sync/test_wallet_sync.py b/tests/wallet/sync/test_wallet_sync.py index c4549a83a1..db0cdc1b40 100644 --- a/tests/wallet/sync/test_wallet_sync.py +++ b/tests/wallet/sync/test_wallet_sync.py @@ -1,11 +1,19 @@ +from typing import List, Optional + import pytest from colorlog import getLogger +from chia.consensus.block_record import BlockRecord from chia.consensus.block_rewards import calculate_base_farmer_reward, calculate_pool_reward from chia.protocols import full_node_protocol +from chia.protocols.wallet_protocol import RequestAdditions, RespondAdditions, SendTransaction +from chia.server.outbound_message import Message from chia.simulator.simulator_protocol import FarmNewBlockProtocol from chia.types.peer_info import PeerInfo -from chia.util.ints import uint16, uint32 +from chia.util.hash import std_hash +from chia.util.ints import uint16, uint32, uint64 +from chia.wallet.transaction_record import TransactionRecord +from chia.wallet.util.wallet_types import AmountWithPuzzlehash from tests.connection_utils import disconnect_all_and_reconnect from tests.pools.test_pool_rpc import wallet_is_synced from tests.setup_nodes import test_constants @@ -302,3 +310,131 @@ class TestWalletSync: await time_out_assert(60, wallet_is_synced, True, wallet_node, full_node_api) await time_out_assert(20, get_tx_count, 2, wallet_node.wallet_state_manager, 1) await time_out_assert(20, wallet.get_confirmed_balance, funds) + + @pytest.mark.asyncio + async def test_request_additions_errors(self, wallet_node_sim_and_wallet, self_hostname): + full_nodes, wallets = wallet_node_sim_and_wallet + wallet_node, wallet_server = wallets[0] + wallet = wallet_node.wallet_state_manager.main_wallet + ph = await wallet.get_new_puzzlehash() + + full_node_api = full_nodes[0] + await wallet_server.start_client(PeerInfo(self_hostname, uint16(full_node_api.full_node.server._port)), None) + + for i in range(2): + await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph)) + + await time_out_assert(20, wallet_is_synced, True, wallet_node, full_node_api) + + last_block: Optional[BlockRecord] = full_node_api.full_node.blockchain.get_peak() + assert last_block is not None + + # Invalid height + with pytest.raises(ValueError): + await full_node_api.request_additions(RequestAdditions(uint64(100), last_block.header_hash, [ph])) + + # Invalid header hash + with pytest.raises(ValueError): + await full_node_api.request_additions(RequestAdditions(last_block.height, std_hash(b""), [ph])) + + # No results + res1: Optional[Message] = await full_node_api.request_additions( + RequestAdditions(last_block.height, last_block.header_hash, [std_hash(b"")]) + ) + assert res1 is not None + response = RespondAdditions.from_bytes(res1.data) + assert response.height == last_block.height + assert response.header_hash == last_block.header_hash + assert len(response.proofs) == 1 + assert len(response.coins) == 1 + + assert response.proofs[0][0] == std_hash(b"") + assert response.proofs[0][1] is not None + assert response.proofs[0][2] is None + + @pytest.mark.asyncio + async def test_request_additions_success(self, wallet_node_sim_and_wallet, self_hostname): + full_nodes, wallets = wallet_node_sim_and_wallet + wallet_node, wallet_server = wallets[0] + wallet = wallet_node.wallet_state_manager.main_wallet + ph = await wallet.get_new_puzzlehash() + + full_node_api = full_nodes[0] + await wallet_server.start_client(PeerInfo(self_hostname, uint16(full_node_api.full_node.server._port)), None) + + for i in range(2): + await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph)) + + await time_out_assert(20, wallet_is_synced, True, wallet_node, full_node_api) + + payees: List[AmountWithPuzzlehash] = [] + for i in range(10): + payee_ph = await wallet.get_new_puzzlehash() + payees.append({"amount": uint64(i + 100), "puzzlehash": payee_ph, "memos": []}) + payees.append({"amount": uint64(i + 200), "puzzlehash": payee_ph, "memos": []}) + + tx: TransactionRecord = await wallet.generate_signed_transaction(uint64(0), ph, primaries=payees) + await full_node_api.send_transaction(SendTransaction(tx.spend_bundle)) + + await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph)) + + last_block: Optional[BlockRecord] = full_node_api.full_node.blockchain.get_peak() + assert last_block is not None + await time_out_assert(20, wallet_is_synced, True, wallet_node, full_node_api) + res2: Optional[Message] = await full_node_api.request_additions( + RequestAdditions( + last_block.height, + None, + [payees[0]["puzzlehash"], payees[2]["puzzlehash"], std_hash(b"1")], + ) + ) + + assert res2 is not None + response = RespondAdditions.from_bytes(res2.data) + assert response.height == last_block.height + assert response.header_hash == last_block.header_hash + assert len(response.proofs) == 3 + + # First two PHs are included + for i in range(2): + assert response.proofs[i][0] in {payees[j]["puzzlehash"] for j in (0, 2)} + assert response.proofs[i][1] is not None + assert response.proofs[i][2] is not None + + # Third PH is not included + assert response.proofs[2][2] is None + + coin_list_dict = {p: coin_list for p, coin_list in response.coins} + + assert len(coin_list_dict) == 3 + for p, coin_list in coin_list_dict.items(): + if p == std_hash(b"1"): + # this is the one that is not included + assert len(coin_list) == 0 + else: + for coin in coin_list: + assert coin.puzzle_hash == p + # The other ones are included + assert len(coin_list) == 2 + + # None for puzzle hashes returns all coins and no proofs + res3: Optional[Message] = await full_node_api.request_additions( + RequestAdditions(last_block.height, last_block.header_hash, None) + ) + + assert res3 is not None + response = RespondAdditions.from_bytes(res3.data) + assert response.height == last_block.height + assert response.header_hash == last_block.header_hash + assert response.proofs is None + assert len(response.coins) == 12 + assert sum([len(c_list) for _, c_list in response.coins]) == 24 + + # [] for puzzle hashes returns nothing + res4: Optional[Message] = await full_node_api.request_additions( + RequestAdditions(last_block.height, last_block.header_hash, []) + ) + assert res4 is not None + response = RespondAdditions.from_bytes(res4.data) + assert response.proofs == [] + assert len(response.coins) == 0