From ad113f53d3766348344b57a5e7eb6fce0c146822 Mon Sep 17 00:00:00 2001 From: Jakub Trllo <43494761+iLLiCiTiT@users.noreply.github.com> Date: Wed, 4 Nov 2026 09:48:50 +0100 Subject: [PATCH] Ping server when websocket receives nothing for a while Server sends heartbeat only when it did not publish any other message for 2 seconds, so while busy a client could receive nothing for longer than the receive timeout and consider the connection lost. Ping the server after 5 seconds without any frame and consider the connection lost only when nothing, not even pong, arrived for 15 seconds. Co-Authored-By: Claude Opus 5.5 --- ayon_api/events.py | 47 ++++++++-- tests/test_event_hub_lifecycle.py | 4 +- tests/test_event_hub_ping.py | 140 ++++++++++++++++++++++++++++++ 3 files changed, 180 insertions(+), 11 deletions(-) create mode 100644 tests/test_event_hub_ping.py diff --git a/ayon_api/events.py b/ayon_api/events.py index b5bc6d11a..5162b73f4 100644 --- a/ayon_api/events.py +++ b/ayon_api/events.py @@ -8,12 +8,13 @@ import os import re import threading +import time import typing from typing import Any, Callable import weakref import requests -from websocket import ABNF +from websocket import ABNF, WebSocketTimeoutException from .exceptions import UnauthorizedError @@ -485,9 +486,14 @@ def _validate_ref(self) -> None: self._partial_func = None -# Server sends 'heartbeat' message to authorized clients when there was no -# other message for ~5 seconds. If nothing arrives for this amount of time -# the connection is considered dead. +# Server sends 'heartbeat' message only when it did not publish any other +# message for 2 seconds, and sends messages only to clients subscribed to +# their topic. So a client might not receive anything for a long time +# while server is busy (e.g. many events during a sync). Ping the server +# when nothing was received for this amount of time. +_PING_INTERVAL: float = 5.0 +# Connection is considered dead if nothing (not even pong) was received +# for this amount of time. _RECV_TIMEOUT: float = 15.0 # Timeout for opening websocket connection and for the auth check request. _CONNECT_TIMEOUT: float = 5.0 @@ -525,6 +531,9 @@ class _LoopState: auth_required: bool = False # Some message was received with current websocket connection received_message: bool = False + # When last frame (message, ping or pong) was received + # ('time.monotonic') + last_received: float = 0.0 # Connection state reported via events # - 'None' means unknown (nothing was reported yet) connected: bool | None = None @@ -845,9 +854,9 @@ def _update_topics(self) -> None: def _create_ws_connection(self) -> WebSocket: ws_con = self._connection.create_websocket("ws") - # Server sends heartbeat periodically, so timeout on receive means - # that the connection is dead. - ws_con.settimeout(_RECV_TIMEOUT) + # Timeout on receive does not mean that the connection is dead, + # the server is pinged on timeout (see '_loop'). + ws_con.settimeout(_PING_INTERVAL) return ws_con def _close_ws_connection(self) -> None: @@ -994,6 +1003,7 @@ def _loop(self, stop_event: threading.Event) -> None: break self._ws_connection = ws_con state.received_message = False + state.last_received = time.monotonic() state.auth_required = True if state.auth_required: @@ -1010,16 +1020,35 @@ def _loop(self, stop_event: threading.Event) -> None: continue try: - op_code, message = ws_con.recv_data() + op_code, message = ws_con.recv_data(control_frame=True) + except WebSocketTimeoutException as exc: + elapsed = time.monotonic() - state.last_received + if elapsed >= _RECV_TIMEOUT: + # Server did not respond to ping + self._on_connection_lost(stop_event, exc) + continue + # Nothing received for a while, server might be busy and + # not sending heartbeat. Pong will confirm it is alive. + try: + ws_con.ping() + except Exception as ping_exc: + self._on_connection_lost(stop_event, ping_exc) + continue except Exception as exc: - # Timeout (no heartbeat), closed connection, socket error... + # Closed connection, socket error... self._on_connection_lost(stop_event, exc) continue + state.last_received = time.monotonic() if op_code == ABNF.OPCODE_CLOSE: self._on_connection_lost(stop_event) continue + # Ping and pong only confirm that the connection is alive, they + # don't tell if the client is authorized. + if op_code in (ABNF.OPCODE_PING, ABNF.OPCODE_PONG): + continue + if not state.received_message: state.received_message = True self._set_connected(True) diff --git a/tests/test_event_hub_lifecycle.py b/tests/test_event_hub_lifecycle.py index 3db7e3444..8ea711ff4 100644 --- a/tests/test_event_hub_lifecycle.py +++ b/tests/test_event_hub_lifecycle.py @@ -54,12 +54,12 @@ def send(self, data): if message["token"] in self._rejected_tokens: self._aborted.set() - def recv_data(self): + def recv_data(self, control_frame=False): if self._aborted.wait(0.02): raise WebSocketConnectionClosedException("closed") if not self.sent: # Server does not send anything before auth message - return self.recv_data() + return self.recv_data(control_frame) return ABNF.OPCODE_TEXT, json.dumps({"topic": "heartbeat"}) def abort(self): diff --git a/tests/test_event_hub_ping.py b/tests/test_event_hub_ping.py new file mode 100644 index 000000000..b14885b89 --- /dev/null +++ b/tests/test_event_hub_ping.py @@ -0,0 +1,140 @@ +"""Tests for detection of lost websocket connection of 'EventHub'.""" + +from __future__ import annotations + +import json +import threading + +import pytest +from websocket import ABNF, WebSocketTimeoutException + +from ayon_api import events as events_module +from ayon_api.events import EventHub + +HEARTBEAT = (ABNF.OPCODE_TEXT, json.dumps({"topic": "heartbeat"})) +PONG = (ABNF.OPCODE_PONG, b"") +SERVER_PING = (ABNF.OPCODE_PING, b"") +TIMEOUT = "timeout" + + +class FakeClock: + def __init__(self) -> None: + self.value = 1000.0 + + def monotonic(self) -> float: + return self.value + + +class FakeWebSocket: + """Replays scripted results of 'recv_data'. + + Timeout moves the clock forward by the socket timeout as a real + socket would. When the script is exhausted the test is stopped. + + """ + def __init__(self, clock, responses, stop_event) -> None: + self._clock = clock + self._responses = list(responses) + self._stop_event = stop_event + self.pings = 0 + + def send(self, data) -> None: + pass + + def ping(self) -> None: + self.pings += 1 + + def abort(self) -> None: + pass + + def recv_data(self, control_frame: bool = False): + if not self._responses: + self._stop_event.set() + raise ConnectionError("End of test") + response = self._responses.pop(0) + if response == TIMEOUT: + self._clock.value += events_module._PING_INTERVAL + raise WebSocketTimeoutException("timed out") + return response + + +class FakeConnection: + access_token = "token" + + +class EventsRecorder: + """Callback storing received topics, callbacks are stored weakly.""" + + def __init__(self) -> None: + self.topics: list[str] = [] + + def on_event(self, event) -> None: + self.topics.append(event.topic) + + +@pytest.fixture +def run_hub(monkeypatch): + """Run event hub loop with scripted websocket responses.""" + clock = FakeClock() + monkeypatch.setattr(events_module, "time", clock) + monkeypatch.setattr(events_module, "_RECONNECT_DELAYS", (0.0,)) + + def _run_hub(responses): + stop_event = threading.Event() + fake_ws = FakeWebSocket(clock, responses, stop_event) + connections = [fake_ws] + + def create_ws_connection(): + if connections: + return connections.pop(0) + stop_event.set() + raise ConnectionError("End of test") + + hub = EventHub(FakeConnection()) + monkeypatch.setattr( + hub, "_create_ws_connection", create_ws_connection + ) + monkeypatch.setattr(hub, "_check_token", lambda: True) + recorder = EventsRecorder() + hub.add_callback("connection.*", recorder.on_event, connect=False) + hub._loop(stop_event) + return recorder.topics, fake_ws + + return _run_hub + + +def test_busy_server_without_heartbeat_is_connected(run_hub): + """Server answering pings is alive even if it does not send heartbeat. + + Server does not send heartbeat while it publishes other messages, + which might take much longer than '_RECV_TIMEOUT'. + + """ + responses = [HEARTBEAT] + [TIMEOUT, PONG] * 10 + topics, fake_ws = run_hub(responses) + + assert topics == ["connection.opened"] + assert fake_ws.pings == 10 + + +def test_server_ping_keeps_connection_alive(run_hub): + responses = [HEARTBEAT] + [TIMEOUT, TIMEOUT, SERVER_PING] * 5 + topics, fake_ws = run_hub(responses) + + assert topics == ["connection.opened"] + + +def test_no_response_to_ping_is_connection_lost(run_hub): + responses = [HEARTBEAT, TIMEOUT, TIMEOUT, TIMEOUT] + topics, fake_ws = run_hub(responses) + + assert topics == ["connection.opened", "connection.closed"] + assert fake_ws.pings == 2 + + +def test_pong_does_not_report_connected(run_hub): + """Pong does not tell if client is authorized.""" + responses = [PONG, PONG] + topics, fake_ws = run_hub(responses) + + assert "connection.opened" not in topics