EventHub: Don't report lost connection while server is busy

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)

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):

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