chia-blockchain 2.5.1rc1__py3-none-any.whl

chia/util/async_pool.py

@@ -0,0 +1,226 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import dataclasses
+import itertools
+import logging
+import traceback
+from collections.abc import AsyncIterator, Iterator
+from typing import Generic, Optional, Protocol, TypeVar, final
+
+import anyio
+
+from chia.util.log_exceptions import log_exceptions
+from chia.util.task_referencer import create_referenced_task
+
+
+class InvalidTargetWorkerCountError(Exception):
+    def __init__(self, o: object) -> None:
+        super().__init__(f"target worker count must be one or greater: {o!r}")
+
+
+class WorkerCallable(Protocol):
+    async def __call__(self, worker_id: int) -> object: ...
+
+
+J = TypeVar("J")
+R = TypeVar("R")
+T = TypeVar("T")
+T_co = TypeVar("T_co", covariant=True)
+T_contra = TypeVar("T_contra", contravariant=True)
+
+
+class QueuedWorkerCallable(Protocol[T, T_co]):
+    async def __call__(self, worker_id: int, job: Job[T]) -> T_co: ...
+
+
+class JobQueueProtocol(Protocol[T_co]):
+    async def get(self) -> T_co: ...
+
+
+class ResultQueueProtocol(Protocol[T_contra]):
+    async def put(self, item: T_contra) -> None: ...
+
+
+# TODO: how does this compare to just using a future
+@dataclasses.dataclass
+class Job(Generic[T]):
+    input: T
+    started: asyncio.Event = dataclasses.field(default_factory=asyncio.Event)
+    done: asyncio.Event = dataclasses.field(default_factory=asyncio.Event)
+    exception: Optional[BaseException] = None
+    task: Optional[asyncio.Task[object]] = None
+    cancelled: bool = False
+
+
+@final
+@dataclasses.dataclass
+class QueuedAsyncPool(Generic[J, R]):
+    name: str
+    job_queue: JobQueueProtocol[Job[J]]
+    result_queue: Optional[ResultQueueProtocol[R]]
+    worker_async_callable: QueuedWorkerCallable[J, R]
+
+    @classmethod
+    @contextlib.asynccontextmanager
+    async def managed(
+        cls,
+        name: str,
+        job_queue: JobQueueProtocol[Job[J]],
+        worker_async_callable: QueuedWorkerCallable[J, R],
+        target_worker_count: int,
+        result_queue: Optional[ResultQueueProtocol[R]] = None,
+        log: logging.Logger = logging.getLogger(__name__),
+    ) -> AsyncIterator[QueuedAsyncPool[J, R]]:
+        self = cls(
+            name=name,
+            job_queue=job_queue,
+            result_queue=result_queue,
+            worker_async_callable=worker_async_callable,
+        )
+
+        async with AsyncPool.managed(
+            name=self.name,
+            worker_async_callable=self.worker,
+            target_worker_count=target_worker_count,
+            log=log,
+        ):
+            yield self
+
+    async def get_job(self) -> Job[J]:
+        while True:
+            job = await self.job_queue.get()
+            if not job.cancelled:
+                # TODO: can the job just be removed from the queue?
+                return job
+
+    async def worker(self, worker_id: int) -> None:
+        job = await self.get_job()
+        job.task = asyncio.current_task()
+        job.started.set()
+
+        try:
+            result = await self.worker_async_callable(worker_id=worker_id, job=job)
+        except BaseException as e:
+            # TODO: can't you not raise the same exception twice so this has to be
+            #       just reference and is all...  well, i dunno.
+            job.exception = e
+            raise
+        else:
+            if self.result_queue is not None:
+                await self.result_queue.put(result)
+        finally:
+            job.done.set()
+
+    def cancel(self, job: Job[J]) -> None:
+        # TODO: should this just be on the Job object?  or can we do something useful
+        #       here like get it out of the queue.
+        job.cancelled = True
+        if job.task is not None:
+            job.task.cancel()
+        # TODO: should this happen only after actual cancellation completes?
+        job.done.set()
+
+
+@final
+@dataclasses.dataclass
+class AsyncPool:
+    name: str
+    log: logging.Logger
+    worker_async_callable: WorkerCallable
+    _target_worker_count: int
+    _workers: dict[asyncio.Task[object], int] = dataclasses.field(init=False, default_factory=dict)
+    _worker_id_counter: Iterator[int] = dataclasses.field(init=False, default_factory=itertools.count)
+    _started: asyncio.Event = dataclasses.field(default_factory=asyncio.Event)
+    _single_use_used: bool = False
+
+    @classmethod
+    @contextlib.asynccontextmanager
+    async def managed(
+        cls,
+        name: str,
+        worker_async_callable: WorkerCallable,
+        target_worker_count: int,
+        log: logging.Logger = logging.getLogger(__name__),
+    ) -> AsyncIterator[AsyncPool]:
+        self = cls(
+            name=name,
+            log=log,
+            worker_async_callable=worker_async_callable,
+            _target_worker_count=target_worker_count,
+        )
+
+        if self._target_worker_count < 1:
+            raise InvalidTargetWorkerCountError(self._target_worker_count)
+
+        task = create_referenced_task(self._run(_check_single_use=False))
+        try:
+            # TODO: should this terminate if the run task ends?
+            await self._started.wait()
+            yield self
+        finally:
+            with anyio.CancelScope(shield=True):
+                task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
+
+    async def _run(self, *, _check_single_use: bool = True) -> None:
+        method_name = f"{type(self).__name__}._run()"
+
+        try:
+            while True:
+                with log_exceptions(
+                    log=self.log,
+                    consume=True,
+                    message=f"exception consumed while looping in {method_name} for {self.name!r}",
+                ):
+                    await self._run_single()
+        finally:
+            with anyio.CancelScope(shield=True):
+                with log_exceptions(
+                    log=self.log,
+                    consume=False,
+                    message=f"exception while tearing down in {method_name} for {self.name!r}",
+                ):
+                    await self._teardown_workers()
+
+    async def _run_single(self) -> None:
+        while len(self._workers) < self._target_worker_count:
+            new_worker_id = next(self._worker_id_counter)
+            new_worker = create_referenced_task(self.worker_async_callable(new_worker_id))
+            self.log.debug(f"{self.name}: adding worker {new_worker_id}")
+            self._workers[new_worker] = new_worker_id
+
+        self._started.set()
+
+        self.log.debug(f"{self.name}: waiting with {len(self._workers)} workers: {list(self._workers.values())}")
+        done_workers, pending_workers = await asyncio.wait(
+            self._workers,
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+        done_workers_by_id = {task: self._workers[task] for task in done_workers}
+        self._workers = {task: self._workers[task] for task in pending_workers}
+
+        for task, id in done_workers_by_id.items():
+            await self._handle_done_worker(task=task, id=id)
+
+    async def _teardown_workers(self) -> None:
+        while True:
+            try:
+                task, id = self._workers.popitem()
+            except KeyError:
+                break
+
+            task.cancel()
+            await self._handle_done_worker(task=task, id=id)
+
+    async def _handle_done_worker(self, task: asyncio.Task[object], id: int) -> None:
+        with anyio.CancelScope(shield=True):
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+            except Exception:
+                error_trace = traceback.format_exc()
+                self.log.error(f"{self.name}: worker {id} raised exception: {error_trace}")

chia/util/augmented_chain.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar, Optional, cast
+
+from chia.consensus.block_record import BlockRecord
+from chia.consensus.blockchain_interface import BlocksProtocol
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.full_block import FullBlock
+from chia.util.errors import Err
+from chia.util.ints import uint32
+
+
+class AugmentedBlockchain:
+    """
+    This class wraps a BlocksProtocol and forwards calls to it, when
+    looking up block records. It allows an in-memory cache of block records to
+    fall back onto in case a block is not available in the underlying
+    BlocksProtocol.
+    This is especially useful when validating blocks in parallel. The batch of
+    blocks will not have been added to the underlying blockchain until they've
+    all been validated, but the validation requires them to be available as-if
+    they were valid.
+    """
+
+    if TYPE_CHECKING:
+        _protocol_check: ClassVar[BlocksProtocol] = cast("AugmentedBlockchain", None)
+
+    _underlying: BlocksProtocol
+    _extra_blocks: dict[bytes32, tuple[FullBlock, BlockRecord]]
+    _height_to_hash: dict[uint32, bytes32]
+
+    def __init__(self, underlying: BlocksProtocol) -> None:
+        self._underlying = underlying
+        self._extra_blocks = {}
+        self._height_to_hash = {}
+
+    def _get_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
+        eb = self._extra_blocks.get(header_hash)
+        if eb is None:
+            return None
+        return eb[1]
+
+    def add_extra_block(self, block: FullBlock, block_record: BlockRecord) -> None:
+        assert block.header_hash == block_record.header_hash
+        self._extra_blocks[block_record.header_hash] = (block, block_record)
+        self._height_to_hash[block_record.height] = block_record.header_hash
+
+    def remove_extra_block(self, hh: bytes32) -> None:
+        if hh in self._extra_blocks:
+            block_record = self._extra_blocks.pop(hh)[1]
+            del self._height_to_hash[block_record.height]
+
+    # BlocksProtocol
+    async def lookup_block_generators(self, header_hash: bytes32, generator_refs: set[uint32]) -> dict[uint32, bytes]:
+        generators: dict[uint32, bytes] = {}
+
+        # traverse the additional blocks (if any) and resolve heights into
+        # generators
+        to_remove = []
+        curr: Optional[tuple[FullBlock, BlockRecord]] = self._extra_blocks.get(header_hash)
+        while curr is not None:
+            b = curr[0]
+            if b.height in generator_refs:
+                if b.transactions_generator is None:
+                    raise ValueError(Err.GENERATOR_REF_HAS_NO_GENERATOR)
+                generators[b.height] = bytes(b.transactions_generator)
+                to_remove.append(b.height)
+            header_hash = b.prev_header_hash
+            curr = self._extra_blocks.get(header_hash)
+        for i in to_remove:
+            generator_refs.remove(i)
+
+        if len(generator_refs) > 0:
+            generators.update(await self._underlying.lookup_block_generators(header_hash, generator_refs))
+        return generators
+
+    async def get_block_record_from_db(self, header_hash: bytes32) -> Optional[BlockRecord]:
+        ret = self._get_block_record(header_hash)
+        if ret is not None:
+            return ret
+        return await self._underlying.get_block_record_from_db(header_hash)
+
+    def add_block_record(self, block_record: BlockRecord) -> None:
+        self._underlying.add_block_record(block_record)
+
+        # now that we're adding the block to the underlying blockchain, we don't
+        # need to keep the extra block around anymore
+        hh = block_record.header_hash
+        if hh in self._extra_blocks:
+            del self._height_to_hash[block_record.height]
+            del self._extra_blocks[hh]
+
+    # BlockRecordsProtocol
+    def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
+        ret = self._get_block_record(header_hash)
+        if ret is not None:
+            return ret
+        return self._underlying.try_block_record(header_hash)
+
+    def block_record(self, header_hash: bytes32) -> BlockRecord:
+        ret = self._get_block_record(header_hash)
+        if ret is not None:
+            return ret
+        return self._underlying.block_record(header_hash)
+
+    def height_to_block_record(self, height: uint32) -> BlockRecord:
+        header_hash = self._height_to_hash.get(height)
+        if header_hash is not None:
+            ret = self._get_block_record(header_hash)
+            if ret is not None:
+                return ret
+        return self._underlying.height_to_block_record(height)
+
+    def height_to_hash(self, height: uint32) -> Optional[bytes32]:
+        ret = self._height_to_hash.get(height)
+        if ret is not None:
+            return ret
+        return self._underlying.height_to_hash(height)
+
+    def contains_block(self, header_hash: bytes32) -> bool:
+        return (header_hash in self._extra_blocks) or self._underlying.contains_block(header_hash)
+
+    def contains_height(self, height: uint32) -> bool:
+        return (height in self._height_to_hash) or self._underlying.contains_height(height)
+
+    async def prev_block_hash(self, header_hashes: list[bytes32]) -> list[bytes32]:
+        ret: list[bytes32] = []
+        for hh in header_hashes:
+            b = self._extra_blocks.get(hh)
+            if b is not None:
+                ret.append(b[1].prev_hash)
+            else:
+                ret.extend(await self._underlying.prev_block_hash([hh]))
+        return ret

chia/util/batches.py

@@ -0,0 +1,42 @@
+# Package: utils
+
+from __future__ import annotations
+
+from collections.abc import Collection, Iterator
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class Batch(Generic[T]):
+    remaining: int
+    entries: list[T]
+
+
+def to_batches(to_split: Collection[T], batch_size: int) -> Iterator[Batch[T]]:
+    if batch_size <= 0:
+        raise ValueError("to_batches: batch_size must be greater than 0.")
+    total_size = len(to_split)
+    if total_size == 0:
+        return
+
+    if isinstance(to_split, list):
+        for batch_start in range(0, total_size, batch_size):
+            batch_end = min(batch_start + batch_size, total_size)
+            yield Batch(total_size - batch_end, to_split[batch_start:batch_end])
+    elif isinstance(to_split, set):
+        processed = 0
+        entries = []
+        for entry in to_split:
+            entries.append(entry)
+            if len(entries) >= batch_size:
+                processed += len(entries)
+                yield Batch(total_size - processed, entries)
+                entries = []
+        if len(entries) > 0:
+            processed += len(entries)
+            yield Batch(total_size - processed, entries)
+    else:
+        raise ValueError(f"to_batches: Unsupported type {type(to_split)}")

chia/util/bech32m.py

@@ -0,0 +1,126 @@
+# Package: utils
+
+# Copyright (c) 2017 Pieter Wuille
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Optional
+
+from chia_rs.sized_bytes import bytes32
+
+# Based on this specification from Pieter Wuille:
+# https://github.com/sipa/bips/blob/bip-bech32m/bip-bech32m.mediawiki
+"""Reference implementation for Bech32m and segwit addresses."""
+
+
+CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
+
+
+def bech32_polymod(values: list[int]) -> int:
+    """Internal function that computes the Bech32 checksum."""
+    generator = [0x3B6A57B2, 0x26508E6D, 0x1EA119FA, 0x3D4233DD, 0x2A1462B3]
+    chk = 1
+    for value in values:
+        top = chk >> 25
+        chk = (chk & 0x1FFFFFF) << 5 ^ value
+        for i in range(5):
+            chk ^= generator[i] if ((top >> i) & 1) else 0
+    return chk
+
+
+def bech32_hrp_expand(hrp: str) -> list[int]:
+    """Expand the HRP into values for checksum computation."""
+    return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp]
+
+
+M = 0x2BC830A3
+
+
+def bech32_verify_checksum(hrp: str, data: list[int]) -> bool:
+    return bech32_polymod(bech32_hrp_expand(hrp) + data) == M
+
+
+def bech32_create_checksum(hrp: str, data: list[int]) -> list[int]:
+    values = bech32_hrp_expand(hrp) + data
+    polymod = bech32_polymod([*values, 0, 0, 0, 0, 0, 0]) ^ M
+    return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]
+
+
+def bech32_encode(hrp: str, data: list[int]) -> str:
+    """Compute a Bech32 string given HRP and data values."""
+    combined = data + bech32_create_checksum(hrp, data)
+    return hrp + "1" + "".join([CHARSET[d] for d in combined])
+
+
+def bech32_decode(bech: str, max_length: int = 90) -> tuple[Optional[str], Optional[list[int]]]:
+    """Validate a Bech32 string, and determine HRP and data."""
+    bech = bech.strip()
+    if (any(ord(x) < 33 or ord(x) > 126 for x in bech)) or (bech.lower() != bech and bech.upper() != bech):
+        return (None, None)
+    bech = bech.lower()
+    pos = bech.rfind("1")
+    if pos < 1 or pos + 7 > len(bech) or len(bech) > max_length:
+        return (None, None)
+    if not all(x in CHARSET for x in bech[pos + 1 :]):
+        return (None, None)
+    hrp = bech[:pos]
+    data = [CHARSET.find(x) for x in bech[pos + 1 :]]
+    if not bech32_verify_checksum(hrp, data):
+        return (None, None)
+    return hrp, data[:-6]
+
+
+def convertbits(data: Iterable[int], frombits: int, tobits: int, pad: bool = True) -> list[int]:
+    """General power-of-2 base conversion."""
+    acc = 0
+    bits = 0
+    ret = []
+    maxv = (1 << tobits) - 1
+    max_acc = (1 << (frombits + tobits - 1)) - 1
+    for value in data:
+        if value < 0 or (value >> frombits):
+            raise ValueError("Invalid Value")
+        acc = ((acc << frombits) | value) & max_acc
+        bits += frombits
+        while bits >= tobits:
+            bits -= tobits
+            ret.append((acc >> bits) & maxv)
+    if pad:
+        if bits:
+            ret.append((acc << (tobits - bits)) & maxv)
+    elif bits >= frombits or ((acc << (tobits - bits)) & maxv):
+        raise ValueError("Invalid bits")
+    return ret
+
+
+def encode_puzzle_hash(puzzle_hash: bytes32, prefix: str) -> str:
+    encoded = bech32_encode(prefix, convertbits(puzzle_hash, 8, 5))
+    return encoded
+
+
+def decode_puzzle_hash(address: str) -> bytes32:
+    _hrpgot, data = bech32_decode(address)
+    if data is None:
+        raise ValueError("Invalid Address")
+    decoded = convertbits(data, 5, 8, False)
+    decoded_bytes = bytes32(decoded)
+    return decoded_bytes

chia/util/beta_metrics.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import platform
+import socket
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+import psutil
+
+from chia.util.config import load_config
+from chia.util.cpu import available_logical_cores
+from chia.util.task_referencer import create_referenced_task
+
+log = logging.getLogger("beta")
+
+metrics_log_interval_default = 5
+metrics_log_interval_min = 1
+metrics_log_interval_max = 30
+
+
+def log_static_info() -> None:
+    log.debug(f"architecture: {platform.architecture()}")
+    log.debug(f"processor: {platform.processor()}")
+    log.debug(f"cpu count: {psutil.cpu_count()}")
+    log.debug(f"available logical cores: {available_logical_cores()}")
+    log.debug(f"machine: {platform.machine()}")
+    log.debug(f"platform: {platform.platform()}")
+
+
+def log_cpu_metrics() -> None:
+    log.debug(
+        f"CPU - percent: {psutil.cpu_percent(percpu=True)}, "
+        f"freq: {psutil.cpu_times(percpu=True)}, "
+        f"freq: {psutil.cpu_freq(percpu=True)}, "
+        f"load_avg: {psutil.getloadavg()}"
+    )
+
+
+def log_memory_metrics() -> None:
+    psutil.disk_io_counters(perdisk=False)
+    log.debug(f"MEMORY - virtual memory: {psutil.virtual_memory()}, swap: {psutil.swap_memory()}")
+
+
+def log_disk_metrics(root_path: Path, plot_dirs: list[str]) -> None:
+    # TODO, Could this spam the logs too much for large farms? Maybe don't log usage of plot dirs and
+    #       set perdisk=False rather for psutil.disk_io_counters? Lets try it with the default interval of 15s for now.
+    log.debug(f"DISK partitions: {psutil.disk_partitions()}")
+    for pot_dir in plot_dirs:
+        try:
+            usage = psutil.disk_usage(pot_dir)
+        except FileNotFoundError:
+            usage = "Directory not found"
+        log.debug(f"DISK - usage {pot_dir}: {usage}")
+    log.debug(f"DISK - usage root: {psutil.disk_usage(str(root_path))}")
+    log.debug(f"DISK - io counters: {psutil.disk_io_counters(perdisk=True)}")
+
+
+def log_port_states(config: dict[str, Any]) -> None:
+    selected_network = config["selected_network"]
+    full_node_port = config["network_overrides"]["config"][selected_network]["default_full_node_port"]
+    test_socket_ipv4 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    port_open_ipv4 = test_socket_ipv4.connect_ex(("127.0.0.1", full_node_port)) == 0
+    log.debug(f"full node port IPv4 [{full_node_port}]: {'open' if port_open_ipv4 else 'closed'}")
+    test_socket_ipv6 = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
+    port_open_ipv6 = test_socket_ipv6.connect_ex(("::1", full_node_port)) == 0
+    log.debug(f"full node port IPv6 [{full_node_port}]: {'open' if port_open_ipv6 else 'closed'}")
+
+
+def log_network_metrics() -> None:
+    log.debug(f"NETWORK: io counters: {psutil.net_io_counters(pernic=False)}")
+
+
+@dataclass
+class BetaMetricsLogger:
+    root_path: Path
+    task: Optional[asyncio.Task[None]] = None
+    stop_task: bool = False
+
+    def start_logging(self) -> None:
+        log.debug("start_logging")
+        log_static_info()
+        if self.task is not None:
+            raise RuntimeError("Already started")
+        self.stop_task = False
+        self.task = create_referenced_task(self.run())
+
+    async def stop_logging(self) -> None:
+        log.debug("stop_logging")
+        if self.task is None:
+            raise RuntimeError("Not yet started")
+
+        self.stop_task = True
+        await self.task
+        self.task = None
+
+    async def run(self) -> None:
+        config = load_config(self.root_path, "config.yaml")
+        interval = min(max(config["beta"]["metrics_log_interval"], metrics_log_interval_min), metrics_log_interval_max)
+        tick = 0
+        while not self.stop_task:
+            try:
+                tick += 1
+                # Log every interval
+                if tick % interval == 0:
+                    log_cpu_metrics()
+                    log_memory_metrics()
+                    log_network_metrics()
+                # Log after 10 intervals passed
+                if tick % (interval * 10) == 0:
+                    log_disk_metrics(self.root_path, config["harvester"]["plot_directories"])
+                    log_port_states(config)
+            except Exception as e:
+                log.warning(f"BetaMetricsLogger run failed: {e}")
+                await asyncio.sleep(10)
+            await asyncio.sleep(1)
+        log.debug("done")