add dynamic send timeout

This commit is contained in:
Timothy Pogue 2022-11-20 14:09:45 -07:00
parent dc79284c54
commit 60a167bdef
2 changed files with 50 additions and 17 deletions

View File

@ -34,7 +34,7 @@ async def channel_broadcaster(channel: WebSocketChannel, message_stream: Message
Iterate over messages in the message stream and send them Iterate over messages in the message stream and send them
""" """
async for message in message_stream: async for message in message_stream:
await channel.send(message) await channel.send(message, timeout=True)
async def _process_consumer_request( async def _process_consumer_request(

View File

@ -1,7 +1,9 @@
import asyncio import asyncio
import logging import logging
import time
from collections import deque
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from typing import Any, Dict, List, Optional, Type, Union from typing import Any, Deque, Dict, List, Optional, Type, Union
from uuid import uuid4 from uuid import uuid4
from freqtrade.rpc.api_server.ws.proxy import WebSocketProxy from freqtrade.rpc.api_server.ws.proxy import WebSocketProxy
@ -29,7 +31,13 @@ class WebSocketChannel:
# Internal event to signify a closed websocket # Internal event to signify a closed websocket
self._closed = asyncio.Event() self._closed = asyncio.Event()
self._send_timeout_high_limit = 2 # The async tasks created for the channel
self._channel_tasks: List[asyncio.Task] = []
# Deque for average send times
self._send_times: Deque[float] = deque([], maxlen=10)
# High limit defaults to 3 to start
self._send_high_limit = 3
# The subscribed message types # The subscribed message types
self._subscriptions: List[str] = [] self._subscriptions: List[str] = []
@ -37,9 +45,6 @@ class WebSocketChannel:
# Wrap the WebSocket in the Serializing class # Wrap the WebSocket in the Serializing class
self._wrapped_ws = serializer_cls(self._websocket) self._wrapped_ws = serializer_cls(self._websocket)
# The async tasks created for the channel
self._channel_tasks: List[asyncio.Task] = []
def __repr__(self): def __repr__(self):
return f"WebSocketChannel({self.channel_id}, {self.remote_addr})" return f"WebSocketChannel({self.channel_id}, {self.remote_addr})"
@ -51,20 +56,48 @@ class WebSocketChannel:
def remote_addr(self): def remote_addr(self):
return self._websocket.remote_addr return self._websocket.remote_addr
async def send(self, message: Union[WSMessageSchemaType, Dict[str, Any]]): def _calc_send_limit(self):
""" """
Send a message on the wrapped websocket Calculate the send high limit for this channel
""" """
await self._wrapped_ws.send(message)
# Without this sleep, messages would send to one channel # Only update if we have enough data
# first then another after the first one finished and prevent if len(self._send_times) == self._send_times.maxlen:
# any normal Rest API calls from processing at the same time. # At least 1s or twice the average of send times
# With the sleep call, it gives control to the event self._send_high_limit = max(
# loop to schedule other channel send methods, and helps (sum(self._send_times) / len(self._send_times)) * 2,
# throttle how fast we send. 1
# 0.005 = 200 messages/second max throughput )
await asyncio.sleep(0.005)
async def send(
self,
message: Union[WSMessageSchemaType, Dict[str, Any]],
timeout: bool = False
):
"""
Send a message on the wrapped websocket. If the sending
takes too long, it will raise a TimeoutError and
disconnect the connection.
:param message: The message to send
:param timeout: Enforce send high limit, defaults to False
"""
try:
_ = time.time()
# If the send times out, it will raise
# a TimeoutError and bubble up to the
# message_endpoint to close the connection
await asyncio.wait_for(
self._wrapped_ws.send(message),
timeout=self._send_high_limit if timeout else None
)
total_time = time.time() - _
self._send_times.append(total_time)
self._calc_send_limit()
except asyncio.TimeoutError:
logger.info(f"Connection for {self} is too far behind, disconnecting")
raise
async def recv(self): async def recv(self):
""" """