chia-blockchain 2.5.1rc1__py3-none-any.whl

chia/server/address_manager_store.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from timeit import default_timer as timer
+from typing import Any, Optional
+
+import aiofiles
+
+from chia.server.address_manager import (
+    BUCKET_SIZE,
+    NEW_BUCKET_COUNT,
+    NEW_BUCKETS_PER_ADDRESS,
+    AddressManager,
+    ExtendedPeerInfo,
+)
+from chia.util.files import write_file_async
+from chia.util.ints import uint64
+from chia.util.streamable import Streamable, streamable
+
+log = logging.getLogger(__name__)
+
+
+@streamable
+@dataclass(frozen=True)
+class PeerDataSerialization(Streamable):
+    """
+    Serializable property bag for the peer data that was previously stored in sqlite.
+    """
+
+    metadata: list[tuple[str, str]]
+    nodes: list[tuple[uint64, str]]
+    new_table: list[tuple[uint64, uint64]]
+
+
+async def makePeerDataSerialization(
+    metadata: list[tuple[str, Any]], nodes: list[tuple[int, ExtendedPeerInfo]], new_table: list[tuple[int, int]]
+) -> bytes:
+    """
+    Create a PeerDataSerialization, adapting the provided collections
+    """
+    transformed_nodes: list[tuple[uint64, str]] = []
+    transformed_new_table: list[tuple[uint64, uint64]] = []
+
+    for index, [node_id, peer_info] in enumerate(nodes):
+        transformed_nodes.append((uint64(node_id), peer_info.to_string()))
+        # Come up to breathe for a moment
+        if index % 1000 == 0:
+            await asyncio.sleep(0)
+
+    for index, [node_id, bucket_id] in enumerate(new_table):
+        transformed_new_table.append((uint64(node_id), uint64(bucket_id)))
+        # Come up to breathe for a moment
+        if index % 1000 == 0:
+            await asyncio.sleep(0)
+
+    serialized_bytes: bytes = bytes(PeerDataSerialization(metadata, transformed_nodes, transformed_new_table))
+    return serialized_bytes
+
+
+class AddressManagerStore:
+    """
+    Metadata table:
+    - private key
+    - new table count
+    - tried table count
+    Nodes table:
+    * Maps entries from new/tried table to unique node ids.
+    - node_id
+    - IP, port, together with the IP, port of the source peer.
+    New table:
+    * Stores node_id, bucket for each occurrence in the new table of an entry.
+    * Once we know the buckets, we can also deduce the bucket positions.
+    Every other information, such as tried_matrix, map_addr, map_info, random_pos,
+    be deduced and it is not explicitly stored, instead it is recalculated.
+    """
+
+    @classmethod
+    async def create_address_manager(cls, peers_file_path: Path) -> AddressManager:
+        """
+        Create an address manager using data deserialized from a peers file.
+        """
+        address_manager: Optional[AddressManager] = None
+        if peers_file_path.exists():
+            try:
+                log.info(f"Loading peers from {peers_file_path}")
+                address_manager = await cls._deserialize(peers_file_path)
+            except Exception:
+                log.exception(f"Unable to create address_manager from {peers_file_path}")
+
+        if address_manager is None:
+            log.info("Creating new address_manager")
+            address_manager = AddressManager()
+
+        return address_manager
+
+    @classmethod
+    async def serialize(cls, address_manager: AddressManager, peers_file_path: Path) -> None:
+        """
+        Serialize the address manager's peer data to a file.
+        """
+        metadata: list[tuple[str, str]] = []
+        nodes: list[tuple[int, ExtendedPeerInfo]] = []
+        new_table_entries: list[tuple[int, int]] = []
+        unique_ids: dict[int, int] = {}
+        count_ids: int = 0
+
+        log.info("Serializing peer data")
+        metadata.append(("key", str(address_manager.key)))
+
+        for node_id, info in address_manager.map_info.items():
+            unique_ids[node_id] = count_ids
+            if info.ref_count > 0:
+                assert count_ids != address_manager.new_count
+                nodes.append((count_ids, info))
+                count_ids += 1
+        metadata.append(("new_count", str(count_ids)))
+
+        tried_ids = 0
+        for node_id, info in address_manager.map_info.items():
+            if info.is_tried:
+                assert info is not None
+                assert tried_ids != address_manager.tried_count
+                nodes.append((count_ids, info))
+                count_ids += 1
+                tried_ids += 1
+        metadata.append(("tried_count", str(tried_ids)))
+
+        for bucket in range(NEW_BUCKET_COUNT):
+            for i in range(BUCKET_SIZE):
+                if address_manager.new_matrix[bucket][i] != -1:
+                    index = unique_ids[address_manager.new_matrix[bucket][i]]
+                    new_table_entries.append((index, bucket))
+
+        try:
+            # Ensure the parent directory exists
+            peers_file_path.parent.mkdir(parents=True, exist_ok=True)
+            start_time = timer()
+            await cls._write_peers(peers_file_path, metadata, nodes, new_table_entries)
+            log.debug(f"Serializing peer data took {timer() - start_time} seconds")
+        except Exception:
+            log.exception(f"Failed to write peer data to {peers_file_path}")
+
+    @classmethod
+    async def _deserialize(cls, peers_file_path: Path) -> AddressManager:
+        """
+        Create an address manager using data deserialized from a peers file.
+        """
+        peer_data: Optional[PeerDataSerialization] = None
+        address_manager = AddressManager()
+        start_time = timer()
+        try:
+            peer_data = await cls._read_peers(peers_file_path)
+        except Exception:
+            log.exception(f"Unable to deserialize peers from {peers_file_path}")
+
+        if peer_data is not None:
+            metadata: dict[str, str] = {key: value for key, value in peer_data.metadata}
+            nodes: list[tuple[int, ExtendedPeerInfo]] = [
+                (node_id, ExtendedPeerInfo.from_string(info_str)) for node_id, info_str in peer_data.nodes
+            ]
+            new_table_entries: list[tuple[int, int]] = [(node_id, bucket) for node_id, bucket in peer_data.new_table]
+            log.debug(f"Deserializing peer data took {timer() - start_time} seconds")
+
+            address_manager.key = int(metadata["key"])
+            address_manager.new_count = int(metadata["new_count"])
+            # address_manager.tried_count = int(metadata["tried_count"])
+            address_manager.tried_count = 0
+
+            new_table_nodes = [(node_id, info) for node_id, info in nodes if node_id < address_manager.new_count]
+            for n, info in new_table_nodes:
+                address_manager.map_addr[info.peer_info.host] = n
+                address_manager.map_info[n] = info
+                info.random_pos = len(address_manager.random_pos)
+                address_manager.random_pos.append(n)
+            address_manager.id_count = len(new_table_nodes)
+            tried_table_nodes = [(node_id, info) for node_id, info in nodes if node_id >= address_manager.new_count]
+            # lost_count = 0
+            for node_id, info in tried_table_nodes:
+                tried_bucket = info.get_tried_bucket(address_manager.key)
+                tried_bucket_pos = info.get_bucket_position(address_manager.key, False, tried_bucket)
+                if address_manager.tried_matrix[tried_bucket][tried_bucket_pos] == -1:
+                    info.random_pos = len(address_manager.random_pos)
+                    info.is_tried = True
+                    id_count = address_manager.id_count
+                    address_manager.random_pos.append(id_count)
+                    address_manager.map_info[id_count] = info
+                    address_manager.map_addr[info.peer_info.host] = id_count
+                    address_manager.tried_matrix[tried_bucket][tried_bucket_pos] = id_count
+                    address_manager.id_count += 1
+                    address_manager.tried_count += 1
+                # else:
+                #    lost_count += 1
+
+            # address_manager.tried_count -= lost_count
+            for node_id, bucket in new_table_entries:
+                if node_id >= 0 and node_id < address_manager.new_count:
+                    info = address_manager.map_info[node_id]
+                    bucket_pos = info.get_bucket_position(address_manager.key, True, bucket)
+                    if (
+                        address_manager.new_matrix[bucket][bucket_pos] == -1
+                        and info.ref_count < NEW_BUCKETS_PER_ADDRESS
+                    ):
+                        info.ref_count += 1
+                        address_manager.new_matrix[bucket][bucket_pos] = node_id
+
+            for node_id, info in list(address_manager.map_info.items()):
+                if not info.is_tried and info.ref_count == 0:
+                    address_manager.delete_new_entry_(node_id)
+
+            address_manager.load_used_table_positions()
+
+        return address_manager
+
+    @classmethod
+    async def _read_peers(cls, peers_file_path: Path) -> PeerDataSerialization:
+        """
+        Read the peers file and return the data as a PeerDataSerialization object.
+        """
+        async with aiofiles.open(peers_file_path, "rb") as f:
+            return PeerDataSerialization.from_bytes(await f.read())
+
+    @classmethod
+    async def _write_peers(
+        cls,
+        peers_file_path: Path,
+        metadata: list[tuple[str, Any]],
+        nodes: list[tuple[int, ExtendedPeerInfo]],
+        new_table: list[tuple[int, int]],
+    ) -> None:
+        """
+        Serializes the given peer data and writes it to the peers file.
+        """
+        serialized_bytes: bytes = await makePeerDataSerialization(metadata, nodes, new_table)
+        await write_file_async(peers_file_path, serialized_bytes, file_mode=0o644)

chia/server/api_protocol.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import functools
+import logging
+from collections.abc import Awaitable
+from dataclasses import dataclass, field
+from logging import Logger
+from typing import Callable, ClassVar, Optional, TypeVar, Union, final, get_type_hints
+
+from typing_extensions import Concatenate, ParamSpec, Protocol
+
+from chia.protocols.protocol_message_types import ProtocolMessageTypes
+from chia.server.outbound_message import Message
+from chia.util.streamable import Streamable
+
+
+class ApiProtocol(Protocol):
+    log: Logger
+    metadata: ClassVar[ApiMetadata]
+
+    def ready(self) -> bool: ...
+
+
+log = logging.getLogger(__name__)
+P = ParamSpec("P")
+R = TypeVar("R", bound=Awaitable[Optional[Message]])
+S = TypeVar("S", bound=Streamable)
+Self = TypeVar("Self")
+api_attribute_name = "_chia_api"
+
+
+@dataclass
+class ApiRequest:
+    request_type: ProtocolMessageTypes
+    message_class: type[Streamable]
+    method: Callable[..., Awaitable[Optional[Message]]]
+    peer_required: bool = False
+    bytes_required: bool = False
+    execute_task: bool = False
+    reply_types: list[ProtocolMessageTypes] = field(default_factory=list)
+
+
+@final
+@dataclass
+class ApiMetadata:
+    message_type_to_request: dict[ProtocolMessageTypes, ApiRequest] = field(default_factory=dict)
+
+    @classmethod
+    def copy(cls, original: ApiMetadata) -> ApiMetadata:
+        return cls(message_type_to_request=dict(original.message_type_to_request))
+
+    @classmethod
+    def from_bound_method(cls, method: Callable[..., Awaitable[Optional[Message]]]) -> ApiRequest:
+        self: ApiMetadata = getattr(method, api_attribute_name)
+        message_type = ProtocolMessageTypes[method.__name__]
+        return self.message_type_to_request[message_type]
+
+    # TODO: This hinting does not express that the returned callable *_bytes parameter
+    #       corresponding to the first parameter name will be filled in by the wrapper.
+    def request(
+        self,
+        peer_required: bool = False,
+        bytes_required: bool = False,
+        execute_task: bool = False,
+        reply_types: Optional[list[ProtocolMessageTypes]] = None,
+        request_type: Optional[ProtocolMessageTypes] = None,
+    ) -> Callable[[Callable[Concatenate[Self, S, P], R]], Callable[Concatenate[Self, Union[bytes, S], P], R]]:
+        non_optional_reply_types: list[ProtocolMessageTypes]
+        if reply_types is None:
+            non_optional_reply_types = []
+        else:
+            non_optional_reply_types = reply_types
+
+        def inner(f: Callable[Concatenate[Self, S, P], R]) -> Callable[Concatenate[Self, Union[bytes, S], P], R]:
+            @functools.wraps(f)
+            def wrapper(self: Self, original: Union[bytes, S], *args: P.args, **kwargs: P.kwargs) -> R:
+                arg: S
+                if isinstance(original, bytes):
+                    if request.bytes_required:
+                        kwargs[message_name_bytes] = original
+                    arg = message_class.from_bytes(original)
+                else:
+                    arg = original
+                    if request.bytes_required:
+                        kwargs[message_name_bytes] = bytes(original)
+
+                return f(self, arg, *args, **kwargs)
+
+            setattr(wrapper, api_attribute_name, self)
+            message_name, message_class = next(
+                (name, hint) for name, hint in get_type_hints(f).items() if name not in {"self", "peer", "return"}
+            )
+            message_name_bytes = f"{message_name}_bytes"
+
+            nonlocal request_type
+            if request_type is None:
+                request_type = ProtocolMessageTypes[f.__name__]
+
+            request = ApiRequest(
+                request_type=request_type,
+                peer_required=peer_required,
+                bytes_required=bytes_required,
+                execute_task=execute_task,
+                reply_types=non_optional_reply_types,
+                message_class=message_class,
+                method=wrapper,
+            )
+
+            if request_type in self.message_type_to_request:
+                raise Exception(f"request type already registered: {request_type}")
+
+            self.message_type_to_request[request_type] = request
+
+            return wrapper
+
+        return inner

chia/server/capabilities.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from chia.protocols.shared_protocol import Capability
+from chia.util.ints import uint16
+
+
+def known_active_capabilities(values: Iterable[tuple[uint16, str]]) -> list[Capability]:
+    # NOTE: order is not guaranteed
+    # TODO: what if there's a claim for both supporting and not?
+    #       presently it considers it supported
+    filtered: set[Capability] = set()
+    for value, state in values:
+        if state != "1":
+            continue
+
+        try:
+            filtered.add(Capability(value))
+        except ValueError:
+            pass
+
+    # TODO: consider changing all uses to sets instead of lists
+    return list(filtered)