Avoid exceptions for flow control

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -2,7 +2,7 @@
 import logging
 from collections.abc import AsyncIterable
 from datetime import timedelta
-from typing import Any, AsyncGenerator, Callable, Coroutine
+from typing import Any, AsyncGenerator, Callable, Coroutine, assert_never
 
 import nanoid
 import websockets
@@ -24,7 +24,7 @@
     parse_transport_msg,
 )
 from replit_river.seq_manager import (
-    IgnoreMessageException,
+    IgnoreMessage,
     InvalidMessageException,
     OutOfOrderMessageException,
 )
@@ -125,11 +125,21 @@ async def _handle_messages_from_ws(self) -> None:
                         # We should not process messages if the websocket is closed.
                         break
                     msg = parse_transport_msg(message)
+                    if isinstance(msg, str):
+                        logger.debug("Ignoring transport message", exc_info=True)
+                        continue
 
                     logger.debug(f"{self._transport_id} got a message %r", msg)
 
                     # Update bookkeeping
-                    await self._seq_manager.check_seq_and_update(msg)
+                    match await self._seq_manager.check_seq_and_update(msg):
+                        case IgnoreMessage():
+                            continue
+                        case None:
+                            pass
+                        case other:
+                            assert_never(other)
+
                     await self._buffer.remove_old_messages(
                         self._seq_manager.receiver_ack,
                     )
@@ -142,9 +152,7 @@ async def _handle_messages_from_ws(self) -> None:
                     if msg.controlFlags & STREAM_OPEN_BIT == 0:
                         if not stream:
                             logger.warning("no stream for %s", msg.streamId)
-                            raise IgnoreMessageException(
-                                "no stream for message, ignoring"
-                            )
+                            continue
 
                         if (
                             msg.controlFlags & STREAM_CLOSED_BIT != 0
@@ -171,9 +179,6 @@ async def _handle_messages_from_ws(self) -> None:
                             stream.close()
                         async with self._stream_lock:
                             del self._streams[msg.streamId]
-                except IgnoreMessageException:
-                    logger.debug("Ignoring transport message", exc_info=True)
-                    continue
                 except OutOfOrderMessageException:
                     logger.exception("Out of order message, closing connection")
                     await ws_wrapper.close()

src/replit_river/client_transport.py

@@ -33,7 +33,6 @@
     TransportMessage,
 )
 from replit_river.seq_manager import (
-    IgnoreMessageException,
     InvalidMessageException,
 )
 from replit_river.session import Session
@@ -296,10 +295,11 @@ async def _get_handshake_response_msg(
                     "Handshake failed, conn closed while waiting for response",
                 ) from e
             try:
-                return parse_transport_msg(data)
-            except IgnoreMessageException:
-                logger.debug("Ignoring transport message", exc_info=True)
-                continue
+                msg = parse_transport_msg(data)
+                if isinstance(msg, str):
+                    logger.debug("Ignoring transport message", exc_info=True)
+                    continue
+                return msg
             except InvalidMessageException as e:
                 raise RiverException(
                     ERROR_HANDSHAKE,

src/replit_river/messages.py

@@ -14,7 +14,6 @@
     TransportMessage,
 )
 from replit_river.seq_manager import (
-    IgnoreMessageException,
     InvalidMessageException,
 )
 
@@ -59,11 +58,9 @@ def formatted_bytes(message: bytes) -> str:
     return " ".join(f"{b:02x}" for b in message)
 
 
-def parse_transport_msg(message: str | bytes) -> TransportMessage:
+def parse_transport_msg(message: str | bytes) -> TransportMessage | str:
     if isinstance(message, str):
-        raise IgnoreMessageException(
-            f"ignored a message beacuse it was a text frame: {message}"
-        )
+        return message
     try:
         # :param int timestamp:
         #     Control how timestamp type is unpacked:

src/replit_river/seq_manager.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from dataclasses import dataclass
 
 from replit_river.rpc import TransportMessage
 
@@ -34,6 +35,11 @@ class SessionStateMismatchException(Exception):
     pass
 
 
+@dataclass
+class IgnoreMessage:
+    pass
+
+
 class SeqManager:
     """Manages the sequence number and ack number for a connection."""
 
@@ -68,14 +74,11 @@ async def get_ack(self) -> int:
         async with self._ack_lock:
             return self.ack
 
-    async def check_seq_and_update(self, msg: TransportMessage) -> None:
+    async def check_seq_and_update(self, msg: TransportMessage) -> IgnoreMessage | None:
         async with self._ack_lock:
             if msg.seq != self.ack:
                 if msg.seq < self.ack:
-                    raise IgnoreMessageException(
-                        f"{msg.from_} received duplicate msg, got {msg.seq}"
-                        f" expected {self.ack}"
-                    )
+                    return IgnoreMessage()
                 else:
                     logger.warn(
                         f"Out of order message received got {msg.seq} expected "

src/replit_river/server_session.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Any, Callable, Coroutine
+from typing import Any, Callable, Coroutine, assert_never
 
 import websockets
 from aiochannel import Channel, ChannelClosed
@@ -12,7 +12,7 @@
     parse_transport_msg,
 )
 from replit_river.seq_manager import (
-    IgnoreMessageException,
+    IgnoreMessage,
     InvalidMessageException,
     OutOfOrderMessageException,
 )
@@ -122,11 +122,20 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                         # We should not process messages if the websocket is closed.
                         break
                     msg = parse_transport_msg(message)
+                    if isinstance(msg, str):
+                        logger.debug("Ignoring transport message", exc_info=True)
+                        continue
 
                     logger.debug(f"{self._transport_id} got a message %r", msg)
 
                     # Update bookkeeping
-                    await self._seq_manager.check_seq_and_update(msg)
+                    match self._seq_manager.check_seq_and_update(msg):
+                        case IgnoreMessage():
+                            continue
+                        case None:
+                            pass
+                        case other:
+                            assert_never(other)
                     await self._buffer.remove_old_messages(
                         self._seq_manager.receiver_ack,
                     )
@@ -135,13 +144,11 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                     if msg.controlFlags & ACK_BIT != 0:
                         continue
                     async with self._stream_lock:
-                        stream = self._streams.get(msg.streamId, None)
+                        stream = self._streams.get(msg.streamId)
                     if msg.controlFlags & STREAM_OPEN_BIT == 0:
                         if not stream:
                             logger.warning("no stream for %s", msg.streamId)
-                            raise IgnoreMessageException(
-                                "no stream for message, ignoring"
-                            )
+                            continue
 
                         if (
                             msg.controlFlags & STREAM_CLOSED_BIT != 0
@@ -160,6 +167,8 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                                 raise InvalidMessageException(e) from e
                     else:
                         _stream = await self._open_stream_and_call_handler(msg, tg)
+                        if isinstance(_stream, IgnoreMessage):
+                            continue
                         if not stream:
                             async with self._stream_lock:
                                 self._streams[msg.streamId] = _stream
@@ -170,9 +179,6 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                             stream.close()
                         async with self._stream_lock:
                             del self._streams[msg.streamId]
-                except IgnoreMessageException:
-                    logger.debug("Ignoring transport message", exc_info=True)
-                    continue
                 except OutOfOrderMessageException:
                     logger.exception("Out of order message, closing connection")
                     await ws_wrapper.close()
@@ -188,17 +194,22 @@ async def _open_stream_and_call_handler(
         self,
         msg: TransportMessage,
         tg: asyncio.TaskGroup | None,
-    ) -> Channel:
+    ) -> Channel | IgnoreMessage:
         if not msg.serviceName or not msg.procedureName:
-            raise IgnoreMessageException(
-                f"Service name or procedure name is missing in the message {msg}"
+            logger.warning(
+                "Service name or procedure name is missing in the message %r",
+                msg,
             )
+            return IgnoreMessage()
         key = (msg.serviceName, msg.procedureName)
         handler = self._handlers.get(key, None)
         if not handler:
-            raise IgnoreMessageException(
-                f"No handler for {key} handlers : {self._handlers.keys()}"
+            logger.warning(
+                "No handler for %r handlers: %r",
+                key,
+                self._handlers.keys(),
             )
+            return IgnoreMessage()
         method_type, handler_func = handler
         is_streaming_output = method_type in (
             "subscription-stream",  # subscription

src/replit_river/server_transport.py

@@ -24,7 +24,6 @@
     TransportMessage,
 )
 from replit_river.seq_manager import (
-    IgnoreMessageException,
     InvalidMessageException,
     SessionStateMismatchException,
 )
@@ -74,12 +73,12 @@ async def handshake_to_get_session(
         async for message in websocket:
             try:
                 msg = parse_transport_msg(message)
+                if isinstance(msg, str):
+                    continue
                 (
                     handshake_request,
                     handshake_response,
                 ) = await self._establish_handshake(msg, websocket)
-            except IgnoreMessageException:
-                continue
             except InvalidMessageException as e:
                 error_msg = "Got invalid transport message, closing connection"
                 raise InvalidMessageException(error_msg) from e

src/replit_river/v2/session.py

@@ -60,7 +60,6 @@
     TransportMessageTracingSetter,
 )
 from replit_river.seq_manager import (
-    IgnoreMessageException,
     InvalidMessageException,
     OutOfOrderMessageException,
 )
@@ -1059,13 +1058,14 @@ async def websocket_closed_callback() -> None:
 
                 try:
                     response_msg = parse_transport_msg(data)
+                    if isinstance(response_msg, str):
+                        logger.debug(
+                            "_do_ensure_connected: Ignoring transport message",
+                            exc_info=True,
+                        )
+                        continue
+
                     break
-                except IgnoreMessageException:
-                    logger.debug(
-                        "_do_ensure_connected: Ignoring transport message",
-                        exc_info=True,
-                    )
-                    continue
                 except InvalidMessageException as e:
                     raise RiverException(
                         ERROR_HANDSHAKE,
@@ -1217,6 +1217,9 @@ async def _serve(
                         transport_id,
                         msg,
                     )
+                    if isinstance(msg, str):
+                        logger.debug("Ignoring transport message", exc_info=True)
+                        continue
 
                     if msg.controlFlags & STREAM_OPEN_BIT != 0:
                         raise InvalidMessageException(