Recover from transient read failures instead of aborting the client

firebase_messaging/fcmpushclient.py

@@ -119,6 +119,15 @@ class FcmPushClientConfig:  # pylint:disable=too-many-instance-attributes
     """Number of sequential errors of the same time to wait before aborting.
         If set to None the client will not abort."""
 
+    abort_on_connection_failure: bool = False
+    """If False (default) the client keeps retrying the connection indefinitely
+        (with backoff) when a reconnect fails, so it can self-heal once
+        connectivity returns. Set to True to shut down after
+        connection_retry_count attempts (legacy behaviour)."""
+
+    reconnect_backoff_max: float = 300
+    """Maximum delay in seconds between transient-read reconnect attempts."""
+
     monitor_interval: float = 1
     """Time in seconds for the monitor task to fire and check for heartbeats,
         stale connections and shut down of the main event loop."""
@@ -171,6 +180,7 @@ def __init__(
         self.do_listen = False
         self.sequential_error_counters: dict[ErrorType, int] = {}
         self.log_warn_counters: dict[str, int] = {}
+        self._transient_error_count = 0
 
         # reset variables
         self.input_stream_id = 0
@@ -234,13 +244,22 @@ async def _reset(self) -> None:
                 await asyncio.sleep(self.config.reset_interval - time_since_last_login)
 
             _logger.debug("Reestablishing connection")
-            if not await self._connect_with_retry():
-                _logger.error(
-                    "Unable to connect to MCS endpoint "
-                    + "after %s tries, shutting down",
+            while self.do_listen and not await self._connect_with_retry():
+                if self.config.abort_on_connection_failure:
+                    _logger.error(
+                        "Unable to connect to MCS endpoint "
+                        + "after %s tries, shutting down",
+                        self.config.connection_retry_count,
+                    )
+                    self._terminate()
+                    return
+                self._log_warn_with_limit(
+                    "Unable to connect to MCS endpoint after %s tries; "
+                    "retrying (will not give up).",
                     self.config.connection_retry_count,
                 )
-                self._terminate()
+                await asyncio.sleep(self.config.reset_interval)
+            if not self.do_listen:
                 return
             _logger.debug("Re-connected to ssl socket")
 
@@ -723,6 +742,7 @@ async def _listen(self) -> None:
                                 self.run_state,
                             )
                     elif msg := await self._receive_msg():
+                        self._transient_error_count = 0
                         await self._handle_message(msg)
 
                 except (OSError, EOFError) as osex:
@@ -751,6 +771,27 @@ async def _listen(self) -> None:
                                 "Expected read error during reset: %s",
                                 type(osex).__name__,
                             )
+                    elif isinstance(osex, (asyncio.IncompleteReadError, TimeoutError)):
+                        # Transient stream read (e.g. a 0-byte read, or "SSL
+                        # shutdown timed out" after a stuck writer close):
+                        # reconnect with exponential backoff, without advancing
+                        # the CONNECTION abort counter.
+                        self._transient_error_count += 1
+                        delay = min(
+                            self.config.reset_interval
+                            * (2 ** min(self._transient_error_count - 1, 8)),
+                            self.config.reconnect_backoff_max,
+                        )
+                        self._log_warn_with_limit(
+                            "Transient read failure (%s), reconnecting in %ss "
+                            "(attempt %s): %s",
+                            type(osex).__name__,
+                            round(delay),
+                            self._transient_error_count,
+                            osex,
+                        )
+                        await asyncio.sleep(delay)
+                        await self._reset()
                     else:
                         _logger.exception("Unexpected exception during read\n")
                         if self._try_increment_error_count(ErrorType.CONNECTION):

tests/test_fcmpushclient.py

@@ -9,7 +9,7 @@
 from http_ece import encrypt
 
 from firebase_messaging import FcmPushClient, FcmRegisterConfig
-from firebase_messaging.fcmpushclient import FcmPushClientRunState
+from firebase_messaging.fcmpushclient import ErrorType, FcmPushClientRunState
 from firebase_messaging.proto.mcs_pb2 import (
     Close,
     DataMessageStanza,
@@ -160,6 +160,88 @@ async def test_terminate(
             assert term_spy.call_count == 1
 
 
+@pytest.mark.parametrize(
+    "error",
+    [TimeoutError("SSL shutdown timed out"), asyncio.IncompleteReadError(b"", 1)],
+)
+async def test_transient_read_failure_reconnects_without_connection_counter(
+    logged_in_push_client, fake_mcs_endpoint, mocker, error
+):
+    # Transient read failures (e.g. SSL shutdown timeout) should reconnect
+    # without advancing the CONNECTION abort counter and without terminating.
+    pr = await logged_in_push_client(
+        None, None, abort_on_sequential_error_count=3, reset_interval=0.05
+    )
+    reset_spy = mocker.spy(pr, "_reset")
+    term_spy = mocker.spy(pr, "_terminate")
+    inc_spy = mocker.spy(pr, "_try_increment_error_count")
+
+    await fake_mcs_endpoint.put_error(error)
+    await asyncio.sleep(0.3)
+
+    assert reset_spy.call_count >= 1
+    assert term_spy.call_count == 0
+    conn_increments = [
+        c
+        for c in inc_spy.call_args_list
+        if c.args and c.args[0] is ErrorType.CONNECTION
+    ]
+    assert not conn_increments
+
+    msg = await fake_mcs_endpoint.get_message()
+    assert isinstance(msg, LoginRequest)
+
+
+async def test_reconnect_self_heals_when_connect_fails(
+    logged_in_push_client, fake_mcs_endpoint, mocker, caplog
+):
+    # When every reconnect attempt fails (sustained outage) the client must keep
+    # retrying with backoff and never give up / terminate (default behaviour).
+    pr = await logged_in_push_client(
+        None,
+        None,
+        reset_interval=0.02,
+        connection_retry_count=2,
+        start_seconds_before_retry_connect=0,
+        supress_disconnect=True,
+    )
+    term_spy = mocker.spy(pr, "_terminate")
+    connect_mock = mocker.patch.object(pr, "_connect", return_value=False)
+
+    await fake_mcs_endpoint.put_error(TimeoutError("SSL shutdown timed out"))
+    await asyncio.sleep(0.5)
+
+    assert term_spy.call_count == 0
+    assert connect_mock.call_count >= 3
+    assert any("will not give up" in record.message for record in caplog.records)
+
+    pr.do_listen = False  # let the self-heal loop exit for a clean teardown
+    await asyncio.sleep(0.1)
+
+
+async def test_abort_on_connection_failure_terminates(
+    logged_in_push_client, fake_mcs_endpoint, mocker
+):
+    # Opting in to abort_on_connection_failure restores the legacy behaviour:
+    # shut down once a reconnect fails.
+    pr = await logged_in_push_client(
+        None,
+        None,
+        reset_interval=0.02,
+        connection_retry_count=2,
+        start_seconds_before_retry_connect=0,
+        abort_on_connection_failure=True,
+        supress_disconnect=True,
+    )
+    term_spy = mocker.spy(pr, "_terminate")
+    mocker.patch.object(pr, "_connect", return_value=False)
+
+    await fake_mcs_endpoint.put_error(TimeoutError("SSL shutdown timed out"))
+    await asyncio.sleep(0.3)
+
+    assert term_spy.call_count == 1
+
+
 async def test_heartbeat_receive(logged_in_push_client, fake_mcs_endpoint, caplog):
     await logged_in_push_client(None, None)