Support customizing reconnection delays.

Fix #1395 (again).

Author: aaugustin

docs/project/changelog.rst

@@ -59,10 +59,17 @@ New features
 * Added the ``--insecure`` option to the ``websockets`` CLI to disable TLS
   certificate validation.
 
+Improvements
+............
+
+* Added the ``reconnect_delays`` argument for customizing the delays between
+  reconnection attempts in :func:`~asyncio.client.connect`, beyond existing
+  ``WEBSOCKETS_BACKOFF_*`` environment variables.
+
 Bug fixes
 .........
 
-* * Restored compatibility of the ``websockets`` CLI with Windows.
+* Restored compatibility of the ``websockets`` CLI with Windows.
 
 .. _16.1:
 

src/websockets/asyncio/client.py

@@ -234,6 +234,8 @@ class connect:
             :obj:`None` disables timeouts.
         close_timeout: Timeout for closing the connection in seconds.
             :obj:`None` disables the timeout.
+        reconnect_delays: Delays in seconds between reconnection attempts.
+            Default is exponential backoff with 5s jitter, capped at 60s.
         max_size: Maximum size of incoming messages in bytes.
             :obj:`None` disables the limit. You may pass a ``(max_message_size,
             max_fragment_size)`` tuple to set different limits for messages and
@@ -314,13 +316,14 @@ def __init__(
         ping_interval: float | None = 20,
         ping_timeout: float | None = 20,
         close_timeout: float | None = 10,
+        reconnect_delays: Callable[[], Generator[float]] = backoff,
         # Limits
         max_size: int | None | tuple[int | None, int | None] = 2**20,
         max_queue: int | None | tuple[int | None, int | None] = 16,
         write_limit: int | tuple[int, int | None] = 2**15,
         # Logging
         logger: LoggerLike | None = None,
-        # Escape hatch for advanced customization
+        # Escape hatches for advanced customization
         create_connection: type[ClientConnection] | None = None,
         # Other keyword arguments are passed to loop.create_connection
         **kwargs: Any,
@@ -368,6 +371,7 @@ def protocol_factory(uri: WebSocketURI) -> ClientConnection:
         self.user_agent_header = user_agent_header
         self.process_exception = process_exception
         self.open_timeout = open_timeout
+        self.reconnect_delays = reconnect_delays
         self.logger = logger
         self.connection_kwargs = kwargs
 
@@ -630,7 +634,7 @@ async def __aiter__(self) -> AsyncIterator[ClientConnection]:
                 # The connection failed with a retryable error.
                 # Start or continue backoff and reconnect.
                 if delays is None:
-                    delays = backoff()
+                    delays = self.reconnect_delays()
                 delay = next(delays)
                 self.logger.info(
                     "connect failed; reconnecting in %.1f seconds: %s",

tests/asyncio/test_client.py

@@ -29,19 +29,14 @@
 from .server import args, get_host_port, get_uri, handler
 
 
-@contextlib.asynccontextmanager
-async def short_backoff_delay():
+def short_backoff():
     defaults = backoff.__defaults__
-    backoff.__defaults__ = (
+    yield from backoff(
         defaults[0] * MS,
         defaults[1] * MS,
         defaults[2] * MS,
         defaults[3],
     )
-    try:
-        yield
-    finally:
-        backoff.__defaults__ = defaults
 
 
 @contextlib.asynccontextmanager
@@ -157,7 +152,6 @@ def create_connection(*args, **kwargs):
             ) as client:
                 self.assertTrue(client.create_connection_ran)
 
-    @short_backoff_delay()
     async def test_reconnect(self):
         """Client reconnects to server."""
         iterations = 0
@@ -179,7 +173,11 @@ async def process_request(connection, request):
 
         async with serve(*args, process_request=process_request) as server:
             with self.assertRaises(InvalidStatus) as raised:
-                async for client in connect(get_uri(server), open_timeout=3 * MS):
+                async for client in connect(
+                    get_uri(server),
+                    open_timeout=3 * MS,
+                    reconnect_delays=short_backoff,
+                ):
                     self.assertEqual(client.protocol.state.name, "OPEN")
                     successful += 1
 
@@ -190,7 +188,6 @@ async def process_request(connection, request):
         self.assertEqual(iterations, 6)
         self.assertEqual(successful, 2)
 
-    @short_backoff_delay()
     async def test_reconnect_with_custom_process_exception(self):
         """Client runs process_exception to tell if errors are retryable or fatal."""
         iteration = 0
@@ -213,7 +210,9 @@ def process_exception(exc):
         async with serve(*args, process_request=process_request) as server:
             with self.assertRaises(Exception) as raised:
                 async for _ in connect(
-                    get_uri(server), process_exception=process_exception
+                    get_uri(server),
+                    process_exception=process_exception,
+                    reconnect_delays=short_backoff,
                 ):
                     self.fail("did not raise")
 
@@ -223,7 +222,6 @@ def process_exception(exc):
             "🫖 💔 ☕️",
         )
 
-    @short_backoff_delay()
     async def test_reconnect_with_custom_process_exception_raising_exception(self):
         """Client supports raising an exception in process_exception."""
 
@@ -238,7 +236,9 @@ def process_exception(exc):
         async with serve(*args, process_request=process_request) as server:
             with self.assertRaises(Exception) as raised:
                 async for _ in connect(
-                    get_uri(server), process_exception=process_exception
+                    get_uri(server),
+                    process_exception=process_exception,
+                    reconnect_delays=short_backoff,
                 ):
                     self.fail("did not raise")