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 ef58e03f56a..6df601c385b 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py @@ -143,11 +143,11 @@ class Database: temp_file_path = temp_file.name self.results_path = Path(temp_file_path).absolute() self.results_file = temp_file_path - if ":" in results_file: - self.results_file = results_file - self.results_path = results_file # type: ignore - kwargs["uri"] = True else: + if ":" in results_file: + self.results_file = results_file + self.results_path = results_file # type: ignore + kwargs["uri"] = True self.results_path = Path(results_file).absolute() self.results_file = results_file diff --git a/openbb_platform/extensions/websockets/openbb_websockets/listen.py b/openbb_platform/core/openbb_core/provider/utils/websockets/listen.py similarity index 100% rename from openbb_platform/extensions/websockets/openbb_websockets/listen.py rename to openbb_platform/core/openbb_core/provider/utils/websockets/listen.py diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py b/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py index 4ecddf97963..3eaa396585b 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py @@ -39,6 +39,22 @@ class MessageQueue: retries = 0 while retries < self.max_retries: + + if self.queue.qsize() / self.queue.maxsize > 0.3: + await sleep(0.000005) + if self.queue.qsize() / self.queue.maxsize > 0.5: + await sleep(0.00005) + if self.queue.qsize() / self.queue.maxsize > 0.55: + await sleep(0.00005) + if self.queue.qsize() / self.queue.maxsize > 0.6: + await sleep(0.00005) + if self.queue.qsize() / self.queue.maxsize > 0.65: + await sleep(0.00005) + if self.queue.qsize() / self.queue.maxsize > 0.7: + await sleep(0.00005) + if self.queue.qsize() / self.queue.maxsize > 0.99: + await sleep(0.00005) + if self.queue.full(): retries += 1 msg = f"Queue is full. Retrying {retries}/{self.max_retries}..." 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 ff50d674cc7..ae4bef7c054 100644 --- a/openbb_platform/core/tests/provider/utils/websockets/test_database.py +++ b/openbb_platform/core/tests/provider/utils/websockets/test_database.py @@ -13,7 +13,7 @@ MOCK_MESSAGES = [ @pytest.fixture(scope="module") def database(): """Return a MessageQueue instance.""" - return Database(table_name="test_table") + return Database(table_name="test") def test_setup_database(database): @@ -61,7 +61,7 @@ def test_multiple_connections(database): ) another_db = Database( results_file=database.results_file, - table_name="other_table", + table_name="other_test", ) assert new_db.fetch_all()[0] == MOCK_MESSAGES[0] database.write_to_db(MOCK_MESSAGES[1]) @@ -87,16 +87,16 @@ def test_query_db(database): assert len(database.query(query)) == 2 query = "json_extract (message, '$.type') == 'quote'" assert len(database.query(query)) == 1 - query = "SELECT message FROM test_table WHERE json_extract (message, '$.type') = 'trade'" + query = "SELECT message FROM test WHERE json_extract (message, '$.type') = 'trade'" assert len(database.query(query)) == 2 - query = "SELECT json_extract (message, '$.symbol') FROM test_table WHERE json_extract (message, '$.type') = 'trade'" + query = "SELECT json_extract (message, '$.symbol') FROM test WHERE json_extract (message, '$.type') = 'trade'" assert database.query(query) == ["test1", "test3"] def test_limit(): """Test if the limit parameter is working and that the auto increment index doesn't reset when cleared.""" database = Database( - table_name="test_limit_table", + table_name="test_limit", limit=2, ) assert database @@ -112,5 +112,5 @@ def test_limit(): database.clear_results() assert database.fetch_all() == [] database.write_to_db(MOCK_MESSAGES[0]) - query = "SELECT id FROM test_limit_table" + query = "SELECT id FROM test_limit" assert database.query(query)[0] > 3 diff --git a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py deleted file mode 100644 index 950a7a41ac0..00000000000 --- a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py +++ /dev/null @@ -1,298 +0,0 @@ -"""Broadcast server for streaming results to connected clients via WebSocket.""" - -import asyncio -import json -import sys -from pathlib import Path -from typing import Optional - -import uvicorn -from fastapi import FastAPI, WebSocket, WebSocketDisconnect -from openbb_core.provider.utils.helpers import run_async -from openbb_core.provider.utils.websockets.database import Database -from openbb_core.provider.utils.websockets.helpers import get_logger, parse_kwargs -from starlette.websockets import WebSocketState - -connected_clients: set = set() - -kwargs = parse_kwargs() - -HOST = kwargs.pop("host", None) or "localhost" -PORT = kwargs.pop("port", None) or 6666 -PORT = int(PORT) - -RESULTS_FILE = kwargs.pop("results_file", None) -TABLE_NAME = kwargs.pop("table_name", None) or "records" -SLEEP_TIME = kwargs.pop("sleep_time", None) or 0.25 -AUTH_TOKEN = kwargs.pop("auth_token", None) - -DATABASE = Database(results_file=RESULTS_FILE, table_name=TABLE_NAME) - -app = FastAPI() - - -async def read_stdin(broadcast_server): - """Read from stdin.""" - while True: - line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) - sys.stdin.flush() - - if not line: - break - - try: - command = ( - json.loads(line.strip()) - if line.strip().startswith("{") or line.strip().startswith("[") - else line.strip() - ) - await broadcast_server.broadcast(json.dumps(command)) - except json.JSONDecodeError: - broadcast_server.logger.error("Invalid JSON received from stdin") - - -@app.websocket("/") -async def websocket_endpoint( # noqa: PLR0915 - websocket: WebSocket, auth_token: Optional[str] = None -): - """Connect to the broadcast server.""" - - broadcast_server = BroadcastServer( - RESULTS_FILE, - TABLE_NAME, - SLEEP_TIME, - str(AUTH_TOKEN), - ) - auth_token = str(auth_token) - - if ( - broadcast_server.auth_token is not None - and auth_token - != broadcast_server._decrypt_value( # pylint: disable=protected-access - broadcast_server.auth_token - ) - ): - await websocket.accept() - await websocket.send_text( - "ERROR: Invalid authentication token. Could not connect to the broadcast." - ) - broadcast_server.logger.error( - "ERROR: Invalid authentication token passed by a client connecting." - ) - await websocket.close(code=1008, reason="Invalid authentication token") - return - - await websocket.accept() - - if RESULTS_FILE is None: - raise ValueError("Results file path is required for WebSocket server.") - - broadcast_server.websocket = websocket - connected_clients.add(broadcast_server) - - stream_task = asyncio.create_task(broadcast_server.stream_results()) - stdin_task = asyncio.create_task(read_stdin(broadcast_server)) - try: - await websocket.receive_text() - - except WebSocketDisconnect: - pass - except Exception as e: # pylint: disable=broad-except - msg = f"Unexpected error: {e.__class__.__name__} -> {e}" - broadcast_server.logger.error(msg) - finally: - if broadcast_server in connected_clients: - connected_clients.remove(broadcast_server) - stream_task.cancel() - stdin_task.cancel() - try: - await stream_task - await stdin_task - except asyncio.CancelledError: - broadcast_server.logger.info("INFO: A listener task was cancelled.") - except Exception as e: # pylint: disable=broad-except - msg = f"Unexpected error while cancelling stream task: {e.__class__.__name__} -> {e}" - broadcast_server.logger.error(msg) - if websocket.client_state != WebSocketState.DISCONNECTED: - try: - await websocket.close() - except RuntimeError as e: - msg = f"Unexpected error while closing websocket: {e.__class__.__name__} -> {e}" - broadcast_server.logger.error(msg) - - -class BroadcastServer: # pylint: disable=too-many-instance-attributes - """Stream new results from a continuously written SQLite database. - - Not intended to be used directly, it is initialized by the server app when it accepts a new connection. - It is responsible for reading the results database and sending new messages to the connected client(s). - """ - - def __init__( - self, - results_file, - table_name, - sleep_time: float = 0.25, - auth_token: Optional[str] = None, - ): - """Initialize the BroadcastServer instance.""" - # pylint: disable=import-outside-toplevel - import os - - self.results_file = results_file - self.table_name = table_name - self.logger = get_logger("openbb.websocket.broadcast_server") - self.sleep_time = sleep_time - self._app = app - self._key = os.urandom(32) - self._iv = os.urandom(16) - self.auth_token = self._encrypt_value(auth_token) if auth_token else None - self.websocket = None - self.kwargs = kwargs - - def _encrypt_value(self, value: str) -> str: - """Encrypt the value for storage.""" - # pylint: disable=import-outside-toplevel - from openbb_core.provider.utils.websockets.helpers import encrypt_value - - return encrypt_value(self._key, self._iv, value) - - def _decrypt_value(self, value: str) -> str: - """Decrypt the value for use.""" - # pylint: disable=import-outside-toplevel - from openbb_core.provider.utils.websockets.helpers import decrypt_value - - return decrypt_value(self._key, self._iv, value) - - async def stream_results( # noqa: PLR0915 # pylint: disable=too-many-branches - self, - ): - """Continuously read the database and send new messages as JSON via WebSocket.""" - # pylint: disable=import-outside-toplevel - import sqlite3 # noqa - from openbb_core.app.model.abstract.error import OpenBBError - - file_path = Path(self.results_file).absolute() - last_id = 0 - - if not file_path.exists(): - self.logger.error("Results file not found: %s", str(file_path)) - return - conn = sqlite3.connect(self.results_file) - cursor = conn.cursor() - cursor.execute(f"SELECT MAX(id) FROM {self.table_name}") # noqa:S608 - last_id = cursor.fetchone()[0] or 0 - conn.close() - - try: # pylint: disable=too-many-nested-blocks - while True: - try: - if file_path.exists(): - conn = sqlite3.connect(self.results_file) - cursor = conn.cursor() - cursor.execute( - f"SELECT * FROM {self.table_name} WHERE id > ?", # noqa:S608 - (last_id,), - ) - rows = cursor.fetchall() - conn.close() - - if rows: - for row in rows: - _, message = row - await self.broadcast(json.dumps(json.loads(message))) - last_id = max(row[0] for row in rows) - else: - self.logger.error("Results file not found: %s", str(file_path)) - break - - await asyncio.sleep(self.sleep_time) - except KeyboardInterrupt: - self.logger.info("\nResults stream cancelled.") - break - except sqlite3.OperationalError as e: - if "no such table" in str(e): - self.logger.error( - "Results file was removed by the parent process." - ) - break - raise OpenBBError(e) from e - except asyncio.CancelledError: - break - except WebSocketDisconnect: - pass - except Exception as e: # pylint: disable=broad-except - msg = f"Unexpected error: {e.__class__.__name__} -> {e}" - self.logger.error(msg) - return - - async def broadcast(self, message: str): - """Broadcast a message to all connected connected clients.""" - disconnected_clients = set() - for client in connected_clients.copy(): - try: - await client.websocket.send_json(message) - except WebSocketDisconnect: - disconnected_clients.add(client) - except Exception as e: # pylint: disable=broad-except - msg = f"Unexpected error: {e.__class__.__name__} -> {e}" - self.logger.error(msg) - disconnected_clients.add(client) - # Remove disconnected connected clients - for client in disconnected_clients: - connected_clients.remove(client) - - def start_app(self, host: str = "127.0.0.1", port: int = 6666): - """Start the FastAPI app with Uvicorn.""" - uvicorn.run( - self._app, - host=host, - port=port, - **kwargs, - ) - - -def create_broadcast_server( - results_file: str, - table_name: str, - sleep_time: float = 0.25, - auth_token: Optional[str] = None, -): - """Create a new BroadcastServer instance.""" - return BroadcastServer(results_file, table_name, sleep_time, auth_token) - - -def main(): - """Run the main function.""" - broadcast_server = create_broadcast_server( - RESULTS_FILE, - TABLE_NAME, - SLEEP_TIME, - str(AUTH_TOKEN), - ) - - try: - broadcast_server.start_app( - host=HOST, - port=PORT, - **kwargs, - ) - except TypeError as e: - msg = ( - f"ERROR: Invalid keyword argument passed to unvicorn. -> {e.args[0]}\n" - ) - broadcast_server.logger.error(msg) - except KeyboardInterrupt: - broadcast_server.logger.info("INFO: Broadcast server terminated.") - finally: - sys.exit(0) - - -if __name__ == "__main__": - if not RESULTS_FILE: - raise ValueError("Results file path is required for Broadcast server.") - - if not Path(RESULTS_FILE).absolute().exists(): - run_async(DATABASE._setup_database) # pylint: disable=protected-access - - main() 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 9ecc4807c50..e8045f82558 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -7,7 +7,6 @@ import signal import sys import websockets -import websockets.exceptions from openbb_core.provider.utils.websockets.database import Database from openbb_core.provider.utils.websockets.helpers import ( get_logger, @@ -34,7 +33,7 @@ URL_MAP = { } logger = get_logger("openbb.websocket.polygon") -queue = MessageQueue(logger=logger) +queue = MessageQueue(logger=logger, backoff_factor=2) command_queue = MessageQueue(logger=logger) kwargs = parse_kwargs() @@ -159,11 +158,15 @@ async def read_stdin(): if not line: break - try: - command = json.loads(line.strip()) - await command_queue.enqueue(command) - except json.JSONDecodeError: - logger.error("Invalid JSON received from stdin") + if "qsize" in line: + logger.info(f"PROVIDER INFO: Queue size: {queue.queue.qsize()}") + + else: + try: + command = json.loads(line.strip()) + await command_queue.enqueue(command) + except json.JSONDecodeError: + logger.error("Invalid JSON received from stdin") async def process_stdin_queue(websocket): @@ -211,9 +214,17 @@ async def process_message(message): async def connect_and_stream(): """Connect to the WebSocket and stream data to file.""" + tasks: set = set() + handler_task = asyncio.create_task( queue.process_queue(lambda message: process_message(message)) ) + tasks.add(handler_task) + for i in range(0, 48): + new_task = asyncio.create_task( + queue.process_queue(lambda message: process_message(message)) + ) + tasks.add(new_task) stdin_task = asyncio.create_task(read_stdin()) try: connect_kwargs = CONNECT_KWARGS.copy() @@ -281,14 +292,19 @@ async def connect_and_stream(): sys.exit(1) except Exception as e: - msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e}" + msg = ( + f"PROVIDER ERROR: Unexpected error -> " + f"{e.__class__.__name__ if hasattr(e, '__class__') else e}: {e}" + ) logger.error(msg) sys.exit(1) finally: - handler_task.cancel() - stdin_task.cancel() - await asyncio.gather(handler_task, stdin_task, return_exceptions=True) + tasks.add(stdin_task) + for task in tasks: + task.cancel() + await task + asyncio.gather(*tasks, return_exceptions=True) sys.exit(0) @@ -296,7 +312,6 @@ if __name__ == "__main__": try: loop = asyncio.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(sig, handle_termination_signal, logger) @@ -315,5 +330,5 @@ if __name__ == "__main__": logger.error(ERR) finally: - loop.stop() + loop.close() sys.exit(0)