Inline ws creation

blast-hardcheese · blast-hardcheese · commit 6d9931690264 · 2025-03-24T14:42:16.000-07:00
diff --git a/src/replit_river/client_transport.py b/src/replit_river/client_transport.py
@@ -117,6 +117,8 @@ async def get_or_create_session(self) -> ClientSession:
                 return existing_session
             else:
                 logger.info("Closing stale session %s", existing_session.session_id)
+                await new_ws.close()  # NB(dstewart): This wasn't there in the
+                #                       v1 transport, were we just leaking WS?
                 await existing_session.close()
                 return await self._create_new_session()
 
diff --git a/src/replit_river/v2/client_transport.py b/src/replit_river/v2/client_transport.py
@@ -13,7 +13,6 @@
 from replit_river.error_schema import (
     ERROR_CODE_STREAM_CLOSED,
     ERROR_HANDSHAKE,
-    ERROR_SESSION,
     RiverException,
 )
 from replit_river.messages import (
@@ -96,31 +95,28 @@ async def get_or_create_session(self) -> Session:
         If we have a disconnected session, attempt to start a new WS and use it.
         """
         async with self._create_session_lock:
-            existing_session = (
-                self._session
-                if self._session and self._session.is_session_open()
-                else None
-            )
-            if existing_session is None:
-                return await self._create_new_session()
-            if existing_session.is_websocket_open():
-                return existing_session
-            new_ws, _, hs_response = await self._establish_new_connection(
-                existing_session
-            )
-            if hs_response.status.sessionId == existing_session.session_id:
-                logger.info(
-                    "Replacing ws connection in session id %s",
-                    existing_session.session_id,
+            existing_session = self._session
+            if not existing_session:
+                logger.info("Creating new session")
+                new_session = Session(
+                    transport_id=self._transport_id,
+                    to_id=self._server_id,
+                    session_id=self.generate_nanoid(),
+                    transport_options=self._transport_options,
+                    close_session_callback=self._delete_session,
+                    retry_connection_callback=self._retry_connection,
                 )
-                await existing_session.replace_with_new_websocket(new_ws)
-                return existing_session
-            else:
-                logger.info("Closing stale session %s", existing_session.session_id)
-                await new_ws.close()  # NB(dstewart): This wasn't there in the
-                #                       v1 transport, were we just leaking WS?
-                await existing_session.close()
-                return await self._create_new_session()
+
+                self._session = new_session
+                existing_session = new_session
+                await existing_session.start_serve_responses()
+
+            await existing_session.ensure_connected(
+                client_id=self._client_id,
+                rate_limiter=self._rate_limiter,
+                uri_and_metadata_factory=self._uri_and_metadata_factory,
+            )
+            return existing_session
 
     async def _establish_new_connection(
         self,
@@ -191,31 +187,6 @@ async def _establish_new_connection(
             f"Failed to create ws after retrying {max_retry} number of times",
         ) from last_error
 
-    async def _create_new_session(
-        self,
-    ) -> Session:
-        logger.info("Creating new session")
-        new_ws, hs_request, hs_response = await self._establish_new_connection()
-        if not hs_response.status.ok:
-            message = hs_response.status.reason
-            raise RiverException(
-                ERROR_SESSION,
-                f"Server did not return OK status on handshake response: {message}",
-            )
-        new_session = Session(
-            transport_id=self._transport_id,
-            to_id=self._server_id,
-            session_id=hs_request.sessionId,
-            websocket=new_ws,
-            transport_options=self._transport_options,
-            close_session_callback=self._delete_session,
-            retry_connection_callback=self._retry_connection,
-        )
-
-        self._session = new_session
-        await new_session.start_serve_responses()
-        return new_session
-
     async def _retry_connection(self) -> Session:
         if not self._transport_options.transparent_reconnect:
             await self._close_session()
diff --git a/src/replit_river/v2/session.py b/src/replit_river/v2/session.py
@@ -15,10 +15,12 @@
 )
 
 import nanoid  # type: ignore
+import websockets.asyncio.client
 from aiochannel import Channel
 from aiochannel.errors import ChannelClosed
 from opentelemetry.trace import Span, use_span
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from pydantic import ValidationError
 from websockets.asyncio.client import ClientConnection
 from websockets.exceptions import ConnectionClosed, ConnectionClosedOK
 from websockets.frames import CloseCode
@@ -33,6 +35,7 @@
 from replit_river.error_schema import (
     ERROR_CODE_CANCEL,
     ERROR_CODE_STREAM_CLOSED,
+    ERROR_HANDSHAKE,
     RiverError,
     RiverException,
     RiverServiceException,
@@ -41,11 +44,17 @@
 )
 from replit_river.messages import (
     FailedSendingMessageException,
+    WebsocketClosedException,
     parse_transport_msg,
+    send_transport_message,
 )
+from replit_river.rate_limiter import LeakyBucketRateLimit
 from replit_river.rpc import (
     ACK_BIT,
     STREAM_OPEN_BIT,
+    ControlMessageHandshakeRequest,
+    ControlMessageHandshakeResponse,
+    ExpectedSessionState,
     TransportMessage,
     TransportMessageTracingSetter,
 )
@@ -55,7 +64,15 @@
     OutOfOrderMessageException,
 )
 from replit_river.task_manager import BackgroundTaskManager
-from replit_river.transport_options import MAX_MESSAGE_BUFFER_SIZE, TransportOptions
+from replit_river.transport_options import (
+    MAX_MESSAGE_BUFFER_SIZE,
+    TransportOptions,
+    UriAndMetadata,
+)
+from replit_river.v2.client_transport import (
+    PROTOCOL_VERSION,
+    HandshakeBudgetExhaustedException,
+)
 
 STREAM_CANCEL_BIT_TYPE = Literal[0b00100]
 STREAM_CANCEL_BIT: STREAM_CANCEL_BIT_TYPE = 0b00100
@@ -107,7 +124,6 @@ def __init__(
         transport_id: str,
         to_id: str,
         session_id: str,
-        websocket: ClientConnection,
         transport_options: TransportOptions,
         close_session_callback: CloseSessionCallback,
         retry_connection_callback: RetryConnectionCallback | None = None,
@@ -123,8 +139,7 @@ def __init__(
         self._close_session_after_time_secs: float | None = None
 
         # ws state
-        self._ws_connected = True
-        self._ws_unwrapped = websocket
+        self._ws_connected = False
         self._heartbeat_misses = 0
         self._retry_connection_callback = retry_connection_callback
 
@@ -187,6 +202,164 @@ def get_next_pending() -> TransportMessage | None:
             )
         )
 
+    async def ensure_connected(
+        self,
+        client_id: str,
+        rate_limiter: LeakyBucketRateLimit,
+        uri_and_metadata_factory: Callable[[], Awaitable[UriAndMetadata]],
+    ) -> None:
+        """
+        Either return immediately or establish a websocket connection and return
+        once we can accept messages
+        """
+        if self._ws_unwrapped and self._ws_connected:
+            return
+        max_retry = self._transport_options.connection_retry_options.max_retry
+        logger.info("Attempting to establish new ws connection")
+
+        last_error: Exception | None = None
+        for i in range(max_retry):
+            if i > 0:
+                logger.info(f"Retrying build handshake number {i} times")
+            if not rate_limiter.has_budget(client_id):
+                logger.debug("No retry budget for %s.", client_id)
+                raise HandshakeBudgetExhaustedException(
+                    ERROR_HANDSHAKE,
+                    "No retry budget",
+                    client_id=client_id,
+                ) from last_error
+
+            rate_limiter.consume_budget(client_id)
+
+            try:
+                uri_and_metadata = await uri_and_metadata_factory()
+                ws = await websockets.asyncio.client.connect(uri_and_metadata["uri"])
+
+                try:
+                    try:
+                        expectedSessionState = ExpectedSessionState(
+                            nextExpectedSeq=0,
+                            nextSentSeq=0,
+                        )
+                        handshake_request = ControlMessageHandshakeRequest[Any](
+                            type="HANDSHAKE_REQ",
+                            protocolVersion=PROTOCOL_VERSION,
+                            sessionId=self.session_id,
+                            metadata=uri_and_metadata["metadata"],
+                            expectedSessionState=expectedSessionState,
+                        )
+                        stream_id = nanoid.generate()
+
+                        async def websocket_closed_callback() -> None:
+                            logger.error("websocket closed before handshake response")
+
+                        try:
+                            payload = handshake_request.model_dump()
+                            await send_transport_message(
+                                TransportMessage(
+                                    from_=self._transport_id,
+                                    to=self._to_id,
+                                    streamId=stream_id,
+                                    controlFlags=0,
+                                    id=nanoid.generate(),
+                                    seq=0,
+                                    ack=0,
+                                    payload=payload,
+                                ),
+                                ws=ws,
+                                websocket_closed_callback=websocket_closed_callback,
+                            )
+                        except (
+                            WebsocketClosedException,
+                            FailedSendingMessageException,
+                        ) as e:  # noqa: E501
+                            raise RiverException(
+                                ERROR_HANDSHAKE,
+                                "Handshake failed, conn closed while sending response",  # noqa: E501
+                            ) from e
+                    except FailedSendingMessageException as e:
+                        raise RiverException(
+                            ERROR_CODE_STREAM_CLOSED,
+                            "Stream closed before response, closing connection",
+                        ) from e
+
+                    startup_grace_deadline_ms = await self._get_current_time() + 60_000
+                    try:
+                        while True:
+                            if (
+                                await self._get_current_time()
+                                >= startup_grace_deadline_ms
+                            ):  # noqa: E501
+                                raise RiverException(
+                                    ERROR_HANDSHAKE,
+                                    "Handshake response timeout, closing connection",  # noqa: E501
+                                )
+                            try:
+                                data = await ws.recv()
+                            except ConnectionClosed as e:
+                                logger.debug(
+                                    "Connection closed during waiting for handshake response",  # noqa: E501
+                                    exc_info=True,
+                                )
+                                raise RiverException(
+                                    ERROR_HANDSHAKE,
+                                    "Handshake failed, conn closed while waiting for response",  # noqa: E501
+                                ) from e
+                            try:
+                                response_msg = parse_transport_msg(data)
+                                break
+                            except IgnoreMessageException:
+                                logger.debug(
+                                    "Ignoring transport message", exc_info=True
+                                )  # noqa: E501
+                                continue
+                            except InvalidMessageException as e:
+                                raise RiverException(
+                                    ERROR_HANDSHAKE,
+                                    "Got invalid transport message, closing connection",
+                                ) from e
+
+                        handshake_response = ControlMessageHandshakeResponse(
+                            **response_msg.payload
+                        )  # noqa: E501
+                        logger.debug("river client waiting for handshake response")
+                    except ValidationError as e:
+                        raise RiverException(
+                            ERROR_HANDSHAKE, "Failed to parse handshake response"
+                        ) from e
+                    except asyncio.TimeoutError as e:
+                        raise RiverException(
+                            ERROR_HANDSHAKE,
+                            "Handshake response timeout, closing connection",  # noqa: E501
+                        ) from e
+
+                    logger.debug(
+                        "river client get handshake response : %r", handshake_response
+                    )  # noqa: E501
+                    if not handshake_response.status.ok:
+                        raise RiverException(
+                            ERROR_HANDSHAKE,
+                            f"Handshake failed with code {handshake_response.status.code}: "  # noqa: E501
+                            + f"{handshake_response.status.reason}",
+                        )
+
+                    rate_limiter.start_restoring_budget(client_id)
+                except RiverException as e:
+                    await ws.close()
+                    raise e
+            except Exception as e:
+                last_error = e
+                backoff_time = rate_limiter.get_backoff_ms(client_id)
+                logger.exception(
+                    f"Error connecting, retrying with {backoff_time}ms backoff"
+                )
+                await asyncio.sleep(backoff_time / 1000)
+
+        raise RiverException(
+            ERROR_HANDSHAKE,
+            f"Failed to create ws after retrying {max_retry} number of times",
+        ) from last_error
+
     def _setup_heartbeats_task(
         self,
         do_close_websocket: Callable[[], Awaitable[None]],
