chia-blockchain 2.5.1rc1__py3-none-any.whl

chia/util/network.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import asyncio
+import ipaddress
+import logging
+import socket
+import ssl
+from collections.abc import Iterable
+from dataclasses import dataclass
+from ipaddress import IPv4Network, IPv6Network, ip_address
+from typing import Any, Literal, Optional, Union
+
+from aiohttp import web
+from aiohttp.log import web_logger
+from typing_extensions import final
+
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.util.ints import uint16
+from chia.util.ip_address import IPAddress
+from chia.util.task_referencer import create_referenced_task
+
+
+@final
+@dataclass
+class WebServer:
+    runner: web.AppRunner
+    hostname: str
+    listen_port: uint16
+    scheme: Literal["http", "https"]
+    _ssl_context: Optional[ssl.SSLContext] = None
+    _close_task: Optional[asyncio.Task[None]] = None
+    _prefer_ipv6: bool = False
+
+    @classmethod
+    async def create(
+        cls,
+        hostname: str,
+        port: uint16,
+        routes: Iterable[web.RouteDef] = (),
+        max_request_body_size: int = 1024**2,  # Default `client_max_size` from web.Application
+        ssl_context: Optional[ssl.SSLContext] = None,
+        keepalive_timeout: int = 75,  # Default from aiohttp.web
+        shutdown_timeout: int = 60,  # Default `shutdown_timeout` from aiohttp.web_runner.BaseRunner
+        prefer_ipv6: bool = False,
+        logger: logging.Logger = web_logger,
+        start: bool = True,
+    ) -> WebServer:
+        app = web.Application(client_max_size=max_request_body_size, logger=logger)
+        runner = web.AppRunner(
+            app,
+            access_log=None,
+            keepalive_timeout=keepalive_timeout,
+            shutdown_timeout=shutdown_timeout,
+        )
+
+        self = cls(
+            runner=runner,
+            hostname=hostname,
+            listen_port=uint16(port),
+            scheme="https" if ssl_context is not None else "http",
+            _ssl_context=ssl_context,
+            _prefer_ipv6=prefer_ipv6,
+        )
+
+        self.add_routes(routes)
+
+        if start:
+            await self.start()
+
+        return self
+
+    async def start(self) -> None:
+        await self.runner.setup()
+        site = web.TCPSite(
+            self.runner,
+            self.hostname,
+            int(self.listen_port),
+            ssl_context=self._ssl_context,
+        )
+        await site.start()
+
+        #
+        # On a dual-stack system, we want to get the (first) IPv4 port unless
+        # prefer_ipv6 is set in which case we use the IPv6 port
+        #
+        if self.listen_port == 0:
+            self.listen_port = select_port(self._prefer_ipv6, self.runner.addresses)
+
+    def add_routes(self, routes: Iterable[web.RouteDef]) -> None:
+        self.runner.app.add_routes(routes)
+
+    def url(self, *segments: str) -> str:
+        path = "/".join(segments)
+        return f"{self.scheme}://{self.hostname}:{self.listen_port}/{path}"
+
+    async def _close(self) -> None:
+        await self.runner.shutdown()
+        await self.runner.cleanup()
+
+    def close(self) -> None:
+        self._close_task = create_referenced_task(self._close())
+
+    async def await_closed(self) -> None:
+        if self._close_task is None:
+            raise RuntimeError("WebServer stop not triggered")
+        await self._close_task
+
+
+def is_in_network(peer_host: str, networks: Iterable[Union[IPv4Network, IPv6Network]]) -> bool:
+    try:
+        peer_host_ip = ip_address(peer_host)
+        return any(peer_host_ip in network for network in networks)
+    except ValueError:
+        return False
+
+
+def is_trusted_cidr(peer_host: str, trusted_cidrs: list[str]) -> bool:
+    try:
+        ip_obj = ipaddress.ip_address(peer_host)
+    except ValueError:
+        return False
+
+    for cidr in trusted_cidrs:
+        network = ipaddress.ip_network(cidr)
+        if ip_obj in network:
+            return True
+
+    return False
+
+
+def is_localhost(peer_host: str) -> bool:
+    return peer_host in {"127.0.0.1", "localhost", "::1", "0:0:0:0:0:0:0:1"}
+
+
+def is_trusted_peer(
+    host: str, node_id: bytes32, trusted_peers: dict[str, Any], trusted_cidrs: list[str], testing: bool = False
+) -> bool:
+    return (
+        (not testing and is_localhost(host)) or node_id.hex() in trusted_peers or is_trusted_cidr(host, trusted_cidrs)
+    )
+
+
+async def resolve(host: str, *, prefer_ipv6: bool = False) -> IPAddress:
+    try:
+        return IPAddress.create(host)
+    except ValueError:
+        pass
+    addrset: list[
+        tuple[socket.AddressFamily, socket.SocketKind, int, str, Union[tuple[str, int], tuple[str, int, int, int]]]
+    ] = await asyncio.get_event_loop().getaddrinfo(host, None)
+    # The list returned by getaddrinfo is never empty, an exception is thrown or data is returned.
+    ips_v4 = []
+    ips_v6 = []
+    for family, _, _, _, ip_port in addrset:
+        ip = IPAddress.create(ip_port[0])
+        if family == socket.AF_INET:
+            ips_v4.append(ip)
+        if family == socket.AF_INET6:
+            ips_v6.append(ip)
+    preferred, alternative = (ips_v6, ips_v4) if prefer_ipv6 else (ips_v4, ips_v6)
+    if len(preferred) > 0:
+        return preferred[0]
+    elif len(alternative) > 0:
+        return alternative[0]
+    else:
+        raise ValueError(f"failed to resolve {host} into an IP address")
+
+
+def select_port(prefer_ipv6: bool, addresses: list[Any]) -> uint16:
+    selected_port: uint16
+    for address_string, port, *_ in addresses:
+        address = ip_address(address_string)
+        if address.version == 6 and prefer_ipv6:
+            selected_port = port
+            break
+        elif address.version == 4 and not prefer_ipv6:
+            selected_port = port
+            break
+    else:
+        selected_port = addresses[0][1]  # no matches, just use the first one in the list
+
+    return selected_port

chia/util/paginator.py

@@ -0,0 +1,48 @@
+# Package: utils
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Sequence
+from math import ceil
+
+
+class InvalidPageSizeLimit(Exception):
+    def __init__(self, page_size_limit: int) -> None:
+        super().__init__(f"Page size limit must be one or more, not: {page_size_limit}")
+
+
+class InvalidPageSizeError(Exception):
+    def __init__(self, page_size: int, page_size_limit: int) -> None:
+        super().__init__(f"Invalid page size {page_size}. Must be between: 1 and {page_size_limit}")
+
+
+class PageOutOfBoundsError(Exception):
+    def __init__(self, page_size: int, max_page_size: int) -> None:
+        super().__init__(f"Page {page_size} out of bounds. Available pages: 0-{max_page_size}")
+
+
+@dataclasses.dataclass
+class Paginator:
+    _source: Sequence[object]
+    _page_size: int
+
+    @classmethod
+    def create(cls, source: Sequence[object], page_size: int, page_size_limit: int = 100) -> Paginator:
+        if page_size_limit < 1:
+            raise InvalidPageSizeLimit(page_size_limit)
+        if page_size > page_size_limit:
+            raise InvalidPageSizeError(page_size, page_size_limit)
+        return cls(source, page_size)
+
+    def page_size(self) -> int:
+        return self._page_size
+
+    def page_count(self) -> int:
+        return max(1, ceil(len(self._source) / self._page_size))
+
+    def get_page(self, page: int) -> Sequence[object]:
+        if page < 0 or page >= self.page_count():
+            raise PageOutOfBoundsError(page, self.page_count() - 1)
+        offset = page * self._page_size
+        return self._source[offset : offset + self._page_size]

chia/util/path.py

@@ -0,0 +1,31 @@
+# Package: utils
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Union
+
+
+def path_from_root(root: Path, path_str: Union[str, Path]) -> Path:
+    """
+    If path is relative, prepend root
+    If path is absolute, return it directly.
+    """
+    root = Path(os.path.expanduser(str(root)))
+    path = Path(path_str)
+    if not path.is_absolute():
+        path = root / path
+    return path.resolve()
+
+
+def make_path_relative(path_str: Union[str, Path], root: Path) -> Path:
+    """
+    Try to make the given path relative, given the default root.
+    """
+    path = Path(path_str)
+    try:
+        path = path.relative_to(root)
+    except ValueError:
+        pass
+    return path

chia/util/permissions.py

@@ -0,0 +1,20 @@
+# Package: utils
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+def verify_file_permissions(path: Path, mask: int) -> tuple[bool, int]:
+    """
+    Check that the file's permissions are properly restricted, as compared to the
+    permission mask
+    """
+    mode = os.stat(path).st_mode & 0o777
+    return (mode & mask == 0, mode)
+
+
+def octal_mode_string(mode: int) -> str:
+    """Yields a permission mode string: e.g. 0644"""
+    return f"0{oct(mode)[-3:]}"

chia/util/prev_transaction_block.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from chia.consensus.block_record import BlockRecord
+from chia.consensus.blockchain_interface import BlockRecordsProtocol
+from chia.util.ints import uint128
+
+
+def get_prev_transaction_block(
+    curr: BlockRecord,
+    blocks: BlockRecordsProtocol,
+    total_iters_sp: uint128,
+) -> tuple[bool, BlockRecord]:
+    prev_transaction_block = curr
+    while not curr.is_transaction_block:
+        curr = blocks.block_record(curr.prev_hash)
+    if total_iters_sp > curr.total_iters:
+        prev_transaction_block = curr
+        is_transaction_block = True
+    else:
+        is_transaction_block = False
+    return is_transaction_block, prev_transaction_block

chia/util/priority_mutex.py

@@ -0,0 +1,95 @@
+# Package: utils
+
+from __future__ import annotations
+
+import asyncio
+import collections
+import contextlib
+import dataclasses
+import logging
+from collections.abc import AsyncIterator
+from enum import IntEnum
+from typing import Generic, Optional, TypeVar
+
+from typing_extensions import final
+
+log = logging.getLogger(__name__)
+
+
+class NestedLockUnsupportedError(Exception):
+    pass
+
+
+_T_Priority = TypeVar("_T_Priority", bound=IntEnum)
+
+
+@dataclasses.dataclass(frozen=True)
+class _Element:
+    task: asyncio.Task[object] = dataclasses.field(compare=False)
+    ready_event: asyncio.Event = dataclasses.field(default_factory=asyncio.Event, compare=False)
+
+
+@final
+@dataclasses.dataclass()
+class PriorityMutex(Generic[_T_Priority]):
+    """
+    The purpose of this class is to be able to control access to a mutex, and give
+    priority to certain requests.  Lower values are given access first.  To use it,
+    create a mutex and use the `.acquire()` context manager method.  In actual uses,
+    a dedicated priority enumeration is likely to be more clear.
+
+    ```
+    my_mutex = PriorityMutex.create(priority_type=int)
+
+    async with my_mutex.acquire(priority=0):
+       ...
+    async with my_mutex.acquire(priority=1):
+       ...
+    ```
+    """
+
+    _deques: dict[_T_Priority, collections.deque[_Element]]
+    _active: Optional[_Element] = None
+
+    @classmethod
+    def create(cls, priority_type: type[_T_Priority]) -> PriorityMutex[_T_Priority]:
+        return cls(
+            _deques={priority: collections.deque() for priority in sorted(priority_type)},
+        )
+
+    @contextlib.asynccontextmanager
+    async def acquire(
+        self,
+        priority: _T_Priority,
+    ) -> AsyncIterator[None]:
+        task = asyncio.current_task()
+        if task is None:
+            raise Exception(f"unable to check current task, got: {task!r}")
+        if self._active is not None and self._active.task is task:
+            raise NestedLockUnsupportedError()
+
+        element = _Element(task=task)
+
+        deque = self._deques[priority]
+        deque.append(element)
+        try:
+            if self._active is None:
+                self._active = element
+            else:
+                await element.ready_event.wait()
+            yield
+        finally:
+            # another element might be active if the wait is cancelled
+            if self._active is element:
+                self._active = None
+            deque.remove(element)
+
+            if self._active is None:
+                for deque in self._deques.values():
+                    if len(deque) == 0:
+                        continue
+
+                    element = deque[0]
+                    self._active = element
+                    element.ready_event.set()
+                    break

chia/util/profiler.py

@@ -0,0 +1,197 @@
+# Package: utils
+
+from __future__ import annotations
+
+import asyncio
+import cProfile
+import logging
+import pathlib
+import tracemalloc
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from typing import Optional
+
+from chia.util.path import path_from_root
+
+# to use the profiler, enable it config file, "enable_profiler"
+# the output will be printed to your chia root path, e.g. ~/.chia/mainnet/profile/
+# to analyze the profile, run:
+
+#   python chia/utils/profiler.py ~/.chia/mainnet/profile | less -r
+
+# this will print CPU usage of the chia full node main thread at 1 second increments.
+# find a time window of interest and analyze the profile file (which are in pstats format).
+
+# for example:
+
+#   python chia/utils/profiler.py ~/.chia/mainnet/profile 10 20
+
+
+async def profile_task(root_path: pathlib.Path, service: str, log: logging.Logger) -> None:
+    profile_dir = path_from_root(root_path, f"profile-{service}")
+    log.info("Starting profiler. saving to %s", profile_dir)
+    profile_dir.mkdir(parents=True, exist_ok=True)
+
+    counter = 0
+
+    while True:
+        pr = cProfile.Profile()
+        pr.enable()
+        # this will throw CancelledError when we're exiting
+        await asyncio.sleep(1)
+        pr.create_stats()
+        pr.dump_stats(profile_dir / f"slot-{counter:05}.profile")
+        log.debug("saving profile %05d", counter)
+        counter += 1
+
+
+if __name__ == "__main__":
+    import io
+    import pstats
+    import sys
+    from subprocess import check_call
+
+    from colorama import Back, Fore, Style, init
+
+    profile_dir = pathlib.Path(sys.argv[1])
+    init(strip=False)
+
+    def analyze_cpu_usage(profile_dir: pathlib.Path) -> None:
+        counter = 0
+        try:
+            while True:
+                f = io.StringIO()
+                st = pstats.Stats(str(profile_dir / f"slot-{counter:05}.profile"), stream=f)
+                st.strip_dirs()
+                st.sort_stats(pstats.SortKey.CUMULATIVE)
+                st.print_stats()
+                f.seek(0)
+                total = 0.0
+                sleep = 0.0
+
+                # output looks like this:
+                # ncalls  tottime  percall  cumtime  percall filename:lineno(function)
+                # 1    0.000    0.000    0.000    0.000 <function>
+                for line in f:
+                    if " function calls " in line and " in " in line and " seconds":
+                        # 304307 function calls (291692 primitive calls) in 1.031 seconds
+                        assert total == 0
+                        total = float(line.split()[-2])
+                        continue
+                    columns = line.split(None, 5)
+                    if len(columns) < 6 or columns[0] == "ncalls":
+                        continue
+
+                    # TODO: to support windows and MacOS, extend this to a list of function known to sleep the process
+                    # e.g. WaitForMultipleObjects or kqueue
+                    if (
+                        "{method 'poll' of 'select.epoll' objects}" in columns[5]
+                        or "method 'control' of 'select.kqueue' objects" in columns[5]
+                        or "method _overlapped.GetQueuedCompletionStatus" in columns[5]
+                    ):
+                        # cumulative time
+                        sleep += float(columns[3])
+
+                if sleep < 0.000001:
+                    percent = 100.0
+                else:
+                    percent = 100.0 * (total - sleep) / total
+
+                if percent > 90:
+                    color = Fore.RED + Style.BRIGHT
+                elif percent > 80:
+                    color = Fore.MAGENTA + Style.BRIGHT
+                elif percent > 70:
+                    color = Fore.YELLOW + Style.BRIGHT
+                elif percent > 60:
+                    color = Style.BRIGHT
+                elif percent < 10:
+                    color = Fore.GREEN
+                else:
+                    color = ""
+
+                quantized = int(percent // 2)
+                print(
+                    ("%05d: " + color + "%3.0f%% CPU " + Back.WHITE + "%s" + Style.RESET_ALL + "%s|")
+                    % (counter, percent, " " * quantized, " " * (50 - quantized))
+                )
+
+                counter += 1
+        except Exception as e:
+            print(e)
+
+    def analyze_slot_range(profile_dir: pathlib.Path, first: int, last: int) -> None:
+        if last < first:
+            print("ERROR: first must be <= last when specifying slot range")
+            return
+
+        files = []
+        for i in range(first, last + 1):
+            files.append(str(profile_dir / f"slot-{i:05}.profile"))
+
+        output_file = f"chia-hotspot-{first}"
+        if first < last:
+            output_file += f"-{last}"
+
+        print(f"generating call tree for slot(s) [{first}, {last}]")
+        check_call(["gprof2dot", "-f", "pstats", "-o", output_file + ".dot", *files])
+        with open(output_file + ".png", "w+") as f:
+            check_call(["dot", "-T", "png", output_file + ".dot"], stdout=f)
+        print(f"output written to: {output_file}.png")
+
+    if len(sys.argv) == 2:
+        # this analyzes the CPU usage at all slots saved to the profiler directory
+        analyze_cpu_usage(profile_dir)
+    elif len(sys.argv) in {3, 4}:
+        # the additional arguments are interpreted as either one slot, or a
+        # slot range (first and last) to analyze
+        first = int(sys.argv[2])
+        last = int(sys.argv[3]) if len(sys.argv) == 4 else first
+        analyze_slot_range(profile_dir, first, last)
+    else:
+        print(
+            """USAGE:
+profiler.py <profile-directory>
+    Analyze CPU usage at each 1 second interval from the profiles in the specified
+    directory. Print colored timeline to stdout
+profiler.py <profile-directory> <slot>
+profiler.py <profile-directory> <first-slot> <last-slot>
+    Analyze a single slot, or a range of time slots, from the profile directory
+"""
+        )
+
+
+async def mem_profile_task(root_path: pathlib.Path, service: str, log: logging.Logger) -> None:
+    profile_dir = path_from_root(root_path, f"memory-profile-{service}") / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    log.info("Starting memory profiler. saving to %s", profile_dir)
+    profile_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        tracemalloc.start(30)
+
+        counter = 0
+
+        while True:
+            # this will throw CancelledError when we're exiting
+            await asyncio.sleep(60)
+            snapshot = tracemalloc.take_snapshot()
+            snapshot.dump(str(profile_dir / f"heap-{counter:05d}.profile"))
+            log.info(f"Heap usage: {tracemalloc.get_traced_memory()[0] / 1000000:0.3f} MB profile {counter:05d}")
+            counter += 1
+    finally:
+        tracemalloc.stop()
+
+
+@asynccontextmanager
+async def enable_profiler(profile: bool) -> AsyncIterator[Optional[cProfile.Profile]]:
+    if not profile:
+        yield None
+        return
+
+    # this is not covered by any unit tests as it's essentially test code
+    # itself. It's exercised manually when investigating performance issues
+    with cProfile.Profile() as pr:  # pragma: no cover
+        pr.enable()
+        yield pr
+        pr.disable()

chia/util/recursive_replace.py

@@ -0,0 +1,24 @@
+# Package: utils
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import Any
+
+
+def recursive_replace(root_obj: Any, replace_str: str, replace_with: Any) -> Any:
+    split_str = replace_str.split(".")
+    if len(split_str) == 1:
+        # This check is here to support native types (implemented in Rust
+        # in chia_rs) that aren't dataclasses. They instead implement a
+        # replace() method in their python bindings.
+        if hasattr(root_obj, "replace"):
+            return root_obj.replace(**{split_str[0]: replace_with})
+        else:
+            return replace(root_obj, **{split_str[0]: replace_with})
+    sub_obj = recursive_replace(getattr(root_obj, split_str[0]), ".".join(split_str[1:]), replace_with)
+    # See comment above
+    if hasattr(root_obj, "replace"):
+        return root_obj.replace(**{split_str[0]: sub_obj})
+    else:
+        return replace(root_obj, **{split_str[0]: sub_obj})

chia/util/safe_cancel_task.py

@@ -0,0 +1,16 @@
+# Package: utils
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Optional
+
+
+def cancel_task_safe(task: Optional[asyncio.Task[None]], log: Optional[logging.Logger] = None) -> None:
+    if task is not None:
+        try:
+            task.cancel()
+        except Exception as e:
+            if log is not None:
+                log.error(f"Error while canceling task.{e} {task}")

chia/util/service_groups.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from collections.abc import Generator, Iterable, KeysView
+
+SERVICES_FOR_GROUP: dict[str, list[str]] = {
+    "all": [
+        "chia_harvester",
+        "chia_timelord_launcher",
+        "chia_timelord",
+        "chia_farmer",
+        "chia_full_node",
+        "chia_wallet",
+        "chia_data_layer",
+        "chia_data_layer_http",
+    ],
+    "daemon": [],
+    # TODO: should this be `data_layer`?
+    "data": ["chia_wallet", "chia_data_layer"],
+    "data_layer_http": ["chia_data_layer_http"],
+    "node": ["chia_full_node"],
+    "harvester": ["chia_harvester"],
+    "farmer": ["chia_harvester", "chia_farmer", "chia_full_node", "chia_wallet"],
+    "farmer-no-wallet": ["chia_harvester", "chia_farmer", "chia_full_node"],
+    "farmer-only": ["chia_farmer"],
+    "timelord": ["chia_timelord_launcher", "chia_timelord", "chia_full_node"],
+    "timelord-only": ["chia_timelord"],
+    "timelord-launcher-only": ["chia_timelord_launcher"],
+    "wallet": ["chia_wallet"],
+    "introducer": ["chia_introducer"],
+    "simulator": ["chia_full_node_simulator"],
+    "crawler": ["chia_crawler"],
+    "seeder": ["chia_crawler", "chia_seeder"],
+    "seeder-only": ["chia_seeder"],
+}
+
+
+def all_groups() -> KeysView[str]:
+    return SERVICES_FOR_GROUP.keys()
+
+
+def services_for_groups(groups: Iterable[str]) -> Generator[str, None, None]:
+    for group in groups:
+        yield from SERVICES_FOR_GROUP[group]
+
+
+def validate_service(service: str) -> bool:
+    return any(service in _ for _ in SERVICES_FOR_GROUP.values())