Optional connection retries (Fixes miguelgrinberg#1306)

miguelgrinberg · miguelgrinberg · commit b4f36148d816 · 2024-02-05T12:55:56.000Z
diff --git a/src/socketio/async_client.py b/src/socketio/async_client.py
@@ -69,7 +69,7 @@ def is_asyncio_based(self):
 
     async def connect(self, url, headers={}, auth=None, transports=None,
                       namespaces=None, socketio_path='socket.io', wait=True,
-                      wait_timeout=1):
+                      wait_timeout=1, retry=False):
         """Connect to a Socket.IO server.
 
         :param url: The URL of the Socket.IO server. It can include custom
@@ -105,6 +105,8 @@ async def connect(self, url, headers={}, auth=None, transports=None,
                              connection. The default is 1 second. This
                              argument is only considered when ``wait`` is set
                              to ``True``.
+        :param retry: Apply the reconnection logic if the initial connection
+                      attempt fails. The default is ``False``.
 
         Note: this method is a coroutine.
 
@@ -147,6 +149,10 @@ async def connect(self, url, headers={}, auth=None, transports=None,
                 await self._trigger_event(
                     'connect_error', n,
                     exc.args[1] if len(exc.args) > 1 else exc.args[0])
+            if retry:  # pragma: no cover
+                await self._handle_reconnect()
+                if self.eio.state == 'connected':
+                    return
             raise exceptions.ConnectionError(exc.args[0]) from None
 
         if wait:
@@ -477,7 +483,8 @@ async def _handle_reconnect(self):
                                    auth=self.connection_auth,
                                    transports=self.connection_transports,
                                    namespaces=self.connection_namespaces,
-                                   socketio_path=self.socketio_path)
+                                   socketio_path=self.socketio_path,
+                                   retry=False)
             except (exceptions.ConnectionError, ValueError):
                 pass
             else:
diff --git a/src/socketio/client.py b/src/socketio/client.py
@@ -69,7 +69,7 @@ class Client(base_client.BaseClient):
     """
     def connect(self, url, headers={}, auth=None, transports=None,
                 namespaces=None, socketio_path='socket.io', wait=True,
-                wait_timeout=1):
+                wait_timeout=1, retry=False):
         """Connect to a Socket.IO server.
 
         :param url: The URL of the Socket.IO server. It can include custom
@@ -105,6 +105,8 @@ def connect(self, url, headers={}, auth=None, transports=None,
                              connection. The default is 1 second. This
                              argument is only considered when ``wait`` is set
                              to ``True``.
+        :param retry: Apply the reconnection logic if the initial connection
+                      attempt fails. The default is ``False``.
 
         Example usage::
 
@@ -145,6 +147,10 @@ def connect(self, url, headers={}, auth=None, transports=None,
                 self._trigger_event(
                     'connect_error', n,
                     exc.args[1] if len(exc.args) > 1 else exc.args[0])
+            if retry:  # pragma: no cover
+                self._handle_reconnect()
+                if self.eio.state == 'connected':
+                    return
             raise exceptions.ConnectionError(exc.args[0]) from None
 
         if wait:
@@ -441,7 +447,8 @@ def _handle_reconnect(self):
                              auth=self.connection_auth,
                              transports=self.connection_transports,
                              namespaces=self.connection_namespaces,
-                             socketio_path=self.socketio_path)
+                             socketio_path=self.socketio_path,
+                             retry=False)
             except (exceptions.ConnectionError, ValueError):
                 pass
             else:
