stable/freqtrade/rpc/api_server/ws/channel.py

349 lines
10 KiB
Python
Raw Normal View History

import asyncio
import logging
2022-11-15 03:27:45 +00:00
from contextlib import asynccontextmanager
from typing import Any, Dict, List, Optional, Type, Union
from uuid import uuid4
from freqtrade.rpc.api_server.ws.proxy import WebSocketProxy
2022-09-02 06:05:36 +00:00
from freqtrade.rpc.api_server.ws.serializer import (HybridJSONWebSocketSerializer,
WebSocketSerializer)
from freqtrade.rpc.api_server.ws.types import WebSocketType
from freqtrade.rpc.api_server.ws_schemas import WSMessageSchemaType
logger = logging.getLogger(__name__)
class WebSocketChannel:
"""
Object to help facilitate managing a websocket connection
"""
def __init__(
self,
websocket: WebSocketType,
channel_id: Optional[str] = None,
2022-09-02 06:05:36 +00:00
serializer_cls: Type[WebSocketSerializer] = HybridJSONWebSocketSerializer
):
self.channel_id = channel_id if channel_id else uuid4().hex[:8]
self._websocket = WebSocketProxy(websocket)
# Internal event to signify a closed websocket
2022-11-02 19:26:27 +00:00
self._closed = asyncio.Event()
2022-11-15 03:27:45 +00:00
# Throttle how fast we send messages
self._throttle = 0.01
# Wrap the WebSocket in the Serializing class
2022-11-13 09:26:27 +00:00
self._wrapped_ws = serializer_cls(self._websocket)
def __repr__(self):
return f"WebSocketChannel({self.channel_id}, {self.remote_addr})"
@property
def raw_websocket(self):
return self._websocket.raw_websocket
@property
def remote_addr(self):
return self._websocket.remote_addr
2022-11-15 03:27:45 +00:00
async def send(self, message: Union[WSMessageSchemaType, Dict[str, Any]]):
"""
2022-11-15 03:27:45 +00:00
Send a message on the wrapped websocket
"""
2022-11-15 03:27:45 +00:00
await asyncio.sleep(self._throttle)
await self._wrapped_ws.send(message)
async def recv(self):
"""
2022-11-15 03:27:45 +00:00
Receive a message on the wrapped websocket
"""
return await self._wrapped_ws.recv()
async def ping(self):
"""
Ping the websocket
"""
return await self._websocket.ping()
2022-11-15 03:27:45 +00:00
async def accept(self):
"""
Accept the underlying websocket connection
"""
return await self._websocket.accept()
async def close(self):
"""
Close the WebSocketChannel
"""
2022-11-02 19:26:27 +00:00
try:
2022-11-15 03:27:45 +00:00
await self._websocket.close()
2022-11-02 19:26:27 +00:00
except Exception:
pass
self._closed.set()
def is_closed(self) -> bool:
"""
Closed flag
"""
2022-11-02 19:26:27 +00:00
return self._closed.is_set()
def set_subscriptions(self, subscriptions: List[str] = []) -> None:
"""
Set which subscriptions this channel is subscribed to
:param subscriptions: List of subscriptions, List[str]
"""
self._subscriptions = subscriptions
def subscribed_to(self, message_type: str) -> bool:
"""
Check if this channel is subscribed to the message_type
:param message_type: The message type to check
"""
return message_type in self._subscriptions
2022-11-15 03:27:45 +00:00
async def __aiter__(self):
"""
2022-11-15 03:27:45 +00:00
Generator for received messages
"""
2022-11-15 03:27:45 +00:00
while True:
try:
2022-11-15 03:27:45 +00:00
yield await self.recv()
except Exception:
break
2022-11-15 03:27:45 +00:00
@asynccontextmanager
async def connect(self):
"""
2022-11-15 03:27:45 +00:00
Context manager for safely opening and closing the websocket connection
"""
2022-11-15 03:27:45 +00:00
try:
await self.accept()
yield self
finally:
await self.close()
# class WebSocketChannel:
# """
# Object to help facilitate managing a websocket connection
# """
# def __init__(
# self,
# websocket: WebSocketType,
# channel_id: Optional[str] = None,
# drain_timeout: int = 3,
# throttle: float = 0.01,
# serializer_cls: Type[WebSocketSerializer] = HybridJSONWebSocketSerializer
# ):
# self.channel_id = channel_id if channel_id else uuid4().hex[:8]
# # The WebSocket object
# self._websocket = WebSocketProxy(websocket)
# self.drain_timeout = drain_timeout
# self.throttle = throttle
# self._subscriptions: List[str] = []
# # 32 is the size of the receiving queue in websockets package
# self.queue: asyncio.Queue[Dict[str, Any]] = asyncio.Queue(maxsize=32)
# self._relay_task = asyncio.create_task(self.relay())
# # Internal event to signify a closed websocket
# self._closed = asyncio.Event()
# # Wrap the WebSocket in the Serializing class
# self._wrapped_ws = serializer_cls(self._websocket)
# def __repr__(self):
# return f"WebSocketChannel({self.channel_id}, {self.remote_addr})"
# @property
# def raw_websocket(self):
# return self._websocket.raw_websocket
# @property
# def remote_addr(self):
# return self._websocket.remote_addr
# async def _send(self, data):
# """
# Send data on the wrapped websocket
# """
# await self._wrapped_ws.send(data)
# async def send(self, data) -> bool:
# """
# Add the data to the queue to be sent.
# :returns: True if data added to queue, False otherwise
# """
# # This block only runs if the queue is full, it will wait
# # until self.drain_timeout for the relay to drain the outgoing queue
# # We can't use asyncio.wait_for here because the queue may have been created with a
# # different eventloop
# start = time.time()
# while self.queue.full():
# await asyncio.sleep(1)
# if (time.time() - start) > self.drain_timeout:
# return False
# # If for some reason the queue is still full, just return False
# try:
# self.queue.put_nowait(data)
# except asyncio.QueueFull:
# return False
# # If we got here everything is ok
# return True
# 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
# """
# try:
# await self.raw_websocket.close()
# except Exception:
# pass
# self._closed.set()
# self._relay_task.cancel()
# def is_closed(self) -> bool:
# """
# Closed flag
# """
# return self._closed.is_set()
# def set_subscriptions(self, subscriptions: List[str] = []) -> None:
# """
# Set which subscriptions this channel is subscribed to
# :param subscriptions: List of subscriptions, List[str]
# """
# self._subscriptions = subscriptions
# def subscribed_to(self, message_type: str) -> bool:
# """
# Check if this channel is subscribed to the message_type
# :param message_type: The message type to check
# """
# return message_type in self._subscriptions
# async def relay(self):
# """
# Relay messages from the channel's queue and send them out. This is started
# as a task.
# """
# while not self._closed.is_set():
# message = await self.queue.get()
# try:
# await self._send(message)
# self.queue.task_done()
# # Limit messages per sec.
# # Could cause problems with queue size if too low, and
# # problems with network traffik if too high.
# # 0.01 = 100/s
# await asyncio.sleep(self.throttle)
# except RuntimeError:
# # The connection was closed, just exit the task
# return
# 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 isinstance(websocket, FastAPIWebSocket):
# 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.info(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 in self.channels.copy().keys():
# await self.on_disconnect(websocket)
# async def broadcast(self, message: WSMessageSchemaType):
# """
# Broadcast a message on all Channels
# :param message: The message to send
# """
# with self._lock:
# for channel in self.channels.copy().values():
# if channel.subscribed_to(message.get('type')):
# await self.send_direct(channel, message)
# async def send_direct(
# self, channel: WebSocketChannel, message: Union[WSMessageSchemaType, Dict[str, Any]]):
# """
# Send a message directly through direct_channel only
# :param direct_channel: The WebSocketChannel object to send the message through
# :param message: The message to send
# """
# if not await channel.send(message):
# await self.on_disconnect(channel.raw_websocket)
# def has_channels(self):
# """
# Flag for more than 0 channels
# """
# return len(self.channels) > 0