diff --git a/benchmarks/blockchains.py b/benchmarks/blockchains.py index afcbbd05887d..caad8858161f 100644 --- a/benchmarks/blockchains.py +++ b/benchmarks/blockchains.py @@ -33,25 +33,25 @@ def enable_profiler(profile: bool, name: str) -> Iterator[None]: async def run_test_chain_benchmark() -> None: with TempKeyring() as keychain: - bt = await create_block_tools_async(constants=test_constants, keychain=keychain) - with enable_profiler(True, "load-test-chain"): - start = time.monotonic() - for version in ["", "_hardfork"]: - for count, name in [ - (400, "test_blocks_400_rc5"), - (1000, "test_blocks_1000_rc5"), - (1000, "pre_genesis_empty_slots_1000_blocksrc5"), - (1500, "test_blocks_1500_rc5"), - (10000, "test_blocks_10000_rc5"), - (758 + 320, "test_blocks_long_reorg_rc5"), - (2000, "test_blocks_2000_compact_rc5"), - (10000, "test_blocks_10000_compact_rc5"), - ]: - persistent_blocks(count, f"{name}{version}.db", bt, seed=b"100") - end = time.monotonic() - KeyringWrapper.cleanup_shared_instance() + async with create_block_tools_async(constants=test_constants, keychain=keychain) as bt: + with enable_profiler(True, "load-test-chain"): + start = time.monotonic() + for version in ["", "_hardfork"]: + for count, name in [ + (400, "test_blocks_400_rc5"), + (1000, "test_blocks_1000_rc5"), + (1000, "pre_genesis_empty_slots_1000_blocksrc5"), + (1500, "test_blocks_1500_rc5"), + (10000, "test_blocks_10000_rc5"), + (758 + 320, "test_blocks_long_reorg_rc5"), + (2000, "test_blocks_2000_compact_rc5"), + (10000, "test_blocks_10000_compact_rc5"), + ]: + persistent_blocks(count, f"{name}{version}.db", bt, seed=b"100") + end = time.monotonic() + KeyringWrapper.cleanup_shared_instance() - print(f"time to load test chains: {end - start:.2f}s") + print(f"time to load test chains: {end - start:.2f}s") if __name__ == "__main__": diff --git a/chia/_tests/blockchain/test_blockchain.py b/chia/_tests/blockchain/test_blockchain.py index 3f9903bcc3b2..c9c3a68a2a13 100644 --- a/chia/_tests/blockchain/test_blockchain.py +++ b/chia/_tests/blockchain/test_blockchain.py @@ -618,14 +618,16 @@ async def test_genesis_no_icc(self, empty_blockchain: Blockchain, bt: BlockTools async def do_test_invalid_icc_sub_slot_vdf( self, keychain: Keychain, db_version: int, constants: ConsensusConstants ) -> None: - bt_high_iters = await create_block_tools_async( - constants=constants.replace( - SUB_SLOT_ITERS_STARTING=uint64(2**12), - DIFFICULTY_STARTING=uint64(2**14), - ), - keychain=keychain, - ) - async with create_blockchain(bt_high_iters.constants, db_version) as (bc1, _): + async with ( + create_block_tools_async( + constants=constants.replace( + SUB_SLOT_ITERS_STARTING=uint64(2**12), + DIFFICULTY_STARTING=uint64(2**14), + ), + keychain=keychain, + ) as bt_high_iters, + create_blockchain(bt_high_iters.constants, db_version) as (bc1, _), + ): blocks = bt_high_iters.get_consecutive_blocks(10) for block in blocks: if ( diff --git a/chia/_tests/conftest.py b/chia/_tests/conftest.py index 6ab186fbca80..46ba31307558 100644 --- a/chia/_tests/conftest.py +++ b/chia/_tests/conftest.py @@ -231,12 +231,14 @@ def blockchain_constants(consensus_mode: ConsensusMode) -> ConsensusConstants: @pytest.fixture(scope="session", name="bt") -async def block_tools_fixture(get_keychain, blockchain_constants, anyio_backend, testrun_uid: str) -> BlockTools: +async def block_tools_fixture( + get_keychain, blockchain_constants, anyio_backend, testrun_uid: str +) -> AsyncIterator[BlockTools]: # Note that this causes a lot of CPU and disk traffic - disk, DB, ports, process creation ... - shared_block_tools = await create_block_tools_async( + async with create_block_tools_async( constants=blockchain_constants, keychain=get_keychain, testrun_uid=testrun_uid - ) - return shared_block_tools + ) as shared_block_tools: + yield shared_block_tools # if you have a system that has an unusual hostname for localhost and you want @@ -952,19 +954,20 @@ def get_temp_keyring(): @pytest.fixture(scope="function") async def get_b_tools_1(get_temp_keyring, testrun_uid: str): - return await create_block_tools_async( + async with create_block_tools_async( constants=test_constants_modified, keychain=get_temp_keyring, testrun_uid=testrun_uid - ) + ) as bt: + yield bt @pytest.fixture(scope="function") async def get_b_tools(get_temp_keyring, testrun_uid): - local_b_tools = await create_block_tools_async( + async with create_block_tools_async( constants=test_constants_modified, keychain=get_temp_keyring, testrun_uid=testrun_uid - ) - new_config = local_b_tools._config - local_b_tools.change_config(new_config) - return local_b_tools + ) as local_b_tools: + new_config = local_b_tools._config + local_b_tools.change_config(new_config) + yield local_b_tools @pytest.fixture(scope="function") @@ -1281,23 +1284,27 @@ async def farmer_harvester_2_simulators_zero_bits_plot_filter( ) async with AsyncExitStack() as async_exit_stack: - bt = await create_block_tools_async( - zero_bit_plot_filter_consts, - keychain=get_temp_keyring, - testrun_uid=testrun_uid, + bt = await async_exit_stack.enter_async_context( + create_block_tools_async( + zero_bit_plot_filter_consts, + keychain=get_temp_keyring, + testrun_uid=testrun_uid, + ) ) config_overrides: dict[str, int] = {"full_node.max_sync_wait": 0} bts = [ - await create_block_tools_async( - zero_bit_plot_filter_consts, - keychain=get_temp_keyring, - num_og_plots=0, - num_pool_plots=0, - num_non_keychain_plots=0, - config_overrides=config_overrides, - testrun_uid=testrun_uid, + await async_exit_stack.enter_async_context( + create_block_tools_async( + zero_bit_plot_filter_consts, + keychain=get_temp_keyring, + num_og_plots=0, + num_pool_plots=0, + num_non_keychain_plots=0, + config_overrides=config_overrides, + testrun_uid=testrun_uid, + ) ) for _ in range(2) ] diff --git a/chia/_tests/core/full_node/stores/test_full_node_store.py b/chia/_tests/core/full_node/stores/test_full_node_store.py index e6edbd62a5f4..840c9fc8e04d 100644 --- a/chia/_tests/core/full_node/stores/test_full_node_store.py +++ b/chia/_tests/core/full_node/stores/test_full_node_store.py @@ -38,7 +38,8 @@ async def custom_block_tools(blockchain_constants: ConsensusConstants) -> AsyncI DISCRIMINANT_SIZE_BITS=uint16(32), SUB_SLOT_ITERS_STARTING=uint64(2**12), ) - yield await create_block_tools_async(constants=patched_constants, keychain=keychain) + async with create_block_tools_async(constants=patched_constants, keychain=keychain) as bt: + yield bt @pytest.fixture(scope="function") diff --git a/chia/_tests/core/full_node/test_full_node.py b/chia/_tests/core/full_node/test_full_node.py index 600fc966e556..4f64362390e5 100644 --- a/chia/_tests/core/full_node/test_full_node.py +++ b/chia/_tests/core/full_node/test_full_node.py @@ -2448,31 +2448,30 @@ async def test_invalid_capability_can_connect( @pytest.mark.anyio async def test_node_start_with_existing_blocks(db_version: int) -> None: with TempKeyring(populate=True) as keychain: - block_tools = await create_block_tools_async(keychain=keychain) - - blocks_per_cycle = 5 - expected_height = 0 - - for cycle in range(2): - async with setup_full_node( - consensus_constants=block_tools.constants, - db_name="node_restart_test.db", - self_hostname=block_tools.config["self_hostname"], - local_bt=block_tools, - simulator=True, - db_version=db_version, - reuse_db=True, - ) as service: - simulator_api = service._api - assert isinstance(simulator_api, FullNodeSimulator) - await simulator_api.farm_blocks_to_puzzlehash(count=blocks_per_cycle) - - expected_height += blocks_per_cycle - assert simulator_api.full_node._blockchain is not None - block_record = simulator_api.full_node._blockchain.get_peak() - - assert block_record is not None, f"block_record is None on cycle {cycle + 1}" - assert block_record.height == expected_height, f"wrong height on cycle {cycle + 1}" + async with create_block_tools_async(keychain=keychain) as block_tools: + blocks_per_cycle = 5 + expected_height = 0 + + for cycle in range(2): + async with setup_full_node( + consensus_constants=block_tools.constants, + db_name="node_restart_test.db", + self_hostname=block_tools.config["self_hostname"], + local_bt=block_tools, + simulator=True, + db_version=db_version, + reuse_db=True, + ) as service: + simulator_api = service._api + assert isinstance(simulator_api, FullNodeSimulator) + await simulator_api.farm_blocks_to_puzzlehash(count=blocks_per_cycle) + + expected_height += blocks_per_cycle + assert simulator_api.full_node._blockchain is not None + block_record = simulator_api.full_node._blockchain.get_peak() + + assert block_record is not None, f"block_record is None on cycle {cycle + 1}" + assert block_record.height == expected_height, f"wrong height on cycle {cycle + 1}" @pytest.mark.anyio diff --git a/chia/_tests/farmer_harvester/test_filter_prefix_bits.py b/chia/_tests/farmer_harvester/test_filter_prefix_bits.py index 0c692f7dd56e..7639764abb5c 100644 --- a/chia/_tests/farmer_harvester/test_filter_prefix_bits.py +++ b/chia/_tests/farmer_harvester/test_filter_prefix_bits.py @@ -55,33 +55,33 @@ async def farmer_harvester_with_filter_size_9( async def have_connections() -> bool: return len(await farmer_rpc_cl.get_connections()) > 0 - local_b_tools = await create_block_tools_async( + async with create_block_tools_async( constants=test_constants.replace(NUMBER_ZERO_BITS_PLOT_FILTER_V1=uint8(9)), keychain=get_temp_keyring - ) - new_config = local_b_tools._config - local_b_tools.change_config(new_config) - async with setup_farmer_solver_multi_harvester( - local_b_tools, 1, tmp_path, local_b_tools.constants, start_services=True - ) as (harvesters, farmer_service, _): - harvester_service = harvesters[0] - assert farmer_service.rpc_server is not None - farmer_rpc_cl = await FarmerRpcClient.create( - self_hostname, farmer_service.rpc_server.listen_port, farmer_service.root_path, farmer_service.config - ) - assert harvester_service.rpc_server is not None - harvester_rpc_cl = await HarvesterRpcClient.create( - self_hostname, - harvester_service.rpc_server.listen_port, - harvester_service.root_path, - harvester_service.config, - ) - await time_out_assert(15, have_connections, True) - yield harvester_service, farmer_service._api + ) as local_b_tools: + new_config = local_b_tools._config + local_b_tools.change_config(new_config) + async with setup_farmer_solver_multi_harvester( + local_b_tools, 1, tmp_path, local_b_tools.constants, start_services=True + ) as (harvesters, farmer_service, _): + harvester_service = harvesters[0] + assert farmer_service.rpc_server is not None + farmer_rpc_cl = await FarmerRpcClient.create( + self_hostname, farmer_service.rpc_server.listen_port, farmer_service.root_path, farmer_service.config + ) + assert harvester_service.rpc_server is not None + harvester_rpc_cl = await HarvesterRpcClient.create( + self_hostname, + harvester_service.rpc_server.listen_port, + harvester_service.root_path, + harvester_service.config, + ) + await time_out_assert(15, have_connections, True) + yield harvester_service, farmer_service._api - farmer_rpc_cl.close() - harvester_rpc_cl.close() - await farmer_rpc_cl.await_closed() - await harvester_rpc_cl.await_closed() + farmer_rpc_cl.close() + harvester_rpc_cl.close() + await farmer_rpc_cl.await_closed() + await harvester_rpc_cl.await_closed() @pytest.mark.parametrize(argnames=["peak_height", "eligible_plots"], argvalues=[(5495999, 0), (5496000, 1)]) diff --git a/chia/_tests/simulation/test_simulation.py b/chia/_tests/simulation/test_simulation.py index ad607df2d60d..95bb29e59597 100644 --- a/chia/_tests/simulation/test_simulation.py +++ b/chia/_tests/simulation/test_simulation.py @@ -58,15 +58,15 @@ @pytest.fixture(scope="function") async def extra_node(self_hostname) -> AsyncIterator[FullNodeAPI | FullNodeSimulator]: with TempKeyring() as keychain: - b_tools = await create_block_tools_async(constants=test_constants_modified, keychain=keychain) - async with setup_full_node( - test_constants_modified, - "blockchain_test_3.db", - self_hostname, - b_tools, - db_version=2, - ) as service: - yield service._api + async with create_block_tools_async(constants=test_constants_modified, keychain=keychain) as b_tools: + async with setup_full_node( + test_constants_modified, + "blockchain_test_3.db", + self_hostname, + b_tools, + db_version=2, + ) as service: + yield service._api class FakeDNSResolver: diff --git a/chia/_tests/solver/test_solver_service.py b/chia/_tests/solver/test_solver_service.py index 2d4257cf0012..a704f7b50933 100644 --- a/chia/_tests/solver/test_solver_service.py +++ b/chia/_tests/solver/test_solver_service.py @@ -18,8 +18,10 @@ @pytest.mark.anyio async def test_solver_api_methods(blockchain_constants: ConsensusConstants, tmp_path: Path) -> None: with TempKeyring(populate=True) as keychain: - bt = await create_block_tools_async(constants=blockchain_constants, keychain=keychain) - async with setup_solver(tmp_path, bt, blockchain_constants) as solver_service: + async with ( + create_block_tools_async(constants=blockchain_constants, keychain=keychain) as bt, + setup_solver(tmp_path, bt, blockchain_constants) as solver_service, + ): solver = solver_service._node solver_api = solver_service._api assert solver_api.ready() is True diff --git a/chia/_tests/util/setup_nodes.py b/chia/_tests/util/setup_nodes.py index 8aa20544c22e..9284f4c02b79 100644 --- a/chia/_tests/util/setup_nodes.py +++ b/chia/_tests/util/setup_nodes.py @@ -93,27 +93,32 @@ async def setup_two_nodes( config_overrides = {"full_node.max_sync_wait": 0, "full_node.log_coins": True} with TempKeyring(populate=True) as keychain1, TempKeyring(populate=True) as keychain2: - bt1 = await create_block_tools_async( - constants=consensus_constants, keychain=keychain1, config_overrides=config_overrides - ) - async with setup_full_node( - consensus_constants, - "blockchain_test.db", - self_hostname, - bt1, - simulator=False, - db_version=db_version, - ) as service1: - async with setup_full_node( + async with ( + create_block_tools_async( + constants=consensus_constants, keychain=keychain1, config_overrides=config_overrides + ) as bt1, + setup_full_node( consensus_constants, - "blockchain_test_2.db", + "blockchain_test.db", self_hostname, - await create_block_tools_async( - constants=consensus_constants, keychain=keychain2, config_overrides=config_overrides - ), + bt1, simulator=False, db_version=db_version, - ) as service2: + ) as service1, + ): + async with ( + create_block_tools_async( + constants=consensus_constants, keychain=keychain2, config_overrides=config_overrides + ) as bt2, + setup_full_node( + consensus_constants, + "blockchain_test_2.db", + self_hostname, + bt2, + simulator=False, + db_version=db_version, + ) as service2, + ): fn1 = service1._api fn2 = service2._api @@ -137,8 +142,10 @@ async def setup_n_nodes( consensus_constants, f"blockchain_test_{i}.db", self_hostname, - await create_block_tools_async( - constants=consensus_constants, keychain=keychain, config_overrides=config_overrides + await async_exit_stack.enter_async_context( + create_block_tools_async( + constants=consensus_constants, keychain=keychain, config_overrides=config_overrides + ) ), simulator=False, db_version=db_version, @@ -255,14 +262,18 @@ async def setup_simulators_and_wallets_inner( config_overrides["full_node.log_coins"] = True async with AsyncExitStack() as async_exit_stack: bt_tools: list[BlockTools] = [ - await create_block_tools_async(consensus_constants, keychain=keychain1, config_overrides=config_overrides) + await async_exit_stack.enter_async_context( + create_block_tools_async(consensus_constants, keychain=keychain1, config_overrides=config_overrides) + ) for _ in range(simulator_count) ] if wallet_count > simulator_count: for _ in range(wallet_count - simulator_count): bt_tools.append( - await create_block_tools_async( - consensus_constants, keychain=keychain2, config_overrides=config_overrides + await async_exit_stack.enter_async_context( + create_block_tools_async( + consensus_constants, keychain=keychain2, config_overrides=config_overrides + ) ) ) @@ -370,18 +381,23 @@ async def setup_full_system_inner( shared_b_tools: BlockTools, ) -> AsyncIterator[FullSystem]: config_overrides = {"full_node.max_sync_wait": 0, "full_node.log_coins": True} - if b_tools is None: - b_tools = await create_block_tools_async( - constants=consensus_constants, keychain=keychain1, config_overrides=config_overrides - ) - if b_tools_1 is None: - b_tools_1 = await create_block_tools_async( - constants=consensus_constants, keychain=keychain2, config_overrides=config_overrides - ) self_hostname = shared_b_tools.config["self_hostname"] async with AsyncExitStack() as async_exit_stack: + if b_tools is None: + b_tools = await async_exit_stack.enter_async_context( + create_block_tools_async( + constants=consensus_constants, keychain=keychain1, config_overrides=config_overrides + ) + ) + if b_tools_1 is None: + b_tools_1 = await async_exit_stack.enter_async_context( + create_block_tools_async( + constants=consensus_constants, keychain=keychain2, config_overrides=config_overrides + ) + ) + vdf1_port = uint16(find_available_listen_port("vdf1")) vdf2_port = uint16(find_available_listen_port("vdf2")) diff --git a/chia/cmds/sim_funcs.py b/chia/cmds/sim_funcs.py index da90ce02df5a..8d7a7303f8bd 100644 --- a/chia/cmds/sim_funcs.py +++ b/chia/cmds/sim_funcs.py @@ -253,18 +253,18 @@ async def generate_plots(config: dict[str, Any], root_path: Path, fingerprint: i os.environ["CHIA_ROOT"] = str(root_path) # change env variable, to make it match what the daemon would set it to # create block tools and use local keychain - bt = BlockTools( + with BlockTools( test_constants, root_path, automated_testing=False, plot_dir=config["simulator"].get("plot_directory", "plots"), keychain=Keychain(), - ) - await bt.setup_keys(fingerprint=fingerprint, reward_ph=farming_puzzle_hash) - existing_plots = await bt.setup_plots( - num_og_plots=PLOTS, num_pool_plots=0, num_non_keychain_plots=0, plot_size=PLOT_SIZE, bitfield=bitfield - ) - print(f"{'New plots generated.' if existing_plots else 'Using Existing Plots'}\n") + ) as bt: + await bt.setup_keys(fingerprint=fingerprint, reward_ph=farming_puzzle_hash) + existing_plots = await bt.setup_plots( + num_og_plots=PLOTS, num_pool_plots=0, num_non_keychain_plots=0, plot_size=PLOT_SIZE, bitfield=bitfield + ) + print(f"{'New plots generated.' if existing_plots else 'Using Existing Plots'}\n") async def get_current_height(root_path: Path) -> int: diff --git a/chia/simulator/block_tools.py b/chia/simulator/block_tools.py index 3919e3998428..7fd064c20a86 100644 --- a/chia/simulator/block_tools.py +++ b/chia/simulator/block_tools.py @@ -10,10 +10,12 @@ import sys import tempfile import time -from collections.abc import Callable, Sequence +from collections.abc import AsyncIterator, Callable, Iterator, Sequence +from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass, replace from pathlib import Path from random import Random +from types import TracebackType from typing import Any import anyio @@ -42,6 +44,7 @@ from chia_rs.sized_bytes import bytes32 from chia_rs.sized_ints import uint8, uint16, uint32, uint64, uint128 from filelock import FileLock +from typing_extensions import Self from chia.consensus.block_creation import create_unfinished_block, unfinished_block_to_full_block from chia.consensus.block_record import BlockRecordProtocol @@ -259,7 +262,7 @@ def __init__( self._tempdir = None if root_path is None: - self._tempdir = tempfile.TemporaryDirectory() + self._tempdir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) root_path = Path(self._tempdir.name) self.root_path = root_path @@ -375,6 +378,18 @@ def test_callback(event: PlotRefreshEvents, update_result: PlotRefreshResult) -> match_str=str(self.plot_dir.relative_to(DEFAULT_ROOT_PATH.parent)) if not automated_testing else None, ) + def __enter__(self) -> Self: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + if self._tempdir is not None: + self._tempdir.cleanup() + def setup_new_gen( self, generator_block_heights: list[uint32], @@ -2162,6 +2177,7 @@ class BlockToolsNewPlotResult: # listen port it uses is to write it to the config file. +@asynccontextmanager async def create_block_tools_async( constants: ConsensusConstants = test_constants, root_path: Path | None = None, @@ -2171,36 +2187,37 @@ async def create_block_tools_async( num_pool_plots: int = 5, num_non_keychain_plots: int = 3, testrun_uid: str | None = None, -) -> BlockTools: +) -> AsyncIterator[BlockTools]: global create_block_tools_async_count create_block_tools_async_count += 1 print(f" create_block_tools_async called {create_block_tools_async_count} times") - bt = BlockTools(constants, root_path, keychain, config_overrides=config_overrides) - await bt.setup_keys() - await bt.setup_plots( - num_og_plots=num_og_plots, - num_pool_plots=num_pool_plots, - num_non_keychain_plots=num_non_keychain_plots, - testrun_uid=testrun_uid, - ) - return bt + with BlockTools(constants, root_path, keychain, config_overrides=config_overrides) as bt: + await bt.setup_keys() + await bt.setup_plots( + num_og_plots=num_og_plots, + num_pool_plots=num_pool_plots, + num_non_keychain_plots=num_non_keychain_plots, + testrun_uid=testrun_uid, + ) + + yield bt +@contextmanager def create_block_tools( constants: ConsensusConstants = test_constants, root_path: Path | None = None, keychain: Keychain | None = None, config_overrides: dict[str, Any] | None = None, -) -> BlockTools: +) -> Iterator[BlockTools]: global create_block_tools_count create_block_tools_count += 1 print(f" create_block_tools called {create_block_tools_count} times") - bt = BlockTools(constants, root_path, keychain, config_overrides=config_overrides) - - asyncio.get_event_loop().run_until_complete(bt.setup_keys()) - asyncio.get_event_loop().run_until_complete(bt.setup_plots()) - return bt + with BlockTools(constants, root_path, keychain, config_overrides=config_overrides) as bt: + asyncio.get_event_loop().run_until_complete(bt.setup_keys()) + asyncio.get_event_loop().run_until_complete(bt.setup_plots()) + yield bt def make_unfinished_block( diff --git a/tools/generate_chain.py b/tools/generate_chain.py index 8109c261d6a3..b9512f496243 100644 --- a/tools/generate_chain.py +++ b/tools/generate_chain.py @@ -77,8 +77,10 @@ def main(length: int, fill_rate: int, profile: bool, block_refs: bool, output: s root_path = Path("./test-chain").resolve() root_path.mkdir(parents=True, exist_ok=True) - with TempKeyring() as keychain: - bt = create_block_tools(constants=test_constants, root_path=root_path, keychain=keychain) + with ( + TempKeyring() as keychain, + create_block_tools(constants=test_constants, root_path=root_path, keychain=keychain) as bt, + ): initialize_logging( "generate_chain", {"log_level": "DEBUG", "log_stdout": False, "log_syslog": False}, root_path=root_path )