mirror of
https://github.com/Chia-Network/chia-blockchain.git
synced 2026-09-03 02:24:18 -05:00
433 lines
13 KiB
Python
433 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import dataclasses
|
|
import enum
|
|
import functools
|
|
import gc
|
|
import logging
|
|
import os
|
|
import pathlib
|
|
import subprocess
|
|
import sys
|
|
from concurrent.futures import Future
|
|
from inspect import getframeinfo, stack
|
|
from statistics import mean
|
|
from textwrap import dedent
|
|
from time import thread_time
|
|
from types import TracebackType
|
|
from typing import Any, Callable, Collection, Iterator, List, Optional, TextIO, Tuple, Type, Union
|
|
|
|
import pytest
|
|
from chia_rs import Coin
|
|
from typing_extensions import Protocol, final
|
|
|
|
import chia
|
|
from chia.full_node.mempool import Mempool
|
|
from chia.types.blockchain_format.sized_bytes import bytes32
|
|
from chia.types.condition_opcodes import ConditionOpcode
|
|
from chia.util.hash import std_hash
|
|
from chia.util.ints import uint32, uint64
|
|
from chia.wallet.util.compute_hints import HintedCoin
|
|
from chia.wallet.wallet_node import WalletNode
|
|
from tests.core.data_layer.util import ChiaRoot
|
|
|
|
|
|
class GcMode(enum.Enum):
|
|
nothing = enum.auto
|
|
precollect = enum.auto
|
|
disable = enum.auto
|
|
enable = enum.auto
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def manage_gc(mode: GcMode) -> Iterator[None]:
|
|
if mode == GcMode.nothing:
|
|
yield
|
|
elif mode == GcMode.precollect:
|
|
gc.collect()
|
|
yield
|
|
elif mode == GcMode.disable:
|
|
was_enabled = gc.isenabled()
|
|
gc.disable()
|
|
try:
|
|
yield
|
|
finally:
|
|
if was_enabled:
|
|
gc.enable()
|
|
elif mode == GcMode.enable:
|
|
was_enabled = gc.isenabled()
|
|
gc.enable()
|
|
try:
|
|
yield
|
|
finally:
|
|
if not was_enabled:
|
|
gc.disable()
|
|
|
|
|
|
def caller_file_and_line(distance: int = 1) -> Tuple[str, int]:
|
|
caller = getframeinfo(stack()[distance + 1][0])
|
|
return caller.filename, caller.lineno
|
|
|
|
|
|
@dataclasses.dataclass(frozen=True)
|
|
class RuntimeResults:
|
|
start: float
|
|
end: float
|
|
duration: float
|
|
entry_file: str
|
|
entry_line: int
|
|
overhead: Optional[float]
|
|
|
|
def block(self, label: str = "") -> str:
|
|
# The entry line is reported starting at the beginning of the line to trigger
|
|
# PyCharm to highlight as a link to the source.
|
|
|
|
return dedent(
|
|
f"""\
|
|
Measuring runtime: {label}
|
|
{self.entry_line}
|
|
run time: {self.duration}
|
|
overhead: {self.overhead if self.overhead is not None else "not measured"}
|
|
"""
|
|
)
|
|
|
|
|
|
@final
|
|
@dataclasses.dataclass(frozen=True)
|
|
class AssertRuntimeResults:
|
|
start: float
|
|
end: float
|
|
duration: float
|
|
entry_file: str
|
|
entry_line: int
|
|
overhead: Optional[float]
|
|
limit: float
|
|
ratio: float
|
|
|
|
@classmethod
|
|
def from_runtime_results(
|
|
cls, results: RuntimeResults, limit: float, entry_file: str, entry_line: int, overhead: Optional[float]
|
|
) -> AssertRuntimeResults:
|
|
return cls(
|
|
start=results.start,
|
|
end=results.end,
|
|
duration=results.duration,
|
|
limit=limit,
|
|
ratio=results.duration / limit,
|
|
entry_file=entry_file,
|
|
entry_line=entry_line,
|
|
overhead=overhead,
|
|
)
|
|
|
|
def block(self, label: str = "") -> str:
|
|
# The entry line is reported starting at the beginning of the line to trigger
|
|
# PyCharm to highlight as a link to the source.
|
|
|
|
return dedent(
|
|
f"""\
|
|
Asserting maximum duration: {label}
|
|
{self.entry_file}:{self.entry_line}
|
|
run time: {self.duration}
|
|
overhead: {self.overhead if self.overhead is not None else "not measured"}
|
|
allowed: {self.limit}
|
|
percent: {self.percent_str()}
|
|
"""
|
|
)
|
|
|
|
def message(self) -> str:
|
|
return f"{self.duration} seconds not less than {self.limit} seconds ( {self.percent_str()} )"
|
|
|
|
def passed(self) -> bool:
|
|
return self.duration < self.limit
|
|
|
|
def percent(self) -> float:
|
|
return self.ratio * 100
|
|
|
|
def percent_str(self) -> str:
|
|
return f"{self.percent():.0f} %"
|
|
|
|
|
|
def measure_overhead(
|
|
manager_maker: Callable[
|
|
[], contextlib.AbstractContextManager[Union[Future[RuntimeResults], Future[AssertRuntimeResults]]]
|
|
],
|
|
cycles: int = 10,
|
|
) -> float:
|
|
times: List[float] = []
|
|
|
|
for _ in range(cycles):
|
|
with manager_maker() as results:
|
|
pass
|
|
|
|
times.append(results.result(timeout=0).duration)
|
|
|
|
overhead = mean(times)
|
|
|
|
return overhead
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def measure_runtime(
|
|
label: str = "",
|
|
clock: Callable[[], float] = thread_time,
|
|
gc_mode: GcMode = GcMode.disable,
|
|
overhead: Optional[float] = None,
|
|
print_results: bool = True,
|
|
) -> Iterator[Future[RuntimeResults]]:
|
|
entry_file, entry_line = caller_file_and_line()
|
|
|
|
results_future: Future[RuntimeResults] = Future()
|
|
|
|
with manage_gc(mode=gc_mode):
|
|
start = clock()
|
|
|
|
try:
|
|
yield results_future
|
|
finally:
|
|
end = clock()
|
|
|
|
duration = end - start
|
|
if overhead is not None:
|
|
duration -= overhead
|
|
|
|
results = RuntimeResults(
|
|
start=start,
|
|
end=end,
|
|
duration=duration,
|
|
entry_file=entry_file,
|
|
entry_line=entry_line,
|
|
overhead=overhead,
|
|
)
|
|
results_future.set_result(results)
|
|
|
|
if print_results:
|
|
print(results.block(label=label))
|
|
|
|
|
|
@final
|
|
@dataclasses.dataclass
|
|
class _AssertRuntime:
|
|
"""Prepare for, measure, and assert about the time taken by code in the context.
|
|
|
|
Defaults are set for single-threaded CPU usage timing without garbage collection.
|
|
|
|
In general, there is no generally correct setup for benchmarking. Only measuring
|
|
a single thread's time using the CPU is not very useful for multithreaded or
|
|
multiprocessed code. Disabling garbage collection, or forcing it ahead of time,
|
|
makes the benchmark not identify any issues the code may introduce in terms of
|
|
actually causing relevant gc slowdowns. And so on...
|
|
|
|
Produces output of the following form.
|
|
|
|
Asserting maximum duration: full block
|
|
/home/altendky/repos/chia-blockchain/tests/core/full_node/test_performance.py:187
|
|
run time: 0.027789528900002837
|
|
allowed: 0.1
|
|
percent: 28 %
|
|
"""
|
|
|
|
# A class is only being used here, to make __tracebackhide__ work.
|
|
# https://github.com/pytest-dev/pytest/issues/2057
|
|
|
|
seconds: float
|
|
label: str = ""
|
|
clock: Callable[[], float] = thread_time
|
|
gc_mode: GcMode = GcMode.disable
|
|
print: bool = True
|
|
overhead: Optional[float] = None
|
|
entry_file: Optional[str] = None
|
|
entry_line: Optional[int] = None
|
|
_results: Optional[AssertRuntimeResults] = None
|
|
runtime_manager: Optional[contextlib.AbstractContextManager[Future[RuntimeResults]]] = None
|
|
runtime_results_callable: Optional[Future[RuntimeResults]] = None
|
|
enable_assertion: bool = True
|
|
record_property: Optional[Callable[[str, object], None]] = None
|
|
|
|
def __enter__(self) -> Future[AssertRuntimeResults]:
|
|
self.entry_file, self.entry_line = caller_file_and_line()
|
|
|
|
self.runtime_manager = measure_runtime(
|
|
clock=self.clock, gc_mode=self.gc_mode, overhead=self.overhead, print_results=False
|
|
)
|
|
self.runtime_results_callable = self.runtime_manager.__enter__()
|
|
self.results_callable: Future[AssertRuntimeResults] = Future()
|
|
|
|
return self.results_callable
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: Optional[Type[BaseException]],
|
|
exc: Optional[BaseException],
|
|
traceback: Optional[TracebackType],
|
|
) -> None:
|
|
if (
|
|
self.entry_file is None
|
|
or self.entry_line is None
|
|
or self.runtime_manager is None
|
|
or self.runtime_results_callable is None
|
|
):
|
|
raise Exception("Context manager must be entered before exiting")
|
|
|
|
self.runtime_manager.__exit__(exc_type, exc, traceback)
|
|
|
|
runtime = self.runtime_results_callable.result(timeout=0)
|
|
results = AssertRuntimeResults.from_runtime_results(
|
|
results=runtime,
|
|
limit=self.seconds,
|
|
entry_file=self.entry_file,
|
|
entry_line=self.entry_line,
|
|
overhead=self.overhead,
|
|
)
|
|
|
|
self.results_callable.set_result(results)
|
|
|
|
if self.print:
|
|
print(results.block(label=self.label))
|
|
|
|
if self.record_property is not None:
|
|
self.record_property(f"duration:{self.label}", results.duration)
|
|
|
|
relative_path_str = (
|
|
pathlib.Path(results.entry_file).relative_to(pathlib.Path(chia.__file__).parent.parent).as_posix()
|
|
)
|
|
|
|
self.record_property(f"path:{self.label}", relative_path_str)
|
|
self.record_property(f"line:{self.label}", results.entry_line)
|
|
self.record_property(f"limit:{self.label}", self.seconds)
|
|
|
|
if exc_type is None and self.enable_assertion:
|
|
__tracebackhide__ = True
|
|
assert runtime.duration < self.seconds, results.message()
|
|
|
|
|
|
@final
|
|
@dataclasses.dataclass
|
|
class BenchmarkRunner:
|
|
enable_assertion: bool = True
|
|
label: Optional[str] = None
|
|
overhead: Optional[float] = None
|
|
record_property: Optional[Callable[[str, object], None]] = None
|
|
|
|
@functools.wraps(_AssertRuntime)
|
|
def assert_runtime(self, *args: Any, **kwargs: Any) -> _AssertRuntime:
|
|
kwargs.setdefault("enable_assertion", self.enable_assertion)
|
|
kwargs.setdefault("overhead", self.overhead)
|
|
kwargs.setdefault("record_property", self.record_property)
|
|
return _AssertRuntime(*args, **kwargs)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def assert_rpc_error(error: str) -> Iterator[None]:
|
|
with pytest.raises(ValueError) as exception_info:
|
|
yield
|
|
assert error in exception_info.value.args[0]["error"]
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def closing_chia_root_popen(chia_root: ChiaRoot, args: List[str]) -> Iterator[subprocess.Popen[Any]]:
|
|
environment = {**os.environ, "CHIA_ROOT": os.fspath(chia_root.path)}
|
|
|
|
with subprocess.Popen(args=args, env=environment) as process:
|
|
try:
|
|
yield process
|
|
finally:
|
|
process.terminate()
|
|
try:
|
|
process.wait(timeout=10)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
|
|
|
|
# https://github.com/pytest-dev/pytest/blob/7.3.1/src/_pytest/mark/__init__.py#L45
|
|
Marks = Union[pytest.MarkDecorator, Collection[Union[pytest.MarkDecorator, pytest.Mark]]]
|
|
|
|
|
|
class DataCase(Protocol):
|
|
marks: Marks
|
|
|
|
@property
|
|
def id(self) -> str:
|
|
...
|
|
|
|
|
|
def datacases(*cases: DataCase, _name: str = "case") -> pytest.MarkDecorator:
|
|
return pytest.mark.parametrize(
|
|
argnames=_name,
|
|
argvalues=[pytest.param(case, id=case.id, marks=case.marks) for case in cases],
|
|
)
|
|
|
|
|
|
class DataCasesDecorator(Protocol):
|
|
def __call__(self, *cases: DataCase, _name: str = "case") -> pytest.MarkDecorator:
|
|
...
|
|
|
|
|
|
def named_datacases(name: str) -> DataCasesDecorator:
|
|
return functools.partial(datacases, _name=name)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class CoinGenerator:
|
|
_seed: int = -1
|
|
|
|
def _get_hash(self) -> bytes32:
|
|
self._seed += 1
|
|
return std_hash(self._seed.to_bytes(length=32, byteorder="big"))
|
|
|
|
def _get_amount(self) -> uint64:
|
|
self._seed += 1
|
|
return uint64(self._seed)
|
|
|
|
def get(self, parent_coin_id: Optional[bytes32] = None, include_hint: bool = True) -> HintedCoin:
|
|
if parent_coin_id is None:
|
|
parent_coin_id = self._get_hash()
|
|
hint = None
|
|
if include_hint:
|
|
hint = self._get_hash()
|
|
return HintedCoin(Coin(parent_coin_id, self._get_hash(), self._get_amount()), hint)
|
|
|
|
|
|
def coin_creation_args(hinted_coin: HintedCoin) -> List[Any]:
|
|
if hinted_coin.hint is not None:
|
|
memos = [hinted_coin.hint]
|
|
else:
|
|
memos = []
|
|
return [ConditionOpcode.CREATE_COIN, hinted_coin.coin.puzzle_hash, hinted_coin.coin.amount, memos]
|
|
|
|
|
|
def create_logger(file: TextIO = sys.stdout) -> logging.Logger:
|
|
logger = logging.getLogger()
|
|
logger.setLevel(level=logging.DEBUG)
|
|
stream_handler = logging.StreamHandler(stream=file)
|
|
log_date_format = "%Y-%m-%dT%H:%M:%S"
|
|
file_log_formatter = logging.Formatter(
|
|
fmt="%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s",
|
|
datefmt=log_date_format,
|
|
)
|
|
stream_handler.setFormatter(file_log_formatter)
|
|
logger.addHandler(hdlr=stream_handler)
|
|
|
|
return logger
|
|
|
|
|
|
def invariant_check_mempool(mempool: Mempool) -> None:
|
|
with mempool._db_conn:
|
|
cursor = mempool._db_conn.execute("SELECT SUM(cost) FROM tx")
|
|
val = cursor.fetchone()[0]
|
|
if val is None:
|
|
val = 0
|
|
assert mempool._total_cost == val
|
|
|
|
with mempool._db_conn:
|
|
cursor = mempool._db_conn.execute("SELECT SUM(fee) FROM tx")
|
|
val = cursor.fetchone()[0]
|
|
if val is None:
|
|
val = 0
|
|
assert mempool._total_fee == val
|
|
|
|
|
|
async def wallet_height_at_least(wallet_node: WalletNode, h: uint32) -> bool:
|
|
height = await wallet_node.wallet_state_manager.blockchain.get_finished_sync_up_to()
|
|
return height == h
|