Remove irrelevant Lock

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -132,7 +132,7 @@ async def _handle_messages_from_ws(self) -> None:
                     logger.debug(f"{self._transport_id} got a message %r", msg)
 
                     # Update bookkeeping
-                    match await self._seq_manager.check_seq_and_update(msg):
+                    match self._seq_manager.check_seq_and_update(msg):
                         case IgnoreMessage():
                             continue
                         case None:

src/replit_river/seq_manager.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 from dataclasses import dataclass
 
@@ -7,12 +6,6 @@
 logger = logging.getLogger(__name__)
 
 
-class IgnoreMessageException(Exception):
-    """Exception to ignore a transport message, but good to continue."""
-
-    pass
-
-
 class InvalidMessageException(Exception):
     """Error processing a transport message, should raise a exception."""
 
@@ -46,53 +39,40 @@ class SeqManager:
     def __init__(
         self,
     ) -> None:
-        self._seq_lock = asyncio.Lock()
         self.seq = 0
-        self._ack_lock = asyncio.Lock()
         self.ack = 0
         self.receiver_ack = 0
 
-    async def get_seq_and_increment(self) -> int:
+    def get_seq_and_increment(self) -> int:
         """Get the current sequence number and increment it.
         This removes one lock acquire than get_seq and increment_seq separately.
         """
-        async with self._seq_lock:
-            current_value = self.seq
-            self.seq += 1
-            return current_value
-
-    async def increment_seq(self) -> int:
-        async with self._seq_lock:
-            self.seq += 1
-            return self.seq
-
-    async def get_seq(self) -> int:
-        async with self._seq_lock:
-            return self.seq
-
-    async def get_ack(self) -> int:
-        async with self._ack_lock:
-            return self.ack
-
-    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:
-                    return IgnoreMessage()
-                else:
-                    logger.warn(
-                        f"Out of order message received got {msg.seq} expected "
-                        f"{self.ack}"
-                    )
-
-                    raise OutOfOrderMessageException(
-                        f"Out of order message received got {msg.seq} expected "
-                        f"{self.ack}"
-                    )
-            self.receiver_ack = msg.ack
-        await self._set_ack(msg.seq + 1)
-
-    async def _set_ack(self, new_ack: int) -> int:
-        async with self._ack_lock:
-            self.ack = new_ack
-            return self.ack
+        current_value = self.seq
+        self.seq += 1
+        return current_value
+
+    def increment_seq(self) -> int:
+        self.seq += 1
+        return self.seq
+
+    def get_seq(self) -> int:
+        return self.seq
+
+    def get_ack(self) -> int:
+        return self.ack
+
+    def check_seq_and_update(self, msg: TransportMessage) -> IgnoreMessage | None:
+        if msg.seq != self.ack:
+            if msg.seq < self.ack:
+                return IgnoreMessage()
+            else:
+                logger.warn(
+                    f"Out of order message received got {msg.seq} expected {self.ack}"
+                )
+
+                raise OutOfOrderMessageException(
+                    f"Out of order message received got {msg.seq} expected {self.ack}"
+                )
+        self.receiver_ack = msg.ack
+        self.ack = msg.seq + 1
+        return None

src/replit_river/server_session.py

@@ -170,8 +170,7 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                         if isinstance(_stream, IgnoreMessage):
                             continue
                         if not stream:
-                            async with self._stream_lock:
-                                self._streams[msg.streamId] = _stream
+                            self._streams[msg.streamId] = _stream
                         stream = _stream
 
                     if msg.controlFlags & STREAM_CLOSED_BIT != 0:
@@ -193,7 +192,7 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
     async def _open_stream_and_call_handler(
         self,
         msg: TransportMessage,
-        tg: asyncio.TaskGroup | None,
+        tg: asyncio.TaskGroup,
     ) -> Channel | IgnoreMessage:
         if not msg.serviceName or not msg.procedureName:
             logger.warning(

src/replit_river/session.py

@@ -225,18 +225,18 @@ async def _send_transport_message(
 
     async def get_next_expected_seq(self) -> int:
         """Get the next expected sequence number from the server."""
-        return await self._seq_manager.get_ack()
+        return self._seq_manager.get_ack()
 
     async def get_next_sent_seq(self) -> int:
         """Get the next sequence number that the client will send."""
         nextMessage = await self._buffer.peek()
         if nextMessage:
             return nextMessage.seq
-        return await self._seq_manager.get_seq()
+        return self._seq_manager.get_seq()
 
     async def get_next_expected_ack(self) -> int:
         """Get the next expected ack that the client expects."""
-        return await self._seq_manager.get_seq()
+        return self._seq_manager.get_seq()
 
     async def send_message(
         self,
@@ -256,8 +256,8 @@ async def send_message(
             id=nanoid.generate(),
             from_=self._transport_id,  # type: ignore
             to=self._to_id,
-            seq=await self._seq_manager.get_seq_and_increment(),
-            ack=await self._seq_manager.get_ack(),
+            seq=self._seq_manager.get_seq_and_increment(),
+            ack=self._seq_manager.get_ack(),
             controlFlags=control_flags,
             payload=payload,
             serviceName=service_name,