add some micro scaling to the queue when it is growing

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
 

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}..."

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

openbb_platform/extensions/websockets/openbb_websockets/broadcast.py

@@ -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()

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)