[Fix] Massive Logs Support to CLI

app/main.py

@@ -20,6 +20,7 @@
 from app.api.api_v1.api import api_router
 from app.core.config import settings
 from app.test_engine.test_script_manager import test_script_manager
+from app.uvicorn_worker import WS_PING_TIMEOUT_S
 
 app = FastAPI(
     title=settings.PROJECT_NAME, openapi_url=f"{settings.API_V1_STR}/openapi.json"
@@ -51,4 +52,9 @@ async def startup_event() -> None:
 
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=80, log_config=None)
+    # ws_ping_timeout raised from uvicorn's 20s default: a large log flush can
+    # keep the event loop busy long enough to miss the keepalive pong and get
+    # the websocket dropped mid-broadcast (see app.uvicorn_worker).
+    uvicorn.run(
+        app, host="0.0.0.0", port=80, log_config=None, ws_ping_timeout=WS_PING_TIMEOUT_S
+    )

app/socket_connection_manager.py

@@ -24,7 +24,7 @@
 from fastapi.websockets import WebSocketDisconnect
 from loguru import logger
 from starlette.websockets import WebSocketState
-from websockets.exceptions import ConnectionClosedOK
+from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK
 
 from app.constants.shared_constants import MessageKeysEnum, MessageTypeEnum
 from app.constants.websockets_constants import (
@@ -129,17 +129,33 @@ async def broadcast(self, message: Union[str, dict, list]) -> None:
         # Convert dictionaries and lists to string using json
         if isinstance(message, dict) or isinstance(message, list):
             message = json.dumps(message, default=pydantic.json.pydantic_encoder)
-        for connection in self.active_connections:
+        # Iterate over a copy: disconnect() below mutates active_connections.
+        for connection in list(self.active_connections):
             if connection.type == WebSocketTypeEnum.MAIN:
                 websocket = connection.websocket
                 try:
                     await websocket.send_text(message)
                 # Starlette raises websockets.exceptions.ConnectionClosedOK
                 # when trying to send to a closed websocket.
                 # https://github.com/encode/starlette/issues/759
-                except ConnectionClosedOK:
+                # ConnectionClosedError is raised on an abrupt drop (e.g. a
+                # missed websocket keepalive pong) rather than a graceful
+                # close - treat it the same way so a dead connection doesn't
+                # keep being retried for the rest of the process's life.
+                except (ConnectionClosedOK, ConnectionClosedError):
                     if websocket.application_state != WebSocketState.DISCONNECTED:
-                        await websocket.close()
+                        try:
+                            await websocket.close()
+                        except Exception as close_error:
+                            # The transport can already be gone by this point
+                            # (that's the whole reason we're here) - closing
+                            # an already-dead socket failing is expected and
+                            # uninteresting, but must not prevent the
+                            # cleanup below from running.
+                            logger.debug(
+                                f"Error closing already-dead websocket: {close_error}"
+                            )
+                    self.disconnect(connection)
                     logger.warning(
                         f'Failed to send message: "{message}"'
                         f' to websocket: "{websocket}", connection closed."'

app/test_engine/test_db_observer.py

@@ -13,8 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import asyncio
 from datetime import datetime
-from queue import Empty, Queue
 from typing import Callable, Generator, Union
 
 from loguru import logger
@@ -30,6 +30,10 @@
 from app.test_engine.models import TestCase, TestRun, TestStep, TestSuite
 from app.test_engine.test_observer import Observer
 
+ExecutionObj = Union[
+    TestCaseExecution, TestStepExecution, TestSuiteExecution, TestRunExecution
+]
+
 
 class TestDBObserver(Observer):
     __test__ = False  # Needed to indicate to PyTest that this is not a "test"
@@ -38,15 +42,19 @@ def __init__(
         self, db_generator: Callable[[], Generator[Session, None, None]] = get_db
     ) -> None:
         self.__db_generator = db_generator
-        self.data_queue: Queue = Queue()
-
-    def apply_updates(self) -> None:
-        while not self.data_queue.empty():
-            try:
-                data = self.data_queue.get(timeout=0.1)
-                self.__save(data)
-            except Empty:
-                pass
+        # Keyed by id(execution_obj) instead of a plain queue: dispatch() can
+        # re-notify the SAME object many times (e.g. once per 0.5s log-flush
+        # tick for the whole run) before apply_updates() ever drains this -
+        # a dict collapses those into a single pending save per distinct
+        # object instead of committing the same (increasingly large) object
+        # redundantly once per notification.
+        self.__pending: dict[int, ExecutionObj] = {}
+
+    async def apply_updates(self) -> None:
+        pending = self.__pending
+        self.__pending = {}
+        for data in pending.values():
+            await self.__save(data)
 
     def dispatch(
         self, observable: Union[TestRun, TestSuite, TestCase, TestStep]
@@ -62,6 +70,9 @@ def dispatch(
             elif isinstance(observable, TestStep):
                 self.__onTestStepUpdate(observable)
 
+    def __enqueue(self, execution_obj: ExecutionObj) -> None:
+        self.__pending[id(execution_obj)] = execution_obj
+
     def __onTestRunUpdate(self, observable: "TestRun") -> None:
         logger.debug("Test Run Observer received", observable)
         test_run_execution = observable.test_run_execution
@@ -74,7 +85,7 @@ def __onTestRunUpdate(self, observable: "TestRun") -> None:
         if self.isCompleted(observable.state):
             test_run_execution.completed_at = datetime.now()
 
-        self.data_queue.put(test_run_execution)
+        self.__enqueue(test_run_execution)
 
     def __onTestSuiteUpdate(self, observable: "TestSuite") -> None:
         logger.debug("Test Suite Observer received", observable)
@@ -89,7 +100,7 @@ def __onTestSuiteUpdate(self, observable: "TestSuite") -> None:
             if self.isCompleted(observable.state):
                 observable.test_suite_execution.completed_at = datetime.now()
 
-            self.data_queue.put(observable.test_suite_execution)
+            self.__enqueue(observable.test_suite_execution)
 
     def __onTestCaseUpdate(self, observable: "TestCase") -> None:
         logger.debug("Test Case Observer received", observable)
@@ -104,7 +115,7 @@ def __onTestCaseUpdate(self, observable: "TestCase") -> None:
             if self.isCompleted(observable.state):
                 observable.test_case_execution.completed_at = datetime.now()
 
-            self.data_queue.put(observable.test_case_execution)
+            self.__enqueue(observable.test_case_execution)
 
     def __onTestStepUpdate(self, observable: "TestStep") -> None:
         logger.debug("Test Step Observer received", observable)
@@ -121,14 +132,9 @@ def __onTestStepUpdate(self, observable: "TestStep") -> None:
             if self.isCompleted(observable.state):
                 observable.test_step_execution.completed_at = datetime.now()
 
-            self.data_queue.put(observable.test_step_execution)
+            self.__enqueue(observable.test_step_execution)
 
-    def __save(
-        self,
-        execution_obj: Union[
-            TestCaseExecution, TestStepExecution, TestSuiteExecution, TestRunExecution
-        ],
-    ) -> None:
+    async def __save(self, execution_obj: ExecutionObj) -> None:
         # We get the session from the model it self to avoid overriding values when
         # using a different session
         insp = inspect(execution_obj)
@@ -139,7 +145,11 @@ def __save(
             session = next(self.__db_generator())
             session.add(execution_obj)
         session.expire_on_commit = False
-        session.commit()
+        # session.commit() is a blocking, synchronous SQLAlchemy call. It can
+        # be a large write (e.g. a run's full log), so keep it off the event
+        # loop rather than stalling every other coroutine (websocket pings,
+        # other requests) for however long it takes.
+        await asyncio.to_thread(session.commit)
         logger.debug(
             f"Saved {execution_obj.__class__} {execution_obj.id}"
             f" with state {execution_obj.state}"

app/test_engine/test_runner.py

@@ -181,7 +181,7 @@ async def run(self) -> None:
             self.test_run.unsubscribe([ui_observer, db_observer])
 
             # Flush all pending DB updates
-            db_observer.apply_updates()
+            await db_observer.apply_updates()
 
             # Ensure all state updates are sent to the frontend
             await ui_observer.complete_tasks()

app/test_engine/test_ui_observer.py

@@ -38,11 +38,30 @@ class TestUpdateTypeEnum(str, Enum):
     TEST_CASE = "Test Case"
 
 
+# Maximum number of log entries broadcast in a single websocket message. A
+# dense burst of log lines (e.g. a large conformance report) can otherwise
+# accumulate thousands of entries into one flush, producing a single
+# multi-MB message whose synchronous JSON serialization (see
+# SocketConnectionManager.broadcast) has no yield point. Splitting into
+# multiple smaller messages, sent in order within a single task (see
+# __handle_test_run_log), lets the event loop interleave other work (e.g.
+# websocket keepalives) between chunks.
+LOG_RECORDS_BROADCAST_CHUNK_SIZE = 200
+
+
 class TestUIObserver(Observer):
     __test__ = False
-    __async_updates: list[Task] = []
-    __last_seen_run_state: Optional[TestStateEnum] = None
-    __last_seen_run_log_len = 0
+
+    def __init__(self) -> None:
+        # Instance state - must not be class-level mutable defaults. A list
+        # declared at class scope is shared across every TestUIObserver
+        # instance until an instance reassigns it; TestRunner creates a
+        # fresh instance per run, so a shared, ever-growing task list would
+        # leak Task references (and their retained coroutine state) across
+        # every run for the life of the process.
+        self.__async_updates: list[Task] = []
+        self.__last_seen_run_state: Optional[TestStateEnum] = None
+        self.__last_seen_run_log_len = 0
 
     def dispatch(
         self, observable: Union[TestRun, TestSuite, TestCase, TestStep]
@@ -81,7 +100,11 @@ def __handle_test_run_log(self, test_run: TestRun) -> None:
         log_len = len(test_run.log)
         if log_len > self.__last_seen_run_log_len:
             new_entries = test_run.log[self.__last_seen_run_log_len :]
-            self.__send_log_records_message(new_entries)
+            chunks = [
+                new_entries[i : i + LOG_RECORDS_BROADCAST_CHUNK_SIZE]
+                for i in range(0, len(new_entries), LOG_RECORDS_BROADCAST_CHUNK_SIZE)
+            ]
+            self.__send_log_records_messages_in_order(chunks)
             self.__last_seen_run_log_len = log_len
 
     def __onTestSuiteUpdate(self, observable: TestSuite) -> None:
@@ -138,13 +161,33 @@ def __send_test_update_message(self, update_payload: dict) -> None:
             }
         )
 
-    def __send_log_records_message(self, log_entries: list[TestRunLogEntry]) -> None:
-        self.__send_message(
-            {
-                MessageKeysEnum.TYPE: MessageTypeEnum.TEST_LOG_RECORDS,
-                MessageKeysEnum.PAYLOAD: log_entries,
-            }
-        )
+    def __send_log_records_messages_in_order(
+        self, chunks: list[list[TestRunLogEntry]]
+    ) -> None:
+        """Broadcast each chunk of one flush as its own TEST_LOG_RECORDS
+        message, in order.
+
+        Scheduled as a single task for the whole flush, not one task per
+        chunk: chunks must be delivered in order, but independently
+        scheduled tasks don't guarantee that - SocketConnectionManager.
+        broadcast()'s websocket.send_text() is genuine async I/O that can
+        yield under backpressure, which could let a later chunk's task
+        complete first. Awaiting each chunk's broadcast sequentially within
+        one task guarantees order while still yielding the event loop
+        between chunks (each await is itself a yield point).
+        """
+
+        async def _send_in_order() -> None:
+            for chunk in chunks:
+                await socket_connection_manager.broadcast(
+                    {
+                        MessageKeysEnum.TYPE: MessageTypeEnum.TEST_LOG_RECORDS,
+                        MessageKeysEnum.PAYLOAD: chunk,
+                    }
+                )
+
+        task = create_task(_send_in_order())
+        self.__async_updates.append(task)
 
     def __send_message(self, message: dict[str, Any]) -> None:
         # enqueue update

app/tests/socket_connection_manager/test_socket_connection_manager.py

@@ -20,7 +20,7 @@
 import pytest
 from fastapi import WebSocket
 from starlette.websockets import WebSocketState
-from websockets.exceptions import ConnectionClosedOK
+from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK
 
 from app.constants.shared_constants import MessageKeysEnum, MessageTypeEnum
 from app.constants.websockets_constants import (
@@ -203,6 +203,39 @@ async def test_broadcast_failed_for_ConnectionClosed() -> None:
     socket_connection_manager.active_connections.clear()
 
 
+@pytest.mark.asyncio
+async def test_broadcast_failed_for_ConnectionClosedError() -> None:
+    """
+    Tests that broadcast() also handles ConnectionClosedError (an abrupt
+    drop, e.g. from a missed keepalive ping under event-loop stall) the same
+    way as ConnectionClosedOK: the socket is closed and the dead connection
+    is removed from active_connections so it isn't retried for the rest of
+    the process's life (regression test for issue #1072).
+    """
+    test_message = "Test"
+    socket_connection_manager.active_connections.clear()
+
+    # Add a websocket object to the "active_connections" list to imitate
+    # an existing active connection
+    socket = mock.MagicMock(spec=WebSocket)
+    socket.application_state = WebSocketState.CONNECTED
+    connection = WebSocketConnection(socket, WebSocketTypeEnum.MAIN)
+
+    socket_connection_manager.active_connections.append(connection)
+    assert len(socket_connection_manager.active_connections) == 1
+
+    # Force an abrupt connection-closed exception
+    socket.send_text.side_effect = ConnectionClosedError(rcvd=None, sent=None)
+
+    await socket_connection_manager.broadcast(message=test_message)
+    socket.send_text.assert_called_once_with(test_message)
+    socket.close.assert_called_once()
+    assert connection not in socket_connection_manager.active_connections
+
+    # Cleanup
+    socket_connection_manager.active_connections.clear()
+
+
 @pytest.mark.asyncio
 async def test_broadcast_failed_for_RuntimeError() -> None:
     """

app/tests/test_engine/test_db_observer.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 import asyncio
+from unittest import mock
 
 import pytest
 from sqlalchemy.orm import Session
@@ -281,3 +282,28 @@ def test_test_db_observer_update(db: Session) -> None:
     db.close()
     test_db_observer.dispatch(test_step)
     assert TestStateEnum.EXECUTING == test_step.test_step_execution.state
+
+
+@pytest.mark.asyncio
+async def test_test_db_observer_dedups_repeated_run_updates(db: Session) -> None:
+    """Repeated dispatch() calls for the same TestRunExecution before
+    apply_updates() runs must collapse into a single commit, instead of one
+    redundant commit per dispatch (regression test for issue #1072's
+    DB-write storm, where ~1,400 redundant commits of the same growing log
+    blocked the event loop for the whole duration of a large test run)."""
+    test_script_manager = TestScriptManager()
+    test_db_observer = TestDBObserver()
+
+    test_run_execution = create_test_run_execution_with_some_test_cases(db=db)
+    test_run = test_script_manager.get_test_run(db, test_run_execution)
+    test_run.state = TestStateEnum.EXECUTING
+
+    with mock.patch.object(Session, "commit") as mock_commit:
+        # Simulate several rapid dispatch() calls for the same object, as
+        # happens once per ~0.5s log-flush tick during a run.
+        for _ in range(5):
+            test_db_observer.dispatch(test_run)
+
+        await test_db_observer.apply_updates()
+
+    assert mock_commit.call_count == 1