Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion app/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
)
24 changes: 20 additions & 4 deletions app/socket_connection_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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."'
Expand Down
52 changes: 31 additions & 21 deletions app/test_engine/test_db_observer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Expand All @@ -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]
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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}"
Expand Down
2 changes: 1 addition & 1 deletion app/test_engine/test_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
65 changes: 54 additions & 11 deletions app/test_engine/test_ui_observer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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:
"""
Expand Down
26 changes: 26 additions & 0 deletions app/tests/test_engine/test_db_observer.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# limitations under the License.
#
import asyncio
from unittest import mock

import pytest
from sqlalchemy.orm import Session
Expand Down Expand Up @@ -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
Loading
Loading