Repository: Chia-Network/chia-blockchain
File: tools/validate_rpcs.py (Python)

Code before:
        peak_height = blockchain_state["peak"].height
        assert peak_height is not None, "Blockchain peak height is None"
        if end_height is None:
            end_height = blockchain_state["peak"]["height"]

        print("Connected to Full Node")

        block_cache_bytearray: bytearray = await get_block_cache_bytearray(
            root_path=root_path,
            config=config,
            peak=peak_height,
        )

        print("Bytearray loaded with block header hashes from height-to-hash file.")

        # set initial segment heights
        start_segment: int = start_height
        end_segment: int = start_height + concurrent_requests
        # measure time for performance measurement per segment
        cycle_start: float = time.time()
        # also measure time for the whole process
        start_time: float = cycle_start

        while end_segment <= end_height:
            # Create an initial list to hold pending tasks
            pending_tasks: list[Coroutine[Any, Any, Any]] = []

            for i in range(start_segment, end_segment):
                block_header_hash: bytes32 = get_block_header_from_height(i, block_cache_bytearray)
                if spends_with_conditions:
                    pending_tasks.append(node_client.get_block_spends_with_conditions(block_header_hash))
                if block_spends:
                    pending_tasks.append(node_client.get_block_spends(block_header_hash))
                if additions_and_removals:
                    pending_tasks.append(node_client.get_additions_and_removals(block_header_hash))
            try:
                results = await asyncio.gather(*pending_tasks)
                for result in results:
                    if result is None:
                        raise ValueError("Received None from RPC call")
            except Exception as e:
                print(f"Error processing block range {start_segment} to {end_segment}: {e}")
                raise e
            # Print status every blocks_per_status blocks
            if end_segment - last_status_height >= blocks_per_status:
                time_taken = time.time() - cycle_start
                time_per_block = time_taken / (end_segment - last_status_height)
                print(
                    f"Time taken for segment"
                    f" {last_status_height} to {end_segment}: {time_taken:.2f} seconds\n"
                    f"Average time per block: {time_per_block:.4f} seconds"

Review comment (refactor):
I was going to suggest using `TaskGroup` here, as a context manager it would make this a lot cleaner. But I suppose that requires python 3.11. That might be fine for a test utility though

Code after:
    block_spends: bool,
    additions_and_removals: bool,
    pipeline_depth: int,
    start_height: int,
    end_height: Optional[int] = None,
) -> None:
    async with get_any_service_client(FullNodeRpcClient, root_path, rpc_port) as (
        node_client,
        config,
    ):
        blockchain_state: dict[str, Any] = await node_client.get_blockchain_state()
        if blockchain_state is None or blockchain_state["peak"] is None:
            # Peak height is required for the cache.
            print("No blockchain found. Exiting.")
            return
        peak_height = blockchain_state["peak"].height
        assert peak_height is not None, "Blockchain peak height is None"
        if end_height is None:
            end_height = peak_height

        print("Connected to Full Node")

        height_to_hash_bytes: bytes = await get_height_to_hash_bytes(root_path=root_path, config=config)

        print("block header hashes loaded from height-to-hash file.")

        # Set initial values for the loop

        pipeline: set[asyncio.Task[None]] = set()
        completed_requests: int = 0
        # measure time for performance measurement per segment
        cycle_start: float = time.monotonic()
        # also measure time for the whole process
        start_time: float = cycle_start

        def add_tasks_for_height(height: int) -> None:
            block_header_hash = get_block_hash_from_height(height, height_to_hash_bytes)
            # Create tasks for each RPC call based on the flags
            if spends_with_conditions:
                pipeline.add(
                    create_referenced_task(node_spends_with_conditions(node_client, block_header_hash, height))
                )
            if block_spends:
                pipeline.add(create_referenced_task(node_block_spends(node_client, block_header_hash, height)))
            if additions_and_removals:
                pipeline.add(create_referenced_task(node_additions_removals(node_client, block_header_hash, height)))

        for i in range(start_height, end_height + 1):
            add_tasks_for_height(height=i)
            # Make Status Updates.
            if len(pipeline) >= pipeline_depth: