From 516a68cf7f6a44d07fa1f18d3a8bc65619f52b36 Mon Sep 17 00:00:00 2001 From: Danglewood <85772166+deeleeramone@users.noreply.github.com> Date: Wed, 1 Jan 2025 23:47:17 -0800 Subject: [PATCH] playing with some settings and update database test --- .../provider/utils/websockets/client.py | 2 +- .../provider/utils/websockets/database.py | 175 ++++++++++-------- .../provider/utils/websockets/helpers.py | 34 +++- .../utils/websockets/test_database.py | 52 +++++- .../websockets/openbb_websockets/helpers.py | 2 + .../openbb_polygon/utils/websocket_client.py | 62 ++++--- .../openbb_tiingo/utils/websocket_client.py | 84 +++++---- 7 files changed, 269 insertions(+), 142 deletions(-) diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/client.py b/openbb_platform/core/openbb_core/provider/utils/websockets/client.py index 659874fb81d..704e3aa7be8 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/client.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/client.py @@ -852,7 +852,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes def __repr__(self): """Return the WebSocketClient representation.""" return ( - f"WebSocketClient(module={[d for d in self.module if "api_key" not in d]}, symbol={self.symbol}, " + f"WebSocketClient(module={[d for d in self.module if 'api_key' not in d]}, symbol={self.symbol}, " f"is_running={self.is_running}, provider_pid: " f"{self._psutil_process.pid if self._psutil_process else ''}, is_broadcasting={self.is_broadcasting}, " f"broadcast_address={self.broadcast_address}, " diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py index 294a02b6879..c307fd94694 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py @@ -8,8 +8,12 @@ from contextlib import asynccontextmanager from datetime import timedelta from typing import TYPE_CHECKING, Any, Iterable, Optional, Union +import uvloop from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.provider.utils.helpers import run_async +from openbb_core.provider.utils.websockets.helpers import kill_thread + +asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) if TYPE_CHECKING: import logging @@ -65,30 +69,6 @@ CHECK_FOR = ( ) -def kill_thread(thread: threading.Thread) -> None: - """Kill thread by setting a stop flag.""" - # pylint: disable=import-outside-toplevel - import ctypes - - if hasattr(thread, "loop") and thread.loop: - for task in asyncio.all_tasks(thread.loop): - task.cancel() - - if not thread.is_alive(): - return - - thread_id = thread.ident - if thread_id is None: - return - - res = ctypes.pythonapi.PyThreadState_SetAsyncExc( - ctypes.c_long(thread_id), ctypes.py_object(SystemExit) - ) - if res > 1: - ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(thread_id), None) - raise SystemError("PyThreadState_SetAsyncExc failed") - - class Database: """ Class to read from, and write to, the SQL file using aiosqlite. @@ -160,7 +140,7 @@ class Database: data_model: Optional["BaseModel"] = None, limit: Optional[int] = None, logger: Optional["logging.Logger"] = None, - loop: Optional["asyncio.AbstractEventLoop"] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, **kwargs, ): """Initialize the ResultsDB class.""" @@ -212,7 +192,7 @@ class Database: try: if self.results_file is not None and os.path.exists(self.results_file): # type: ignore - async with self.get_connection("read") as conn: + async with self.get_connection("write") as conn: try: cursor = await conn.execute( "SELECT name FROM sqlite_master WHERE type='table';" @@ -238,7 +218,8 @@ class Database: ) pragmas = [ "PRAGMA journal_mode=WAL", - "PRAGMA synchronous=OFF", + "PRAGMA synchronous=off", + "PRAGMA temp_store=memory", ] for pragma in pragmas: @@ -288,6 +269,15 @@ class Database: if name not in self._connections: conn = await aiosqlite.connect(results_file, **conn_kwargs) + pragmas = [ + "PRAGMA journal_mode=WAL", + "PRAGMA synchronous=off", + "PRAGMA temp_store=memory", + ] + for pragma in pragmas: + await conn.execute(pragma) + + await conn.commit() self._connections[name] = conn yield self._connections[name] @@ -373,30 +363,39 @@ class Database: params = () async with conn.execute(query, params) as cursor: async for row in cursor: - rows.append(await self._deserialize_row(row)) + rows.append(await self._deserialize_row(row, cursor)) return rows except Exception as e: # pylint: disable=broad-except raise OpenBBError(e) from e - async def _deserialize_row(self, row) -> dict: - """Deserialize a row from the SQLite database.""" + async def _deserialize_row(self, row, cursor) -> dict: + """Deserialize a row from the SQLite database. + Handles both single message column and multiple extracted fields.""" # pylint: disable=import-outside-toplevel import json try: - return ( - json.loads(row[0]) - if ( - ( + if len(row) == 1: + # Single column case (full message) + return ( + json.loads(row[0]) + if ( isinstance(row[0], str) and (row[0].startswith("{") or row[0].startswith("[")) ) or isinstance(row[0], bytes) + else row[0] ) - else row[0] - ) + else: + # Multiple column case (extracted fields) + return {cursor.description[i][0]: row[i] for i in range(len(row))} + + except (json.JSONDecodeError, AttributeError) as e: + self.logger.error(f"Failed to deserialize row: {e}") + return row[0] if len(row) == 1 else dict(enumerate(row)) + except Exception as e: # pylint: disable=broad-except msg = ( "Unexpected error while deserializing row -> " @@ -453,14 +452,12 @@ class Database: async def _query_db(self, sql, parameters: Optional[Iterable[Any]] = None) -> list: """Query the SQLite database.""" - # pylint: disable=import-outside-toplevel - import json if not sql or sql in ("", "''"): raise OpenBBError("Empty query not allowed.") query = ( sql - if sql.startswith("SELECT") + if sql.strip().startswith("SELECT") else f"SELECT message FROM {self.table_name} WHERE {sql}" # noqa ) if not query.endswith(";"): @@ -479,17 +476,7 @@ class Database: query, parameters ) as cursor: async for row in cursor: - rows.append( - json.loads(row[0]) - if ( - ( - isinstance(row[0], str) - and (row[0].startswith("{") or row[0].startswith("[")) - ) - or isinstance(row[0], bytes) - ) - else row[0] - ) + rows.append(await self._deserialize_row(row, cursor)) return rows except Exception as e: # pylint: disable=broad-except raise OpenBBError(e) from e @@ -563,7 +550,7 @@ class Database: def create_writer( self, queue: Optional["MessageQueue"] = None, - batch_size: int = 25000, + batch_size: int = 200, prune_interval: Optional[int] = None, export_directory: Optional[Union[str, "Path"]] = None, export_interval: Optional[int] = None, @@ -624,7 +611,7 @@ class DatabaseWriter: self, database: Database, queue: Optional["MessageQueue"] = None, - batch_size: int = 25000, + batch_size: int = 200, prune_interval: Optional[int] = None, export_directory: Optional[Union[str, "Path"]] = None, export_interval: Optional[int] = None, @@ -687,7 +674,8 @@ class DatabaseWriter: async def start_writer(self): """Start writing tasks.""" if not self.writer_running: - asyncio.get_event_loop().run_in_executor(None, self.batch_processor.start()) + if not self.batch_processor.is_alive(): + self.batch_processor.start() await self._create_connection() for _ in range(self.num_workers): task = asyncio.create_task(self._process_queue()) @@ -700,7 +688,7 @@ class DatabaseWriter: await asyncio.gather(*self.write_tasks, return_exceptions=True) if self._conn: await self._conn.close() - self.batch_processor.running = False + self.batch_processor.stop() kill_thread(self.batch_processor) async def _process_queue(self): @@ -718,7 +706,7 @@ class DatabaseWriter: except asyncio.TimeoutError: break if batch: - self._write_batch(batch) + await self._write_batch(batch) batch = [] else: await asyncio.sleep(0.001) @@ -734,9 +722,9 @@ class DatabaseWriter: while not self.queue.queue.empty() and len(batch) < self.batch_size: batch.append(await self.queue.dequeue()) - self._write_batch(batch) + await self._write_batch(batch) - def _write_batch(self, batch): + async def _write_batch(self, batch): """Write the batch of messages to the database.""" if not batch: return @@ -757,6 +745,7 @@ class DatabaseWriter: chunk_size = 20000 minutes = self.export_interval or 5 latest_date = None + earliest_date = None if not self._export_running or not self.export_thread: return try: @@ -767,6 +756,12 @@ class DatabaseWriter: ORDER BY json_extract(message, '$.date') DESC LIMIT 1 """ # noqa + earliest_query = f""" + SELECT json_extract(message, '$.date') + FROM {self.database.table_name} + ORDER BY json_extract(message, '$.date') ASC LIMIT 1 + """ # noqa + async with self.database.get_connection("read") as conn: try: async with conn.execute(latest_query) as cursor: @@ -774,19 +769,29 @@ class DatabaseWriter: if not latest_date: return + async with conn.execute(earliest_query) as cursor: + earliest_date = await cursor.fetchone() + if not latest_date: + return + except asyncio.InvalidStateError as e: self.database.logger.error(f"Database connection state error: {e}") self._export_running = False return latest_timestamp = to_datetime(latest_date[0]) - + earliest_timestamp = to_datetime(earliest_date[0]) # Round down to nearest interval cutoff_time = latest_timestamp - timedelta( minutes=latest_timestamp.minute % minutes, seconds=latest_timestamp.second, microseconds=latest_timestamp.microsecond, ) + earliest_time = earliest_timestamp - timedelta( + minutes=latest_timestamp.minute % minutes, + seconds=latest_timestamp.second, + microseconds=latest_timestamp.microsecond, + ) # If we have processed data before, use that as reference if self._last_processed_timestamp: @@ -794,8 +799,9 @@ class DatabaseWriter: else: start_time = cutoff_time - timedelta(minutes=minutes) - cutoff_time = self._last_processed_timestamp + timedelta(minutes=minutes) - start_time = self._last_processed_timestamp + if start_time < earliest_timestamp: + start_time = earliest_time + results_file = ( self.export_directory + "/" @@ -1042,9 +1048,7 @@ class DatabaseWriter: count = await self.database._query_db(query, (cutoff_timestamp,)) if count[0] > 0: - # task = asyncio.create_task(self._export_database()) - task = await asyncio.to_thread(self._export_database) - await asyncio.gather(task) + await self._export_database() else: await asyncio.sleep(minutes * 60) # Stagger slightly so things don't happen exactly on the minute. @@ -1058,13 +1062,11 @@ class DatabaseWriter: class BatchProcessor(threading.Thread): """ - Batch processor for writing messages to the SQLite database in a new thread. - This class is a thread intended for use as a subprocess and is called by `DatabaseWriter.start_writer()`. """ def __init__( - self, database_writer: DatabaseWriter, num_workers=120, collection_time=0.35 + self, database_writer: DatabaseWriter, num_workers=120, collection_time=0.25 ): """Initialize the BatchProcessor class.""" # pylint: disable=import-outside-toplevel @@ -1078,20 +1080,42 @@ class BatchProcessor(threading.Thread): self.num_workers = num_workers self.workers: list = [] self.collection_time = collection_time + self._shutdown = threading.Event() def run(self): """Run the batch processor as tasks.""" - - self.loop = asyncio.new_event_loop() - asyncio.set_event_loop(self.loop) try: + self.loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.loop) # Create worker tasks - for _ in range(self.num_workers): - worker = self.loop.create_task(self._worker()) - self.workers.append(worker) - # Run workers - self.loop.run_until_complete(asyncio.gather(*self.workers)) + while self.running and not self._shutdown.is_set(): + try: + self.loop.run_until_complete(self._worker()) + except (SystemExit, KeyboardInterrupt): + self.running = False + break + except Exception as e: + self.writer.database.logger.error(f"Batch processing error: {e}") + break finally: + self._cleanup() + + def stop(self): + """Signal thread to stop gracefully.""" + self.running = False + self._shutdown.set() + if self.loop and self.loop.is_running(): + self.loop.call_soon_threadsafe(self.loop.stop) + + def _cleanup(self): + """Clean up resources on shutdown""" + if self.loop: + pending = asyncio.all_tasks(self.loop) + for task in pending: + task.cancel() + self.loop.run_until_complete( + asyncio.gather(*pending, return_exceptions=True) + ) self.loop.close() async def _worker(self): @@ -1101,7 +1125,7 @@ class BatchProcessor(threading.Thread): batch_size = 200 while self.running: try: - batch = [] + batch: list = [] total = 0 collection_start = time.time() @@ -1136,6 +1160,7 @@ class BatchProcessor(threading.Thread): for b in batch: values.extend([(msg,) for msg in b]) async with self.writer.database.get_connection("write") as conn: + await conn.execute("PRAGMA temp_store=memory") async with conn.cursor() as cursor: await cursor.executemany(query, values) await conn.commit() diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py b/openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py index 23ea8f2f060..df856d64cfe 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py @@ -4,9 +4,13 @@ import logging import re +from typing import TYPE_CHECKING from pydantic import ValidationError +if TYPE_CHECKING: + import threading + AUTH_TOKEN_FILTER = re.compile( r"(auth_token=)([^&]*)", re.IGNORECASE | re.MULTILINE, @@ -73,13 +77,10 @@ def decrypt_value(key, iv, encrypted_value): def handle_termination_signal(logger): """Handle termination signals to ensure graceful shutdown.""" - # pylint: disable=import-outside-toplevel - import sys - logger.info( "PROVIDER INFO: Termination signal received. WebSocket connection closed." ) - sys.exit(0) + raise SystemExit("Termination signal received.") def parse_kwargs() -> dict: @@ -127,3 +128,28 @@ def parse_kwargs() -> dict: _kwargs[key] = True return _kwargs + + +def kill_thread(thread: "threading.Thread") -> None: + """Kill thread by setting a stop flag.""" + # pylint: disable=import-outside-toplevel + import asyncio + import ctypes + + if hasattr(thread, "loop") and thread.loop: + for task in asyncio.all_tasks(thread.loop): + task.cancel() + + if not thread.is_alive(): + return + + thread_id = thread.ident + if thread_id is None: + return + + res = ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_long(thread_id), ctypes.py_object(SystemExit) + ) + if res > 1: + ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(thread_id), None) + raise SystemError("PyThreadState_SetAsyncExc failed") diff --git a/openbb_platform/core/tests/provider/utils/websockets/test_database.py b/openbb_platform/core/tests/provider/utils/websockets/test_database.py index ae4bef7c054..5e2532df7e7 100644 --- a/openbb_platform/core/tests/provider/utils/websockets/test_database.py +++ b/openbb_platform/core/tests/provider/utils/websockets/test_database.py @@ -1,7 +1,11 @@ """Unit Tests For Database Operations.""" +import asyncio +import json + import pytest -from openbb_core.provider.utils.websockets.database import Database +from openbb_core.provider.utils.websockets.database import Database, DatabaseWriter +from openbb_core.provider.utils.websockets.message_queue import MessageQueue MOCK_MESSAGES = [ {"type": "trade", "symbol": "test1", "price": 100}, @@ -16,6 +20,19 @@ def database(): return Database(table_name="test") +@pytest.fixture +def database_writer(): + """Return a MessageQueue instance.""" + writer = DatabaseWriter(database=Database(table_name="test"), queue=MessageQueue()) + yield writer + + +@pytest.fixture +def message_queue(): + """Return a MessageQueue instance.""" + return MessageQueue() + + def test_setup_database(database): """Test if the database was setup.""" assert database @@ -114,3 +131,36 @@ def test_limit(): database.write_to_db(MOCK_MESSAGES[0]) query = "SELECT id FROM test_limit" assert database.query(query)[0] > 3 + + +def test_batch_process_thread(database_writer): + """Test if the batch process thread starts and stops.""" + database_writer.batch_processor.start() + assert database_writer.batch_processor.is_alive() + database_writer.batch_processor.stop() + database_writer.batch_processor.join(timeout=1) + assert not database_writer.batch_processor.is_alive() + + +@pytest.mark.asyncio +async def test_database_writer(database_writer): + """Test if the database writer is working.""" + assert database_writer + writer = database_writer + assert isinstance(writer, DatabaseWriter) + assert len(writer.database.fetch_all()) == 0 + await writer.start_writer() + assert not writer._export_running + assert not writer._prune_running + + for message in MOCK_MESSAGES: + await writer.queue.enqueue(json.dumps(message)) + + await asyncio.sleep(0.5) + + await writer.stop_writer() + + messages = writer.database.fetch_all() + assert len(messages) == 3 + + assert not writer.batch_processor.is_alive() diff --git a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py index 0d9a46cac41..5a1ecd3f21b 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py @@ -66,8 +66,10 @@ class StdOutSink: def write(self, message): """Write to stdout.""" # pylint: disable=import-outside-toplevel + import os import sys + os.set_blocking(sys.__stdout__.fileno(), False) # type: ignore cleaned_message = AUTH_TOKEN_FILTER.sub(r"\1********", message) if cleaned_message != message: cleaned_message = f"{cleaned_message}\n" diff --git a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py index 39a23bf165d..d12153f3eea 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -36,9 +36,9 @@ import os import signal import sys import time -from concurrent.futures import ThreadPoolExecutor import orjson +import uvloop import websockets from openbb_core.provider.utils.websockets.database import Database, DatabaseWriter from openbb_core.provider.utils.websockets.helpers import ( @@ -54,6 +54,9 @@ from openbb_polygon.models.websocket_connection import ( ) from pydantic import ValidationError from websockets.asyncio.client import connect +from websockets.extensions import permessage_deflate + +asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) URL_MAP = { "stock": "wss://socket.polygon.io/stocks", @@ -78,7 +81,6 @@ ASSET_TYPE = kwargs.pop("asset_type", "crypto") kwargs["results_file"] = os.path.abspath(kwargs.get("results_file")) URL = URL_MAP.get(ASSET_TYPE) SUBSCRIBED_SYMBOLS: set = set() -message_thread_pool = ThreadPoolExecutor(max_workers=16) if not kwargs.get("api_key"): raise ValueError("No API key provided.") @@ -86,6 +88,7 @@ if not kwargs.get("api_key"): if not URL: raise ValueError("Invalid asset type provided.") + DATABASE = DatabaseWriter( database=Database( results_file=kwargs["results_file"], @@ -158,8 +161,8 @@ async def login(websocket): login_event = f'{{"action":"auth","params":"{kwargs["api_key"]}"}}' try: await websocket.send(login_event) - res = await websocket.recv() - response = json.loads(res) + res = await websocket.recv(decode=False) + response = orjson.loads(res) messages = response if isinstance(response, list) else [response] for msg in messages: if msg.get("status") == "connected": @@ -213,7 +216,7 @@ async def subscribe(websocket, symbol, event): async def read_stdin(): """Read from stdin and queue commands.""" while True: - line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) + line = await asyncio.to_thread(sys.stdin.readline) sys.stdin.flush() if line: try: @@ -280,22 +283,21 @@ async def connect_and_stream(): conn_kwargs.update( { - "ping_interval": None, "ping_timeout": None, "close_timeout": 1, - "max_size": 2**63, + "max_size": 32768, "max_queue": None, } ) async def message_receiver(websocket): """Message receiver.""" - try: - async for message in websocket: + while True: + try: + message = await websocket.recv(decode=False) input_queue.queue.put_nowait(message) - - except Exception as e: - raise e from e + except Exception as e: + raise e from e async def process_input_messages(message): """Process the messages offloaded from the websocket.""" @@ -325,7 +327,8 @@ async def connect_and_stream(): # Run processing in thread process = asyncio.to_thread(_process_in_thread) - asyncio.create_task(process) + process_task = asyncio.create_task(process) + tasks.add(process_task) try: handler_task = asyncio.create_task( @@ -337,15 +340,25 @@ async def connect_and_stream(): await DATABASE.start_writer() - for i in range(16): - processor_task = asyncio.create_task( - input_queue.process_queue( - lambda message: process_input_messages(message) - ) - ) - tasks.add(processor_task) + processor_task = asyncio.create_task( + input_queue.process_queue(lambda message: process_input_messages(message)) + ) + tasks.add(processor_task) - async for websocket in connect(URL, **conn_kwargs): + conn_kwargs["extensions"] = [ + permessage_deflate.ClientPerMessageDeflateFactory( + server_max_window_bits=11, + client_max_window_bits=11, + compress_settings={ + "memLevel": 1, + "level": 1, + }, + ), + ] + async for websocket in connect( + URL, + **conn_kwargs, + ): try: if not any( task.name == "cmd_task" for task in tasks if hasattr(task, "name") @@ -357,10 +370,9 @@ async def connect_and_stream(): await login(websocket) - response = await websocket.recv() - messages = orjson.loads(response) + response = await websocket.recv(decode=False) - await process_message(messages) + await process_message(orjson.loads(response)) await subscribe(websocket, kwargs["symbol"], "subscribe") @@ -406,7 +418,7 @@ async def connect_and_stream(): if __name__ == "__main__": try: - loop = asyncio.new_event_loop() + loop = uvloop.new_event_loop() asyncio.set_event_loop(loop) loop.set_exception_handler(lambda loop, context: None) diff --git a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py index 9e04862a93d..9f88a667934 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -1,12 +1,13 @@ """FMP WebSocket server.""" import asyncio -import json import os import signal import sys -from datetime import datetime, UTC +from datetime import UTC, datetime +import orjson +import uvloop import websockets from openbb_core.provider.utils.errors import UnauthorizedError from openbb_core.provider.utils.websockets.database import Database, DatabaseWriter @@ -21,6 +22,8 @@ from openbb_tiingo.models.websocket_connection import TiingoWebSocketData from pandas import to_datetime from pydantic import ValidationError +asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + URL_MAP = { "stock": "wss://api.tiingo.com/iex", "fx": "wss://api.tiingo.com/fx", @@ -97,7 +100,6 @@ DATABASE = DatabaseWriter( logger=logger, ), queue=db_queue, - batch_size=100, ) @@ -120,9 +122,9 @@ async def update_symbols(symbol, event): } async with websockets.connect(URL) as websocket: - await websocket.send(json.dumps(update_event)) - response = await websocket.recv() - message = json.loads(response) + await websocket.send(orjson.dumps(update_event)) + response = await websocket.recv(decode=False) + message = orjson.loads(response) if "tickers" in message.get("data", {}): tickers = message["data"]["tickers"] threshold_level = message["data"].get("thresholdLevel") @@ -155,7 +157,7 @@ async def read_stdin_and_update_symbols(): f" Database Queue : {db_queue.queue.qsize()}" ) else: - line = json.loads(line.strip()) + line = orjson.loads(line.strip()) if line: symbol = line.get("symbol") @@ -167,7 +169,7 @@ async def process_message(message): # pylint: disable=too-many-branches """Process the message and write to the database.""" result: dict = {} data_message: dict = {} - message = message if isinstance(message, (dict, list)) else json.loads(message) + message = message if isinstance(message, (dict, list)) else orjson.loads(message) msg: str = "" if message.get("messageType") == "E": response = message.get("response", {}) @@ -232,7 +234,7 @@ async def process_message(message): # pylint: disable=too-many-branches raise e from e if result: - await db_queue.enqueue(result) + db_queue.queue.put_nowait(result) async def connect_and_stream(): @@ -247,9 +249,7 @@ async def connect_and_stream(): { "ping_interval": 8, "ping_timeout": 8, - "read_limit": 2**256, "close_timeout": 1, - "max_queue": 10000, } ) @@ -268,38 +268,50 @@ async def connect_and_stream(): async def message_receiver(websocket): """Receive messages from the WebSocket.""" while True: - message = await websocket.recv() - await input_queue.enqueue(message) + message = await websocket.recv(decode=False) + input_queue.queue.put_nowait(orjson.loads(message)) stdin_task = asyncio.create_task(read_stdin_and_update_symbols()) tasks.add(stdin_task) try: await DATABASE.start_writer() - websocket = await websockets.connect(URL, **conn_kwargs) - receiver_task = asyncio.create_task(message_receiver(websocket)) - tasks.add(receiver_task) - await websocket.send(json.dumps(subscribe_event)) - logger.info("PROVIDER INFO: WebSocket connection established.") - for _ in range(9): - process_task = asyncio.create_task( - input_queue.process_queue(lambda message: process_message(message)) - ) - tasks.add(process_task) - await asyncio.gather(*tasks, return_exceptions=True) + async for websocket in websockets.connect(URL, **conn_kwargs): + try: + if not any( + task.name == "receiver_task" + for task in tasks + if hasattr(task, "name") + ): + receiver_task = asyncio.create_task( + message_receiver(websocket), name="receiver_task" + ) + tasks.add(receiver_task) - except UnauthorizedError as e: - logger.error(str(e)) - sys.exit(1) + await websocket.send(orjson.dumps(subscribe_event)) + logger.info("PROVIDER INFO: WebSocket connection established.") + for _ in range(9): + process_task = asyncio.create_task( + input_queue.process_queue( + lambda message: process_message(message) + ) + ) + tasks.add(process_task) - except websockets.ConnectionClosed as e: - msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}" - logger.info(msg) - # Attempt to reopen the connection - logger.info("PROVIDER INFO: Attempting to reconnect...") - await asyncio.sleep(1) - await connect_and_stream() + await asyncio.gather(*tasks, return_exceptions=True) + + except UnauthorizedError as e: + logger.error(str(e)) + sys.exit(1) + + except websockets.ConnectionClosed as e: + msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}" + logger.info(msg) + # Attempt to reopen the connection + logger.info("PROVIDER INFO: Attempting to reconnect...") + await asyncio.sleep(1) + continue except websockets.WebSocketException as e: logger.info(str(e)) @@ -311,7 +323,6 @@ async def connect_and_stream(): sys.exit(1) finally: - await websocket.close() for task in tasks: task.cancel() await task @@ -322,9 +333,10 @@ async def connect_and_stream(): if __name__ == "__main__": try: - loop = asyncio.new_event_loop() + loop = uvloop.new_event_loop() asyncio.set_event_loop(loop) loop.set_exception_handler(lambda loop, context: None) + for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(signal.SIGTERM, handle_termination_signal, logger)