royalnet/royalherald/link.py

import asyncio
import websockets
import uuid
import functools
import logging as _logging
import typing
from .package import Package
from .request import Request
from .response import Response, ResponseSuccess, ResponseFailure
from .broadcast import Broadcast
from .errors import ConnectionClosedError, InvalidServerResponseError


log = _logging.getLogger(__name__)


class PendingRequest:
    def __init__(self, *, loop: asyncio.AbstractEventLoop = None):
        if loop is None:
            self.loop = asyncio.get_event_loop()
        else:
            self.loop = loop
        self.event: asyncio.Event = asyncio.Event(loop=loop)
        self.data: typing.Optional[dict] = None

    def __repr__(self):
        if self.event.is_set():
            return f"<{self.__class__.__qualname__}: {self.data.__class__.__name__}>"
        return f"<{self.__class__.__qualname__}>"

    def set(self, data):
        self.data = data
        self.event.set()


def requires_connection(func):
    @functools.wraps(func)
    async def new_func(self, *args, **kwargs):
        await self.connect_event.wait()
        return await func(self, *args, **kwargs)
    return new_func


def requires_identification(func):
    @functools.wraps(func)
    async def new_func(self, *args, **kwargs):
        await self.identify_event.wait()
        return await func(self, *args, **kwargs)
    return new_func


class Link:
    def __init__(self, master_uri: str, secret: str, link_type: str, request_handler, *,
                 loop: asyncio.AbstractEventLoop = None):
        if ":" in link_type:
            raise ValueError("Link types cannot contain colons.")
        self.master_uri: str = master_uri
        self.link_type: str = link_type
        self.nid: str = str(uuid.uuid4())
        self.secret: str = secret
        self.websocket: typing.Optional[websockets.WebSocketClientProtocol] = None
        # Not sure on the type annotation here
        self.request_handler: typing.Callable[[Request], typing.Awaitable[Response]] = request_handler
        self._pending_requests: typing.Dict[str, PendingRequest] = {}
        if loop is None:
            self._loop = asyncio.get_event_loop()
        else:
            self._loop = loop
        self.error_event: asyncio.Event = asyncio.Event(loop=self._loop)
        self.connect_event: asyncio.Event = asyncio.Event(loop=self._loop)
        self.identify_event: asyncio.Event = asyncio.Event(loop=self._loop)

    def __repr__(self):
        if self.identify_event.is_set():
            return f"<{self.__class__.__qualname__} (identified)>"
        elif self.connect_event.is_set():
            return f"<{self.__class__.__qualname__} (connected)>"
        elif self.error_event.is_set():
            return f"<{self.__class__.__qualname__} (error)>"
        else:
            return f"<{self.__class__.__qualname__} (disconnected)>"

    async def connect(self):
        """Connect to the :py:class:`royalnet.network.NetworkServer` at ``self.master_uri``."""
        log.info(f"Connecting to {self.master_uri}...")
        self.websocket = await websockets.connect(self.master_uri, loop=self._loop)
        self.connect_event.set()
        log.info(f"Connected!")

    @requires_connection
    async def receive(self) -> Package:
        """Recieve a :py:class:`Package` from the :py:class:`Server`.

        Raises:
            :py:exc:`royalnet.network.royalnetlink.ConnectionClosedError` if the connection closes."""
        try:
            jbytes: bytes = await self.websocket.recv()
            package: Package = Package.from_json_bytes(jbytes)
        except websockets.ConnectionClosed:
            self.error_event.set()
            self.connect_event.clear()
            self.identify_event.clear()
            log.info(f"Connection to {self.master_uri} was closed.")
            # What to do now? Let's just reraise.
            raise ConnectionClosedError()
        if self.identify_event.is_set() and package.destination != self.nid:
            raise InvalidServerResponseError("Package is not addressed to this NetworkLink.")
        log.debug(f"Received package: {package}")
        return package

    @requires_connection
    async def identify(self) -> None:
        log.info(f"Identifying to {self.master_uri}...")
        await self.websocket.send(f"Identify {self.nid}:{self.link_type}:{self.secret}")
        response: Package = await self.receive()
        if not response.source == "<server>":
            raise InvalidServerResponseError("Received a non-service package before identification.")
        if "type" not in response.data:
            raise InvalidServerResponseError("Missing 'type' in response data")
        if response.data["type"] == "error":
            raise ConnectionClosedError(f"Identification error: {response.data['type']}")
        assert response.data["type"] == "success"
        self.identify_event.set()
        log.info(f"Identified successfully!")

    @requires_identification
    async def send(self, package: Package):
        await self.websocket.send(package.to_json_bytes())
        log.debug(f"Sent package: {package}")

    @requires_identification
    async def broadcast(self, destination: str, broadcast: Broadcast) -> None:
        package = Package(broadcast.to_dict(), source=self.nid, destination=destination)
        await self.send(package)
        log.debug(f"Sent broadcast: {broadcast}")

    @requires_identification
    async def request(self, destination: str, request: Request) -> Response:
        if destination.startswith("*"):
            raise ValueError("requests cannot have multiple destinations")
        package = Package(request.to_dict(), source=self.nid, destination=destination)
        request = PendingRequest(loop=self._loop)
        self._pending_requests[package.source_conv_id] = request
        await self.send(package)
        log.debug(f"Sent request to {destination}: {request}")
        await request.event.wait()
        if request.data["type"] == "ResponseSuccess":
            response: Response = ResponseSuccess.from_dict(request.data)
        elif request.data["type"] == "ResponseFailure":
            response: Response = ResponseFailure.from_dict(request.data)
        else:
            raise TypeError("Unknown response type")
        log.debug(f"Received from {destination}: {request} -> {response}")
        return response

    async def run(self):
        """Blockingly run the Link."""
        log.debug(f"Running main client loop for {self.nid}.")
        if self.error_event.is_set():
            raise ConnectionClosedError("RoyalnetLinks can't be rerun after an error.")
        while True:
            if not self.connect_event.is_set():
                await self.connect()
            if not self.identify_event.is_set():
                await self.identify()
            package: Package = await self.receive()
            # Package is a response
            if package.destination_conv_id in self._pending_requests:
                request = self._pending_requests[package.destination_conv_id]
                request.set(package.data)
                continue
            # Package is a request
            assert isinstance(package, Package)
            log.debug(f"Received request {package.source_conv_id}: {package}")
            response: Response = await self.request_handler(Request.from_dict(package.data))
            response_package: Package = package.reply(response.to_dict())
            await self.send(response_package)
            log.debug(f"Replied to request {response_package.source_conv_id}: {response_package}")