Add max retries to websocket reconnect in Python SDK (#379)

This should prevent from sending a lot of request to our infrastructure when the sandbox is already closed
e2b-dev · May 31, 2024 · b4ca523 · b4ca523
2 parents 88b4d59 + f0cbe76
commit b4ca523
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 3 deletions.
diff --git a/.changeset/small-oranges-roll.md b/.changeset/small-oranges-roll.md
@@ -0,0 +1,5 @@
+---
+"@e2b/python-sdk": patch
+---
+
+Add max retries for websocket connection
diff --git a/packages/python-sdk/e2b/sandbox/websocket_client.py b/packages/python-sdk/e2b/sandbox/websocket_client.py
@@ -2,15 +2,18 @@
 
 import asyncio
 import logging
+import random
 
 from queue import Queue
 from threading import Event
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, List, Optional, AsyncIterator
 
-from websockets.legacy.client import WebSocketClientProtocol, connect
+from websockets.legacy.client import WebSocketClientProtocol, Connect
 from websockets.exceptions import ConnectionClosed
 from websockets.typing import Data
 
+from e2b.sandbox.exception import SandboxException
+
 logger = logging.getLogger(__name__)
 
 
@@ -73,7 +76,7 @@ async def _connect(self):
         ws_logger = logger.getChild("websockets.client")
         ws_logger.setLevel(logging.ERROR)
 
-        websocket_connector = connect(
+        websocket_connector = E2BConnect(
             self.url,
             max_size=None,
             max_queue=None,
@@ -118,3 +121,43 @@ async def close(self):
 
         if self._ws:
             await self._ws.close()
+
+
+class E2BConnect(Connect):
+    async def __aiter__(self) -> AsyncIterator[WebSocketClientProtocol]:
+        retries = 0
+        max_retries = 12
+        backoff_delay = 0.1
+        while True:
+            try:
+                async with self as protocol:
+                    yield protocol
+            except Exception:
+                retries += 1
+                if retries >= max_retries:
+                    raise SandboxException("Failed to connect to the server")
+                # Add a random initial delay between 0 and 5 seconds.
+                # See 7.2.3. Recovering from Abnormal Closure in RFC 6544.
+                if backoff_delay == 0.1:
+                    initial_delay = random.random()
+                    self.logger.info(
+                        "! connect failed; reconnecting in %.1f seconds",
+                        initial_delay,
+                        exc_info=True,
+                    )
+                    await asyncio.sleep(initial_delay)
+                else:
+                    self.logger.info(
+                        "! connect failed again; retrying in %d seconds",
+                        int(backoff_delay),
+                        exc_info=True,
+                    )
+                    await asyncio.sleep(int(backoff_delay))
+                # Increase delay with truncated exponential backoff.
+                if retries > 4:
+                    backoff_delay = backoff_delay * 1.2
+                backoff_delay = min(backoff_delay, 10)
+                continue
+            else:
+                # Connection succeeded - reset backoff delay
+                backoff_delay = 0.1