Repository: aaugustin/websockets (python)

Issue:
Common pattern for reconnections
In several projects of mines I do something like the following, in order to handle scenario that may happen with connection errors and reconnection attempts:

```
    async def listen_forever(self):
        while True:
        # outer loop restarted every time the connection fails
            try:
                async with websockets.connect(self.url) as ws:
                    while True:
                    # listener loop
                        try:
                            reply = await asyncio.wait_for(ws.recv(), timeout=***)
                        except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
                            try:
                                pong = await ws.ping()
                                await asyncio.wait_for(pong, timeout=self.ping_timeout)
                                logger.debug('Ping OK, keeping connection alive...')
                                continue
                            except:
                                await asyncio.sleep(self.sleep_time)
                                break  # inner loop
                        # do stuff with reply object
            except socket.gaierror:
                # log something
                continue
            except ConnectionRefusedError:
                # log something else
                continue
```

and I was wondering whether (1) this makes sense at all, and (2) there is any shortcut already provided in `websockets` to what seems quite a recurrent behavior (assuming that is correct!)

Pull request:
Add support for reconnecting automatically.
Fix #414.

Patch:
diff --git a/docs/howto/logging.rst b/docs/howto/logging.rst
index f69ee47..8248129 100644
--- a/docs/howto/logging.rst
+++ b/docs/howto/logging.rst
@@ -210,6 +210,7 @@ Here's what websockets logs at each level.
 
 * Server starting and stopping
 * Server establishing and closing connections
+* Client reconnecting automatically
 
 :attr:`~logging.DEBUG`
 ......................
diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst
index 6034665..a0cd8f0 100644
--- a/docs/project/changelog.rst
+++ b/docs/project/changelog.rst
@@ -51,6 +51,9 @@ They may change at any time.
 
 * Added compatibility with Python 3.10.
 
+* Added support for reconnecting automatically by using
+  :func:`~legacy.client.connect` as an asynchronous iterator.
+
 * Added ``open_timeout`` to :func:`~legacy.client.connect`.
 
 * Improved logging.
@@ -69,7 +72,11 @@ They may change at any time.
 
     **Version 9.1 fixes a security issue introduced in version 8.0.**
 
-    Version 8.0 was vulnerable to timing attacks on HTTP Basic Auth passwords.
+    Version 8.0 was vulnerable to timing attacks on HTTP Basic Auth passwords
+    (`CVE-2021-33880`_).
+
+    .. _CVE-2021-33880: https://nvd.nist.gov/vuln/detail/CVE-2021-33880
+
 
 9.0.2
 .....
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index b460ef0..8346ace 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -6,6 +6,7 @@ autoscaler
 awaitable
 aymeric
 backend
+backoff
 backpressure
 balancer
 balancers
@@ -52,6 +53,7 @@ pong
 pongs
 proxying
 pythonic
+reconnection
 redis
 retransmit
 runtime
diff --git a/src/websockets/legacy/client.py b/src/websockets/legacy/client.py
index df1a2f5..63fdbf9 100644
--- a/src/websockets/legacy/client.py
+++ b/src/websockets/legacy/client.py
@@ -9,7 +9,18 @@ import functools
 import logging
 import warnings
 from types import TracebackType
-from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Type, cast
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Generator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    cast,
+)
 
 from ..datastructures import Headers, HeadersLike
 from ..exceptions import (
@@ -33,7 +44,6 @@ from ..headers import (
 from ..http import USER_AGENT, build_host
 from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
 from ..uri import WebSocketURI, parse_uri
-from .compatibility import asyncio_get_running_loop
 from .handshake import build_request, check_response
 from .http import read_response
 from .protocol import WebSocketCommonProtocol
@@ -413,12 +423,23 @@ class Connect:
     Awaiting :func:`connect` yields a :class:`WebSocketClientProtocol` which
     can then be used to send and receive messages.
 
-    :func:`connect` can also be used as a asynchronous context manager::
+    :func:`connect` can be used as a asynchronous context manager::
 
         async with connect(...) as websocket:
             ...
 
-    In that case, the connection is closed when exiting the context.
+    The connection is closed automatically when exiting the context.
+
+    :func:`connect` can be used as an infinite asynchronous iterator to
+    reconnect automatically on errors::
+
+        async for websocket in connect(...):
+            ...
+
+    You must catch all exceptions, or else you will exit the loop prematurely.
+    As above, connections are closed automatically. Connection attempts are
+    delayed with exponential backoff, starting at three seconds and
+    increasing up to one minute.
 
     :func:`connect` is a wrapper around the event loop's
     :meth:`~asyncio.loop.create_connection` method. Unknown keyword arguments
@@ -515,10 +536,11 @@ class Connect:
         legacy_recv: bool = kwargs.pop("legacy_recv", False)
 
         # Backwards compatibility: the loop parameter used to be supported.
-        loop: Optional[asyncio.AbstractEventLoop] = kwargs.pop("loop", None)
-        if loop is None:
-            loop = asyncio_get_running_loop()
+        _loop: Optional[asyncio.AbstractEventLoop] = kwargs.pop("loop", None)
+        if _loop is None:
+            loop = asyncio.get_event_loop()
         else:
+            loop = _loop
             warnings.warn("remove loop argument", DeprecationWarning)
 
         wsuri = parse_uri(uri)
@@ -544,7 +566,7 @@ class Connect:
             max_queue=max_queue,
             read_limit=read_limit,
             write_limit=write_limit,
-            loop=loop,
+            loop=_loop,
             host=wsuri.host,
             port=wsuri.port,
             secure=wsuri.secure,
@@ -577,6 +599,10 @@ class Connect:
             )
 
         self.open_timeout = open_timeout
+        if logger is None:
+            logger = logging.getLogger("websockets.client")
+        self.logger = logger
+
         # This is a coroutine function.
         self._create_connection = create_connection
         self._wsuri = wsuri
@@ -615,7 +641,38 @@ class Connect:
         # Set the new WebSocket URI. This suffices for same-origin redirects.
         self._wsuri = new_wsuri
 
-    # async with connect(...)
+    # async for ... in connect(...):
+
+    BACKOFF_MIN = 2.0
+    BACKOFF_MAX = 60.0
+    BACKOFF_FACTOR = 1.5
+
+    async def __aiter__(self) -> AsyncIterator[WebSocketClientProtocol]:
+        backoff_delay = self.BACKOFF_MIN
+        while True:
+            try:
+                async with self as protocol:
+                    yield protocol
+            # Remove this branch when dropping support for Python < 3.8
+            # because CancelledError no longer inherits Exception.
+            except asyncio.CancelledError:  # pragma: no cover
+                raise
+            except Exception:
+                # Connection timed out - increase backoff delay
+                backoff_delay = backoff_delay * self.BACKOFF_FACTOR
+                backoff_delay = min(backoff_delay, self.BACKOFF_MAX)
+                self.logger.info(
+                    "! connect failed; retrying in %d seconds",
+                    int(backoff_delay),
+                    exc_info=True,
+                )
+                await asyncio.sleep(backoff_delay)
+                continue
+            else:
+                # Connection succeeded - reset backoff delay
+                backoff_delay = self.BACKOFF_MIN
+
+    # async with connect(...) as ...:
 
     async def __aenter__(self) -> WebSocketClientProtocol:
         return await self
@@ -628,7 +685,7 @@ class Connect:
     ) -> None:
         await self.protocol.close()
 
-    # await connect(...)
+    # ... = await connect(...)
 
     def __await__(self) -> Generator[Any, None, WebSocketClientProtocol]:
         # Create a suitable iterator by calling __await__ on a coroutine.
@@ -665,7 +722,7 @@ class Connect:
         else:
             raise SecurityError("too many redirects")
 
-    # yield from connect(...)
+    # ... = yield from connect(...)
 
     __iter__ = __await__
 
diff --git a/src/websockets/legacy/compatibility.py b/src/websockets/legacy/compatibility.py
index 86f6715..96df028 100644
--- a/src/websockets/legacy/compatibility.py
+++ b/src/websockets/legacy/compatibility.py
@@ -9,14 +9,3 @@ def loop_if_py_lt_38(loop: asyncio.AbstractEventLoop) -> Dict[str, Any]:
 
     """
     return {"loop": loop} if sys.version_info[:2] < (3, 8) else {}
-
-
-def asyncio_get_running_loop() -> asyncio.AbstractEventLoop:
-    """
-    Helper for the deprecation of get_event_loop in Python 3.10.
-
-    """
-    if sys.version_info[:2] < (3, 10):  # pragma: no cover
-        return asyncio.get_event_loop()
-    else:  # pragma: no cover
-        return asyncio.get_running_loop()
diff --git a/src/websockets/legacy/protocol.py b/src/websockets/legacy/protocol.py
index 353317a..ca60761 100644
--- a/src/websockets/legacy/protocol.py
+++ b/src/websockets/legacy/protocol.py
@@ -124,6 +124,12 @@ class WebSocketCommonProtocol(asyncio.Protocol):
         if close_timeout is None:
             close_timeout = timeout
 
+        # Backwards compatibility: the loop parameter used to be supported.
+        if loop is None:
+            loop = asyncio.get_event_loop()
+        else:
+            warnings.warn("remove loop argument", DeprecationWarning)
+
         self.ping_interval = ping_interval
         self.ping_timeout = ping_timeout
         self.close_timeout = close_timeout
@@ -145,8 +151,6 @@ class WebSocketCommonProtocol(asyncio.Protocol):
         # Track if DEBUG is enabled. Shortcut logging calls if it isn't.
         self.debug = logger.isEnabledFor(logging.DEBUG)
 
-        assert loop is not None
-        # Remove when dropping Python < 3.10 - use get_running_loop instead.
         self.loop = loop
 
         self._host = host
diff --git a/src/websockets/legacy/server.py b/src/websockets/legacy/server.py
index 1704ae0..9ace9a2 100644
--- a/src/websockets/legacy/server.py
+++ b/src/websockets/legacy/server.py
@@ -42,7 +42,7 @@ from ..extensions.permessage_deflate import enable_server_permessage_deflate
 from ..headers import build_extension, parse_extension, parse_subprotocol
 from ..http import USER_AGENT
 from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
-from .compatibility import asyncio_get_running_loop, loop_if_py_lt_38
+from .compatibility import loop_if_py_lt_38
 from .handshake import build_response, check_request
 from .http import read_request
 from .protocol import WebSocketCommonProtocol
@@ -904,7 +904,7 @@ class Serve:
     :exc:`~websockets.exceptions.ConnectionClosedOK` exception on their
     current or next interaction with the WebSocket connection.
 
-    :func:`serve` can also be used as an asynchronous context manager::
+    :func:`serve` can be used as an asynchronous context manager::
 
         stop = asyncio.Future()  # set this future to exit the server
 
@@ -1024,10 +1024,11 @@ class Serve:
         legacy_recv: bool = kwargs.pop("legacy_recv", False)
 
         # Backwards compatibility: the loop parameter used to be supported.
-        loop: Optional[asyncio.AbstractEventLoop] = kwargs.pop("loop", None)
-        if loop is None:
-            loop = asyncio_get_running_loop()
+        _loop: Optional[asyncio.AbstractEventLoop] = kwargs.pop("loop", None)
+        if _loop is None:
+            loop = asyncio.get_event_loop()
         else:
+            loop = _loop
             warnings.warn("remove loop argument", DeprecationWarning)
 
         ws_server = WebSocketServer(logger=logger, loop=loop)
@@ -1053,7 +1054,7 @@ class Serve:
             max_queue=max_queue,
             read_limit=read_limit,
             write_limit=write_limit,
-            loop=loop,
+            loop=_loop,
             legacy_recv=legacy_recv,
             origins=origins,
             extensions=extensions,


Test patch:
diff --git a/tests/legacy/test_client_server.py b/tests/legacy/test_client_server.py
index 60c0a14..2f754aa 100644
--- a/tests/legacy/test_client_server.py
+++ b/tests/legacy/test_client_server.py
@@ -213,17 +213,17 @@ class ClientServerTestsMixin:
         kwargs.setdefault("compression", None)
         # Disable pings by default in tests.
         kwargs.setdefault("ping_interval", None)
-        # Python 3.10 dislikes not having a running event loop
-        if sys.version_info[:2] >= (3, 10):  # pragma: no cover
-            kwargs.setdefault("loop", self.loop)
 
         with warnings.catch_warnings(record=True) as recorded_warnings:
             start_server = serve(handler, "localhost", 0, **kwargs)
             self.server = self.loop.run_until_complete(start_server)
 
         expected_warnings = [] if deprecation_warnings is None else deprecation_warnings
-        if sys.version_info[:2] >= (3, 10):  # pragma: no cover
-            expected_warnings += ["remove loop argument"]
+        if (
+            sys.version_info[:2] >= (3, 10)
+            and "remove loop argument" not in expected_warnings
+        ):  # pragma: no cover
+            expected_warnings += ["There is no current event loop"]
         self.assertDeprecationWarnings(recorded_warnings, expected_warnings)
 
     def start_redirecting_server(
@@ -234,10 +234,6 @@ class ClientServerTestsMixin:
         deprecation_warnings=None,
         **kwargs,
     ):
-        # Python 3.10 dislikes not having a running event loop
-        if sys.version_info[:2] >= (3, 10):  # pragma: no cover
-            kwargs.setdefault("loop", self.loop)
-
         async def process_request(path, headers):
             server_uri = get_server_uri(self.server, self.secure, path)
             if force_insecure:
@@ -259,8 +255,11 @@ class ClientServerTestsMixin:
             self.redirecting_server = self.loop.run_until_complete(start_server)
 
         expected_warnings = [] if deprecation_warnings is None else deprecation_warnings
-        if sys.version_info[:2] >= (3, 10):  # pragma: no cover
-            expected_warnings += ["remove loop argument"]
+        if (
+            sys.version_info[:2] >= (3, 10)
+            and "remove loop argument" not in expected_warnings
+        ):  # pragma: no cover
+            expected_warnings += ["There is no current event loop"]
         self.assertDeprecationWarnings(recorded_warnings, expected_warnings)
 
     def start_client(
@@ -270,9 +269,6 @@ class ClientServerTestsMixin:
         kwargs.setdefault("compression", None)
         # Disable pings by default in tests.
         kwargs.setdefault("ping_interval", None)
-        # Python 3.10 dislikes not having a running event loop
-        if sys.version_info[:2] >= (3, 10):  # pragma: no cover
-            kwargs.setdefault("loop", self.loop)
 
         secure = kwargs.get("ssl") is not None
         try:
@@ -286,8 +282,11 @@ class ClientServerTestsMixin:
             self.client = self.loop.run_until_complete(start_client)
 
         expected_warnings = [] if deprecation_warnings is None else deprecation_warnings
-        if sys.version_info[:2] >= (3, 10):  # pragma: no cover
-            expected_warnings += ["remove loop argument"]
+        if (
+            sys.version_info[:2] >= (3, 10)
+            and "remove loop argument" not in expected_warnings
+        ):  # pragma: no cover
+            expected_warnings += ["There is no current event loop"]
         self.assertDeprecationWarnings(recorded_warnings, expected_warnings)
 
     def stop_client(self):
@@ -409,8 +408,6 @@ class CommonClientServerTests:
         with temp_test_redirecting_server(
             self,
             http.HTTPStatus.FOUND,
-            loop=self.loop,
-            deprecation_warnings=["remove loop argument"],
         ):
             self.server = self.redirecting_server
             with self.assertRaises(InvalidHandshake):
@@ -430,7 +427,7 @@ class CommonClientServerTests:
                 with temp_test_client(self):
                     self.fail("Did not raise")  # pragma: no cover
 
-    def test_explicit_event_loop(self):
+    def test_loop_backwards_compatibility(self):
         with self.temp_server(
             loop=self.loop, deprecation_warnings=["remove loop argument"]
         ):
@@ -1474,6 +1471,85 @@ class AsyncIteratorTests(ClientServerTestsMixin, AsyncioTestCase):
         self.assertEqual(messages, self.MESSAGES)
 
 
+class ReconnectionTests(ClientServerTestsMixin, AsyncioTestCase):
+    async def echo_handler(ws, path):
+        async for msg in ws:
+            await ws.send(msg)
+
+    service_available = True
+
+    async def maybe_service_unavailable(path, headers):
+        if not ReconnectionTests.service_available:
+            return http.HTTPStatus.SERVICE_UNAVAILABLE, [], b""
+
+    async def disable_server(self, duration):
+        ReconnectionTests.service_available = False
+        await asyncio.sleep(duration)
+        ReconnectionTests.service_available = True
+
+    @with_server(handler=echo_handler, process_request=maybe_service_unavailable)
+    def test_reconnect(self):
+        # Big, ugly integration test :-(
+
+        async def run_client():
+            iteration = 0
+            connect_inst = connect(get_server_uri(self.server))
+            connect_inst.BACKOFF_MIN = 10 * MS
+            connect_inst.BACKOFF_MAX = 200 * MS
+            async for ws in connect_inst:
+                await ws.send("spam")
+                msg = await ws.recv()
+                self.assertEqual(msg, "spam")
+
+                iteration += 1
+                if iteration == 1:
+                    # Exit block normally.
+                    pass
+                elif iteration == 2:
+                    # Disable server for a little bit
+                    asyncio.create_task(self.disable_server(70 * MS))
+                    await asyncio.sleep(0)
+                elif iteration == 3:
+                    # Exit block after catching connection error.
+                    server_ws = next(iter(self.server.websockets))
+                    await server_ws.close()
+                    with self.assertRaises(ConnectionClosed):
+                        await ws.recv()
+                else:
+                    # Exit block with an exception.
+                    raise Exception("BOOM!")
+
+        with self.assertLogs("websockets", logging.INFO) as logs:
+            with self.assertRaisesRegex(Exception, "BOOM!"):
+                self.loop.run_until_complete(run_client())
+
+        self.assertEqual(
+            [record.getMessage() for record in logs.records],
+            [
+                # Iteration 1
+                "connection open",
+                "connection closed",
+                # Iteration 2
+                "connection open",
+                "connection closed",
+                # Iteration 3
+                "connection failed (503 Service Unavailable)",
+                "connection closed",
+                "! connect failed; retrying in 0 seconds",
+                "connection failed (503 Service Unavailable)",
+                "connection closed",
+                "! connect failed; retrying in 0 seconds",
+                "connection failed (503 Service Unavailable)",
+                "connection closed",
+                "! connect failed; retrying in 0 seconds",
+                "connection open",
+                "connection closed",
+                # Iteration 4
+                "connection open",
+            ],
+        )
+
+
 class LoggerTests(ClientServerTestsMixin, AsyncioTestCase):
     def test_logger_client(self):
         with self.assertLogs("test.server", logging.DEBUG) as server_logs:
diff --git a/tests/legacy/test_protocol.py b/tests/legacy/test_protocol.py
index c1948c8..ccbbffe 100644
--- a/tests/legacy/test_protocol.py
+++ b/tests/legacy/test_protocol.py
@@ -1,5 +1,6 @@
 import asyncio
 import contextlib
+import sys
 import unittest
 import unittest.mock
 import warnings
@@ -86,8 +87,9 @@ class CommonTests:
 
     def setUp(self):
         super().setUp()
-        # Disable pings to make it easier to test what frames are sent exactly.
-        self.protocol = WebSocketCommonProtocol(ping_interval=None, loop=self.loop)
+        with warnings.catch_warnings(record=True):
+            # Disable pings to make it easier to test what frames are sent exactly.
+            self.protocol = WebSocketCommonProtocol(ping_interval=None)
         self.transport = TransportMock()
         self.transport.setup_mock(self.loop, self.protocol)
 
@@ -309,14 +311,29 @@ class CommonTests:
     # Test constructor.
 
     def test_timeout_backwards_compatibility(self):
-        with warnings.catch_warnings(record=True) as recorded_warnings:
-            protocol = WebSocketCommonProtocol(timeout=5, loop=self.loop)
+        with warnings.catch_warnings(record=True) as recorded:
+            protocol = WebSocketCommonProtocol(timeout=5)
 
         self.assertEqual(protocol.close_timeout, 5)
 
-        self.assertDeprecationWarnings(
-            recorded_warnings, ["rename timeout to close_timeout"]
-        )
+        expected = ["rename timeout to close_timeout"]
+        if sys.version_info[:2] >= (3, 10):  # pragma: no cover
+            expected += ["There is no current event loop"]
+
+        self.assertDeprecationWarnings(recorded, expected)
+
+    def test_loop_backwards_compatibility(self):
+        loop = asyncio.new_event_loop()
+        self.addCleanup(loop.close)
+
+        with warnings.catch_warnings(record=True) as recorded:
+            protocol = WebSocketCommonProtocol(loop=loop)
+
+        self.assertEqual(protocol.loop, loop)
+
+        expected = ["remove loop argument"]
+
+        self.assertDeprecationWarnings(recorded, expected)
 
     # Test public attributes.
 
@@ -1116,11 +1133,11 @@ class CommonTests:
         self.transport.close()
         self.loop.run_until_complete(self.protocol.close())
         # copied from setUp, but enables keepalive pings
-        self.protocol = WebSocketCommonProtocol(
-            ping_interval=ping_interval,
-            ping_timeout=ping_timeout,
-            loop=self.loop,
-        )
+        with warnings.catch_warnings(record=True):
+            self.protocol = WebSocketCommonProtocol(
+                ping_interval=ping_interval,
+                ping_timeout=ping_timeout,
+            )
         self.transport = TransportMock()
         self.transport.setup_mock(self.loop, self.protocol)
         self.protocol.is_client = initial_protocol.is_client