Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Optimize update_pool_config() #17969

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Optimize update_pool_config() #17969

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fadad89
optimize PoolWallet.update_pool_config()
arvidn May 3, 2024
96d9692
print timing to update_pool_config
arvidn May 4, 2024
25bdc49
fix monkeypatching in test
arvidn May 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
102 changes: 51 additions & 51 deletions chia/_tests/pools/test_pool_wallet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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]
Expand All @@ -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()
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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()
Expand All @@ -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:
Expand All @@ -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))
66 changes: 47 additions & 19 deletions chia/pools/pool_wallet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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:
"""
Expand Down