freqtrade_origin/freqtrade/rpc/replicate/__init__.py

"""
This module manages replicate mode communication
"""
import asyncio
import logging
import secrets
import socket
from threading import Event, Thread
from typing import Any, Coroutine, Dict, Union

import websockets
from fastapi import Depends
from fastapi import WebSocket as FastAPIWebSocket
from fastapi import WebSocketDisconnect, status

from freqtrade.enums import LeaderMessageType, ReplicateModeType, RPCMessageType
from freqtrade.rpc import RPC, RPCHandler
from freqtrade.rpc.replicate.channel import ChannelManager
from freqtrade.rpc.replicate.thread_queue import Queue as ThreadedQueue
from freqtrade.rpc.replicate.utils import is_websocket_alive


logger = logging.getLogger(__name__)


class ReplicateController(RPCHandler):
    """  This class handles all websocket communication """

    def __init__(
        self,
        rpc: RPC,
        config: Dict[str, Any],
        api_server: Union[Any, None] = None
    ) -> None:
        """
        Init the ReplicateRPC class, and init the super class RPCHandler
        :param rpc: instance of RPC Helper class
        :param config: Configuration object
        :return: None
        """
        super().__init__(rpc, config)

        self.freqtrade = rpc._freqtrade
        self.api_server = api_server

        if not self.api_server:
            raise RuntimeError("The API server must be enabled for replicate to work")

        self._loop = None
        self._running = False
        self._thread = None
        self._queue = None

        self._stop_event = Event()
        self._follower_tasks = None

        self.channel_manager = ChannelManager()

        self.replicate_config = config.get('replicate', {})

        # What the config should look like
        # "replicate": {
        #     "enabled": true,
        #     "mode": "follower",
        #     "leaders": [
        #       {
        #         "url": "ws://localhost:8080/replicate/ws",
        #         "token": "test"
        #       }
        #     ]
        # }

        # "replicate": {
        #     "enabled": true,
        #     "mode": "leader",
        #     "api_key": "test"
        # }

        self.mode = ReplicateModeType[self.replicate_config.get('mode', 'leader').lower()]

        self.leaders_list = self.replicate_config.get('leaders', [])
        self.push_throttle_secs = self.replicate_config.get('push_throttle_secs', 1)

        self.reply_timeout = self.replicate_config.get('follower_reply_timeout', 10)
        self.ping_timeout = self.replicate_config.get('follower_ping_timeout', 2)
        self.sleep_time = self.replicate_config.get('follower_sleep_time', 1)

        if self.mode == ReplicateModeType.follower and len(self.leaders_list) == 0:
            raise ValueError("You must specify at least 1 leader in follower mode.")

        # This is only used by the leader, the followers use the tokens specified
        # in each of the leaders
        # If you do not specify an API key in the config, one will be randomly
        # generated and logged on startup
        default_api_key = secrets.token_urlsafe(16)
        self.secret_api_key = self.replicate_config.get('api_key', default_api_key)

        self.start_threaded_loop()

        self.start()

    def start_threaded_loop(self):
        """
        Start the main internal loop in another thread to run coroutines
        """
        self._loop = asyncio.new_event_loop()

        if not self._thread:
            self._thread = Thread(target=self._loop.run_forever)
            self._thread.start()
            self._running = True
        else:
            raise RuntimeError("A loop is already running")

    def submit_coroutine(self, coroutine: Coroutine):
        """
        Submit a coroutine to the threaded loop
        """
        if not self._running:
            raise RuntimeError("Cannot schedule new futures after shutdown")

        if not self._loop or not self._loop.is_running():
            raise RuntimeError("Loop must be started before any function can"
                               " be submitted")

        try:
            return asyncio.run_coroutine_threadsafe(coroutine, self._loop)
        except Exception as e:
            logger.error(f"Error running coroutine - {str(e)}")
            return None

    async def main_loop(self):
        """
        Main loop coro

        Start the loop based on what mode we're in
        """
        try:
            if self.mode == ReplicateModeType.leader:
                await self.leader_loop()
            elif self.mode == ReplicateModeType.follower:
                await self.follower_loop()

        except asyncio.CancelledError:
            pass
        finally:
            self._loop.stop()

    def start(self):
        """
        Start the controller main loop
        """
        self.submit_coroutine(self.main_loop())

    def cleanup(self) -> None:
        """
        Cleanup pending module resources.
        """
        if self._thread:
            if self._loop.is_running():

                self._running = False

                # Tell all coroutines submitted to the loop they're cancelled
                pending = asyncio.all_tasks(loop=self._loop)
                for task in pending:
                    task.cancel()

                self._loop.call_soon_threadsafe(self.channel_manager.disconnect_all)

            self._thread.join()

    def send_msg(self, msg: Dict[str, Any]) -> None:
        """
        Support RPC calls
        """
        if msg["type"] == RPCMessageType.EMIT_DATA:
            self.send_message(
                {
                    "data_type": msg.get("data_type"),
                    "data": msg.get("data")
                }
            )

    def send_message(self, msg: Dict[str, Any]) -> None:
        """ Push message through """

        # We should probably do some type of schema validation here

        if self.channel_manager.has_channels():
            self._send_message(msg)
        else:
            logger.debug("No listening followers, skipping...")
            pass

    def _send_message(self, msg: Dict[Any, Any]):
        """
        Add data to the internal queue to be broadcasted. This func will block
        if the queue is full. This is meant to be called in the main thread.
        """

        if self._queue:
            queue = self._queue.sync_q
            queue.put(msg)  # This will block if the queue is full
        else:
            logger.warning("Can not send data, leader loop has not started yet!")

    def is_leader(self):
        """
        Leader flag
        """
        return self.enabled() and self.mode == ReplicateModeType.leader

    def enabled(self):
        """
        Enabled flag
        """
        return self.replicate_config.get('enabled', False)

    # ----------------------- LEADER LOGIC ------------------------------

    async def leader_loop(self):
        """
        Main leader coroutine

        This starts all of the leader coros and registers the endpoint on
        the ApiServer
        """

        logger.info("Running rpc.replicate in Leader mode")
        logger.info("-" * 15)
        logger.info(f"API_KEY: {self.secret_api_key}")
        logger.info("-" * 15)

        self.register_leader_endpoint()

        try:
            await self._broadcast_queue_data()
        except Exception as e:
            logger.error("Exception occurred in Leader loop: ")
            logger.exception(e)

    async def _broadcast_queue_data(self):
        """
        Loop over queue data and broadcast it
        """
        # Instantiate the queue in this coroutine so it's attached to our loop
        self._queue = ThreadedQueue()
        async_queue = self._queue.async_q

        try:
            while self._running:
                # Get data from queue
                data = await async_queue.get()

                logger.info(f"Found data - broadcasting: {data}")

                # Broadcast it to everyone
                await self.channel_manager.broadcast(data)

                # Sleep
                await asyncio.sleep(self.push_throttle_secs)
        except asyncio.CancelledError:
            # Silently stop
            pass

    async def get_api_token(
        self,
        websocket: FastAPIWebSocket,
        token: Union[str, None] = None
    ):
        """
        Extract the API key from query param. Must match the
        set secret_api_key or the websocket connection will be closed.
        """
        if token == self.secret_api_key:
            return token
        else:
            logger.info("Denying websocket request...")
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)

    def register_leader_endpoint(self, path: str = "/replicate/ws"):
        """
        Attach and start the main leader loop to the ApiServer

        :param path: The endpoint path
        """

        if not self.api_server:
            raise RuntimeError("The leader needs the ApiServer to be active")

        # The endpoint function for running the main leader loop
        @self.api_server.app.websocket(path)
        async def leader_endpoint(
            websocket: FastAPIWebSocket,
            api_key: str = Depends(self.get_api_token)
        ):
            await self.leader_endpoint_loop(websocket)

    async def leader_endpoint_loop(self, websocket: FastAPIWebSocket):
        """
        The WebSocket endpoint served by the ApiServer. This handles connections,
        and adding them to the channel manager.
        """
        try:
            if is_websocket_alive(websocket):
                logger.info(f"Follower connected - {websocket.client}")
                channel = await self.channel_manager.on_connect(websocket)

                # Send initial data here
                # Data is being broadcasted right away as soon as startup,
                # we may not have to send initial data at all. Further testing
                # required.

                # Keep connection open until explicitly closed, and sleep
                try:
                    while not channel.is_closed():
                        await channel.recv()

                except WebSocketDisconnect:
                    # Handle client disconnects
                    logger.info(f"Follower disconnected - {websocket.client}")
                    await self.channel_manager.on_disconnect(websocket)
                except Exception as e:
                    logger.info(f"Follower connection failed - {websocket.client}")
                    logger.exception(e)
                    # Handle cases like -
                    # RuntimeError('Cannot call "send" once a closed message has been sent')
                    await self.channel_manager.on_disconnect(websocket)

        except Exception:
            logger.error(f"Failed to serve - {websocket.client}")
            await self.channel_manager.on_disconnect(websocket)

    # -------------------------------FOLLOWER LOGIC----------------------------

    async def follower_loop(self):
        """
        Main follower coroutine

        This starts all of the follower connection coros
        """

        logger.info("Starting rpc.replicate in Follower mode")

        try:
            results = await self._connect_to_leaders()
        except Exception as e:
            logger.error("Exception occurred in Follower loop: ")
            logger.exception(e)
        finally:
            for result in results:
                if isinstance(result, Exception):
                    logger.debug(f"Exception in Follower loop: {result}")

    async def _connect_to_leaders(self):
        """
        For each leader in `self.leaders_list` create a connection and
        listen for data.
        """
        rpc_lock = asyncio.Lock()

        logger.info("Starting connections to Leaders...")

        self.follower_tasks = [
            self._loop.create_task(self._handle_leader_connection(leader, rpc_lock))
            for leader in self.leaders_list
        ]
        return await asyncio.gather(*self.follower_tasks, return_exceptions=True)

    async def _handle_leader_connection(self, leader, lock):
        """
        Given a leader, connect and wait on data. If connection is lost,
        it will attempt to reconnect.
        """
        try:
            url, token = leader["url"], leader["token"]

            websocket_url = f"{url}?token={token}"

            logger.info(f"Attempting to connect to Leader at: {url}")
            # TODO: limit the amount of connection retries
            while True:
                try:
                    async with websockets.connect(websocket_url) as ws:
                        channel = await self.channel_manager.on_connect(ws)
                        while True:
                            try:
                                data = await asyncio.wait_for(
                                    channel.recv(),
                                    timeout=self.reply_timeout
                                )
                            except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
                                # We haven't received data yet. Check the connection and continue.
                                try:
                                    # ping
                                    ping = await channel.ping()
                                    await asyncio.wait_for(ping, timeout=self.ping_timeout)
                                    logger.debug(f"Connection to {url} still alive...")
                                    continue
                                except Exception:
                                    logger.info(
                                        f"Ping error {url} - retrying in {self.sleep_time}s")
                                    asyncio.sleep(self.sleep_time)
                                    break

                            async with lock:
                                # Acquire lock so only 1 coro handling at a time
                                # as we might call the RPC module in the main thread
                                await self._handle_leader_message(data)

                except socket.gaierror:
                    logger.info(f"Socket error - retrying connection in {self.sleep_time}s")
                    await asyncio.sleep(self.sleep_time)
                    continue
                except ConnectionRefusedError:
                    logger.info(f"Connection Refused - retrying connection in {self.sleep_time}s")
                    await asyncio.sleep(self.sleep_time)
                    continue

        except asyncio.CancelledError:
            pass

    async def _handle_leader_message(self, message: Dict[str, Any]):
        type = message.get('data_type')
        data = message.get('data')

        logger.info(f"Received message from Leader: {type} - {data}")

        if type == LeaderMessageType.pairlist:
            # Add the data to the ExternalPairlist
            self.freqtrade.pairlists._pairlist_handlers[0].add_pairlist_data(data)