Add support for reconnecting automatically.

Fix #414.

Author: aaugustin

docs/howto/logging.rst

@@ -170,8 +170,9 @@ Disable logging
 
 If your application doesn't configure :mod:`logging`, Python outputs messages
 of severity :data:`~logging.WARNING` and higher to :data:`~sys.stderr`. As a
-consequence, you will see a message and a stack trace if a connection handler
-coroutine crashes or if you hit a bug in websockets.
+consequence, you will see a message and a stack trace if a server crashes in
+a connection handler coroutine, if a client crashes and relies on automatic
+reconnection, or if you hit a bug in websockets.
 
 If you want to disable this behavior for websockets, you can add
 a :class:`~logging.NullHandler`::
@@ -203,8 +204,10 @@ Here's what websockets logs at each level.
 ......................
 
 * Exceptions raised by connection handler coroutines in servers
+* Exceptions swallowed by automatic reconnection in clients
 * Exceptions resulting from bugs in websockets
 
+
 :attr:`~logging.INFO`
 .....................
 

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.

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

src/websockets/legacy/client.py

@@ -9,10 +9,22 @@
 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 (
+    ConnectionClosed,
     InvalidHandshake,
     InvalidHeader,
     InvalidMessage,
@@ -413,12 +425,22 @@ 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(...):
+            ...
+
+    As above, connections are closed automatically. Connection attempts are
+    delayed with exponential backoff, starting at two 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
@@ -577,6 +599,10 @@ def __init__(
             )
 
         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,45 @@ def handle_redirect(self, uri: str) -> None:
         # Set the new WebSocket URI. This suffices for same-origin redirects.
         self._wsuri = new_wsuri
 
-    # async with connect(...)
+    BACKOFF_MIN = 2
+    BACKOFF_MAX = 60
+    BACKOFF_FACTOR = 1.5
+
+    # async for ... in connect(...):
+
+    async def __aiter__(self) -> AsyncIterator[WebSocketClientProtocol]:
+        backoff_delay = self.BACKOFF_MIN
+        while True:
+            try:
+                protocol = await self
+            except Exception:
+                # Connection timed out - increase backoff delay
+                backoff_delay = min(int(1.5 * backoff_delay), self.BACKOFF_MAX)
+                self.logger.error(
+                    "! connect failed; retrying in %d seconds", backoff_delay
+                )
+                await asyncio.sleep(backoff_delay)
+                continue
+            else:
+                # Connection succeeded - reset backoff delay
+                backoff_delay = self.BACKOFF_MIN
+
+            try:
+                yield protocol
+            except GeneratorExit:
+                raise
+            except ConnectionClosed:
+                self.logger.debug("! connection closed; reconnecting", exc_info=True)
+            # Remove this branch when dropping support for Python < 3.8
+            # because CancelledError no longer inherits Exception.
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                self.logger.error("! an error occurred; reconnecting", exc_info=True)
+            finally:
+                await protocol.close()
+
+    # async with connect(...) as ...:
 
     async def __aenter__(self) -> WebSocketClientProtocol:
         return await self
@@ -628,7 +692,7 @@ async def __aexit__(
     ) -> 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 +729,7 @@ async def __await_impl__(self) -> WebSocketClientProtocol:
         else:
             raise SecurityError("too many redirects")
 
-    # yield from connect(...)
+    # ... = yield from connect(...)
 
     __iter__ = __await__
 

src/websockets/legacy/server.py

@@ -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