Communicate handshake errors back to callers as well

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -62,6 +62,15 @@ async def buffered_message_sender(
     commit: Callable[[TransportMessage], Awaitable[None]],
     get_state: Callable[[], SessionState],
 ) -> None:
+    """
+    buffered_message_sender runs in a task and consumes from a queue, emitting
+    messages over the websocket as quickly as it can.
+
+    One of the design goals is to keep the message queue as short as possible to permit
+    quickly cancelling streams or acking heartbeats, so to that end it is wise to
+    incorporate backpressure into the lifecycle of get_next_pending/commit.
+    """
+
     our_task = asyncio.current_task()
     while our_task and not our_task.cancelling() and not our_task.cancelled():
         while get_state() in ConnectingStates:

src/replit_river/v2/session.py

@@ -239,13 +239,13 @@ def get_next_sent_seq() -> int:
                 return self._send_buffer[0].seq
             return self.seq
 
-        def do_close() -> None:
+        def close_session(reason: Exception | None) -> None:
             # Avoid closing twice
             if self._terminating_task is None:
                 # We can't just call self.close() directly because
                 # we're inside a thread that will eventually be awaited
                 # during the cleanup procedure.
-                self._terminating_task = asyncio.create_task(self.close())
+                self._terminating_task = asyncio.create_task(self.close(reason))
 
         def transition_connecting() -> None:
             if self._state in TerminalStates:
@@ -296,7 +296,7 @@ def unbind_connecting_task() -> None:
                     close_ws_in_background=close_ws_in_background,
                     transition_connected=transition_connected,
                     unbind_connecting_task=unbind_connecting_task,
-                    do_close=do_close,
+                    close_session=close_session,
                 )
             )
 
@@ -433,6 +433,26 @@ async def close(self, reason: Exception | None = None) -> None:
     def _start_buffered_message_sender(
         self,
     ) -> None:
+        """
+        Building on buffered_message_sender's documentation, we implement backpressure
+        per-stream by way of self._streams'
+
+            error_channel: Channel[Exception | None]
+
+        This is accomplished via the following strategy:
+        - If buffered_message_sender encounters an error, we transition back to
+          connecting and attempt to handshake.
+
+          If the handshake fails, we close the session with an informative error that
+          gets emitted to all backpressured client methods.
+
+        -  Alternately, if buffered_message_sender successfully writes back to the
+
+        - Finally, if _recv_from_ws encounters an error (transport or deserialization),
+          we emit an informative error to close_session which gets emitted to all
+          backpressured client methods.
+        """
+
         async def commit(msg: TransportMessage) -> None:
             pending = self._send_buffer.popleft()
             if msg.seq != pending.seq:
@@ -935,7 +955,7 @@ async def _do_ensure_connected[HandshakeMetadata](
     close_ws_in_background: Callable[[ClientConnection], None],
     transition_connected: Callable[[ClientConnection], None],
     unbind_connecting_task: Callable[[], None],
-    do_close: Callable[[], None],
+    close_session: Callable[[Exception | None], None],
 ) -> None:
     logger.info("Attempting to establish new ws connection")
 
@@ -1040,15 +1060,16 @@ async def websocket_closed_callback() -> None:
 
             logger.debug("river client get handshake response : %r", handshake_response)
             if not handshake_response.status.ok:
-                if handshake_response.status.code == ERROR_CODE_SESSION_STATE_MISMATCH:
-                    do_close()
-
-                raise RiverException(
+                err = RiverException(
                     ERROR_HANDSHAKE,
                     f"Handshake failed with code {handshake_response.status.code}: {
                         handshake_response.status.reason
                     }",
                 )
+                if handshake_response.status.code == ERROR_CODE_SESSION_STATE_MISMATCH:
+                    close_session(err)
+
+                raise err
 
             # We did it! We're connected!
             last_error = None
@@ -1069,7 +1090,7 @@ async def websocket_closed_callback() -> None:
 
     if last_error is not None:
         logger.debug("Handshake attempts exhausted, terminating")
-        do_close()
+        close_session(last_error)
         raise RiverException(
             ERROR_HANDSHAKE,
             f"Failed to create ws after retrying {attempt_count} number of times",