initial rework separating server and client impl

freqtrade/rpc/api_server/ws/channel.py

@@ -0,0 +1,146 @@
+import logging
+from threading import RLock
+from typing import Type
+
+from freqtrade.rpc.api_server.ws.proxy import WebSocketProxy
+from freqtrade.rpc.api_server.ws.serializer import ORJSONWebSocketSerializer, WebSocketSerializer
+from freqtrade.rpc.api_server.ws.types import WebSocketType
+
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketChannel:
+    """
+    Object to help facilitate managing a websocket connection
+    """
+
+    def __init__(
+        self,
+        websocket: WebSocketType,
+        serializer_cls: Type[WebSocketSerializer] = ORJSONWebSocketSerializer
+    ):
+        # The WebSocket object
+        self._websocket = WebSocketProxy(websocket)
+        # The Serializing class for the WebSocket object
+        self._serializer_cls = serializer_cls
+
+        # Internal event to signify a closed websocket
+        self._closed = False
+
+        # Wrap the WebSocket in the Serializing class
+        self._wrapped_ws = self._serializer_cls(self._websocket)
+
+    async def send(self, data):
+        """
+        Send data on the wrapped websocket
+        """
+        # logger.info(f"Serialized Send - {self._wrapped_ws._serialize(data)}")
+        await self._wrapped_ws.send(data)
+
+    async def recv(self):
+        """
+        Receive data on the wrapped websocket
+        """
+        return await self._wrapped_ws.recv()
+
+    async def ping(self):
+        """
+        Ping the websocket
+        """
+        return await self._websocket.ping()
+
+    async def close(self):
+        """
+        Close the WebSocketChannel
+        """
+
+        self._closed = True
+
+    def is_closed(self):
+        return self._closed
+
+
+class ChannelManager:
+    def __init__(self):
+        self.channels = dict()
+        self._lock = RLock()  # Re-entrant Lock
+
+    async def on_connect(self, websocket: WebSocketType):
+        """
+        Wrap websocket connection into Channel and add to list
+
+        :param websocket: The WebSocket object to attach to the Channel
+        """
+        if hasattr(websocket, "accept"):
+            try:
+                await websocket.accept()
+            except RuntimeError:
+                # The connection was closed before we could accept it
+                return
+
+        ws_channel = WebSocketChannel(websocket)
+
+        with self._lock:
+            self.channels[websocket] = ws_channel
+
+        return ws_channel
+
+    async def on_disconnect(self, websocket: WebSocketType):
+        """
+        Call close on the channel if it's not, and remove from channel list
+
+        :param websocket: The WebSocket objet attached to the Channel
+        """
+        with self._lock:
+            channel = self.channels.get(websocket)
+            if channel:
+                logger.debug(f"Disconnecting channel - {channel}")
+
+                if not channel.is_closed():
+                    await channel.close()
+
+                del self.channels[websocket]
+
+    async def disconnect_all(self):
+        """
+        Disconnect all Channels
+        """
+        with self._lock:
+            for websocket, channel in self.channels.items():
+                if not channel.is_closed():
+                    await channel.close()
+
+            self.channels = dict()
+
+    async def broadcast(self, data):
+        """
+        Broadcast data on all Channels
+
+        :param data: The data to send
+        """
+        with self._lock:
+            logger.debug(f"Broadcasting data: {data}")
+            for websocket, channel in self.channels.items():
+                try:
+                    await channel.send(data)
+                except RuntimeError:
+                    # Handle cannot send after close cases
+                    await self.on_disconnect(websocket)
+
+    async def send_direct(self, channel, data):
+        """
+        Send data directly through direct_channel only
+
+        :param direct_channel: The WebSocketChannel object to send data through
+        :param data: The data to send
+        """
+        # We iterate over the channels to get reference to the websocket object
+        # so we can disconnect incase of failure
+        await channel.send(data)
+
+    def has_channels(self):
+        """
+        Flag for more than 0 channels
+        """
+        return len(self.channels) > 0

freqtrade/rpc/api_server/ws/proxy.py

@@ -0,0 +1,61 @@
+from typing import Union
+
+from fastapi import WebSocket as FastAPIWebSocket
+from websockets import WebSocketClientProtocol as WebSocket
+
+from freqtrade.rpc.api_server.ws.types import WebSocketType
+
+
+class WebSocketProxy:
+    """
+    WebSocketProxy object to bring the FastAPIWebSocket and websockets.WebSocketClientProtocol
+    under the same API
+    """
+
+    def __init__(self, websocket: WebSocketType):
+        self._websocket: Union[FastAPIWebSocket, WebSocket] = websocket
+
+    async def send(self, data):
+        """
+        Send data on the wrapped websocket
+        """
+        if isinstance(data, str):
+            data = data.encode()
+
+        if hasattr(self._websocket, "send_bytes"):
+            await self._websocket.send_bytes(data)
+        else:
+            await self._websocket.send(data)
+
+    async def recv(self):
+        """
+        Receive data on the wrapped websocket
+        """
+        if hasattr(self._websocket, "receive_bytes"):
+            return await self._websocket.receive_bytes()
+        else:
+            return await self._websocket.recv()
+
+    async def ping(self):
+        """
+        Ping the websocket, not supported by FastAPI WebSockets
+        """
+        if hasattr(self._websocket, "ping"):
+            return await self._websocket.ping()
+        return False
+
+    async def close(self, code: int = 1000):
+        """
+        Close the websocket connection, only supported by FastAPI WebSockets
+        """
+        if hasattr(self._websocket, "close"):
+            return await self._websocket.close(code)
+        pass
+
+    async def accept(self):
+        """
+        Accept the WebSocket connection, only support by FastAPI WebSockets
+        """
+        if hasattr(self._websocket, "accept"):
+            return await self._websocket.accept()
+        pass

freqtrade/rpc/api_server/ws/serializer.py

@@ -0,0 +1,65 @@
+import json
+import logging
+from abc import ABC, abstractmethod
+
+import msgpack
+import orjson
+
+from freqtrade.rpc.api_server.ws.proxy import WebSocketProxy
+
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketSerializer(ABC):
+    def __init__(self, websocket: WebSocketProxy):
+        self._websocket: WebSocketProxy = websocket
+
+    @abstractmethod
+    def _serialize(self, data):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _deserialize(self, data):
+        raise NotImplementedError()
+
+    async def send(self, data: bytes):
+        await self._websocket.send(self._serialize(data))
+
+    async def recv(self) -> bytes:
+        data = await self._websocket.recv()
+
+        return self._deserialize(data)
+
+    async def close(self, code: int = 1000):
+        await self._websocket.close(code)
+
+# Going to explore using MsgPack as the serialization,
+# as that might be the best method for sending pandas
+# dataframes over the wire
+
+
+class JSONWebSocketSerializer(WebSocketSerializer):
+    def _serialize(self, data):
+        return json.dumps(data)
+
+    def _deserialize(self, data):
+        return json.loads(data)
+
+
+class ORJSONWebSocketSerializer(WebSocketSerializer):
+    ORJSON_OPTIONS = orjson.OPT_NAIVE_UTC | orjson.OPT_SERIALIZE_NUMPY
+
+    def _serialize(self, data):
+        return orjson.dumps(data, option=self.ORJSON_OPTIONS)
+
+    def _deserialize(self, data):
+        return orjson.loads(data, option=self.ORJSON_OPTIONS)
+
+
+class MsgPackWebSocketSerializer(WebSocketSerializer):
+    def _serialize(self, data):
+        return msgpack.packb(data, use_bin_type=True)
+
+    def _deserialize(self, data):
+        return msgpack.unpackb(data, raw=False)

freqtrade/rpc/api_server/ws/types.py

@@ -0,0 +1,8 @@
+from typing import Any, Dict, TypeVar
+
+from fastapi import WebSocket as FastAPIWebSocket
+from websockets import WebSocketClientProtocol as WebSocket
+
+
+WebSocketType = TypeVar("WebSocketType", FastAPIWebSocket, WebSocket)
+MessageType = Dict[str, Any]

freqtrade/rpc/api_server/ws/utils.py

@@ -0,0 +1,12 @@
+from fastapi import WebSocket
+# fastapi does not make this available through it, so import directly from starlette
+from starlette.websockets import WebSocketState
+
+
+async def is_websocket_alive(ws: WebSocket) -> bool:
+    if (
+        ws.application_state == WebSocketState.CONNECTED and
+        ws.client_state == WebSocketState.CONNECTED
+    ):
+        return True
+    return False