diff --git a/chia/_tests/pools/test_pool_wallet.py b/chia/_tests/pools/test_pool_wallet.py index e2d45a9db7ed..34b966ec8fb9 100644 --- a/chia/_tests/pools/test_pool_wallet.py +++ b/chia/_tests/pools/test_pool_wallet.py @@ -1,8 +1,9 @@ from __future__ import annotations +import contextlib from dataclasses import dataclass from pathlib import Path -from typing import Any, List, Optional, cast +from typing import Any, Dict, Iterator, Optional, Union, cast from unittest.mock import MagicMock import pytest @@ -11,6 +12,7 @@ from chia._tests.util.benchmarks import rand_g1, rand_hash from chia.pools.pool_wallet import PoolWallet from chia.types.blockchain_format.sized_bytes import bytes32 +from chia.util.streamable import recurse_jsonify @dataclass @@ -26,16 +28,6 @@ class MockWalletStateManager: root_path: Optional[Path] = None -@dataclass -class MockPoolWalletConfig: - launcher_id: bytes32 - pool_url: str - payout_instructions: str - target_puzzle_hash: bytes32 - p2_singleton_puzzle_hash: bytes32 - owner_public_key: G1Element - - @dataclass class MockPoolState: pool_url: Optional[str] @@ -56,7 +48,7 @@ async def test_update_pool_config_new_config(monkeypatch: Any) -> None: Test that PoolWallet can create a new pool config """ - updated_configs: List[MockPoolWalletConfig] = [] + updated_configs: Dict[str, Any] = {} payout_instructions_ph = rand_hash() launcher_id: bytes32 = rand_hash() p2_singleton_puzzle_hash: bytes32 = rand_hash() @@ -75,19 +67,22 @@ async def test_update_pool_config_new_config(monkeypatch: Any) -> None: ) # No config data - def mock_load_pool_config(root_path: Path) -> List[MockPoolWalletConfig]: - return [] + @contextlib.contextmanager + def mock_lock_and_load_config( + root_path: Path, + filename: Union[str, Path], + fill_missing_services: bool = False, + ) -> Iterator[Dict[str, Any]]: + yield {"pool": {"pool_list": []}} - monkeypatch.setattr("chia.pools.pool_wallet.load_pool_config", mock_load_pool_config) + monkeypatch.setattr("chia.util.config.lock_and_load_config", mock_lock_and_load_config) # Mock pool_config.update_pool_config to capture the updated configs - async def mock_pool_config_update_pool_config( - root_path: Path, pool_config_list: List[MockPoolWalletConfig] - ) -> None: + def mock_save_config(root_path: Path, filename: Union[str, Path], config_data: Any) -> None: nonlocal updated_configs - updated_configs = pool_config_list + updated_configs = config_data - monkeypatch.setattr("chia.pools.pool_wallet.update_pool_config", mock_pool_config_update_pool_config) + monkeypatch.setattr("chia.util.config.save_config", mock_save_config) # Mock PoolWallet.get_current_state to return our canned state async def mock_get_current_state(self: Any) -> Any: @@ -106,13 +101,14 @@ async def mock_get_current_state(self: Any) -> Any: await wallet.update_pool_config() - assert len(updated_configs) == 1 - assert updated_configs[0].launcher_id == launcher_id - assert updated_configs[0].pool_url == pool_url - assert updated_configs[0].payout_instructions == payout_instructions_ph.hex() - assert updated_configs[0].target_puzzle_hash == target_puzzle_hash - assert updated_configs[0].p2_singleton_puzzle_hash == p2_singleton_puzzle_hash - assert updated_configs[0].owner_public_key == owner_pubkey + pools = updated_configs["pool"]["pool_list"] + assert len(pools) == 1 + assert pools[0]["launcher_id"] == recurse_jsonify(launcher_id) + assert pools[0]["pool_url"] == pool_url + assert pools[0]["payout_instructions"] == payout_instructions_ph.hex() + assert pools[0]["target_puzzle_hash"] == recurse_jsonify(target_puzzle_hash) + assert pools[0]["p2_singleton_puzzle_hash"] == recurse_jsonify(p2_singleton_puzzle_hash) + assert pools[0]["owner_public_key"] == recurse_jsonify(bytes(owner_pubkey)) @pytest.mark.anyio @@ -121,7 +117,7 @@ async def test_update_pool_config_existing_payout_instructions(monkeypatch: Any) Test that PoolWallet will retain existing payout_instructions when updating the pool config. """ - updated_configs: List[MockPoolWalletConfig] = [] + updated_configs: Dict[str, Any] = {} payout_instructions_ph = rand_hash() launcher_id: bytes32 = rand_hash() p2_singleton_puzzle_hash: bytes32 = rand_hash() @@ -147,30 +143,33 @@ async def test_update_pool_config_existing_payout_instructions(monkeypatch: Any) existing_target_puzzle_hash: bytes32 = rand_hash() existing_p2_singleton_puzzle_hash: bytes32 = rand_hash() existing_owner_pubkey: G1Element = rand_g1() - existing_config: MockPoolWalletConfig = MockPoolWalletConfig( - launcher_id=existing_launcher_id, - pool_url=existing_pool_url, - payout_instructions=existing_payout_instructions_ph.hex(), - target_puzzle_hash=existing_target_puzzle_hash, - p2_singleton_puzzle_hash=existing_p2_singleton_puzzle_hash, - owner_public_key=existing_owner_pubkey, - ) + existing_config: Dict[str, Any] = { + "launcher_id": recurse_jsonify(existing_launcher_id), + "pool_url": existing_pool_url, + "payout_instructions": existing_payout_instructions_ph.hex(), + "target_puzzle_hash": recurse_jsonify(existing_target_puzzle_hash), + "p2_singleton_puzzle_hash": recurse_jsonify(existing_p2_singleton_puzzle_hash), + "owner_public_key": recurse_jsonify(existing_owner_pubkey), + } # No config data - def mock_load_pool_config(root_path: Path) -> List[MockPoolWalletConfig]: + @contextlib.contextmanager + def mock_lock_and_load_config( + root_path: Path, + filename: Union[str, Path], + fill_missing_services: bool = False, + ) -> Iterator[Dict[str, Any]]: nonlocal existing_config - return [existing_config] + yield {"pool": {"pool_list": [existing_config]}} - monkeypatch.setattr("chia.pools.pool_wallet.load_pool_config", mock_load_pool_config) + monkeypatch.setattr("chia.util.config.lock_and_load_config", mock_lock_and_load_config) # Mock pool_config.update_pool_config to capture the updated configs - async def mock_pool_config_update_pool_config( - root_path: Path, pool_config_list: List[MockPoolWalletConfig] - ) -> None: + def mock_save_config(root_path: Path, filename: Union[str, Path], config_data: Any) -> None: nonlocal updated_configs - updated_configs = pool_config_list + updated_configs = config_data - monkeypatch.setattr("chia.pools.pool_wallet.update_pool_config", mock_pool_config_update_pool_config) + monkeypatch.setattr("chia.util.config.save_config", mock_save_config) # Mock PoolWallet.get_current_state to return our canned state async def mock_get_current_state(self: Any) -> Any: @@ -189,13 +188,14 @@ async def mock_get_current_state(self: Any) -> Any: await wallet.update_pool_config() - assert len(updated_configs) == 1 - assert updated_configs[0].launcher_id == launcher_id - assert updated_configs[0].pool_url == pool_url + pools = updated_configs["pool"]["pool_list"] + assert len(pools) == 1 + assert pools[0]["launcher_id"] == recurse_jsonify(launcher_id) + assert pools[0]["pool_url"] == pool_url # payout_instructions should still point to existing_payout_instructions_ph - assert updated_configs[0].payout_instructions == existing_payout_instructions_ph.hex() + assert pools[0]["payout_instructions"] == existing_payout_instructions_ph.hex() - assert updated_configs[0].target_puzzle_hash == target_puzzle_hash - assert updated_configs[0].p2_singleton_puzzle_hash == p2_singleton_puzzle_hash - assert updated_configs[0].owner_public_key == owner_pubkey + assert pools[0]["target_puzzle_hash"] == recurse_jsonify(target_puzzle_hash) + assert pools[0]["p2_singleton_puzzle_hash"] == recurse_jsonify(p2_singleton_puzzle_hash) + assert pools[0]["owner_public_key"] == recurse_jsonify(bytes(owner_pubkey)) diff --git a/chia/pools/pool_wallet.py b/chia/pools/pool_wallet.py index e81cebd634ee..4be54aa3d5b6 100644 --- a/chia/pools/pool_wallet.py +++ b/chia/pools/pool_wallet.py @@ -8,8 +8,9 @@ from chia_rs import G1Element, G2Element, PrivateKey from typing_extensions import final +import chia from chia.clvm.singleton import SINGLETON_LAUNCHER -from chia.pools.pool_config import PoolWalletConfig, load_pool_config, update_pool_config +from chia.pools.pool_config import PoolWalletConfig from chia.pools.pool_puzzles import ( create_absorb_spend, create_full_puzzle, @@ -236,26 +237,53 @@ async def get_tip(self) -> Tuple[uint32, CoinSpend]: return (await self.wallet_state_manager.pool_store.get_spends_for_wallet(self.wallet_id))[-1] async def update_pool_config(self) -> None: + start_time = time.monotonic() current_state: PoolWalletInfo = await self.get_current_state() - pool_config_list: List[PoolWalletConfig] = load_pool_config(self.wallet_state_manager.root_path) - pool_config_dict: Dict[bytes32, PoolWalletConfig] = {c.launcher_id: c for c in pool_config_list} - existing_config: Optional[PoolWalletConfig] = pool_config_dict.get(current_state.launcher_id, None) - payout_instructions: str = existing_config.payout_instructions if existing_config is not None else "" - - if len(payout_instructions) == 0: - payout_instructions = (await self.standard_wallet.get_new_puzzlehash()).hex() - self.log.info(f"New config entry. Generated payout_instructions puzzle hash: {payout_instructions}") - - new_config: PoolWalletConfig = PoolWalletConfig( - current_state.launcher_id, - current_state.current.pool_url if current_state.current.pool_url else "", - payout_instructions, - current_state.current.target_puzzle_hash, - current_state.p2_singleton_puzzle_hash, - current_state.current.owner_pubkey, + with chia.util.config.lock_and_load_config(self.wallet_state_manager.root_path, "config.yaml") as config: + start_time2 = time.monotonic() + pool_list = config["pool"].get("pool_list", []) + existing_config: int = -1 + for idx, c in enumerate(pool_list): + try: + launcher_id = bytes32.from_hexstr(c["launcher_id"]) + if launcher_id != current_state.launcher_id: + continue + existing_config = idx + break + except Exception as e: + self.log.error(f"Exception loading config: {c} {e}") + continue + + payout_instructions: str = ( + pool_list[existing_config].get("payout_instructions", "") if existing_config >= 0 else "" + ) + + if len(payout_instructions) == 0: + payout_instructions = (await self.standard_wallet.get_new_puzzlehash()).hex() + self.log.info(f"New config entry. Generated payout_instructions puzzle hash: {payout_instructions}") + + new_config: PoolWalletConfig = PoolWalletConfig( + current_state.launcher_id, + current_state.current.pool_url if current_state.current.pool_url else "", + payout_instructions, + current_state.current.target_puzzle_hash, + current_state.p2_singleton_puzzle_hash, + current_state.current.owner_pubkey, + ) + + if existing_config >= 0: + pool_list[existing_config] = new_config.to_json_dict() + else: + pool_list.append(new_config.to_json_dict()) + + config["pool"]["pool_list"] = pool_list + chia.util.config.save_config(self.wallet_state_manager.root_path, "config.yaml", config) + + end_time = time.monotonic() + self.log.info( + f"update_pool_config time: {end_time - start_time2:0.2f}s " + f"(waited for lock: {start_time2-start_time:0.2f})" ) - pool_config_dict[new_config.launcher_id] = new_config - await update_pool_config(self.wallet_state_manager.root_path, list(pool_config_dict.values())) async def apply_state_transition(self, new_state: CoinSpend, block_height: uint32) -> bool: """