fix: refine events websocket reconnect error handling and logging

After the first successful connection, delegate to the default `websockets` transient/fatal classification
instead of retrying every error. Log graceful and abnormal closes (with close code and reason) as well as
reconnect success, and cover both close paths with parametrized tests.

Author: vdusek

src/apify/events/_apify_event_manager.py

@@ -92,8 +92,8 @@ async def __aexit__(
         exc_value: BaseException | None,
         exc_traceback: TracebackType | None,
     ) -> None:
-        # Cancel the message-processing task first so that closing the websocket below is not treated
-        # as a dropped connection and followed by a reconnect attempt.
+        # Cancel the task before closing the websocket so that the closed connection is not treated as a drop
+        # and followed by a reconnect attempt.
         if self._process_platform_messages_task and not self._process_platform_messages_task.done():
             self._process_platform_messages_task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
@@ -104,21 +104,28 @@ async def __aexit__(
 
         await super().__aexit__(exc_type, exc_value, exc_traceback)
 
-    async def _process_platform_messages(self, ws_url: str) -> None:
-        def process_exception(exc: Exception) -> Exception | None:
-            # Until the first connection succeeds, treat every error as fatal so that `__aenter__` fails fast.
-            # Afterwards, treat every error as transient — the reconnect iterator keeps retrying with backoff
-            # so that platform events (e.g. `MIGRATING`) are not missed for the rest of the run.
-            if self._connected_to_platform_websocket is None or not self._connected_to_platform_websocket.done():
-                return exc
-            return None
+    def _process_connection_exception(self, exc: Exception) -> Exception | None:
+        """Decide whether a failed connection attempt to the platform websocket should be retried.
+
+        Before the first successful connection, every error is fatal so that `__aenter__` fails fast. After that,
+        the default `websockets` behavior decides which errors are transient and retried with exponential backoff.
+        """
+        if self._connected_to_platform_websocket and self._connected_to_platform_websocket.done():
+            return websockets.asyncio.client.process_exception(exc)
+        return exc
 
+    async def _process_platform_messages(self, ws_url: str) -> None:
         try:
-            async for websocket in websockets.asyncio.client.connect(ws_url, process_exception=process_exception):
+            # Used as an async iterator, `connect` reconnects with exponential backoff whenever a connection
+            # attempt fails with a transient error.
+            async for websocket in websockets.asyncio.client.connect(
+                ws_url, process_exception=self._process_connection_exception
+            ):
                 self._platform_events_websocket = websocket
-                connected_future = self._connected_to_platform_websocket
-                if connected_future is not None and not connected_future.done():
-                    connected_future.set_result(True)
+                if self._connected_to_platform_websocket and not self._connected_to_platform_websocket.done():
+                    self._connected_to_platform_websocket.set_result(True)
+                else:
+                    logger.info('Reconnected to the platform events websocket.')
 
                 try:
                     async for message in websocket:
@@ -150,12 +157,15 @@ def process_exception(exc: Exception) -> Exception | None:
                         except Exception:
                             logger.exception('Cannot parse Actor event', extra={'raw_message': message})
                 except websockets.exceptions.ConnectionClosed:
-                    pass
-
-                logger.warning(
-                    f'Connection to platform events websocket was closed '
-                    f'(code={websocket.close_code}, reason={websocket.close_reason!r}), reconnecting...'
-                )
+                    logger.warning(
+                        f'Connection to platform events websocket was lost '
+                        f'(code={websocket.close_code}, reason={websocket.close_reason!r}), reconnecting...'
+                    )
+                else:
+                    logger.info(
+                        f'Connection to platform events websocket was closed '
+                        f'(code={websocket.close_code}, reason={websocket.close_reason!r}), reconnecting...'
+                    )
         except Exception:
             logger.exception('Error in websocket connection')
             if self._connected_to_platform_websocket is not None and not self._connected_to_platform_websocket.done():

tests/unit/events/test_apify_event_manager.py

@@ -25,6 +25,18 @@
     from collections.abc import AsyncGenerator, Callable
 
 
+DUMMY_SYSTEM_INFO = {
+    'memAvgBytes': 19328860.328293584,
+    'memCurrentBytes': 65171456,
+    'memMaxBytes': 65171456,
+    'cpuAvgUsage': 2.0761105633130397,
+    'cpuMaxUsage': 53.941134593993326,
+    'cpuCurrentUsage': 8.45549815498155,
+    'isCpuOverloaded': False,
+    'createdAt': '2024-08-09T16:04:16.161Z',
+}
+
+
 @contextlib.asynccontextmanager
 async def _platform_ws_server(
     monkeypatch: pytest.MonkeyPatch,
@@ -188,17 +200,7 @@ async def send_platform_event(event_name: Event, data: Any = None) -> None:
 
             websockets.broadcast(connected_ws_clients, json.dumps(message))
 
-        dummy_system_info = {
-            'memAvgBytes': 19328860.328293584,
-            'memCurrentBytes': 65171456,
-            'memMaxBytes': 65171456,
-            'cpuAvgUsage': 2.0761105633130397,
-            'cpuMaxUsage': 53.941134593993326,
-            'cpuCurrentUsage': 8.45549815498155,
-            'isCpuOverloaded': False,
-            'createdAt': '2024-08-09T16:04:16.161Z',
-        }
-        SystemInfoEventData.model_validate(dummy_system_info)
+        SystemInfoEventData.model_validate(DUMMY_SYSTEM_INFO)
 
         async with ApifyEventManager(Configuration.get_global_configuration()) as event_manager:
             await client_connected.wait()
@@ -210,7 +212,7 @@ def listener(data: Any) -> None:
             event_manager.on(event=Event.SYSTEM_INFO, listener=listener)
 
             # Test sending event with data
-            await send_platform_event(Event.SYSTEM_INFO, dummy_system_info)
+            await send_platform_event(Event.SYSTEM_INFO, DUMMY_SYSTEM_INFO)
             await poll_until_condition(lambda: len(event_calls) == 1, poll_interval=0.05)
             assert len(event_calls) == 1
             assert event_calls[0] is not None
@@ -319,32 +321,54 @@ def migrating_listener(data: Any) -> None:
         assert len(migration_persist_events) >= 1
 
 
-async def test_websocket_reconnects_after_connection_drop(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Test that after a mid-stream websocket drop, the manager reconnects and keeps receiving platform events."""
+@pytest.mark.parametrize(
+    ('close_code', 'expected_log'),
+    [
+        pytest.param(1000, 'Connection to platform events websocket was closed (code=1000', id='graceful_close'),
+        pytest.param(1011, 'Connection to platform events websocket was lost (code=1011', id='abnormal_close'),
+    ],
+)
+async def test_websocket_reconnects_after_connection_drop(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, close_code: int, expected_log: str
+) -> None:
+    """Test that the event manager logs a websocket drop, reconnects, and keeps receiving platform events.
+
+    Also a regression test for the resolved `_connected_to_platform_websocket` future: a mid-stream disconnect
+    must not kill the message-processing task with `InvalidStateError`.
+    """
+    caplog.set_level(logging.INFO, logger='apify')
     async with (
         _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected),
         ApifyEventManager(Configuration.get_global_configuration()) as event_manager,
     ):
         await client_connected.wait()
-        aborting_calls: list[Any] = []
-
-        def listener(data: Any) -> None:
-            aborting_calls.append(data)
+        assert len(connected_ws_clients) == 1
 
-        event_manager.on(event=Event.ABORTING, listener=listener)
+        event_calls: list[Any] = []
+        event_manager.on(event=Event.SYSTEM_INFO, listener=event_calls.append)
 
-        # Drop the connection abnormally from the server side.
+        # Drop the connection from the server side and wait for the client to reconnect.
         client_connected.clear()
         for ws in list(connected_ws_clients):
-            await ws.close(code=1011, reason='Simulated server error')
+            await ws.close(code=close_code, reason='Simulated connection drop')
+        await asyncio.wait_for(client_connected.wait(), timeout=10)
+        # Poll because the old server-side handler may not have deregistered its connection yet.
+        await poll_until_condition(lambda: len(connected_ws_clients) == 1, poll_interval=0.05)
+        assert len(connected_ws_clients) == 1
+
+        # The message-processing task must have survived the drop.
+        task = event_manager._process_platform_messages_task
+        assert task is not None
+        assert not task.done()
 
-        # The event manager should reconnect on its own.
-        await asyncio.wait_for(client_connected.wait(), timeout=5.0)
+        # Events sent over the new connection must still be emitted.
+        websockets.broadcast(connected_ws_clients, json.dumps({'name': 'systemInfo', 'data': DUMMY_SYSTEM_INFO}))
+        await poll_until_condition(lambda: len(event_calls) == 1, poll_interval=0.05)
+        assert len(event_calls) == 1
 
-        # Events sent over the new connection must still be received.
-        websockets.broadcast(connected_ws_clients, json.dumps({'name': 'aborting'}))
-        await poll_until_condition(lambda: bool(aborting_calls), poll_interval=0.05)
-        assert len(aborting_calls) == 1
+        # Both the drop and the successful reconnect must be logged.
+        assert expected_log in caplog.text
+        assert 'Reconnected to the platform events websocket.' in caplog.text
 
 
 async def test_malformed_message_logs_exception(